feat(rocke/conv): add transpose load for bwd weight and bwd data

dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py

@@ -455,6 +455,7 @@ class ConvGroupedSpec:
     dtype: str
     arch: str
     split_k: int = 1  # wgrad only
+    lds_k_outer: bool = False  # wgrad only
     name: str = "rocke_conv_grouped"
 
     def kernel_name(self) -> str:
@@ -532,6 +533,7 @@ def to_wgrad_spec(self, problem: "ConvProblem") -> "WgradConvSpec":
         return WgradConvSpec(
             problem=problem,
             name=self.name,
+            lds_k_outer=self.lds_k_outer,
             data=ConvDataSpec(
                 dtype_a=self.dtype,
                 dtype_b=self.dtype,
@@ -1038,6 +1040,37 @@ def select(req: OperatorRequest) -> ConvGroupedSpec:
 # ---------------------------------------------------------------------------
 
 
+def _wgrad_lds_k_outer(req: "ConvGroupedRequest", warp_tile_mn: int) -> bool:
+    """Whether the gfx950 wgrad candidate should use the K-outer LDS layout.
+
+    wgrad's stride-1 global axis is the GEMM *free* axis, so an M-outer LDS tile
+    forces a transpose on store: one ``ds_write_b16`` per element, all of them
+    bank-conflicting because the row stride is a multiple of the LDS bank period.
+    The K-outer layout stores the tile in global order (one wide store) and lets
+    ``ds_read_b64_tr_b16`` transpose on the read side instead.
+
+    K-outer lowers the LDS instruction count and removes the row-stride bank
+    conflicts the M-outer transpose-on-store incurs. See the ``lds_k_outer``
+    field docs on ``WgradConvSpec``.
+
+    Gated to exactly what the transpose-read lane mapping is validated for:
+    gfx950 (the instruction does not exist on gfx942), wave64 MFMA, 16-bit A/B
+    operands, and a 16- or 32-wide atom edge.
+
+    Delegates so dispatch and the sweep driver cannot drift: this used to be a
+    second copy of the gate that compared the module constant against a tuple
+    (constant-true regardless of the request) and never checked wave_size.
+    """
+    return WgradConvSpec.default_lds_k_outer(
+        arch=req.arch,
+        dtype_a=req.dtype.lower(),
+        dtype_b=req.dtype.lower(),
+        warp_tile_m=warp_tile_mn,
+        warp_tile_n=warp_tile_mn,
+        wave_size=ArchTarget.from_gfx(req.arch).wave_size,
+    )
+
+
 def _make_gfx950_wgrad_candidate() -> KernelCandidate:
     """Backward-weight conv for gfx950: 64×64×64, 2×2, 32×32×16 MFMA.
 
@@ -1066,6 +1099,7 @@ def _build_instance_spec(req: ConvGroupedRequest) -> WgradConvSpec:
         return WgradConvSpec(
             problem=_problem(req),
             name=name,
+            lds_k_outer=_wgrad_lds_k_outer(req, wtmn),
             data=_data_spec(req),
             tile_m=tm,
             tile_n=tn,
@@ -1116,6 +1150,7 @@ def select(req: OperatorRequest) -> ConvGroupedSpec:
             warp_tile_k=wtk,
             pipeline=_PIPELINE,
             epilogue=_ep,
+            lds_k_outer=_wgrad_lds_k_outer(req, wtmn),
             dtype=req.dtype.lower(),
             arch=req.arch,
             split_k=_sk,

dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/loads.cpp

@@ -509,7 +509,7 @@ void rocke_coalesced_tile_loader_store_lds(rocke_ir_builder_t* b,
  * ========================================================================== */
 
 rocke_status_t rocke_async_tile_loader_choose_dwords(
-    int tile_rows, int tile_cols, int block_size, int max_dwords, int* out)
+    int tile_rows, int tile_cols, int block_size, int max_dwords, int contig_cols, int* out)
 {
     /* Python:
      *   if max_dwords > 4: max_dwords = 4
@@ -539,6 +539,13 @@ rocke_status_t rocke_async_tile_loader_choose_dwords(
             continue;
         }
         halves = d * 2;
+        /* A chunk must lie entirely inside one contiguous run of columns
+         * (Python: contig_cols is not None and (elems > contig_cols or
+         * contig_cols % elems != 0) -> skip). contig_cols == 0 means None. */
+        if(contig_cols > 0 && (halves > contig_cols || contig_cols % halves != 0))
+        {
+            continue;
+        }
         if(tile_cols % halves != 0)
         {
             continue;
@@ -548,6 +555,14 @@ rocke_status_t rocke_async_tile_loader_choose_dwords(
         {
             continue;
         }
+        /* Every pass fires block_size chunks; chunks past chunks_total still
+         * issue (OOB sentinel, zero-filled), so a partial last pass writes past
+         * this tile's allocation into the next smem_alloc. Require whole
+         * passes -- a narrower width is tried next. Mirrors the Python. */
+        if(chunks % block_size != 0)
+        {
+            continue;
+        }
         if(out != NULL)
         {
             *out = d;
@@ -562,6 +577,7 @@ rocke_status_t rocke_async_tile_loader_from_tile(int tile_rows,
                                                  int block_size,
                                                  int wave_size,
                                                  int max_dwords,
+                                                 int contig_cols,
                                                  rocke_async_tile_loader_t* out)
 {
     int d;
@@ -578,7 +594,8 @@ rocke_status_t rocke_async_tile_loader_from_tile(int tile_rows,
      *   return cls(... dwords=d, chunks_total=chunks,
      *              chunks_per_pass=block_size, passes=passes)
      */
-    st = rocke_async_tile_loader_choose_dwords(tile_rows, tile_cols, block_size, max_dwords, &d);
+    st = rocke_async_tile_loader_choose_dwords(
+        tile_rows, tile_cols, block_size, max_dwords, contig_cols, &d);
     if(st != ROCKE_OK)
     {
         return st;
@@ -596,6 +613,7 @@ rocke_status_t rocke_async_tile_loader_from_tile(int tile_rows,
         out->chunks_total = chunks;
         out->chunks_per_pass = block_size;
         out->passes = passes;
+        out->contig_cols = contig_cols;
     }
     return ROCKE_OK;
 }
@@ -612,10 +630,21 @@ int rocke_async_tile_loader_bytes_per_chunk(const rocke_async_tile_loader_t* sel
     return (self == NULL) ? 0 : self->dwords * 4;
 }
 
-int rocke_async_tile_loader_cols_per_chunk(const rocke_async_tile_loader_t* self)
+int rocke_async_tile_loader_chunks_per_row(const rocke_async_tile_loader_t* self)
 {
-    /* Python: return self.halves_per_chunk */
-    return rocke_async_tile_loader_halves_per_chunk(self);
+    /* Python: return self.tile_cols // self.elems_per_chunk
+     *
+     * Chunks needed to cover one tile row.  NOT elems_per_chunk: the two
+     * coincide only when tile_cols == elems_per_chunk**2 (tile_cols=64 for a
+     * 2-byte dtype at dwords=4), which is why 64 used to be the only tile
+     * width the async path loaded correctly.  choose_dwords() guarantees
+     * tile_cols % elems_per_chunk == 0, so this division is exact. */
+    if(self == NULL)
+    {
+        return 0;
+    }
+    const int elems = rocke_async_tile_loader_halves_per_chunk(self);
+    return (elems == 0) ? 0 : self->tile_cols / elems;
 }
 
 int rocke_async_tile_loader_wave_bytes(const rocke_async_tile_loader_t* self)
@@ -697,7 +726,7 @@ void rocke_async_tile_loader_slot_issue(rocke_ir_builder_t* b,
     rocke_value_t* c_half_bytes;
     rocke_value_t* c_oob;
     rocke_value_t* c0;
-    rocke_value_t* c_cols_per_chunk;
+    rocke_value_t* c_chunks_per_row;
     int p;
 
     if(b != NULL && b->status != ROCKE_OK)
@@ -715,13 +744,13 @@ void rocke_async_tile_loader_slot_issue(rocke_ir_builder_t* b,
      *   c_half_bytes     = b.const_i32(2)
      *   c_oob            = b.const_i32(oob_sentinel)
      *   c0               = b.const_i32(0)
-     *   c_cols_per_chunk = b.const_i32(L.cols_per_chunk)
+     *   c_chunks_per_row = b.const_i32(L.chunks_per_row)
      */
     L = &self->loader;
     c_half_bytes = rocke_b_const_i32(b, 2);
     c_oob = rocke_b_const_i32(b, oob_sentinel);
     c0 = rocke_b_const_i32(b, 0);
-    c_cols_per_chunk = rocke_b_const_i32(b, rocke_async_tile_loader_cols_per_chunk(L));
+    c_chunks_per_row = rocke_b_const_i32(b, rocke_async_tile_loader_chunks_per_row(L));
 
     for(p = 0; p < L->passes; ++p)
     {
@@ -759,13 +788,13 @@ void rocke_async_tile_loader_slot_issue(rocke_ir_builder_t* b,
 
         /* Python:
          *   chunk_idx = b.add(tid, b.const_i32(p * L.block_size))
-         *   row   = b.div(chunk_idx, c_cols_per_chunk)
-         *   col_v = b.mod(chunk_idx, c_cols_per_chunk)
+         *   row   = b.div(chunk_idx, c_chunks_per_row)
+         *   col_v = b.mod(chunk_idx, c_chunks_per_row)
          *   col   = b.mul(col_v, b.const_i32(L.halves_per_chunk))
          */
         chunk_idx = rocke_b_add(b, tid, rocke_b_const_i32(b, p * L->block_size));
-        row = rocke_b_div(b, chunk_idx, c_cols_per_chunk);
-        col_v = rocke_b_mod(b, chunk_idx, c_cols_per_chunk);
+        row = rocke_b_div(b, chunk_idx, c_chunks_per_row);
+        col_v = rocke_b_mod(b, chunk_idx, c_chunks_per_row);
         col = rocke_b_mul(
             b, col_v, rocke_b_const_i32(b, rocke_async_tile_loader_halves_per_chunk(L)));
 

dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.loads.h

@@ -17,7 +17,7 @@
  *     .from_tile     (classmethod)    rocke_async_tile_loader_from_tile()
  *     .halves_per_chunk (property)    rocke_async_tile_loader_halves_per_chunk()
  *     .bytes_per_chunk  (property)    rocke_async_tile_loader_bytes_per_chunk()
- *     .cols_per_chunk   (property)    rocke_async_tile_loader_cols_per_chunk()
+ *     .chunks_per_row   (property)    rocke_async_tile_loader_chunks_per_row()
  *     .wave_bytes       (property)    rocke_async_tile_loader_wave_bytes()
  *     .pass_bytes       (property)    rocke_async_tile_loader_pass_bytes()
  *     .bind          (method)         rocke_async_tile_loader_bind()
@@ -238,6 +238,12 @@ typedef struct rocke_async_tile_loader
     int chunks_total; /* tile_rows * tile_cols / (dwords * 2) */
     int chunks_per_pass; /* = block_size */
     int passes; /* ceil(chunks_total / block_size) */
+    /* Python `contig_cols: Optional[int] = None`, spelled here as 0 => None.
+     * Longest run of tile columns that is contiguous in global memory. When the
+     * col axis is composite -- e.g. the conv reduction index (y, x, c), where
+     * only c is stride-1 -- a chunk crossing one of those boundaries would fetch
+     * the wrong elements silently. */
+    int contig_cols;
 } rocke_async_tile_loader_t;
 
 /* Bound AsyncTileLoader (rocke.helpers.loads.AsyncTileLoaderSlot). */
@@ -255,7 +261,7 @@ typedef struct rocke_async_tile_loader_slot
  * ValueError path. max_dwords default 4 (values > 4 are clamped to 4, matching
  * Python). */
 rocke_status_t rocke_async_tile_loader_choose_dwords(
-    int tile_rows, int tile_cols, int block_size, int max_dwords, int* out);
+    int tile_rows, int tile_cols, int block_size, int max_dwords, int contig_cols, int* out);
 
 /* AsyncTileLoader.from_tile classmethod.
  *
@@ -267,12 +273,13 @@ rocke_status_t rocke_async_tile_loader_from_tile(int tile_rows,
                                                  int block_size,
                                                  int wave_size,
                                                  int max_dwords,
+                                                 int contig_cols, /* 0 => None */
                                                  rocke_async_tile_loader_t* out);
 
 /* AsyncTileLoader properties (pure int arithmetic). */
 int rocke_async_tile_loader_halves_per_chunk(const rocke_async_tile_loader_t* self);
 int rocke_async_tile_loader_bytes_per_chunk(const rocke_async_tile_loader_t* self);
-int rocke_async_tile_loader_cols_per_chunk(const rocke_async_tile_loader_t* self);
+int rocke_async_tile_loader_chunks_per_row(const rocke_async_tile_loader_t* self);
 int rocke_async_tile_loader_wave_bytes(const rocke_async_tile_loader_t* self);
 int rocke_async_tile_loader_pass_bytes(const rocke_async_tile_loader_t* self);
 

dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_dgrad.h

@@ -98,6 +98,11 @@ typedef struct rocke_dgrad_conv_spec
     bool async_dma; /* default false */
     bool unroll_k; /* default false */
 
+    /* Mirrors DgradConvSpec.lds_k_outer. Strictly additive: default false, and
+     * the B tile only. A (dY, NHWK) already has a stride-1 reduction axis, so
+     * only B pays a transpose-on-store. */
+    bool lds_k_outer; /* default false */
+
     bool has_lds_k_pad; /* false => Python None */
     int lds_k_pad;
     void* lds_layout; /* NULL => Python None */
@@ -129,6 +134,11 @@ typedef struct rocke_dgrad_conv_spec
 } rocke_dgrad_conv_spec_t;
 
 /* Default-constructed spec (every field == Python dataclass default). */
+/* Row stride pad, in elements, for the K-outer B tile. Mirrors _KOUTER_PAD in
+ * conv_implicit_gemm_dgrad.py: the stride must not be a multiple of the
+ * 32-dword LDS bank period or the transpose read degenerates. */
+#define ROCKE_DGRAD_KOUTER_PAD 8
+
 rocke_dgrad_conv_spec_t rocke_dgrad_conv_spec_default(void);
 
 /* ---- DgradConvSpec @property analogues (pure int arithmetic) ---- */

dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_internal.h

@@ -284,6 +284,30 @@ typedef struct rocke_conv_build_ctx
      * pipeline="wavelet": a_wavelet_loader/b_wavelet_loader are
      *   CoalescedTileLoader sized to num_load_waves * wave_size threads. */
     bool async_dma; /* spec.async_dma (cached)         */
+    /* spec.lds_k_outer (wgrad only). Zeroed by the ctx memset, so the forward
+     * and dgrad drivers keep their exact previous behaviour. */
+    bool lds_k_outer;
+    /* Optional override for the A-operand descriptor on the ASYNC load path.
+     * The sync path already has a_load_override; the async slot had no hook and
+     * always called rocke_conv_a_descriptor, which is the forward descriptor.
+     * NULL keeps that behaviour, so fwd/dgrad are unaffected. */
+    rocke_loads_descriptor_fn a_descriptor_fn;
+    /* Base added to every k offset the k-loop drivers hand to the load phase.
+     * The forward conv passes a bare const_i32(it*block_k); wgrad's Python
+     * emits b.add(k_lo, const_i32(...)) because its slice starts at the
+     * split-K lower bound. NULL keeps the bare-const form, so fwd/dgrad emit
+     * exactly what they emitted before. */
+    rocke_value_t* kloop_k_lo;
+    /* Iteration count for the unrolled / async k-loop drivers. 0 = derive from
+     * the problem's full K_gemm, which is what the forward conv wants. wgrad
+     * must override it: its stub problem carries the UNSLICED wg_K, while the
+     * Python emitter sizes the loop from wg_K_padded()/split_k. They agree only
+     * at split_k == 1, so without this the C engine emits a full-range
+     * reduction under split-K atomics. */
+    int kloop_num_iters;
+    /* Atom edge + fragment length used by the K-outer transpose-read feed. */
+    rocke_value_t* tr_lane_mod4;
+    rocke_value_t* tr_grp16;
     rocke_async_tile_loader_t a_loader; /* async A loader (valid iff async)*/
     rocke_async_tile_loader_t b_loader; /* async B loader                  */
     bool have_async_loaders; /* true => a_loader/b_loader valid */
@@ -367,6 +391,21 @@ rocke_value_t* rocke_conv_emit_smem_load(
 /* _emit_frag_smem_load(b, src, mn_in_atom, k_in_atom, atom_mn_base, k_tile_base,
  * frag_len): one frag_len-wide operand fragment from a row-major LDS tile
  * (8-wide chunked + vec_concat for wide WMMA frags). */
+/* K-outer transpose-read fragment feed. Takes lane and the two hoisted lane
+ * constants directly rather than a build ctx, so both the shared compute phase
+ * (wgrad, A and B) and dgrad's own operand fetch (B only) can call it without
+ * duplicating the lane mapping. See conv_implicit_gemm_conv_compute_phase.cpp. */
+rocke_value_t* rocke_conv_tr_frag(rocke_ir_builder_t* b,
+                                  rocke_value_t* lane,
+                                  rocke_value_t* tr_lane_mod4,
+                                  rocke_value_t* tr_grp16,
+                                  rocke_value_t* smem,
+                                  rocke_value_t* mn_base,
+                                  rocke_value_t* k_base,
+                                  int mn_atom,
+                                  int n,
+                                  const rocke_type_t* dtype);
+
 rocke_value_t* rocke_conv_emit_frag_smem_load(rocke_ir_builder_t* b,
                                               rocke_value_t* src,
                                               rocke_value_t* mn_in_atom,

dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h

@@ -47,6 +47,16 @@
 #define ROCKE_INSTANCE_CONV_IMPLICIT_GEMM_WGRAD_H
 
 #include <stdbool.h>
+
+/* Column pad (in elements) for the K-outer wgrad LDS tile. Keeps the row stride
+ * off a multiple of the LDS bank period while staying 16-byte aligned for the
+ * wide store. Mirrors _KOUTER_PAD in conv_implicit_gemm_wgrad.py. */
+#define ROCKE_WGRAD_KOUTER_PAD 8
+
+/* The K-outer LDS tile is fed by ds_read_tr16_b64, a CDNA4 transpose read.
+ * Emitting it for an older target produces IR the assembler will reject.
+ * Mirrors _LDS_K_OUTER_ARCH in conv_implicit_gemm_wgrad.py. */
+#define ROCKE_WGRAD_LDS_K_OUTER_ARCH "gfx950"
 #include <stddef.h>
 
 #include "rocke/helper_rocke.instances.common.conv_implicit_gemm.h" /* rocke_conv_problem_t */
@@ -114,6 +124,11 @@ typedef struct rocke_implicit_gemm_conv_wgrad_spec
     const char* epilogue; /* default "default" */
     bool async_dma; /* default false */
     bool unroll_k; /* default false */
+    /* Store the A/B tiles K-outer (LDS[k][m] / LDS[k][n]) and feed the MFMA with
+     * gfx950 ds_read_b64_tr_b16 transpose reads instead of transposing on store.
+     * Mirrors WgradConvSpec.lds_k_outer. Default false: strictly additive, so
+     * every existing config emits byte-identical IR. */
+    bool lds_k_outer; /* default false */
 
     bool has_lds_k_pad; /* false => Python None */
     int lds_k_pad;
@@ -168,6 +183,11 @@ int rocke_wgrad_conv_spec_wg_N(const rocke_implicit_gemm_conv_wgrad_spec_t* s);
 /* spec.wg_K: output spatial positions (N * Ho * Wo [* Do]). */
 int rocke_wgrad_conv_spec_wg_K(const rocke_implicit_gemm_conv_wgrad_spec_t* s);
 
+/* Max K iterations the Python-unrolled loops (pipeline="basic" and async_dma)
+ * may unroll to. Build-practicality bound (code size / compile time), not a
+ * hardware limit. Mirrors _MAX_UNROLLED_K_ITERS in conv_implicit_gemm_wgrad.py. */
+#define ROCKE_MAX_UNROLLED_K_ITERS 128
+
 /* spec.wg_K_padded(): wg_K rounded up to tile_k * split_k. */
 int rocke_wgrad_conv_spec_wg_K_padded(const rocke_implicit_gemm_conv_wgrad_spec_t* s);