refactor(rocke): deduplicate conv/attention emitter clone groups

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/_conv_implicit_gemm_common.py

@@ -40,6 +40,7 @@
     Value,
 )
 from rocke.helpers.spec import choose_load_vec
+from rocke.helpers.tensor_view import make_buffer_resource
 from rocke.helpers.transforms import (
     TensorDescriptor,
     DynamicTensorDescriptor,
@@ -1309,6 +1310,38 @@ def make_d_descriptor_dynamic(
     ).transform(*transforms)
 
 
+def _implicit_gemm_output_config(op, spec):
+    """Shared output layout for the implicit-GEMM epilogue copies."""
+    c_map = op.c_layout()
+    _fp32_out = spec.data.dtype_d == "fp32"
+    _bf16_out = spec.data.dtype_d == "bf16"
+    _elem_bytes = 4 if _fp32_out else 2
+    flat = 0
+    return c_map, _fp32_out, _bf16_out, _elem_bytes, flat
+
+
+def _c_lane_rows_cols(b, atom, lane, kc_m1, kc_nlane, c_per_lane):
+    """Decode the per-lane (row, col) pairs for a warp tile's C accumulator slots."""
+    from rocke.helpers.atoms import make_c_warp_dstr_encoding
+    from rocke.helpers.distribution import make_static_tile_distribution
+
+    c_dist = make_static_tile_distribution(make_c_warp_dstr_encoding(atom))
+
+    c_nlane = b.const_i32(kc_nlane)
+    n_in_atom = b.mod(lane, c_nlane)
+    m_blk = b.div(lane, c_nlane)
+    p_lane = [m_blk, n_in_atom]
+
+    rows: List[Value] = []
+    cols: List[Value] = []
+    for i in range(c_per_lane):
+        ys = [b.const_i32(i // kc_m1), b.const_i32(i % kc_m1)]
+        x_row, x_col = c_dist.calculate_x(b, ys=ys, ps=[p_lane])
+        rows.append(x_row)
+        cols.append(x_col)
+    return rows, cols
+
+
 def emit_direct_epilogue_wmma(
     b,
     spec,
@@ -1335,12 +1368,9 @@ def emit_direct_epilogue_wmma(
     mfmas_n = spec.mfmas_per_warp_n
     warp_m_off = b.mul(warp_m_idx, b.const_i32(mfmas_m * spec.warp_tile_m))
     warp_n_off = b.mul(warp_n_idx, b.const_i32(mfmas_n * spec.warp_tile_n))
-    c_map = op.c_layout()
-    _fp32_out = spec.data.dtype_d == "fp32"
-    _bf16_out = spec.data.dtype_d == "bf16"
-    _elem_bytes = 4 if _fp32_out else 2
-
-    flat = 0
+    c_map, _fp32_out, _bf16_out, _elem_bytes, flat = _implicit_gemm_output_config(
+        op, spec
+    )
     for mi in range(mfmas_m):
         for ni in range(mfmas_n):
             acc = accs[flat]
@@ -1373,3 +1403,16 @@ def emit_direct_epilogue_wmma(
                     )
                 else:
                     b.buffer_store_f16(d_rsrc, safe_off, c0, b.trunc_f32_to_f16(v_f32))
+
+
+def _conv_io_buffer_rsrcs(b, a, a_bytes, b_ptr, b_bytes, d, d_bytes):
+    """Buffer-resource views over the conv IO tensors.  Returns
+    ``(a_rsrc, b_rsrc, d_rsrc)`` — the raw :class:`BufferResource` handles.
+    """
+    a_buf_rsrc = make_buffer_resource(b, a, num_bytes=a_bytes)
+    b_buf_rsrc = make_buffer_resource(b, b_ptr, num_bytes=b_bytes)
+    d_buf_rsrc = make_buffer_resource(b, d, num_bytes=d_bytes)
+    a_rsrc = a_buf_rsrc.rsrc
+    b_rsrc = b_buf_rsrc.rsrc
+    d_rsrc = d_buf_rsrc.rsrc
+    return a_rsrc, b_rsrc, d_rsrc

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/_fmha_common.py

@@ -508,3 +508,26 @@ def _row_base(self, tensor_name: str, tok: Value, hd: Value) -> Value:
             self.b.mul(tok, tok_stride),
             self.b.mul(hd, head_stride),
         )
+
+
+def _declare_fmha_fwd_params(kb):
+    """Declare the shared FMHA-forward kernel ABI: tensors, scalars, strides."""
+    kb.add_tensor("Q", readonly=True)
+    kb.add_tensor("K", readonly=True)
+    kb.add_tensor("V", readonly=True)
+    kb.add_tensor("O", readonly=False, writeonly=True)
+    kb.add_scalar("scale_log2", "f32")
+    kb.add_scalar("seqlen_q", "i32")
+    kb.add_scalar("seqlen_k", "i32")
+    kb.add_strides("q", "k", "v", "o")
+
+
+def _sparse_builder_setup(kb):
+    """Decode the grid and read the attention tensor handles for sparse builders."""
+    kb.decode_grid()
+    b = kb.builder
+    Q = kb.tensor("Q")
+    K = kb.tensor("K")
+    V = kb.tensor("V")
+    O = kb.tensor("O")  # noqa: E741 - standard attention notation
+    return b, Q, K, V, O

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/_fmha_warp_body.py

@@ -81,6 +81,17 @@
 WARP_SIZE = 64
 
 
+def _online_softmax_update(b, partial, scale_log2, k_idx, m, lse, apply_mask_and_score):
+    """Warp-reduce the partial QK dot, mask/score it, then do the online-softmax update."""
+    dot = warp_xor_reduce_sum(b, partial, stages=6)
+    score_log2 = apply_mask_and_score(b.fmul(dot, scale_log2), k_idx)
+    m_new = b.fmax(m, score_log2)
+    alpha = b.exp2(b.fsub(m, m_new))
+    p = b.exp2(b.fsub(score_log2, m_new))
+    lse_new = b.fadd(b.fmul(lse, alpha), p)
+    return m_new, alpha, p, lse_new
+
+
 def fmha_warp_fwd_inner_body(
     b: IRBuilder,
     *,
@@ -272,13 +283,9 @@ def _apply_mask_and_score(score_log2: Value, k_idx: Value) -> Value:
             for slot in range(ept):
                 partial = b.fadd(partial, b.fmul(q_f32_list[slot], k_f32_list[slot]))
 
-            dot = warp_xor_reduce_sum(b, partial, stages=6)
-            score_log2 = _apply_mask_and_score(b.fmul(dot, scale_log2), k_idx)
-
-            m_new = b.fmax(m, score_log2)
-            alpha = b.exp2(b.fsub(m, m_new))
-            p = b.exp2(b.fsub(score_log2, m_new))
-            lse_new = b.fadd(b.fmul(lse, alpha), p)
+            m_new, alpha, p, lse_new = _online_softmax_update(
+                b, partial, scale_log2, k_idx, m, lse, _apply_mask_and_score
+            )
 
             new_yields = [m_new, lse_new]
             for slot in range(ept):
@@ -328,13 +335,9 @@ def _apply_mask_and_score(score_log2: Value, k_idx: Value) -> Value:
             partial = b.fmul(q_scalar, kd)
 
             # Warp-wide butterfly reduce: 64-lane wave64 sum = 6 stages.
-            dot = warp_xor_reduce_sum(b, partial, stages=6)
-            score_log2 = _apply_mask_and_score(b.fmul(dot, scale_log2), k_idx)
-
-            m_new = b.fmax(m, score_log2)
-            alpha = b.exp2(b.fsub(m, m_new))
-            p = b.exp2(b.fsub(score_log2, m_new))
-            lse_new = b.fadd(b.fmul(lse, alpha), p)
+            m_new, alpha, p, lse_new = _online_softmax_update(
+                b, partial, scale_log2, k_idx, m, lse, _apply_mask_and_score
+            )
             new_acc0 = b.fadd(b.fmul(acc0, alpha), b.fmul(p, vd))
             b.scf_yield(m_new, lse_new, new_acc0)
 
@@ -389,13 +392,9 @@ def _apply_mask_and_score(score_log2: Value, k_idx: Value) -> Value:
         partial = b.vector_sum(b.vector_mul(q_vec_f32, k_vec_f32))
 
         # Warp-wide butterfly reduce: 64-lane wave64 sum = 6 stages.
-        dot = warp_xor_reduce_sum(b, partial, stages=6)
-        score_log2 = _apply_mask_and_score(b.fmul(dot, scale_log2), k_idx)
-
-        m_new = b.fmax(m, score_log2)
-        alpha = b.exp2(b.fsub(m, m_new))
-        p = b.exp2(b.fsub(score_log2, m_new))
-        lse_new = b.fadd(b.fmul(lse, alpha), p)
+        m_new, alpha, p, lse_new = _online_softmax_update(
+            b, partial, scale_log2, k_idx, m, lse, _apply_mask_and_score
+        )
 
         # PV accumulator update: acc = acc * alpha + p * v.
         # The splats reduce to a single VGPR broadcast each on AMDGPU