Skip to content

EmitCUDA: migrate wmma store_matrix_sync conversion to __choreo_v* virtual-index convention #4

Description

@Garfee023

Context

tools/coir/lib/CodeGen/GPU/EmitCUDA.cpp (the CUDA emitter) now generates
named virtual-index preludes following the project convention:

  • __choreo_vtid_x - thread/lane index (threadIdx.x, or % 32 / % 128
    under GROUP / GROUPx4)
  • __choreo_vgid_x - warp-group index (threadIdx.x / 32)
  • __choreo_vg4id_x - quad-group index (threadIdx.x / 128)

emitParallel was recently migrated to this convention, but one remaining
codegen path still hardcodes the warp/lane decomposition inline.

Problem

In the wmma::store_matrix_sync MMA-store conversion path, two virtual
indices are still written as raw threadIdx.x arithmetic:

  1. Warp index (EmitCUDA.cpp:2512):

    float* __mma_cvt_<idx>_local = __mma_cvt_<idx> + (threadIdx.x / 32) * <tileElems>;

    threadIdx.x / 32 is the warp index and should use __choreo_vgid_x.

  2. Lane index (EmitCUDA.cpp:2520):

    for (int _c = threadIdx.x % 32; _c < <tileN>; _c += 32)

    threadIdx.x % 32 is the lane index and should use __choreo_vtid_x.

Why it matters

  • Consistency: the warp/lane decomposition is defined once in
    emitParallel (mirroring the native backend); this path duplicates it
    with magic numbers.
  • Correctness under GROUPx4: GROUPx4 uses 128-lane warp groups, so the
    hardcoded / 32 and % 32 are only correct for plain GROUP (32-lane
    warps). Under GROUPx4 they compute the wrong indices and the _local
    offset can run past the shared buffer.
  • The surrounding code already walks up to the enclosing GROUP/GROUPx4
    parallel to compute numWarps; that same context can select the correct
    named prelude.

Suggested fix

Reuse the enclosing-parallel context (already computed for numWarps) to
emit the matching named prelude, e.g. __choreo_vgid_x (or
__choreo_vg4id_x) for the warp index and __choreo_vtid_x for the lane
index, with the loop stride 32 or 128 matching the group level.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions