From 31b968685cf96c56925fd0f0e4b51f591d2a0569 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Mon, 5 Oct 2026 10:43:32 +0000 Subject: [PATCH 01/20] feat(rocke): add cross-workgroup fence and flag IR ops Add three leaf ops that let one workgroup hand data to another through global memory, needed by the deterministic stream-K fixups: - memref.fence: LLVM fence at workgroup / agent / system scope. - memref.global_flag_store: scoped atomic i32 store of a flag. - memref.global_flag_wait_eq: self-contained spin loop around an acquire atomic load until the flag matches. The protocol is expressed through the AMDGPU memory model rather than raw cache-policy asm, so an agent-scope release writes back and an acquire invalidates the per-XCD L2 on multi-XCD CDNA parts. The spin is a leaf op, so structured control flow stays scf.for / scf.if and no region op or recipe ABI change is needed. Both engines implement the builders, verifier attrs, LLVM lowering and HIP lowering. The new opcodes are appended after the control-flow block so no existing opcode index moves. TestNewTargetIntrinsics and its C++ twin pin the emitted text and cross-check the engines. ISSUE ID : AICK-1943 Co-Authored-By: Claude Opus 5.5 --- .../rocke/platform/cpp/core/ir/core_types.cpp | 14 ++- .../rocke/platform/cpp/core/ir/ir_mem.cpp | 118 ++++++++++++++++++ .../cpp/core/lower_hip/lower_hip_mem.cpp | 86 +++++++++++++ .../platform/cpp/core/lower_llvm/mem.cpp | 82 ++++++++++++ .../rocke/platform/cpp/core/verify.cpp | 9 ++ .../rocke/platform/cpp/include/rocke/ir.h | 27 ++++ .../primitives/intrinsics_and_primitives.md | 26 ++++ .../platform/dsl_docs/reference/api_index.md | 5 + .../dsl_docs/reference/op_vocabulary.md | 3 + .../rocke/platform/python/rocke/core/ir.py | 78 ++++++++++++ .../platform/python/rocke/core/lower_hip.py | 39 ++++++ .../platform/python/rocke/core/lower_llvm.py | 57 +++++++++ .../platform/python/rocke/core/verify.py | 3 + .../tests/core/future_intrinsic_lowering.cpp | 54 ++++++++ .../rocke/platform/tests/test_rocke.py | 74 +++++++++++ 15 files changed, 673 insertions(+), 2 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp index cd34e3b984c0..7c2230c5c50e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp @@ -668,7 +668,12 @@ static const char* const rocke_opcode_names[ROCKE_OP__COUNT] = { "scf.if", "scf.if_else", "scf.yield", - "cf.return"}; + "cf.return", + + /* memref.* cross-workgroup synchronisation */ + "memref.fence", + "memref.global_flag_store", + "memref.global_flag_wait_eq"}; const char* rocke_opcode_name(rocke_opcode_t op) { @@ -914,7 +919,12 @@ static const bool rocke_opcode_pure[ROCKE_OP__COUNT] = { /* scf.if */ false, /* scf.if_else */ false, /* scf.yield */ false, - /* cf.return */ false}; + /* cf.return */ false, + + /* memref.* cross-workgroup synchronisation (effectful) */ + /* memref.fence */ false, + /* memref.global_flag_store */ false, + /* memref.global_flag_wait_eq */ false}; bool rocke_opcode_is_pure(rocke_opcode_t op) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp index 05608fbafc5d..a146d21a7486 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp @@ -429,6 +429,124 @@ static bool rocke_i_ordering_ok(const char* ordering) || !strcmp(ordering, "seq_cst")); } +/* IRBuilder._FENCE_SCOPES */ +static bool rocke_i_fence_scope_ok(const char* scope) +{ + return scope + && (!strcmp(scope, "workgroup") || !strcmp(scope, "agent") || !strcmp(scope, "system")); +} + +void rocke_b_fence(rocke_ir_builder_t* b, const char* scope, const char* ordering) +{ + rocke_attr_map_t a; + if(!rocke_i_live(b)) + return; + if(scope == NULL) + scope = "agent"; + if(ordering == NULL) + ordering = "acq_rel"; + if(!rocke_i_fence_scope_ok(scope)) + { + (void)rocke_i_set_err(b, ROCKE_ERR_VALUE, "unknown fence scope '%s'", scope); + return; + } + if(strcmp(ordering, "acquire") && strcmp(ordering, "release") && strcmp(ordering, "acq_rel") + && strcmp(ordering, "seq_cst")) + { + (void)rocke_i_set_err(b, ROCKE_ERR_VALUE, "unknown fence ordering '%s'", ordering); + return; + } + a = rocke_i_attrs(b); + rocke_attr_set_str(b, &a, "scope", scope); + rocke_attr_set_str(b, &a, "ordering", ordering); + (void)rocke_i_op0(b, ROCKE_OP_MEMREF_FENCE, NULL, 0, &a); +} + +void rocke_b_global_flag_store(rocke_ir_builder_t* b, + rocke_value_t* ptr, + rocke_value_t* idx, + rocke_value_t* value, + const char* scope, + const char* ordering) +{ + rocke_value_t* ops[3]; + rocke_attr_map_t a; + if(!rocke_i_live(b)) + return; + if(!ptr || !idx || !value) + { + (void)rocke_i_set_err(b, ROCKE_ERR_VALUE, "global_flag_store: null operand"); + return; + } + if(scope == NULL) + scope = "agent"; + if(ordering == NULL) + ordering = "release"; + if(!rocke_i_type_is(value->type, "i32")) + { + (void)rocke_i_set_err(b, + ROCKE_ERR_VALUE, + "global_flag_store expects an i32 value, got %s", + value->type->name); + return; + } + if(!rocke_i_fence_scope_ok(scope)) + { + (void)rocke_i_set_err(b, ROCKE_ERR_VALUE, "unknown flag scope '%s'", scope); + return; + } + if(strcmp(ordering, "monotonic") && strcmp(ordering, "release") && strcmp(ordering, "seq_cst")) + { + (void)rocke_i_set_err(b, ROCKE_ERR_VALUE, "unknown flag store ordering '%s'", ordering); + return; + } + ops[0] = ptr; + ops[1] = idx; + ops[2] = value; + a = rocke_i_attrs(b); + rocke_attr_set_str(b, &a, "scope", scope); + rocke_attr_set_str(b, &a, "ordering", ordering); + (void)rocke_i_op0(b, ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, ops, 3, &a); +} + +void rocke_b_global_flag_wait_eq(rocke_ir_builder_t* b, + rocke_value_t* ptr, + rocke_value_t* idx, + rocke_value_t* expect, + const char* scope) +{ + rocke_value_t* ops[3]; + rocke_attr_map_t a; + if(!rocke_i_live(b)) + return; + if(!ptr || !idx || !expect) + { + (void)rocke_i_set_err(b, ROCKE_ERR_VALUE, "global_flag_wait_eq: null operand"); + return; + } + if(scope == NULL) + scope = "agent"; + if(!rocke_i_type_is(expect->type, "i32")) + { + (void)rocke_i_set_err(b, + ROCKE_ERR_VALUE, + "global_flag_wait_eq expects an i32 value, got %s", + expect->type->name); + return; + } + if(!rocke_i_fence_scope_ok(scope)) + { + (void)rocke_i_set_err(b, ROCKE_ERR_VALUE, "unknown flag scope '%s'", scope); + return; + } + ops[0] = ptr; + ops[1] = idx; + ops[2] = expect; + a = rocke_i_attrs(b); + rocke_attr_set_str(b, &a, "scope", scope); + (void)rocke_i_op0(b, ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ, ops, 3, &a); +} + rocke_value_t* rocke_b_global_atomic_add(rocke_ir_builder_t* b, rocke_value_t* ptr, rocke_value_t* idx, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp index ced6af4ad0a9..2d34472e106a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp @@ -697,6 +697,89 @@ static rocke_status_t _op_memref_global_atomic_add(rocke_h_lowerer_t* lw, const return lw->status; } +/* Python _HIP_ATOMIC_ORDER / _HIP_FENCE_SCOPE / _HIP_MEMORY_SCOPE. */ +static const char* hip_atomic_order(const char* o) +{ + if(!strcmp(o, "monotonic")) + return "__ATOMIC_RELAXED"; + if(!strcmp(o, "acquire")) + return "__ATOMIC_ACQUIRE"; + if(!strcmp(o, "release")) + return "__ATOMIC_RELEASE"; + if(!strcmp(o, "acq_rel")) + return "__ATOMIC_ACQ_REL"; + return "__ATOMIC_SEQ_CST"; +} + +static const char* hip_fence_scope(const char* s) +{ + return !strcmp(s, "system") ? "" : s; +} + +static const char* hip_memory_scope(const char* s) +{ + if(!strcmp(s, "workgroup")) + return "__HIP_MEMORY_SCOPE_WORKGROUP"; + if(!strcmp(s, "agent")) + return "__HIP_MEMORY_SCOPE_AGENT"; + return "__HIP_MEMORY_SCOPE_SYSTEM"; +} + +/* Python _op_memref_fence */ +static rocke_status_t _op_memref_fence(rocke_h_lowerer_t* lw, const rocke_op_t* op) +{ + if(!rocke_h_live(lw)) + { + return lw->status; + } + rocke_h_emitf(lw, + "__builtin_amdgcn_fence(%s, \"%s\");", + hip_atomic_order(mem_attr_str(op, "ordering", "acq_rel")), + hip_fence_scope(mem_attr_str(op, "scope", "agent"))); + return lw->status; +} + +/* Python _op_memref_global_flag_store */ +static rocke_status_t _op_memref_global_flag_store(rocke_h_lowerer_t* lw, const rocke_op_t* op) +{ + if(!rocke_h_live(lw)) + { + return lw->status; + } + if(op->num_operands < 3) + { + return rocke_h_fail(lw, ROCKE_ERR_VALUE, "memref.global_flag_store: too few operands"); + } + rocke_h_emitf(lw, + "__hip_atomic_store(&%s[%s], %s, %s, %s);", + rocke_h_name(lw, op->operands[0]), + rocke_h_name(lw, op->operands[1]), + rocke_h_name(lw, op->operands[2]), + hip_atomic_order(mem_attr_str(op, "ordering", "release")), + hip_memory_scope(mem_attr_str(op, "scope", "agent"))); + return lw->status; +} + +/* Python _op_memref_global_flag_wait_eq */ +static rocke_status_t _op_memref_global_flag_wait_eq(rocke_h_lowerer_t* lw, const rocke_op_t* op) +{ + if(!rocke_h_live(lw)) + { + return lw->status; + } + if(op->num_operands < 3) + { + return rocke_h_fail(lw, ROCKE_ERR_VALUE, "memref.global_flag_wait_eq: too few operands"); + } + rocke_h_emitf(lw, + "while (__hip_atomic_load(&%s[%s], __ATOMIC_ACQUIRE, %s) != %s) {}", + rocke_h_name(lw, op->operands[0]), + rocke_h_name(lw, op->operands[1]), + hip_memory_scope(mem_attr_str(op, "scope", "agent")), + rocke_h_name(lw, op->operands[2])); + return lw->status; +} + /* Python _op_memref_global_atomic_add_f32 */ static rocke_status_t _op_memref_global_atomic_add_f32(rocke_h_lowerer_t* lw, const rocke_op_t* op) { @@ -1517,6 +1600,9 @@ const rocke_h_handler_entry_t* rocke_h_handlers_mem(void) {ROCKE_OP_MEMREF_GLOBAL_ATOMIC_ADD_F32, _op_memref_global_atomic_add_f32}, {ROCKE_OP_MEMREF_GLOBAL_ATOMIC_ADD_PK_BF16, _op_memref_global_atomic_add_pk_bf16}, {ROCKE_OP_MEMREF_GLOBAL_ATOMIC_ADD_PK_F16, _op_memref_global_atomic_add_pk_f16}, + {ROCKE_OP_MEMREF_FENCE, _op_memref_fence}, + {ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, _op_memref_global_flag_store}, + {ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ, _op_memref_global_flag_wait_eq}, {ROCKE_OP_MEMREF_COOPERATIVE_GLOBAL_STORE, _op_memref_cooperative_global_store}, /* global pointer arithmetic + buffer rsrc */ {ROCKE_OP_TILE_GLOBAL_PTR_ADD, _op_tile_global_ptr_add}, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp index ab59c3061643..b9c81a457900 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp @@ -183,6 +183,85 @@ static void op_memref_global_store_typed(rocke_lower_t* L, const rocke_op_t* op) (long long)align); } +/* Python _llvm_syncscope: "system" is LLVM's default scope (no clause). */ +static const char* ll_syncscope(rocke_lower_t* L, const char* scope) +{ + if(strcmp(scope, "system") == 0) + return ""; + return rocke_arena_printf(&L->arena, " syncscope(\"%s\")", scope); +} + +/* Python _op_memref_fence. */ +static void op_memref_fence(rocke_lower_t* L, const rocke_op_t* op) +{ + rocke_ll_emitf(L, + " fence%s %s", + ll_syncscope(L, ll_attr_str(op, "scope", "agent")), + ll_attr_str(op, "ordering", "acq_rel")); +} + +/* Python _op_memref_global_flag_store. */ +static void op_memref_global_flag_store(rocke_lower_t* L, const rocke_op_t* op) +{ + const rocke_value_t* ptr = op->operands[0]; + const rocke_value_t* idx = op->operands[1]; + const rocke_value_t* val = op->operands[2]; + const char* gep = rocke_ll_fresh(L, "gep"); + rocke_ll_emitf(L, + " %s = getelementptr inbounds i32, ptr addrspace(1) %s, i32 %s", + gep, + rocke_ll_operand(L, ptr), + rocke_ll_operand(L, idx)); + rocke_ll_emitf(L, + " store atomic i32 %s, ptr addrspace(1) %s%s %s, align 4", + rocke_ll_operand(L, val), + gep, + ll_syncscope(L, ll_attr_str(op, "scope", "agent")), + ll_attr_str(op, "ordering", "release")); +} + +/* Python _op_memref_global_flag_wait_eq: a self-contained spin loop around an + * acquire atomic load; the fall-through block becomes current. */ +static void op_memref_global_flag_wait_eq(rocke_lower_t* L, const rocke_op_t* op) +{ + const rocke_value_t* ptr = op->operands[0]; + const rocke_value_t* idx = op->operands[1]; + const rocke_value_t* expect = op->operands[2]; + const char* gep = rocke_ll_fresh(L, "gep"); + rocke_ll_block_t* cur = rocke_ll_current(L); + rocke_ll_block_t* spin; + rocke_ll_block_t* done; + const char* val; + const char* hit; + rocke_ll_block_emitf(L, + cur, + " %s = getelementptr inbounds i32, ptr addrspace(1) %s, i32 %s", + gep, + rocke_ll_operand(L, ptr), + rocke_ll_operand(L, idx)); + spin = rocke_ll_new_block(L, "flag.wait"); + if(!cur || !spin) + return; + rocke_ll_block_emitf(L, cur, " br label %%%s", spin->label); + cur->terminated = true; + val = rocke_ll_fresh(L, "flag"); + hit = rocke_ll_fresh(L, "flag.hit"); + rocke_ll_block_emitf(L, + spin, + " %s = load atomic i32, ptr addrspace(1) %s%s acquire, align 4", + val, + gep, + ll_syncscope(L, ll_attr_str(op, "scope", "agent"))); + rocke_ll_block_emitf( + L, spin, " %s = icmp eq i32 %s, %s", hit, val, rocke_ll_operand(L, expect)); + done = rocke_ll_new_block(L, "flag.done"); + if(!done) + return; + rocke_ll_block_emitf( + L, spin, " br i1 %s, label %%%s, label %%%s", hit, done->label, spin->label); + spin->terminated = true; +} + static void op_memref_global_atomic_add(rocke_lower_t* L, const rocke_op_t* op) { const rocke_value_t* ptr = op->operands[0]; @@ -1635,6 +1714,9 @@ void rocke_ll_register_mem(void) rocke_ll_set_handler(ROCKE_OP_MEMREF_GLOBAL_STORE_TYPED, op_memref_global_store_typed); rocke_ll_set_handler(ROCKE_OP_MEMREF_GLOBAL_STORE_VN, op_memref_global_store_vN); rocke_ll_set_handler(ROCKE_OP_MEMREF_GLOBAL_ATOMIC_ADD, op_memref_global_atomic_add); + rocke_ll_set_handler(ROCKE_OP_MEMREF_FENCE, op_memref_fence); + rocke_ll_set_handler(ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, op_memref_global_flag_store); + rocke_ll_set_handler(ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ, op_memref_global_flag_wait_eq); rocke_ll_set_handler(ROCKE_OP_MEMREF_GLOBAL_ATOMIC_ADD_F32, op_memref_global_atomic_add_f32); rocke_ll_set_handler(ROCKE_OP_MEMREF_GLOBAL_ATOMIC_ADD_PK_BF16, op_memref_global_atomic_add_pk_bf16); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp index d30fcbb53c6a..a50e15a3b514 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp @@ -291,6 +291,8 @@ static const char* const* required_attrs(rocke_opcode_t o, int* n) static const char* mma_keys[] = {"op_id"}; static const char* yield_keys[] = {"num"}; static const char* asm_keys[] = {"template", "constraints"}; + static const char* scope_order_keys[] = {"scope", "ordering"}; + static const char* scope_keys[] = {"scope"}; switch(o) { case ROCKE_OP_ARITH_CONSTANT: @@ -309,6 +311,13 @@ static const char* const* required_attrs(rocke_opcode_t o, int* n) case ROCKE_OP_TILE_INLINE_ASM: *n = 2; return asm_keys; + case ROCKE_OP_MEMREF_FENCE: + case ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE: + *n = 2; + return scope_order_keys; + case ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ: + *n = 1; + return scope_keys; default: *n = 0; return NULL; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h index 8c6b09ed045a..360a67d8f289 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h @@ -423,6 +423,12 @@ typedef enum rocke_opcode ROCKE_OP_SCF_YIELD, ROCKE_OP_CF_RETURN, + /* memref.* cross-workgroup synchronisation (stream-K fixup). Appended + * after the control-flow block so no existing opcode index shifts. */ + ROCKE_OP_MEMREF_FENCE, + ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, + ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ, + ROCKE_OP__COUNT } rocke_opcode_t; @@ -796,6 +802,27 @@ rocke_value_t* rocke_b_cvt_pk_fp8_f32x4(rocke_ir_builder_t* b, rocke_value_t* v) rocke_value_t* rocke_b_cvt_pk_bf8_f32x4(rocke_ir_builder_t* b, rocke_value_t* v); rocke_value_t* rocke_b_cvt_pk_i8_f32x4(rocke_ir_builder_t* b, rocke_value_t* v); +/* ----- cross-workgroup synchronisation (stream-K fixup) ----- */ +/* IRBuilder.fence: scope in {workgroup, agent, system}; ordering in + * {acquire, release, acq_rel, seq_cst}. NULL selects the Python defaults + * (agent / acq_rel). */ +void rocke_b_fence(rocke_ir_builder_t* b, const char* scope, const char* ordering); +/* IRBuilder.global_flag_store: atomic i32 ptr[idx] = value. NULL scope / + * ordering select agent / release. */ +void rocke_b_global_flag_store(rocke_ir_builder_t* b, + rocke_value_t* ptr, + rocke_value_t* idx, + rocke_value_t* value, + const char* scope, + const char* ordering); +/* IRBuilder.global_flag_wait_eq: spin until ptr[idx] == expect (acquire). + * NULL scope selects agent. */ +void rocke_b_global_flag_wait_eq(rocke_ir_builder_t* b, + rocke_value_t* ptr, + rocke_value_t* idx, + rocke_value_t* expect, + const char* scope); + /* ----- atomics ----- */ rocke_value_t* rocke_b_global_atomic_add(rocke_ir_builder_t* b, rocke_value_t* ptr, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/intrinsics_and_primitives.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/intrinsics_and_primitives.md index 50813d2efb07..6c5aa7f7f667 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/intrinsics_and_primitives.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/intrinsics_and_primitives.md @@ -249,6 +249,32 @@ b.global_atomic_add(Hist, eid, b.const_i32(1)) local_off = b.global_atomic_add(Counter, eid, b.const_i32(1)) ``` +### Cross-workgroup flags and fences + +Stream-K fixups hand a partial tile from one workgroup to another through +global memory, guarded by a per-workgroup i32 flag. Three leaf ops cover +the protocol: + +```text +fence(*, scope="agent", ordering="acq_rel") +global_flag_store(ptr, idx, value, *, scope="agent", ordering="release") +global_flag_wait_eq(ptr, idx, expect, *, scope="agent") +``` + +`scope` is one of `workgroup`, `agent`, `system` (LLVM `syncscope`; +`system` has no clause). The producer stores its partial tile, every wave +issues `fence(ordering="release")`, the workgroup barriers, and then the +flag is stored. The consumer calls `global_flag_wait_eq`, which lowers to a +self-contained loop around an `acquire` atomic load, so its later global +loads observe the released partial. At `agent` scope on multi-XCD CDNA +parts the release writes back and the acquire invalidates the per-XCD L2. + +| IR op | LLVM | HIP | +|---|---|---| +| `memref.fence` | `fence syncscope("agent") release` | `__builtin_amdgcn_fence` | +| `memref.global_flag_store` | `store atomic i32 ... release` | `__hip_atomic_store` | +| `memref.global_flag_wait_eq` | `load atomic i32 ... acquire` spin | `while (__hip_atomic_load(...) != e)` | + Block-level histogram + scan helpers in `helpers/scan.py` wrap the common patterns: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md index 2856aad6821c..9af3c05f444b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md @@ -337,6 +337,11 @@ b.clamp_f32(v, lo, hi) # fmin(hi, fmax(lo, v)) -- folds to v_med3_f32 b.global_atomic_add(ptr, idx, val, *, ordering="monotonic") b.lds_atomic_add(smem, [idx], val, *, ordering="monotonic") +# : cross-workgroup synchronisation (stream-K fixup) +b.fence(*, scope="agent", ordering="acq_rel") # memref.fence +b.global_flag_store(ptr, idx, val_i32, *, scope="agent", ordering="release") +b.global_flag_wait_eq(ptr, idx, expect_i32, *, scope="agent") # acquire spin + # : bitwise integer helpers (also used by i4 unpack) b.xor(a, b) # arith.xor; same type on both operands b.shl(a, b) # arith.shl; logical left shift diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/op_vocabulary.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/op_vocabulary.md index da95fd84bb0d..4730a9f37f04 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/op_vocabulary.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/op_vocabulary.md @@ -94,6 +94,9 @@ Complete reference of operations recognized by `core/ir.py` and lowered to AMDGP | `memref.global_store_typed` | `global_store` | typed `store addrspace(1)` | | `memref.global_store_vN` | `global_store_vN` | vector `store addrspace(1)` | | `memref.global_atomic_add_f32`| `global_atomic_add_f32` | `atomicrmw fadd addrspace(1)` | +| `memref.fence` | `fence` | `fence syncscope(...) ` | +| `memref.global_flag_store` | `global_flag_store` | `store atomic i32 ... syncscope(...) ` | +| `memref.global_flag_wait_eq` | `global_flag_wait_eq` | spin loop around `load atomic i32 ... acquire` | `masked_global_load(ptr, idx, mask, other, dtype)` clamps the false-lane index to 0 before loading, then selects `other` afterward. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py index d08e3a6b54e8..5ea993b385a7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py @@ -100,6 +100,10 @@ def dtype_to_ir_type(dtype: str) -> Type: CACHE_STREAM = 2 # SLC set — streaming hint (don't evict useful lines). NON_TEMPORAL = 3 # GLC + SLC — bypass cache hierarchy entirely. +# LLVM ``syncscope`` values accepted by the cross-workgroup fence / flag ops. +# ``system`` renders as the empty (default) scope. +_FENCE_SCOPES = ("workgroup", "agent", "system") + # ----- target-neutral MMA metadata --------------------------------------- # @@ -1423,6 +1427,80 @@ def global_atomic_add_pk_f16( result_name_hint="atom_f16", ).result + # ----- cross-workgroup synchronisation (stream-K fixup) ----- + + def fence(self, *, scope: str = "agent", ordering: str = "acq_rel") -> None: + """Memory fence at ``scope`` (LLVM ``fence syncscope(...) ``). + + A ``release`` fence orders every prior global store of the issuing + wave before any later flag store; an ``acquire`` fence orders a prior + flag load before every later global load. At ``agent`` scope on + multi-XCD CDNA parts the backend also writes back / invalidates the + per-XCD L2, which is what makes a producer's partial tile visible to + a consumer workgroup running on a different XCD. + """ + if scope not in _FENCE_SCOPES: + raise ValueError(f"unknown fence scope {scope!r}") + if ordering not in ("acquire", "release", "acq_rel", "seq_cst"): + raise ValueError(f"unknown fence ordering {ordering!r}") + self._op("memref.fence", [], attrs={"scope": scope, "ordering": ordering}) + + def global_flag_store( + self, + ptr: Value, + idx: Value, + value: Value, + *, + scope: str = "agent", + ordering: str = "release", + ) -> None: + """Atomic i32 store of a cross-workgroup flag (``ptr[idx] = value``). + + Lowers to ``store atomic i32 ... syncscope(scope) ordering``. Pair + with :meth:`global_flag_wait_eq` on the consumer side. + """ + if value.type.name != "i32": + raise ValueError( + f"global_flag_store expects an i32 value, got {value.type.name}" + ) + if scope not in _FENCE_SCOPES: + raise ValueError(f"unknown flag scope {scope!r}") + if ordering not in ("monotonic", "release", "seq_cst"): + raise ValueError(f"unknown flag store ordering {ordering!r}") + self._op( + "memref.global_flag_store", + [ptr, idx, value], + attrs={"scope": scope, "ordering": ordering}, + ) + + def global_flag_wait_eq( + self, + ptr: Value, + idx: Value, + expect: Value, + *, + scope: str = "agent", + ) -> None: + """Spin until the i32 flag ``ptr[idx]`` equals ``expect``. + + Lowers to a self-contained loop around an ``acquire`` atomic load at + ``scope``, so global loads issued after the wait observe every store + the producer released before setting the flag. The loop is the only + data-dependent trip count in the IR; it is a leaf op so structured + control flow stays ``scf.for`` / ``scf.if`` only. + """ + if expect.type.name != "i32": + raise ValueError( + f"global_flag_wait_eq expects an i32 value, got {expect.type.name}" + ) + if scope not in _FENCE_SCOPES: + raise ValueError(f"unknown flag scope {scope!r}") + self._op( + "memref.global_flag_wait_eq", + [ptr, idx, expect], + attrs={"scope": scope}, + ) + def fp16_zero(self) -> Value: return self._op( "arith.constant", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py index 32249f619793..c24c37e5da89 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py @@ -263,6 +263,21 @@ def _encode_waitcnt_gfx11(vmcnt: int, expcnt: int, lgkmcnt: int) -> int: # Default arch for the HIP path. gfx950 is the historical CDNA target whose # emitted source is the byte-identical baseline; callers that don't pass an # arch (e.g. the in-tree coverage tests) keep getting exactly that output. +# Fence / flag-op scope and ordering spellings for the HIP atomic builtins. +_HIP_ATOMIC_ORDER = { + "monotonic": "__ATOMIC_RELAXED", + "acquire": "__ATOMIC_ACQUIRE", + "release": "__ATOMIC_RELEASE", + "acq_rel": "__ATOMIC_ACQ_REL", + "seq_cst": "__ATOMIC_SEQ_CST", +} +_HIP_FENCE_SCOPE = {"workgroup": "workgroup", "agent": "agent", "system": ""} +_HIP_MEMORY_SCOPE = { + "workgroup": "__HIP_MEMORY_SCOPE_WORKGROUP", + "agent": "__HIP_MEMORY_SCOPE_AGENT", + "system": "__HIP_MEMORY_SCOPE_SYSTEM", +} + _DEFAULT_HIP_ARCH = "gfx950" @@ -1710,6 +1725,30 @@ def _op_memref_global_atomic_add(self, op: Op) -> None: f"{_name(val)});" ) + def _op_memref_fence(self, op: Op) -> None: + """Lower ``memref.fence`` via ``__builtin_amdgcn_fence``.""" + self._emit( + f"__builtin_amdgcn_fence({_HIP_ATOMIC_ORDER[op.attrs['ordering']]}, " + f"\"{_HIP_FENCE_SCOPE[op.attrs['scope']]}\");" + ) + + def _op_memref_global_flag_store(self, op: Op) -> None: + """Lower ``memref.global_flag_store`` via ``__hip_atomic_store``.""" + ptr, idx, val = op.operands + self._emit( + f"__hip_atomic_store(&{_name(ptr)}[{_name(idx)}], {_name(val)}, " + f"{_HIP_ATOMIC_ORDER[op.attrs['ordering']]}, " + f"{_HIP_MEMORY_SCOPE[op.attrs['scope']]});" + ) + + def _op_memref_global_flag_wait_eq(self, op: Op) -> None: + """Lower ``memref.global_flag_wait_eq`` to an acquire-load spin loop.""" + ptr, idx, expect = op.operands + self._emit( + f"while (__hip_atomic_load(&{_name(ptr)}[{_name(idx)}], __ATOMIC_ACQUIRE, " + f"{_HIP_MEMORY_SCOPE[op.attrs['scope']]}) != {_name(expect)}) {{}}" + ) + def _op_memref_global_atomic_add_pk_bf16(self, op: Op) -> None: """Lower the packed-bf16 atomic add via the AMDGPU builtin. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py index ea3b0b47c9ee..d489b75e811e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py @@ -1226,6 +1226,16 @@ def _param_llvm_type(p: Param) -> str: return _llvm_type(p.type) +def _llvm_syncscope(scope: str) -> str: + """Render a fence / flag-op scope as an LLVM ``syncscope`` clause. + + ``system`` is LLVM's default scope and has no textual clause. + """ + if scope == "system": + return "" + return f' syncscope("{scope}")' + + def _escape_llvm_asm_string(s: str) -> str: r"""Escape a Python string for an LLVM IR asm/string literal. @@ -3033,6 +3043,53 @@ def _op_memref_global_atomic_add(self, op: Op) -> None: f"{elem_ty} {self._operand(val)} {ordering}{md}" ) + def _op_memref_fence(self, op: Op) -> None: + """Lower ``memref.fence`` to an LLVM ``fence`` at the op's scope.""" + self._current().emit( + f" fence{_llvm_syncscope(op.attrs['scope'])} {op.attrs['ordering']}" + ) + + def _op_memref_global_flag_store(self, op: Op) -> None: + """Lower ``memref.global_flag_store`` to a scoped ``store atomic i32``.""" + ptr, idx, val = op.operands + gep = self._fresh("gep") + self._current().emit( + f" {gep} = getelementptr inbounds i32, ptr addrspace(1) " + f"{self._operand(ptr)}, i32 {self._operand(idx)}" + ) + self._current().emit( + f" store atomic i32 {self._operand(val)}, ptr addrspace(1) {gep}" + f"{_llvm_syncscope(op.attrs['scope'])} {op.attrs['ordering']}, align 4" + ) + + def _op_memref_global_flag_wait_eq(self, op: Op) -> None: + """Lower ``memref.global_flag_wait_eq`` to a self-contained spin loop. + + The loop re-issues an ``acquire`` atomic load until the flag matches; + an atomic load cannot be hoisted out of the loop, and the acquire + orders every later global load after the observed flag value. + """ + ptr, idx, expect = op.operands + gep = self._fresh("gep") + cur = self._current() + cur.emit( + f" {gep} = getelementptr inbounds i32, ptr addrspace(1) " + f"{self._operand(ptr)}, i32 {self._operand(idx)}" + ) + spin = self._new_block("flag.wait") + cur.emit(f" br label %{spin.label}") + cur.terminated = True + val = self._fresh("flag") + hit = self._fresh("flag.hit") + spin.emit( + f" {val} = load atomic i32, ptr addrspace(1) {gep}" + f"{_llvm_syncscope(op.attrs['scope'])} acquire, align 4" + ) + spin.emit(f" {hit} = icmp eq i32 {val}, {self._operand(expect)}") + done = self._new_block("flag.done") + spin.emit(f" br i1 {hit}, label %{done.label}, label %{spin.label}") + spin.terminated = True + def _op_memref_global_atomic_add_pk_bf16(self, op: Op) -> None: """Lower the packed-bf16 atomic add via a generic ``atomicrmw fadd`` on a ``<2 x bfloat>``. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/verify.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/verify.py index a099becd169a..2618acebc66f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/verify.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/verify.py @@ -120,6 +120,9 @@ def __str__(self) -> str: "tile.mma": ["op_id"], "scf.yield": ["num"], "tile.inline_asm": ["template", "constraints"], + "memref.fence": ["scope", "ordering"], + "memref.global_flag_store": ["scope", "ordering"], + "memref.global_flag_wait_eq": ["scope"], } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp index d2f9209d7039..f423b3ebec04 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp @@ -607,6 +607,57 @@ void case_s_wqm() EXPECT_IR(i64_ir, "call i64 @llvm.amdgcn.s.wqm.i64.i64(i64 15)"); } +/* ---- cross-workgroup fence / flag ops (stream-K fixup) ---- */ +void case_flag_ops() +{ + const std::string ir = lower_one("flagwt", [](rocke_ir_builder_t* b) { + rocke_value_t* f = global_ptr_param(b, "flags", rocke_i32()); + rocke_b_global_flag_store( + b, f, rocke_b_const_i32(b, 3), rocke_b_const_i32(b, 1), NULL, NULL); + rocke_b_fence(b, "agent", "release"); + rocke_b_global_flag_wait_eq(b, f, rocke_b_const_i32(b, 2), rocke_b_const_i32(b, 1), NULL); + }); + EXPECT_IR(ir, + " %gep.1 = getelementptr inbounds i32, ptr addrspace(1) %flags, i32 3\n" + " store atomic i32 1, ptr addrspace(1) %gep.1 syncscope(\"agent\") release, " + "align 4\n" + " fence syncscope(\"agent\") release\n" + " %gep.2 = getelementptr inbounds i32, ptr addrspace(1) %flags, i32 2\n" + " br label %flag.wait.1\n" + "flag.wait.1:\n" + " %flag.3 = load atomic i32, ptr addrspace(1) %gep.2 syncscope(\"agent\") " + "acquire, align 4\n" + " %flag.hit.4 = icmp eq i32 %flag.3, 1\n" + " br i1 %flag.hit.4, label %flag.done.2, label %flag.wait.1\n" + "flag.done.2:\n"); +} + +void case_fence_system_scope() +{ + const std::string ir = lower_one( + "fencesys", [](rocke_ir_builder_t* b) { rocke_b_fence(b, "system", "acquire"); }); + EXPECT_IR(ir, " fence acquire\n"); + EXPECT_NO_IR(ir, "syncscope"); +} + +void case_flag_ops_hip() +{ + const std::string hip = lower_one_hip( + "flaghip", + [](rocke_ir_builder_t* b) { + rocke_value_t* f = global_ptr_param(b, "flags", rocke_i32()); + rocke_b_global_flag_store( + b, f, rocke_b_const_i32(b, 0), rocke_b_const_i32(b, 1), NULL, NULL); + rocke_b_fence(b, "agent", "acquire"); + rocke_b_global_flag_wait_eq( + b, f, rocke_b_const_i32(b, 0), rocke_b_const_i32(b, 1), NULL); + }, + "gfx950"); + EXPECT_IR(hip, "__ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_AGENT);"); + EXPECT_IR(hip, "__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, \"agent\");"); + EXPECT_IR(hip, "while (__hip_atomic_load("); +} + /* ---- av.load / av.store (agent-scope 128-bit vector mem) ---- */ void case_av_load_b128() { @@ -1019,6 +1070,9 @@ const TestCase k_cases[] = { {"hip_zext_uses_unsigned_source_cast", case_hip_zext_uses_unsigned_source_cast}, {"gfx1250_scaled_wmma", case_gfx1250_scaled_wmma}, {"gfx1250_standalone_bridge", case_gfx1250_standalone_bridge}, + {"flag_ops", case_flag_ops}, + {"fence_system_scope", case_fence_system_scope}, + {"flag_ops_hip", case_flag_ops_hip}, {"opcode_names_are_aligned", case_opcode_names_are_aligned}, }; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py index 13d540571328..3be30150e6a3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py @@ -3285,6 +3285,80 @@ def _lower(self, name, build, **kw): _assert_ir_assembles(self, ll, name) return ll + # ---- cross-workgroup fence / flag ops (stream-K fixup) ---- + @staticmethod + def _flag_kernel(b): + from rocke.core.ir import I32, PtrType + + return b.param("flags", PtrType(I32, "global")) + + def test_fence_renders_scope_and_ordering(self): + ll = self._lower("fence", lambda b: b.fence(scope="agent", ordering="release")) + self.assertIn(' fence syncscope("agent") release', ll) + + def test_fence_system_scope_has_no_syncscope_clause(self): + ll = self._lower( + "fencesys", lambda b: b.fence(scope="system", ordering="acquire") + ) + self.assertIn(" fence acquire", ll) + self.assertNotIn("syncscope", ll) + + def test_global_flag_store_is_scoped_atomic_store(self): + def build(b): + flags = self._flag_kernel(b) + b.global_flag_store(flags, b.const_i32(3), b.const_i32(1)) + + ll = self._lower("flagst", build) + self.assertRegex( + ll, + r"store atomic i32 1, ptr addrspace\(1\) %gep\.\d+ " + r'syncscope\("agent"\) release, align 4', + ) + + def test_global_flag_wait_eq_is_acquire_spin_loop(self): + def build(b): + flags = self._flag_kernel(b) + b.global_flag_wait_eq(flags, b.const_i32(2), b.const_i32(1)) + + ll = self._lower("flagwt", build) + self.assertRegex( + ll, + r"(?s)br label %(flag\.wait\.\d+)\n\1:\n" + r" (%flag\.\d+) = load atomic i32, ptr addrspace\(1\) %gep\.\d+ " + r'syncscope\("agent"\) acquire, align 4\n' + r" (%flag\.hit\.\d+) = icmp eq i32 \2, 1\n" + r" br i1 \3, label %flag\.done\.\d+, label %\1\n", + ) + + def test_flag_ops_lower_to_hip_atomics(self): + from rocke.core.lower_hip import lower_kernel_to_hip + + b = self._builder("flaghip") + flags = self._flag_kernel(b) + b.global_flag_store(flags, b.const_i32(0), b.const_i32(1)) + b.fence(scope="agent", ordering="acquire") + b.global_flag_wait_eq(flags, b.const_i32(0), b.const_i32(1)) + hip = lower_kernel_to_hip(b.kernel) + self.assertIn("__ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_AGENT);", hip) + self.assertIn('__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "agent");', hip) + self.assertIn("while (__hip_atomic_load(", hip) + + def test_flag_ops_reject_bad_scope_ordering_and_type(self): + b = self._builder("flagbad") + flags = self._flag_kernel(b) + with self.assertRaisesRegex(ValueError, "unknown fence scope"): + b.fence(scope="device") + with self.assertRaisesRegex(ValueError, "unknown fence ordering"): + b.fence(ordering="monotonic") + with self.assertRaisesRegex(ValueError, "expects an i32 value"): + b.global_flag_store(flags, b.const_i32(0), b.const_f32(1.0)) + with self.assertRaisesRegex(ValueError, "unknown flag store ordering"): + b.global_flag_store( + flags, b.const_i32(0), b.const_i32(1), ordering="acquire" + ) + with self.assertRaisesRegex(ValueError, "unknown flag scope"): + b.global_flag_wait_eq(flags, b.const_i32(0), b.const_i32(1), scope="gpu") + # ---- ds_swizzle (raw offset + XOR-butterfly encoding) ---- def test_ds_swizzle_passes_raw_offset_immediate(self): ll = self._lower("dssw", lambda b: b.ds_swizzle(b.const_i32(1), 0x041F)) From bcd26d58f236a3a3045840cac78dddc2cd41f3b8 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Mon, 5 Oct 2026 10:43:32 +0000 Subject: [PATCH 02/20] feat(rocke): add CK stream-K iteration partitioner helper Add StreamKIterPartition, the CK Tile StreamKTilePartitioner resolved at build time: output tiles that divide the CTA pool run data-parallel and the remainder is spread over max_active_wgs stream-K CTAs, each owning a contiguous, balanced range of MAC iterations that may straddle output tiles. The existing StreamKPartition used by streamk_gemm is unchanged. The helper exposes the host-side plan (DP/SK split, grid size, Linear partner and Tree round bounds, flag-table size), start/end iteration of a stream-K CTA, and SSA emitters that map a CTA onto its iteration range. The C port returns the resolved plan by value so no derived quantity can be recomputed with different rounding. Reference vectors and an invariant sweep (contiguous coverage, at most one non-owner tile per stream-K CTA, contributor bounds) run in both engines; the C test is registered as rocke_streamk_iter_partition. ISSUE ID : AICK-1943 Co-Authored-By: Claude Opus 5.5 --- ...ROCKE_ENGINE_test_categories_external.yaml | 2 + .../rocke/platform/cpp/helpers/streamk.cpp | 166 +++++++++++++ .../rocke/helper_rocke.helpers.streamk.h | 76 ++++++ .../platform/python/rocke/helpers/__init__.py | 12 + .../platform/python/rocke/helpers/streamk.py | 222 ++++++++++++++++++ .../rocke/platform/tests/CMakeLists.txt | 20 ++ .../instances/streamk_iter_partition.cpp | 150 ++++++++++++ .../instances/test_streamk_iter_partition.py | 135 +++++++++++ 8 files changed, 783 insertions(+) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/streamk_iter_partition.cpp create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_streamk_iter_partition.py diff --git a/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml b/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml index 3c7996e6b099..9079d828529a 100644 --- a/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml +++ b/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml @@ -28,6 +28,7 @@ test_categories: - "rocke_arch_target_identity" - "rocke_ir_serialize_roundtrip" - "rocke_mma_frag_ssot" + - "rocke_streamk_iter_partition" - "rocke_tiled_attention_2d_reentrancy" - "rocke_python_import_smoke" - "rocke_golden_static" @@ -47,6 +48,7 @@ test_categories: - "rocke_arch_target_identity" - "rocke_ir_serialize_roundtrip" - "rocke_mma_frag_ssot" + - "rocke_streamk_iter_partition" - "rocke_tiled_attention_2d_reentrancy" - "rocke_python_import_smoke" - "rocke_golden_static" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/streamk.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/streamk.cpp index 60abbab8d85b..2c459f20e74a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/streamk.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/streamk.cpp @@ -138,3 +138,169 @@ rocke_streamk_decoded_tile_t rocke_emit_streamk_decode(rocke_ir_builder_t* b, return res; } + +/* ------------------------------------------------------------------------ + * Iteration-balanced stream-K (Python StreamKIterPartition and friends). + * ------------------------------------------------------------------------ */ +static int sk_ceil_div(int a, int c) +{ + return (a + c - 1) / c; +} + +rocke_streamk_iter_partition_t rocke_streamk_iter_partition_make(int m_tiles, + int n_tiles, + int iters_per_tile, + int max_active_wgs, + bool persistent, + rocke_status_t* out_status) +{ + rocke_streamk_iter_partition_t p; + p.m_tiles = 0; + p.n_tiles = 0; + p.iters_per_tile = 0; + p.max_active_wgs = 0; + p.persistent = false; + if(m_tiles <= 0 || n_tiles <= 0 || iters_per_tile <= 0 || max_active_wgs <= 0) + { + if(out_status) + *out_status = ROCKE_ERR_VALUE; + return p; + } + p.m_tiles = m_tiles; + p.n_tiles = n_tiles; + p.iters_per_tile = iters_per_tile; + p.max_active_wgs = max_active_wgs; + p.persistent = persistent; + if(out_status) + *out_status = ROCKE_OK; + return p; +} + +rocke_streamk_iter_plan_t rocke_streamk_iter_plan(const rocke_streamk_iter_partition_t* p) +{ + rocke_streamk_iter_plan_t r; + int rem; + int contributors; + int rest; + r.num_tiles = p->m_tiles * p->n_tiles; + /* sk_tiles: CK "DP + 2-tile SK", reverting to all-DP when the remainder + * cannot give every stream-K CTA one iteration. */ + rem = p->max_active_wgs > 0 ? r.num_tiles % p->max_active_wgs : 0; + if(rem == 0) + r.sk_tiles = 0; + else + { + r.sk_tiles = r.num_tiles > p->max_active_wgs ? p->max_active_wgs + rem : r.num_tiles; + if(r.sk_tiles * p->iters_per_tile < p->max_active_wgs) + r.sk_tiles = 0; + } + r.sk_ctas = r.sk_tiles ? p->max_active_wgs : 0; + r.total_sk_iters = r.sk_tiles * p->iters_per_tile; + r.iters_per_sk_cta = r.sk_ctas ? r.total_sk_iters / r.sk_ctas : 0; + r.extra_iters = r.sk_ctas ? r.total_sk_iters % r.sk_ctas : 0; + r.dp_tiles = r.num_tiles - r.sk_tiles; + r.total_dp_iters = r.dp_tiles * p->iters_per_tile; + r.grid_size = p->persistent ? p->max_active_wgs : r.dp_tiles + r.sk_ctas; + if(!r.sk_ctas) + r.max_linear_partners = 0; + else + r.max_linear_partners = sk_ceil_div(p->iters_per_tile - 1, r.iters_per_sk_cta); + /* tree_rounds = (contributors - 1).bit_length() */ + contributors = r.max_linear_partners + 1; + r.tree_rounds = 0; + for(rest = contributors - 1; rest > 0; rest >>= 1) + ++r.tree_rounds; + r.flags_bytes = sk_ceil_div(4 * r.sk_ctas, 256) * 256; + return r; +} + +int rocke_streamk_start_iter(const rocke_streamk_iter_partition_t* p, int sk_cta) +{ + const rocke_streamk_iter_plan_t r = rocke_streamk_iter_plan(p); + return r.total_dp_iters + sk_cta * r.iters_per_sk_cta + + (sk_cta < r.extra_iters ? sk_cta : r.extra_iters); +} + +int rocke_streamk_end_iter(const rocke_streamk_iter_partition_t* p, int sk_cta) +{ + const rocke_streamk_iter_plan_t r = rocke_streamk_iter_plan(p); + return rocke_streamk_start_iter(p, sk_cta) + r.iters_per_sk_cta + + (sk_cta < r.extra_iters ? 1 : 0); +} + +/* Python emit_streamk_sk_start_iter: + * c_dp_iters = b.const_i32(part.total_dp_iters) + * c_per_cta = b.const_i32(part.iters_per_sk_cta) + * c_extra = b.const_i32(part.extra_iters) + * body = b.mul(sk_cta, c_per_cta) + * lead = b.smin(sk_cta, c_extra) + * return b.add(c_dp_iters, b.add(body, lead)) + */ +rocke_value_t* rocke_emit_streamk_sk_start_iter(rocke_ir_builder_t* b, + rocke_value_t* sk_cta, + const rocke_streamk_iter_partition_t* p) +{ + const rocke_streamk_iter_plan_t r = rocke_streamk_iter_plan(p); + rocke_value_t* c_dp_iters = rocke_b_const_i32(b, (int64_t)r.total_dp_iters); + rocke_value_t* c_per_cta = rocke_b_const_i32(b, (int64_t)r.iters_per_sk_cta); + rocke_value_t* c_extra = rocke_b_const_i32(b, (int64_t)r.extra_iters); + rocke_value_t* body = rocke_b_mul(b, sk_cta, c_per_cta); + rocke_value_t* lead = rocke_b_smin(b, sk_cta, c_extra); + rocke_value_t* inner = rocke_b_add(b, body, lead); + return rocke_b_add(b, c_dp_iters, inner); +} + +/* Python emit_streamk_iter_range; every nested builder call is pinned to a + * temporary in the Python evaluation order. */ +rocke_streamk_iter_range_t rocke_emit_streamk_iter_range(rocke_ir_builder_t* b, + rocke_value_t* cta, + const rocke_streamk_iter_partition_t* p, + bool dp) +{ + const rocke_streamk_iter_plan_t r = rocke_streamk_iter_plan(p); + rocke_streamk_iter_range_t res; + rocke_value_t* c_dp_tiles = NULL; + rocke_value_t* sk_cta; + rocke_value_t* sk_start; + rocke_value_t* c_per_cta; + rocke_value_t* c_extra; + rocke_value_t* c_one; + rocke_value_t* c_zero; + rocke_value_t* has_extra; + rocke_value_t* extra; + rocke_value_t* sk_len; + rocke_value_t* start; + rocke_value_t* length; + if(dp) + { + c_dp_tiles = rocke_b_const_i32(b, (int64_t)r.dp_tiles); + sk_cta = rocke_b_sub(b, cta, c_dp_tiles); + } + else + sk_cta = cta; + sk_start = rocke_emit_streamk_sk_start_iter(b, sk_cta, p); + c_per_cta = rocke_b_const_i32(b, (int64_t)r.iters_per_sk_cta); + c_extra = rocke_b_const_i32(b, (int64_t)r.extra_iters); + c_one = rocke_b_const_i32(b, 1); + c_zero = rocke_b_const_i32(b, 0); + has_extra = rocke_b_cmp_lt(b, sk_cta, c_extra); + extra = rocke_b_select(b, has_extra, c_one, c_zero); + sk_len = rocke_b_add(b, c_per_cta, extra); + if(dp) + { + rocke_value_t* c_ipt = rocke_b_const_i32(b, (int64_t)p->iters_per_tile); + rocke_value_t* is_dp = rocke_b_cmp_lt(b, cta, c_dp_tiles); + rocke_value_t* dp_start = rocke_b_mul(b, cta, c_ipt); + start = rocke_b_select(b, is_dp, dp_start, sk_start); + length = rocke_b_select(b, is_dp, c_ipt, sk_len); + } + else + { + start = sk_start; + length = sk_len; + } + res.start = rocke_b_to_sgpr_u32(b, start); + res.end = rocke_b_to_sgpr_u32(b, rocke_b_add(b, res.start, length)); + res.sk_cta = sk_cta; + return res; +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.streamk.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.streamk.h index 9a43f6ef2587..50ca10de8774 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.streamk.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.streamk.h @@ -141,6 +141,82 @@ rocke_streamk_decoded_tile_t rocke_emit_streamk_decode(rocke_ir_builder_t* b, rocke_value_t* linear_id, const rocke_streamk_partition_t* spec); +/* ------------------------------------------------------------------ * + * Iteration-balanced stream-K (CK Tile StreamKTilePartitioner) + * ------------------------------------------------------------------ * + * Python C99 (this header) + * ----------------------------------- ----------------------------------- + * class StreamKIterPartition rocke_streamk_iter_partition_t + * (its @property values) + rocke_streamk_iter_plan_t, resolved + * once by rocke_streamk_iter_plan() + * streamk_iter_partition(...) rocke_streamk_iter_partition_make(...) + * streamk_start_iter / streamk_end_iter rocke_streamk_start_iter / _end_iter + * emit_streamk_sk_start_iter(...) rocke_emit_streamk_sk_start_iter(...) + * emit_streamk_iter_range(...) rocke_emit_streamk_iter_range(...) + * + * The derived quantities are returned as one resolved plan by value rather + * than as per-property accessors, so the C side cannot recompute any of them + * with different rounding than the Python @property it mirrors. + */ +typedef struct rocke_streamk_iter_partition +{ + int m_tiles; + int n_tiles; + int iters_per_tile; + int max_active_wgs; + bool persistent; +} rocke_streamk_iter_partition_t; + +typedef struct rocke_streamk_iter_plan +{ + int num_tiles; + int sk_tiles; + int sk_ctas; + int total_sk_iters; + int iters_per_sk_cta; + int extra_iters; + int dp_tiles; + int total_dp_iters; + int grid_size; + int max_linear_partners; + int tree_rounds; + int flags_bytes; +} rocke_streamk_iter_plan_t; + +/* streamk_iter_partition: every count must be > 0. On violation *out_status + * (if non-NULL) is ROCKE_ERR_VALUE and the returned partition is zeroed. */ +rocke_streamk_iter_partition_t rocke_streamk_iter_partition_make(int m_tiles, + int n_tiles, + int iters_per_tile, + int max_active_wgs, + bool persistent, + rocke_status_t* out_status); + +/* Every StreamKIterPartition @property, resolved once. */ +rocke_streamk_iter_plan_t rocke_streamk_iter_plan(const rocke_streamk_iter_partition_t* p); + +int rocke_streamk_start_iter(const rocke_streamk_iter_partition_t* p, int sk_cta); +int rocke_streamk_end_iter(const rocke_streamk_iter_partition_t* p, int sk_cta); + +/* emit_streamk_sk_start_iter: SSA start iteration of a runtime SK CTA index. */ +rocke_value_t* rocke_emit_streamk_sk_start_iter(rocke_ir_builder_t* b, + rocke_value_t* sk_cta, + const rocke_streamk_iter_partition_t* p); + +/* emit_streamk_iter_range result (Python _IterRange). */ +typedef struct rocke_streamk_iter_range +{ + rocke_value_t* start; /* i32 global first iteration (SGPR) */ + rocke_value_t* end; /* i32 one past the last iteration (SGPR) */ + rocke_value_t* sk_cta; /* i32 stream-K CTA index; meaningless for a DP CTA */ +} rocke_streamk_iter_range_t; + +/* emit_streamk_iter_range(b, cta, part, *, dp). */ +rocke_streamk_iter_range_t rocke_emit_streamk_iter_range(rocke_ir_builder_t* b, + rocke_value_t* cta, + const rocke_streamk_iter_partition_t* p, + bool dp); + #ifdef __cplusplus } /* extern "C" */ #endif diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py index 4e015d6d7411..52b72370c896 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py @@ -346,13 +346,19 @@ ) from .schedule import SchedulePolicy from .streamk import ( + StreamKIterPartition, StreamKPartition, StreamKReductionStrategy, compute_streamk_grid_size, emit_streamk_decode, + emit_streamk_iter_range, emit_streamk_partial_load_accumulate, emit_streamk_partial_store, + emit_streamk_sk_start_iter, + streamk_end_iter, + streamk_iter_partition, streamk_num_macro_tiles, + streamk_start_iter, ) from .spec import ( IOSpecRule, @@ -613,12 +619,18 @@ "load_sorted_topk_weight", "scatter_token_offset", # StreamK partitioner ( StreamK GEMM) + "StreamKIterPartition", "StreamKPartition", "StreamKReductionStrategy", "compute_streamk_grid_size", "emit_streamk_decode", + "emit_streamk_iter_range", "emit_streamk_partial_load_accumulate", "emit_streamk_partial_store", + "emit_streamk_sk_start_iter", + "streamk_end_iter", + "streamk_iter_partition", + "streamk_start_iter", "mfma_attention_bwd_dot_do_o_inner_body", "mfma_attention_bwd_dq_dk_dv_inner_body", "streamk_num_macro_tiles", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/streamk.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/streamk.py index 587bffd19c93..e1cb605687b0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/streamk.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/streamk.py @@ -55,13 +55,19 @@ __all__ = [ + "StreamKIterPartition", "StreamKPartition", "StreamKReductionStrategy", "compute_streamk_grid_size", "emit_streamk_decode", + "emit_streamk_iter_range", "emit_streamk_partial_load_accumulate", "emit_streamk_partial_store", + "emit_streamk_sk_start_iter", + "streamk_end_iter", + "streamk_iter_partition", "streamk_num_macro_tiles", + "streamk_start_iter", ] @@ -233,3 +239,219 @@ def emit_streamk_partial_load_accumulate( (``WaitStorePartialDone``). """ return b.global_load_f32(workspace, workspace_off) + + +# --------------------------------------------------------------------------- +# Iteration-balanced stream-K (CK Tile StreamKTilePartitioner) +# --------------------------------------------------------------------------- +# +# The partition above hands every CTA exactly one ``tile_k`` slice. CK Tile's +# production partitioner instead balances *iterations*: output tiles that +# divide evenly across the CTA pool run data-parallel (one CTA per tile, the +# full K loop), and only the remainder is spread over ``max_active_wgs`` +# stream-K CTAs, each owning a contiguous range of MAC-loop iterations that +# may straddle output-tile boundaries. rocKE compiles one kernel per shape, so +# the whole partition is resolved host-side and folded into the IR as +# constants; the device only maps its linear CTA id onto an iteration range. + + +def _ceil_div(a: int, b: int) -> int: + return -(-a // b) + + +@dataclass(frozen=True) +class StreamKIterPartition: + """CK Tile ``StreamKTilePartitionerBase``, resolved at build time. + + ``m_tiles`` / ``n_tiles`` count output tiles (a caller folding a batch + dimension into M passes the folded count). ``iters_per_tile`` is the + number of ``tile_k`` MAC iterations per output tile. ``max_active_wgs`` + is the CTA pool the stream-K remainder is spread over; every stream-K + CTA must be co-resident with the CTAs it waits on, so it must not exceed + the number of workgroups the device can hold at once. + """ + + m_tiles: int + n_tiles: int + iters_per_tile: int + max_active_wgs: int + persistent: bool = False + + @property + def num_tiles(self) -> int: + return self.m_tiles * self.n_tiles + + @property + def sk_tiles(self) -> int: + """Output tiles handled by stream-K CTAs (0 means all-DP). + + CK's "DP + 2-tile SK": when the tiles do not divide the CTA pool, + one full wave of tiles plus the remainder go stream-K so the last + wave is balanced; a remainder too small to give every stream-K CTA + at least one iteration falls back to all-DP. + """ + rem = self.num_tiles % self.max_active_wgs + if rem == 0: + return 0 + if self.num_tiles > self.max_active_wgs: + tiles = self.max_active_wgs + rem + else: + tiles = self.num_tiles + if tiles * self.iters_per_tile < self.max_active_wgs: + return 0 + return tiles + + @property + def sk_ctas(self) -> int: + return self.max_active_wgs if self.sk_tiles else 0 + + @property + def total_sk_iters(self) -> int: + return self.sk_tiles * self.iters_per_tile + + @property + def iters_per_sk_cta(self) -> int: + return self.total_sk_iters // self.sk_ctas if self.sk_ctas else 0 + + @property + def extra_iters(self) -> int: + """Stream-K CTAs ``[0, extra_iters)`` own one extra iteration.""" + return self.total_sk_iters % self.sk_ctas if self.sk_ctas else 0 + + @property + def dp_tiles(self) -> int: + return self.num_tiles - self.sk_tiles + + @property + def total_dp_iters(self) -> int: + return self.dp_tiles * self.iters_per_tile + + @property + def grid_size(self) -> int: + """Launch grid (x): the CTA pool when persistent, else DP + SK CTAs.""" + if self.persistent: + return self.max_active_wgs + return self.dp_tiles + self.sk_ctas + + @property + def max_linear_partners(self) -> int: + """Upper bound on the CTAs a tile owner has to fold in (Linear). + + The owner is whichever CTA covers the tile's first iteration; it may + have started in the previous tile, so it is only guaranteed one + iteration here. Every later contributor covers at least + ``iters_per_sk_cta`` iterations until the tile ends. + """ + if not self.sk_ctas: + return 0 + return _ceil_div(self.iters_per_tile - 1, self.iters_per_sk_cta) + + @property + def tree_rounds(self) -> int: + """Pairwise fan-in rounds bounding the Tree fixup: ``ceil(log2)`` of + the most CTAs that can contribute to one tile.""" + contributors = self.max_linear_partners + 1 + return (contributors - 1).bit_length() + + @property + def flags_bytes(self) -> int: + """One i32 flag per stream-K CTA, padded to 256 bytes (CK layout).""" + return _ceil_div(4 * self.sk_ctas, 256) * 256 + + +def streamk_iter_partition( + *, + m_tiles: int, + n_tiles: int, + iters_per_tile: int, + max_active_wgs: int, + persistent: bool = False, +) -> StreamKIterPartition: + """Validated constructor for :class:`StreamKIterPartition`.""" + for name, v in ( + ("m_tiles", m_tiles), + ("n_tiles", n_tiles), + ("iters_per_tile", iters_per_tile), + ("max_active_wgs", max_active_wgs), + ): + if v <= 0: + raise ValueError(f"stream-K partition: {name} must be > 0 (got {v})") + return StreamKIterPartition( + m_tiles=m_tiles, + n_tiles=n_tiles, + iters_per_tile=iters_per_tile, + max_active_wgs=max_active_wgs, + persistent=persistent, + ) + + +def streamk_start_iter(part: StreamKIterPartition, sk_cta: int) -> int: + """Global first MAC iteration of stream-K CTA ``sk_cta``.""" + return ( + part.total_dp_iters + + sk_cta * part.iters_per_sk_cta + + min(sk_cta, part.extra_iters) + ) + + +def streamk_end_iter(part: StreamKIterPartition, sk_cta: int) -> int: + """One past the last MAC iteration of stream-K CTA ``sk_cta``.""" + return ( + streamk_start_iter(part, sk_cta) + + part.iters_per_sk_cta + + (1 if sk_cta < part.extra_iters else 0) + ) + + +def emit_streamk_sk_start_iter( + b: IRBuilder, sk_cta: Value, part: StreamKIterPartition +) -> Value: + """SSA :func:`streamk_start_iter` for a runtime stream-K CTA index.""" + c_dp_iters = b.const_i32(part.total_dp_iters) + c_per_cta = b.const_i32(part.iters_per_sk_cta) + c_extra = b.const_i32(part.extra_iters) + body = b.mul(sk_cta, c_per_cta) + lead = b.smin(sk_cta, c_extra) + return b.add(c_dp_iters, b.add(body, lead)) + + +class _IterRange(NamedTuple): + """The SSA bundle returned by :func:`emit_streamk_iter_range`.""" + + start: Value # i32 global first iteration (SGPR) + end: Value # i32 one past the last iteration (SGPR) + sk_cta: Value # i32 stream-K CTA index; meaningless for a DP CTA + + +def emit_streamk_iter_range( + b: IRBuilder, cta: Value, part: StreamKIterPartition, *, dp: bool +) -> _IterRange: + """Map a CTA onto its contiguous range of global MAC iterations. + + ``dp=True`` (non-persistent launch): CTAs ``[0, dp_tiles)`` each own one + whole tile and the rest are stream-K CTAs, selected branch-free so the + GEMM body is emitted once. ``dp=False``: ``cta`` already is a stream-K + CTA index (the persistent launch runs its DP sweep separately). + """ + if dp: + c_dp_tiles = b.const_i32(part.dp_tiles) + sk_cta = b.sub(cta, c_dp_tiles) + else: + sk_cta = cta + sk_start = emit_streamk_sk_start_iter(b, sk_cta, part) + c_per_cta = b.const_i32(part.iters_per_sk_cta) + c_extra = b.const_i32(part.extra_iters) + c_one = b.const_i32(1) + c_zero = b.const_i32(0) + sk_len = b.add(c_per_cta, b.select(b.cmp_lt(sk_cta, c_extra), c_one, c_zero)) + if dp: + c_ipt = b.const_i32(part.iters_per_tile) + is_dp = b.cmp_lt(cta, c_dp_tiles) + start = b.select(is_dp, b.mul(cta, c_ipt), sk_start) + length = b.select(is_dp, c_ipt, sk_len) + else: + start = sk_start + length = sk_len + start = b.to_sgpr_u32(start) + end = b.to_sgpr_u32(b.add(start, length)) + return _IterRange(start=start, end=end, sk_cta=sk_cta) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt b/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt index 87d32004f5d2..26c71f9f0b40 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt @@ -246,6 +246,26 @@ if(BUILD_TESTING) rocke_moe_fused_mega_spec_guards) endif() + # Host-side stream-K partition math: the launch grid it sizes is invisible to + # the byte-identity gate, so the C port is pinned to the Python reference + # vectors here (mirrors tests/instances/test_streamk_iter_partition.py). + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/instances/streamk_iter_partition.cpp") + add_executable(rocke_streamk_iter_partition instances/streamk_iter_partition.cpp) + target_link_libraries(rocke_streamk_iter_partition PRIVATE rocke_core) + if(UNIX) + target_link_libraries(rocke_streamk_iter_partition PRIVATE m) + endif() + add_test(NAME rocke_streamk_iter_partition COMMAND rocke_streamk_iter_partition) + set_tests_properties(rocke_streamk_iter_partition PROPERTIES LABELS "ckc;unit;host") + rocke_cxx_test_installed_name(_rocke_installed_name rocke_streamk_iter_partition) + set_target_properties(rocke_streamk_iter_partition + PROPERTIES OUTPUT_NAME "${_rocke_installed_name}") + install(TARGETS rocke_streamk_iter_partition + RUNTIME DESTINATION ${ROCKE_CXX_TEST_INSTALL_DIR}) + set_property(GLOBAL APPEND PROPERTY ROCKE_INSTALL_TEST_TARGETS + rocke_streamk_iter_partition) + endif() + # Apply YAML-driven quick/standard/comprehensive/full tier labels to the # tests above. Must run in THIS directory scope: CMake (pre-3.28, # which this project's cmake_minimum_required predates) only allows diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/streamk_iter_partition.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/streamk_iter_partition.cpp new file mode 100644 index 000000000000..392ef616ac65 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/streamk_iter_partition.cpp @@ -0,0 +1,150 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* + * tests/instances/streamk_iter_partition.cpp -- C parity for + * tests/instances/test_streamk_iter_partition.py. + * + * The iteration-balanced stream-K partition sizes the launch grid, which the + * byte-identity gate never sees, so a C/Python divergence here would surface + * as a silently wrong launch rather than an IR diff. The reference vectors are + * the Python test's, verbatim; the sweep checks the same invariants the fixup + * relies on. + */ +#include + +#include "rocke/helper_rocke.helpers.streamk.h" + +namespace +{ + +int g_failures = 0; + +#define CHECK(cond) \ + do \ + { \ + if(!(cond)) \ + { \ + fprintf(stderr, "FAIL: %s (%s:%d)\n", #cond, __FILE__, __LINE__); \ + ++g_failures; \ + } \ + } while(0) + +struct Vector +{ + int m, n, ipt, w; + bool persistent; + int sk_tiles, sk_ctas, ipsc, extra, dp_tiles, grid, partners, rounds, flags_bytes; +}; + +const Vector k_vectors[] = { + {4, 4, 8, 16, false, 0, 0, 0, 0, 16, 16, 0, 0, 0}, + {2, 3, 10, 8, false, 6, 8, 7, 4, 0, 8, 2, 2, 256}, + {5, 7, 12, 16, false, 19, 16, 14, 4, 16, 32, 1, 1, 256}, + {1, 3, 2, 8, false, 0, 0, 0, 0, 3, 3, 0, 0, 0}, + {5, 7, 12, 16, true, 19, 16, 14, 4, 16, 16, 1, 1, 256}, + {3, 3, 64, 4, false, 5, 4, 80, 0, 4, 8, 1, 1, 256}, + {7, 5, 3, 16, false, 19, 16, 3, 9, 16, 32, 1, 1, 256}, +}; + +rocke_streamk_iter_partition_t make(int m, int n, int ipt, int w, bool persistent) +{ + rocke_status_t st = ROCKE_OK; + rocke_streamk_iter_partition_t p + = rocke_streamk_iter_partition_make(m, n, ipt, w, persistent, &st); + CHECK(st == ROCKE_OK); + return p; +} + +void check_vectors() +{ + for(const Vector& v : k_vectors) + { + const rocke_streamk_iter_partition_t p = make(v.m, v.n, v.ipt, v.w, v.persistent); + const rocke_streamk_iter_plan_t r = rocke_streamk_iter_plan(&p); + CHECK(r.sk_tiles == v.sk_tiles); + CHECK(r.sk_ctas == v.sk_ctas); + CHECK(r.iters_per_sk_cta == v.ipsc); + CHECK(r.extra_iters == v.extra); + CHECK(r.dp_tiles == v.dp_tiles); + CHECK(r.grid_size == v.grid); + CHECK(r.max_linear_partners == v.partners); + CHECK(r.tree_rounds == v.rounds); + CHECK(r.flags_bytes == v.flags_bytes); + } + const rocke_streamk_iter_partition_t p = make(5, 7, 12, 16, false); + const int starts[] = {192, 207, 222, 237, 252}; + const int ends[] = {207, 222, 237, 252, 266}; + for(int c = 0; c < 5; ++c) + { + CHECK(rocke_streamk_start_iter(&p, c) == starts[c]); + CHECK(rocke_streamk_end_iter(&p, c) == ends[c]); + } +} + +void check_rejects_non_positive() +{ + rocke_status_t st = ROCKE_OK; + (void)rocke_streamk_iter_partition_make(0, 2, 4, 8, false, &st); + CHECK(st == ROCKE_ERR_VALUE); + st = ROCKE_OK; + (void)rocke_streamk_iter_partition_make(2, 2, 4, 0, false, &st); + CHECK(st == ROCKE_ERR_VALUE); +} + +void check_invariants(int m, int n, int ipt, int w) +{ + const rocke_streamk_iter_partition_t p = make(m, n, ipt, w, false); + const rocke_streamk_iter_plan_t r = rocke_streamk_iter_plan(&p); + CHECK(r.dp_tiles + r.sk_tiles == r.num_tiles); + if(!r.sk_ctas) + return; + CHECK(r.iters_per_sk_cta >= 1); + int cursor = r.total_dp_iters; + int prev_first_tile = -1; + int run = 0; /* contributors to the current tile */ + for(int c = 0; c < r.sk_ctas; ++c) + { + const int lo = rocke_streamk_start_iter(&p, c); + const int hi = rocke_streamk_end_iter(&p, c); + CHECK(lo == cursor); + CHECK(hi > lo); + cursor = hi; + const int first = lo / ipt; + const int last = (hi - 1) / ipt; + /* Contributors to `first` are this CTA plus the run carried from the + * previous CTA when it ended inside the same tile. */ + run = (first == prev_first_tile) ? run + 1 : 1; + CHECK(run - 1 <= r.max_linear_partners); + CHECK(run <= (1 << r.tree_rounds)); + if(last != first) + { + run = 1; + prev_first_tile = last; + } + else + prev_first_tile = first; + } + CHECK(cursor == r.total_dp_iters + r.total_sk_iters); +} + +} // namespace + +int main(void) +{ + check_vectors(); + check_rejects_non_positive(); + const int ipts[] = {1, 2, 3, 5, 8, 13}; + const int ws[] = {1, 2, 3, 4, 7, 8, 16}; + for(int m = 1; m < 7; ++m) + for(int n = 1; n < 5; ++n) + for(int ipt : ipts) + for(int w : ws) + check_invariants(m, n, ipt, w); + if(g_failures) + { + fprintf(stderr, "streamk_iter_partition: %d failure(s)\n", g_failures); + return 1; + } + printf("streamk_iter_partition: OK\n"); + return 0; +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_streamk_iter_partition.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_streamk_iter_partition.py new file mode 100644 index 000000000000..c345efa7afb5 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_streamk_iter_partition.py @@ -0,0 +1,135 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Host-side math of the iteration-balanced stream-K partitioner. + +The partition drives the launch grid, which lives outside the IR, so the +byte-identity gate cannot see a disagreement here. The reference vectors +below are mirrored verbatim by ``streamk_iter_partition.cpp`` for the C +engine; the invariant sweep checks the properties the fixup relies on. +""" + +from __future__ import annotations + +import unittest + +from rocke.helpers.streamk import ( + StreamKIterPartition, + streamk_end_iter, + streamk_iter_partition, + streamk_start_iter, +) + +# (m_tiles, n_tiles, iters_per_tile, max_active_wgs, persistent) -> +# (sk_tiles, sk_ctas, iters_per_sk_cta, extra_iters, dp_tiles, grid_size, +# max_linear_partners, tree_rounds, flags_bytes) +_VECTORS = ( + ((4, 4, 8, 16, False), (0, 0, 0, 0, 16, 16, 0, 0, 0)), # tiles % W == 0 + ((2, 3, 10, 8, False), (6, 8, 7, 4, 0, 8, 2, 2, 256)), # tiles < W + ((5, 7, 12, 16, False), (19, 16, 14, 4, 16, 32, 1, 1, 256)), # tiles > W + ((1, 3, 2, 8, False), (0, 0, 0, 0, 3, 3, 0, 0, 0)), # too little SK work + ((5, 7, 12, 16, True), (19, 16, 14, 4, 16, 16, 1, 1, 256)), # persistent + ((3, 3, 64, 4, False), (5, 4, 80, 0, 4, 8, 1, 1, 256)), + ((7, 5, 3, 16, False), (19, 16, 3, 9, 16, 32, 1, 1, 256)), +) + + +def _plan(p: StreamKIterPartition) -> tuple: + return ( + p.sk_tiles, + p.sk_ctas, + p.iters_per_sk_cta, + p.extra_iters, + p.dp_tiles, + p.grid_size, + p.max_linear_partners, + p.tree_rounds, + p.flags_bytes, + ) + + +def _make(m, n, ipt, w, persistent=False) -> StreamKIterPartition: + return streamk_iter_partition( + m_tiles=m, + n_tiles=n, + iters_per_tile=ipt, + max_active_wgs=w, + persistent=persistent, + ) + + +class TestStreamKIterPartitionVectors(unittest.TestCase): + def test_reference_vectors(self): + for args, want in _VECTORS: + with self.subTest(args=args): + self.assertEqual(_plan(_make(*args)), want) + + def test_start_end_iter_vectors(self): + p = _make(5, 7, 12, 16) + self.assertEqual( + [streamk_start_iter(p, c) for c in range(5)], [192, 207, 222, 237, 252] + ) + self.assertEqual( + [streamk_end_iter(p, c) for c in range(5)], [207, 222, 237, 252, 266] + ) + + def test_rejects_non_positive_counts(self): + for kw in ("m_tiles", "n_tiles", "iters_per_tile", "max_active_wgs"): + args = { + "m_tiles": 2, + "n_tiles": 2, + "iters_per_tile": 4, + "max_active_wgs": 8, + } + args[kw] = 0 + with ( + self.subTest(field=kw), + self.assertRaisesRegex(ValueError, f"{kw} must be > 0"), + ): + streamk_iter_partition(**args) + + +class TestStreamKIterPartitionInvariants(unittest.TestCase): + """Exhaustive sweep of small shapes against the fixup's assumptions.""" + + def _check(self, p: StreamKIterPartition) -> None: + ipt = p.iters_per_tile + self.assertEqual(p.dp_tiles + p.sk_tiles, p.num_tiles) + if not p.sk_ctas: + self.assertEqual(p.sk_tiles, 0) + return + # Every SK CTA owns >= 1 iteration and the ranges tile the SK span + # contiguously. + self.assertGreaterEqual(p.iters_per_sk_cta, 1) + cursor = p.total_dp_iters + contributors = {} + for c in range(p.sk_ctas): + lo, hi = streamk_start_iter(p, c), streamk_end_iter(p, c) + self.assertEqual(lo, cursor) + self.assertGreater(hi, lo) + cursor = hi + tiles = range(lo // ipt, (hi - 1) // ipt + 1) + # A CTA is a non-owner in at most its first tile: one partial slot + # per SK CTA is enough. + non_owner = [t for t in tiles if lo > t * ipt] + self.assertLessEqual(len(non_owner), 1) + for t in tiles: + contributors.setdefault(t, []).append(c) + self.assertEqual(cursor, p.total_dp_iters + p.total_sk_iters) + for t, ctas in contributors.items(): + # Contributors to one tile are consecutive SK CTAs. + self.assertEqual(ctas, list(range(ctas[0], ctas[0] + len(ctas)))) + self.assertLessEqual(len(ctas) - 1, p.max_linear_partners) + self.assertLessEqual(len(ctas), 1 << p.tree_rounds) + + def test_sweep(self): + for m in range(1, 7): + for n in range(1, 5): + for ipt in (1, 2, 3, 5, 8, 13): + for w in (1, 2, 3, 4, 7, 8, 16): + with self.subTest(m=m, n=n, ipt=ipt, w=w): + self._check(_make(m, n, ipt, w)) + + +if __name__ == "__main__": + unittest.main() From e52ebbddcb04e472ab027aaa84c323cda6bb51cc Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Mon, 5 Oct 2026 11:36:36 +0000 Subject: [PATCH 03/20] feat(rocke): add stream-K to implicit-GEMM conv wgrad Port CK Tile's stream-K to the backward-weight convolution in both engines. Split-K gives every output tile the same number of K slices, which leaves CUs idle on a wgrad: the output is filter-sized while the reduction N*Ho*Wo is long, so the tile count can be smaller than the device no matter how the slices are cut. Stream-K balances MAC iterations instead: tiles that divide the CTA pool run data-parallel with no reduction at all, and only the remainder is spread over the pool, each CTA owning a contiguous, balanced iteration range that may straddle tile boundaries. The whole partition is resolved on the host, since rocke compiles one kernel per shape, so the device only maps its linear CTA index onto an iteration range and walks it with one bounded loop. Spec: streamk ("off" | "dp_sk" | "persistent"), streamk_reduction ("linear" | "tree" | "atomic" | "workspace") and streamk_ctas. All default off or to the existing behaviour, so every current config keeps its kernel name and emits byte-identical IR. linear and tree are deterministic: a contributor publishes its partial tile, releases, and sets a per-wave flag; the consumer acquires that flag and folds in a fixed order, so dW needs no pre-zeroing. atomic and workspace reuse the existing split-K and two-stage epilogues. Conv groups fold into GEMM-M rather than riding block_id_z. Defining the tile count per group keeps the fold exact when the per-group M is not a multiple of the tile, which the upstream fold does not cover. This is the one path on which the C port accepts groups > 1; the grid-per-group path stays rejected there as before. ISSUE ID : AICK-1943 Co-Authored-By: Claude Opus 5 --- .../common/conv_implicit_gemm_wgrad.py | 544 ++++++++++++- .../platform/cpp/bindings/rocke_engine.cpp | 37 + .../instance_conv_implicit_gemm_internal.h | 7 + .../rocke/instance_conv_implicit_gemm_wgrad.h | 40 + .../common/conv_implicit_gemm_wgrad.cpp | 769 +++++++++++++++++- 5 files changed, 1361 insertions(+), 36 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py index 2130cd4015cd..c3ec86a23cfd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py @@ -477,6 +477,35 @@ class WgradConvSpec: # is a knob, not a constant: another arch or a different filter shape can # want a different contention-versus-footprint tradeoff. ws_replicas: int = _DEFAULT_WS_REPLICAS + # Stream-K execution (CK Tile StreamKTilePartitioner). "off" keeps the + # classic grid-per-tile kernel. "dp_sk" launches one CTA per data-parallel + # tile plus ``streamk_ctas`` stream-K CTAs, each owning a balanced, + # contiguous range of MAC iterations that may straddle output tiles. + # "persistent" launches exactly ``streamk_ctas`` CTAs that round-robin the + # data-parallel tiles and then run their stream-K range. + # + # Default "off": strictly additive, so every existing config emits + # byte-identical IR and keeps its kernel name. + streamk: str = "off" + # How the CTAs that share an output tile combine their partial K-sums. + # "linear" -- the tile owner waits on each later contributor's flag + # and folds its partial in order, then stores dW once. + # Deterministic; dW needs no pre-zeroing. + # "tree" -- pairwise fan-in with doubling stride; the owner stores. + # Deterministic; dW needs no pre-zeroing. + # "atomic" -- every contributor f32-atomic-adds into dW (dtype_d fp32 + # only). The caller zeroes dW before every launch. + # "workspace" -- every contributor f32-atomic-adds into the two-stage + # scratch; the Stage 2 reduce kernel casts it into dW. + # linear/tree hand partials over a flag table; their workspace is + # ``wgrad_streamk_workspace_nbytes`` and only its flag region has to be + # zeroed before each launch. + streamk_reduction: str = "linear" + # The CTA pool the stream-K remainder is spread over (CK max_active_wgs). + # -1 = one per CU of the target, resolved by the builder. Every stream-K + # CTA must be resident alongside the CTAs it waits on, so an explicit value + # must not exceed what the device can hold at once. + streamk_ctas: int = -1 @property def block_size(self) -> int: @@ -607,6 +636,14 @@ def kernel_name(self) -> str: # reads. Tracks the `twostage` flag above so the two move # together. f"wsr{self.ws_replicas}": self.two_stage and self.ws_replicas > 1, + # Stream-K changes the grid, the K ranges and the epilogue; the + # pool size and the reduction both change the emitted body. + f"sk{self.streamk_ctas}": self.streamk != "off" + and self.streamk_ctas > 0, + "skauto": self.streamk != "off" and self.streamk_ctas == -1, + "skpers": self.streamk == "persistent", + f"skr{self.streamk_reduction}": self.streamk != "off" + and self.streamk_reduction != "linear", }, ) @@ -642,6 +679,9 @@ def validate(self) -> None: f"ws_replicas must be >= 1 (got {self.ws_replicas}); it is the " f"number of scratch slabs a group's K-slices spread over" ) + _sk_ok, _sk_why = wgrad_streamk_available(self) + if not _sk_ok: + raise ValueError(_sk_why) # Delegate rather than re-derive -- is_valid_wgrad_spec calls the same # predicate, so the two cannot drift apart. _gm_ok, _gm_why = wgrad_group_merge_available(self) @@ -989,6 +1029,202 @@ def dw_addr(b_: IRBuilder, m_val: Value, n_val: Value): return dw_addr +# Stream-K is MFMA-only: the flag/fence fixup and the per-CTA partial slots +# are validated on these targets. +_STREAMK_ARCHES = ("gfx942", "gfx950") +_STREAMK_MODES = ("off", "dp_sk", "persistent") +_STREAMK_REDUCTIONS = ("linear", "tree", "atomic", "workspace") + + +def wgrad_streamk_available(spec: "WgradConvSpec") -> Tuple[bool, str]: + """Arch-independent stream-K gates, shared by ``validate()`` and + :func:`is_valid_wgrad_spec` so the two cannot drift apart. + + Stream-K owns the K range and the tile mapping, so every knob that also + partitions K, remaps tiles, or needs a compile-time K trip count is + rejected alongside it. + """ + if spec.streamk not in _STREAMK_MODES: + return False, ( + f"streamk must be 'off', 'dp_sk' or 'persistent' (got {spec.streamk!r})" + ) + if spec.streamk == "off": + return True, "ok" + if spec.streamk_reduction not in _STREAMK_REDUCTIONS: + return False, ( + "streamk_reduction must be 'linear', 'tree', 'atomic' or 'workspace' " + f"(got {spec.streamk_reduction!r})" + ) + if spec.streamk_ctas == 0 or spec.streamk_ctas < -1: + return False, f"streamk_ctas must be -1 (auto) or > 0 (got {spec.streamk_ctas})" + if spec.wave_size != 64: + return ( + False, + f"stream-K wgrad is MFMA wave64 only (got wave_size={spec.wave_size})", + ) + if spec.split_k != 1: + return False, ( + f"streamk and split_k both partition K_wg; use split_k=1 with streamk " + f"(got split_k={spec.split_k})" + ) + if spec.two_stage: + return False, ( + "two_stage is the split-K scratch path; use " + "streamk_reduction='workspace' for the stream-K equivalent" + ) + if spec.async_dma or spec.unroll_k: + return False, ( + "streamk needs a runtime K trip count per tile; async_dma/unroll_k " + "lay out a compile-time one" + ) + if spec.epilogue != "default": + return False, ( + f"streamk requires epilogue='default' (got {spec.epilogue!r}): the " + "epilogue runs once per tile inside the stream-K tile loop" + ) + if spec.chiplet_swizzle: + return False, ( + "chiplet_swizzle remaps the 2-D tile grid; stream-K decodes its " + "tiles from a linear CTA index instead" + ) + if spec.group_merge > 1: + return False, ( + f"streamk folds conv groups into GEMM-M; use group_merge=1 " + f"(got {spec.group_merge})" + ) + if spec.streamk_reduction == "atomic" and spec.data.dtype_d != "fp32": + return False, ( + f"streamk_reduction='atomic' needs dtype_d='fp32' (got " + f"{spec.data.dtype_d!r}); use 'linear', 'tree' or 'workspace'" + ) + if not spec.acc_epilogue.is_identity(): + return False, ( + "streamk supports only the identity acc_epilogue: partial K-sums " + "are combined after the per-tile accumulation" + ) + return True, "ok" + + +def wgrad_streamk_default_ctas(arch: str) -> int: + """Default stream-K CTA pool: one CTA per CU of ``arch``. + + One resident workgroup per CU is always available, which is what keeps a + stream-K CTA that waits on a later one from deadlocking. + """ + from rocke.helpers.split_k import _ARCH_NUM_CUS, _DEFAULT_NUM_CUS + + return _ARCH_NUM_CUS.get(arch, _DEFAULT_NUM_CUS) + + +def wgrad_streamk_partition(spec: "WgradConvSpec", *, arch: str = "gfx950"): + """The :class:`~rocke.helpers.streamk.StreamKIterPartition` of ``spec``. + + Output tiles are ``ceil(wg_M / tile_m) * groups`` by ``ceil(wg_N / + tile_n)``; the conv groups fold into GEMM-M per group, so a partial last + M tile of one group never spills into the next. + """ + from rocke.helpers.streamk import streamk_iter_partition + + if spec.streamk == "off": + raise ValueError("wgrad_streamk_partition requires streamk != 'off'") + p = spec.problem + ctas = ( + spec.streamk_ctas if spec.streamk_ctas > 0 else wgrad_streamk_default_ctas(arch) + ) + return streamk_iter_partition( + m_tiles=-(-_wg_M(p) // spec.tile_m) * p.groups, + n_tiles=-(-_wg_N(p) // spec.tile_n), + iters_per_tile=-(-_wg_K(p) // spec.tile_k), + max_active_wgs=ctas, + persistent=spec.streamk == "persistent", + ) + + +def wgrad_streamk_grid( + spec: "WgradConvSpec", *, arch: str = "gfx950" +) -> Tuple[int, int, int]: + """Launch grid of a stream-K wgrad kernel (one-dimensional).""" + return (wgrad_streamk_partition(spec, arch=arch).grid_size, 1, 1) + + +def wgrad_streamk_workspace_layout( + spec: "WgradConvSpec", *, arch: str = "gfx950" +) -> Tuple[int, int]: + """``(flags_bytes, partials_bytes)`` of a linear/tree stream-K kernel. + + One i32 flag per (stream-K CTA, wave), padded to 256 bytes, then one f32 + tile of partial sums per stream-K CTA. Pass the base as ``sk_flags`` and + ``base + flags_bytes`` as ``sk_partials``. Only the flag region has to be + zeroed before each launch; partials are always written before being read. + """ + if spec.streamk_reduction not in ("linear", "tree"): + return 0, 0 + part = wgrad_streamk_partition(spec, arch=arch) + flags = 4 * part.sk_ctas * spec.warp_m * spec.warp_n + flags_bytes = -(-flags // 256) * 256 + partials_bytes = part.sk_ctas * spec.tile_m * spec.tile_n * 4 + return flags_bytes, partials_bytes + + +def wgrad_streamk_workspace_nbytes( + spec: "WgradConvSpec", *, arch: str = "gfx950" +) -> int: + """Device workspace a stream-K wgrad launch needs, in bytes. + + linear/tree: the flag table plus the partial slots. workspace: the + two-stage f32 scratch (``groups * ws_replicas * wg_M * wg_N * 4``) that the + Stage 2 reduce kernel consumes. atomic: none. + """ + if spec.streamk == "off" or spec.streamk_reduction == "atomic": + return 0 + if spec.streamk_reduction == "workspace": + p = spec.problem + return p.groups * spec.ws_replicas * _wg_M(p) * _wg_N(p) * 4 + flags_bytes, partials_bytes = wgrad_streamk_workspace_layout(spec, arch=arch) + return flags_bytes + partials_bytes + + +def wgrad_streamk_signature(spec: "WgradConvSpec") -> list: + """Launch signature of a stream-K wgrad kernel. + + The standard conv ABI (A/B/D + byte sizes), plus ``ws_ptr``/``ws_bytes`` + for the workspace reduction or ``sk_flags``/``sk_partials`` for the + linear/tree fixups. Binding is positional, as for the two-stage kernel. + """ + _ir = {"fp16": "f16", "bf16": "bf16", "fp32": "f32"} + sig = [ + { + "name": "A", + "type": f"ptr<{_ir[spec.data.dtype_a]}, global>", + "size_bytes": 8, + }, + { + "name": "B", + "type": f"ptr<{_ir[spec.data.dtype_b]}, global>", + "size_bytes": 8, + }, + { + "name": "D", + "type": f"ptr<{_ir[spec.data.dtype_d]}, global>", + "size_bytes": 8, + }, + {"name": "A_bytes", "type": "i32", "size_bytes": 4}, + {"name": "B_bytes", "type": "i32", "size_bytes": 4}, + {"name": "D_bytes", "type": "i32", "size_bytes": 4}, + ] + if spec.streamk_reduction == "workspace": + sig += [ + {"name": "ws_ptr", "type": "ptr", "size_bytes": 8}, + {"name": "ws_bytes", "type": "i32", "size_bytes": 4}, + ] + elif spec.streamk_reduction in ("linear", "tree"): + sig += [ + {"name": "sk_flags", "type": "ptr", "size_bytes": 8}, + {"name": "sk_partials", "type": "ptr", "size_bytes": 8}, + ] + return sig + + def wgrad_group_merge_available( spec: WgradConvSpec, arch: str = "gfx950" ) -> Tuple[bool, str]: @@ -1095,6 +1331,14 @@ def is_valid_wgrad_spec(spec: WgradConvSpec, arch: str = "gfx950") -> Tuple[bool return False, ( f"spec wave_size {spec.wave_size} != {arch} wave_size {target.wave_size}" ) + if spec.streamk != "off" and arch not in _STREAMK_ARCHES: + return False, ( + f"stream-K wgrad supports only {', '.join(_STREAMK_ARCHES)} " + f"(got {arch})" + ) + _sk_ok, _sk_why = wgrad_streamk_available(spec) + if not _sk_ok: + return False, _sk_why sk = spec.split_k if sk < -1: @@ -1433,6 +1677,11 @@ def build_implicit_gemm_conv_wgrad( ) spec = _dc_replace(spec, split_k=decision.split_k) + # Resolve streamk_ctas=-1 (auto) the same way, so the kernel name and the + # partition both carry the concrete pool size. + if spec.streamk != "off" and spec.streamk_ctas == -1: + spec = dc_replace(spec, streamk_ctas=wgrad_streamk_default_ctas(arch)) + spec.validate() ok, why = is_valid_wgrad_spec(spec, arch=arch) if not ok: @@ -1457,6 +1706,10 @@ def build_implicit_gemm_conv_wgrad( _is_split_k = spec.split_k > 1 or spec.split_k == 0 _is_two_stage = spec.split_k > 1 and spec.two_stage + _is_streamk = spec.streamk != "off" + # Stream-K reductions that accumulate atomically instead of storing once. + _streamk_atomic = _is_streamk and spec.streamk_reduction in ("atomic", "workspace") + _streamk_flags = _is_streamk and spec.streamk_reduction in ("linear", "tree") _split_k_runtime = spec.split_k == 0 # ks passed as kernel arg at launch # At group_merge == groups the merged problem has a single group, but the # kernel still has to decode the K-slice off z and the epilogue still has to @@ -1482,7 +1735,9 @@ def build_implicit_gemm_conv_wgrad( # split_k>1 two_stage: dW is still the final output written by Stage 2; # Stage 1 (this kernel) never touches dW, but we keep it in the ABI so # the signature is identical between the atomic and two-stage variants. - _dw_writeonly = not _is_split_k + # Stream-K linear/tree: the tile owner stores dW once (writeonly); the + # atomic and workspace reductions read-modify-write their target. + _dw_writeonly = not _is_split_k and not _streamk_atomic dW = b.param( "dW", PtrType(ir_dtype_d, "global"), @@ -1499,11 +1754,18 @@ def build_implicit_gemm_conv_wgrad( # landing in its own region. Not present in the 16-bit atomic ABI. # # Not ``writeonly``: an atomicrmw reads its target. - if _is_two_stage: + if _is_two_stage or (_is_streamk and spec.streamk_reduction == "workspace"): ws_ptr = b.param("ws_ptr", PtrType(F32, "global"), noalias=True, align=16) _ws_bytes = b.param( "ws_bytes", I32 ) # noqa: F841 (side-effect: adds param to kernel signature) + # Stream-K linear/tree: per-wave i32 ready flags and per-CTA f32 partial + # slots. Neither is writeonly: both are read back by the tile owner. + if _streamk_flags: + sk_flags = b.param("sk_flags", PtrType(I32, "global"), noalias=True, align=16) + sk_partials = b.param( + "sk_partials", PtrType(F32, "global"), noalias=True, align=16 + ) # Runtime split-K: ks = slice width per CTA, computed and passed by the launcher. # Only present when split_k == 0; fixed-degree kernels bake ks as a constant. # ks_count = number of slices; only needed when grouped+runtime so the kernel @@ -1582,7 +1844,8 @@ def build_implicit_gemm_conv_wgrad( slice_v = None if grouped: c_kpg = b.const_i32(p_load.kpg) # kpg: dY output-channel slab stride - if not _is_split_k: + # Stream-K folds the group into its linear tile index instead. + if not _is_split_k and not _is_streamk: group_v = b.to_sgpr_u32(b.block_id_z()) # Split-K K-slice bounds. K_wg is padded so every slice is exactly ks @@ -2146,6 +2409,267 @@ def emit_mfma_phase( return new_accs + def _emit_streamk_tiles() -> None: + """Stream-K body: a loop over this CTA's work items. + + A work item is one output tile plus the slice of its MAC iterations this + CTA owns. A non-persistent CTA walks the tiles its iteration range + touches (exactly one whole tile for a data-parallel CTA); a persistent + CTA first round-robins the data-parallel tiles and then walks its + stream-K range. Every item runs the ordinary K loop over its slice and + then the configured fixup, so the GEMM body is emitted once. + + The tile mapping and group fold rebind the offsets the descriptor + closures read, exactly as the split-K path binds its K slice. + """ + nonlocal block_m_off_v, block_n_off_v, group_v, grid + from rocke.helpers.streamk import ( + emit_streamk_iter_range, + emit_streamk_sk_start_iter, + ) + + part = wgrad_streamk_partition(spec, arch=arch) + reduction = spec.streamk_reduction + c_one = b.const_i32(1) + c_ipt = b.const_i32(part.iters_per_tile) + cta = b.to_sgpr_u32(b.block_id_x()) + if part.persistent: + c_pool = b.const_i32(part.max_active_wgs) + # DP tiles cta, cta + pool, ... : ceil((dp_tiles - cta) / pool). + n_dp = b.div( + b.add( + b.sub(b.const_i32(part.dp_tiles), cta), + b.const_i32(part.max_active_wgs - 1), + ), + c_pool, + ) + if part.sk_ctas: + rng = emit_streamk_iter_range(b, cta, part, dp=False) + first_tile = b.div(rng.start, c_ipt) + last_tile = b.div(b.sub(rng.end, c_one), c_ipt) + n_sk = b.add(b.sub(last_tile, first_tile), c_one) + n_items = b.to_sgpr_u32(b.add(n_dp, n_sk)) + else: + n_items = b.to_sgpr_u32(n_dp) + sk_cta = cta + else: + rng = emit_streamk_iter_range(b, cta, part, dp=True) + first_tile = b.div(rng.start, c_ipt) + last_tile = b.div(b.sub(rng.end, c_one), c_ipt) + n_items = b.to_sgpr_u32(b.add(b.sub(last_tile, first_tile), c_one)) + sk_cta = rng.sk_cta + + n_waves = spec.warp_m * spec.warp_n + n_slots = len(accs) * c_per_lane + + def _flag_idx(sk_idx: Value) -> Value: + # One flag per (stream-K CTA, wave): a consumer lane only ever + # reads the partial its own wave's lane wrote, so each wave can + # release and acquire independently with no workgroup barrier. + return b.add(b.mul(sk_idx, b.const_i32(n_waves)), warp_id) + + def _slot_base(sk_idx: Value) -> Value: + # Partials are [sk_cta][slot][tid] f32: lanes of one slot are + # adjacent, so every partial store and load is fully coalesced. + return b.add(b.mul(sk_idx, b.const_i32(n_slots * threads)), tid) + + def _store_partial(sk_idx: Value, vals: Sequence[Value]) -> None: + base = _slot_base(sk_idx) + for ai, acc in enumerate(vals): + for i in range(c_per_lane): + off = b.add(base, b.const_i32((ai * c_per_lane + i) * threads)) + b.global_store(sk_partials, off, b.vec_extract(acc, i), align=4) + b.fence(scope="agent", ordering="release") + b.global_flag_store(sk_flags, _flag_idx(sk_idx), c_one) + + def _load_partial(sk_idx: Value) -> List[Value]: + base = _slot_base(sk_idx) + out = [] + for ai in range(len(accs)): + comps = [ + b.global_load_f32( + sk_partials, + b.add(base, b.const_i32((ai * c_per_lane + i) * threads)), + ) + for i in range(c_per_lane) + ] + out.append(b.vec_pack(comps, F32)) + return out + + def _sk_cta_of(it: Value) -> Value: + # The stream-K CTA owning global iteration ``it``: CTAs below + # extra_iters own iters_per_sk_cta + 1 iterations, the rest one fewer. + rel = b.sub(it, b.const_i32(part.total_dp_iters)) + wide = part.iters_per_sk_cta + 1 + boundary = b.const_i32(part.extra_iters * wide) + lead = b.div(rel, b.const_i32(wide)) + tail = b.add( + b.const_i32(part.extra_iters), + b.div(b.sub(rel, boundary), b.const_i32(part.iters_per_sk_cta)), + ) + return b.select(b.cmp_lt(rel, boundary), lead, tail) + + def _store_tile(vals: Sequence[Value]) -> None: + _emit_wgrad_direct_epilogue(b, spec, vals, grid, dw_rsrc, group=group_v) + + loop = b.scf_for(c0, n_items, c_one, iv_name="sk_item") + with loop as item: + if part.persistent and part.sk_ctas: + is_dp = b.cmp_lt(item, n_dp) + dp_tile = b.add(cta, b.mul(item, c_pool)) + sk_tile = b.add(first_tile, b.sub(item, n_dp)) + tile = b.to_sgpr_u32(b.select(is_dp, dp_tile, sk_tile)) + dp_start = b.mul(tile, c_ipt) + item_start = b.select(is_dp, dp_start, rng.start) + item_end = b.select(is_dp, b.add(dp_start, c_ipt), rng.end) + elif part.persistent: + tile = b.to_sgpr_u32(b.add(cta, b.mul(item, c_pool))) + item_start = b.mul(tile, c_ipt) + item_end = b.add(item_start, c_ipt) + else: + tile = b.to_sgpr_u32(b.add(first_tile, item)) + item_start = rng.start + item_end = rng.end + tile_start = b.mul(tile, c_ipt) + tile_end = b.add(tile_start, c_ipt) + lo = b.sub(b.smax(item_start, tile_start), tile_start) + hi = b.sub(b.smin(item_end, tile_end), tile_start) + + # Tile -> (group, m_tile, n_tile). Groups are folded into GEMM-M + # group-fastest, CK's i_g = tile_m % GemmBatch; m_tiles is defined + # per group, so the fold is exact for any wg_M. + c_n_tiles = b.const_i32(part.n_tiles) + n_tile = b.mod(tile, c_n_tiles) + m_idx = b.div(tile, c_n_tiles) + if p.groups > 1: + c_groups = b.const_i32(p.groups) + group_v = b.to_sgpr_u32(b.mod(m_idx, c_groups)) + m_tile = b.div(m_idx, c_groups) + else: + m_tile = m_idx + block_m_off_v = b.to_sgpr_u32(b.mul(m_tile, b.const_i32(block_m))) + block_n_off_v = b.to_sgpr_u32(b.mul(n_tile, b.const_i32(block_n))) + grid = dc_replace( + grid, block_m_off=block_m_off_v, block_n_off=block_n_off_v + ) + + k_lo_t = b.to_sgpr_u32(b.mul(lo, c_block_k)) + k_hi_t = b.to_sgpr_u32(b.mul(hi, c_block_k)) + for_op = b.scf_for_iter(k_lo_t, k_hi_t, c_block_k, accs, iv_name="k0") + with for_op as (k0, iter_vars): + emit_load_phase(k0, A_smem, B_smem) + b.sync() + new_accs = emit_mfma_phase(A_smem, B_smem, iter_vars) + b.sync() + b.scf_yield(*new_accs) + tile_accs = for_op.results + + if reduction == "atomic": + _emit_wgrad_split_k_epilogue( + b, + spec, + atom, + tile_accs, + warp_m_idx, + warp_n_idx, + lane, + block_m_off_v, + block_n_off_v, + dW, + c_per_lane, + group=group_v, + ) + elif reduction == "workspace": + _emit_wgrad_workspace_store_epilogue( + b, + spec, + atom, + tile_accs, + warp_m_idx, + warp_n_idx, + lane, + block_m_off_v, + block_n_off_v, + ws_ptr, + c_per_lane, + gm_group=group_v, + replica_sel=cta, + ) + elif not part.sk_ctas: + # All-DP partition: every item is a whole tile. + _store_tile(tile_accs) + elif reduction == "linear": + started = b.cmp_eq(lo, c0) + ended = b.cmp_eq(hi, c_ipt) + with b.scf_if_else(started) as (owner, contributor): + with owner: + # Fold every later contributor in order, then store once. + last = _sk_cta_of(b.sub(tile_end, c_one)) + n_partners = b.select(ended, c0, b.sub(last, sk_cta)) + fold = b.scf_for_iter( + c_one, + b.to_sgpr_u32(b.add(n_partners, c_one)), + c_one, + [(f"sk_{nm}", v) for (nm, _), v in zip(accs, tile_accs)], + iv_name="sk_p", + ) + with fold as (pi, folded): + partner = b.add(sk_cta, pi) + b.global_flag_wait_eq(sk_flags, _flag_idx(partner), c_one) + parts = _load_partial(partner) + b.scf_yield( + *[b.vector_add(v, q) for v, q in zip(folded, parts)] + ) + _store_tile(fold.results) + with contributor: + _store_partial(sk_cta, tile_accs) + else: + # Tree: round r pairs CTAs 2^r apart within the tile. A receiver + # folds its partner in; the first round in which a contributor + # is not a receiver it publishes its sum and drops out. Unrolled + # over the compile-time round bound with uniform predicates, so + # loop-carried values stay plain SSA. + started = b.cmp_eq(lo, c0) + ended = b.cmp_eq(hi, c_ipt) + not_started = b.lnot(started) + local = b.sub(sk_cta, _sk_cta_of(tile_start)) + zero_acc = b.zero_vec_f32(c_per_lane) + vals = list(tile_accs) + done = None + for r in range(part.tree_rounds): + stride = 1 << r + partner = b.add(sk_cta, b.const_i32(stride)) + in_tile = b.land( + b.lnot(ended), + b.cmp_lt( + emit_streamk_sk_start_iter(b, partner, part), tile_end + ), + ) + is_recv = b.cmp_eq(b.mod(local, b.const_i32(2 * stride)), c0) + recv = b.land(is_recv, in_tile) + if done is not None: + recv = b.land(recv, b.lnot(done)) + with b.scf_if(recv): + b.global_flag_wait_eq(sk_flags, _flag_idx(partner), c_one) + # Non-receivers read slot 0 (always in bounds) and discard it. + parts = _load_partial(b.select(recv, partner, c0)) + vals = [ + b.vector_add(v, b.select(recv, q, zero_acc)) + for v, q in zip(vals, parts) + ] + send = b.land(not_started, b.lnot(is_recv)) + if done is not None: + send = b.land(send, b.lnot(done)) + with b.scf_if(send): + _store_partial(sk_cta, vals) + done = send if done is None else b.lor(done, send) + with b.scf_if(started): + _store_tile(vals) + + if _is_streamk: + _emit_streamk_tiles() + return b.kernel + # ---- K loop ---- # k_lo / k_hi select the slice this CTA processes: # split_k == 1: k_lo=0, k_hi=None → full [0, wg_K) @@ -2871,6 +3395,7 @@ def _emit_wgrad_workspace_store_epilogue( ws_ptr: Value, c_per_lane: int, gm_group: Optional[Value] = None, + replica_sel: Optional[Value] = None, ) -> None: """Two-stage Stage 1 epilogue: f32 atomic-add into the scratch accumulator. @@ -2924,18 +3449,25 @@ def _emit_wgrad_workspace_store_epilogue( # does no folding, so an unconditional constant would renumber every # downstream SSA value and break byte-identity on the default path. # Scratch total size = groups * R * wg_M * wg_N (f32 elements). + # + # ``replica_sel`` replaces blockIdx.z as the hash input: a stream-K launch + # is one-dimensional, so it passes its linear CTA index instead. reps = spec.ws_replicas slab_elems = wg_M * wg_N _slab_idx = None + + def _replica_src() -> Value: + return replica_sel if replica_sel is not None else b.block_id_z() + if gm_group is not None and reps > 1: _slab_idx = b.add( b.mul(gm_group, b.const_i32(reps)), - b.mod(b.block_id_z(), b.const_i32(reps)), + b.mod(_replica_src(), b.const_i32(reps)), ) elif gm_group is not None: _slab_idx = gm_group elif reps > 1: - _slab_idx = b.mod(b.block_id_z(), b.const_i32(reps)) + _slab_idx = b.mod(_replica_src(), b.const_i32(reps)) slab_off = ( b.mul(_slab_idx, b.const_i32(slab_elems)) if _slab_idx is not None else None ) @@ -3015,7 +3547,7 @@ def _emit_wgrad_workspace_store_epilogue( if reps > 1: group_true = b.add( b.mul(group_true, c_reps_v), - b.mod(b.block_id_z(), c_reps_v), + b.mod(_replica_src(), c_reps_v), ) n_true = b.mul(yx, c_cpg_v) if p.cpg > 1 else yx if p.cpg > 1: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp index 720da7886a36..e2263708e866 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp @@ -1457,6 +1457,7 @@ void fill_conv_problem(rocke_conv_problem_t* p, const py::dict& d) p->pW = dict_int(d, "pW", p->pW); p->dH = dict_int(d, "dH", p->dH); p->dW = dict_int(d, "dW", p->dW); + p->groups = dict_int(d, "groups", p->groups); p->is_3d = dict_bool(d, "is_3d", p->is_3d); p->Di = dict_int(d, "Di", p->Di); p->Z = dict_int(d, "Z", p->Z); @@ -2171,10 +2172,22 @@ rocke_implicit_gemm_conv_wgrad_spec_t conv_wgrad_build_spec(const py::dict& d, s.split_k = dict_int(d, "split_k", s.split_k); s.two_stage = dict_bool(d, "two_stage", s.two_stage); s.ws_replicas = dict_int(d, "ws_replicas", s.ws_replicas); + s.streamk_ctas = dict_int(d, "streamk_ctas", s.streamk_ctas); + s.async_dma = dict_bool(d, "async_dma", s.async_dma); + s.unroll_k = dict_bool(d, "unroll_k", s.unroll_k); + s.lds_k_outer = dict_bool(d, "lds_k_outer", s.lds_k_outer); + s.chiplet_swizzle = dict_bool(d, "chiplet_swizzle", s.chiplet_swizzle); + s.chiplet_wgm = dict_int(d, "chiplet_wgm", s.chiplet_wgm); + s.chiplet_num_xcds = dict_int(d, "chiplet_num_xcds", s.chiplet_num_xcds); + s.chiplet_chunk_size = dict_int(d, "chiplet_chunk_size", s.chiplet_chunk_size); { std::string v; if(dict_str(d, "name", v)) s.name = keep(v); + if(dict_str(d, "streamk", v)) + s.streamk = keep(v); + if(dict_str(d, "streamk_reduction", v)) + s.streamk_reduction = keep(v); if(dict_str(d, "dtype_a", v)) s.dtype_a = keep(v); if(dict_str(d, "dtype_b", v)) @@ -3766,6 +3779,30 @@ PYBIND11_MODULE(rocke_engine, m) "Return workspace bytes for the two-stage wgrad path.\n" "Formula: groups * ws_replicas * wg_M * wg_N * 4 (always f32).\n" "Returns 0 when two_stage=false, or split_k <= 1."); + m.def( + "conv_wgrad_streamk_host", + [](const py::dict& d) -> py::dict { + std::deque store; + rocke_implicit_gemm_conv_wgrad_spec_t s = conv_wgrad_build_spec(d, store); + int grid[3] = {0, 0, 0}; + if(rocke_wgrad_conv_streamk_grid(&s, grid) != ROCKE_OK) + throw std::runtime_error( + "rocke_engine.conv_wgrad_streamk_host: needs streamk != 'off' and " + "streamk_ctas > 0"); + size_t flags_bytes = 0; + size_t partials_bytes = 0; + rocke_wgrad_conv_streamk_workspace_layout(&s, &flags_bytes, &partials_bytes); + py::dict out; + out["grid"] = py::make_tuple(grid[0], grid[1], grid[2]); + out["flags_bytes"] = flags_bytes; + out["partials_bytes"] = partials_bytes; + out["workspace_bytes"] = rocke_wgrad_conv_streamk_workspace_bytes(&s); + out["is_deterministic"] = rocke_wgrad_conv_spec_is_deterministic(&s); + return out; + }, + py::arg("spec"), + "Stream-K wgrad host plan of the C port: launch grid, workspace layout\n" + "and determinism, for cross-checking the Python wgrad_streamk_* helpers."); /* ---- attention families (separate TU; shared fmha/tiled struct tags) ---- */ register_attention(m); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_internal.h index bbd8be27bcb5..5345fc247d2b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_internal.h @@ -205,6 +205,13 @@ typedef struct rocke_conv_build_ctx rocke_value_t* group_idx; rocke_value_t* k_out_group_base; + /* ---- wgrad stream-K state (NULL off the stream-K path) ---- * + * wgrad_c_kpg = b.const_i32(kpg) (grouped dY k_out fold, Python c_kpg) + * wgrad_replica_sel = linear CTA index hashed into the workspace replica in + * place of block_id_z (Python replica_sel). */ + rocke_value_t* wgrad_c_kpg; + rocke_value_t* wgrad_replica_sel; + /* ---- LDS plan + smem handles ---- */ rocke_conv_lds_layout_t lds_layout; /* spec.effective_lds_layout() */ bool double_buffer; /* compv4 || async_dma || unroll_k */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h index de9e2b444197..db80315c6402 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h @@ -61,6 +61,7 @@ #define ROCKE_WGRAD_LDS_K_OUTER_ARCH "gfx950" #include +#include "rocke/helper_rocke.helpers.streamk.h" /* rocke_streamk_iter_partition_t */ #include "rocke/helper_rocke.instances.common.conv_implicit_gemm.h" /* rocke_conv_problem_t */ #include "rocke/ir.h" #include "rocke/lower_llvm.h" @@ -172,6 +173,19 @@ typedef struct rocke_implicit_gemm_conv_wgrad_spec * cut that R-fold and Stage 2 folds them back with a fixed unrolled add. * Scratch size is groups * R * wg_M * wg_N * 4 bytes. */ int ws_replicas; /* default 8 */ + + /* Stream-K execution (CK Tile StreamKTilePartitioner). Mirrors + * WgradConvSpec.streamk / streamk_reduction / streamk_ctas: + * streamk "off" | "dp_sk" | "persistent" + * streamk_reduction "linear" | "tree" | "atomic" | "workspace" + * streamk_ctas CTA pool the stream-K remainder is spread over; + * -1 (auto) is rejected by the builder, like split_k=-1. + * "off" keeps every existing config byte-identical. Stream-K folds the + * conv groups into GEMM-M, so it is the one path on which this port + * accepts groups > 1. */ + const char* streamk; /* default "off" */ + const char* streamk_reduction; /* default "linear" */ + int streamk_ctas; /* default -1 */ } rocke_implicit_gemm_conv_wgrad_spec_t; /* Default-constructed spec (every field == Python dataclass default). */ @@ -208,6 +222,10 @@ int rocke_wgrad_conv_spec_wg_N(const rocke_implicit_gemm_conv_wgrad_spec_t* s); * Stage 2's ordered fold over those slabs cannot un-reorder * sums that were already reordered inside one. * + * streamk "linear"/"tree" -> true. The owner of each tile folds the other + * contributors' partials in a fixed order and stores once. + * streamk "atomic"/"workspace" -> false, for the split-K reasons above. + * * So a deterministic wgrad is still available -- ask for split_k <= 1 -- but a * split-K wgrad, two-stage or not, is not one. Hosts that need bit-exactness * should gate on this predicate rather than on two_stage. */ @@ -219,6 +237,28 @@ bool rocke_wgrad_conv_spec_is_deterministic(const rocke_implicit_gemm_conv_wgrad * Analogous to rocke_streamk_gemm_workspace_bytes / rocke_moe_fused_workspace_bytes. */ size_t rocke_wgrad_conv_workspace_bytes(const rocke_implicit_gemm_conv_wgrad_spec_t* s); +/* ---- stream-K host-side helpers (Python wgrad_streamk_*) ---- */ + +/* Python wgrad_streamk_partition(spec): the iteration partition the kernel and + * its launch grid share. Requires streamk != "off" and streamk_ctas > 0; + * otherwise *out_status (if non-NULL) is ROCKE_ERR_VALUE and the result is + * zeroed. */ +rocke_streamk_iter_partition_t + rocke_wgrad_conv_streamk_partition(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + rocke_status_t* out_status); + +/* Python wgrad_streamk_grid(spec): one-dimensional launch grid. */ +rocke_status_t rocke_wgrad_conv_streamk_grid(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + int out[3]); + +/* Python wgrad_streamk_workspace_layout(spec): (flags_bytes, partials_bytes) of + * a linear/tree kernel; both 0 for the atomic and workspace reductions. */ +rocke_status_t rocke_wgrad_conv_streamk_workspace_layout( + const rocke_implicit_gemm_conv_wgrad_spec_t* s, size_t* flags_bytes, size_t* partials_bytes); + +/* Python wgrad_streamk_workspace_nbytes(spec). */ +size_t rocke_wgrad_conv_streamk_workspace_bytes(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); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp index 2fda3ff14fe6..b9e19d1ade90 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp @@ -73,9 +73,23 @@ rocke_implicit_gemm_conv_wgrad_spec_t rocke_implicit_gemm_conv_wgrad_spec_defaul s.chiplet_chunk_size = 64; s.split_k = 1; s.ws_replicas = ROCKE_WGRAD_DEFAULT_WS_REPLICAS; + s.streamk = "off"; + s.streamk_reduction = "linear"; + s.streamk_ctas = -1; return s; } +/* Python: spec.streamk != "off". A NULL field is the default. */ +static bool wgrad_streamk_on(const rocke_implicit_gemm_conv_wgrad_spec_t* s) +{ + return s->streamk != NULL && strcmp(s->streamk, "off") != 0; +} + +static const char* wgrad_streamk_reduction(const rocke_implicit_gemm_conv_wgrad_spec_t* s) +{ + return s->streamk_reduction ? s->streamk_reduction : "linear"; +} + int rocke_wgrad_conv_spec_block_size(const rocke_implicit_gemm_conv_wgrad_spec_t* s) { return s->warp_m * s->warp_n * s->wave_size; @@ -96,18 +110,18 @@ int rocke_wgrad_conv_spec_mfmas_per_warp_n(const rocke_implicit_gemm_conv_wgrad_ return s->tile_n / (s->warp_n * s->warp_tile_n); } -/* wg_M = K (output channels, groups=1 always for wgrad) */ +/* wg_M = kpg (output channels per group; K when groups == 1) */ int rocke_wgrad_conv_spec_wg_M(const rocke_implicit_gemm_conv_wgrad_spec_t* s) { - return s->problem.K; + return rocke_conv_problem_kpg(&s->problem); } -/* wg_N = Z*Y*X*C (filter spatial x input channel) */ +/* wg_N = Z*Y*X*cpg (filter spatial x input channels per group) */ int rocke_wgrad_conv_spec_wg_N(const rocke_implicit_gemm_conv_wgrad_spec_t* s) { const rocke_conv_problem_t* p = &s->problem; int z = p->is_3d ? p->Z : 1; - return z * p->Y * p->X * p->C; + return z * p->Y * p->X * rocke_conv_problem_cpg(p); } /* wg_K = N*Ho*Wo (output spatial positions) */ @@ -139,11 +153,95 @@ bool rocke_wgrad_conv_spec_is_deterministic(const rocke_implicit_gemm_conv_wgrad * argument and the epilogue is always packed atomics. Treating 0 as * "<= 1" here would tell a host that an atomic kernel produces * reproducible dW. */ + if(wgrad_streamk_on(s)) + { + const char* red = wgrad_streamk_reduction(s); + return strcmp(red, "linear") == 0 || strcmp(red, "tree") == 0; + } if(s->split_k == 0) return false; return s->split_k <= 1; } +/* Python wgrad_streamk_partition. */ +rocke_streamk_iter_partition_t + rocke_wgrad_conv_streamk_partition(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + rocke_status_t* out_status) +{ + rocke_streamk_iter_partition_t zero; + memset(&zero, 0, sizeof(zero)); + if(s == NULL || !wgrad_streamk_on(s) || s->streamk_ctas <= 0) + { + if(out_status) + *out_status = ROCKE_ERR_VALUE; + return zero; + } + const int groups = s->problem.groups > 0 ? s->problem.groups : 1; + const int wg_M = rocke_wgrad_conv_spec_wg_M(s); + const int wg_N = rocke_wgrad_conv_spec_wg_N(s); + const int wg_K = rocke_wgrad_conv_spec_wg_K(s); + return rocke_streamk_iter_partition_make((wg_M + s->tile_m - 1) / s->tile_m * groups, + (wg_N + s->tile_n - 1) / s->tile_n, + (wg_K + s->tile_k - 1) / s->tile_k, + s->streamk_ctas, + strcmp(s->streamk, "persistent") == 0, + out_status); +} + +rocke_status_t rocke_wgrad_conv_streamk_grid(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + int out[3]) +{ + rocke_status_t st = ROCKE_OK; + rocke_streamk_iter_partition_t part = rocke_wgrad_conv_streamk_partition(s, &st); + if(st != ROCKE_OK || out == NULL) + return ROCKE_ERR_VALUE; + out[0] = rocke_streamk_iter_plan(&part).grid_size; + out[1] = 1; + out[2] = 1; + return ROCKE_OK; +} + +rocke_status_t rocke_wgrad_conv_streamk_workspace_layout( + const rocke_implicit_gemm_conv_wgrad_spec_t* s, size_t* flags_bytes, size_t* partials_bytes) +{ + if(flags_bytes == NULL || partials_bytes == NULL) + return ROCKE_ERR_VALUE; + *flags_bytes = 0; + *partials_bytes = 0; + const char* red = wgrad_streamk_reduction(s); + if(strcmp(red, "linear") != 0 && strcmp(red, "tree") != 0) + return ROCKE_OK; + rocke_status_t st = ROCKE_OK; + rocke_streamk_iter_partition_t part = rocke_wgrad_conv_streamk_partition(s, &st); + if(st != ROCKE_OK) + return st; + const rocke_streamk_iter_plan_t plan = rocke_streamk_iter_plan(&part); + const size_t flags = (size_t)4 * (size_t)plan.sk_ctas * (size_t)(s->warp_m * s->warp_n); + *flags_bytes = (flags + 255) / 256 * 256; + *partials_bytes = (size_t)plan.sk_ctas * (size_t)s->tile_m * (size_t)s->tile_n * sizeof(float); + return ROCKE_OK; +} + +size_t rocke_wgrad_conv_streamk_workspace_bytes(const rocke_implicit_gemm_conv_wgrad_spec_t* s) +{ + if(s == NULL || !wgrad_streamk_on(s)) + return 0; + const char* red = wgrad_streamk_reduction(s); + if(strcmp(red, "atomic") == 0) + return 0; + if(strcmp(red, "workspace") == 0) + { + const int groups = s->problem.groups > 0 ? s->problem.groups : 1; + return (size_t)groups * (size_t)s->ws_replicas * (size_t)rocke_wgrad_conv_spec_wg_M(s) + * (size_t)rocke_wgrad_conv_spec_wg_N(s) * sizeof(float); + } + size_t flags_bytes = 0; + size_t partials_bytes = 0; + if(rocke_wgrad_conv_streamk_workspace_layout(s, &flags_bytes, &partials_bytes) != ROCKE_OK) + return 0; + return flags_bytes + partials_bytes; +} + size_t rocke_wgrad_conv_workspace_bytes(const rocke_implicit_gemm_conv_wgrad_spec_t* s) { if(!s->two_stage) @@ -218,8 +316,8 @@ rocke_status_t rocke_wgrad_conv_spec_kernel_name(const rocke_implicit_gemm_conv_ /* flags: async, kouter, pad{N}, spk{N}/spkauto/spkrt, twostage, wsr{N} */ char spk_flag[32] = {0}; char pad_flag[32] = {0}; - const char* flag_names[6]; - int flag_on[6]; + const char* flag_names[10]; + int flag_on[10]; int n_flags = 0; flag_names[n_flags] = "async"; @@ -290,6 +388,41 @@ rocke_status_t rocke_wgrad_conv_spec_kernel_name(const rocke_implicit_gemm_conv_ n_flags++; } + /* Stream-K tags, in the Python flag order: sk | skauto, skpers, + * skr (only when not the default "linear"). */ + char sk_buf[32]; + char skr_buf[48]; + if(wgrad_streamk_on(s)) + { + const char* red = wgrad_streamk_reduction(s); + if(s->streamk_ctas > 0) + { + snprintf(sk_buf, sizeof(sk_buf), "sk%d", s->streamk_ctas); + flag_names[n_flags] = sk_buf; + flag_on[n_flags] = 1; + n_flags++; + } + else if(s->streamk_ctas == -1) + { + flag_names[n_flags] = "skauto"; + flag_on[n_flags] = 1; + n_flags++; + } + if(strcmp(s->streamk, "persistent") == 0) + { + flag_names[n_flags] = "skpers"; + flag_on[n_flags] = 1; + n_flags++; + } + if(strcmp(red, "linear") != 0) + { + snprintf(skr_buf, sizeof(skr_buf), "skr%s", red); + flag_names[n_flags] = skr_buf; + flag_on[n_flags] = 1; + n_flags++; + } + } + return rocke_kernel_name_join( s->name, parts, 5, flag_names, flag_on, n_flags, out, out_cap, NULL); } @@ -298,6 +431,64 @@ rocke_status_t rocke_wgrad_conv_spec_kernel_name(const rocke_implicit_gemm_conv_ // is_valid_wgrad_spec // --------------------------------------------------------------------------- +/* Python is_valid_wgrad_spec's arch gate plus wgrad_streamk_available: the + * stream-K gates, with byte-identical reason strings. */ +static bool wgrad_streamk_gates(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + const char* arch, + char* reason, + size_t reason_cap) +{ +#define SK_REJECT(...) \ + do \ + { \ + if(reason && reason_cap) \ + snprintf(reason, reason_cap, __VA_ARGS__); \ + return false; \ + } while(0) + const char* mode = s->streamk ? s->streamk : "off"; + if(strcmp(mode, "off") != 0 && strcmp(arch, "gfx942") != 0 && strcmp(arch, "gfx950") != 0) + SK_REJECT("stream-K wgrad supports only gfx942, gfx950 (got %s)", arch); + if(strcmp(mode, "off") != 0 && strcmp(mode, "dp_sk") != 0 && strcmp(mode, "persistent") != 0) + SK_REJECT("streamk must be 'off', 'dp_sk' or 'persistent' (got '%s')", mode); + if(strcmp(mode, "off") == 0) + return true; + const char* red = wgrad_streamk_reduction(s); + if(strcmp(red, "linear") != 0 && strcmp(red, "tree") != 0 && strcmp(red, "atomic") != 0 + && strcmp(red, "workspace") != 0) + SK_REJECT("streamk_reduction must be 'linear', 'tree', 'atomic' or 'workspace' " + "(got '%s')", + red); + if(s->streamk_ctas == 0 || s->streamk_ctas < -1) + SK_REJECT("streamk_ctas must be -1 (auto) or > 0 (got %d)", s->streamk_ctas); + if(s->wave_size != 64) + SK_REJECT("stream-K wgrad is MFMA wave64 only (got wave_size=%d)", s->wave_size); + if(s->split_k != 1) + SK_REJECT("streamk and split_k both partition K_wg; use split_k=1 with streamk " + "(got split_k=%d)", + s->split_k); + if(s->two_stage) + SK_REJECT("two_stage is the split-K scratch path; use " + "streamk_reduction='workspace' for the stream-K equivalent"); + if(s->async_dma || s->unroll_k) + SK_REJECT("streamk needs a runtime K trip count per tile; async_dma/unroll_k " + "lay out a compile-time one"); + const char* epi = s->epilogue ? s->epilogue : "default"; + if(strcmp(epi, "default") != 0) + SK_REJECT("streamk requires epilogue='default' (got '%s'): the epilogue runs once " + "per tile inside the stream-K tile loop", + epi); + if(s->chiplet_swizzle) + SK_REJECT("chiplet_swizzle remaps the 2-D tile grid; stream-K decodes its tiles " + "from a linear CTA index instead"); + const char* dd = s->dtype_d ? s->dtype_d : "fp16"; + if(strcmp(red, "atomic") == 0 && strcmp(dd, "fp32") != 0) + SK_REJECT("streamk_reduction='atomic' needs dtype_d='fp32' (got '%s'); use " + "'linear', 'tree' or 'workspace'", + dd); + return true; +#undef SK_REJECT +} + bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv_wgrad_spec_t* s, const char* arch, char* reason, @@ -319,8 +510,10 @@ bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv if(arch == NULL) arch = "gfx950"; - /* groups > 1 is not supported for wgrad (Python validate() raises on it). */ - if(s->problem.groups != 1) + /* groups > 1 is only ported on the stream-K path, which folds the group + * into its linear tile index; the grid-per-group path stays Python-only. */ + const bool sk_on = wgrad_streamk_on(s); + if(s->problem.groups != 1 && !sk_on) { if(reason && reason_cap) snprintf(reason, @@ -329,6 +522,14 @@ bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv s->problem.groups); return false; } + if(!wgrad_streamk_gates(s, arch, reason, reason_cap)) + return false; + if(s->problem.groups > 1 && rocke_conv_problem_is_pointwise(&s->problem)) + { + if(reason && reason_cap) + snprintf(reason, reason_cap, "grouped pointwise (1x1) wgrad is not yet supported"); + return false; + } /* geometry */ if(s->tile_m % (s->warp_m * s->warp_tile_m)) @@ -823,6 +1024,17 @@ static rocke_tensor_descriptor_t* wgrad_make_dy_descriptor(rocke_ir_builder_t* b * Simple approach: build the DAG manually mirroring make_a_descriptor but * substituting "m"->"k_gemm" and "k"->"k_out" throughout. */ +/* Grouped X channel decode (Python _a_channel_decode, groups > 1): + * embed(["group", "c_in_group"] -> "c", strides=[cpg, 1], offset=0, lo=0, hi=C) + * so the group selects its [g*cpg, (g+1)*cpg) slab of the NHWC channel axis. */ +static const rocke_transform_t* wgrad_group_channel_embed(rocke_ir_builder_t* b, + const rocke_conv_problem_t* p) +{ + const char* up[2] = {"group", "c_in_group"}; + int strides[2] = {rocke_conv_problem_cpg(p), 1}; + return rocke_embed_bounded(b, up, 2, "c", strides, 0, 0, p->C); +} + static rocke_tensor_descriptor_t* wgrad_make_x_descriptor(rocke_ir_builder_t* b, const rocke_conv_problem_t* p) { @@ -884,13 +1096,21 @@ static rocke_tensor_descriptor_t* wgrad_make_x_descriptor(rocke_ir_builder_t* b, return NULL; n_x++; - /* unmerge_magic("k_out" -> [z,y,x,c]) */ - const char* into_k[4] = {"z", "y", "x", "c"}; - int dims_k[4] = {p->Z, p->Y, p->X, p->C}; + /* unmerge_magic("k_out" -> [z,y,x,c]); grouped: Python _a_channel_decode */ + const bool grouped3 = p->groups > 1; + const char* into_k[4] = {"z", "y", "x", grouped3 ? "c_in_group" : "c"}; + int dims_k[4] = {p->Z, p->Y, p->X, grouped3 ? rocke_conv_problem_cpg(p) : p->C}; xforms[n_x] = rocke_unmerge_magic(b, "k_out", into_k, 4, dims_k); if(xforms[n_x] == NULL) return NULL; n_x++; + if(grouped3) + { + xforms[n_x] = wgrad_group_channel_embed(b, p); + if(xforms[n_x] == NULL) + return NULL; + n_x++; + } xforms[n_x] = rocke_pad(b, "z", 0, p->Z); if(xforms[n_x] == NULL) @@ -940,13 +1160,21 @@ static rocke_tensor_descriptor_t* wgrad_make_x_descriptor(rocke_ir_builder_t* b, return NULL; n_x++; - /* unmerge_magic("k_out" -> [y, x, c]) */ - const char* into_k[3] = {"y", "x", "c"}; - int dims_k[3] = {p->Y, p->X, p->C}; + /* unmerge_magic("k_out" -> [y, x, c]); grouped: Python _a_channel_decode */ + const bool grouped = p->groups > 1; + const char* into_k[3] = {"y", "x", grouped ? "c_in_group" : "c"}; + int dims_k[3] = {p->Y, p->X, grouped ? rocke_conv_problem_cpg(p) : p->C}; xforms[n_x] = rocke_unmerge_magic(b, "k_out", into_k, 3, dims_k); if(xforms[n_x] == NULL) return NULL; n_x++; + if(grouped) + { + xforms[n_x] = wgrad_group_channel_embed(b, p); + if(xforms[n_x] == NULL) + return NULL; + n_x++; + } xforms[n_x] = rocke_pad(b, "y", 0, p->Y); if(xforms[n_x] == NULL) @@ -975,17 +1203,19 @@ static rocke_tensor_descriptor_t* wgrad_make_x_descriptor(rocke_ir_builder_t* b, static rocke_tensor_descriptor_t* wgrad_make_dw_descriptor(rocke_ir_builder_t* b, const rocke_conv_problem_t* p) { + /* Grouped dW is packed per group: [K, [Z,] Y, X, cpg] (Python cdim). */ + const int cdim = p->groups > 1 ? rocke_conv_problem_cpg(p) : p->C; if(p->is_3d) { /* naive coords: first dim "m" = K (output channels), rest are spatial */ - int lengths[5] = {p->K, p->Z, p->Y, p->X, p->C}; + int lengths[5] = {p->K, p->Z, p->Y, p->X, cdim}; const char* coords[5] = {"k_out", "z", "y", "x", "c"}; rocke_tensor_descriptor_t* desc = rocke_tensor_descriptor_naive(b, "dW_kzyxc", lengths, 5, NULL, coords, 5); if(desc == NULL) return NULL; const char* into[4] = {"z", "y", "x", "c"}; - int dims[4] = {p->Z, p->Y, p->X, p->C}; + int dims[4] = {p->Z, p->Y, p->X, cdim}; const rocke_transform_t* xforms[4]; int n_x = 0; xforms[n_x] = rocke_unmerge_magic(b, "n_wg", into, 4, dims); @@ -1008,14 +1238,14 @@ static rocke_tensor_descriptor_t* wgrad_make_dw_descriptor(rocke_ir_builder_t* b } /* 2-D: naive("dW_kyxc", [K,Y,X,C], coords=["k_out","y","x","c"]) */ - int lengths[4] = {p->K, p->Y, p->X, p->C}; + int lengths[4] = {p->K, p->Y, p->X, cdim}; const char* coords[4] = {"k_out", "y", "x", "c"}; rocke_tensor_descriptor_t* desc = rocke_tensor_descriptor_naive(b, "dW_kyxc", lengths, 4, NULL, coords, 4); if(desc == NULL) return NULL; const char* into[3] = {"y", "x", "c"}; - int dims[3] = {p->Y, p->X, p->C}; + int dims[3] = {p->Y, p->X, cdim}; const rocke_transform_t* xforms[4]; int n_x = 0; xforms[n_x] = rocke_unmerge_magic(b, "n_wg", into, 3, dims); @@ -1074,6 +1304,13 @@ static rocke_value_t* wgrad_dy_descriptor(rocke_ir_builder_t* b, rocke_conv_build_ctx_t* ctx = (rocke_conv_build_ctx_t*)ctx_user; /* k_out = block_m_off + row (= output channel, m_val) -- computed FIRST */ rocke_value_t* m_val = rocke_b_add(b, ctx->block_m_off_v, row); + if(ctx->wgrad_c_kpg != NULL) + { + /* Grouped: group g owns the dY output-channel slab [g*kpg, (g+1)*kpg). + * Python: k_out = b_.add(k_out, b_.mul(group_v, c_kpg)) */ + rocke_value_t* g_base = rocke_b_mul(b, ctx->group_idx, ctx->wgrad_c_kpg); + m_val = rocke_b_add(b, m_val, g_base); + } /* k_wg_red = k_off + col (= output position, k_val) -- computed SECOND */ rocke_value_t* k_val = rocke_b_add(b, ctx->k_off_capture, col); @@ -1140,13 +1377,15 @@ static rocke_value_t* wgrad_x_descriptor(rocke_ir_builder_t* b, return off; } - /* Full descriptor path: X_desc.offset(k_gemm=m_val, k_out=k_val) - * X descriptor top-level coords: "k_gemm" (K_wg reduction) and "k_out" (N_wg). */ - const char* names[2] = {"k_gemm", "k_out"}; - rocke_value_t* vals[2] = {m_val, k_val}; + /* Full descriptor path: X_desc.offset(k_gemm=m_val, k_out=k_val[, group=g]) + * X descriptor top-level coords: "k_gemm" (K_wg reduction) and "k_out" (N_wg); + * grouped adds "group", which selects the channel slab. */ + const char* names[3] = {"k_gemm", "k_out", "group"}; + rocke_value_t* vals[3] = {m_val, k_val, ctx->group_idx}; + const int n_coords = ctx->wgrad_c_kpg != NULL ? 3 : 2; rocke_value_t* off = NULL; rocke_value_t* valid = NULL; - rocke_transforms_descriptor_offset(b, ctx->B_desc, names, vals, 2, &off, &valid); + rocke_transforms_descriptor_offset(b, ctx->B_desc, names, vals, n_coords, &off, &valid); if(out_valid) *out_valid = valid; return off; @@ -1257,11 +1496,13 @@ static void wgrad_emit_workspace_store_epilogue(rocke_ir_builder_t* b, const int slab_elems = wg_M * wg_N; rocke_value_t* group_v = ctx->group_idx; rocke_value_t* slab_idx = NULL; + /* Python replica_sel: a stream-K launch is one-dimensional and hashes its + * linear CTA index into the replica instead of block_id_z. */ if(group_v != NULL && reps > 1) { rocke_value_t* c_reps_a = rocke_b_const_i32(b, reps); rocke_value_t* grp_term = rocke_b_mul(b, group_v, c_reps_a); - rocke_value_t* z = rocke_b_block_id_z(b); + rocke_value_t* z = ctx->wgrad_replica_sel ? ctx->wgrad_replica_sel : rocke_b_block_id_z(b); rocke_value_t* c_reps_b = rocke_b_const_i32(b, reps); rocke_value_t* rep_term = rocke_b_mod(b, z, c_reps_b); slab_idx = rocke_b_add(b, grp_term, rep_term); @@ -1272,7 +1513,7 @@ static void wgrad_emit_workspace_store_epilogue(rocke_ir_builder_t* b, } else if(reps > 1) { - rocke_value_t* z = rocke_b_block_id_z(b); + rocke_value_t* z = ctx->wgrad_replica_sel ? ctx->wgrad_replica_sel : rocke_b_block_id_z(b); rocke_value_t* c_reps = rocke_b_const_i32(b, reps); slab_idx = rocke_b_mod(b, z, c_reps); } @@ -1379,6 +1620,18 @@ static void wgrad_emit_workspace_store_epilogue(rocke_ir_builder_t* b, * so accumulation error grows with split_k. For large split_k consider using * an fp32 workspace and downcasting after the kernel completes. */ +/* Python _row_addr: c_m, or c_m + group*kpg on the grouped path. */ +static rocke_value_t* wgrad_row_addr(rocke_ir_builder_t* b, + rocke_value_t* c_m, + rocke_value_t* group, + rocke_value_t* c_kpg) +{ + if(group == NULL) + return c_m; + rocke_value_t* g_base = rocke_b_mul(b, group, c_kpg); + return rocke_b_add(b, c_m, g_base); +} + static void wgrad_emit_split_k_epilogue_f32(rocke_ir_builder_t* b, const rocke_conv_build_ctx_t* ctx, const rocke_implicit_gemm_conv_wgrad_spec_t* spec, @@ -1441,6 +1694,12 @@ static void wgrad_emit_split_k_epilogue_f32(rocke_ir_builder_t* b, /* Python creates wg_M_v / wg_N_v after the slot decode loop. */ rocke_value_t* wg_M_v = rocke_b_const_i32(b, wg_M); rocke_value_t* wg_N_v = rocke_b_const_i32(b, wg_N); + /* Grouped: Python _c_kpg = b.const_i32(p.kpg); the atomic row lands in the + * group's absolute output-channel slab (_row_addr) while the bounds stay on + * the per-group row. */ + rocke_value_t* c_kpg_e = ctx->group_idx != NULL + ? rocke_b_const_i32(b, rocke_conv_problem_kpg(&spec->problem)) + : NULL; int flat = 0; for(int mi = 0; mi < mfmas_m; ++mi) @@ -1467,7 +1726,8 @@ static void wgrad_emit_split_k_epilogue_f32(rocke_ir_builder_t* b, * c_off = add(mul(c_m, wg_N_v), c_n) * ok = land(cmp_lt(c_m, wg_M_v), cmp_lt(c_n, wg_N_v)) * scf_if(ok): global_atomic_add(ptr, c_off, val) */ - rocke_value_t* c_off = rocke_b_add(b, rocke_b_mul(b, c_m, wg_N_v), c_n); + rocke_value_t* row = wgrad_row_addr(b, c_m, ctx->group_idx, c_kpg_e); + rocke_value_t* c_off = rocke_b_add(b, rocke_b_mul(b, row, wg_N_v), c_n); rocke_value_t* m_ok = rocke_b_cmp_lt(b, c_m, wg_M_v); rocke_value_t* n_ok = rocke_b_cmp_lt(b, c_n, wg_N_v); rocke_value_t* ok = rocke_b_land(b, m_ok, n_ok); @@ -1507,8 +1767,9 @@ static void wgrad_emit_split_k_epilogue_f32(rocke_ir_builder_t* b, rocke_value_t* is_odd = rocke_b_cmp_ne(b, c_n_is_odd, rocke_b_const_i32(b, 0)); rocke_value_t* c_n_even = rocke_b_sub(b, c_n, c_n_is_odd); + rocke_value_t* row = wgrad_row_addr(b, c_m, ctx->group_idx, c_kpg_e); rocke_value_t* c_off_even - = rocke_b_add(b, rocke_b_mul(b, c_m, wg_N_v), c_n_even); + = rocke_b_add(b, rocke_b_mul(b, row, wg_N_v), c_n_even); rocke_value_t* v_even = rocke_b_select(b, is_odd, zero, val); rocke_value_t* v_odd = rocke_b_select(b, is_odd, val, zero); rocke_value_t* elems[2] = {v_even, v_odd}; @@ -1594,6 +1855,9 @@ static void wgrad_emit_split_k_cshuffle_epilogue(rocke_ir_builder_t* b, struct WgradDwAddrCtx { rocke_tensor_descriptor_t* dW_desc; + /* Grouped: the per-group row m_val gains group*kpg (Python dw_addr). */ + rocke_value_t* group; + rocke_value_t* c_kpg; }; static rocke_value_t* wgrad_dw_addr(rocke_ir_builder_t* b, @@ -1603,6 +1867,12 @@ static rocke_value_t* wgrad_dw_addr(rocke_ir_builder_t* b, void* user) { WgradDwAddrCtx* wc = static_cast(user); + if(wc->group != NULL) + { + /* Python: m_g = b_.add(m_val, b_.mul(group, _c_kpg)) */ + rocke_value_t* g_base = rocke_b_mul(b, wc->group, wc->c_kpg); + m_global = rocke_b_add(b, m_global, g_base); + } /* Python: dW_desc.offset(b_, k_out=m_val, n_wg=n_val) * dW_desc top-level coords: ("k_out" = output channel, "n_wg" = filter+chan). * The epilogue calls addr_fn(m_global=output_channel, n_global=filter+channel). */ @@ -1672,6 +1942,11 @@ static void wgrad_emit_direct_epilogue(rocke_ir_builder_t* b, { WgradDwAddrCtx addr_ctx; addr_ctx.dW_desc = dW_desc; + addr_ctx.group = ctx->group_idx; + /* Python emits _c_kpg before the bounds constants. */ + addr_ctx.c_kpg = ctx->group_idx != NULL + ? rocke_b_const_i32(b, rocke_conv_problem_kpg(&spec->problem)) + : NULL; rocke_value_t* bound_m = rocke_b_const_i32(b, wg_M); rocke_value_t* bound_n = rocke_b_const_i32(b, wg_N); rocke_direct_epilogue_store(b, @@ -1831,7 +2106,7 @@ static bool wgrad_build_ctx_init(rocke_conv_build_ctx_t* ctx, return false; } { - int wg_M_val = spec->problem.K; /* groups==1 */ + int wg_M_val = rocke_wgrad_conv_spec_wg_M(spec); /* kpg */ int wg_N_val = rocke_wgrad_conv_spec_wg_N(spec); /* N=wg_M, Hi=1, Wi=1, C=wg_K, K=wg_N, Y=1, X=1 -> Ho=Wo=1, all positive */ *stub_p = rocke_conv_problem_default(wg_M_val, 1, 1, wg_K, wg_N_val, 1, 1); @@ -2000,6 +2275,11 @@ static bool wgrad_build_ctx_init(rocke_conv_build_ctx_t* ctx, ctx->c_block_k = rocke_b_const_i32(b, ctx->block_k); /* c_wg_K: always created (occupies SSA slot even for split_k>1). */ rocke_value_t* c_wg_K = rocke_b_const_i32(b, wg_K); + /* Grouped (stream-K only in this port): Python + * if grouped: c_kpg = b.const_i32(p_load.kpg) + * The group index itself comes from the stream-K tile decode. */ + if(spec->problem.groups > 1) + ctx->wgrad_c_kpg = rocke_b_const_i32(b, rocke_conv_problem_kpg(&spec->problem)); rocke_value_t* k_lo; rocke_value_t* k_hi_v; /* NULL => loop runs to c_wg_K */ @@ -2310,6 +2590,393 @@ static bool wgrad_build_ctx_init(rocke_conv_build_ctx_t* ctx, return rocke_ir_builder_ok(b); } +// --------------------------------------------------------------------------- +// Stream-K body (Python _emit_streamk_tiles in build_implicit_gemm_conv_wgrad) +// --------------------------------------------------------------------------- +// +// Every helper below sequences each builder call into its own local: C++ leaves +// the evaluation order of call arguments unspecified and the builder numbers +// SSA values in call order, so the statement order here is Python's +// left-to-right argument evaluation, call for call. + +typedef struct wgrad_streamk_env +{ + rocke_conv_build_ctx_t* ctx; + const rocke_implicit_gemm_conv_wgrad_spec_t* spec; + rocke_streamk_iter_partition_t part; + rocke_streamk_iter_plan_t plan; + rocke_value_t* sk_flags; + rocke_value_t* sk_partials; + rocke_value_t* c_one; + int n_waves; + int n_slots; + int threads; +} wgrad_streamk_env_t; + +/* _flag_idx: add(mul(sk_idx, const(n_waves)), warp_id) */ +static rocke_value_t* wgrad_sk_flag_idx(wgrad_streamk_env_t* e, rocke_value_t* sk_idx) +{ + rocke_ir_builder_t* b = e->ctx->b; + rocke_value_t* c_waves = rocke_b_const_i32(b, e->n_waves); + rocke_value_t* scaled = rocke_b_mul(b, sk_idx, c_waves); + return rocke_b_add(b, scaled, e->ctx->warp_id); +} + +/* _slot_base: add(mul(sk_idx, const(n_slots * threads)), tid) */ +static rocke_value_t* wgrad_sk_slot_base(wgrad_streamk_env_t* e, rocke_value_t* sk_idx) +{ + rocke_ir_builder_t* b = e->ctx->b; + rocke_value_t* c_stride = rocke_b_const_i32(b, e->n_slots * e->threads); + rocke_value_t* scaled = rocke_b_mul(b, sk_idx, c_stride); + return rocke_b_add(b, scaled, e->ctx->tid); +} + +/* _store_partial: coalesced scalar stores, release fence, per-wave flag. */ +static void wgrad_sk_store_partial(wgrad_streamk_env_t* e, + rocke_value_t* sk_idx, + rocke_value_t* const* vals) +{ + rocke_ir_builder_t* b = e->ctx->b; + const int c_per_lane = e->ctx->c_per_lane; + rocke_value_t* base = wgrad_sk_slot_base(e, sk_idx); + for(int ai = 0; ai < e->ctx->num_accs; ++ai) + { + for(int i = 0; i < c_per_lane; ++i) + { + rocke_value_t* c_slot = rocke_b_const_i32(b, (ai * c_per_lane + i) * e->threads); + rocke_value_t* off = rocke_b_add(b, base, c_slot); + rocke_value_t* v = rocke_b_vec_extract(b, vals[ai], i); + rocke_b_global_store(b, e->sk_partials, off, v, 4); + } + } + rocke_b_fence(b, "agent", "release"); + rocke_value_t* flag = wgrad_sk_flag_idx(e, sk_idx); + rocke_b_global_flag_store(b, e->sk_flags, flag, e->c_one, "agent", "release"); +} + +/* _load_partial: one per accumulator. */ +static void + wgrad_sk_load_partial(wgrad_streamk_env_t* e, rocke_value_t* sk_idx, rocke_value_t** out) +{ + rocke_ir_builder_t* b = e->ctx->b; + const int c_per_lane = e->ctx->c_per_lane; + rocke_value_t* base = wgrad_sk_slot_base(e, sk_idx); + rocke_value_t* comps[64]; + for(int ai = 0; ai < e->ctx->num_accs; ++ai) + { + for(int i = 0; i < c_per_lane; ++i) + { + rocke_value_t* c_slot = rocke_b_const_i32(b, (ai * c_per_lane + i) * e->threads); + rocke_value_t* off = rocke_b_add(b, base, c_slot); + comps[i] = rocke_b_global_load_f32(b, e->sk_partials, off, 4); + } + out[ai] = rocke_b_vec_pack(b, comps, c_per_lane, rocke_f32()); + } +} + +/* _sk_cta_of: the stream-K CTA owning global iteration `it`. */ +static rocke_value_t* wgrad_sk_cta_of(wgrad_streamk_env_t* e, rocke_value_t* it) +{ + rocke_ir_builder_t* b = e->ctx->b; + const rocke_streamk_iter_plan_t* r = &e->plan; + const int wide = r->iters_per_sk_cta + 1; + rocke_value_t* c_dp_iters = rocke_b_const_i32(b, r->total_dp_iters); + rocke_value_t* rel = rocke_b_sub(b, it, c_dp_iters); + rocke_value_t* boundary = rocke_b_const_i32(b, r->extra_iters * wide); + rocke_value_t* c_wide = rocke_b_const_i32(b, wide); + rocke_value_t* lead = rocke_b_div(b, rel, c_wide); + rocke_value_t* c_extra = rocke_b_const_i32(b, r->extra_iters); + rocke_value_t* past = rocke_b_sub(b, rel, boundary); + rocke_value_t* c_per = rocke_b_const_i32(b, r->iters_per_sk_cta); + rocke_value_t* past_ctas = rocke_b_div(b, past, c_per); + rocke_value_t* tail = rocke_b_add(b, c_extra, past_ctas); + rocke_value_t* in_lead = rocke_b_cmp_lt(b, rel, boundary); + return rocke_b_select(b, in_lead, lead, tail); +} + +/* _store_tile: the plain dW store of a finished tile. */ +static void wgrad_sk_store_tile(wgrad_streamk_env_t* e, + rocke_value_t* const* vals, + rocke_tensor_descriptor_t* dW_desc, + int wg_M, + int wg_N) +{ + rocke_conv_build_ctx_t* ctx = e->ctx; + for(int i = 0; i < ctx->num_accs; ++i) + ctx->final_accs[i] = vals[i]; + ctx->num_final_accs = ctx->num_accs; + wgrad_emit_direct_epilogue(ctx->b, ctx, e->spec, dW_desc, ctx->d_rsrc, wg_M, wg_N); +} + +static void wgrad_emit_streamk(rocke_conv_build_ctx_t* ctx, + const rocke_implicit_gemm_conv_wgrad_spec_t* spec, + rocke_value_t* dW, + rocke_tensor_descriptor_t* dW_desc, + rocke_value_t* ws_ptr, + rocke_value_t* sk_flags, + rocke_value_t* sk_partials, + int wg_M, + int wg_N) +{ + rocke_ir_builder_t* b = ctx->b; + wgrad_streamk_env_t e; + rocke_status_t st = ROCKE_OK; + memset(&e, 0, sizeof(e)); + e.ctx = ctx; + e.spec = spec; + e.part = rocke_wgrad_conv_streamk_partition(spec, &st); + if(st != ROCKE_OK) + { + rocke_i_set_err(b, ROCKE_ERR_VALUE, "wgrad: invalid stream-K partition"); + return; + } + e.plan = rocke_streamk_iter_plan(&e.part); + e.sk_flags = sk_flags; + e.sk_partials = sk_partials; + e.n_waves = spec->warp_m * spec->warp_n; + e.n_slots = ctx->num_accs * ctx->c_per_lane; + e.threads = ctx->threads; + if(ctx->c_per_lane > 64) + { + rocke_i_set_err(b, ROCKE_ERR_VALUE, "wgrad: stream-K partial too wide"); + return; + } + const char* reduction = wgrad_streamk_reduction(spec); + const rocke_streamk_iter_plan_t* r = &e.plan; + const bool persistent = e.part.persistent; + const int groups = spec->problem.groups > 0 ? spec->problem.groups : 1; + + e.c_one = rocke_b_const_i32(b, 1); + rocke_value_t* c_one = e.c_one; + rocke_value_t* c_ipt = rocke_b_const_i32(b, e.part.iters_per_tile); + rocke_value_t* cta = rocke_b_to_sgpr_u32(b, rocke_b_block_id_x(b)); + rocke_value_t* c_pool = NULL; + rocke_value_t* n_dp = NULL; + rocke_value_t* first_tile = NULL; + rocke_value_t* n_items = NULL; + rocke_value_t* sk_cta = NULL; + rocke_streamk_iter_range_t rng; + memset(&rng, 0, sizeof(rng)); + if(persistent) + { + c_pool = rocke_b_const_i32(b, e.part.max_active_wgs); + rocke_value_t* c_dp_tiles = rocke_b_const_i32(b, r->dp_tiles); + rocke_value_t* left = rocke_b_sub(b, c_dp_tiles, cta); + rocke_value_t* c_round = rocke_b_const_i32(b, e.part.max_active_wgs - 1); + rocke_value_t* rounded = rocke_b_add(b, left, c_round); + n_dp = rocke_b_div(b, rounded, c_pool); + if(r->sk_ctas) + { + rng = rocke_emit_streamk_iter_range(b, cta, &e.part, false); + first_tile = rocke_b_div(b, rng.start, c_ipt); + rocke_value_t* last_it = rocke_b_sub(b, rng.end, c_one); + rocke_value_t* last_tile = rocke_b_div(b, last_it, c_ipt); + rocke_value_t* span = rocke_b_sub(b, last_tile, first_tile); + rocke_value_t* n_sk = rocke_b_add(b, span, c_one); + n_items = rocke_b_to_sgpr_u32(b, rocke_b_add(b, n_dp, n_sk)); + } + else + n_items = rocke_b_to_sgpr_u32(b, n_dp); + sk_cta = cta; + } + else + { + rng = rocke_emit_streamk_iter_range(b, cta, &e.part, true); + first_tile = rocke_b_div(b, rng.start, c_ipt); + rocke_value_t* last_it = rocke_b_sub(b, rng.end, c_one); + rocke_value_t* last_tile = rocke_b_div(b, last_it, c_ipt); + rocke_value_t* span = rocke_b_sub(b, last_tile, first_tile); + n_items = rocke_b_to_sgpr_u32(b, rocke_b_add(b, span, c_one)); + sk_cta = rng.sk_cta; + } + + rocke_value_t* c0 = ctx->c0; /* k_lo of the classic path == const 0 */ + rocke_for_t loop = rocke_b_scf_for(b, c0, n_items, c_one, "sk_item"); + rocke_b_region_enter(b, loop.body); + { + rocke_value_t* item = loop.iv; + rocke_value_t* tile; + rocke_value_t* item_start; + rocke_value_t* item_end; + if(persistent && r->sk_ctas) + { + rocke_value_t* is_dp = rocke_b_cmp_lt(b, item, n_dp); + rocke_value_t* stride = rocke_b_mul(b, item, c_pool); + rocke_value_t* dp_tile = rocke_b_add(b, cta, stride); + rocke_value_t* sk_off = rocke_b_sub(b, item, n_dp); + rocke_value_t* sk_tile = rocke_b_add(b, first_tile, sk_off); + tile = rocke_b_to_sgpr_u32(b, rocke_b_select(b, is_dp, dp_tile, sk_tile)); + rocke_value_t* dp_start = rocke_b_mul(b, tile, c_ipt); + item_start = rocke_b_select(b, is_dp, dp_start, rng.start); + rocke_value_t* dp_end = rocke_b_add(b, dp_start, c_ipt); + item_end = rocke_b_select(b, is_dp, dp_end, rng.end); + } + else if(persistent) + { + rocke_value_t* stride = rocke_b_mul(b, item, c_pool); + tile = rocke_b_to_sgpr_u32(b, rocke_b_add(b, cta, stride)); + item_start = rocke_b_mul(b, tile, c_ipt); + item_end = rocke_b_add(b, item_start, c_ipt); + } + else + { + tile = rocke_b_to_sgpr_u32(b, rocke_b_add(b, first_tile, item)); + item_start = rng.start; + item_end = rng.end; + } + rocke_value_t* tile_start = rocke_b_mul(b, tile, c_ipt); + rocke_value_t* tile_end = rocke_b_add(b, tile_start, c_ipt); + rocke_value_t* lo_max = rocke_b_smax(b, item_start, tile_start); + rocke_value_t* lo = rocke_b_sub(b, lo_max, tile_start); + rocke_value_t* hi_min = rocke_b_smin(b, item_end, tile_end); + rocke_value_t* hi = rocke_b_sub(b, hi_min, tile_start); + + /* Tile -> (group, m_tile, n_tile), groups folded into GEMM-M. */ + rocke_value_t* c_n_tiles = rocke_b_const_i32(b, e.part.n_tiles); + rocke_value_t* n_tile = rocke_b_mod(b, tile, c_n_tiles); + rocke_value_t* m_idx = rocke_b_div(b, tile, c_n_tiles); + rocke_value_t* m_tile = m_idx; + if(groups > 1) + { + rocke_value_t* c_groups = rocke_b_const_i32(b, groups); + ctx->group_idx = rocke_b_to_sgpr_u32(b, rocke_b_mod(b, m_idx, c_groups)); + m_tile = rocke_b_div(b, m_idx, c_groups); + } + rocke_value_t* c_bm = rocke_b_const_i32(b, ctx->block_m); + ctx->block_m_off_v = rocke_b_to_sgpr_u32(b, rocke_b_mul(b, m_tile, c_bm)); + rocke_value_t* c_bn = rocke_b_const_i32(b, ctx->block_n); + ctx->block_n_off_v = rocke_b_to_sgpr_u32(b, rocke_b_mul(b, n_tile, c_bn)); + ctx->grid.block_m_off = ctx->block_m_off_v; + ctx->grid.block_n_off = ctx->block_n_off_v; + + ctx->c0 = rocke_b_to_sgpr_u32(b, rocke_b_mul(b, lo, ctx->c_block_k)); + ctx->c_K_gemm = rocke_b_to_sgpr_u32(b, rocke_b_mul(b, hi, ctx->c_block_k)); + rocke_conv_emit_kloop_simple(ctx); + rocke_value_t* tile_accs[ROCKE_CONV_MAX_ACCS]; + for(int i = 0; i < ctx->num_accs; ++i) + tile_accs[i] = ctx->final_accs[i]; + + if(strcmp(reduction, "atomic") == 0) + { + wgrad_emit_split_k_epilogue_f32(b, ctx, spec, dW, wg_M, wg_N); + } + else if(strcmp(reduction, "workspace") == 0) + { + ctx->wgrad_replica_sel = cta; + wgrad_emit_workspace_store_epilogue(b, ctx, spec, ws_ptr, wg_M, wg_N); + } + else if(!r->sk_ctas) + { + wgrad_sk_store_tile(&e, tile_accs, dW_desc, wg_M, wg_N); + } + else if(strcmp(reduction, "linear") == 0) + { + rocke_value_t* started = rocke_b_cmp_eq(b, lo, c0); + rocke_value_t* ended = rocke_b_cmp_eq(b, hi, c_ipt); + rocke_if_else_t ie = rocke_b_scf_if_else(b, started); + rocke_b_region_enter(b, ie.then_region); + { + rocke_value_t* last_it = rocke_b_sub(b, tile_end, c_one); + rocke_value_t* last = wgrad_sk_cta_of(&e, last_it); + rocke_value_t* span = rocke_b_sub(b, last, sk_cta); + rocke_value_t* n_partners = rocke_b_select(b, ended, c0, span); + rocke_iter_arg_t iter_args[ROCKE_CONV_MAX_ACCS]; + for(int i = 0; i < ctx->num_accs; ++i) + { + iter_args[i].name = rocke_arena_printf(&b->arena, "sk_%s", ctx->acc_names[i]); + iter_args[i].init = tile_accs[i]; + } + rocke_value_t* fold_hi = rocke_b_to_sgpr_u32(b, rocke_b_add(b, n_partners, c_one)); + rocke_for_t fold = rocke_b_scf_for_iter( + b, c_one, fold_hi, c_one, iter_args, ctx->num_accs, "sk_p", false, true); + rocke_b_region_enter(b, fold.body); + { + rocke_value_t* partner = rocke_b_add(b, sk_cta, fold.iv); + rocke_value_t* flag = wgrad_sk_flag_idx(&e, partner); + rocke_b_global_flag_wait_eq(b, sk_flags, flag, c_one, "agent"); + rocke_value_t* parts[ROCKE_CONV_MAX_ACCS]; + wgrad_sk_load_partial(&e, partner, parts); + rocke_value_t* sums[ROCKE_CONV_MAX_ACCS]; + for(int i = 0; i < ctx->num_accs; ++i) + sums[i] = rocke_b_vector_add(b, fold.iter_vars[i], parts[i]); + rocke_b_scf_yield(b, sums, ctx->num_accs); + } + rocke_b_region_leave(b); + wgrad_sk_store_tile(&e, fold.op->results, dW_desc, wg_M, wg_N); + } + rocke_b_region_leave(b); + rocke_b_region_enter(b, ie.else_region); + wgrad_sk_store_partial(&e, sk_cta, tile_accs); + rocke_b_region_leave(b); + } + else + { + /* Tree: unrolled pairwise fan-in (see the Python comment). */ + rocke_value_t* started = rocke_b_cmp_eq(b, lo, c0); + rocke_value_t* ended = rocke_b_cmp_eq(b, hi, c_ipt); + rocke_value_t* not_started = rocke_b_lnot(b, started); + rocke_value_t* owner = wgrad_sk_cta_of(&e, tile_start); + rocke_value_t* local = rocke_b_sub(b, sk_cta, owner); + rocke_value_t* zero_acc = rocke_b_zero_vec_f32(b, ctx->c_per_lane); + rocke_value_t* vals[ROCKE_CONV_MAX_ACCS]; + for(int i = 0; i < ctx->num_accs; ++i) + vals[i] = tile_accs[i]; + rocke_value_t* done = NULL; + for(int rr = 0; rr < r->tree_rounds; ++rr) + { + const int stride = 1 << rr; + rocke_value_t* c_stride = rocke_b_const_i32(b, stride); + rocke_value_t* partner = rocke_b_add(b, sk_cta, c_stride); + rocke_value_t* not_ended = rocke_b_lnot(b, ended); + rocke_value_t* p_start = rocke_emit_streamk_sk_start_iter(b, partner, &e.part); + rocke_value_t* p_before_end = rocke_b_cmp_lt(b, p_start, tile_end); + rocke_value_t* in_tile = rocke_b_land(b, not_ended, p_before_end); + rocke_value_t* c_pair = rocke_b_const_i32(b, 2 * stride); + rocke_value_t* rem = rocke_b_mod(b, local, c_pair); + rocke_value_t* is_recv = rocke_b_cmp_eq(b, rem, c0); + rocke_value_t* recv = rocke_b_land(b, is_recv, in_tile); + if(done != NULL) + { + rocke_value_t* not_done = rocke_b_lnot(b, done); + recv = rocke_b_land(b, recv, not_done); + } + rocke_if_t wait_if = rocke_b_scf_if(b, recv); + rocke_b_region_enter(b, wait_if.then_region); + { + rocke_value_t* flag = wgrad_sk_flag_idx(&e, partner); + rocke_b_global_flag_wait_eq(b, sk_flags, flag, c_one, "agent"); + } + rocke_b_region_leave(b); + rocke_value_t* src = rocke_b_select(b, recv, partner, c0); + rocke_value_t* parts[ROCKE_CONV_MAX_ACCS]; + wgrad_sk_load_partial(&e, src, parts); + for(int i = 0; i < ctx->num_accs; ++i) + { + rocke_value_t* q = rocke_b_select(b, recv, parts[i], zero_acc); + vals[i] = rocke_b_vector_add(b, vals[i], q); + } + rocke_value_t* not_recv = rocke_b_lnot(b, is_recv); + rocke_value_t* send = rocke_b_land(b, not_started, not_recv); + if(done != NULL) + { + rocke_value_t* not_done = rocke_b_lnot(b, done); + send = rocke_b_land(b, send, not_done); + } + rocke_if_t send_if = rocke_b_scf_if(b, send); + rocke_b_region_enter(b, send_if.then_region); + wgrad_sk_store_partial(&e, sk_cta, vals); + rocke_b_region_leave(b); + done = (done == NULL) ? send : rocke_b_lor(b, done, send); + } + rocke_if_t store_if = rocke_b_scf_if(b, started); + rocke_b_region_enter(b, store_if.then_region); + wgrad_sk_store_tile(&e, vals, dW_desc, wg_M, wg_N); + rocke_b_region_leave(b); + } + } + rocke_b_region_leave(b); +} + // --------------------------------------------------------------------------- // rocke_build_implicit_gemm_conv_wgrad // --------------------------------------------------------------------------- @@ -2346,6 +3013,23 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv_wgrad( return NULL; } + /* streamk_ctas=-1 (auto) is rejected for the same reason: Python resolves + * it from the arch CU table, which this port does not carry. */ + const bool is_streamk = wgrad_streamk_on(spec); + if(is_streamk && spec->streamk_ctas == -1) + { + rocke_i_set_err(b, + ROCKE_ERR_VALUE, + "wgrad: streamk_ctas=-1 (auto) is not supported in the C port; " + "resolve it from the arch CU count and pass the explicit pool size"); + return NULL; + } + const char* sk_red = wgrad_streamk_reduction(spec); + const bool streamk_atomic + = is_streamk && (strcmp(sk_red, "atomic") == 0 || strcmp(sk_red, "workspace") == 0); + const bool streamk_flags + = is_streamk && (strcmp(sk_red, "linear") == 0 || strcmp(sk_red, "tree") == 0); + bool effective_two_stage = spec->two_stage; bool is_split_k = (split_k > 1 || split_k == 0); @@ -2385,7 +3069,7 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv_wgrad( /* split_k>1 or split_k==0: dW is read+write (atomic); split_k=1: writeonly. * Caller MUST zero-init dW before launch for atomic paths -- the kernel only * issues atomic-adds. See the header contract note for details. */ - d_opts.writeonly = !is_split_k; + d_opts.writeonly = !is_split_k && !streamk_atomic; d_opts.writeonly_set = true; d_opts.align = 16; d_opts.align_set = true; @@ -2420,7 +3104,7 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv_wgrad( * Only present when two_stage=true && split_k>1. */ bool is_two_stage = is_split_k && effective_two_stage; rocke_value_t* ws_ptr = NULL; - if(is_two_stage) + if(is_two_stage || (is_streamk && strcmp(sk_red, "workspace") == 0)) { rocke_param_opts_t ws_opts; memset(&ws_opts, 0, sizeof(ws_opts)); @@ -2432,6 +3116,21 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv_wgrad( ws_ptr = rocke_b_param(b, "ws_ptr", rocke_ptr_type(b, rocke_f32(), "global"), &ws_opts); rocke_b_param(b, "ws_bytes", rocke_i32(), NULL); /* consumed by host */ } + /* Stream-K linear/tree: per-wave ready flags and per-CTA partial slots. */ + rocke_value_t* sk_flags = NULL; + rocke_value_t* sk_partials = NULL; + if(streamk_flags) + { + rocke_param_opts_t sk_opts; + memset(&sk_opts, 0, sizeof(sk_opts)); + sk_opts.noalias = true; + sk_opts.noalias_set = true; + sk_opts.align = 16; + sk_opts.align_set = true; + sk_flags = rocke_b_param(b, "sk_flags", rocke_ptr_type(b, rocke_i32(), "global"), &sk_opts); + sk_partials + = rocke_b_param(b, "sk_partials", rocke_ptr_type(b, rocke_f32(), "global"), &sk_opts); + } /* --- build wgrad ctx (with correct param names) --- */ rocke_conv_build_ctx_t ctx; @@ -2553,6 +3252,14 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv_wgrad( if(!rocke_ir_builder_ok(b)) return NULL; + if(is_streamk) + { + wgrad_emit_streamk(&ctx, spec, dW, dW_desc, ws_ptr, sk_flags, sk_partials, wg_M, wg_N); + if(!rocke_ir_builder_ok(b)) + return NULL; + return b->kernel; + } + /* --- K-loop --- */ if(spec->unroll_k) rocke_conv_emit_kloop_unroll(&ctx); @@ -2605,6 +3312,8 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv_wgrad( WgradDwAddrCtx dw_addr_ctx; dw_addr_ctx.dW_desc = dW_desc; + dw_addr_ctx.group = NULL; + dw_addr_ctx.c_kpg = NULL; if(spec->epilogue && strcmp(spec->epilogue, "cshuffle") == 0) { From f8f8200c7b746de2598f4027504df4c4e39d531d Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Mon, 5 Oct 2026 11:58:40 +0000 Subject: [PATCH 04/20] test(rocke): gate stream-K conv wgrad IR in both engines Add parity configs 18-28 covering both launch modes, all four reductions, both MFMA arches, grouped shapes with a partial last M tile per group, the all-data-parallel partition and the K-outer tile. The byte-identity gate is GREEN for all 29 wgrad configs at llvm20, llvm22 and llvm23. Add three representative-IR cases and re-bless that golden; the diff is append-only, no existing hash moves. The L5 differential anchor is deliberately left alone. It carries no conv_implicit_gemm_wgrad entry today, and on a box whose clang is a different LLVM vintage than the one it was blessed on it already reports most of its families as changed -- so entries recorded here would be anchored against a different toolchain state than the rest of the file. Bless that family from the canonical toolchain instead. ISSUE ID : AICK-1943 Co-Authored-By: Claude Opus 5 --- .../parity/conv_implicit_gemm_wgrad_emit.c | 78 ++++++++ .../parity/conv_implicit_gemm_wgrad_emit.py | 172 +++++++++++++++++- .../rocke_representative_ir_sha256.json | 96 ++++++++-- .../instances/rocke_ir_parity_harness.py | 68 +++++++ 4 files changed, 401 insertions(+), 13 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c index ccb4ecf1230e..b12c9274f118 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c @@ -29,6 +29,17 @@ * 15 N8H56W56C64_K64Y3X3, t64x64x64, w2x2, a32x32x16, mem/default, gfx950, split_k=4 two_stage fp16 * 16 N8H56W56C64_K64Y3X3, t64x64x64, w2x2, a16x16x16, mem/default, gfx942, split_k=4 two_stage fp16 * 17 N8H56W56C64_K64Y3X3 pad1, t32x32x32, w1x1, a16x16x32, mem/default, gfx1250 (WMMA w32), lds_k_outer fp32 out + * 18 N8H56W56C64_K64Y3X3, a32x32x16, gfx950, streamk dp_sk linear sk16 fp16 + * 19 N8H56W56C64_K64Y3X3, a32x32x16, gfx950, streamk dp_sk tree sk16 fp16 + * 20 N8H56W56C64_K64Y3X3, a32x32x16, gfx950, streamk dp_sk atomic sk4 fp32 (DP + SK) + * 21 N8H56W56C64_K64Y3X3, a32x32x16, gfx950, streamk dp_sk workspace sk4 fp16 + * 22 N8H56W56C64_K64Y3X3, a16x16x16, gfx942, streamk persistent linear sk4 fp16 + * 23 N4H28W28C96_K120Y3X3G3 pad1, a32x32x16, gfx950, streamk dp_sk linear sk6 fp16 + * 24 N4H28W28C64_K128Y3X3G2 pad1, a32x32x16, gfx950, streamk persistent tree sk5 bf16 dW + * 25 N4H28W28C64_K128Y3X3G2 pad1, a32x32x16, gfx950, streamk persistent atomic sk3 fp32 + * 26 N8H56W56C64_K64Y3X3, a16x16x16, gfx942, streamk persistent workspace sk6 fp16 + * 27 N8H56W56C64_K64Y3X3, a32x32x16, gfx950, streamk dp_sk linear sk9 fp16 (all DP) + * 28 N4H14W14C64_K96Y3X3G2 pad1, a32x32x16, gfx950, streamk dp_sk linear sk7 fp16, lds_k_outer */ #include #include @@ -40,6 +51,42 @@ #include "rocke/lower_llvm.h" #include "rocke/verify.h" +/* Stream-K fields of configs 18-27 (Python _STREAMK_CFGS). */ +static int streamk_cfg(int idx, rocke_implicit_gemm_conv_wgrad_spec_t* spec, const char** arch) +{ + /* idx: dtype a/b, dtype d, atom edge, streamk, reduction, ctas, arch */ + static const struct + { + const char *dt_ab, *dt_d; + int edge; + const char *mode, *reduction; + int ctas; + const char* arch; + } cfgs[10] = { + {"fp16", "fp16", 32, "dp_sk", "linear", 16, "gfx950"}, + {"fp16", "fp16", 32, "dp_sk", "tree", 16, "gfx950"}, + {"fp16", "fp32", 32, "dp_sk", "atomic", 4, "gfx950"}, + {"fp16", "fp16", 32, "dp_sk", "workspace", 4, "gfx950"}, + {"fp16", "fp16", 16, "persistent", "linear", 4, "gfx942"}, + {"fp16", "fp16", 32, "dp_sk", "linear", 6, "gfx950"}, + {"fp16", "bf16", 32, "persistent", "tree", 5, "gfx950"}, + {"fp16", "fp32", 32, "persistent", "atomic", 3, "gfx950"}, + {"fp16", "fp16", 16, "persistent", "workspace", 6, "gfx942"}, + {"fp16", "fp16", 32, "dp_sk", "linear", 9, "gfx950"}, + }; + const int i = idx - 18; + spec->dtype_a = cfgs[i].dt_ab; + spec->dtype_b = cfgs[i].dt_ab; + spec->dtype_d = cfgs[i].dt_d; + spec->warp_tile_m = cfgs[i].edge; + spec->warp_tile_n = cfgs[i].edge; + spec->streamk = cfgs[i].mode; + spec->streamk_reduction = cfgs[i].reduction; + spec->streamk_ctas = cfgs[i].ctas; + *arch = cfgs[i].arch; + return 0; +} + /* Fill the config for index `idx`. Returns 0 on success, -1 if unknown. * On success sets *spec and *arch. */ static int make_cfg(int idx, rocke_implicit_gemm_conv_wgrad_spec_t* spec, const char** arch) @@ -213,6 +260,37 @@ static int make_cfg(int idx, rocke_implicit_gemm_conv_wgrad_spec_t* spec, const spec->lds_k_outer = true; *arch = "gfx1250"; return 0; + case 18: + case 19: + case 20: + case 21: + case 22: + case 26: + case 27: + spec->problem = rocke_conv_problem_default(8, 56, 56, 64, 64, 3, 3); + return streamk_cfg(idx, spec, arch); + case 23: + spec->problem = rocke_conv_problem_default(4, 28, 28, 96, 120, 3, 3); + spec->problem.pH = 1; + spec->problem.pW = 1; + spec->problem.groups = 3; + return streamk_cfg(idx, spec, arch); + case 24: + case 25: + spec->problem = rocke_conv_problem_default(4, 28, 28, 64, 128, 3, 3); + spec->problem.pH = 1; + spec->problem.pW = 1; + spec->problem.groups = 2; + return streamk_cfg(idx, spec, arch); + case 28: + spec->problem = rocke_conv_problem_default(4, 14, 14, 64, 96, 3, 3); + spec->problem.pH = 1; + spec->problem.pW = 1; + spec->problem.groups = 2; + streamk_cfg(23, spec, arch); + spec->streamk_ctas = 7; + spec->lds_k_outer = true; + return 0; default: return -1; } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py index c92b8aa4a5e4..ad88da051ec7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.py @@ -28,6 +28,17 @@ # 15 -- split-K=4, two_stage=True (workspace-store epilogue), fp16, gfx950 # 16 -- split-K=4, two_stage=True (workspace-store epilogue), fp16, gfx942 # 17 -- gfx1250 wave32 WMMA 16x16x32 K-outer (ds_load_tr16_b128 transpose reads) +# 18 -- stream-K dp_sk, linear fixup, fp16, gfx950 +# 19 -- stream-K dp_sk, tree fixup, fp16, gfx950 +# 20 -- stream-K dp_sk, atomic reduction, fp32, DP + SK mix, gfx950 +# 21 -- stream-K dp_sk, workspace reduction, fp16, gfx950 +# 22 -- stream-K persistent, linear fixup, 16x16x16 atom, gfx942 +# 23 -- stream-K dp_sk, linear, groups=3 with a partial last M tile, gfx950 +# 24 -- stream-K persistent, tree, bf16 dW, groups=2, gfx950 +# 25 -- stream-K persistent, atomic, fp32, groups=2, gfx950 +# 26 -- stream-K persistent, workspace, fp16, gfx942 +# 27 -- stream-K dp_sk, linear, all-DP partition (tiles divide the pool), gfx950 +# 28 -- stream-K dp_sk, linear, K-outer LDS tile, groups=2, gfx950 # (async_dma omitted: C++ async load path does not yet honour the wgrad A-descriptor # override, so it would produce different IR and break the byte-identity gate) # @@ -35,8 +46,13 @@ # 100 -- odd C with fp16 split-K (must raise ValueError) # 102 -- split_k > 1 on RDNA gfx1151 (must raise ValueError) # 103 -- two_stage=True with split_k=1 (must raise ValueError) +# 104 -- streamk with split_k=4 (must raise ValueError) +# 105 -- streamk on gfx1250 (must raise ValueError) +# 106 -- streamk with the cshuffle epilogue (must raise ValueError) +# 107 -- streamk with unroll_k (must raise ValueError) +# 108 -- streamk atomic reduction into fp16 dW (must raise ValueError) # (These illustrate the validator contract. The C emitter defines only cases -# 0-17, so run_diff.py stops at the shared END before reaching 100+; these +# 0-28, so run_diff.py stops at the shared END before reaching 100+; these # configs are not exercised by the differential gate.) from kernels.common.conv_implicit_gemm_wgrad import ( WgradConvSpec, @@ -459,6 +475,24 @@ def _spec(idx: int): "gfx1250", ) + if 18 <= idx <= 27: + return _streamk_spec(idx) + if idx == 28: + from dataclasses import replace + + spec, arch = _streamk_spec(23) + return ( + replace( + spec, + problem=ConvProblem( + N=4, Hi=14, Wi=14, C=64, K=96, Y=3, X=3, pH=1, pW=1, groups=2 + ), + streamk_ctas=7, + lds_k_outer=True, + ), + arch, + ) + # ---------------------------------------------------------------- # Negative cases: these specs must be REJECTED by the validator. # The harness (run_emit) expects a ValueError / SystemExit when @@ -526,9 +560,145 @@ def _spec(idx: int): ), "gfx950", ) + if 104 <= idx <= 108: + from dataclasses import replace + + from kernels.common._conv_implicit_gemm_common import ConvDataSpec + + base, arch = _streamk_spec(18) + bad = { + 104: (dict(split_k=4), arch), + 105: (dict(), "gfx1250"), + 106: (dict(epilogue="cshuffle"), arch), + 107: (dict(unroll_k=True), arch), + 108: (dict(streamk_reduction="atomic", data=ConvDataSpec()), arch), + }[idx] + return replace(base, **bad[0]), bad[1] raise SystemExit(f"unknown config index {idx}") +# (problem kwargs, data dtypes (a/b, d), atom edge, streamk, reduction, ctas, arch) +_STREAMK_CFGS = { + 18: ( + dict(N=8, Hi=56, Wi=56, C=64, K=64, Y=3, X=3), + ("fp16", "fp16"), + 32, + "dp_sk", + "linear", + 16, + "gfx950", + ), + 19: ( + dict(N=8, Hi=56, Wi=56, C=64, K=64, Y=3, X=3), + ("fp16", "fp16"), + 32, + "dp_sk", + "tree", + 16, + "gfx950", + ), + 20: ( + dict(N=8, Hi=56, Wi=56, C=64, K=64, Y=3, X=3), + ("fp16", "fp32"), + 32, + "dp_sk", + "atomic", + 4, + "gfx950", + ), + 21: ( + dict(N=8, Hi=56, Wi=56, C=64, K=64, Y=3, X=3), + ("fp16", "fp16"), + 32, + "dp_sk", + "workspace", + 4, + "gfx950", + ), + 22: ( + dict(N=8, Hi=56, Wi=56, C=64, K=64, Y=3, X=3), + ("fp16", "fp16"), + 16, + "persistent", + "linear", + 4, + "gfx942", + ), + 23: ( + dict(N=4, Hi=28, Wi=28, C=96, K=120, Y=3, X=3, pH=1, pW=1, groups=3), + ("fp16", "fp16"), + 32, + "dp_sk", + "linear", + 6, + "gfx950", + ), + 24: ( + dict(N=4, Hi=28, Wi=28, C=64, K=128, Y=3, X=3, pH=1, pW=1, groups=2), + ("fp16", "bf16"), + 32, + "persistent", + "tree", + 5, + "gfx950", + ), + 25: ( + dict(N=4, Hi=28, Wi=28, C=64, K=128, Y=3, X=3, pH=1, pW=1, groups=2), + ("fp16", "fp32"), + 32, + "persistent", + "atomic", + 3, + "gfx950", + ), + 26: ( + dict(N=8, Hi=56, Wi=56, C=64, K=64, Y=3, X=3), + ("fp16", "fp16"), + 16, + "persistent", + "workspace", + 6, + "gfx942", + ), + 27: ( + dict(N=8, Hi=56, Wi=56, C=64, K=64, Y=3, X=3), + ("fp16", "fp16"), + 32, + "dp_sk", + "linear", + 9, + "gfx950", + ), +} + + +def _streamk_spec(idx: int): + """Stream-K configs; mirrored field-for-field by the C emitter.""" + from kernels.common._conv_implicit_gemm_common import ConvDataSpec + + prob, (dt_ab, dt_d), edge, mode, reduction, ctas, arch = _STREAMK_CFGS[idx] + return ( + WgradConvSpec( + problem=ConvProblem(**prob), + data=ConvDataSpec(dtype_a=dt_ab, dtype_b=dt_ab, dtype_d=dt_d), + tile_m=64, + tile_n=64, + tile_k=64, + warp_m=2, + warp_n=2, + warp_tile_m=edge, + warp_tile_n=edge, + warp_tile_k=16, + pipeline="mem", + epilogue="default", + streamk=mode, + streamk_reduction=reduction, + streamk_ctas=ctas, + ), + arch, + ) + + def main() -> int: return run_emit( _spec, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json index c8ec389b798d..bc73cb54136e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json @@ -746,6 +746,14 @@ "sha256": "52900a534a609c6c422ec342a5d55fc42b883f703d772495b8ec9908ae2af865", "status": "ok" }, + "conv_wgrad/gfx942/n1h8c64k64r3_g4_sk_tree_pers": { + "arch": "gfx942", + "bytes": 39488, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_942_g4_sk_tree_pers_N1H8W8C64_K64Y3X3G4_t64x32x16_w2x2_a16x16x16_mem_default_sk6_skpers_skrtree", + "sha256": "71077f557f73bd7dc7b5879dd88371ea6bfb1ce9c70e0f291c57c1e13036e131", + "status": "ok" + }, "conv_wgrad/gfx942/n1h8c64k64r3_g4_spk4": { "arch": "gfx942", "bytes": 25573, @@ -762,6 +770,22 @@ "sha256": "9583afa07deb02fd74da899ec796898f345b20634fd555804bd631c773ab4408", "status": "ok" }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_atomic": { + "arch": "gfx950", + "bytes": 27823, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_atomic_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk4_skratomic", + "sha256": "b9d2343a6ee63e7e4eab25739233874ee6164e15cf3d1fe99e6be58292c8c0e6", + "status": "ok" + }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_linear": { + "arch": "gfx950", + "bytes": 54812, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_linear_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk4", + "sha256": "4dbfc25a7ba984c0d3261a98e42c74bec09c363f8e2313c67dd5da2d3f0d024f", + "status": "ok" + }, "conv_wgrad/gfx950/n1h8c16k32r3_spk4": { "arch": "gfx950", "bytes": 24812, @@ -1228,8 +1252,8 @@ "gfx1201": 7, "gfx1250": 4, "gfx90a": 5, - "gfx942": 32, - "gfx950": 91 + "gfx942": 33, + "gfx950": 93 }, "expected_failures": 4, "families": { @@ -1241,7 +1265,7 @@ "conv_dgrad": 11, "conv_direct": 11, "conv_direct_dgrad": 5, - "conv_wgrad": 13, + "conv_wgrad": 16, "conv_wgrad_reduce": 3, "conv_wgrad_two_stage": 2, "deep_fused_conv": 9, @@ -1250,7 +1274,7 @@ "kda_chunkwise": 13, "moe": 9 }, - "ok_cases": 150 + "ok_cases": 153 } }, "llvm22": { @@ -1999,6 +2023,14 @@ "sha256": "da3b6b47e391747863cc5ffabdb6e59f8c1e4111a0074b709e27dbbdc703b127", "status": "ok" }, + "conv_wgrad/gfx942/n1h8c64k64r3_g4_sk_tree_pers": { + "arch": "gfx942", + "bytes": 39627, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_942_g4_sk_tree_pers_N1H8W8C64_K64Y3X3G4_t64x32x16_w2x2_a16x16x16_mem_default_sk6_skpers_skrtree", + "sha256": "f1f7c75f607f8251ac12813743eabe4ed05873624fef3fc7e23bcd8432ff798f", + "status": "ok" + }, "conv_wgrad/gfx942/n1h8c64k64r3_g4_spk4": { "arch": "gfx942", "bytes": 25712, @@ -2015,6 +2047,22 @@ "sha256": "da8dc905a5d13734652fc50c17e19292d8d93dcc025a357d8c3c2984e3f2f639", "status": "ok" }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_atomic": { + "arch": "gfx950", + "bytes": 27956, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_atomic_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk4_skratomic", + "sha256": "b5468f91ae49e458a986624fb24e47b0b04d69aada5338e2b3e106749b73d805", + "status": "ok" + }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_linear": { + "arch": "gfx950", + "bytes": 54951, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_linear_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk4", + "sha256": "7efa530595cf1278ecd9c1a6d834786cc69f65965a7742b7a6ad83e752e4dd5f", + "status": "ok" + }, "conv_wgrad/gfx950/n1h8c16k32r3_spk4": { "arch": "gfx950", "bytes": 24951, @@ -2481,8 +2529,8 @@ "gfx1201": 7, "gfx1250": 4, "gfx90a": 5, - "gfx942": 32, - "gfx950": 91 + "gfx942": 33, + "gfx950": 93 }, "expected_failures": 4, "families": { @@ -2494,7 +2542,7 @@ "conv_dgrad": 11, "conv_direct": 11, "conv_direct_dgrad": 5, - "conv_wgrad": 13, + "conv_wgrad": 16, "conv_wgrad_reduce": 3, "conv_wgrad_two_stage": 2, "deep_fused_conv": 9, @@ -2503,7 +2551,7 @@ "kda_chunkwise": 13, "moe": 9 }, - "ok_cases": 150 + "ok_cases": 153 } }, "llvm23": { @@ -3252,6 +3300,14 @@ "sha256": "472def6cd2a3cf4779957596112eb429bea3babad2909eebf6d29e6e8eb42f59", "status": "ok" }, + "conv_wgrad/gfx942/n1h8c64k64r3_g4_sk_tree_pers": { + "arch": "gfx942", + "bytes": 39631, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_942_g4_sk_tree_pers_N1H8W8C64_K64Y3X3G4_t64x32x16_w2x2_a16x16x16_mem_default_sk6_skpers_skrtree", + "sha256": "6a88e42497095a3747b2c6014a64611b7fa9185e8efb802950ac7c2e8354cfdd", + "status": "ok" + }, "conv_wgrad/gfx942/n1h8c64k64r3_g4_spk4": { "arch": "gfx942", "bytes": 25716, @@ -3268,6 +3324,22 @@ "sha256": "41165edfcbfa49787abc2e51641866c101d2a58cf50402ac87354d2e65540dcb", "status": "ok" }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_atomic": { + "arch": "gfx950", + "bytes": 27960, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_atomic_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk4_skratomic", + "sha256": "3c3928e6bb08240c45440dd34c760a1d11a4fe5ea038182438cf4fc5a1fcdd0b", + "status": "ok" + }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_linear": { + "arch": "gfx950", + "bytes": 54955, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_linear_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk4", + "sha256": "92be3e857e53dcf1dd8303e10c33086bee5afab1ffd939cb960f540944b06241", + "status": "ok" + }, "conv_wgrad/gfx950/n1h8c16k32r3_spk4": { "arch": "gfx950", "bytes": 24955, @@ -3749,8 +3821,8 @@ "gfx1201": 7, "gfx1250": 8, "gfx90a": 5, - "gfx942": 32, - "gfx950": 91 + "gfx942": 33, + "gfx950": 93 }, "expected_failures": 0, "families": { @@ -3762,7 +3834,7 @@ "conv_dgrad": 11, "conv_direct": 11, "conv_direct_dgrad": 5, - "conv_wgrad": 13, + "conv_wgrad": 16, "conv_wgrad_reduce": 3, "conv_wgrad_two_stage": 2, "deep_fused_conv": 9, @@ -3771,7 +3843,7 @@ "kda_chunkwise": 13, "moe": 9 }, - "ok_cases": 154 + "ok_cases": 157 } } }, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py index e5115172cd03..e02fbab23637 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py @@ -175,6 +175,9 @@ def build_conv_wgrad( split_k=1, dtype_d="fp16", two_stage=False, + streamk="off", + streamk_reduction="linear", + streamk_ctas=-1, ): def _build(): from kernels.common.conv_implicit_gemm_wgrad import ( @@ -204,6 +207,9 @@ def _build(): epilogue=epilogue, split_k=split_k, two_stage=two_stage, + streamk=streamk, + streamk_reduction=streamk_reduction, + streamk_ctas=streamk_ctas, ) return build_implicit_gemm_conv_wgrad(spec, arch=arch) @@ -1757,6 +1763,68 @@ def add(family, case_id, arch, build): epilogue="cshuffle", ), ) + # Stream-K wgrad: the linear and tree flag fixups, the atomic reduction, + # the persistent launch and the group fold into GEMM-M. MFMA-only. + add( + "conv_wgrad", + "conv_wgrad/gfx950/n1h8c16k32r3_sk_linear", + "gfx950", + build_conv_wgrad( + "irhash_wgrad_950_sk_linear", + "gfx950", + wgrad1, + wave_size=64, + wtm=32, + wtn=32, + wtk=16, + tile_m=64, + tile_n=64, + tile_k=16, + streamk="dp_sk", + streamk_ctas=4, + ), + ) + add( + "conv_wgrad", + "conv_wgrad/gfx950/n1h8c16k32r3_sk_atomic", + "gfx950", + build_conv_wgrad( + "irhash_wgrad_950_sk_atomic", + "gfx950", + wgrad1, + wave_size=64, + wtm=32, + wtn=32, + wtk=16, + tile_m=64, + tile_n=64, + tile_k=16, + dtype_d="fp32", + streamk="dp_sk", + streamk_reduction="atomic", + streamk_ctas=4, + ), + ) + add( + "conv_wgrad", + "conv_wgrad/gfx942/n1h8c64k64r3_g4_sk_tree_pers", + "gfx942", + build_conv_wgrad( + "irhash_wgrad_942_g4_sk_tree_pers", + "gfx942", + wgrad_g4, + wave_size=64, + wtm=16, + wtn=16, + wtk=16, + tile_m=64, + tile_n=32, + tile_k=16, + streamk="persistent", + streamk_reduction="tree", + streamk_ctas=6, + ), + ) # Grouped + cshuffle epilogue (MFMA): the LDS-staged store threads the # per-group k_out fold (group*kpg) and derives its store-vec from cpg. add( From 6cf88cac70586ce47bb8a4cb7626ef63e4fb5652 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Mon, 5 Oct 2026 11:58:40 +0000 Subject: [PATCH 05/20] test(rocke): add CPU and GPU tests for stream-K conv wgrad CPU: the spec surface and kernel-name tags, every validator gate with validate() and is_valid_wgrad_spec agreeing on the reason string, the C port rejecting with that same string, admits-implies-builds over mode x reduction x arch x groups, the host launch plan cross-checked against the C port through a new binding, and byte-identical lowering between the engines. GPU: numerics against a float32 torch reference on shapes chosen to exercise the partitioner rather than the GEMM -- up to four CTAs per tile (so the linear fold and both tree rounds run), a data-parallel plus stream-K mix, and a grouped problem whose last M tile is partial. A row-coverage guard reports whole unwritten output channels as such instead of as an opaque tolerance miss. The linear and tree fixups are launched repeatedly and must be bit-identical, which is what would catch a missing release or acquire on the partial hand-off. Both files are registered in the CTest lists so CI runs them. ISSUE ID : AICK-1943 Co-Authored-By: Claude Opus 5 --- .../library/tests/test_conv_wgrad_streamk.py | 325 +++++++++++++++ .../tests/test_conv_wgrad_streamk_gpu.py | 373 ++++++++++++++++++ .../rocke/platform/CMakeLists.txt | 2 + 3 files changed, 700 insertions(+) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk_gpu.py diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py new file mode 100644 index 000000000000..328882a6d61b --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py @@ -0,0 +1,325 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Host-side tests for stream-K implicit-GEMM backward-weight convolution. + +No GPU needed. Covers the spec surface (defaults, kernel-name tags), the +validator gates and their agreement between ``validate()``, +``is_valid_wgrad_spec`` and the C++ port, admits-implies-builds over the +mode x reduction x arch x groups matrix, the host-side launch plan (grid, +workspace layout, determinism) against the C port, and byte-identical lowering +between the two engines through the ``rocke_engine`` binding. + +GPU numerics live in ``test_conv_wgrad_streamk_gpu.py``. +""" + +from __future__ import annotations + +import dataclasses +import unittest +from dataclasses import replace + +from kernels.common._conv_implicit_gemm_common import ConvDataSpec, ConvProblem +from kernels.common.conv_implicit_gemm_wgrad import ( + WgradConvSpec, + build_implicit_gemm_conv_wgrad, + is_valid_wgrad_spec, + wgrad_streamk_grid, + wgrad_streamk_partition, + wgrad_streamk_signature, + wgrad_streamk_workspace_layout, + wgrad_streamk_workspace_nbytes, +) + +_DENSE = ConvProblem(N=2, Hi=12, Wi=12, C=32, K=96, Y=3, X=3, pH=1, pW=1) +# kpg = 40 is not a multiple of tile_m, so the last M tile of every group is +# partial: the case CK's own GemmM * GemmBatch fold leaves untested. +_GROUPED = ConvProblem(N=2, Hi=10, Wi=10, C=48, K=120, Y=3, X=3, pH=1, pW=1, groups=3) +# Depthwise: the only shape group merging admits, so its stream-K gate is reached. +_DW = ConvProblem(N=2, Hi=8, Wi=8, C=16, K=16, Y=3, X=3, pH=1, pW=1, groups=16) + + +def _spec(problem=_DENSE, *, dtype_d="fp16", edge=32, **kw) -> WgradConvSpec: + base = dict( + problem=problem, + data=ConvDataSpec(dtype_a="fp16", dtype_b="fp16", dtype_d=dtype_d), + tile_m=64, + tile_n=64, + tile_k=32, + warp_m=2, + warp_n=2, + warp_tile_m=edge, + warp_tile_n=edge, + warp_tile_k=16, + streamk="dp_sk", + streamk_ctas=6, + ) + base.update(kw) + return WgradConvSpec(**base) + + +def _engine(): + """The ``rocke_engine`` binding, or skip the calling test.""" + from rocke.core.backend import BackendError, _import_engine + + try: + return _import_engine() + except BackendError as e: + raise unittest.SkipTest(f"C++ engine not importable: {e}") from None + + +def _spec_dict(spec: WgradConvSpec) -> dict: + """Flatten a WgradConvSpec into the dict ``conv_wgrad_build_spec`` reads.""" + d = { + f.name: getattr(spec, f.name) + for f in dataclasses.fields(spec) + if f.name not in ("problem", "data", "acc_epilogue", "lds_layout") + } + d = {k: v for k, v in d.items() if v is not None} + d["problem"] = { + k: v for k, v in dataclasses.asdict(spec.problem).items() if v is not None + } + d["dtype_a"] = spec.data.dtype_a + d["dtype_b"] = spec.data.dtype_b + d["dtype_d"] = spec.data.dtype_d + return d + + +class TestStreamKSpecSurface(unittest.TestCase): + def test_default_is_off_and_name_unchanged(self): + plain = replace(_spec(), streamk="off", streamk_ctas=-1) + self.assertNotIn("_sk", plain.kernel_name()) + self.assertEqual( + plain.kernel_name(), + WgradConvSpec( + **{ + f.name: getattr(plain, f.name) + for f in dataclasses.fields(plain) + if not f.name.startswith("streamk") + } + ).kernel_name(), + ) + + def test_name_tags_every_knob_that_changes_the_body(self): + names = { + _spec().kernel_name(), + _spec(streamk_ctas=8).kernel_name(), + _spec(streamk="persistent").kernel_name(), + _spec(streamk_reduction="tree").kernel_name(), + _spec(streamk_reduction="workspace").kernel_name(), + _spec(dtype_d="fp32", streamk_reduction="atomic").kernel_name(), + _spec(streamk_ctas=-1).kernel_name(), + } + self.assertEqual(len(names), 7, names) + self.assertTrue(_spec().kernel_name().endswith("_sk6")) + self.assertIn("_skauto", _spec(streamk_ctas=-1).kernel_name()) + + def test_builder_resolves_auto_pool_into_the_name(self): + k = build_implicit_gemm_conv_wgrad(_spec(streamk_ctas=-1), arch="gfx950") + self.assertTrue(k.name.endswith("_sk256"), k.name) + + +class TestStreamKValidator(unittest.TestCase): + """Every gate rejects in ``validate()`` and ``is_valid_wgrad_spec`` alike.""" + + _REJECTS = [ + (dict(streamk="sideways"), "streamk must be"), + (dict(streamk_reduction="ring"), "streamk_reduction must be"), + (dict(streamk_ctas=0), "streamk_ctas must be"), + (dict(streamk_ctas=-3), "streamk_ctas must be"), + (dict(split_k=4), "both partition K_wg"), + (dict(split_k=-1), "both partition K_wg"), + (dict(two_stage=True, split_k=2), "both partition K_wg"), + (dict(unroll_k=True), "runtime K trip count"), + (dict(epilogue="cshuffle"), "epilogue='default'"), + (dict(chiplet_swizzle=True), "chiplet_swizzle"), + (dict(streamk_reduction="atomic"), "needs dtype_d='fp32'"), + ] + + def test_rejects(self): + for kw, needle in self._REJECTS: + spec = _spec(**kw) + with self.subTest(kw=kw): + ok, why = is_valid_wgrad_spec(spec, "gfx950") + self.assertFalse(ok) + self.assertIn(needle, why) + with self.assertRaises(ValueError) as cm: + spec.validate() + self.assertEqual(str(cm.exception), why) + + def test_group_merge_and_acc_epilogue_rejected(self): + from kernels.common._conv_implicit_gemm_common import ConvAccumulatorEpilogue + + for spec, needle in ( + (_spec(_DW, group_merge=2), "group_merge=1"), + (_spec(acc_epilogue=ConvAccumulatorEpilogue(relu=True)), "acc_epilogue"), + ): + with self.subTest(needle=needle): + ok, why = is_valid_wgrad_spec(spec, "gfx950") + self.assertFalse(ok) + self.assertIn(needle, why) + + def test_arch_gate(self): + ok, why = is_valid_wgrad_spec(_spec(), "gfx90a") + self.assertFalse(ok) + self.assertIn("stream-K wgrad supports only gfx942, gfx950", why) + # Wave32 targets already fail the wave-size check; either way, rejected. + for arch in ("gfx1151", "gfx1201", "gfx1250"): + with self.subTest(arch=arch): + self.assertFalse(is_valid_wgrad_spec(_spec(), arch)[0]) + + def test_cpp_rejects_with_the_same_reason(self): + eng = _engine() + for kw, _needle in self._REJECTS: + spec = _spec(**kw) + _ok, why = is_valid_wgrad_spec(spec, "gfx950") + with self.subTest(kw=kw), self.assertRaises(RuntimeError) as cm: + eng.conv_wgrad_lower_llvm(_spec_dict(spec), "gfx950") + self.assertIn(why, str(cm.exception)) + + +class TestStreamKAdmitsImpliesBuilds(unittest.TestCase): + def test_matrix(self): + built = 0 + for arch, edge in (("gfx950", 32), ("gfx942", 16)): + for problem in (_DENSE, _GROUPED): + for mode in ("dp_sk", "persistent"): + for reduction, dtype_d in ( + ("linear", "fp16"), + ("tree", "bf16"), + ("atomic", "fp32"), + ("workspace", "fp16"), + ): + for ctas in (1, 4, 7, 64): + spec = _spec( + problem, + dtype_d=dtype_d, + edge=edge, + streamk=mode, + streamk_reduction=reduction, + streamk_ctas=ctas, + ) + ok, why = is_valid_wgrad_spec(spec, arch) + with self.subTest( + arch=arch, + g=problem.groups, + mode=mode, + r=reduction, + w=ctas, + ): + self.assertTrue(ok, why) + build_implicit_gemm_conv_wgrad(spec, arch=arch) + built += 1 + self.assertEqual(built, 2 * 2 * 2 * 4 * 4) + + +class TestStreamKHostPlan(unittest.TestCase): + def test_grid_matches_partition(self): + for mode in ("dp_sk", "persistent"): + spec = _spec(streamk=mode, streamk_ctas=4) + part = wgrad_streamk_partition(spec) + with self.subTest(mode=mode): + want = 4 if mode == "persistent" else part.dp_tiles + part.sk_ctas + self.assertEqual(wgrad_streamk_grid(spec), (want, 1, 1)) + + def test_group_fold_tile_count(self): + # m_tiles is per group: ceil(kpg / tile_m) * groups, never + # ceil(kpg * groups / tile_m). + part = wgrad_streamk_partition(_spec(_GROUPED)) + self.assertEqual(part.m_tiles, 3) # ceil(40 / 64) * 3 + self.assertEqual(part.n_tiles, 3) # ceil(3*3*16 / 64) + + def test_workspace_layout(self): + spec = _spec(streamk_ctas=7) + self.assertEqual(wgrad_streamk_partition(spec).sk_ctas, 7) + flags, partials = wgrad_streamk_workspace_layout(spec) + self.assertEqual(flags, 256) # 7 CTAs * 4 waves * 4 B, 256-aligned + self.assertEqual(partials, 7 * 64 * 64 * 4) + self.assertEqual(wgrad_streamk_workspace_nbytes(spec), flags + partials) + # 10 tiles over a pool of 5 split evenly: all data-parallel, no slots. + self.assertEqual(wgrad_streamk_workspace_layout(_spec(streamk_ctas=5)), (0, 0)) + atomic = _spec(dtype_d="fp32", streamk_reduction="atomic") + self.assertEqual(wgrad_streamk_workspace_nbytes(atomic), 0) + ws = _spec(streamk_reduction="workspace") + self.assertEqual(wgrad_streamk_workspace_nbytes(ws), 8 * 96 * 288 * 4) + + def test_signature_tracks_reduction(self): + names = lambda s: [p["name"] for p in wgrad_streamk_signature(s)] # noqa: E731 + self.assertEqual(names(_spec())[-2:], ["sk_flags", "sk_partials"]) + self.assertEqual( + names(_spec(streamk_reduction="workspace"))[-2:], ["ws_ptr", "ws_bytes"] + ) + self.assertEqual( + len(names(_spec(dtype_d="fp32", streamk_reduction="atomic"))), 6 + ) + + def test_cpp_host_plan_agrees(self): + eng = _engine() + for problem in (_DENSE, _GROUPED): + for mode in ("dp_sk", "persistent"): + for reduction, dtype_d in ( + ("linear", "fp16"), + ("tree", "fp16"), + ("atomic", "fp32"), + ("workspace", "fp16"), + ): + spec = _spec( + problem, + dtype_d=dtype_d, + streamk=mode, + streamk_reduction=reduction, + streamk_ctas=7, + ) + got = eng.conv_wgrad_streamk_host(_spec_dict(spec)) + flags, partials = wgrad_streamk_workspace_layout(spec) + with self.subTest(g=problem.groups, mode=mode, r=reduction): + self.assertEqual(tuple(got["grid"]), wgrad_streamk_grid(spec)) + self.assertEqual(got["flags_bytes"], flags) + self.assertEqual(got["partials_bytes"], partials) + self.assertEqual( + got["workspace_bytes"], wgrad_streamk_workspace_nbytes(spec) + ) + self.assertEqual( + got["is_deterministic"], reduction in ("linear", "tree") + ) + + +class TestStreamKDualEngine(unittest.TestCase): + """The C++ builder lowers every stream-K variant to the Python engine's bytes.""" + + def test_lowered_ir_is_byte_identical(self): + from rocke.core.lower_llvm import _lower_kernel_to_llvm_python + + eng = _engine() + cases = [ + ("gfx950", _DENSE, "dp_sk", "linear", "fp16", 32), + ("gfx950", _DENSE, "dp_sk", "tree", "fp16", 32), + ("gfx950", _GROUPED, "persistent", "linear", "bf16", 32), + ("gfx950", _GROUPED, "dp_sk", "atomic", "fp32", 32), + ("gfx942", _GROUPED, "persistent", "tree", "fp16", 16), + ("gfx942", _DENSE, "persistent", "workspace", "fp16", 16), + ] + for arch, problem, mode, reduction, dtype_d, edge in cases: + spec = _spec( + problem, + dtype_d=dtype_d, + edge=edge, + streamk=mode, + streamk_reduction=reduction, + ) + with self.subTest(arch=arch, g=problem.groups, mode=mode, r=reduction): + py_ll = _lower_kernel_to_llvm_python( + build_implicit_gemm_conv_wgrad(spec, arch=arch), arch=arch + ) + self.assertEqual( + eng.conv_wgrad_lower_llvm(_spec_dict(spec), arch), py_ll + ) + + def test_cpp_rejects_auto_pool(self): + eng = _engine() + with self.assertRaises(RuntimeError) as cm: + eng.conv_wgrad_lower_llvm(_spec_dict(_spec(streamk_ctas=-1)), "gfx950") + self.assertIn("streamk_ctas=-1 (auto)", str(cm.exception)) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk_gpu.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk_gpu.py new file mode 100644 index 000000000000..33b21acdf5a2 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk_gpu.py @@ -0,0 +1,373 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""GPU numerics for stream-K implicit-GEMM backward-weight convolution. + +Every case builds, compiles and launches a stream-K wgrad kernel, then checks +dW against a float32 ``torch.nn.grad.conv2d_weight`` reference computed on the +CPU. Before the tolerance check a row-coverage guard reports whole output +channels that were never written (a dropped group or tile reports as such +rather than as an opaque numeric miss). + +The shapes are chosen to exercise the partitioner, not the GEMM: several +stream-K CTAs per tile (multi-round linear/tree fixups), a data-parallel plus +stream-K mix, a grouped problem whose last M tile is partial, and a small CTA +pool so the stream-K path is taken at these sizes. The linear and tree fixups +are also run repeatedly and must be bit-identical across launches -- the only +check that catches a missing release/acquire on the partial hand-off. + +Requires a gfx942/gfx950 GPU and torch; skips cleanly otherwise. Host-side +coverage lives in ``test_conv_wgrad_streamk.py``. +""" + +from __future__ import annotations + +import ctypes +import importlib.util +import unittest + +from rocke.runtime.hip_module import get_device_arch + +_HAS_TORCH = importlib.util.find_spec("torch") is not None +GPU_ARCH = get_device_arch(0) + + +def _skip_reason() -> str: + if not GPU_ARCH: + return "no ROCm GPU detected" + if not _HAS_TORCH: + return "torch not importable" + if GPU_ARCH not in ("gfx942", "gfx950"): + return f"stream-K wgrad is gfx942/gfx950 only; got {GPU_ARCH}" + return "" + + +_SKIP_REASON = _skip_reason() + +# Relative max-error bound by dW dtype: the reference is fp32 over fp16/bf16 +# inputs, so the bound is set by the output rounding. +_TOL = {"fp32": 1e-4, "fp16": 2e-3, "bf16": 2e-2} +_REDUCTIONS = ( + ("linear", "fp16"), + ("tree", "fp16"), + ("atomic", "fp32"), + ("workspace", "bf16"), +) + + +def _u8(t): + return (ctypes.c_uint8 * t.nbytes).from_address(t.data_ptr()) + + +def _reference(X, dY, p): + import torch + + w = torch.nn.grad.conv2d_weight( + X.permute(0, 3, 1, 2).contiguous(), + weight_size=(p.K, p.C // p.groups, p.Y, p.X), + grad_output=dY.permute(0, 3, 1, 2).contiguous(), + stride=(p.sH, p.sW), + padding=(p.pH, p.pW), + dilation=(p.dH, p.dW), + groups=p.groups, + ) + return w.permute(0, 2, 3, 1).contiguous() + + +def _run(spec, *, launches: int = 1, grid=None): + """Build, launch ``launches`` times, return (dW outputs, fp32 reference).""" + import torch + + from rocke import compile_kernel + from rocke.runtime.hip_module import Runtime + from rocke.runtime.launcher import KernelLauncher, LaunchConfig + from kernels.common.conv_implicit_gemm_wgrad import ( + build_implicit_gemm_conv_wgrad, + wgrad_streamk_grid, + wgrad_streamk_signature, + wgrad_streamk_workspace_layout, + wgrad_streamk_workspace_nbytes, + ) + from kernels.common.conv_wgrad_workspace_reduce import ( + WgradReduceSpec, + build_conv_wgrad_workspace_reduce, + wgrad_reduce_grid, + wgrad_reduce_signature, + ) + + arch = GPU_ARCH + p = spec.problem + art = compile_kernel(build_implicit_gemm_conv_wgrad(spec, arch=arch), arch=arch) + launcher = KernelLauncher( + hsaco=art.hsaco, + kernel_name=art.kernel_name, + signature=wgrad_streamk_signature(spec), + ) + dtypes = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32} + torch.manual_seed(0) + X = torch.empty(p.N, p.Hi, p.Wi, p.C).uniform_(-1, 1).to(dtypes[spec.data.dtype_b]) + dY = torch.empty(p.N, p.Ho, p.Wo, p.K).uniform_(-1, 1).to(dtypes[spec.data.dtype_a]) + dW = torch.empty(p.K, p.Y, p.X, p.C // p.groups, dtype=dtypes[spec.data.dtype_d]) + ref = _reference(X.float(), dY.float(), p) + + rt = Runtime() + ws_bytes = wgrad_streamk_workspace_nbytes(spec, arch=arch) + bufs = [rt.alloc(dY.nbytes), rt.alloc(X.nbytes), rt.alloc(dW.nbytes)] + ws = rt.alloc(max(ws_bytes, 4)) + bufs.append(ws) + rt.memcpy_h2d(bufs[0], _u8(dY), dY.nbytes) + rt.memcpy_h2d(bufs[1], _u8(X), X.nbytes) + values = { + "A": bufs[0], + "B": bufs[1], + "D": bufs[2], + "A_bytes": dY.nbytes, + "B_bytes": X.nbytes, + "D_bytes": dW.nbytes, + } + reduction = spec.streamk_reduction + flags_bytes, _ = wgrad_streamk_workspace_layout(spec, arch=arch) + if reduction == "workspace": + values.update(ws_ptr=ws, ws_bytes=ws_bytes) + rspec = WgradReduceSpec( + problem=p, + dtype_d=spec.data.dtype_d, + groups=p.groups, + ws_replicas=spec.ws_replicas, + ) + rart = compile_kernel( + build_conv_wgrad_workspace_reduce(rspec, arch=arch), arch=arch + ) + reducer = KernelLauncher( + hsaco=rart.hsaco, + kernel_name=rart.kernel_name, + signature=wgrad_reduce_signature(rspec), + ) + rvalues = { + "ws_ptr": ws, + "dw_ptr": bufs[2], + "wg_M": spec.wg_M, + "wg_N": spec.wg_N, + "ws_bytes": ws_bytes, + "dw_bytes": dW.nbytes, + "groups": p.groups, + } + elif reduction in ("linear", "tree"): + values.update(sk_flags=ws, sk_partials=ws + flags_bytes) + + outs = [] + try: + for _ in range(launches): + # The documented zeroing contract, per reduction: atomic accumulates + # into dW, workspace into the scratch, linear/tree only reset flags. + if reduction == "atomic": + rt.memset(bufs[2], 0, dW.nbytes) + elif reduction == "workspace": + rt.memset(ws, 0, ws_bytes) + elif flags_bytes: + rt.memset(ws, 0, flags_bytes) + launcher( + values, + config=LaunchConfig( + grid=grid or wgrad_streamk_grid(spec, arch=arch), + block=(spec.block_size, 1, 1), + fence=True, + ), + ) + if reduction == "workspace": + reducer( + rvalues, + config=LaunchConfig( + grid=wgrad_reduce_grid(rspec), + block=(rspec.block_size, 1, 1), + fence=True, + ), + ) + out = torch.empty_like(dW) + rt.memcpy_d2h(_u8(out), bufs[2], dW.nbytes) + outs.append(out) + finally: + for buf in bufs: + rt.free(buf) + return outs, ref + + +def _check(test, spec, *, launches: int = 1, grid=None): + outs, ref = _run(spec, launches=launches, grid=grid) + p = spec.problem + out = outs[0].float() + written = out.reshape(p.K, -1).abs().sum(dim=1) > 0 + expected = ref.reshape(p.K, -1).abs().sum(dim=1) > 0 + missing = int((expected & ~written).sum()) + test.assertEqual( + missing, 0, f"{missing}/{p.K} dW output-channel rows were never written" + ) + rel = float((out - ref).abs().max() / ref.abs().max().clamp(min=1.0)) + test.assertLess(rel, _TOL[spec.data.dtype_d], f"rel_err={rel:.3e}") + return outs + + +def _spec(problem, *, mode, reduction, dtype_d, ctas): + from kernels.common._conv_implicit_gemm_common import ConvDataSpec + from kernels.common.conv_implicit_gemm_wgrad import WgradConvSpec + + ab = "bf16" if dtype_d == "bf16" else "fp16" + edge = 32 if GPU_ARCH == "gfx950" else 16 + return WgradConvSpec( + problem=problem, + data=ConvDataSpec(dtype_a=ab, dtype_b=ab, dtype_d=dtype_d), + tile_m=64, + tile_n=64, + tile_k=64, + warp_m=2, + warp_n=2, + warp_tile_m=edge, + warp_tile_n=edge, + warp_tile_k=16, + streamk=mode, + streamk_reduction=reduction, + streamk_ctas=ctas, + ) + + +def _problems(): + from kernels.common._conv_implicit_gemm_common import ConvProblem + + return { + # 6 tiles, 32 iterations each, over 16 CTAs: up to 3 partners per tile, + # two tree rounds. + "multi_partner": ( + ConvProblem(N=8, Hi=16, Wi=16, C=16, K=96, Y=3, X=3, pH=1, pW=1), + 16, + ), + # 25 tiles over 7 CTAs: 14 data-parallel tiles plus 11 stream-K tiles. + "dp_plus_sk": ( + ConvProblem(N=2, Hi=8, Wi=8, C=32, K=320, Y=3, X=3, pH=1, pW=1), + 7, + ), + "strided": ( + ConvProblem( + N=3, Hi=13, Wi=11, C=24, K=72, Y=3, X=3, sH=2, sW=2, pH=1, pW=1 + ), + 5, + ), + } + + +@unittest.skipUnless(not _SKIP_REASON, _SKIP_REASON or "no GPU") +class TestConvWgradStreamKNumerics(unittest.TestCase): + def test_modes_and_reductions(self): + for name, (problem, ctas) in _problems().items(): + for mode in ("dp_sk", "persistent"): + for reduction, dtype_d in _REDUCTIONS: + spec = _spec( + problem, + mode=mode, + reduction=reduction, + dtype_d=dtype_d, + ctas=ctas, + ) + with self.subTest(shape=name, mode=mode, reduction=reduction): + _check(self, spec) + + def test_default_pool(self): + # streamk_ctas=-1 resolves to one CTA per CU of the device. + problem, _ = _problems()["dp_plus_sk"] + spec = _spec(problem, mode="dp_sk", reduction="linear", dtype_d="fp16", ctas=-1) + _check(self, spec) + + +@unittest.skipUnless(not _SKIP_REASON, _SKIP_REASON or "no GPU") +class TestConvWgradStreamKGrouped(unittest.TestCase): + """Groups folded into GEMM-M, including a partial last M tile per group.""" + + def test_grouped(self): + from kernels.common._conv_implicit_gemm_common import ConvProblem + + problems = { + "g3_partial_m": ( + ConvProblem( + N=4, Hi=10, Wi=10, C=48, K=120, Y=3, X=3, pH=1, pW=1, groups=3 + ), + 6, + ), + "g8": ( + ConvProblem( + N=2, Hi=12, Wi=12, C=64, K=64, Y=3, X=3, pH=1, pW=1, groups=8 + ), + 5, + ), + } + for name, (problem, ctas) in problems.items(): + for mode in ("dp_sk", "persistent"): + for reduction, dtype_d in _REDUCTIONS: + spec = _spec( + problem, + mode=mode, + reduction=reduction, + dtype_d=dtype_d, + ctas=ctas, + ) + with self.subTest(shape=name, mode=mode, reduction=reduction): + _check(self, spec) + + +@unittest.skipUnless(not _SKIP_REASON, _SKIP_REASON or "no GPU") +class TestConvWgradStreamKViaDispatch(unittest.TestCase): + """The shipped selector reaches the stream-K kernel and its launch grid.""" + + def test_dispatch_spec_and_grid(self): + from dispatch.grouped_convolution import ( + ConvGroupedRequest, + _problem, + dispatch_conv_grouped, + ) + + for mode in ("dp_sk", "persistent"): + for reduction in ("linear", "tree", "workspace"): + req = ConvGroupedRequest( + N=4, + C=64, + K=96, + Hi=12, + Wi=12, + Y=3, + X=3, + pad_h=1, + pad_w=1, + G=2, + arch=GPU_ARCH, + direction="wgrad", + streamk=mode, + streamk_reduction=reduction, + ) + r = dispatch_conv_grouped(req) + spec = r.spec.to_wgrad_spec(_problem(req)) + with self.subTest(mode=mode, reduction=reduction): + self.assertEqual(spec.streamk, mode) + _check(self, spec, grid=r.grid) + + +@unittest.skipUnless(not _SKIP_REASON, _SKIP_REASON or "no GPU") +class TestConvWgradStreamKDeterminism(unittest.TestCase): + """Linear and tree fold partials in a fixed order: dW is bit-exact.""" + + def test_repeat_launches_are_bit_identical(self): + import torch + + problem, ctas = _problems()["multi_partner"] + for mode in ("dp_sk", "persistent"): + for reduction in ("linear", "tree"): + spec = _spec( + problem, mode=mode, reduction=reduction, dtype_d="fp16", ctas=ctas + ) + with self.subTest(mode=mode, reduction=reduction): + outs = _check(self, spec, launches=20) + for i, out in enumerate(outs[1:], start=1): + self.assertTrue( + torch.equal(out, outs[0]), + f"launch {i} differs from launch 0", + ) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt index 1ec5d0f77825..d0a3a5414d02 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt @@ -462,6 +462,7 @@ if(ROCKE_INSTALL_TESTS) # prefix on each, matching the other installed tests. set(_ROCKE_LIBRARY_CONV_TEST_RELPATHS "tests/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py" + "tests/library/tests/test_conv_wgrad_streamk.py" "tests/library/tests/test_conv_build_signatures.py" "tests/library/tests/test_builder_signature_contract.py" "tests/library/tests/test_library_layering.py") @@ -511,6 +512,7 @@ if(ROCKE_INSTALL_TESTS) set(_ROCKE_LIBRARY_GPU_TEST_RELPATHS "tests/library/tests/test_conv_fwd_correctness.py" "tests/library/tests/test_conv_wgrad_correctness.py" + "tests/library/tests/test_conv_wgrad_streamk_gpu.py" "tests/library/tests/test_conv_dgrad_correctness.py" "tests/library/tests/test_direct_conv_correctness.py" "tests/library/tests/test_attention_dense_gfx942_numeric.py" From 7c605807c34c72f378ea52c23039a5431974d44e Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Mon, 5 Oct 2026 11:58:40 +0000 Subject: [PATCH 06/20] feat(rocke): select stream-K conv wgrad through grouped conv dispatch Make stream-K reachable from the shipped selector instead of only from a hand-built spec: ConvGroupedRequest gains streamk and streamk_reduction, which pin split_k=1 and the direct epilogue and flow into the instance spec. The launch grid routes through the instance's own stream-K grid, so the spec the kernel is built from and the grid it is launched with cannot disagree about the CTA pool. The gfx1250 candidate turns a stream-K request away rather than silently serving the grid-per-tile kernel. The grouped-conv benchmark gains --streamk / --streamk-reduction and resets per launch exactly what each reduction's contract requires. ISSUE ID : AICK-1943 Co-Authored-By: Claude Opus 5 --- .../grouped_conv/benchmark_grouped_conv.py | 160 +++++++++++++++++- .../library/dispatch/grouped_convolution.py | 83 ++++++++- .../test_grouped_conv_wgrad_dispatch.py | 77 +++++++++ 3 files changed, 315 insertions(+), 5 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/grouped_conv/benchmark_grouped_conv.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/grouped_conv/benchmark_grouped_conv.py index 664851558c58..1554a949a809 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/grouped_conv/benchmark_grouped_conv.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/grouped_conv/benchmark_grouped_conv.py @@ -540,6 +540,26 @@ def main() -> int: choices=["fp16", "bf16", "fp32"], help="data type (default: fp16)", ) + parser.add_argument( + "--streamk", + default="off", + choices=["off", "dp_sk", "persistent"], + help=( + "wgrad only: ask the dispatcher for the stream-K kernel " + "(dp_sk = data-parallel + stream-K CTAs, persistent = one CTA per CU " + "round-robining tiles); off keeps the grid-per-tile kernel (default)" + ), + ) + parser.add_argument( + "--streamk-reduction", + default="linear", + dest="streamk_reduction", + choices=["linear", "tree", "workspace"], + help=( + "wgrad stream-K fixup: linear/tree (deterministic flag hand-off) or " + "workspace (f32 scratch atomics + reduce kernel) (default: linear)" + ), + ) parser.add_argument( "--warmup", type=int, default=3, help="warmup iterations (default: 3)" ) @@ -823,6 +843,10 @@ def _u8(t): dtype=dtype, direction=args.direction, ) + if args.streamk != "off": + req_base.update( + streamk=args.streamk, streamk_reduction=args.streamk_reduction + ) if problem.is_3d: req_base.update( Di=problem.Di, @@ -1105,6 +1129,114 @@ def _make(*shape): return 0 +def _streamk_launcher( + *, + rt, + spec, + kernel, + arch: str, + grid, + block, + stream, + ws_dev, + ws_nbytes: int, + buffers, + nbytes, + compile_kernel, + KernelLauncher, + LaunchConfig, +): + """``(launcher, kernel_name, launch(fence))`` for a stream-K wgrad spec. + + Each launch resets exactly what the reduction's contract requires before + the kernel runs: the flag table for linear/tree (dW is stored once by the + tile owner), the f32 scratch for the workspace reduction, which is then + cast into dW by the Stage 2 reduce kernel on the same stream. + """ + from kernels.common.conv_implicit_gemm_wgrad import ( + wgrad_streamk_signature, + wgrad_streamk_workspace_layout, + ) + from kernels.common.conv_wgrad_workspace_reduce import ( + WgradReduceSpec, + build_conv_wgrad_workspace_reduce, + wgrad_reduce_grid, + wgrad_reduce_signature, + ) + + dY_dev, X_dev, dW_dev = buffers + dY_nb, X_nb, dW_nb = nbytes + artifact = compile_kernel(kernel, arch=arch) + launcher = KernelLauncher( + hsaco=artifact.hsaco, + kernel_name=artifact.kernel_name, + signature=wgrad_streamk_signature(spec), + ) + values = { + "A": dY_dev, + "B": X_dev, + "D": dW_dev, + "A_bytes": dY_nb, + "B_bytes": X_nb, + "D_bytes": dW_nb, + } + reducer = None + if spec.streamk_reduction == "workspace": + values.update(ws_ptr=ws_dev, ws_bytes=ws_nbytes) + p = spec.problem + rspec = WgradReduceSpec( + problem=p, + dtype_d=spec.data.dtype_d, + groups=p.groups, + ws_replicas=spec.ws_replicas, + ) + rart = compile_kernel( + build_conv_wgrad_workspace_reduce(rspec, arch=arch), arch=arch + ) + reducer = KernelLauncher( + hsaco=rart.hsaco, + kernel_name=rart.kernel_name, + signature=wgrad_reduce_signature(rspec), + ) + rvalues = { + "ws_ptr": ws_dev, + "dw_ptr": dW_dev, + "wg_M": spec.wg_M, + "wg_N": spec.wg_N, + "ws_bytes": ws_nbytes, + "dw_bytes": dW_nb, + "groups": p.groups, + } + reset_bytes = ws_nbytes + else: + flags_bytes, _ = wgrad_streamk_workspace_layout(spec, arch=arch) + values.update(sk_flags=ws_dev, sk_partials=ws_dev + flags_bytes) + reset_bytes = flags_bytes + + def _launch(fence: bool): + if reset_bytes: + rt.memset(ws_dev, 0, reset_bytes) + launcher( + values, + config=LaunchConfig( + grid=grid, block=block, stream=stream, fence=fence and reducer is None + ), + ) + if reducer is not None: + reducer( + rvalues, + config=LaunchConfig( + grid=wgrad_reduce_grid(rspec), + block=(rspec.block_size, 1, 1), + stream=stream, + fence=fence, + ), + ) + + name = artifact.kernel_name + ("+cast" if reducer is not None else "") + return launcher, name, _launch + + def _run_wgrad( *, args, @@ -1132,6 +1264,9 @@ def _run_wgrad( _DEFAULT_TILE_M as _WGRAD_REDUCE_TILE_M, _DEFAULT_TILE_N as _WGRAD_REDUCE_TILE_N, ) + from kernels.common.conv_implicit_gemm_wgrad import ( + wgrad_streamk_workspace_nbytes, + ) _u8 = u8 p = problem @@ -1224,6 +1359,9 @@ def _make(*shape): launcher, ws_nbytes = build_implicit_gemm_conv_wgrad_two_stage( instance_spec, arch=arch ) + elif instance_spec.streamk != "off": + kernel = build_implicit_gemm_conv_wgrad(instance_spec, arch=arch) + ws_nbytes = wgrad_streamk_workspace_nbytes(instance_spec, arch=arch) else: kernel = build_implicit_gemm_conv_wgrad(instance_spec, arch=arch) ws_nbytes = 0 @@ -1242,7 +1380,27 @@ def _make(*shape): # factor from z entirely, so every grouped launch covered group 0 only. grid = _wgrad_grid(dspec, req) - if instance_spec.two_stage: + if instance_spec.streamk != "off": + if ws_dev is not None: + rt.free(ws_dev) + ws_dev = rt.alloc(max(ws_nbytes, 4)) + launcher, kernel_name, _launch = _streamk_launcher( + rt=rt, + spec=instance_spec, + kernel=kernel, + arch=arch, + grid=grid, + block=block, + stream=stream, + ws_dev=ws_dev, + ws_nbytes=ws_nbytes, + buffers=(dY_dev, X_dev, dW_dev), + nbytes=(dY_t.nbytes, X_t.nbytes, dW_t.nbytes), + compile_kernel=compile_kernel, + KernelLauncher=KernelLauncher, + LaunchConfig=LaunchConfig, + ) + elif instance_spec.two_stage: if ws_dev is not None: rt.free(ws_dev) ws_dev = rt.alloc(ws_nbytes) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py index d26a3cb510d0..bc30a068fd23 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py @@ -156,6 +156,7 @@ _DEFAULT_WS_REPLICAS as _WGRAD_WS_REPLICAS, is_valid_wgrad_spec as _wgrad_is_valid_spec, wgrad_atomic_epilogue_available as _wgrad_atomic_epilogue_available, + wgrad_streamk_grid as _wgrad_streamk_grid, ) from kernels.common.conv_implicit_gemm_dgrad import ( DgradConvSpec, @@ -275,6 +276,11 @@ class ConvGroupedRequest(OperatorRequest): dilation_d: Optional[int] = None # optional vec_size_c override; None = let the candidate decide vec_size_c: Optional[int] = None + # wgrad only: opt into stream-K ("dp_sk" | "persistent"); "off" keeps the + # grid-per-tile kernel with auto split-K. streamk_reduction picks how the + # CTAs sharing an output tile combine their partial sums. + streamk: str = "off" + streamk_reduction: str = "linear" op: str = "conv_grouped" algorithm: str = "auto" spec_id: str = "auto" @@ -369,6 +375,18 @@ def _request_errors(req: OperatorRequest) -> list[str]: errors.append(f"unsupported dtype {req.dtype!r}; fp16 or bf16 only") if req.layout.upper() != "NHWC": errors.append(f"unsupported layout {req.layout!r}; NHWC only") + if req.streamk not in ("off", "dp_sk", "persistent"): + errors.append( + f"streamk must be 'off', 'dp_sk' or 'persistent', got {req.streamk!r}" + ) + elif req.streamk != "off" and req.direction != "wgrad": + errors.append(f"streamk is wgrad-only, got direction={req.direction!r}") + if req.streamk_reduction not in ("linear", "tree", "workspace"): + # 'atomic' needs an fp32 dW, which a fp16/bf16 request never has. + errors.append( + "streamk_reduction must be 'linear', 'tree' or 'workspace', got " + f"{req.streamk_reduction!r}" + ) try: ArchTarget.from_gfx(req.arch) except KeyError as e: @@ -447,9 +465,29 @@ def _wgrad_grouped_overrides(req: ConvGroupedRequest) -> Tuple[str, int]: (WMMA grouped is handled by its own candidate, which forces default@split_k=1 -- WMMA has neither a cshuffle nor a split-K path.) """ + if req.streamk != "off": + # Stream-K owns the K range (no split-K) and stores each tile through + # the direct epilogue once per tile inside its tile loop. + return "default", 1 return _epilogue_for(req), -1 +def _wgrad_streamk_fields(req: ConvGroupedRequest) -> dict: + """The WgradConvSpec stream-K fields a request asks for. + + ``streamk_ctas`` stays -1 (one CTA per CU of the target, resolved by the + instance) so the spec, the kernel and the launch grid all derive the pool + from the same place. + """ + if req.streamk == "off": + return {} + return { + "streamk": req.streamk, + "streamk_reduction": req.streamk_reduction, + "streamk_ctas": -1, + } + + def _data_spec(req: ConvGroupedRequest) -> ConvDataSpec: return ConvDataSpec( dtype_a=req.dtype.lower(), @@ -493,6 +531,8 @@ class ConvGroupedSpec: arch: str split_k: int = 1 # wgrad only lds_k_outer: bool = False # wgrad only + streamk: str = "off" # wgrad only + streamk_reduction: str = "linear" # wgrad only name: str = "rocke_conv_grouped" def kernel_name(self) -> str: @@ -516,6 +556,9 @@ def kernel_name(self) -> str: # fetch rather than a transpose-on-store. if self.direction in ("wgrad", "dgrad") and self.lds_k_outer: parts.append("kouter") + # streamk: a different grid, K ranges and fixup epilogue. + if self.direction == "wgrad" and self.streamk != "off": + parts.append(f"sk_{self.streamk}_{self.streamk_reduction}") return kernel_name_join(self.name, *parts) def to_fwd_spec(self, problem: "ConvProblem") -> "ImplicitGemmConvSpec": @@ -553,7 +596,18 @@ def to_wgrad_spec(self, problem: "ConvProblem") -> "WgradConvSpec": """ assert self.direction == "wgrad", "to_wgrad_spec is only valid for wgrad specs" target = ArchTarget.from_gfx(self.arch) - resolved_split_k, two_stage, _ = _resolve_wgrad_split_k(self, problem) + if self.streamk != "off": + # Stream-K owns the K range: no split-K, no two-stage, and the + # direct epilogue runs per tile. + resolved_split_k, two_stage = 1, False + streamk = { + "streamk": self.streamk, + "streamk_reduction": self.streamk_reduction, + "streamk_ctas": -1, + } + else: + resolved_split_k, two_stage, _ = _resolve_wgrad_split_k(self, problem) + streamk = {} return WgradConvSpec( problem=problem, name=self.name, @@ -575,7 +629,8 @@ def to_wgrad_spec(self, problem: "ConvProblem") -> "WgradConvSpec": pipeline=self.pipeline, # two_stage writes one f32 per element via workspace store; cshuffle # is not used and would produce an invalid spec (validator rejects it). - epilogue="default" if two_stage else self.epilogue, + # Stream-K likewise stores through the direct epilogue. + epilogue="default" if (two_stage or streamk) else self.epilogue, split_k=resolved_split_k, two_stage=two_stage, # Pin the replica count explicitly rather than inheriting the @@ -583,6 +638,7 @@ def to_wgrad_spec(self, problem: "ConvProblem") -> "WgradConvSpec": # the i32 ws_bytes ABI using this same constant, and a cap computed # over a different R than the spec carries bounds nothing. ws_replicas=_WGRAD_WS_REPLICAS, + **streamk, ) def to_dgrad_spec(self, problem: "ConvProblem") -> "DgradConvSpec": @@ -722,6 +778,10 @@ def _resolve_wgrad_split_k( def _wgrad_grid(spec: ConvGroupedSpec, req: OperatorRequest) -> Tuple[int, int, int]: assert isinstance(req, ConvGroupedRequest) p = _problem(req) + if spec.streamk != "off": + # One-dimensional stream-K launch, sized from the same spec the kernel + # is built from so the two cannot disagree on the CTA pool. + return _wgrad_streamk_grid(spec.to_wgrad_spec(p), arch=spec.arch) spatial = (p.Z if p.is_3d else 1) * p.Y * p.X # Per-group GEMM dims. For the ungrouped path G==1 these reduce to wg_M=K, # wg_N=spatial*C. @@ -1126,6 +1186,7 @@ def _build_instance_spec(req: ConvGroupedRequest) -> WgradConvSpec: pipeline=_PIPELINE, epilogue=_ep, split_k=_sk, + **_wgrad_streamk_fields(req), ) def support(req: OperatorRequest) -> Tuple[bool, str]: @@ -1166,6 +1227,8 @@ def select(req: OperatorRequest) -> ConvGroupedSpec: dtype=req.dtype.lower(), arch=req.arch, split_k=_sk, + streamk=req.streamk, + streamk_reduction=req.streamk_reduction, name=name, ) @@ -1318,6 +1381,7 @@ def _build_instance_spec(req: ConvGroupedRequest) -> WgradConvSpec: pipeline=_pipeline(req), epilogue=_ep, split_k=_sk, + **_wgrad_streamk_fields(req), ) def support(req: OperatorRequest) -> Tuple[bool, str]: @@ -1359,6 +1423,8 @@ def select(req: OperatorRequest) -> ConvGroupedSpec: dtype=req.dtype.lower(), arch=req.arch, split_k=_sk, + streamk=req.streamk, + streamk_reduction=req.streamk_reduction, name=name, ) @@ -1531,7 +1597,9 @@ def _make_gfx1250_wgrad_candidate() -> KernelCandidate: test_gfx1250_grouped_wgrad_dual_engine). Group merging (``group_merge > 1``) is MFMA-only and additionally needs split_k > 1 on the two-stage path, so it is unreachable here on both counts: is_valid_wgrad_spec rejects it for - wave32 and for the split_k=1 this candidate forces. + wave32 and for the split_k=1 this candidate forces. Stream-K is MFMA-only + too, so a stream-K request is turned away rather than silently served by + the grid-per-tile kernel. """ name = "implicit_gemm_conv_wgrad_gfx1250" spec_id = "igemm_conv_wgrad_gfx1250_32x32" @@ -1578,6 +1646,8 @@ def support(req: OperatorRequest) -> Tuple[bool, str]: return False, f"gfx1250 candidate requires arch=gfx1250 (got {req.arch!r})" if req.direction != "wgrad": return False, f"candidate handles 'wgrad', got direction={req.direction!r}" + if req.streamk != "off": + return False, "stream-K wgrad is MFMA-only (gfx942, gfx950)" ok, why = selector_matches(req, candidate) if not ok: return False, why @@ -1770,7 +1840,12 @@ def dispatch_conv_grouped( f"spec_hash={kid.spec_hash}", f"request_hash={kid.request_hash}", ] - if req.direction == "wgrad": + if req.direction == "wgrad" and spec.streamk != "off": + explanation.append( + f"streamk={spec.streamk} reduction={spec.streamk_reduction} " + "(split-K off; one CTA per CU in the stream-K pool)" + ) + elif req.direction == "wgrad": split_k, _two_stage, requested = _resolve_wgrad_split_k(spec, _problem(req)) if split_k != requested: explanation.append( diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py index e494b70243a8..cd081b42adcb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py @@ -422,5 +422,82 @@ def test_ws_replicas_changes_the_instance_name(self): ) +class TestStreamKSelection(unittest.TestCase): + """Stream-K is opt-in on the request and reaches the instance spec intact.""" + + def test_streamk_reaches_the_instance_spec(self): + for arch in ("gfx942", "gfx950"): + for mode in ("dp_sk", "persistent"): + for reduction in ("linear", "tree", "workspace"): + r = dispatch_conv_grouped( + _wgrad(arch, G=4, streamk=mode, streamk_reduction=reduction) + ) + ws = r.spec.to_wgrad_spec(_problem(r.request)) + with self.subTest(arch=arch, mode=mode, reduction=reduction): + self.assertEqual(ws.streamk, mode) + self.assertEqual(ws.streamk_reduction, reduction) + self.assertEqual(ws.split_k, 1) + self.assertFalse(ws.two_stage) + self.assertEqual(ws.epilogue, "default") + + def test_default_request_is_unchanged(self): + r = dispatch_conv_grouped(_wgrad("gfx950", G=4)) + self.assertEqual(r.spec.streamk, "off") + self.assertEqual(r.spec.to_wgrad_spec(_problem(r.request)).streamk, "off") + + def test_streamk_builds_from_dispatch(self): + from kernels.common.conv_implicit_gemm_wgrad import ( + build_implicit_gemm_conv_wgrad, + ) + + r = dispatch_conv_grouped(_wgrad("gfx950", G=2, streamk="dp_sk")) + k = build_implicit_gemm_conv_wgrad( + r.spec.to_wgrad_spec(_problem(r.request)), arch="gfx950" + ) + self.assertIn("_sk256", k.name) + + def test_gfx1250_turns_streamk_away(self): + with self.assertRaises(ValueError) as cm: + dispatch_conv_grouped(_wgrad("gfx1250", streamk="dp_sk")) + self.assertIn("stream-K wgrad is MFMA-only", str(cm.exception)) + + def test_request_validation(self): + from dispatch.grouped_convolution import _request_errors + + self.assertTrue(_request_errors(_wgrad("gfx950", streamk="sideways"))) + self.assertTrue(_request_errors(_wgrad("gfx950", streamk_reduction="atomic"))) + fwd = ConvGroupedRequest( + N=2, C=64, K=64, Hi=8, Wi=8, Y=3, X=3, arch="gfx950", streamk="dp_sk" + ) + self.assertTrue(_request_errors(fwd)) + + +class TestStreamKGridShape(unittest.TestCase): + """The dispatch grid is the instance's stream-K grid, one-dimensional.""" + + def test_grid_matches_instance_partition(self): + from kernels.common.conv_implicit_gemm_wgrad import wgrad_streamk_grid + + for arch in ("gfx942", "gfx950"): + for mode in ("dp_sk", "persistent"): + r = dispatch_conv_grouped(_wgrad(arch, G=4, streamk=mode)) + ws = r.spec.to_wgrad_spec(_problem(r.request)) + with self.subTest(arch=arch, mode=mode): + self.assertEqual(r.grid, wgrad_streamk_grid(ws, arch=arch)) + self.assertEqual(r.grid[1:], (1, 1)) + + def test_streamk_changes_the_dispatch_name(self): + off = dispatch_conv_grouped(_wgrad("gfx950", G=4)).spec.kernel_name() + names = { + dispatch_conv_grouped( + _wgrad("gfx950", G=4, streamk=m, streamk_reduction=red) + ).spec.kernel_name() + for m in ("dp_sk", "persistent") + for red in ("linear", "tree") + } + self.assertEqual(len(names), 4) + self.assertNotIn(off, names) + + if __name__ == "__main__": unittest.main() From ed7bdf99453fabb5e87eaf284e9a6ac9e34f88e8 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Mon, 5 Oct 2026 11:58:40 +0000 Subject: [PATCH 07/20] docs(rocke): document stream-K conv wgrad Describe why wgrad is the direction that wants stream-K, the two launch modes, the four reductions and their determinism/zeroing contracts, the workspace layout, the validator constraints and the group fold. Add the three knobs to the runbook catalog, record the support matrix rows, and point the forward-conv future-work note at what landed. ISSUE ID : AICK-1943 Co-Authored-By: Claude Opus 5 --- .../common/README_conv_implicit_gemm.md | 7 ++ .../common/README_conv_implicit_gemm_wgrad.md | 66 +++++++++++++++++++ .../optimization/optimization_runbook.md | 3 + .../python/rocke/instances/SUPPORT_MATRIX.md | 2 + 4 files changed, 78 insertions(+) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm.md b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm.md index a02ac3da23ab..ccf11e0897f4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm.md @@ -233,6 +233,13 @@ the K loop across CTAs (similar to split-K but without the atomic reduction overhead), potentially delivering better utilisation for small or irregular shapes. +This has landed for **wgrad** (`streamk` on `WgradConvSpec`; see +`README_conv_implicit_gemm_wgrad.md`), where the lopsided `K_wg = N*Ho*Wo` +reduction makes it most valuable. The forward direction still maps one CTA per +output tile; the partitioner (`rocke.helpers.streamk.StreamKIterPartition`) and +the flag-based fixups are family-agnostic, so porting them here is wiring rather +than new machinery. + ### Wider vector loads for A The A descriptor emits one element offset per thread per load. Enabling vector diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md index 552458a92f48..d8286fcadca2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md @@ -245,6 +245,72 @@ Both loops that Python-unroll the K iteration — `pipeline="basic"` and `async_dma` — are bounded by `_MAX_UNROLLED_K_ITERS` (128). Over the cap the spec is rejected; raise `split_k` or `tile_k` rather than the constant. +### Stream-K (`streamk`, `streamk_reduction`, `streamk_ctas`) + +Split-K cuts `K_wg` into `split_k` equal slices and gives every output tile the +same number of them. That is the right shape when the tiles outnumber the CUs, +but a wgrad is usually the opposite: `wg_M x wg_N` is filter-sized while +`K_wg = N*Ho*Wo` is huge, so the tile count can be a handful and the leftover +CUs idle no matter how the slices are cut. + +Stream-K (CK Tile's `StreamKTilePartitioner`) balances *MAC iterations* instead. +Output tiles that divide the CTA pool evenly run data-parallel — one CTA, the +whole K loop, no reduction at all — and only the remainder is spread over +`streamk_ctas` stream-K CTAs, each owning a contiguous, balanced range of +iterations that may straddle tile boundaries. rocke builds one kernel per shape, +so the whole partition is resolved host-side +(`rocke.helpers.streamk.StreamKIterPartition`) and folded in as constants; the +device only maps `blockIdx.x` onto its iteration range. + +`streamk="dp_sk"` launches `dp_tiles + sk_ctas` CTAs. `streamk="persistent"` +launches only the pool: each CTA round-robins the data-parallel tiles +(`tile = cta + t*pool`, CK's static `StreamKDispatch`, no work queue) and then +runs its stream-K range. Both walk their work with one bounded `scf.for` whose +trip count is CTA-uniform, so the barriers inside the GEMM body stay uniform and +the body is emitted once. + +`streamk_reduction` picks how the CTAs sharing a tile combine partial sums: + +| | Determinism | dW zeroing | Workspace | +|---|---|---|---| +| `linear` | yes | not needed | flags + one f32 tile per SK CTA | +| `tree` | yes | not needed | same | +| `atomic` | no | required | none (fp32 dW only) | +| `workspace` | no | not needed | the two-stage f32 scratch + Stage 2 reduce | + +`linear` and `tree` hand partials over explicitly: a non-owner stores its tile +to its own slot, issues a release fence and sets a per-wave flag; the consumer +spins on an acquire load of that flag and then adds. Both are deterministic +because the fold order is fixed by the partition, not by arrival. `linear` has +the owner walk its later contributors in turn; `tree` fans in pairwise over +`ceil(log2)` rounds. The flag protocol is three core IR ops (`memref.fence`, +`memref.global_flag_store`, `memref.global_flag_wait_eq`) rather than raw cache +hints, so the AMDGPU memory model supplies the L2 writeback/invalidate that +makes a partial visible across XCDs. + +Size the buffer with `wgrad_streamk_workspace_nbytes`; split it with +`wgrad_streamk_workspace_layout` (`sk_flags` at the base, `sk_partials` at +`base + flags_bytes`). **Only the flag region has to be zeroed before each +launch** — partials are always written before they are read. + +Constraints, all enforced by `validate()` / `is_valid_wgrad_spec` in both +engines: gfx942/gfx950 (MFMA wave64) only; `split_k=1` and no `two_stage` +(stream-K owns the K range); `epilogue="default"`; no `async_dma` / `unroll_k` +(the per-tile trip count is a runtime value); no `chiplet_swizzle` or +`group_merge` (both remap tiles, which stream-K owns). Conv groups are folded +into GEMM-M instead of riding `block_id_z`: `m_tiles = ceil(wg_M/tile_m) * +groups`, decoded group-fastest. Defining the tile count per group makes the fold +exact for any `wg_M`, avoiding CK's `GemmM % MPerBlock` caveat. + +`streamk_ctas=-1` resolves to one CTA per CU of the target. The pool must not +exceed what the device can hold at once: a stream-K CTA waits on CTAs with +higher indices, so one that is not resident would deadlock it. + +Reachable from dispatch: `ConvGroupedRequest(streamk=..., streamk_reduction=...)`, +which pins `split_k=1` and the direct epilogue and routes `_wgrad_grid` through +the instance's own `wgrad_streamk_grid`, so the spec and the launch grid cannot +disagree about the pool. + ## Next steps ### K0-M-K1 LDS layout diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md index cd95558cfc26..b819d35a38eb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md @@ -1991,6 +1991,9 @@ Per-arch facts belong in the arch references — see §21. | `reduction` | `StreamKGemmSpec` | `Atomic` | `StreamKReductionStrategy.{Atomic, Reduction, AtomicWithFixup}` | | split-K | GEMM (not yet a spec field; available primitives) | — | `b.global_atomic_add_f32` for atomic split-K; `helpers/streamk.py` for the Stream-K macro tile decode | | Grouped GEMM persistent | `instances/grouped_gemm.py` | per-group launches | v2 persistent variant is a documented follow-up | +| `streamk` | `WgradConvSpec` | `"off"` | `"dp_sk"` launches one CTA per data-parallel tile plus `streamk_ctas` stream-K CTAs; `"persistent"` launches only the pool and round-robins the data-parallel tiles. Both balance MAC iterations across CTAs instead of tiles, so a shape with few tiles and a long `wg_K` still fills the device. Mutually exclusive with `split_k` | +| `streamk_reduction` | `WgradConvSpec` | `"linear"` | How CTAs sharing a tile combine partials. `"linear"`: the owner waits on each later contributor in turn and stores once. `"tree"`: pairwise fan-in, `ceil(log2)` rounds. Both are deterministic and need no dW pre-zeroing. `"atomic"`: f32 atomics into dW (fp32 only). `"workspace"`: f32 scratch atomics plus the Stage 2 reduce kernel | +| `streamk_ctas` | `WgradConvSpec` | `-1` | The CTA pool the stream-K remainder is spread over (CK `max_active_wgs`). `-1` resolves to one CTA per CU of the target. A pool larger than the device can hold at once would let a waiting CTA block one it waits on | #### 12.1.J Quantization diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md index ce43548f0038..fc98ee807914 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/SUPPORT_MATRIX.md @@ -73,6 +73,8 @@ See the [packed FP6 input contract](../examples/gfx1250/gemm/FP6.md). | `conv_implicit_gemm_auto` | ✅ | ✅ | ❌ | MFMA-specialized autotuned path (raw `MfmaAtom`, K=32 kpack); not ported to WMMA | | `direct_conv_16c` | ❌ | ✅ | ❌ | `fold_k32` needs 16x16x32 atom (CDNA4) | | `direct_conv_4c` | ✅ | ✅ | ❌ | 4x4x4 MFMA atom not in WMMA catalog | +| `conv_implicit_gemm_wgrad` | ✅ | ✅ | ✅ | gfx1151/gfx1201/gfx1250: WMMA, `default` epilogue, `split_k=1` | +| `conv_implicit_gemm_wgrad` stream-K | ✅ | ✅ | ❌ | CK iteration-balanced partition; `linear`/`tree` deterministic fixups need the cross-workgroup flag ops | --- From 91d94614781ba0085d07b86e4b09a482f96b8b24 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Mon, 5 Oct 2026 14:41:55 +0000 Subject: [PATCH 08/20] perf(rocke): size the stream-K wgrad CTA pool per reduction Resolve streamk_ctas=-1 from the reduction instead of always using one CTA per CU. A CU-sized pool leaves a single workgroup per CU, so a stream-K CTA stalled on memory has nothing to hide behind, while the split-K grid it competes with keeps several resident. The workspace and atomic reductions never wait on another CTA, so their pool can be sized for throughput: the same waves-per-CU target the split-K heuristic already sizes its grid for, divided by the waves per CTA. linear and tree keep one CTA per CU: a CTA there waits on later CTAs, so the whole pool must be co-resident, and those fixups get slower with more contributors per tile because the owner folds each one. Only the -1 default changes. Explicit pools, every parity config and every golden are unaffected. ISSUE ID : AICK-1943 Co-Authored-By: Claude Opus 5.5 --- .../common/README_conv_implicit_gemm_wgrad.md | 31 ++++++++++---- .../common/conv_implicit_gemm_wgrad.py | 41 ++++++++++++++----- .../library/tests/test_conv_wgrad_streamk.py | 17 ++++++++ .../tests/test_conv_wgrad_streamk_gpu.py | 19 +++++++-- .../optimization/optimization_runbook.md | 3 +- 5 files changed, 87 insertions(+), 24 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md index d8286fcadca2..6854045a08ab 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md @@ -302,14 +302,29 @@ into GEMM-M instead of riding `block_id_z`: `m_tiles = ceil(wg_M/tile_m) * groups`, decoded group-fastest. Defining the tile count per group makes the fold exact for any `wg_M`, avoiding CK's `GemmM % MPerBlock` caveat. -`streamk_ctas=-1` resolves to one CTA per CU of the target. The pool must not -exceed what the device can hold at once: a stream-K CTA waits on CTAs with -higher indices, so one that is not resident would deadlock it. - -Reachable from dispatch: `ConvGroupedRequest(streamk=..., streamk_reduction=...)`, -which pins `split_k=1` and the direct epilogue and routes `_wgrad_grid` through -the instance's own `wgrad_streamk_grid`, so the spec and the launch grid cannot -disagree about the pool. +`streamk_ctas=-1` resolves per reduction (`wgrad_streamk_default_ctas`): + +* `workspace` / `atomic` never wait on another CTA, so their pool is sized for + throughput: `num_cus * WGRAD_TARGET_WAVES_PER_CU / waves_per_cta`, the same + occupancy target the split-K heuristic sizes its grid for (4 CTAs per CU for + a 4-wave CTA). One CTA per CU leaves a single workgroup per CU with nothing + to hide memory latency behind, which is what made a CU-sized pool lose to + split-K on most shapes. +* `linear` / `tree` stay at one CTA per CU. A CTA there waits on CTAs with + higher indices, so the whole pool must be resident at once; and these fixups + get slower with a larger pool anyway, because the tile owner folds every + additional contributor. + +Reachable from dispatch: `ConvGroupedRequest.streamk` defaults to `"auto"`, +which selects stream-K (with the `workspace` reduction) for few output tiles on +a long reduction -- at most `_WGRAD_STREAMK_MAX_TILES` tiles of at least +`_WGRAD_STREAMK_MIN_ITERS_PER_TILE` K iterations each -- and the grid-per-tile +split-K kernel otherwise. Elsewhere split-K's equal slices already balance the +grid and stream-K's per-tile fixup is pure overhead. `"off"` never selects it; +`"dp_sk"` / `"persistent"` force it. A stream-K choice pins `split_k=1` and the +direct epilogue and routes `_wgrad_grid` through the instance's own +`wgrad_streamk_grid`, so the spec and the launch grid cannot disagree about the +pool. ## Next steps diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py index c3ec86a23cfd..3c6d53fbc28b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py @@ -502,9 +502,10 @@ class WgradConvSpec: # zeroed before each launch. streamk_reduction: str = "linear" # The CTA pool the stream-K remainder is spread over (CK max_active_wgs). - # -1 = one per CU of the target, resolved by the builder. Every stream-K - # CTA must be resident alongside the CTAs it waits on, so an explicit value - # must not exceed what the device can hold at once. + # -1 = auto, resolved by the builder via wgrad_streamk_default_ctas: sized + # to the occupancy target for the reductions that never wait, one per CU + # for linear/tree. Those two wait on later CTAs, so with them an explicit + # value must not exceed what the device can hold at once. streamk_ctas: int = -1 @property @@ -1105,15 +1106,31 @@ def wgrad_streamk_available(spec: "WgradConvSpec") -> Tuple[bool, str]: return True, "ok" -def wgrad_streamk_default_ctas(arch: str) -> int: - """Default stream-K CTA pool: one CTA per CU of ``arch``. +def wgrad_streamk_default_ctas(spec: "WgradConvSpec", arch: str) -> int: + """Default stream-K CTA pool for ``spec`` on ``arch`` (``streamk_ctas=-1``). - One resident workgroup per CU is always available, which is what keeps a - stream-K CTA that waits on a later one from deadlocking. + The reductions that never wait on another CTA (``workspace``, ``atomic``) + are sized for throughput: enough CTAs to reach the same waves-per-CU target + the split-K heuristic sizes its grid for, i.e. ``num_cus * + WGRAD_TARGET_WAVES_PER_CU / waves_per_cta``. One CTA per CU leaves a single + workgroup per CU with nothing to hide memory latency behind. + + ``linear`` and ``tree`` wait on later CTAs, so every CTA in the pool must be + resident at once or a waiting CTA can block the one it waits on. One CTA + per CU is always resident, and those fixups get slower, not faster, with + more contributors per tile (the owner folds every partial). """ - from rocke.helpers.split_k import _ARCH_NUM_CUS, _DEFAULT_NUM_CUS + from rocke.helpers.split_k import ( + _ARCH_NUM_CUS, + _DEFAULT_NUM_CUS, + WGRAD_TARGET_WAVES_PER_CU, + ) - return _ARCH_NUM_CUS.get(arch, _DEFAULT_NUM_CUS) + num_cus = _ARCH_NUM_CUS.get(arch, _DEFAULT_NUM_CUS) + if spec.streamk_reduction in ("linear", "tree"): + return num_cus + waves_per_cta = max(1, spec.block_size // spec.wave_size) + return num_cus * max(1, WGRAD_TARGET_WAVES_PER_CU // waves_per_cta) def wgrad_streamk_partition(spec: "WgradConvSpec", *, arch: str = "gfx950"): @@ -1129,7 +1146,9 @@ def wgrad_streamk_partition(spec: "WgradConvSpec", *, arch: str = "gfx950"): raise ValueError("wgrad_streamk_partition requires streamk != 'off'") p = spec.problem ctas = ( - spec.streamk_ctas if spec.streamk_ctas > 0 else wgrad_streamk_default_ctas(arch) + spec.streamk_ctas + if spec.streamk_ctas > 0 + else wgrad_streamk_default_ctas(spec, arch) ) return streamk_iter_partition( m_tiles=-(-_wg_M(p) // spec.tile_m) * p.groups, @@ -1680,7 +1699,7 @@ def build_implicit_gemm_conv_wgrad( # Resolve streamk_ctas=-1 (auto) the same way, so the kernel name and the # partition both carry the concrete pool size. if spec.streamk != "off" and spec.streamk_ctas == -1: - spec = dc_replace(spec, streamk_ctas=wgrad_streamk_default_ctas(arch)) + spec = dc_replace(spec, streamk_ctas=wgrad_streamk_default_ctas(spec, arch)) spec.validate() ok, why = is_valid_wgrad_spec(spec, arch=arch) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py index 328882a6d61b..2beea55691e8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py @@ -113,6 +113,23 @@ def test_name_tags_every_knob_that_changes_the_body(self): self.assertTrue(_spec().kernel_name().endswith("_sk6")) self.assertIn("_skauto", _spec(streamk_ctas=-1).kernel_name()) + def test_auto_pool_depends_on_the_reduction(self): + from kernels.common.conv_implicit_gemm_wgrad import wgrad_streamk_default_ctas + + # linear/tree wait on later CTAs: one per CU, always resident. + for red in ("linear", "tree"): + spec = _spec(streamk_ctas=-1, streamk_reduction=red) + self.assertEqual(wgrad_streamk_default_ctas(spec, "gfx950"), 256) + self.assertEqual(wgrad_streamk_default_ctas(spec, "gfx942"), 304) + # workspace/atomic never wait: sized to 16 waves per CU (4-wave CTAs). + for red, dd in (("workspace", "fp16"), ("atomic", "fp32")): + spec = _spec(streamk_ctas=-1, streamk_reduction=red, dtype_d=dd) + self.assertEqual(wgrad_streamk_default_ctas(spec, "gfx950"), 1024) + self.assertEqual(wgrad_streamk_default_ctas(spec, "gfx942"), 1216) + # An 8-wave CTA halves the CTA count for the same occupancy target. + wide = _spec(streamk_ctas=-1, streamk_reduction="workspace", warp_m=4) + self.assertEqual(wgrad_streamk_default_ctas(wide, "gfx950"), 512) + def test_builder_resolves_auto_pool_into_the_name(self): k = build_implicit_gemm_conv_wgrad(_spec(streamk_ctas=-1), arch="gfx950") self.assertTrue(k.name.endswith("_sk256"), k.name) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk_gpu.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk_gpu.py index 33b21acdf5a2..89df9ddd45d6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk_gpu.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk_gpu.py @@ -270,10 +270,21 @@ def test_modes_and_reductions(self): _check(self, spec) def test_default_pool(self): - # streamk_ctas=-1 resolves to one CTA per CU of the device. - problem, _ = _problems()["dp_plus_sk"] - spec = _spec(problem, mode="dp_sk", reduction="linear", dtype_d="fp16", ctas=-1) - _check(self, spec) + # streamk_ctas=-1: one CTA per CU for linear/tree, the occupancy-sized + # pool (several CTAs per CU) for the reductions that never wait. + for name in ("dp_plus_sk", "multi_partner"): + problem, _ = _problems()[name] + for mode in ("dp_sk", "persistent"): + for reduction, dtype_d in _REDUCTIONS: + spec = _spec( + problem, + mode=mode, + reduction=reduction, + dtype_d=dtype_d, + ctas=-1, + ) + with self.subTest(shape=name, mode=mode, reduction=reduction): + _check(self, spec) @unittest.skipUnless(not _SKIP_REASON, _SKIP_REASON or "no GPU") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md index b819d35a38eb..9d33b9ea3fe5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md @@ -1993,7 +1993,8 @@ Per-arch facts belong in the arch references — see §21. | Grouped GEMM persistent | `instances/grouped_gemm.py` | per-group launches | v2 persistent variant is a documented follow-up | | `streamk` | `WgradConvSpec` | `"off"` | `"dp_sk"` launches one CTA per data-parallel tile plus `streamk_ctas` stream-K CTAs; `"persistent"` launches only the pool and round-robins the data-parallel tiles. Both balance MAC iterations across CTAs instead of tiles, so a shape with few tiles and a long `wg_K` still fills the device. Mutually exclusive with `split_k` | | `streamk_reduction` | `WgradConvSpec` | `"linear"` | How CTAs sharing a tile combine partials. `"linear"`: the owner waits on each later contributor in turn and stores once. `"tree"`: pairwise fan-in, `ceil(log2)` rounds. Both are deterministic and need no dW pre-zeroing. `"atomic"`: f32 atomics into dW (fp32 only). `"workspace"`: f32 scratch atomics plus the Stage 2 reduce kernel | -| `streamk_ctas` | `WgradConvSpec` | `-1` | The CTA pool the stream-K remainder is spread over (CK `max_active_wgs`). `-1` resolves to one CTA per CU of the target. A pool larger than the device can hold at once would let a waiting CTA block one it waits on | +| `streamk_ctas` | `WgradConvSpec` | `-1` | The CTA pool the stream-K remainder is spread over (CK `max_active_wgs`). `-1` resolves per reduction: `workspace`/`atomic` (no waiting) get `num_cus * WGRAD_TARGET_WAVES_PER_CU / waves_per_cta`; `linear`/`tree` get one CTA per CU, because a pool larger than the device can hold at once would let a waiting CTA block one it waits on | +| `streamk` (dispatch) | `ConvGroupedRequest` | `"auto"` | Picks stream-K + `workspace` when the candidate's tiling yields at most `_WGRAD_STREAMK_MAX_TILES` output tiles of at least `_WGRAD_STREAMK_MIN_ITERS_PER_TILE` K iterations; split-K otherwise. `"off"` disables, `"dp_sk"`/`"persistent"` force | #### 12.1.J Quantization From d738d30346e29c601bdfebefd530fa0e94849b97 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Mon, 5 Oct 2026 14:41:56 +0000 Subject: [PATCH 09/20] feat(rocke): let grouped conv dispatch choose stream-K for wgrad ConvGroupedRequest.streamk now defaults to "auto". The dispatcher picks stream-K when the candidate's tiling yields few output tiles over a long reduction, and keeps the grid-per-tile split-K kernel everywhere else. In that regime split-K can only cut the reduction into equal slices per tile, while stream-K balances iterations across the whole CTA pool. With many tiles or a short reduction, split-K is already balanced and stream-K's per-tile fixup is overhead. Both thresholds are coarse on purpose, so the rule does not encode the corpus it was fitted on. streamk_reduction defaults to "auto", which resolves to the workspace reduction: it never waits on another CTA, so it gets the occupancy-sized pool. "off" restores the previous selection exactly; "dp_sk"/"persistent" and explicit reductions still force a choice. One resolver feeds every wgrad candidate's instance spec, dispatch spec and grid, so they cannot disagree on whether the kernel is stream-K. gfx1250 treats "auto" as off and still rejects a forced stream-K request. The benchmark follows the dispatcher default and accepts the new values. ISSUE ID : AICK-1943 Co-Authored-By: Claude Opus 5.5 --- .../grouped_conv/benchmark_grouped_conv.py | 19 +-- .../library/dispatch/grouped_convolution.py | 138 +++++++++++++----- .../test_grouped_conv_wgrad_dispatch.py | 44 +++++- 3 files changed, 154 insertions(+), 47 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/grouped_conv/benchmark_grouped_conv.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/grouped_conv/benchmark_grouped_conv.py index 1554a949a809..41f1f225ff24 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/grouped_conv/benchmark_grouped_conv.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/grouped_conv/benchmark_grouped_conv.py @@ -542,22 +542,23 @@ def main() -> int: ) parser.add_argument( "--streamk", - default="off", - choices=["off", "dp_sk", "persistent"], + default="auto", + choices=["auto", "off", "dp_sk", "persistent"], help=( - "wgrad only: ask the dispatcher for the stream-K kernel " - "(dp_sk = data-parallel + stream-K CTAs, persistent = one CTA per CU " - "round-robining tiles); off keeps the grid-per-tile kernel (default)" + "wgrad only: auto (default) lets the dispatcher pick stream-K where it " + "is expected to win; off never uses it; dp_sk (data-parallel + " + "stream-K CTAs) / persistent (a fixed CTA pool round-robining tiles) " + "force it" ), ) parser.add_argument( "--streamk-reduction", - default="linear", + default="auto", dest="streamk_reduction", - choices=["linear", "tree", "workspace"], + choices=["auto", "linear", "tree", "workspace"], help=( "wgrad stream-K fixup: linear/tree (deterministic flag hand-off) or " - "workspace (f32 scratch atomics + reduce kernel) (default: linear)" + "workspace (f32 scratch atomics + reduce kernel); auto = workspace" ), ) parser.add_argument( @@ -843,7 +844,7 @@ def _u8(t): dtype=dtype, direction=args.direction, ) - if args.streamk != "off": + if args.direction == "wgrad": req_base.update( streamk=args.streamk, streamk_reduction=args.streamk_reduction ) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py index bc30a068fd23..92aea8abf5d5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py @@ -276,11 +276,13 @@ class ConvGroupedRequest(OperatorRequest): dilation_d: Optional[int] = None # optional vec_size_c override; None = let the candidate decide vec_size_c: Optional[int] = None - # wgrad only: opt into stream-K ("dp_sk" | "persistent"); "off" keeps the - # grid-per-tile kernel with auto split-K. streamk_reduction picks how the - # CTAs sharing an output tile combine their partial sums. - streamk: str = "off" - streamk_reduction: str = "linear" + # wgrad only: "auto" (default) picks stream-K for the shapes it is measured + # to win on (_wgrad_streamk_auto) and the grid-per-tile kernel with auto + # split-K otherwise; "off" never uses stream-K; "dp_sk" / "persistent" + # force it. streamk_reduction picks how the CTAs sharing an output tile + # combine their partial sums; "auto" is the f32 workspace reduction. + streamk: str = "auto" + streamk_reduction: str = "auto" op: str = "conv_grouped" algorithm: str = "auto" spec_id: str = "auto" @@ -375,17 +377,18 @@ def _request_errors(req: OperatorRequest) -> list[str]: errors.append(f"unsupported dtype {req.dtype!r}; fp16 or bf16 only") if req.layout.upper() != "NHWC": errors.append(f"unsupported layout {req.layout!r}; NHWC only") - if req.streamk not in ("off", "dp_sk", "persistent"): + if req.streamk not in ("off", "auto", "dp_sk", "persistent"): errors.append( - f"streamk must be 'off', 'dp_sk' or 'persistent', got {req.streamk!r}" + "streamk must be 'off', 'auto', 'dp_sk' or 'persistent', got " + f"{req.streamk!r}" ) - elif req.streamk != "off" and req.direction != "wgrad": + elif req.streamk in ("dp_sk", "persistent") and req.direction != "wgrad": errors.append(f"streamk is wgrad-only, got direction={req.direction!r}") - if req.streamk_reduction not in ("linear", "tree", "workspace"): + if req.streamk_reduction not in ("auto", "linear", "tree", "workspace"): # 'atomic' needs an fp32 dW, which a fp16/bf16 request never has. errors.append( - "streamk_reduction must be 'linear', 'tree' or 'workspace', got " - f"{req.streamk_reduction!r}" + "streamk_reduction must be 'auto', 'linear', 'tree' or 'workspace', " + f"got {req.streamk_reduction!r}" ) try: ArchTarget.from_gfx(req.arch) @@ -454,7 +457,9 @@ def _epilogue_for(req: ConvGroupedRequest) -> str: return "cshuffle" if _vec_size_c(req) > 1 else "default" -def _wgrad_grouped_overrides(req: ConvGroupedRequest) -> Tuple[str, int]: +def _wgrad_grouped_overrides( + req: ConvGroupedRequest, streamk: str = "off" +) -> Tuple[str, int]: """(epilogue, split_k) for a wgrad spec. Grouping is orthogonal to the epilogue and split-K: every epilogue (direct, @@ -464,28 +469,79 @@ def _wgrad_grouped_overrides(req: ConvGroupedRequest) -> Tuple[str, int]: same vec-derived epilogue and auto split-K formula (split_k=-1) as ungrouped. (WMMA grouped is handled by its own candidate, which forces default@split_k=1 -- WMMA has neither a cshuffle nor a split-K path.) + + ``streamk`` is the resolved mode from :func:`_wgrad_streamk_choice`. """ - if req.streamk != "off": + if streamk != "off": # Stream-K owns the K range (no split-K) and stores each tile through # the direct epilogue once per tile inside its tile loop. return "default", 1 return _epilogue_for(req), -1 -def _wgrad_streamk_fields(req: ConvGroupedRequest) -> dict: - """The WgradConvSpec stream-K fields a request asks for. +# Auto-selection thresholds for stream-K wgrad, in units of the candidate's +# own tiles. Stream-K pays off when a handful of output tiles sit on a long +# reduction: split-K can only cut that reduction into equal slices per tile, +# while stream-K balances iterations across the whole CTA pool. With many +# tiles, or a short reduction, the grid-per-tile kernel is already balanced and +# stream-K's per-tile fixup is pure overhead. Fitted on a benchmark corpus of +# production wgrad shapes on gfx950 and kept deliberately coarse (two +# thresholds) so it does not encode that corpus. +_WGRAD_STREAMK_MAX_TILES = 128 +_WGRAD_STREAMK_MIN_ITERS_PER_TILE = 1024 + + +def _wgrad_streamk_auto( + req: ConvGroupedRequest, tile_m: int, tile_n: int, tile_k: int +) -> bool: + """Whether ``streamk="auto"`` should pick stream-K for this request.""" + p = _problem(req) + spatial = (p.Z if p.is_3d else 1) * p.Y * p.X + wg_M = p.K // p.groups + wg_N = spatial * (p.C // p.groups) + wg_K = p.N * p.Ho * p.Wo * (p.Do if p.is_3d else 1) + tiles = -(-wg_M // tile_m) * -(-wg_N // tile_n) * p.groups + iters_per_tile = -(-wg_K // tile_k) + return ( + tiles <= _WGRAD_STREAMK_MAX_TILES + and iters_per_tile >= _WGRAD_STREAMK_MIN_ITERS_PER_TILE + ) + + +def _wgrad_streamk_choice( + req: ConvGroupedRequest, tile_m: int, tile_n: int, tile_k: int +) -> Tuple[str, str]: + """Resolve the request's ``(streamk, streamk_reduction)`` to concrete values. + + The one resolver every wgrad candidate's instance spec, dispatch spec and + grid go through, so they cannot disagree on whether the kernel is stream-K. + """ + mode = req.streamk + if mode == "auto": + mfma = req.arch in ("gfx942", "gfx950") + mode = ( + "dp_sk" + if mfma and _wgrad_streamk_auto(req, tile_m, tile_n, tile_k) + else "off" + ) + # The f32 workspace reduction never waits on another CTA, so its pool can be + # sized for occupancy; it is also the fastest of the reductions measured. + reduction = ( + "workspace" if req.streamk_reduction == "auto" else req.streamk_reduction + ) + return mode, reduction - ``streamk_ctas`` stays -1 (one CTA per CU of the target, resolved by the - instance) so the spec, the kernel and the launch grid all derive the pool + +def _wgrad_streamk_fields(mode: str, reduction: str) -> dict: + """The WgradConvSpec stream-K fields for a resolved choice. + + ``streamk_ctas`` stays -1 (sized by the instance from the arch and the + reduction) so the spec, the kernel and the launch grid all derive the pool from the same place. """ - if req.streamk == "off": + if mode == "off": return {} - return { - "streamk": req.streamk, - "streamk_reduction": req.streamk_reduction, - "streamk_ctas": -1, - } + return {"streamk": mode, "streamk_reduction": reduction, "streamk_ctas": -1} def _data_spec(req: ConvGroupedRequest) -> ConvDataSpec: @@ -1169,7 +1225,8 @@ def _tile(req: ConvGroupedRequest): def _build_instance_spec(req: ConvGroupedRequest) -> WgradConvSpec: tm, tn, tk, wm, wn, wtmn, wtk = _tile(req) - _ep, _sk = _wgrad_grouped_overrides(req) + _skm, _skr = _wgrad_streamk_choice(req, tm, tn, tk) + _ep, _sk = _wgrad_grouped_overrides(req, _skm) return WgradConvSpec( problem=_problem(req), name=name, @@ -1186,7 +1243,7 @@ def _build_instance_spec(req: ConvGroupedRequest) -> WgradConvSpec: pipeline=_PIPELINE, epilogue=_ep, split_k=_sk, - **_wgrad_streamk_fields(req), + **_wgrad_streamk_fields(_skm, _skr), ) def support(req: OperatorRequest) -> Tuple[bool, str]: @@ -1212,7 +1269,8 @@ def select(req: OperatorRequest) -> ConvGroupedSpec: raise ValueError(f"{name} does not support request: {why}") assert isinstance(req, ConvGroupedRequest) tm, tn, tk, wm, wn, wtmn, wtk = _tile(req) - _ep, _sk = _wgrad_grouped_overrides(req) + _skm, _skr = _wgrad_streamk_choice(req, tm, tn, tk) + _ep, _sk = _wgrad_grouped_overrides(req, _skm) return ConvGroupedSpec( direction="wgrad", tile_m=tm, @@ -1227,8 +1285,8 @@ def select(req: OperatorRequest) -> ConvGroupedSpec: dtype=req.dtype.lower(), arch=req.arch, split_k=_sk, - streamk=req.streamk, - streamk_reduction=req.streamk_reduction, + streamk=_skm, + streamk_reduction=_skr, name=name, ) @@ -1363,7 +1421,8 @@ def _pipeline(req: ConvGroupedRequest) -> str: def _build_instance_spec(req: ConvGroupedRequest) -> WgradConvSpec: tm, tn, tk, wm, wn, wtmn, wtk = _tile(req) - _ep, _sk = _wgrad_grouped_overrides(req) + _skm, _skr = _wgrad_streamk_choice(req, tm, tn, tk) + _ep, _sk = _wgrad_grouped_overrides(req, _skm) return WgradConvSpec( problem=_problem(req), name=name, @@ -1381,7 +1440,7 @@ def _build_instance_spec(req: ConvGroupedRequest) -> WgradConvSpec: pipeline=_pipeline(req), epilogue=_ep, split_k=_sk, - **_wgrad_streamk_fields(req), + **_wgrad_streamk_fields(_skm, _skr), ) def support(req: OperatorRequest) -> Tuple[bool, str]: @@ -1407,7 +1466,8 @@ def select(req: OperatorRequest) -> ConvGroupedSpec: raise ValueError(f"{name} does not support request: {why}") assert isinstance(req, ConvGroupedRequest) tm, tn, tk, wm, wn, wtmn, wtk = _tile(req) - _ep, _sk = _wgrad_grouped_overrides(req) + _skm, _skr = _wgrad_streamk_choice(req, tm, tn, tk) + _ep, _sk = _wgrad_grouped_overrides(req, _skm) return ConvGroupedSpec( direction="wgrad", tile_m=tm, @@ -1423,8 +1483,8 @@ def select(req: OperatorRequest) -> ConvGroupedSpec: dtype=req.dtype.lower(), arch=req.arch, split_k=_sk, - streamk=req.streamk, - streamk_reduction=req.streamk_reduction, + streamk=_skm, + streamk_reduction=_skr, name=name, ) @@ -1646,7 +1706,7 @@ def support(req: OperatorRequest) -> Tuple[bool, str]: return False, f"gfx1250 candidate requires arch=gfx1250 (got {req.arch!r})" if req.direction != "wgrad": return False, f"candidate handles 'wgrad', got direction={req.direction!r}" - if req.streamk != "off": + if req.streamk in ("dp_sk", "persistent"): return False, "stream-K wgrad is MFMA-only (gfx942, gfx950)" ok, why = selector_matches(req, candidate) if not ok: @@ -1841,9 +1901,15 @@ def dispatch_conv_grouped( f"request_hash={kid.request_hash}", ] if req.direction == "wgrad" and spec.streamk != "off": + why = ( + f"auto: <= {_WGRAD_STREAMK_MAX_TILES} tiles over >= " + f"{_WGRAD_STREAMK_MIN_ITERS_PER_TILE} K iterations each" + if req.streamk == "auto" + else "requested" + ) explanation.append( - f"streamk={spec.streamk} reduction={spec.streamk_reduction} " - "(split-K off; one CTA per CU in the stream-K pool)" + f"streamk={spec.streamk} reduction={spec.streamk_reduction} ({why}; " + "split-K off)" ) elif req.direction == "wgrad": split_k, _two_stage, requested = _resolve_wgrad_split_k(spec, _problem(req)) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py index cd081b42adcb..bda754adba19 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py @@ -440,11 +440,48 @@ def test_streamk_reaches_the_instance_spec(self): self.assertFalse(ws.two_stage) self.assertEqual(ws.epilogue, "default") - def test_default_request_is_unchanged(self): + def test_auto_keeps_short_k_shapes_on_split_k(self): + # Default request: streamk="auto". A short reduction (wg_K = 2*14*14) + # never reaches the iteration threshold, so the split-K kernel stays. r = dispatch_conv_grouped(_wgrad("gfx950", G=4)) + self.assertEqual(r.request.streamk, "auto") self.assertEqual(r.spec.streamk, "off") self.assertEqual(r.spec.to_wgrad_spec(_problem(r.request)).streamk, "off") + def test_auto_picks_streamk_for_few_tiles_long_k(self): + # Few output tiles (K=96, C=24, 3x3) on a long reduction (N*Ho*Wo is + # 128*96*96): the shape stream-K is selected for. + for arch in ("gfx942", "gfx950"): + r = dispatch_conv_grouped(_wgrad(arch, N=128, C=24, K=96, Hi=96, Wi=96)) + with self.subTest(arch=arch): + self.assertEqual(r.spec.streamk, "dp_sk") + self.assertEqual(r.spec.streamk_reduction, "workspace") + self.assertEqual(r.spec.split_k, 1) + self.assertTrue(any("streamk=dp_sk" in e for e in r.explanation)) + + def test_auto_respects_thresholds(self): + from dispatch.grouped_convolution import ( + _WGRAD_STREAMK_MAX_TILES, + _WGRAD_STREAMK_MIN_ITERS_PER_TILE, + _wgrad_streamk_auto, + ) + + long_k = _wgrad("gfx950", N=128, C=24, K=96, Hi=96, Wi=96) + self.assertTrue(_wgrad_streamk_auto(long_k, 64, 64, 64)) + # Same shape with a tile_k that leaves fewer iterations than the floor. + big_tk = -(-(128 * 96 * 96) // (_WGRAD_STREAMK_MIN_ITERS_PER_TILE - 1)) + self.assertFalse(_wgrad_streamk_auto(long_k, 64, 64, big_tk)) + # Many tiles: 16 M tiles x 144 N tiles is far above the ceiling. + many = _wgrad("gfx950", N=128, C=1024, K=1024, Hi=96, Wi=96) + self.assertGreater(16 * 144, _WGRAD_STREAMK_MAX_TILES) + self.assertFalse(_wgrad_streamk_auto(many, 64, 64, 64)) + + def test_off_and_gfx1250_never_pick_streamk(self): + req = _wgrad("gfx950", N=128, C=24, K=96, Hi=96, Wi=96, streamk="off") + self.assertEqual(dispatch_conv_grouped(req).spec.streamk, "off") + r = dispatch_conv_grouped(_wgrad("gfx1250", N=128, C=24, K=96, Hi=96, Wi=96)) + self.assertEqual(r.spec.streamk, "off") + def test_streamk_builds_from_dispatch(self): from kernels.common.conv_implicit_gemm_wgrad import ( build_implicit_gemm_conv_wgrad, @@ -454,7 +491,9 @@ def test_streamk_builds_from_dispatch(self): k = build_implicit_gemm_conv_wgrad( r.spec.to_wgrad_spec(_problem(r.request)), arch="gfx950" ) - self.assertIn("_sk256", k.name) + # streamk_reduction defaults to the workspace reduction, whose pool is + # sized to the occupancy target: 256 CUs x (16 waves / 4 per CTA). + self.assertIn("_sk1024_skrworkspace", k.name) def test_gfx1250_turns_streamk_away(self): with self.assertRaises(ValueError) as cm: @@ -466,6 +505,7 @@ def test_request_validation(self): self.assertTrue(_request_errors(_wgrad("gfx950", streamk="sideways"))) self.assertTrue(_request_errors(_wgrad("gfx950", streamk_reduction="atomic"))) + self.assertFalse(_request_errors(_wgrad("gfx950", streamk="auto"))) fwd = ConvGroupedRequest( N=2, C=64, K=64, Hi=8, Wi=8, Y=3, X=3, arch="gfx950", streamk="dp_sk" ) From 7484126ccfaf36e20e2ca9f218aa9286b1d55f1a Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Mon, 5 Oct 2026 15:16:53 +0000 Subject: [PATCH 10/20] perf(rocke): count the stream-K wgrad pool in CTAs per CU Size the default pool of the workspace and atomic reductions at four CTAs per CU instead of deriving it from a waves-per-CU target. Each resident CTA streams its own K range, so what hides memory latency is the number of independent reduction streams per CU, not the number of waves. A wider CTA adds no streams, and the waves-based rule halved the pool for 8-wave CTAs, where the larger pool is measurably better. For 4-wave CTAs the two rules agree. linear and tree keep their residency-safe pool of one CTA per CU. ISSUE ID : AICK-1943 Co-Authored-By: Claude Opus 5.5 --- .../common/README_conv_implicit_gemm_wgrad.md | 11 ++++---- .../common/conv_implicit_gemm_wgrad.py | 25 +++++++++++-------- .../test_grouped_conv_wgrad_dispatch.py | 2 +- .../library/tests/test_conv_wgrad_streamk.py | 6 ++--- .../optimization/optimization_runbook.md | 2 +- 5 files changed, 25 insertions(+), 21 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md index 6854045a08ab..a3b4d3262207 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md @@ -305,11 +305,12 @@ exact for any `wg_M`, avoiding CK's `GemmM % MPerBlock` caveat. `streamk_ctas=-1` resolves per reduction (`wgrad_streamk_default_ctas`): * `workspace` / `atomic` never wait on another CTA, so their pool is sized for - throughput: `num_cus * WGRAD_TARGET_WAVES_PER_CU / waves_per_cta`, the same - occupancy target the split-K heuristic sizes its grid for (4 CTAs per CU for - a 4-wave CTA). One CTA per CU leaves a single workgroup per CU with nothing - to hide memory latency behind, which is what made a CU-sized pool lose to - split-K on most shapes. + throughput: four CTAs per CU (`_STREAMK_CTAS_PER_CU`). Each resident CTA + streams its own K range, so this counts independent reduction streams per CU, + and is counted in CTAs rather than waves: a wider CTA adds no streams. One + CTA per CU leaves a single stream per CU with nothing to hide memory latency + behind, which is what made a CU-sized pool lose to split-K on most shapes; + past four the CTAs mostly contend for the same caches. * `linear` / `tree` stay at one CTA per CU. A CTA there waits on CTAs with higher indices, so the whole pool must be resident at once; and these fixups get slower with a larger pool anyway, because the tile owner folds every diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py index 3c6d53fbc28b..75086d8e7251 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py @@ -1106,31 +1106,34 @@ def wgrad_streamk_available(spec: "WgradConvSpec") -> Tuple[bool, str]: return True, "ok" +# CTAs per CU in the default pool of the stream-K reductions that never wait. +# Each resident CTA streams its own K range, so this is the number of +# independent reduction streams a CU keeps in flight to hide memory latency. +# Counted in CTAs rather than waves: a wider CTA does not add streams. One per +# CU leaves a single stream per CU; past four the CTAs mostly contend for the +# same caches. +_STREAMK_CTAS_PER_CU = 4 + + def wgrad_streamk_default_ctas(spec: "WgradConvSpec", arch: str) -> int: """Default stream-K CTA pool for ``spec`` on ``arch`` (``streamk_ctas=-1``). The reductions that never wait on another CTA (``workspace``, ``atomic``) - are sized for throughput: enough CTAs to reach the same waves-per-CU target - the split-K heuristic sizes its grid for, i.e. ``num_cus * - WGRAD_TARGET_WAVES_PER_CU / waves_per_cta``. One CTA per CU leaves a single - workgroup per CU with nothing to hide memory latency behind. + are sized for throughput: ``_STREAMK_CTAS_PER_CU`` CTAs per CU. Nothing + there requires the pool to be resident at once, so a CTA that does not fit + only runs later. ``linear`` and ``tree`` wait on later CTAs, so every CTA in the pool must be resident at once or a waiting CTA can block the one it waits on. One CTA per CU is always resident, and those fixups get slower, not faster, with more contributors per tile (the owner folds every partial). """ - from rocke.helpers.split_k import ( - _ARCH_NUM_CUS, - _DEFAULT_NUM_CUS, - WGRAD_TARGET_WAVES_PER_CU, - ) + from rocke.helpers.split_k import _ARCH_NUM_CUS, _DEFAULT_NUM_CUS num_cus = _ARCH_NUM_CUS.get(arch, _DEFAULT_NUM_CUS) if spec.streamk_reduction in ("linear", "tree"): return num_cus - waves_per_cta = max(1, spec.block_size // spec.wave_size) - return num_cus * max(1, WGRAD_TARGET_WAVES_PER_CU // waves_per_cta) + return num_cus * _STREAMK_CTAS_PER_CU def wgrad_streamk_partition(spec: "WgradConvSpec", *, arch: str = "gfx950"): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py index bda754adba19..9e3a7de7d79a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py @@ -492,7 +492,7 @@ def test_streamk_builds_from_dispatch(self): r.spec.to_wgrad_spec(_problem(r.request)), arch="gfx950" ) # streamk_reduction defaults to the workspace reduction, whose pool is - # sized to the occupancy target: 256 CUs x (16 waves / 4 per CTA). + # sized at four CTAs per CU: 256 CUs x 4. self.assertIn("_sk1024_skrworkspace", k.name) def test_gfx1250_turns_streamk_away(self): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py index 2beea55691e8..892f22465e78 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py @@ -121,14 +121,14 @@ def test_auto_pool_depends_on_the_reduction(self): spec = _spec(streamk_ctas=-1, streamk_reduction=red) self.assertEqual(wgrad_streamk_default_ctas(spec, "gfx950"), 256) self.assertEqual(wgrad_streamk_default_ctas(spec, "gfx942"), 304) - # workspace/atomic never wait: sized to 16 waves per CU (4-wave CTAs). + # workspace/atomic never wait: four CTAs per CU. for red, dd in (("workspace", "fp16"), ("atomic", "fp32")): spec = _spec(streamk_ctas=-1, streamk_reduction=red, dtype_d=dd) self.assertEqual(wgrad_streamk_default_ctas(spec, "gfx950"), 1024) self.assertEqual(wgrad_streamk_default_ctas(spec, "gfx942"), 1216) - # An 8-wave CTA halves the CTA count for the same occupancy target. + # Counted in CTAs, not waves: an 8-wave CTA gets the same pool. wide = _spec(streamk_ctas=-1, streamk_reduction="workspace", warp_m=4) - self.assertEqual(wgrad_streamk_default_ctas(wide, "gfx950"), 512) + self.assertEqual(wgrad_streamk_default_ctas(wide, "gfx950"), 1024) def test_builder_resolves_auto_pool_into_the_name(self): k = build_implicit_gemm_conv_wgrad(_spec(streamk_ctas=-1), arch="gfx950") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md index 9d33b9ea3fe5..c5da7100bf33 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md @@ -1993,7 +1993,7 @@ Per-arch facts belong in the arch references — see §21. | Grouped GEMM persistent | `instances/grouped_gemm.py` | per-group launches | v2 persistent variant is a documented follow-up | | `streamk` | `WgradConvSpec` | `"off"` | `"dp_sk"` launches one CTA per data-parallel tile plus `streamk_ctas` stream-K CTAs; `"persistent"` launches only the pool and round-robins the data-parallel tiles. Both balance MAC iterations across CTAs instead of tiles, so a shape with few tiles and a long `wg_K` still fills the device. Mutually exclusive with `split_k` | | `streamk_reduction` | `WgradConvSpec` | `"linear"` | How CTAs sharing a tile combine partials. `"linear"`: the owner waits on each later contributor in turn and stores once. `"tree"`: pairwise fan-in, `ceil(log2)` rounds. Both are deterministic and need no dW pre-zeroing. `"atomic"`: f32 atomics into dW (fp32 only). `"workspace"`: f32 scratch atomics plus the Stage 2 reduce kernel | -| `streamk_ctas` | `WgradConvSpec` | `-1` | The CTA pool the stream-K remainder is spread over (CK `max_active_wgs`). `-1` resolves per reduction: `workspace`/`atomic` (no waiting) get `num_cus * WGRAD_TARGET_WAVES_PER_CU / waves_per_cta`; `linear`/`tree` get one CTA per CU, because a pool larger than the device can hold at once would let a waiting CTA block one it waits on | +| `streamk_ctas` | `WgradConvSpec` | `-1` | The CTA pool the stream-K remainder is spread over (CK `max_active_wgs`). `-1` resolves per reduction: `workspace`/`atomic` (no waiting) get four CTAs per CU (`_STREAMK_CTAS_PER_CU`, counted in CTAs, not waves); `linear`/`tree` get one CTA per CU, because a pool larger than the device can hold at once would let a waiting CTA block one it waits on | | `streamk` (dispatch) | `ConvGroupedRequest` | `"auto"` | Picks stream-K + `workspace` when the candidate's tiling yields at most `_WGRAD_STREAMK_MAX_TILES` output tiles of at least `_WGRAD_STREAMK_MIN_ITERS_PER_TILE` K iterations; split-K otherwise. `"off"` disables, `"dp_sk"`/`"persistent"` force | #### 12.1.J Quantization From e69bb4c50761fb91058701f7e6cb31aa7696f0f7 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Tue, 6 Oct 2026 10:57:13 +0000 Subject: [PATCH 11/20] style(rocke): clang-format stream-K wgrad sources ISSUE ID : AICK-1943 Co-Authored-By: Claude Opus 5.5 --- .../rocke/platform/cpp/instances/common/conv_abi.cpp | 4 ++-- .../cpp/instances/common/conv_implicit_gemm_wgrad.cpp | 9 +++------ 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_abi.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_abi.cpp index 9e0e8165115c..76ae1370edcb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_abi.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_abi.cpp @@ -169,8 +169,8 @@ const int kNumStreamKScalars = (int)(sizeof(kStreamKScalars) / sizeof(kStreamKSc bool streamk_reduction_known(const char* red) { - return strcmp(red, "atomic") == 0 || strcmp(red, "workspace") == 0 - || strcmp(red, "linear") == 0 || strcmp(red, "tree") == 0; + return strcmp(red, "atomic") == 0 || strcmp(red, "workspace") == 0 || strcmp(red, "linear") == 0 + || strcmp(red, "tree") == 0; } void wgrad_arg_names(rocke_conv_arg_list_t* out, bool is_3d, bool two_stage, const char* streamk) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp index 69becddb2624..a7bf710f5ff3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp @@ -2707,10 +2707,8 @@ static rocke_value_t* wgrad_sk_cta_of(wgrad_streamk_env_t* e, rocke_value_t* it) } /* _store_tile: the plain dW store of a finished tile. */ -static void wgrad_sk_store_tile(wgrad_streamk_env_t* e, - rocke_value_t* const* vals, - int wg_M, - int wg_N) +static void + wgrad_sk_store_tile(wgrad_streamk_env_t* e, rocke_value_t* const* vals, int wg_M, int wg_N) { rocke_conv_build_ctx_t* ctx = e->ctx; for(int i = 0; i < ctx->num_accs; ++i) @@ -2971,8 +2969,7 @@ static void wgrad_emit_streamk(rocke_conv_build_ctx_t* ctx, for(int i = 0; i < n_acc; ++i) yields[i] = vals[i]; yields[n_acc] = rocke_b_mul(b, stride, c_two); - yields[n_acc + 1] - = rocke_b_select(b, send, c_one, rounds.iter_vars[n_acc + 1]); + yields[n_acc + 1] = rocke_b_select(b, send, c_one, rounds.iter_vars[n_acc + 1]); rocke_b_scf_yield(b, yields, n_acc + 2); } rocke_b_region_leave(b); From 082f58426070c759084e98083d5ed10de2d56d9a Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Tue, 6 Oct 2026 12:36:02 +0000 Subject: [PATCH 12/20] fix(rocke): harden stream-K wgrad validation and HIP flag lowering Address review feedback on the stream-K wgrad port: - Reject a stream-K launch shape that changes what the kernel baked in (grouping, 2-D/3-D, pointwise addressing, channel vector widths), including on the workspace-size path. - Carry group_merge into the C++ spec and gate it with the same reason as Python, so both engines reject group_merge > 1 alike. - Make the C++ HIP lowerer reject unknown fence/flag scope and ordering values with ROCKE_ERR_VALUE instead of substituting a default. - Correct the docs: a linear/tree pool need not be resident, since CTAs wait only on later CTAs and publish before waiting; and the C builder accepts streamk_ctas=-1. ISSUE ID : AICK-1943 Co-Authored-By: Claude Opus 5.5 --- .../common/README_conv_implicit_gemm_wgrad.md | 11 ++-- .../common/conv_implicit_gemm_wgrad.py | 64 +++++++++++++++--- .../library/tests/test_conv_wgrad_streamk.py | 28 ++++++++ .../platform/cpp/bindings/rocke_engine.cpp | 1 + .../cpp/core/lower_hip/lower_hip_mem.cpp | 66 +++++++++++++++---- .../rocke/instance_conv_implicit_gemm_wgrad.h | 8 ++- .../common/conv_implicit_gemm_wgrad.cpp | 13 ++++ .../optimization/optimization_runbook.md | 2 +- .../platform/python/rocke/helpers/streamk.py | 7 +- .../tests/core/future_intrinsic_lowering.cpp | 42 ++++++++++++ 10 files changed, 214 insertions(+), 28 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md index 701ebf1c971b..d261bbc98bcf 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md @@ -315,10 +315,13 @@ exact for any `wg_M`, avoiding CK's `GemmM % MPerBlock` caveat. CTA per CU leaves a single stream per CU with nothing to hide memory latency behind, which is what made a CU-sized pool lose to split-K on most shapes; past four the CTAs mostly contend for the same caches. -* `linear` / `tree` stay at one CTA per CU. A CTA there waits on CTAs with - higher indices, so the whole pool must be resident at once; and these fixups - get slower with a larger pool anyway, because the tile owner folds every - additional contributor. +* `linear` / `tree` stay at one CTA per CU, because these fixups get slower + with a larger pool: the tile owner folds every additional contributor. The + pool does not have to be resident at once. A CTA waits only on CTAs with + higher indices and publishes its own partial before it waits, so with + workgroups dispatched in ascending order only the CTAs of the one tile that + straddles the dispatch front (its owner for `linear`, about one per round + for `tree`) can wait on a CTA that has not started yet. Reachable from dispatch: `ConvGroupedRequest.streamk` defaults to `"auto"`, which selects stream-K (with the `workspace` reduction) for few output tiles on diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py index cd9cc6eff181..a7b4e99edb18 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py @@ -594,8 +594,9 @@ class WgradConvSpec: # The CTA pool the stream-K remainder is spread over (CK max_active_wgs). # -1 = auto, resolved by the builder via wgrad_streamk_default_ctas: sized # to the occupancy target for the reductions that never wait, one per CU - # for linear/tree. Those two wait on later CTAs, so with them an explicit - # value must not exceed what the device can hold at once. + # for linear/tree. Any explicit value > 0 is safe: those two wait only on + # later CTAs, which ascending dispatch always starts (see + # wgrad_streamk_default_ctas). streamk_ctas: int = -1 @property @@ -1202,10 +1203,13 @@ def wgrad_streamk_default_ctas(spec: "WgradConvSpec", arch: str) -> int: there requires the pool to be resident at once, so a CTA that does not fit only runs later. - ``linear`` and ``tree`` wait on later CTAs, so every CTA in the pool must be - resident at once or a waiting CTA can block the one it waits on. One CTA - per CU is always resident, and those fixups get slower, not faster, with - more contributors per tile (the owner folds every partial). + ``linear`` and ``tree`` get one CTA per CU: those fixups get slower, not + faster, with more contributors per tile (the owner folds every partial). + Their pool need not be resident at once. A CTA waits only on higher-index + CTAs and publishes its own partial before it waits, so with workgroups + dispatched in ascending order only the CTAs of the tile straddling the + dispatch front can wait on one that has not started, and they hold a + bounded number of slots (one for linear, about one per round for tree). """ from rocke.helpers.split_k import _ARCH_NUM_CUS, _DEFAULT_NUM_CUS @@ -1238,11 +1242,55 @@ def wgrad_streamk_launch(spec: "WgradConvSpec", *, arch: str = "gfx950"): ) +def _wgrad_streamk_launch_problem(spec: WgradConvSpec, problem=None) -> ConvProblem: + """The launch shape, checked against what the kernel baked in. + + One stream-K binary serves any shape that keeps its build-time modes: + grouping, 2-D/3-D (the kernarg list), pointwise addressing, and channel + runs at least as wide as the built load/store vectors. Anything else would + launch the wrong code, so it is rejected rather than silently mis-run. + """ + built = spec.problem + if problem is None: + return built + + def _vec(chan: int, dtype: str, cap: int | None) -> int: + widths = (8, 4, 2, 1) if dtype != "fp32" else (4, 2, 1) + auto = next(v for v in widths if chan % v == 0) + return auto if cap is None else min(auto, cap) + + checks = ( + ("grouping (groups > 1)", built.groups > 1, problem.groups > 1), + ("is_3d", built.is_3d, problem.is_3d), + ("is_pointwise", built.is_pointwise, problem.is_pointwise), + ) + for what, want, got in checks: + if want != got: + raise ValueError( + f"launch problem is incompatible with the built kernel: {what} " + f"differs (built {want}, launch {got}); build a kernel for this shape" + ) + chans = ( + ("kpg", built.kpg, problem.kpg, spec.data.dtype_a, spec.vector_size_a), + ("cpg", built.cpg, problem.cpg, spec.data.dtype_b, spec.vector_size_b), + ("cpg", built.cpg, problem.cpg, spec.data.dtype_d, spec.vector_size_c), + ) + for what, want, got, dtype, cap in chans: + width = _vec(want, dtype, cap) + if got % width: + raise ValueError( + f"launch problem is incompatible with the built kernel: {what}=" + f"{got} is not a multiple of the built {width}-wide vector " + f"({what}={want}); build a kernel for this shape" + ) + return problem + + def _wgrad_streamk_args(spec: "WgradConvSpec", problem=None): from kernels.common.conv_args import ConvArgs return ConvArgs.from_problem( - problem if problem is not None else spec.problem, + _wgrad_streamk_launch_problem(spec, problem), direction="wgrad", tile_m=spec.tile_m, tile_n=spec.tile_n, @@ -1304,7 +1352,7 @@ def wgrad_streamk_workspace_nbytes( if spec.streamk == "off" or spec.streamk_reduction == "atomic": return 0 if spec.streamk_reduction == "workspace": - p = problem if problem is not None else spec.problem + p = _wgrad_streamk_launch_problem(spec, problem) return p.groups * spec.ws_replicas * _wg_M(p) * _wg_N(p) * 4 flags_bytes, partials_bytes = wgrad_streamk_workspace_layout( spec, arch=arch, problem=problem diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py index 61bef58945ba..c39116f1a240 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py @@ -217,6 +217,14 @@ def test_cpp_rejects_with_the_same_reason(self): eng.conv_wgrad_lower_llvm(_spec_dict(spec), "gfx950") self.assertIn(why, str(cm.exception)) + def test_cpp_rejects_group_merge_with_the_same_reason(self): + eng = _engine() + spec = _spec(_DW, group_merge=2) + _ok, why = is_valid_wgrad_spec(spec, "gfx950") + with self.assertRaises(RuntimeError) as cm: + eng.conv_wgrad_lower_llvm(_spec_dict(spec), "gfx950") + self.assertIn(why, str(cm.exception)) + class TestStreamKAdmitsImpliesBuilds(unittest.TestCase): def test_matrix(self): @@ -422,6 +430,26 @@ def test_launch_shape_overrides_the_spec_problem(self): wgrad_streamk_partition(spec, problem=other), wgrad_streamk_partition(spec) ) + def test_launch_shape_must_keep_the_built_modes(self): + # The kernel bakes grouping, 2-D/3-D, pointwise addressing and its + # channel vector widths; a launch shape that changes any is rejected. + base = {"N": 2, "Hi": 12, "Wi": 12, "Y": 3, "X": 3, "pH": 1, "pW": 1} + spec = _spec() # _DENSE: groups=1, 2-D, 3x3, kpg=96, cpg=32 + for needle, other in ( + ("grouping", ConvProblem(C=32, K=96, groups=2, **base)), + ("is_3d", ConvProblem(C=32, K=96, Di=4, Z=1, sD=1, pD=0, dD=1, **base)), + ("is_pointwise", ConvProblem(N=2, Hi=12, Wi=12, C=32, K=96, Y=1, X=1)), + ("kpg=36", ConvProblem(C=32, K=36, **base)), + ("cpg=20", ConvProblem(C=20, K=96, **base)), + ): + with self.subTest(needle=needle): + for helper in (wgrad_streamk_partition, wgrad_streamk_workspace_nbytes): + with self.assertRaises(ValueError) as cm: + helper(spec, problem=other) + self.assertIn(needle, str(cm.exception)) + # Wider channel runs than the build keep the same code path. + wgrad_streamk_partition(spec, problem=ConvProblem(C=64, K=192, **base)) + class TestStreamKDualEngine(unittest.TestCase): """The C++ builder lowers every stream-K variant to the Python engine's bytes.""" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp index e2263708e866..2292097ce3a5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp @@ -2173,6 +2173,7 @@ rocke_implicit_gemm_conv_wgrad_spec_t conv_wgrad_build_spec(const py::dict& d, s.two_stage = dict_bool(d, "two_stage", s.two_stage); s.ws_replicas = dict_int(d, "ws_replicas", s.ws_replicas); s.streamk_ctas = dict_int(d, "streamk_ctas", s.streamk_ctas); + s.group_merge = dict_int(d, "group_merge", s.group_merge); s.async_dma = dict_bool(d, "async_dma", s.async_dma); s.unroll_k = dict_bool(d, "unroll_k", s.unroll_k); s.lds_k_outer = dict_bool(d, "lds_k_outer", s.lds_k_outer); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp index 2d34472e106a..95b5315c954d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp @@ -697,7 +697,8 @@ static rocke_status_t _op_memref_global_atomic_add(rocke_h_lowerer_t* lw, const return lw->status; } -/* Python _HIP_ATOMIC_ORDER / _HIP_FENCE_SCOPE / _HIP_MEMORY_SCOPE. */ +/* Python _HIP_ATOMIC_ORDER / _HIP_FENCE_SCOPE / _HIP_MEMORY_SCOPE. Unknown + * spellings return NULL so the ops reject them, like Python's table lookup. */ static const char* hip_atomic_order(const char* o) { if(!strcmp(o, "monotonic")) @@ -708,12 +709,20 @@ static const char* hip_atomic_order(const char* o) return "__ATOMIC_RELEASE"; if(!strcmp(o, "acq_rel")) return "__ATOMIC_ACQ_REL"; - return "__ATOMIC_SEQ_CST"; + if(!strcmp(o, "seq_cst")) + return "__ATOMIC_SEQ_CST"; + return NULL; } static const char* hip_fence_scope(const char* s) { - return !strcmp(s, "system") ? "" : s; + if(!strcmp(s, "workgroup")) + return "workgroup"; + if(!strcmp(s, "agent")) + return "agent"; + if(!strcmp(s, "system")) + return ""; + return NULL; } static const char* hip_memory_scope(const char* s) @@ -722,26 +731,39 @@ static const char* hip_memory_scope(const char* s) return "__HIP_MEMORY_SCOPE_WORKGROUP"; if(!strcmp(s, "agent")) return "__HIP_MEMORY_SCOPE_AGENT"; - return "__HIP_MEMORY_SCOPE_SYSTEM"; + if(!strcmp(s, "system")) + return "__HIP_MEMORY_SCOPE_SYSTEM"; + return NULL; } /* Python _op_memref_fence */ static rocke_status_t _op_memref_fence(rocke_h_lowerer_t* lw, const rocke_op_t* op) { + const char *ord_s, *scp_s, *ord, *scp; if(!rocke_h_live(lw)) { return lw->status; } - rocke_h_emitf(lw, - "__builtin_amdgcn_fence(%s, \"%s\");", - hip_atomic_order(mem_attr_str(op, "ordering", "acq_rel")), - hip_fence_scope(mem_attr_str(op, "scope", "agent"))); + ord_s = mem_attr_str(op, "ordering", "acq_rel"); + scp_s = mem_attr_str(op, "scope", "agent"); + ord = hip_atomic_order(ord_s); + scp = hip_fence_scope(scp_s); + if(!ord) + { + return rocke_h_fail(lw, ROCKE_ERR_VALUE, "memref.fence: unknown ordering '%s'", ord_s); + } + if(!scp) + { + return rocke_h_fail(lw, ROCKE_ERR_VALUE, "memref.fence: unknown scope '%s'", scp_s); + } + rocke_h_emitf(lw, "__builtin_amdgcn_fence(%s, \"%s\");", ord, scp); return lw->status; } /* Python _op_memref_global_flag_store */ static rocke_status_t _op_memref_global_flag_store(rocke_h_lowerer_t* lw, const rocke_op_t* op) { + const char *ord_s, *scp_s, *ord, *scp; if(!rocke_h_live(lw)) { return lw->status; @@ -750,19 +772,34 @@ static rocke_status_t _op_memref_global_flag_store(rocke_h_lowerer_t* lw, const { return rocke_h_fail(lw, ROCKE_ERR_VALUE, "memref.global_flag_store: too few operands"); } + ord_s = mem_attr_str(op, "ordering", "release"); + scp_s = mem_attr_str(op, "scope", "agent"); + ord = hip_atomic_order(ord_s); + scp = hip_memory_scope(scp_s); + if(!ord) + { + return rocke_h_fail( + lw, ROCKE_ERR_VALUE, "memref.global_flag_store: unknown ordering '%s'", ord_s); + } + if(!scp) + { + return rocke_h_fail( + lw, ROCKE_ERR_VALUE, "memref.global_flag_store: unknown scope '%s'", scp_s); + } rocke_h_emitf(lw, "__hip_atomic_store(&%s[%s], %s, %s, %s);", rocke_h_name(lw, op->operands[0]), rocke_h_name(lw, op->operands[1]), rocke_h_name(lw, op->operands[2]), - hip_atomic_order(mem_attr_str(op, "ordering", "release")), - hip_memory_scope(mem_attr_str(op, "scope", "agent"))); + ord, + scp); return lw->status; } /* Python _op_memref_global_flag_wait_eq */ static rocke_status_t _op_memref_global_flag_wait_eq(rocke_h_lowerer_t* lw, const rocke_op_t* op) { + const char *scp_s, *scp; if(!rocke_h_live(lw)) { return lw->status; @@ -771,11 +808,18 @@ static rocke_status_t _op_memref_global_flag_wait_eq(rocke_h_lowerer_t* lw, cons { return rocke_h_fail(lw, ROCKE_ERR_VALUE, "memref.global_flag_wait_eq: too few operands"); } + scp_s = mem_attr_str(op, "scope", "agent"); + scp = hip_memory_scope(scp_s); + if(!scp) + { + return rocke_h_fail( + lw, ROCKE_ERR_VALUE, "memref.global_flag_wait_eq: unknown scope '%s'", scp_s); + } rocke_h_emitf(lw, "while (__hip_atomic_load(&%s[%s], __ATOMIC_ACQUIRE, %s) != %s) {}", rocke_h_name(lw, op->operands[0]), rocke_h_name(lw, op->operands[1]), - hip_memory_scope(mem_attr_str(op, "scope", "agent")), + scp, rocke_h_name(lw, op->operands[2])); return lw->status; } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h index b3a138b82152..f79474ca8797 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h @@ -183,13 +183,19 @@ typedef struct rocke_implicit_gemm_conv_wgrad_spec * streamk "off" | "dp_sk" | "persistent" * streamk_reduction "linear" | "tree" | "atomic" | "workspace" * streamk_ctas CTA pool the stream-K remainder is spread over; - * -1 (auto) is rejected by the builder, like split_k=-1. + * a launch parameter, not part of the kernel, so the + * builder accepts -1 (auto) (unlike split_k=-1); the + * host helpers need it resolved to > 0. * "off" keeps every existing config byte-identical. Stream-K folds the * conv groups into GEMM-M, so it is the one path on which this port * accepts groups > 1. */ const char* streamk; /* default "off" */ const char* streamk_reduction; /* default "linear" */ int streamk_ctas; /* default -1 */ + + /* WgradConvSpec.group_merge. Only 1 is ported: stream-K folds the + * groups into GEMM-M itself, and the grid-per-group path is Python-only. */ + int group_merge; /* default 1 */ } rocke_implicit_gemm_conv_wgrad_spec_t; /* Default-constructed spec (every field == Python dataclass default). */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp index a7bf710f5ff3..1e4da42a7667 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp @@ -77,6 +77,7 @@ rocke_implicit_gemm_conv_wgrad_spec_t rocke_implicit_gemm_conv_wgrad_spec_defaul s.streamk = "off"; s.streamk_reduction = "linear"; s.streamk_ctas = -1; + s.group_merge = 1; return s; } @@ -473,6 +474,9 @@ static bool wgrad_streamk_gates(const rocke_implicit_gemm_conv_wgrad_spec_t* s, if(s->chiplet_swizzle) SK_REJECT("chiplet_swizzle remaps the 2-D tile grid; stream-K decodes its tiles " "from a linear CTA index instead"); + if(s->group_merge > 1) + SK_REJECT("streamk folds conv groups into GEMM-M; use group_merge=1 (got %d)", + s->group_merge); const char* dd = s->dtype_d ? s->dtype_d : "fp16"; if(strcmp(red, "atomic") == 0 && strcmp(dd, "fp32") != 0) SK_REJECT("streamk_reduction='atomic' needs dtype_d='fp32' (got '%s'); use " @@ -517,6 +521,15 @@ bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv } if(!wgrad_streamk_gates(s, arch, reason, reason_cap)) return false; + if(s->group_merge != 1 && !sk_on) + { + if(reason && reason_cap) + snprintf(reason, + reason_cap, + "group_merge=%d is not supported for wgrad in the C port", + s->group_merge); + return false; + } if(s->problem.groups > 1 && rocke_conv_problem_is_pointwise(&s->problem)) { if(reason && reason_cap) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md index 791a13e5e582..e215fbf7bc96 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md @@ -1993,7 +1993,7 @@ Per-arch facts belong in the arch references — see §21. | Grouped GEMM persistent | `instances/grouped_gemm.py` | per-group launches | v2 persistent variant is a documented follow-up | | `streamk` | `WgradConvSpec` | `"off"` | `"dp_sk"` launches one CTA per data-parallel tile plus `streamk_ctas` stream-K CTAs; `"persistent"` launches only the pool and round-robins the data-parallel tiles. Both balance MAC iterations across CTAs instead of tiles, so a shape with few tiles and a long `wg_K` still fills the device. Mutually exclusive with `split_k` | | `streamk_reduction` | `WgradConvSpec` | `"linear"` | How CTAs sharing a tile combine partials. `"linear"`: the owner waits on each later contributor in turn and stores once. `"tree"`: pairwise fan-in, `ceil(log2)` rounds. Both are deterministic and need no dW pre-zeroing. `"atomic"`: f32 atomics into dW (fp32 only). `"workspace"`: f32 scratch atomics plus the Stage 2 reduce kernel | -| `streamk_ctas` | `WgradConvSpec` | `-1` | The CTA pool the stream-K remainder is spread over (CK `max_active_wgs`). A launch parameter, not part of the kernel: it sizes the host partition that arrives as the `sk_*` kernargs. `-1` resolves per reduction: `workspace`/`atomic` (no waiting) get four CTAs per CU (`_STREAMK_CTAS_PER_CU`, counted in CTAs, not waves); `linear`/`tree` get one CTA per CU, because a pool larger than the device can hold at once would let a waiting CTA block one it waits on | +| `streamk_ctas` | `WgradConvSpec` | `-1` | The CTA pool the stream-K remainder is spread over (CK `max_active_wgs`). A launch parameter, not part of the kernel: it sizes the host partition that arrives as the `sk_*` kernargs. `-1` resolves per reduction: `workspace`/`atomic` (no waiting) get four CTAs per CU (`_STREAMK_CTAS_PER_CU`, counted in CTAs, not waves); `linear`/`tree` get one CTA per CU, because their fixups get slower as more CTAs share a tile. The pool need not be resident: a CTA waits only on higher-index CTAs and publishes its own partial first, so ascending workgroup dispatch keeps the waits from starving their producers | | `streamk` (dispatch) | `ConvGroupedRequest` | `"auto"` | Picks stream-K + `workspace` when the candidate's tiling yields at most `_WGRAD_STREAMK_MAX_TILES` output tiles of at least `_WGRAD_STREAMK_MIN_ITERS_PER_TILE` K iterations; split-K otherwise. `"off"` disables, `"dp_sk"`/`"persistent"` force | #### 12.1.J Quantization diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/streamk.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/streamk.py index 72c70744b872..85cbb69a9ee4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/streamk.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/streamk.py @@ -267,9 +267,10 @@ class StreamKIterPartition: ``m_tiles`` / ``n_tiles`` count output tiles (a caller folding a batch dimension into M passes the folded count). ``iters_per_tile`` is the number of ``tile_k`` MAC iterations per output tile. ``max_active_wgs`` - is the CTA pool the stream-K remainder is spread over; every stream-K - CTA must be co-resident with the CTAs it waits on, so it must not exceed - the number of workgroups the device can hold at once. + is the CTA pool the stream-K remainder is spread over. It need not fit on + the device at once: a CTA only waits on higher-index CTAs and publishes + its own partial before it waits, so with workgroups dispatched in + ascending order the waiting CTAs cannot starve their producers. """ m_tiles: int diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp index f423b3ebec04..e8cd74faecd8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp @@ -658,6 +658,47 @@ void case_flag_ops_hip() EXPECT_IR(hip, "while (__hip_atomic_load("); } +/* Raw op with a bad scope/ordering, as serialized IR could carry past the + * builder's checks: the HIP lowerer must reject it, not pick a default. */ +void expect_flag_attr_rejected(rocke_opcode_t kind, const char* key, const char* value) +{ + rocke_ir_builder_t b; + rocke_ir_builder_init(&b, "flag_bad_attr_hip"); + rocke_value_t* operands[3]; + int n = 0; + if(kind != ROCKE_OP_MEMREF_FENCE) + { + operands[n++] = global_ptr_param(&b, "flags", rocke_i32()); + operands[n++] = rocke_b_const_i32(&b, 0); + operands[n++] = rocke_b_const_i32(&b, 1); + } + rocke_attr_map_t attrs; + rocke_attr_map_init(&attrs); + rocke_attr_set_str(&b, &attrs, key, value); + rocke_b_op(&b, kind, operands, n, nullptr, 0, &attrs, nullptr, 0, nullptr, nullptr); + rocke_b_ret(&b); + + rocke_strbuf_t out; + rocke_strbuf_init(&out, 256); + rocke_lower_hip_opts_t opts{}; + opts.arch = "gfx950"; + const rocke_status_t st + = rocke_lower_kernel_to_hip(&b, rocke_ir_builder_kernel(&b), &opts, &out); + if(st != ROCKE_ERR_VALUE) + fail("HIP lowering must reject an unknown flag/fence scope or ordering", __LINE__); + rocke_strbuf_free(&out); + rocke_ir_builder_free(&b); +} + +void case_flag_ops_hip_reject_bad_attrs() +{ + expect_flag_attr_rejected(ROCKE_OP_MEMREF_FENCE, "ordering", "bogus"); + expect_flag_attr_rejected(ROCKE_OP_MEMREF_FENCE, "scope", "bogus"); + expect_flag_attr_rejected(ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, "ordering", "bogus"); + expect_flag_attr_rejected(ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, "scope", "bogus"); + expect_flag_attr_rejected(ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ, "scope", "bogus"); +} + /* ---- av.load / av.store (agent-scope 128-bit vector mem) ---- */ void case_av_load_b128() { @@ -1073,6 +1114,7 @@ const TestCase k_cases[] = { {"flag_ops", case_flag_ops}, {"fence_system_scope", case_fence_system_scope}, {"flag_ops_hip", case_flag_ops_hip}, + {"flag_ops_hip_reject_bad_attrs", case_flag_ops_hip_reject_bad_attrs}, {"opcode_names_are_aligned", case_opcode_names_are_aligned}, }; From cb6a47b8e1541eb89feb03f68bd3056e91c97901 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Tue, 6 Oct 2026 13:08:10 +0000 Subject: [PATCH 13/20] fix(rocke): verify flag-op contracts and skip all-DP stream-K auto picks Address the second round of review feedback on the stream-K wgrad port: - Give memref.fence, memref.global_flag_store and memref.global_flag_wait_eq verifier contracts in both engines: operand/result counts, a global flag pointer, i32 index and value, and scope/ordering values from the builder's own sets (now shared constants). Messages match between Python and C++. - Make the C++ LLVM lowerer reject flag ops without exactly three operands, since portable-IR import can reach it without the verifier. - Have streamk="auto" build the partition the launch will use and keep split-K when it has no stream-K CTAs. On a larger default pool a single long tile would otherwise run its whole reduction on one CTA. The default pool is now one shared helper for selection and launch. ISSUE ID : AICK-1943 Co-Authored-By: Claude Opus 5.5 --- .../library/dispatch/grouped_convolution.py | 44 +++++++++--- .../common/conv_implicit_gemm_wgrad.py | 11 ++- .../test_grouped_conv_wgrad_dispatch.py | 17 +++++ .../platform/cpp/core/lower_llvm/mem.cpp | 4 ++ .../rocke/platform/cpp/core/verify.cpp | 62 ++++++++++++++++ .../rocke/platform/python/rocke/core/ir.py | 7 +- .../platform/python/rocke/core/verify.py | 36 ++++++++++ .../tests/core/future_intrinsic_lowering.cpp | 70 +++++++++++++++++++ .../rocke/platform/tests/test_rocke.py | 33 +++++++++ 9 files changed, 270 insertions(+), 14 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py index d2c375eae975..401af0882e25 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py @@ -156,6 +156,7 @@ _DEFAULT_WS_REPLICAS as _WGRAD_WS_REPLICAS, is_valid_wgrad_spec as _wgrad_is_valid_spec, wgrad_atomic_epilogue_available as _wgrad_atomic_epilogue_available, + wgrad_streamk_default_pool as _wgrad_streamk_default_pool, wgrad_streamk_grid as _wgrad_streamk_grid, wgrad_streamk_launch_values as _wgrad_streamk_launch_values, ) @@ -493,9 +494,22 @@ def _wgrad_grouped_overrides( def _wgrad_streamk_auto( - req: ConvGroupedRequest, tile_m: int, tile_n: int, tile_k: int + req: ConvGroupedRequest, + tile_m: int, + tile_n: int, + tile_k: int, + reduction: str = "workspace", ) -> bool: - """Whether ``streamk="auto"`` should pick stream-K for this request.""" + """Whether ``streamk="auto"`` should pick stream-K for this request. + + Besides the two thresholds, the partition the launch will use must give + the stream-K CTAs work: a pool larger than ``tiles * iters_per_tile`` + (one tile on a bigger arch's pool) falls back to all-DP, which runs each + tile's whole reduction on one CTA, slower than the split-K path it would + replace. + """ + from rocke.helpers.streamk import streamk_iter_partition + p = _problem(req) spatial = (p.Z if p.is_3d else 1) * p.Y * p.X wg_M = p.K // p.groups @@ -503,10 +517,18 @@ def _wgrad_streamk_auto( wg_K = p.N * p.Ho * p.Wo * (p.Do if p.is_3d else 1) tiles = -(-wg_M // tile_m) * -(-wg_N // tile_n) * p.groups iters_per_tile = -(-wg_K // tile_k) - return ( - tiles <= _WGRAD_STREAMK_MAX_TILES - and iters_per_tile >= _WGRAD_STREAMK_MIN_ITERS_PER_TILE + if ( + tiles > _WGRAD_STREAMK_MAX_TILES + or iters_per_tile < _WGRAD_STREAMK_MIN_ITERS_PER_TILE + ): + return False + part = streamk_iter_partition( + m_tiles=tiles, + n_tiles=1, + iters_per_tile=iters_per_tile, + max_active_wgs=_wgrad_streamk_default_pool(reduction, req.arch), ) + return part.sk_ctas > 0 def _wgrad_streamk_choice( @@ -517,19 +539,19 @@ def _wgrad_streamk_choice( The one resolver every wgrad candidate's instance spec, dispatch spec and grid go through, so they cannot disagree on whether the kernel is stream-K. """ + # The f32 workspace reduction never waits on another CTA, so its pool can be + # sized for occupancy; it is also the fastest of the reductions measured. + reduction = ( + "workspace" if req.streamk_reduction == "auto" else req.streamk_reduction + ) mode = req.streamk if mode == "auto": mfma = req.arch in ("gfx942", "gfx950") mode = ( "dp_sk" - if mfma and _wgrad_streamk_auto(req, tile_m, tile_n, tile_k) + if mfma and _wgrad_streamk_auto(req, tile_m, tile_n, tile_k, reduction) else "off" ) - # The f32 workspace reduction never waits on another CTA, so its pool can be - # sized for occupancy; it is also the fastest of the reductions measured. - reduction = ( - "workspace" if req.streamk_reduction == "auto" else req.streamk_reduction - ) return mode, reduction diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py index a7b4e99edb18..92f2dadda5fb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py @@ -1211,10 +1211,19 @@ def wgrad_streamk_default_ctas(spec: "WgradConvSpec", arch: str) -> int: dispatch front can wait on one that has not started, and they hold a bounded number of slots (one for linear, about one per round for tree). """ + return wgrad_streamk_default_pool(spec.streamk_reduction, arch) + + +def wgrad_streamk_default_pool(reduction: str, arch: str) -> int: + """:func:`wgrad_streamk_default_ctas` for a bare ``reduction`` name. + + Lets selection policy size the pool a launch will use before any spec + exists. + """ from rocke.helpers.split_k import _ARCH_NUM_CUS, _DEFAULT_NUM_CUS num_cus = _ARCH_NUM_CUS.get(arch, _DEFAULT_NUM_CUS) - if spec.streamk_reduction in ("linear", "tree"): + if reduction in ("linear", "tree"): return num_cus return num_cus * _STREAMK_CTAS_PER_CU diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py index 652266be4c23..b4dfa8b824cf 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py @@ -476,6 +476,23 @@ def test_auto_respects_thresholds(self): self.assertGreater(16 * 144, _WGRAD_STREAMK_MAX_TILES) self.assertFalse(_wgrad_streamk_auto(many, 64, 64, 64)) + def test_auto_skips_an_all_dp_partition(self): + # One output tile with 1024..1215 iterations clears both thresholds, + # but on gfx942 that is below the 1216-CTA workspace pool, so the + # partition would be all-DP: one CTA running the whole reduction. + # Auto keeps split-K there; gfx950's 1024-CTA pool still splits it. + for N, arch, mode in ( + (64, "gfx942", "off"), + (72, "gfx942", "off"), + (76, "gfx942", "dp_sk"), + (64, "gfx950", "dp_sk"), + ): + req = _wgrad( + arch, N=N, Hi=32, Wi=32, Y=1, X=1, pad_h=0, pad_w=0, dtype="bf16" + ) + with self.subTest(arch=arch, N=N): + self.assertEqual(dispatch_conv_grouped(req).spec.streamk, mode) + def test_off_and_gfx1250_never_pick_streamk(self): req = _wgrad("gfx950", N=128, C=24, K=96, Hi=96, Wi=96, streamk="off") self.assertEqual(dispatch_conv_grouped(req).spec.streamk, "off") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp index b9c81a457900..6c66abbebe26 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp @@ -203,6 +203,8 @@ static void op_memref_fence(rocke_lower_t* L, const rocke_op_t* op) /* Python _op_memref_global_flag_store. */ static void op_memref_global_flag_store(rocke_lower_t* L, const rocke_op_t* op) { + if(op->num_operands != 3) + rocke_ll_fail(L, ROCKE_ERR_VALUE, "memref.global_flag_store expects three operands"); const rocke_value_t* ptr = op->operands[0]; const rocke_value_t* idx = op->operands[1]; const rocke_value_t* val = op->operands[2]; @@ -224,6 +226,8 @@ static void op_memref_global_flag_store(rocke_lower_t* L, const rocke_op_t* op) * acquire atomic load; the fall-through block becomes current. */ static void op_memref_global_flag_wait_eq(rocke_lower_t* L, const rocke_op_t* op) { + if(op->num_operands != 3) + rocke_ll_fail(L, ROCKE_ERR_VALUE, "memref.global_flag_wait_eq expects three operands"); const rocke_value_t* ptr = op->operands[0]; const rocke_value_t* idx = op->operands[1]; const rocke_value_t* expect = op->operands[2]; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp index a50e15a3b514..bb3ac250de7d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp @@ -329,6 +329,31 @@ static void check_op(verifier_t* v, const rocke_op_t* op); /* ---- contracts ---- */ +/* IRBuilder._FENCE_SCOPES / _FENCE_ORDERINGS / _FLAG_STORE_ORDERINGS. */ +static const char* const k_fence_scopes[] = {"workgroup", "agent", "system", NULL}; +static const char* const k_fence_orderings[] = {"acquire", "release", "acq_rel", "seq_cst", NULL}; +static const char* const k_flag_store_orderings[] = {"monotonic", "release", "seq_cst", NULL}; + +static bool in_set(const char* s, const char* const* set) +{ + for(; *set; ++set) + if(!strcmp(s, *set)) + return true; + return false; +} + +/* Python Verifier._check_sync_attrs: scope / ordering values of a fence or + * flag op (presence is checked by the required-attrs pass). */ +static void check_sync_attrs(verifier_t* v, const rocke_op_t* op, const char* const* orderings) +{ + const char* scope = rocke_attr_get_str(&op->attrs, "scope"); + const char* ordering = rocke_attr_get_str(&op->attrs, "ordering"); + if(scope && !in_set(scope, k_fence_scopes)) + v_errf(v, op, "%s: unknown scope '%s'", op->name, scope); + if(orderings && ordering && !in_set(ordering, orderings)) + v_errf(v, op, "%s: unknown ordering '%s'", op->name, ordering); +} + static void check_contract(verifier_t* v, const rocke_op_t* op) { const char* tf32_error = rocke_tf32_op_error(op); @@ -483,6 +508,43 @@ static void check_contract(verifier_t* v, const rocke_op_t* op) v_errf(v, op, "%s: must have no operands and no results", name); } } + else if(o == ROCKE_OP_MEMREF_FENCE) + { + if(op->num_operands || op->num_results) + { + v_errf(v, op, "%s: must have no operands and no results", name); + } + check_sync_attrs(v, op, k_fence_orderings); + } + else if(o == ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE || o == ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ) + { + if(op->num_operands != 3 || op->num_results) + { + v_errf(v, + op, + "%s: expected 3 operands / 0 results, got %d / %d", + name, + op->num_operands, + op->num_results); + return; + } + const rocke_type_t* pt = op->operands[0]->type; + if(!pt || pt->kind != ROCKE_TYPE_PTR || !pt->space || strcmp(pt->space, "global")) + { + v_errf(v, op, "%s: expected a global pointer, got %s", name, tn(pt)); + } + const char* what[] = {"index", "value"}; + for(int i = 0; i < 2; ++i) + { + const rocke_type_t* t = op->operands[i + 1]->type; + if(strcmp(tn(t), "i32")) + { + v_errf(v, op, "%s: %s must be i32, got %s", name, what[i], tn(t)); + } + } + check_sync_attrs( + v, op, o == ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE ? k_flag_store_orderings : NULL); + } } /* ---- scf.for ---- */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py index 5ea993b385a7..8b079ae4401a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py @@ -103,6 +103,9 @@ def dtype_to_ir_type(dtype: str) -> Type: # LLVM ``syncscope`` values accepted by the cross-workgroup fence / flag ops. # ``system`` renders as the empty (default) scope. _FENCE_SCOPES = ("workgroup", "agent", "system") +# Orderings each op accepts; the verifier checks serialized IR against these. +_FENCE_ORDERINGS = ("acquire", "release", "acq_rel", "seq_cst") +_FLAG_STORE_ORDERINGS = ("monotonic", "release", "seq_cst") # ----- target-neutral MMA metadata --------------------------------------- @@ -1441,7 +1444,7 @@ def fence(self, *, scope: str = "agent", ordering: str = "acq_rel") -> None: """ if scope not in _FENCE_SCOPES: raise ValueError(f"unknown fence scope {scope!r}") - if ordering not in ("acquire", "release", "acq_rel", "seq_cst"): + if ordering not in _FENCE_ORDERINGS: raise ValueError(f"unknown fence ordering {ordering!r}") self._op("memref.fence", [], attrs={"scope": scope, "ordering": ordering}) @@ -1465,7 +1468,7 @@ def global_flag_store( ) if scope not in _FENCE_SCOPES: raise ValueError(f"unknown flag scope {scope!r}") - if ordering not in ("monotonic", "release", "seq_cst"): + if ordering not in _FLAG_STORE_ORDERINGS: raise ValueError(f"unknown flag store ordering {ordering!r}") self._op( "memref.global_flag_store", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/verify.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/verify.py index 2618acebc66f..df2e2091e549 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/verify.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/verify.py @@ -35,6 +35,9 @@ from typing import Dict, List, Optional, Set from .ir import ( + _FENCE_ORDERINGS, + _FENCE_SCOPES, + _FLAG_STORE_ORDERINGS, KernelDef, Op, Region, @@ -384,6 +387,39 @@ def _check_contract(self, op: Op, scope: Dict[str, Value]) -> None: elif name == "cf.return": if op.operands or op.results: self.err(f"{name}: must have no operands and no results", op) + elif name == "memref.fence": + if op.operands or op.results: + self.err(f"{name}: must have no operands and no results", op) + self._check_sync_attrs(op, _FENCE_ORDERINGS) + elif name in ("memref.global_flag_store", "memref.global_flag_wait_eq"): + if len(op.operands) != 3 or op.results: + self.err( + f"{name}: expected 3 operands / 0 results, got " + f"{len(op.operands)} / {len(op.results)}", + op, + ) + return + ptr, idx, val = op.operands + if not isinstance(ptr.type, PtrType) or ptr.type.space != "global": + self.err(f"{name}: expected a global pointer, got {ptr.type.name}", op) + for what, v in (("index", idx), ("value", val)): + if v.type.name != "i32": + self.err(f"{name}: {what} must be i32, got {v.type.name}", op) + orderings = ( + _FLAG_STORE_ORDERINGS if name == "memref.global_flag_store" else None + ) + self._check_sync_attrs(op, orderings) + + def _check_sync_attrs(self, op: Op, orderings) -> None: + """Scope / ordering values of a fence or flag op (presence is checked + by the required-attrs pass, so a missing key is not reported twice).""" + scope = op.attrs.get("scope") + if isinstance(scope, str) and scope not in _FENCE_SCOPES: + self.err(f"{op.name}: unknown scope '{scope}'", op) + ordering = op.attrs.get("ordering") + if orderings is not None and isinstance(ordering, str): + if ordering not in orderings: + self.err(f"{op.name}: unknown ordering '{ordering}'", op) # ---- scf.for ---- diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp index e8cd74faecd8..98d03e9bc0ef 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp @@ -24,6 +24,7 @@ #include "rocke/lower_hip.h" #include "rocke/lower_llvm.h" #include "rocke/strbuf.h" +#include "rocke/verify.h" namespace { @@ -690,6 +691,74 @@ void expect_flag_attr_rejected(rocke_opcode_t kind, const char* key, const char* rocke_ir_builder_free(&b); } +/* The verifier rejects what the builder would have refused (serialized IR), + * and the LLVM lowerer refuses a flag op it cannot unpack. */ +void case_flag_ops_verify_and_llvm_arity() +{ + rocke_ir_builder_t b; + rocke_ir_builder_init(&b, "flag_verify"); + rocke_value_t* f = global_ptr_param(&b, "flags", rocke_i32()); + rocke_value_t* two[] = {f, rocke_b_const_i32(&b, 0)}; + rocke_attr_map_t fence_attrs; + rocke_attr_map_init(&fence_attrs); + rocke_attr_set_str(&b, &fence_attrs, "scope", "agent"); + rocke_attr_set_str(&b, &fence_attrs, "ordering", "monotonic"); + rocke_b_op(&b, + ROCKE_OP_MEMREF_FENCE, + nullptr, + 0, + nullptr, + 0, + &fence_attrs, + nullptr, + 0, + nullptr, + nullptr); + rocke_attr_map_t store_attrs; + rocke_attr_map_init(&store_attrs); + rocke_attr_set_str(&b, &store_attrs, "scope", "agent"); + rocke_attr_set_str(&b, &store_attrs, "ordering", "release"); + rocke_b_op(&b, + ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, + two, + 2, + nullptr, + 0, + &store_attrs, + nullptr, + 0, + nullptr, + nullptr); + rocke_b_ret(&b); + + rocke_diag_t* diags = nullptr; + size_t n = 0; + rocke_verify(rocke_ir_builder_kernel(&b), &diags, &n); + bool saw_ordering = false, saw_arity = false; + for(size_t i = 0; i < n; ++i) + { + const std::string m = diags[i].message ? diags[i].message : ""; + saw_ordering |= m.find("memref.fence: unknown ordering 'monotonic'") != std::string::npos; + saw_arity |= m.find("memref.global_flag_store: expected 3 operands / 0 results, got 2 / 0") + != std::string::npos; + } + rocke_diags_free(diags, n); + if(!saw_ordering) + fail("verifier must reject a fence ordering outside the builder's set", __LINE__); + if(!saw_arity) + fail("verifier must reject a two-operand flag store", __LINE__); + + char* ll = nullptr; + char err[ROCKE_ERR_MSG_CAP]; + err[0] = '\0'; + const rocke_status_t st = rocke_lower_kernel_to_llvm_ex( + rocke_ir_builder_kernel(&b), ROCKE_LLVM_FLAVOR_AUTO, "gfx950", &ll, err, sizeof(err)); + if(st != ROCKE_ERR_VALUE) + fail("LLVM lowering must reject a two-operand flag store", __LINE__); + std::free(ll); + rocke_ir_builder_free(&b); +} + void case_flag_ops_hip_reject_bad_attrs() { expect_flag_attr_rejected(ROCKE_OP_MEMREF_FENCE, "ordering", "bogus"); @@ -1115,6 +1184,7 @@ const TestCase k_cases[] = { {"fence_system_scope", case_fence_system_scope}, {"flag_ops_hip", case_flag_ops_hip}, {"flag_ops_hip_reject_bad_attrs", case_flag_ops_hip_reject_bad_attrs}, + {"flag_ops_verify_and_llvm_arity", case_flag_ops_verify_and_llvm_arity}, {"opcode_names_are_aligned", case_opcode_names_are_aligned}, }; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py index d2270eafb9e8..ded4a46d54c1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py @@ -3372,6 +3372,39 @@ def test_flag_ops_reject_bad_scope_ordering_and_type(self): with self.assertRaisesRegex(ValueError, "unknown flag scope"): b.global_flag_wait_eq(flags, b.const_i32(0), b.const_i32(1), scope="gpu") + def test_verifier_rejects_malformed_flag_ops(self): + # Serialized IR can carry flag ops the builder would have refused; the + # verifier has to catch them before a lowerer unpacks the operands. + from rocke.core.ir import F32 + from rocke.core.verify import ERROR, verify + + b = self._builder("flagverify") + flags = self._flag_kernel(b) + b.fence(scope="agent", ordering="release") + b.global_flag_store(flags, b.const_i32(0), b.const_i32(1)) + b.global_flag_wait_eq(flags, b.const_i32(0), b.const_i32(1)) + b.ret() + ops = {op.name: op for op in b.kernel.body.ops} + self.assertEqual( + [d for d in verify(b.kernel) if d.severity == ERROR], [], "well-formed" + ) + ops["memref.fence"].attrs["ordering"] = "monotonic" + ops["memref.global_flag_store"].attrs["scope"] = "gpu" + wait = ops["memref.global_flag_wait_eq"] + wait.operands[2] = b.const_f32(1.0) + msgs = [d.message for d in verify(b.kernel) if d.severity == ERROR] + for needle in ( + "memref.fence: unknown ordering 'monotonic'", + "memref.global_flag_store: unknown scope 'gpu'", + f"memref.global_flag_wait_eq: value must be i32, got {F32.name}", + ): + self.assertTrue(any(needle in m for m in msgs), (needle, msgs)) + del wait.operands[2] + msgs = [d.message for d in verify(b.kernel) if d.severity == ERROR] + self.assertTrue( + any("expected 3 operands / 0 results, got 2 / 0" in m for m in msgs), msgs + ) + # ---- ds_swizzle (raw offset + XOR-butterfly encoding) ---- def test_ds_swizzle_passes_raw_offset_immediate(self): ll = self._lower("dssw", lambda b: b.ds_swizzle(b.const_i32(1), 0x041F)) From 6c6bf35004f0781bc0da151848eb6755d74eecae Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Thu, 8 Oct 2026 16:24:02 +0000 Subject: [PATCH 14/20] fix(rocke): close stream-K wgrad review gaps in both engines - Clamp StreamKIterPartition.iters_per_sk_cta to >= 1 at the source (Python and C++) so an all-data-parallel partition never feeds a zero divisor to the kernel's owner lookup. - Make flags_bytes take the wave count (one flag per CTA and wave) and use it from both wgrad workspace layouts instead of a duplicated formula. - Remove the unused emit_streamk_partial_store/_load_accumulate helpers; the wgrad fixup is the reference flag/partial protocol. - Port the streamk_ctas == -1 CTA-pool auto-sizing to C++ and bind the C++ validity gate; test both against Python. - Lock all eight stream-K mode x reduction combos in the golden IR. - Poison partials and dW before each launch in the GPU repeat test so a missing release/acquire can actually fail it. - Build the linear/tree started/ended predicates once and name the operand in the vector-width launch errors. ISSUE ID : AICK-1943 Co-Authored-By: Claude Opus 5.5 --- .../rocke/library/kernels/common/conv_args.py | 4 +- .../common/conv_implicit_gemm_wgrad.py | 33 ++-- .../library/tests/test_conv_wgrad_streamk.py | 137 ++++++++++++++++- .../tests/test_conv_wgrad_streamk_gpu.py | 71 +++------ .../platform/cpp/bindings/rocke_engine.cpp | 31 +++- .../rocke/platform/cpp/helpers/streamk.cpp | 22 ++- .../rocke/helper_rocke.helpers.streamk.h | 15 +- .../rocke/instance_conv_implicit_gemm_wgrad.h | 19 ++- .../common/conv_implicit_gemm_wgrad.cpp | 79 ++++++++-- .../platform/python/rocke/helpers/__init__.py | 4 - .../platform/python/rocke/helpers/streamk.py | 115 +++++--------- .../rocke_representative_ir_sha256.json | 144 ++++++++++++++++-- .../instances/rocke_ir_parity_harness.py | 97 ++++++++++++ .../instances/streamk_iter_partition.cpp | 38 ++++- .../instances/test_streamk_iter_partition.py | 25 ++- 15 files changed, 641 insertions(+), 193 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_args.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_args.py index 3ff0488b7332..3b89b31f031a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_args.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_args.py @@ -583,8 +583,8 @@ def _streamk_values(self, sk: StreamKLaunch) -> Dict[str, int]: "sk_dp_tiles": part.dp_tiles, "sk_ctas": part.sk_ctas, # A divisor in the kernel's owner lookup, which an all-DP launch - # still evaluates; nothing reads it as a length when sk_ctas == 0. - "sk_iters_per_cta": max(part.iters_per_sk_cta, 1), + # still evaluates; the partition clamps it to >= 1 for that case. + "sk_iters_per_cta": part.iters_per_sk_cta, "sk_extra_iters": part.extra_iters, "sk_pool": part.max_active_wgs, "sk_tree_rounds": part.tree_rounds, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py index 92f2dadda5fb..dc0397223f6f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py @@ -1279,18 +1279,20 @@ def _vec(chan: int, dtype: str, cap: int | None) -> int: f"launch problem is incompatible with the built kernel: {what} " f"differs (built {want}, launch {got}); build a kernel for this shape" ) + # Both X (operand B) and dW (operand D) run along cpg, with their own + # vector widths, so each check names the operand it guards. chans = ( - ("kpg", built.kpg, problem.kpg, spec.data.dtype_a, spec.vector_size_a), - ("cpg", built.cpg, problem.cpg, spec.data.dtype_b, spec.vector_size_b), - ("cpg", built.cpg, problem.cpg, spec.data.dtype_d, spec.vector_size_c), + ("dY", "kpg", built.kpg, problem.kpg, spec.data.dtype_a, spec.vector_size_a), + ("X", "cpg", built.cpg, problem.cpg, spec.data.dtype_b, spec.vector_size_b), + ("dW", "cpg", built.cpg, problem.cpg, spec.data.dtype_d, spec.vector_size_c), ) - for what, want, got, dtype, cap in chans: + for operand, what, want, got, dtype, cap in chans: width = _vec(want, dtype, cap) if got % width: raise ValueError( f"launch problem is incompatible with the built kernel: {what}=" - f"{got} is not a multiple of the built {width}-wide vector " - f"({what}={want}); build a kernel for this shape" + f"{got} is not a multiple of the built {width}-wide {operand} " + f"vector ({what}={want}); build a kernel for this shape" ) return problem @@ -1343,8 +1345,7 @@ def wgrad_streamk_workspace_layout( if spec.streamk_reduction not in ("linear", "tree"): return 0, 0 part = wgrad_streamk_partition(spec, arch=arch, problem=problem) - flags = 4 * part.sk_ctas * spec.warp_m * spec.warp_n - flags_bytes = -(-flags // 256) * 256 + flags_bytes = part.flags_bytes(spec.warp_m * spec.warp_n) partials_bytes = part.sk_ctas * spec.tile_m * spec.tile_n * 4 return flags_bytes, partials_bytes @@ -2804,6 +2805,14 @@ def _store_tile(vals: Sequence[Value]) -> None: b.scf_yield(*new_accs) tile_accs = for_op.results + if reduction in ("linear", "tree"): + # Whether this item covers the start / the end of its tile. A + # data-parallel tile (and every tile of an all-DP partition) + # has both. Built only for the flag-handoff fixups, so the + # atomic and workspace bodies carry no dead compares. + started = b.cmp_eq(lo, c0) + ended = b.cmp_eq(hi, c_ipt) + if reduction == "atomic": _emit_wgrad_split_k_epilogue( b, @@ -2838,10 +2847,8 @@ def _store_tile(vals: Sequence[Value]) -> None: replica_sel=cta, ) elif reduction == "linear": - # A data-parallel tile (and every tile of an all-DP partition) - # has started and ended, so it folds no partners and stores. - started = b.cmp_eq(lo, c0) - ended = b.cmp_eq(hi, c_ipt) + # A data-parallel tile has started and ended, so it folds no + # partners and stores. with b.scf_if_else(started) as (owner, contributor): with owner: # Fold every later contributor in order, then store once. @@ -2871,8 +2878,6 @@ def _store_tile(vals: Sequence[Value]) -> None: # round count is a kernarg (0 for an all-DP partition), so the # rounds are a runtime loop carrying the partial sums, the pair # stride and a dropped-out flag; every predicate is uniform. - started = b.cmp_eq(lo, c0) - ended = b.cmp_eq(hi, c_ipt) not_started = b.lnot(started) local = b.sub(sk_cta, _sk_cta_of(tile_start)) zero_acc = b.zero_vec_f32(c_per_lane) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py index c39116f1a240..970694d35e15 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py @@ -217,6 +217,66 @@ def test_cpp_rejects_with_the_same_reason(self): eng.conv_wgrad_lower_llvm(_spec_dict(spec), "gfx950") self.assertIn(why, str(cm.exception)) + # (label, problem, arch, spec overrides). Each rejected config trips one + # gate only, so the engines' differing check order cannot pick a different + # reason. + _AGREEMENT = [ + ("ok_dp_sk_linear", _DENSE, "gfx950", {}), + ( + "ok_pers_tree", + _GROUPED, + "gfx950", + dict(streamk="persistent", streamk_reduction="tree"), + ), + ( + "ok_atomic_fp32", + _DENSE, + "gfx950", + dict(dtype_d="fp32", streamk_reduction="atomic"), + ), + ("ok_workspace", _GROUPED, "gfx950", dict(streamk_reduction="workspace")), + ("ok_gfx942", _DENSE, "gfx942", dict(edge=16, streamk="persistent")), + ("ok_auto_ctas", _DENSE, "gfx950", dict(streamk_ctas=-1)), + ("ok_off", _DENSE, "gfx950", dict(streamk="off")), + ("split_k", _DENSE, "gfx950", dict(split_k=4)), + ("split_k_auto", _DENSE, "gfx950", dict(split_k=-1)), + ("two_stage", _DENSE, "gfx950", dict(two_stage=True, split_k=2)), + ("cshuffle", _DENSE, "gfx950", dict(epilogue="cshuffle")), + ("unroll_k", _DENSE, "gfx950", dict(unroll_k=True)), + ("chiplet", _DENSE, "gfx950", dict(chiplet_swizzle=True)), + ("atomic_fp16", _DENSE, "gfx950", dict(streamk_reduction="atomic")), + ( + "atomic_bf16", + _DENSE, + "gfx950", + dict(dtype_d="bf16", streamk_reduction="atomic"), + ), + ("ctas_zero", _DENSE, "gfx950", dict(streamk_ctas=0)), + ("ctas_minus_two", _DENSE, "gfx950", dict(streamk_ctas=-2)), + ("bad_mode", _DENSE, "gfx950", dict(streamk="sideways")), + ("bad_reduction", _DENSE, "gfx950", dict(streamk_reduction="sideways")), + ("gfx90a", _DENSE, "gfx90a", {}), + # wave32 so Python's wave-size check passes and the arch gate decides. + ("gfx1250", _DENSE, "gfx1250", dict(wave_size=32, warp_tile_k=32)), + ("group_merge", _DW, "gfx950", dict(group_merge=2)), + ] + + def test_cpp_validity_agrees(self): + eng = _engine() + for label, problem, arch, kw in self._AGREEMENT: + spec = _spec(problem, **kw) + want = is_valid_wgrad_spec(spec, arch) + with self.subTest(label): + self.assertEqual(label.startswith("ok_"), want[0], want[1]) + got = eng.conv_wgrad_is_valid_spec(_spec_dict(spec), arch) + self.assertEqual(tuple(got), want) + # A wave64 spec on gfx1250 trips Python's wave-size check before the + # stream-K arch gate, which the C port has no counterpart of; the + # reasons differ but both engines reject it. + spec = _spec() + self.assertFalse(is_valid_wgrad_spec(spec, "gfx1250")[0]) + self.assertFalse(eng.conv_wgrad_is_valid_spec(_spec_dict(spec), "gfx1250")[0]) + def test_cpp_rejects_group_merge_with_the_same_reason(self): eng = _engine() spec = _spec(_DW, group_merge=2) @@ -332,6 +392,62 @@ def test_cpp_host_plan_agrees(self): got["is_deterministic"], reduction in ("linear", "tree") ) + def test_cpp_default_pool_agrees(self): + from kernels.common.conv_implicit_gemm_wgrad import wgrad_streamk_default_pool + from rocke.helpers.split_k import _ARCH_NUM_CUS + + eng = _engine() + # Every arch of the Python CU table, plus one it does not list. + for arch in [*_ARCH_NUM_CUS, "gfx90a"]: + for reduction in ("linear", "tree", "atomic", "workspace"): + with self.subTest(arch=arch, r=reduction): + self.assertEqual( + eng.conv_wgrad_streamk_default_pool(reduction, arch), + wgrad_streamk_default_pool(reduction, arch), + ) + + def test_cpp_host_plan_agrees_auto_pool(self): + # streamk_ctas=-1 (the spec default) resolves to the arch's CU-based + # pool in both engines. + eng = _engine() + for arch, edge in (("gfx950", 32), ("gfx942", 16)): + for problem in (_DENSE, _GROUPED): + for mode in ("dp_sk", "persistent"): + for reduction, dtype_d in ( + ("linear", "fp16"), + ("tree", "fp16"), + ("atomic", "fp32"), + ("workspace", "fp16"), + ): + spec = _spec( + problem, + dtype_d=dtype_d, + edge=edge, + streamk=mode, + streamk_reduction=reduction, + streamk_ctas=-1, + ) + got = eng.conv_wgrad_streamk_host(_spec_dict(spec), arch) + flags, partials = wgrad_streamk_workspace_layout( + spec, arch=arch + ) + with self.subTest( + arch=arch, g=problem.groups, mode=mode, r=reduction + ): + self.assertEqual( + tuple(got["grid"]), wgrad_streamk_grid(spec, arch=arch) + ) + self.assertEqual(got["flags_bytes"], flags) + self.assertEqual(got["partials_bytes"], partials) + self.assertEqual( + got["workspace_bytes"], + wgrad_streamk_workspace_nbytes(spec, arch=arch), + ) + # Without an arch both engines default to gfx950. + spec = _spec(streamk="persistent", streamk_ctas=-1) + got = eng.conv_wgrad_streamk_host(_spec_dict(spec)) + self.assertEqual(tuple(got["grid"]), wgrad_streamk_grid(spec)) + class TestStreamKLaunchAbi(unittest.TestCase): """The stream-K kernarg ABI and the launch values that fill it.""" @@ -376,9 +492,7 @@ def test_values_are_the_partition(self): self.assertEqual(v["sk_iters_per_tile"], part.iters_per_tile) self.assertEqual(v["sk_dp_tiles"], part.dp_tiles) self.assertEqual(v["sk_ctas"], part.sk_ctas) - self.assertEqual( - v["sk_iters_per_cta"], max(part.iters_per_sk_cta, 1) - ) + self.assertEqual(v["sk_iters_per_cta"], part.iters_per_sk_cta) self.assertEqual(v["sk_extra_iters"], part.extra_iters) self.assertEqual(v["sk_pool"], pool) self.assertEqual(v["sk_tree_rounds"], part.tree_rounds) @@ -439,14 +553,27 @@ def test_launch_shape_must_keep_the_built_modes(self): ("grouping", ConvProblem(C=32, K=96, groups=2, **base)), ("is_3d", ConvProblem(C=32, K=96, Di=4, Z=1, sD=1, pD=0, dD=1, **base)), ("is_pointwise", ConvProblem(N=2, Hi=12, Wi=12, C=32, K=96, Y=1, X=1)), - ("kpg=36", ConvProblem(C=32, K=36, **base)), - ("cpg=20", ConvProblem(C=20, K=96, **base)), + ( + "kpg=36 is not a multiple of the built 8-wide dY", + ConvProblem(C=32, K=36, **base), + ), + ( + "cpg=20 is not a multiple of the built 8-wide X", + ConvProblem(C=20, K=96, **base), + ), ): with self.subTest(needle=needle): for helper in (wgrad_streamk_partition, wgrad_streamk_workspace_nbytes): with self.assertRaises(ValueError) as cm: helper(spec, problem=other) self.assertIn(needle, str(cm.exception)) + # X and dW both run along cpg; the error names the operand that failed. + narrow_b = _spec(vector_size_b=4) + with self.assertRaises(ValueError) as cm: + wgrad_streamk_partition(narrow_b, problem=ConvProblem(C=20, K=96, **base)) + self.assertIn( + "cpg=20 is not a multiple of the built 8-wide dW", str(cm.exception) + ) # Wider channel runs than the build keep the same code path. wgrad_streamk_partition(spec, problem=ConvProblem(C=64, K=192, **base)) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk_gpu.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk_gpu.py index 3035ed27d018..bed52a41e7f2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk_gpu.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk_gpu.py @@ -13,7 +13,11 @@ stream-K mix, a grouped problem whose last M tile is partial, and a small CTA pool so the stream-K path is taken at these sizes. The linear and tree fixups are also run repeatedly and must be bit-identical across launches -- the only -check that catches a missing release/acquire on the partial hand-off. +check that catches a missing release/acquire on the partial hand-off. Before +every launch the partial slots and dW are poisoned with NaN bytes (only the +flags are zeroed, as the contract requires), so a consumer that reads a slot +before its producer's release lands sees NaN rather than the previous +launch's bit-identical partial. Requires a gfx942/gfx950 GPU and torch; skips cleanly otherwise. Host-side coverage lives in ``test_conv_wgrad_streamk.py``. @@ -73,7 +77,7 @@ def _reference(X, dY, p): return w.permute(0, 2, 3, 1).contiguous() -def _run(spec, *, launches: int = 1, grid=None, problem=None, artifact=None): +def _run(spec, *, launches: int = 1, problem=None, artifact=None): """Build, launch ``launches`` times, return (dW outputs, fp32 reference). ``problem`` is the launch shape (default ``spec.problem``): the kernel is @@ -125,7 +129,9 @@ def _run(spec, *, launches: int = 1, grid=None, problem=None, artifact=None): rt.memcpy_h2d(bufs[0], _u8(dY), dY.nbytes) rt.memcpy_h2d(bufs[1], _u8(X), X.nbytes) reduction = spec.streamk_reduction - flags_bytes, _ = wgrad_streamk_workspace_layout(spec, arch=arch, problem=p) + flags_bytes, partials_bytes = wgrad_streamk_workspace_layout( + spec, arch=arch, problem=p + ) extra = {} if reduction == "workspace": extra = dict(ws_ptr=ws, ws_bytes=ws_bytes) @@ -173,16 +179,25 @@ def _run(spec, *, launches: int = 1, grid=None, problem=None, artifact=None): for _ in range(launches): # The documented zeroing contract, per reduction: atomic accumulates # into dW, workspace into the scratch, linear/tree only reset flags. + # Everything the contract leaves alone is poisoned (0xFF bytes are + # NaN in every float dtype): a stale value from the previous + # launch would otherwise be bit-identical to the right one and + # hide an early read of a partial or an unwritten dW element. if reduction == "atomic": rt.memset(bufs[2], 0, dW.nbytes) - elif reduction == "workspace": + else: + rt.memset(bufs[2], 0xFF, dW.nbytes) + if reduction == "workspace": rt.memset(ws, 0, ws_bytes) - elif flags_bytes: - rt.memset(ws, 0, flags_bytes) + elif reduction in ("linear", "tree"): + if partials_bytes: + rt.memset(ws + flags_bytes, 0xFF, partials_bytes) + if flags_bytes: + rt.memset(ws, 0, flags_bytes) launcher( values, config=LaunchConfig( - grid=grid or wgrad_streamk_grid(spec, arch=arch, problem=p), + grid=wgrad_streamk_grid(spec, arch=arch, problem=p), block=(spec.block_size, 1, 1), fence=True, ), @@ -205,10 +220,8 @@ def _run(spec, *, launches: int = 1, grid=None, problem=None, artifact=None): return outs, ref -def _check(test, spec, *, launches: int = 1, grid=None, problem=None, artifact=None): - outs, ref = _run( - spec, launches=launches, grid=grid, problem=problem, artifact=artifact - ) +def _check(test, spec, *, launches: int = 1, problem=None, artifact=None): + outs, ref = _run(spec, launches=launches, problem=problem, artifact=artifact) p = problem if problem is not None else spec.problem out = outs[0].float() written = out.reshape(p.K, -1).abs().sum(dim=1) > 0 @@ -386,42 +399,6 @@ def test_one_binary_many_shapes(self): _check(self, spec, problem=problem, artifact=art) -@unittest.skipUnless(not _SKIP_REASON, _SKIP_REASON or "no GPU") -class TestConvWgradStreamKViaDispatch(unittest.TestCase): - """The shipped selector reaches the stream-K kernel and its launch grid.""" - - def test_dispatch_spec_and_grid(self): - from dispatch.grouped_convolution import ( - ConvGroupedRequest, - _problem, - dispatch_conv_grouped, - ) - - for mode in ("dp_sk", "persistent"): - for reduction in ("linear", "tree", "workspace"): - req = ConvGroupedRequest( - N=4, - C=64, - K=96, - Hi=12, - Wi=12, - Y=3, - X=3, - pad_h=1, - pad_w=1, - G=2, - arch=GPU_ARCH, - direction="wgrad", - streamk=mode, - streamk_reduction=reduction, - ) - r = dispatch_conv_grouped(req) - spec = r.spec.to_wgrad_spec(_problem(req)) - with self.subTest(mode=mode, reduction=reduction): - self.assertEqual(spec.streamk, mode) - _check(self, spec, grid=r.grid) - - @unittest.skipUnless(not _SKIP_REASON, _SKIP_REASON or "no GPU") class TestConvWgradStreamKDeterminism(unittest.TestCase): """Linear and tree fold partials in a fixed order: dW is bit-exact.""" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp index 51fedd4faefa..4d91ddb8198b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp @@ -3876,16 +3876,42 @@ PYBIND11_MODULE(rocke_engine, m) "Return workspace bytes for the two-stage wgrad path.\n" "Formula: groups * ws_replicas * wg_M * wg_N * 4 (always f32).\n" "Returns 0 when two_stage=false, or split_k <= 1."); + m.def( + "conv_wgrad_is_valid_spec", + [](const py::dict& d, const std::string& arch) -> py::tuple { + std::deque store; + rocke_implicit_gemm_conv_wgrad_spec_t s = conv_wgrad_build_spec(d, store); + char reason[1024] = {0}; + const bool ok = rocke_implicit_gemm_conv_wgrad_is_valid_spec( + &s, arch_or_default(arch), reason, sizeof(reason)); + return py::make_tuple(ok, ok ? std::string("ok") : std::string(reason)); + }, + py::arg("spec"), + py::arg("arch") = "gfx950", + "C port of is_valid_wgrad_spec: (ok, reason), reason 'ok' when accepted."); + m.def( + "conv_wgrad_streamk_default_pool", + [](const std::string& reduction, const std::string& arch) -> int { + return rocke_wgrad_conv_streamk_default_pool(reduction.c_str(), arch_or_default(arch)); + }, + py::arg("reduction"), + py::arg("arch") = "gfx950", + "C port of wgrad_streamk_default_pool: the CTA pool of a streamk_ctas=-1 launch."); m.def( "conv_wgrad_streamk_host", - [](const py::dict& d) -> py::dict { + [](const py::dict& d, const std::string& arch) -> py::dict { std::deque store; rocke_implicit_gemm_conv_wgrad_spec_t s = conv_wgrad_build_spec(d, store); + /* Resolve the auto pool (-1) for `arch` up front, as Python + * wgrad_streamk_launch does, so every helper below sees it. */ + const int pool = rocke_wgrad_conv_streamk_ctas(&s, arch_or_default(arch)); + if(pool > 0) + s.streamk_ctas = pool; int grid[3] = {0, 0, 0}; if(rocke_wgrad_conv_streamk_grid(&s, grid) != ROCKE_OK) throw std::runtime_error( "rocke_engine.conv_wgrad_streamk_host: needs streamk != 'off' and " - "streamk_ctas > 0"); + "streamk_ctas == -1 or > 0"); size_t flags_bytes = 0; size_t partials_bytes = 0; rocke_wgrad_conv_streamk_workspace_layout(&s, &flags_bytes, &partials_bytes); @@ -3898,6 +3924,7 @@ PYBIND11_MODULE(rocke_engine, m) return out; }, py::arg("spec"), + py::arg("arch") = "gfx950", "Stream-K wgrad host plan of the C port: launch grid, workspace layout\n" "and determinism, for cross-checking the Python wgrad_streamk_* helpers."); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/streamk.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/streamk.cpp index 2ad371193725..3f5011d01cf6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/streamk.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/streamk.cpp @@ -196,7 +196,8 @@ rocke_streamk_iter_plan_t rocke_streamk_iter_plan(const rocke_streamk_iter_parti } r.sk_ctas = r.sk_tiles ? p->max_active_wgs : 0; r.total_sk_iters = r.sk_tiles * p->iters_per_tile; - r.iters_per_sk_cta = r.sk_ctas ? r.total_sk_iters / r.sk_ctas : 0; + /* Clamped to 1 for an all-DP plan: kernels divide by it unconditionally. */ + r.iters_per_sk_cta = r.sk_ctas ? r.total_sk_iters / r.sk_ctas : 1; r.extra_iters = r.sk_ctas ? r.total_sk_iters % r.sk_ctas : 0; r.dp_tiles = r.num_tiles - r.sk_tiles; r.total_dp_iters = r.dp_tiles * p->iters_per_tile; @@ -210,10 +211,27 @@ rocke_streamk_iter_plan_t rocke_streamk_iter_plan(const rocke_streamk_iter_parti r.tree_rounds = 0; for(rest = contributors - 1; rest > 0; rest >>= 1) ++r.tree_rounds; - r.flags_bytes = sk_ceil_div(4 * r.sk_ctas, 256) * 256; return r; } +/* Python StreamKIterPartition.flags_bytes(n_waves): one i32 flag per + * (stream-K CTA, wave), padded to 256 bytes. */ +size_t rocke_streamk_flags_bytes(const rocke_streamk_iter_plan_t* r, + int n_waves, + rocke_status_t* out_status) +{ + if(r == NULL || n_waves <= 0) + { + if(out_status) + *out_status = ROCKE_ERR_VALUE; + return 0; + } + if(out_status) + *out_status = ROCKE_OK; + const size_t flags = (size_t)4 * (size_t)r->sk_ctas * (size_t)n_waves; + return (flags + 255) / 256 * 256; +} + int rocke_streamk_start_iter(const rocke_streamk_iter_partition_t* p, int sk_cta) { const rocke_streamk_iter_plan_t r = rocke_streamk_iter_plan(p); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.streamk.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.streamk.h index 9e20e2883ee3..71d29644048f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.streamk.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.streamk.h @@ -150,6 +150,7 @@ rocke_streamk_decoded_tile_t rocke_emit_streamk_decode(rocke_ir_builder_t* b, * (its @property values) + rocke_streamk_iter_plan_t, resolved * once by rocke_streamk_iter_plan() * streamk_iter_partition(...) rocke_streamk_iter_partition_make(...) + * StreamKIterPartition.flags_bytes() rocke_streamk_flags_bytes(...) * streamk_start_iter / streamk_end_iter rocke_streamk_start_iter / _end_iter * StreamKIterArgs rocke_streamk_iter_args_t * emit_streamk_sk_start_iter(...) rocke_emit_streamk_sk_start_iter(...) @@ -174,14 +175,13 @@ typedef struct rocke_streamk_iter_plan int sk_tiles; int sk_ctas; int total_sk_iters; - int iters_per_sk_cta; + int iters_per_sk_cta; /* >= 1; clamped to 1 for an all-DP plan (a divisor) */ int extra_iters; int dp_tiles; int total_dp_iters; int grid_size; int max_linear_partners; int tree_rounds; - int flags_bytes; } rocke_streamk_iter_plan_t; /* streamk_iter_partition: every count must be > 0. On violation *out_status @@ -196,6 +196,15 @@ rocke_streamk_iter_partition_t rocke_streamk_iter_partition_make(int m_tiles, /* Every StreamKIterPartition @property, resolved once. */ rocke_streamk_iter_plan_t rocke_streamk_iter_plan(const rocke_streamk_iter_partition_t* p); +/* StreamKIterPartition.flags_bytes(n_waves): bytes of the stream-K fixup flag + * table, one i32 flag per (stream-K CTA, wave) padded to 256 bytes; 0 for an + * all-DP plan. n_waves is the waves per CTA that signal independently + * (warp_m * warp_n for wgrad). n_waves <= 0 (Python ValueError) sets + * *out_status (if non-NULL) to ROCKE_ERR_VALUE and returns 0. */ +size_t rocke_streamk_flags_bytes(const rocke_streamk_iter_plan_t* r, + int n_waves, + rocke_status_t* out_status); + int rocke_streamk_start_iter(const rocke_streamk_iter_partition_t* p, int sk_cta); int rocke_streamk_end_iter(const rocke_streamk_iter_partition_t* p, int sk_cta); @@ -207,7 +216,7 @@ typedef struct rocke_streamk_iter_args rocke_value_t* iters_per_tile; /* i32 MAC iterations per output tile */ rocke_value_t* dp_tiles; /* i32 tiles owned whole by data-parallel CTAs */ rocke_value_t* total_dp_iters; /* i32 dp_tiles * iters_per_tile */ - rocke_value_t* iters_per_sk_cta; /* i32 base iterations per stream-K CTA (> 0) */ + rocke_value_t* iters_per_sk_cta; /* i32 base iterations per stream-K CTA (>= 1) */ rocke_value_t* extra_iters; /* i32 stream-K CTAs that own one more */ } rocke_streamk_iter_args_t; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h index f79474ca8797..cfabdc546498 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h @@ -250,10 +250,23 @@ size_t rocke_wgrad_conv_workspace_bytes(const rocke_implicit_gemm_conv_wgrad_spe /* ---- stream-K host-side helpers (Python wgrad_streamk_*) ---- */ +/* Python wgrad_streamk_default_pool(reduction, arch): the CTA pool an auto + * (streamk_ctas == -1) launch uses -- one CTA per CU for linear/tree, four per + * CU for atomic/workspace. CU counts come from a static per-arch table that + * mirrors rocke.helpers.split_k._ARCH_NUM_CUS (256 for an unknown arch). + * A NULL arch is gfx950, Python's default. */ +int rocke_wgrad_conv_streamk_default_pool(const char* reduction, const char* arch); + +/* Python wgrad_streamk_launch(spec, arch=arch).pool: streamk_ctas when > 0, + * the arch default pool when -1, and 0 for streamk == "off" or an invalid + * count. A NULL arch is gfx950. */ +int rocke_wgrad_conv_streamk_ctas(const rocke_implicit_gemm_conv_wgrad_spec_t* s, const char* arch); + /* Python wgrad_streamk_partition(spec): the iteration partition the kernel and - * its launch grid share. Requires streamk != "off" and streamk_ctas > 0; - * otherwise *out_status (if non-NULL) is ROCKE_ERR_VALUE and the result is - * zeroed. */ + * its launch grid share. Requires streamk != "off" and streamk_ctas == -1 + * (auto, resolved for gfx950 like Python's default arch; resolve it with + * rocke_wgrad_conv_streamk_ctas for another arch) or > 0; otherwise + * *out_status (if non-NULL) is ROCKE_ERR_VALUE and the result is zeroed. */ rocke_streamk_iter_partition_t rocke_wgrad_conv_streamk_partition(const rocke_implicit_gemm_conv_wgrad_spec_t* s, rocke_status_t* out_status); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp index 1e4da42a7667..2cdba4ee88a7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp @@ -162,14 +162,62 @@ bool rocke_wgrad_conv_spec_is_deterministic(const rocke_implicit_gemm_conv_wgrad return s->split_k <= 1; } -/* Python wgrad_streamk_partition. */ +/* Per-arch CU counts. Mirrors the static table rocke.helpers.split_k + * _ARCH_NUM_CUS / _DEFAULT_NUM_CUS that Python wgrad_streamk_default_pool + * reads; test_conv_wgrad_streamk.py checks the two agree key by key. */ +static int wgrad_arch_num_cus(const char* arch) +{ + static const struct + { + const char* gfx; + int num_cus; + } table[] = { + {"gfx942", 304}, /* MI300X: 8 XCDs * 38 CUs */ + {"gfx950", 256}, /* MI350X: 8 XCDs * 32 CUs */ + {"gfx1151", 64}, /* RDNA gfx1151 (WMMA) */ + }; + if(arch != NULL) + for(const auto& e : table) + if(strcmp(arch, e.gfx) == 0) + return e.num_cus; + return 256; /* _DEFAULT_NUM_CUS */ +} + +/* Python _STREAMK_CTAS_PER_CU: CTAs per CU in the default pool of the + * reductions that never wait (atomic, workspace). */ +#define ROCKE_WGRAD_STREAMK_CTAS_PER_CU 4 + +/* Python wgrad_streamk_default_pool. */ +int rocke_wgrad_conv_streamk_default_pool(const char* reduction, const char* arch) +{ + const int num_cus = wgrad_arch_num_cus(arch ? arch : "gfx950"); + const char* red = reduction ? reduction : "linear"; + if(strcmp(red, "linear") == 0 || strcmp(red, "tree") == 0) + return num_cus; + return num_cus * ROCKE_WGRAD_STREAMK_CTAS_PER_CU; +} + +/* Python wgrad_streamk_launch(spec, arch=arch).pool. */ +int rocke_wgrad_conv_streamk_ctas(const rocke_implicit_gemm_conv_wgrad_spec_t* s, const char* arch) +{ + if(s == NULL || !wgrad_streamk_on(s)) + return 0; + if(s->streamk_ctas > 0) + return s->streamk_ctas; + if(s->streamk_ctas == -1) + return rocke_wgrad_conv_streamk_default_pool(wgrad_streamk_reduction(s), arch); + return 0; +} + +/* Python wgrad_streamk_partition (arch defaults to gfx950, as in Python). */ rocke_streamk_iter_partition_t rocke_wgrad_conv_streamk_partition(const rocke_implicit_gemm_conv_wgrad_spec_t* s, rocke_status_t* out_status) { rocke_streamk_iter_partition_t zero; memset(&zero, 0, sizeof(zero)); - if(s == NULL || !wgrad_streamk_on(s) || s->streamk_ctas <= 0) + const int pool = rocke_wgrad_conv_streamk_ctas(s, NULL); + if(pool <= 0) { if(out_status) *out_status = ROCKE_ERR_VALUE; @@ -182,7 +230,7 @@ rocke_streamk_iter_partition_t return rocke_streamk_iter_partition_make((wg_M + s->tile_m - 1) / s->tile_m * groups, (wg_N + s->tile_n - 1) / s->tile_n, (wg_K + s->tile_k - 1) / s->tile_k, - s->streamk_ctas, + pool, strcmp(s->streamk, "persistent") == 0, out_status); } @@ -215,8 +263,9 @@ rocke_status_t rocke_wgrad_conv_streamk_workspace_layout( if(st != ROCKE_OK) return st; const rocke_streamk_iter_plan_t plan = rocke_streamk_iter_plan(&part); - const size_t flags = (size_t)4 * (size_t)plan.sk_ctas * (size_t)(s->warp_m * s->warp_n); - *flags_bytes = (flags + 255) / 256 * 256; + *flags_bytes = rocke_streamk_flags_bytes(&plan, s->warp_m * s->warp_n, &st); + if(st != ROCKE_OK) + return st; *partials_bytes = (size_t)plan.sk_ctas * (size_t)s->tile_m * (size_t)s->tile_n * sizeof(float); return ROCKE_OK; } @@ -627,9 +676,8 @@ bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv * any grouped conv it disagreed with Python (which tests cpg); * - it did not exempt the two-stage path, which stores f32 to a * workspace and emits no atomic at all. - * The local per-group computation is deliberate: the shared - * rocke_wgrad_conv_spec_wg_N() helper still returns the dense Z*Y*X*C - * and is used for workspace sizing, so it is not interchangeable here. */ + * The product is recomputed locally in 64-bit rather than through + * rocke_wgrad_conv_spec_wg_N() (the same per-group value, as an int). */ const bool effective_two_stage_gate = s->two_stage && sk > 1; const char* dt = s->dtype_d ? s->dtype_d : "fp16"; if(!effective_two_stage_gate && (strcmp(dt, "fp16") == 0 || strcmp(dt, "bf16") == 0)) @@ -2865,6 +2913,17 @@ static void wgrad_emit_streamk(rocke_conv_build_ctx_t* ctx, for(int i = 0; i < ctx->num_accs; ++i) tile_accs[i] = ctx->final_accs[i]; + /* Whether this item covers the start / the end of its tile. Built only + * for the flag-handoff fixups, so the atomic and workspace bodies carry + * no dead compares. Mirrors Python. */ + rocke_value_t* started = NULL; + rocke_value_t* ended = NULL; + if(strcmp(reduction, "linear") == 0 || strcmp(reduction, "tree") == 0) + { + started = rocke_b_cmp_eq(b, lo, c0); + ended = rocke_b_cmp_eq(b, hi, c_ipt); + } + if(strcmp(reduction, "atomic") == 0) { wgrad_emit_split_k_epilogue_f32(b, ctx, spec, dW, wg_M, wg_N); @@ -2876,8 +2935,6 @@ static void wgrad_emit_streamk(rocke_conv_build_ctx_t* ctx, } else if(strcmp(reduction, "linear") == 0) { - rocke_value_t* started = rocke_b_cmp_eq(b, lo, c0); - rocke_value_t* ended = rocke_b_cmp_eq(b, hi, c_ipt); rocke_if_else_t ie = rocke_b_scf_if_else(b, started); rocke_b_region_enter(b, ie.then_region); { @@ -2919,8 +2976,6 @@ static void wgrad_emit_streamk(rocke_conv_build_ctx_t* ctx, /* Tree: the round count is a kernarg, so the rounds are a loop * carrying the partial sums, the pair stride and a dropped-out * flag (see the Python comment). */ - rocke_value_t* started = rocke_b_cmp_eq(b, lo, c0); - rocke_value_t* ended = rocke_b_cmp_eq(b, hi, c_ipt); rocke_value_t* not_started = rocke_b_lnot(b, started); rocke_value_t* owner = wgrad_sk_cta_of(&e, tile_start); rocke_value_t* local = rocke_b_sub(b, sk_cta, owner); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py index ec9eb54522d9..42c27763368d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py @@ -353,8 +353,6 @@ compute_streamk_grid_size, emit_streamk_decode, emit_streamk_iter_range, - emit_streamk_partial_load_accumulate, - emit_streamk_partial_store, emit_streamk_sk_start_iter, streamk_end_iter, streamk_iter_partition, @@ -625,8 +623,6 @@ "compute_streamk_grid_size", "emit_streamk_decode", "emit_streamk_iter_range", - "emit_streamk_partial_load_accumulate", - "emit_streamk_partial_store", "emit_streamk_sk_start_iter", "streamk_end_iter", "streamk_iter_partition", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/streamk.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/streamk.py index 85cbb69a9ee4..3fb1428fdae6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/streamk.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/streamk.py @@ -21,11 +21,12 @@ and is the path CK Tile defaults to. This module ships the partitioner math (a pure-Python view) and the -IR-side glue that -``instances/streamk_gemm.py`` uses. The kernel-shape primitives stay +IR-side glue that ``instances/streamk_gemm.py`` and the stream-K +implicit-GEMM wgrad conv kernel use. The kernel-shape primitives stay in :mod:`rocke.helpers.atoms` / :mod:`rocke.helpers.loads` / :mod:`rocke.helpers.epilogues`; this file is purely about -*partitioning* the work, not about running the per-tile GEMM. +*partitioning* the work, not about running the per-tile GEMM or its +cross-CTA fixup. What we ship today: @@ -37,12 +38,17 @@ * :func:`emit_streamk_decode` -- IR-side decode of a linear partition id into the SSA ``(m_tile, n_tile, k_iter, is_first, is_last)`` bundle the kernel needs. - -v1 implements the ``Atomic`` strategy end-to-end and the -``Reduction`` strategy's *decode* surface (so the spec layer can -emit the right kernel shape); the per-tile reduction pass that the -``Reduction`` strategy needs lands with the StreamK GEMM kernel -itself in . +* :class:`StreamKIterPartition` and friends -- CK Tile's + iteration-balanced (DP + stream-K) partitioner, plus the SSA emitters + that map a CTA onto its MAC-iteration range. + +The ``Reduction`` strategy's cross-CTA fixup (partial slots + a +per-(CTA, wave) flag table, Linear or Tree fan-in) is kernel-specific +and lives with its kernel; the reference implementation is the nested +``_emit_streamk_tiles`` of ``build_implicit_gemm_conv_wgrad`` in +``library/kernels/common/conv_implicit_gemm_wgrad.py`` (C++ twin: +``wgrad_emit_streamk`` in ``conv_implicit_gemm_wgrad.cpp``). +:meth:`StreamKIterPartition.flags_bytes` sizes its flag table. """ from __future__ import annotations @@ -62,8 +68,6 @@ "compute_streamk_grid_size", "emit_streamk_decode", "emit_streamk_iter_range", - "emit_streamk_partial_load_accumulate", - "emit_streamk_partial_store", "emit_streamk_sk_start_iter", "streamk_end_iter", "streamk_iter_partition", @@ -181,67 +185,6 @@ def emit_streamk_decode( ) -def emit_streamk_partial_store( - b: IRBuilder, - *, - workspace: Value, - flag_table: Value, - workspace_off: Value, - flag_off: Value, - value: Value, - is_first: Value, -) -> None: - """Cooperative ``Reduction`` strategy partial-store + signal. - - P37: each contributing CTA stores its partial K-sum at the - ``(m_tile, n_tile, k_iter)`` slot in the workspace and bumps the - flag-table counter for ``(m_tile, n_tile)`` via ``atomic_add(+1)``. - The first contributor's store is unconditionally a fresh write - (``is_first=True``); later contributors atomic-add into the - workspace too so the order doesn't matter. - - Reference: CK Tile ``streamk_common.hpp::SignalStorePartialDone`` - + ``streamk_gemm_kernel.hpp:448-504`` (Tree-reduction). - """ - # Initial store: lane-0-only write (the workspace is per-lane sized - # already in the caller's view) — the consumer in - # :func:`emit_streamk_partial_load_accumulate` reads each slot once. - with b.scf_if(is_first): - b.global_store(workspace, workspace_off, value, align=4) - # Non-first contributors: atomic add into the same f32 slot (the - # ``is_last`` reducer reads the converged value). - with b.scf_if(b.lnot(is_first)): - b.global_atomic_add(workspace, workspace_off, value) - # Bump the flag counter so the last-finishing CTA can detect "I'm - # the reducer" via flag == k_iters_per_tile. - b.global_atomic_add(flag_table, flag_off, b.const_i32(1)) - - -def emit_streamk_partial_load_accumulate( - b: IRBuilder, - *, - workspace: Value, - workspace_off: Value, - is_last: Value, -) -> Value: - """Read the converged f32 partial sum at ``workspace_off`` once - every contributor has signalled (``is_last == True`` is the - caller's responsibility). - - The flag-table-based wait protocol is a busy-loop on - ``atomic_load(flag_table[tile_id]) == k_iters_per_tile``; we expose - it as a separate helper because the busy-loop shape differs - between Linear / Tree reduction. For now the Atomic strategy is - the canonical path; this helper plus - :func:`emit_streamk_partial_store` give callers the surface - needed to opt into the Reduction strategy. - - Reference: CK Tile ``streamk_common.hpp:34-73`` - (``WaitStorePartialDone``). - """ - return b.global_load_f32(workspace, workspace_off) - - # --------------------------------------------------------------------------- # Iteration-balanced stream-K (CK Tile StreamKTilePartitioner) # --------------------------------------------------------------------------- @@ -313,7 +256,16 @@ def total_sk_iters(self) -> int: @property def iters_per_sk_cta(self) -> int: - return self.total_sk_iters // self.sk_ctas if self.sk_ctas else 0 + """Base MAC iterations per stream-K CTA; always >= 1. + + Kernels divide by this value (the owner lookup) on every path, + including an all-DP launch, so it is clamped to 1 when there are no + stream-K CTAs. With ``sk_ctas > 0`` it is already >= 1 because + :attr:`sk_tiles` falls back to all-DP otherwise. + """ + if not self.sk_ctas: + return 1 + return self.total_sk_iters // self.sk_ctas @property def extra_iters(self) -> int: @@ -355,10 +307,17 @@ def tree_rounds(self) -> int: contributors = self.max_linear_partners + 1 return (contributors - 1).bit_length() - @property - def flags_bytes(self) -> int: - """One i32 flag per stream-K CTA, padded to 256 bytes (CK layout).""" - return _ceil_div(4 * self.sk_ctas, 256) * 256 + def flags_bytes(self, n_waves: int) -> int: + """Bytes of the fixup flag table: one i32 flag per (stream-K CTA, + wave), padded to 256 bytes (CK layout). + + ``n_waves`` is the waves per CTA that publish/wait independently + (``warp_m * warp_n`` for the wgrad kernel); pass 1 for one flag per + CTA. 0 for an all-DP partition. + """ + if n_waves <= 0: + raise ValueError(f"stream-K flags: n_waves must be > 0 (got {n_waves})") + return _ceil_div(4 * self.sk_ctas * n_waves, 256) * 256 def streamk_iter_partition( @@ -416,7 +375,7 @@ class StreamKIterArgs(NamedTuple): iters_per_tile: Value # i32 MAC iterations per output tile dp_tiles: Value # i32 tiles owned whole by data-parallel CTAs total_dp_iters: Value # i32 dp_tiles * iters_per_tile - iters_per_sk_cta: Value # i32 base iterations per stream-K CTA (> 0) + iters_per_sk_cta: Value # i32 base iterations per stream-K CTA (>= 1) extra_iters: Value # i32 stream-K CTAs that own one more iteration diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json index b4822a68b1ae..d5dd4370d9d9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json @@ -898,6 +898,14 @@ "sha256": "5c0a33026f958a437bea9ad18b33d3e9dbcc9821923b2373e203f3edb1cb0222", "status": "ok" }, + "conv_wgrad/gfx942/n1h8c16k32r3_sk_atomic_pers": { + "arch": "gfx942", + "bytes": 25156, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_942_sk_atomic_pers_N1H8W8C16_K32Y3X3_t64x32x16_w2x2_a16x16x16_mem_default_sk_skpers_skratomic", + "sha256": "1d70097054f9a71c691f4b89a87c5490e8b2371cc3c243cb7a725f8640744d20", + "status": "ok" + }, "conv_wgrad/gfx942/n1h8c64k64r3_g4": { "arch": "gfx942", "bytes": 22466, @@ -914,6 +922,14 @@ "sha256": "4c3f744ae6f1ae5198f47461f06f90611d4d15fac5322b87147501133956c0c2", "status": "ok" }, + "conv_wgrad/gfx942/n1h8c64k64r3_g4_sk_workspace": { + "arch": "gfx942", + "bytes": 25798, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_942_g4_sk_workspace_N1H8W8C64_K64Y3X3G4_t64x32x16_w2x2_a16x16x16_mem_default_sk_skrworkspace", + "sha256": "21bee95d9ae4d04fde79971cb14764bc3207d0c63d590a45eb2e1e6340a01b0e", + "status": "ok" + }, "conv_wgrad/gfx942/n1h8c64k64r3_g4_spk4": { "arch": "gfx942", "bytes": 26896, @@ -946,6 +962,22 @@ "sha256": "1819f9e2134239476b7315a2a059a7c2014004af3829fd5a764dbb06fea4b05d", "status": "ok" }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_tree": { + "arch": "gfx950", + "bytes": 60379, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_tree_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk_skrtree", + "sha256": "1d2674baa67db756ba52747582b5551d663aaf9e3d56b8069bd83665de9a5485", + "status": "ok" + }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_workspace_pers": { + "arch": "gfx950", + "bytes": 31327, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_workspace_pers_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk_skpers_skrworkspace", + "sha256": "9d25da3d386646386e6e4075eec6794d76c3eda4da66121976d59a15ee4f0cde", + "status": "ok" + }, "conv_wgrad/gfx950/n1h8c16k32r3_spk4": { "arch": "gfx950", "bytes": 26145, @@ -970,6 +1002,14 @@ "sha256": "7ec4ccded2d060144039b3c5a644dcc8d4fd36c3f95c8a79384d205d687a5ba6", "status": "ok" }, + "conv_wgrad/gfx950/n1h8c64k64r3_g4_sk_linear_pers": { + "arch": "gfx950", + "bytes": 61458, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_g4_sk_linear_pers_N1H8W8C64_K64Y3X3G4_t64x64x16_w2x2_a32x32x16_mem_default_sk_skpers", + "sha256": "7e23d8219721828777ad6b20909fc105514a4304d1cb638f0779dadafa35eb7f", + "status": "ok" + }, "conv_wgrad/gfx950/n1h8c64k64r3_g4_spk4": { "arch": "gfx950", "bytes": 38655, @@ -1428,8 +1468,8 @@ "gfx1201": 7, "gfx1250": 4, "gfx90a": 5, - "gfx942": 37, - "gfx950": 111 + "gfx942": 39, + "gfx950": 114 }, "expected_failures": 4, "families": { @@ -1443,7 +1483,7 @@ "conv_direct_depthwise_tiled": 3, "conv_direct_dgrad": 5, "conv_direct_nongrouped": 7, - "conv_wgrad": 16, + "conv_wgrad": 21, "conv_wgrad_reduce": 3, "conv_wgrad_two_stage": 2, "deep_fused_conv": 9, @@ -1453,7 +1493,7 @@ "moe": 9, "nontemporal": 2 }, - "ok_cases": 175 + "ok_cases": 180 } }, "llvm22": { @@ -2354,6 +2394,14 @@ "sha256": "a4d404a44b5f8a351ce2f2d24957c976c79d0ee13887ff15c7c7a47438575b82", "status": "ok" }, + "conv_wgrad/gfx942/n1h8c16k32r3_sk_atomic_pers": { + "arch": "gfx942", + "bytes": 25289, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_942_sk_atomic_pers_N1H8W8C16_K32Y3X3_t64x32x16_w2x2_a16x16x16_mem_default_sk_skpers_skratomic", + "sha256": "f4fa2ad553aa4a1180c5a4a67d3d13cd86f0ee6671bfad4954dd55c93fe67e3d", + "status": "ok" + }, "conv_wgrad/gfx942/n1h8c64k64r3_g4": { "arch": "gfx942", "bytes": 22599, @@ -2370,6 +2418,14 @@ "sha256": "2ef49bc86988a2fe9177d29325c801bb9ec0e7631a3255cc5a6d0fc2c86179be", "status": "ok" }, + "conv_wgrad/gfx942/n1h8c64k64r3_g4_sk_workspace": { + "arch": "gfx942", + "bytes": 25931, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_942_g4_sk_workspace_N1H8W8C64_K64Y3X3G4_t64x32x16_w2x2_a16x16x16_mem_default_sk_skrworkspace", + "sha256": "266d9829bcf2d91378ebbf7d2840fce771486b21d08c460fda3bbc1f206309e2", + "status": "ok" + }, "conv_wgrad/gfx942/n1h8c64k64r3_g4_spk4": { "arch": "gfx942", "bytes": 27035, @@ -2402,6 +2458,22 @@ "sha256": "a18cb4d993ab0e837f5e4cbda76848fce46e76a8393900e755b3cdaed05235f8", "status": "ok" }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_tree": { + "arch": "gfx950", + "bytes": 60518, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_tree_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk_skrtree", + "sha256": "efe578b01d1d931ed9ee46a7974a5167cb5032222f20718d28c8e546046d577a", + "status": "ok" + }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_workspace_pers": { + "arch": "gfx950", + "bytes": 31460, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_workspace_pers_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk_skpers_skrworkspace", + "sha256": "f2b8b993253d99f9d3b796085d5ac5c6ef3d15ae453ec0bd0aee0d6397591c28", + "status": "ok" + }, "conv_wgrad/gfx950/n1h8c16k32r3_spk4": { "arch": "gfx950", "bytes": 26284, @@ -2426,6 +2498,14 @@ "sha256": "e24daa9ac08990da4b1d1ad4ccbeece698279dcda588cb28c785ee2d4e21cd38", "status": "ok" }, + "conv_wgrad/gfx950/n1h8c64k64r3_g4_sk_linear_pers": { + "arch": "gfx950", + "bytes": 61597, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_g4_sk_linear_pers_N1H8W8C64_K64Y3X3G4_t64x64x16_w2x2_a32x32x16_mem_default_sk_skpers", + "sha256": "f978aa87c489a5989e562217a201cc15675aa1b8fd38a816846881a2697d55cb", + "status": "ok" + }, "conv_wgrad/gfx950/n1h8c64k64r3_g4_spk4": { "arch": "gfx950", "bytes": 38794, @@ -2884,8 +2964,8 @@ "gfx1201": 7, "gfx1250": 4, "gfx90a": 5, - "gfx942": 37, - "gfx950": 111 + "gfx942": 39, + "gfx950": 114 }, "expected_failures": 4, "families": { @@ -2899,7 +2979,7 @@ "conv_direct_depthwise_tiled": 3, "conv_direct_dgrad": 5, "conv_direct_nongrouped": 7, - "conv_wgrad": 16, + "conv_wgrad": 21, "conv_wgrad_reduce": 3, "conv_wgrad_two_stage": 2, "deep_fused_conv": 9, @@ -2909,7 +2989,7 @@ "moe": 9, "nontemporal": 2 }, - "ok_cases": 175 + "ok_cases": 180 } }, "llvm23": { @@ -3810,6 +3890,14 @@ "sha256": "24c106cf812be8655f0e9ed1a272346ecccc29dd2e893342eada4b6dbe9dd227", "status": "ok" }, + "conv_wgrad/gfx942/n1h8c16k32r3_sk_atomic_pers": { + "arch": "gfx942", + "bytes": 25353, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_942_sk_atomic_pers_N1H8W8C16_K32Y3X3_t64x32x16_w2x2_a16x16x16_mem_default_sk_skpers_skratomic", + "sha256": "2863b1a0e57941df6e773f9e981c40f1843856237554d7615e92731db6ff06c4", + "status": "ok" + }, "conv_wgrad/gfx942/n1h8c64k64r3_g4": { "arch": "gfx942", "bytes": 22663, @@ -3826,6 +3914,14 @@ "sha256": "2558f2ea4588e91e12e697b23a095805e6a83b20d9c01d9a7f78be69a7e688ce", "status": "ok" }, + "conv_wgrad/gfx942/n1h8c64k64r3_g4_sk_workspace": { + "arch": "gfx942", + "bytes": 25995, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_942_g4_sk_workspace_N1H8W8C64_K64Y3X3G4_t64x32x16_w2x2_a16x16x16_mem_default_sk_skrworkspace", + "sha256": "7c6ce1e064d7d92ebcfb0110334c89e674a828bb65bd784d566e5e9b58129721", + "status": "ok" + }, "conv_wgrad/gfx942/n1h8c64k64r3_g4_spk4": { "arch": "gfx942", "bytes": 27099, @@ -3858,6 +3954,22 @@ "sha256": "65640f53273c121fa0e1127e49cce824a2e7e4f7ec8a6ccb53725570d419b2e9", "status": "ok" }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_tree": { + "arch": "gfx950", + "bytes": 60582, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_tree_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk_skrtree", + "sha256": "e50b150b37785afabbc8c8f136c069bdd88ffbbff332d14dbe0f9db82b88c975", + "status": "ok" + }, + "conv_wgrad/gfx950/n1h8c16k32r3_sk_workspace_pers": { + "arch": "gfx950", + "bytes": 31524, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_sk_workspace_pers_N1H8W8C16_K32Y3X3_t64x64x16_w2x2_a32x32x16_mem_default_sk_skpers_skrworkspace", + "sha256": "3dbc6ebcdf5be7498ae4f73eeb9b99932a9cea7fe80df7077e4d629959f137dc", + "status": "ok" + }, "conv_wgrad/gfx950/n1h8c16k32r3_spk4": { "arch": "gfx950", "bytes": 26348, @@ -3882,6 +3994,14 @@ "sha256": "822ab7ae23cca12ac3351b10bec5ee0cc5a18945c210f3f39e1c8f38b76252d0", "status": "ok" }, + "conv_wgrad/gfx950/n1h8c64k64r3_g4_sk_linear_pers": { + "arch": "gfx950", + "bytes": 61661, + "family": "conv_wgrad", + "kernel_name": "irhash_wgrad_950_g4_sk_linear_pers_N1H8W8C64_K64Y3X3G4_t64x64x16_w2x2_a32x32x16_mem_default_sk_skpers", + "sha256": "1b892d27711fae5606fa0bf1eca307a105ab3a9db98a9860d46e72db11937992", + "status": "ok" + }, "conv_wgrad/gfx950/n1h8c64k64r3_g4_spk4": { "arch": "gfx950", "bytes": 38858, @@ -4355,8 +4475,8 @@ "gfx1201": 7, "gfx1250": 8, "gfx90a": 5, - "gfx942": 37, - "gfx950": 111 + "gfx942": 39, + "gfx950": 114 }, "expected_failures": 0, "families": { @@ -4370,7 +4490,7 @@ "conv_direct_depthwise_tiled": 3, "conv_direct_dgrad": 5, "conv_direct_nongrouped": 7, - "conv_wgrad": 16, + "conv_wgrad": 21, "conv_wgrad_reduce": 3, "conv_wgrad_two_stage": 2, "deep_fused_conv": 9, @@ -4380,7 +4500,7 @@ "moe": 9, "nontemporal": 2 }, - "ok_cases": 179 + "ok_cases": 184 } } }, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py index 67bb248b8932..b15e19993d40 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py @@ -1979,6 +1979,103 @@ def add(family, case_id, arch, build): streamk_reduction="tree", ), ) + # The remaining {dp_sk, persistent} x {linear, tree, atomic, workspace} + # combinations, so every fixup is pinned under both launch modes. + add( + "conv_wgrad", + "conv_wgrad/gfx950/n1h8c16k32r3_sk_tree", + "gfx950", + build_conv_wgrad( + "irhash_wgrad_950_sk_tree", + "gfx950", + wgrad1, + wave_size=64, + wtm=32, + wtn=32, + wtk=16, + tile_m=64, + tile_n=64, + tile_k=16, + streamk="dp_sk", + streamk_reduction="tree", + ), + ) + add( + "conv_wgrad", + "conv_wgrad/gfx942/n1h8c64k64r3_g4_sk_workspace", + "gfx942", + build_conv_wgrad( + "irhash_wgrad_942_g4_sk_workspace", + "gfx942", + wgrad_g4, + wave_size=64, + wtm=16, + wtn=16, + wtk=16, + tile_m=64, + tile_n=32, + tile_k=16, + streamk="dp_sk", + streamk_reduction="workspace", + ), + ) + add( + "conv_wgrad", + "conv_wgrad/gfx950/n1h8c64k64r3_g4_sk_linear_pers", + "gfx950", + build_conv_wgrad( + "irhash_wgrad_950_g4_sk_linear_pers", + "gfx950", + wgrad_g4, + wave_size=64, + wtm=32, + wtn=32, + wtk=16, + tile_m=64, + tile_n=64, + tile_k=16, + streamk="persistent", + ), + ) + add( + "conv_wgrad", + "conv_wgrad/gfx942/n1h8c16k32r3_sk_atomic_pers", + "gfx942", + build_conv_wgrad( + "irhash_wgrad_942_sk_atomic_pers", + "gfx942", + wgrad1, + wave_size=64, + wtm=16, + wtn=16, + wtk=16, + tile_m=64, + tile_n=32, + tile_k=16, + dtype_d="fp32", + streamk="persistent", + streamk_reduction="atomic", + ), + ) + add( + "conv_wgrad", + "conv_wgrad/gfx950/n1h8c16k32r3_sk_workspace_pers", + "gfx950", + build_conv_wgrad( + "irhash_wgrad_950_sk_workspace_pers", + "gfx950", + wgrad1, + wave_size=64, + wtm=32, + wtn=32, + wtk=16, + tile_m=64, + tile_n=64, + tile_k=16, + streamk="persistent", + streamk_reduction="workspace", + ), + ) # Grouped + cshuffle epilogue (MFMA): the LDS-staged store threads the # per-group k_out fold (group*kpg) and derives its store-vec from cpg. add( diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/streamk_iter_partition.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/streamk_iter_partition.cpp index 392ef616ac65..1195cb0989a9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/streamk_iter_partition.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/streamk_iter_partition.cpp @@ -33,14 +33,15 @@ struct Vector { int m, n, ipt, w; bool persistent; + /* flags_bytes is for n_waves = 1; ipsc is clamped to 1 when all-DP. */ int sk_tiles, sk_ctas, ipsc, extra, dp_tiles, grid, partners, rounds, flags_bytes; }; const Vector k_vectors[] = { - {4, 4, 8, 16, false, 0, 0, 0, 0, 16, 16, 0, 0, 0}, + {4, 4, 8, 16, false, 0, 0, 1, 0, 16, 16, 0, 0, 0}, {2, 3, 10, 8, false, 6, 8, 7, 4, 0, 8, 2, 2, 256}, {5, 7, 12, 16, false, 19, 16, 14, 4, 16, 32, 1, 1, 256}, - {1, 3, 2, 8, false, 0, 0, 0, 0, 3, 3, 0, 0, 0}, + {1, 3, 2, 8, false, 0, 0, 1, 0, 3, 3, 0, 0, 0}, {5, 7, 12, 16, true, 19, 16, 14, 4, 16, 16, 1, 1, 256}, {3, 3, 64, 4, false, 5, 4, 80, 0, 4, 8, 1, 1, 256}, {7, 5, 3, 16, false, 19, 16, 3, 9, 16, 32, 1, 1, 256}, @@ -69,7 +70,7 @@ void check_vectors() CHECK(r.grid_size == v.grid); CHECK(r.max_linear_partners == v.partners); CHECK(r.tree_rounds == v.rounds); - CHECK(r.flags_bytes == v.flags_bytes); + CHECK(rocke_streamk_flags_bytes(&r, 1, NULL) == (size_t)v.flags_bytes); } const rocke_streamk_iter_partition_t p = make(5, 7, 12, 16, false); const int starts[] = {192, 207, 222, 237, 252}; @@ -81,6 +82,34 @@ void check_vectors() } } +void check_flags_bytes_per_wave() +{ + /* One i32 flag per (stream-K CTA, wave), padded to 256 bytes. */ + struct + { + int m, n, ipt, w, n_waves; + size_t want; + } const cases[] = { + {5, 7, 12, 16, 4, 256}, + {5, 7, 12, 16, 8, 512}, + {2, 3, 10, 8, 9, 512}, + {4, 4, 8, 16, 4, 0}, + }; + for(const auto& c : cases) + { + const rocke_streamk_iter_partition_t p = make(c.m, c.n, c.ipt, c.w, false); + const rocke_streamk_iter_plan_t r = rocke_streamk_iter_plan(&p); + rocke_status_t st = ROCKE_ERR_VALUE; + CHECK(rocke_streamk_flags_bytes(&r, c.n_waves, &st) == c.want); + CHECK(st == ROCKE_OK); + } + const rocke_streamk_iter_partition_t p = make(5, 7, 12, 16, false); + const rocke_streamk_iter_plan_t r = rocke_streamk_iter_plan(&p); + rocke_status_t st = ROCKE_OK; + CHECK(rocke_streamk_flags_bytes(&r, 0, &st) == 0); + CHECK(st == ROCKE_ERR_VALUE); +} + void check_rejects_non_positive() { rocke_status_t st = ROCKE_OK; @@ -96,9 +125,9 @@ void check_invariants(int m, int n, int ipt, int w) const rocke_streamk_iter_partition_t p = make(m, n, ipt, w, false); const rocke_streamk_iter_plan_t r = rocke_streamk_iter_plan(&p); CHECK(r.dp_tiles + r.sk_tiles == r.num_tiles); + CHECK(r.iters_per_sk_cta >= 1); /* a kernel divisor even when all-DP */ if(!r.sk_ctas) return; - CHECK(r.iters_per_sk_cta >= 1); int cursor = r.total_dp_iters; int prev_first_tile = -1; int run = 0; /* contributors to the current tile */ @@ -132,6 +161,7 @@ void check_invariants(int m, int n, int ipt, int w) int main(void) { check_vectors(); + check_flags_bytes_per_wave(); check_rejects_non_positive(); const int ipts[] = {1, 2, 3, 5, 8, 13}; const int ws[] = {1, 2, 3, 4, 7, 8, 16}; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_streamk_iter_partition.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_streamk_iter_partition.py index c345efa7afb5..7471cfca79e6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_streamk_iter_partition.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_streamk_iter_partition.py @@ -22,12 +22,13 @@ # (m_tiles, n_tiles, iters_per_tile, max_active_wgs, persistent) -> # (sk_tiles, sk_ctas, iters_per_sk_cta, extra_iters, dp_tiles, grid_size, -# max_linear_partners, tree_rounds, flags_bytes) +# max_linear_partners, tree_rounds, flags_bytes(n_waves=1)) +# iters_per_sk_cta is clamped to 1 for an all-DP partition (a kernel divisor). _VECTORS = ( - ((4, 4, 8, 16, False), (0, 0, 0, 0, 16, 16, 0, 0, 0)), # tiles % W == 0 + ((4, 4, 8, 16, False), (0, 0, 1, 0, 16, 16, 0, 0, 0)), # tiles % W == 0 ((2, 3, 10, 8, False), (6, 8, 7, 4, 0, 8, 2, 2, 256)), # tiles < W ((5, 7, 12, 16, False), (19, 16, 14, 4, 16, 32, 1, 1, 256)), # tiles > W - ((1, 3, 2, 8, False), (0, 0, 0, 0, 3, 3, 0, 0, 0)), # too little SK work + ((1, 3, 2, 8, False), (0, 0, 1, 0, 3, 3, 0, 0, 0)), # too little SK work ((5, 7, 12, 16, True), (19, 16, 14, 4, 16, 16, 1, 1, 256)), # persistent ((3, 3, 64, 4, False), (5, 4, 80, 0, 4, 8, 1, 1, 256)), ((7, 5, 3, 16, False), (19, 16, 3, 9, 16, 32, 1, 1, 256)), @@ -44,7 +45,7 @@ def _plan(p: StreamKIterPartition) -> tuple: p.grid_size, p.max_linear_partners, p.tree_rounds, - p.flags_bytes, + p.flags_bytes(1), ) @@ -73,6 +74,19 @@ def test_start_end_iter_vectors(self): [streamk_end_iter(p, c) for c in range(5)], [207, 222, 237, 252, 266] ) + def test_flags_bytes_per_wave(self): + # One i32 flag per (stream-K CTA, wave), padded to 256 bytes. + for args, n_waves, want in ( + ((5, 7, 12, 16), 4, 256), # 4 * 16 * 4 = 256 + ((5, 7, 12, 16), 8, 512), + ((2, 3, 10, 8), 9, 512), # 4 * 8 * 9 = 288 + ((4, 4, 8, 16), 4, 0), # all-DP: no flags + ): + with self.subTest(args=args, n_waves=n_waves): + self.assertEqual(_make(*args).flags_bytes(n_waves), want) + with self.assertRaisesRegex(ValueError, "n_waves must be > 0"): + _make(5, 7, 12, 16).flags_bytes(0) + def test_rejects_non_positive_counts(self): for kw in ("m_tiles", "n_tiles", "iters_per_tile", "max_active_wgs"): args = { @@ -95,12 +109,13 @@ class TestStreamKIterPartitionInvariants(unittest.TestCase): def _check(self, p: StreamKIterPartition) -> None: ipt = p.iters_per_tile self.assertEqual(p.dp_tiles + p.sk_tiles, p.num_tiles) + # A kernel divisor even when all-DP. + self.assertGreaterEqual(p.iters_per_sk_cta, 1) if not p.sk_ctas: self.assertEqual(p.sk_tiles, 0) return # Every SK CTA owns >= 1 iteration and the ranges tile the SK span # contiguously. - self.assertGreaterEqual(p.iters_per_sk_cta, 1) cursor = p.total_dp_iters contributors = {} for c in range(p.sk_ctas): From 4c3159605506171cc7aa8a91b86afec701d7f61e Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Thu, 8 Oct 2026 16:24:02 +0000 Subject: [PATCH 15/20] refactor(rocke): move wgrad stream-K from the dispatcher to the sweep The grouped-conv dispatcher and its benchmark are due for a refactor, so this change no longer extends them: the dispatcher stream-K selection and the grouped-conv benchmark flags are reverted. Stream-K wgrad candidates are instead opt-in in the implicit-GEMM sweep (--streamk, --streamk-reductions, --streamk-ctas) for the JIT, AOT-cache and run-from-cache paths. The cache identity and kernel label keep their existing values when stream-K is off. ISSUE ID : AICK-1943 Co-Authored-By: Claude Opus 5.5 --- .../common/benchmark_conv_compare.py | 35 + .../common/benchmark_implicit_gemm_conv.py | 795 +++++++++++++++--- .../grouped_conv/benchmark_grouped_conv.py | 167 +--- .../library/benchmarks/common/kernel_cache.py | 9 + .../library/benchmarks/common/kernel_label.py | 2 + .../library/benchmarks/common/kernel_sweep.py | 178 +++- .../library/dispatch/grouped_convolution.py | 219 +---- .../common/README_conv_implicit_gemm_wgrad.md | 24 +- .../test_grouped_conv_wgrad_dispatch.py | 180 ---- .../optimization/optimization_runbook.md | 2 +- 10 files changed, 922 insertions(+), 689 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_conv_compare.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_conv_compare.py index ca8a82705a8a..ab7c62d3d5b0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_conv_compare.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_conv_compare.py @@ -230,6 +230,12 @@ def _build_implicit_only_args(args) -> list[str]: a += ["--split-k", str(args.split_k)] if args.split_k_prune is not None: a += ["--split-k-prune", str(args.split_k_prune)] + if args.streamk != "off": + a += ["--streamk", args.streamk] + if args.streamk_reductions is not None: + a += ["--streamk-reductions", args.streamk_reductions] + if args.streamk_ctas != -1: + a += ["--streamk-ctas", str(args.streamk_ctas)] return a @@ -426,6 +432,35 @@ def main() -> int: "Only effective with --split-k 0. Forwarded to implicit-GEMM only." ), ) + implicit_grp.add_argument( + "--streamk", + default="off", + choices=["off", "dp_sk", "persistent", "all"], + help=( + "wgrad stream-K leg (off, dp_sk, persistent or all launch modes). " + "Forwarded to implicit-GEMM only." + ), + ) + implicit_grp.add_argument( + "--streamk-reductions", + default=None, + dest="streamk_reductions", + metavar="LIST", + help=( + "comma-separated stream-K reductions (linear,tree,workspace,atomic) " + "for --streamk. Forwarded to implicit-GEMM only." + ), + ) + implicit_grp.add_argument( + "--streamk-ctas", + type=int, + default=-1, + dest="streamk_ctas", + metavar="N", + help=( + "stream-K CTA pool (-1 = arch default). Forwarded to implicit-GEMM " "only." + ), + ) miopen_grp = parser.add_argument_group( "MIOpen input", diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py index 77688e97211d..f95008aebae8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py @@ -84,6 +84,9 @@ CACHE_TILE_MN as _TILE_MN, CACHE_WARP_MN as _WARP_MN, CACHE_WARP_TILE_MN as _WARP_TILE_MN, + CACHE_WGRAD_STREAMK_MODES as _STREAMK_MODES, + CACHE_WGRAD_STREAMK_REDUCTIONS as _STREAMK_REDUCTIONS, + _AOT_ALIAS_PIPELINES, ) # gfx1250-only: wider CUs allow larger tiles and more warp concurrency. @@ -148,6 +151,10 @@ class WgradCombo(NamedTuple): # 16-bit atomics and allocates no scratch), which is why the axis is only # expanded on the two-stage work list. ws_replicas: int = 0 + # Stream-K launch mode ("dp_sk" / "persistent") and reduction. "off" on + # every other leg; the stream-K leg keeps split_k = 1 (stream-K owns K). + streamk: str = "off" + streamk_reduction: str = "" @dataclass @@ -182,6 +189,10 @@ class Result: # argument as group_merge -- an --ws-replicas sweep prints one row per R # and they must be tellable apart. ws_replicas: int = 0 + # Stream-K leg only: launch mode, reduction and the resolved CTA pool. + streamk: str = "off" + streamk_reduction: str = "" + streamk_ctas: int = 0 # --------------------------------------------------------------------------- @@ -927,6 +938,43 @@ def main() -> int: ).replace("%(sweep)s", str(list(_WS_REPLICAS_SWEEP))), ) + parser.add_argument( + "--streamk", + default="off", + choices=["off", "dp_sk", "persistent", "all"], + dest="streamk", + help=( + "wgrad stream-K kernels (gfx942/gfx950): off (default) = none; " + "dp_sk / persistent = that launch mode; all = both. The JIT sweep " + "adds a stream-K leg (default epilogue, split_k=1) over the " + "--streamk-reductions; --compile-all adds the stream-K kernels to " + "the cache; --run-from-cache runs the cached ones that match." + ), + ) + parser.add_argument( + "--streamk-reductions", + default=",".join(_STREAMK_REDUCTIONS), + dest="streamk_reductions", + metavar="LIST", + help=( + "comma-separated stream-K reductions to sweep with --streamk " + "(default: %(default)s). 'atomic' accumulates into dW with f32 " + "atomics, so it only applies when dW is fp32." + ), + ) + parser.add_argument( + "--streamk-ctas", + type=int, + default=-1, + dest="streamk_ctas", + metavar="N", + help=( + "stream-K CTA pool, a launch parameter (default: -1 = the arch " + "default: one CTA per CU for linear/tree, several per CU for " + "workspace/atomic)." + ), + ) + parser.add_argument( "--csv-top", type=int, @@ -1143,6 +1191,28 @@ def main() -> int: print(f"--csv-top must be >= 1 (got {args.csv_top})", file=sys.stderr) return 2 + args.streamk_modes = { + "off": (), + "all": _STREAMK_MODES, + }.get(args.streamk, (args.streamk,)) + args.streamk_reductions = tuple( + r.strip() for r in args.streamk_reductions.split(",") if r.strip() + ) + for r in args.streamk_reductions: + if r not in _STREAMK_REDUCTIONS: + print( + f"--streamk-reductions: unknown reduction {r!r} " + f"(expected some of {', '.join(_STREAMK_REDUCTIONS)})", + file=sys.stderr, + ) + return 2 + if args.streamk_ctas == 0 or args.streamk_ctas < -1: + print( + f"--streamk-ctas must be -1 (auto) or > 0 (got {args.streamk_ctas})", + file=sys.stderr, + ) + return 2 + if args.miopen_cmd is None and args.miopen_file is None and args.json_file is None: if args.Di is not None and args.Z is None: print("--Z (filter depth) is required when --Di is set", file=sys.stderr) @@ -1808,6 +1878,87 @@ def _build_wgrad_two_stage_one(args_tuple): return combo, spec, resolved_split_k, s1_kernel, s2_kernel +def _build_wgrad_streamk_one(args_tuple): + """Top-level picklable worker: validate + build IR for one stream-K combo. + + The combo's ``streamk`` / ``streamk_reduction`` pick the kernel; the CTA + pool is left at -1 (it is a kernarg, resolved at launch). Returns + ``(combo, spec, 1, kernel)`` on success, or ``None``. + """ + combo, problem, dtype, arch = args_tuple + combo = WgradCombo(*combo) + + from rocke.core.arch import ArchTarget + from kernels.common.conv_implicit_gemm import ConvDataSpec + from kernels.common.conv_implicit_gemm_wgrad import ( + WgradConvSpec, + build_implicit_gemm_conv_wgrad, + is_valid_wgrad_spec, + ) + + target = ArchTarget.from_gfx(arch) + _mma_family = "wmma" if target.wave_size == 32 else "mma" + atom = target.mma.select_largest_k( + family=_mma_family, + a_dtype=dtype, + b_dtype=dtype, + c_dtype="fp32", + m=combo.warp_tile_mn, + n=combo.warp_tile_mn, + k_max=combo.tile_k, + ) + if atom is None: + return None + spec = WgradConvSpec( + problem=problem, + name="rocke_bench_igemm_wgrad", + data=ConvDataSpec(dtype_a=dtype, dtype_b=dtype, dtype_d=dtype), + tile_m=combo.tile_m, + tile_n=combo.tile_n, + tile_k=combo.tile_k, + warp_m=combo.warp_m, + warp_n=combo.warp_n, + warp_tile_m=combo.warp_tile_mn, + warp_tile_n=combo.warp_tile_mn, + warp_tile_k=atom.k, + wave_size=target.wave_size, + pipeline=combo.pipeline, + epilogue=combo.epilogue, + split_k=1, + # Same per-combo K-outer gate as the other legs, so the timings stay + # comparable. + lds_k_outer=WgradConvSpec.default_lds_k_outer( + arch=arch, + dtype_a=dtype, + dtype_b=dtype, + warp_tile_m=combo.warp_tile_mn, + warp_tile_n=combo.warp_tile_mn, + wave_size=target.wave_size, + ), + streamk=combo.streamk, + streamk_reduction=combo.streamk_reduction, + ) + ok, _ = is_valid_wgrad_spec(spec, arch) + if not ok: + return None + try: + kernel = build_implicit_gemm_conv_wgrad(spec, arch=arch) + except ValueError: + return None + return combo, spec, 1, kernel + + +def _build_and_compile_wgrad_streamk_one(args_tuple): + """Merged worker: :func:`_build_wgrad_streamk_one` + compile.""" + result = _build_wgrad_streamk_one(args_tuple) + if result is None: + return None + combo, spec, split_k, kernel = result + from rocke import compile_kernel as _compile_kernel + + return combo, spec, split_k, _compile_kernel(kernel, arch=args_tuple[3]) + + def _build_dgrad_one(args_tuple): """Top-level picklable worker: validate + build IR for one dgrad combo. @@ -2716,6 +2867,53 @@ def _degrees(pipeline: str, async_dma: bool) -> tuple: else: pending_2s = [] + # Stream-K leg (--streamk): one kernel per geometry x pipeline x launch + # mode x reduction. Stream-K runs the default epilogue on a plain K loop + # and owns the K partition, so there is no cshuffle, async or split-K + # axis here; the pipelines that alias another's code are skipped as in + # the AOT grid. The CTA pool is a kernarg (--streamk-ctas). + _sk_reductions = tuple( + r for r in args.streamk_reductions if r != "atomic" or dtype == "fp32" + ) + pending_sk = [] + if args.streamk_modes and arch not in ("gfx942", "gfx950"): + print(f" --streamk: stream-K wgrad is gfx942/gfx950 only; skipped on {arch}.") + elif args.streamk_modes and _sk_reductions: + sk_combos = [ + WgradCombo( + *_geom, + _pipeline, + "default", + False, + 1, + streamk=_mode, + streamk_reduction=_red, + ) + for _geom in itertools.product( + _TILE_MN, _TILE_MN, _TILE_K, _WARP_MN, _WARP_MN, _WARP_TILE_MN + ) + if _geom[3] * _geom[5] <= _geom[0] and _geom[4] * _geom[5] <= _geom[1] + for _pipeline in _PIPELINES + if _pipeline not in _AOT_ALIAS_PIPELINES["wgrad"] + for _mode in args.streamk_modes + for _red in _sk_reductions + ] + if args.sample is not None: + sk_combos = _sample_combos(sk_combos, args.sample, args.seed + case_idx) + print( + f" stream-K: building {len(sk_combos)} combos " + f"(modes {list(args.streamk_modes)} x reductions " + f"{list(_sk_reductions)}) ...", + flush=True, + ) + pending_sk = _build_ir_parallel( + [(c, problem, dtype, arch) for c in sk_combos], + _build_and_compile_wgrad_streamk_one, + jobs, + ) + n_skipped += len(sk_combos) - len(pending_sk) + n_built += len(pending_sk) + # --------------------------------------------------------------------------- # Phase 3 – GPU run: load modules and time each kernel serially. # --------------------------------------------------------------------------- @@ -3198,6 +3396,133 @@ def _launch_two_stage( rt2.free(X_dev2) rt2.free(dW_dev2) + # --------------------------------------------------------------------------- + # Phase 3c – stream-K GPU run. Each launch zeroes what its reduction + # accumulates into (dW for atomic, the scratch for workspace, the flag + # table for linear/tree) and the workspace reduction runs its Stage 2. + # --------------------------------------------------------------------------- + if pending_sk: + from dataclasses import replace as _dc_replace + + from kernels.common.conv_implicit_gemm_wgrad import wgrad_streamk_signature + + print(f"\nSweeping {len(pending_sk)} stream-K wgrad configs ...", flush=True) + rt3 = Runtime() + dY_dev3 = rt3.alloc(dY_t.nbytes) + X_dev3 = rt3.alloc(X_t.nbytes) + dW_dev3 = rt3.alloc(dW_t.nbytes) + rt3.memcpy_h2d(dY_dev3, _u8(dY_t), dY_t.nbytes) + rt3.memcpy_h2d(X_dev3, _u8(X_t), X_t.nbytes) + _sk_stage2: dict = {} + _sk_stopped = False + try: + for combo, spec, _, artifact in pending_sk: + spec = _dc_replace(spec, streamk_ctas=args.streamk_ctas) + _mode = _streamk_mode_label(spec.streamk, spec.streamk_reduction) + _cfg_str = ( + f"tile={spec.tile_m}x{spec.tile_n}x{spec.tile_k} " + f"warp={spec.warp_m}x{spec.warp_n} " + f"atom={spec.warp_tile_m}x{spec.warp_tile_n}x{spec.warp_tile_k} " + f"{spec.pipeline}/{spec.epilogue:9s}" + ) + try: + launcher = KernelLauncher( + hsaco=artifact.hsaco, + kernel_name=artifact.kernel_name, + signature=wgrad_streamk_signature(spec), + ) + run, free_ws, pool = _streamk_launch( + rt=rt3, + launcher=launcher, + spec=spec, + problem=p, + arch=arch, + ptrs=(int(dY_dev3), int(X_dev3), int(dW_dev3)), + sizes=(dY_t.nbytes, X_t.nbytes, dW_t.nbytes), + stage2=_sk_stage2, + KernelLauncher=KernelLauncher, + LaunchConfig=LaunchConfig, + ) + except (HipError, ValueError) as e: + n_skipped += 1 + print(f"[skip] stream-K {_cfg_str} {_mode}: {e}", flush=True) + continue + try: + kernel_passed = None + if args.verify or args.dump_fail: + _sk_stopped, kernel_passed = _verify_kernel( + rt=rt3, + launch_fn=run, + out_dev=dW_dev3, + out_t=dW_t, + zero_init_out=True, + ref_out=ref_out, + kernel_name=artifact.kernel_name, + dump_fail=args.dump_fail, + extra_tensors={"dY": dY_t, "X": X_t}, + u8=_u8, + arch=arch, + compute_dtype=dtype, + ) + if _sk_stopped: + break + ms = _stop.measure( + run, + warmup=args.warmup, + iters=args.iters, + stream=0, + passed=kernel_passed, + ) + if ms is None: + _stop.report(f"{artifact.kernel_name} {_mode}") + continue + synchronize_and_release(0) + finally: + free_ws() + + cur_tflops = (flop / ms) * 1e-9 + cur_gbps = (bytes_xfer / ms) * 1e-6 + n_run += 1 + _va, _vb, _vc = WgradConvSpec.default_vector_sizes( + p.C, p.K, dtype, split_k=1 + ) + results.append( + Result( + kernel_name=artifact.kernel_name, + tile_m=spec.tile_m, + tile_n=spec.tile_n, + tile_k=spec.tile_k, + warp_m=spec.warp_m, + warp_n=spec.warp_n, + warp_tile_mn=spec.warp_tile_m, + warp_tile_k=spec.warp_tile_k, + pipeline=spec.pipeline, + epilogue=spec.epilogue, + split_k=1, + ms=ms, + tflops=cur_tflops, + gbps=cur_gbps, + vec_a=_va, + vec_b=_vb, + vec_c=_vc, + passed=kernel_passed, + streamk=spec.streamk, + streamk_reduction=spec.streamk_reduction, + streamk_ctas=pool, + ) + ) + print( + f"[{n_run:4d}] {_cfg_str} {_mode:<10s} ctas={pool:<5d} " + f"{cur_tflops:6.1f} TFLOPS {ms:.3f} ms", + flush=True, + ) + finally: + rt3.free(dY_dev3) + rt3.free(X_dev3) + rt3.free(dW_dev3) + if _sk_stopped: + return 1, [] + print(f"\nWgrad sweep done: {n_built} compiled, {n_skipped} skipped.", flush=True) if not results: @@ -3217,6 +3542,8 @@ def _launch_two_stage( print("-" * 92) for rank, r in enumerate(results[:top_n], 1): mode = f"spk{r.split_k}2s" if r.two_stage else f"spk{r.split_k}" + if r.streamk != "off": + mode = _streamk_mode_label(r.streamk, r.streamk_reduction) if r.group_merge > 1: mode += f" gm{r.group_merge}" if r.ws_replicas > 0: @@ -3619,6 +3946,8 @@ def _cache_dispatch(args, arch, target, cases) -> int: directions=directions or args.compile_directions, jobs=max(1, int(args.jobs or 1)), limit=args.limit, + streamk_modes=args.streamk_modes, + streamk_reductions=args.streamk_reductions, ), ) return rc @@ -3708,6 +4037,27 @@ def _u8(t): flush=True, ) continue + if direction == "wgrad": + # Stream-K kernels run only when asked for (--streamk), and + # only the requested modes/reductions, as in the JIT sweep. + n_before = len(candidates) + candidates = [ + c + for c in candidates + if c[0].streamk == "off" + or ( + c[0].streamk in args.streamk_modes + and c[0].streamk_reduction in args.streamk_reductions + ) + ] + if not candidates: + print( + f"Case {case_idx} {problem.short()} {case_dtype} wgrad: " + f"no compatible kernel ({n_before} stream-K kernels " + f"filtered out; see --streamk / --streamk-reductions)", + flush=True, + ) + continue n_compatible = len(candidates) n_vec_pruned = 0 if not args.all_vector_sizes: @@ -3814,124 +4164,156 @@ def _u8(t): degrees = (None,) for split_k in degrees: + free_ws = None try: - extras, grid = _extras_and_grid( - direction, problem, ident, rt, case_dtype, split_k=split_k - ) - except ValueError as e: - print(f" [skip] {ident.label()}: {e}", flush=True) - continue - values = _launch_values_for( - direction, - problem, - ident, - (A_dev, B_dev, D_dev), - (A_t.nbytes, B_t.nbytes, D_t.nbytes), - extras, - ) - cfg = LaunchConfig( - grid=grid, - block=( - meta.get("block_size") - or (ident.warp_m * ident.warp_n * ident.wave_size), - 1, - 1, - ), - stream=0, - ) - run = launcher.bind(values, config=cfg) - if direction == "wgrad" and ident.two_stage: - # Two-stage is three steps and all of them are the - # algorithm's cost: zero the scratch (Stage 1 - # accumulates), Stage 1, then the fold/cast into dW. - from kernels.common.conv_wgrad_workspace_reduce import ( - wgrad_reduce_grid, - ) - - s2_launcher, s2_spec = _stage2_reduce( - problem, - ident.dtype_d, - arch, - stage2_kernels, - KernelLauncher, - _ws_replicas_of(ident), - ) - s2_values = { - "ws_ptr": extras["ws_ptr"], - "dw_ptr": D_dev, - "wg_M": s2_spec.wg_M, - "wg_N": s2_spec.wg_N, - "ws_bytes": extras["ws_bytes"], - "dw_bytes": D_t.nbytes, - "groups": problem.groups, - } - s2_cfg = LaunchConfig( - grid=wgrad_reduce_grid(s2_spec), - block=(s2_spec.block_size, 1, 1), - stream=0, - ) - - def run( - values=values, - cfg=cfg, - extras=extras, - s2_launcher=s2_launcher, - s2_values=s2_values, - s2_cfg=s2_cfg, - ): - rt.memset(extras["ws_ptr"], 0, extras["ws_bytes"]) - launcher(values, config=cfg) - s2_launcher(s2_values, config=s2_cfg) - - label = ident.label() - if split_k is not None: - label += f" @split_k={split_k}" - if ref_out is not None: - # The output is zeroed first: the atomic split-K - # epilogues accumulate into it. A kernel that fails is - # reported and not timed, so it can never rank. - stopped, passed = _verify_kernel( - rt=rt, - launch_fn=run, - out_dev=D_dev, - out_t=D_t, - zero_init_out=True, - ref_out=ref_out, - kernel_name=label, - dump_fail=args.dump_fail, - u8=_u8, - arch=arch, - compute_dtype=case_dtype, - ) - if stopped: - return 1 - if not passed: - n_failed += 1 - overall_rc = 1 + if direction == "wgrad" and ident.streamk != "off": + # Stream-K: the kernel partitions K itself over a CTA pool + # (a kernarg) and has its own ABI, workspace and zeroing + # contract per reduction. + try: + run, free_ws, label = _cached_streamk_run( + args, + problem, + ident, + rt, + arch, + launcher, + (A_dev, B_dev, D_dev), + (A_t.nbytes, B_t.nbytes, D_t.nbytes), + stage2_kernels, + KernelLauncher, + LaunchConfig, + ) + except ValueError as e: + print(f" [skip] {ident.label()}: {e}", flush=True) + continue + else: + try: + extras, grid = _extras_and_grid( + direction, + problem, + ident, + rt, + case_dtype, + split_k=split_k, + ) + except ValueError as e: + print(f" [skip] {ident.label()}: {e}", flush=True) + continue + values = _launch_values_for( + direction, + problem, + ident, + (A_dev, B_dev, D_dev), + (A_t.nbytes, B_t.nbytes, D_t.nbytes), + extras, + ) + cfg = LaunchConfig( + grid=grid, + block=( + meta.get("block_size") + or (ident.warp_m * ident.warp_n * ident.wave_size), + 1, + 1, + ), + stream=0, + ) + run = launcher.bind(values, config=cfg) + if direction == "wgrad" and ident.two_stage: + # Two-stage is three steps and all of them are the + # algorithm's cost: zero the scratch (Stage 1 + # accumulates), Stage 1, then the fold/cast into dW. + from kernels.common.conv_wgrad_workspace_reduce import ( + wgrad_reduce_grid, + ) + + s2_launcher, s2_spec = _stage2_reduce( + problem, + ident.dtype_d, + arch, + stage2_kernels, + KernelLauncher, + _ws_replicas_of(ident), + ) + s2_values = { + "ws_ptr": extras["ws_ptr"], + "dw_ptr": D_dev, + "wg_M": s2_spec.wg_M, + "wg_N": s2_spec.wg_N, + "ws_bytes": extras["ws_bytes"], + "dw_bytes": D_t.nbytes, + "groups": problem.groups, + } + s2_cfg = LaunchConfig( + grid=wgrad_reduce_grid(s2_spec), + block=(s2_spec.block_size, 1, 1), + stream=0, + ) + + def run( + values=values, + cfg=cfg, + extras=extras, + s2_launcher=s2_launcher, + s2_values=s2_values, + s2_cfg=s2_cfg, + ): + rt.memset(extras["ws_ptr"], 0, extras["ws_bytes"]) + launcher(values, config=cfg) + s2_launcher(s2_values, config=s2_cfg) + + label = ident.label() + if split_k is not None: + label += f" @split_k={split_k}" + if ref_out is not None: + # The output is zeroed first: the atomic split-K + # epilogues accumulate into it. A kernel that fails is + # reported and not timed, so it can never rank. + stopped, passed = _verify_kernel( + rt=rt, + launch_fn=run, + out_dev=D_dev, + out_t=D_t, + zero_init_out=True, + ref_out=ref_out, + kernel_name=label, + dump_fail=args.dump_fail, + u8=_u8, + arch=arch, + compute_dtype=case_dtype, + ) + if stopped: + return 1 + if not passed: + n_failed += 1 + overall_rc = 1 + continue + try: + ms = _stop.measure( + run, + warmup=args.warmup, + iters=args.iters, + stream=0, + ) + except (HipError, RuntimeError) as e: + print(f" [skip] {ident.label()}: {e}", flush=True) continue - try: - ms = _stop.measure( - run, - warmup=args.warmup, - iters=args.iters, - stream=0, + if ms is None: + _stop.report(label) + continue + results.append( + (flop / (ms * 1e9), bytes_xfer / (ms * 1e6), ms, label) ) - except (HipError, RuntimeError) as e: - print(f" [skip] {ident.label()}: {e}", flush=True) - continue - if ms is None: - _stop.report(label) - continue - results.append( - (flop / (ms * 1e9), bytes_xfer / (ms * 1e6), ms, label) - ) - # One line per measured launch, as the JIT sweep prints, - # so a long run shows its progress and partial results. - print( - f" [{cand_idx:>{len(str(n_cand))}}/{n_cand}] " - f"{results[-1][0]:7.1f} TFLOPS {ms:8.3f} ms {label}", - flush=True, - ) + # One line per measured launch, as the JIT sweep prints, + # so a long run shows its progress and partial results. + print( + f" [{cand_idx:>{len(str(n_cand))}}/{n_cand}] " + f"{results[-1][0]:7.1f} TFLOPS {ms:8.3f} ms {label}", + flush=True, + ) + finally: + if free_ws is not None: + free_ws() synchronize_and_release() if ref_out is not None: @@ -3991,6 +4373,8 @@ def kind(ident): ident.split_k, ident.two_stage, ident.group_merge, + ident.streamk, + ident.streamk_reduction, ) def rank(ident): @@ -4007,7 +4391,14 @@ def rank(ident): # kernels are dropped. Where nothing stores wider (two-stage wgrad is # default-only, or the problem's output run is odd) they stay. def store_kind(ident): - return (ident.direction, ident.split_k, ident.two_stage, ident.group_merge) + return ( + ident.direction, + ident.split_k, + ident.two_stage, + ident.group_merge, + ident.streamk, + ident.streamk_reduction, + ) widest_store = {} for ident, *_ in candidates: @@ -4059,7 +4450,11 @@ def _signature_for_identity(ident, dtype): if ident.direction == "wgrad": return conv_args_signature( - dtype, direction="wgrad", is_3d=ident.is_3d, two_stage=ident.two_stage + dtype, + direction="wgrad", + is_3d=ident.is_3d, + two_stage=ident.two_stage, + streamk=ident.streamk_reduction if ident.streamk != "off" else None, ) if ident.direction == "dgrad": return conv_args_signature(dtype, direction="dgrad", is_3d=ident.is_3d) @@ -4112,6 +4507,172 @@ def _stage2_reduce(problem, dtype_d, arch, compiled, KernelLauncher, ws_replicas return compiled[key], spec +def _streamk_launch( + *, + rt, + launcher, + spec, + problem, + arch: str, + ptrs, + sizes, + stage2: dict, + KernelLauncher, + LaunchConfig, +): + """Set up one stream-K wgrad launch per the kernel's zeroing contract. + + ``spec`` is the stream-K :class:`WgradConvSpec` (its ``streamk_ctas`` is + the CTA pool, -1 = the arch default); ``problem`` the launch shape; + ``ptrs`` / ``sizes`` the dY, X, dW buffers. Allocates the reduction's + workspace and returns ``(run, free, pool)``: ``run()`` zeroes what the + reduction accumulates into -- dW for ``atomic``, the f32 scratch for + ``workspace``, only the flag table for ``linear``/``tree`` -- launches the + kernel and, for ``workspace``, the Stage 2 reduce into dW; ``free()`` + releases the workspace; ``pool`` is the resolved CTA pool. + """ + from kernels.common.conv_implicit_gemm_wgrad import ( + wgrad_streamk_grid, + wgrad_streamk_launch, + wgrad_streamk_launch_values, + wgrad_streamk_workspace_layout, + wgrad_streamk_workspace_nbytes, + ) + + reduction = spec.streamk_reduction + ws_bytes = wgrad_streamk_workspace_nbytes(spec, arch=arch, problem=problem) + flags_bytes, _ = wgrad_streamk_workspace_layout(spec, arch=arch, problem=problem) + grid = wgrad_streamk_grid(spec, arch=arch, problem=problem) + pool = wgrad_streamk_launch(spec, arch=arch).pool + ws = rt.alloc(ws_bytes) if ws_bytes else None + try: + extra = {} + if reduction == "workspace": + extra = dict(ws_ptr=ws, ws_bytes=ws_bytes) + elif reduction in ("linear", "tree"): + extra = dict(sk_flags=ws, sk_partials=int(ws) + flags_bytes) + values = wgrad_streamk_launch_values( + spec, *ptrs, *sizes, arch=arch, problem=problem, **extra + ) + cfg = LaunchConfig(grid=grid, block=(spec.block_size, 1, 1), stream=0) + s2 = None + if reduction == "workspace": + from kernels.common.conv_wgrad_workspace_reduce import wgrad_reduce_grid + + s2_launcher, s2_spec = _stage2_reduce( + problem, + spec.data.dtype_d, + arch, + stage2, + KernelLauncher, + spec.ws_replicas, + ) + s2 = ( + s2_launcher, + { + "ws_ptr": ws, + "dw_ptr": ptrs[2], + "wg_M": s2_spec.wg_M, + "wg_N": s2_spec.wg_N, + "ws_bytes": ws_bytes, + "dw_bytes": sizes[2], + "groups": problem.groups, + }, + LaunchConfig( + grid=wgrad_reduce_grid(s2_spec), + block=(s2_spec.block_size, 1, 1), + stream=0, + ), + ) + except Exception: + if ws is not None: + rt.free(ws) + raise + + def run(): + if reduction == "atomic": + rt.memset(ptrs[2], 0, sizes[2]) + elif reduction == "workspace": + rt.memset(ws, 0, ws_bytes) + elif flags_bytes: + rt.memset(ws, 0, flags_bytes) + launcher(values, config=cfg) + if s2 is not None: + s2[0](s2[1], config=s2[2]) + + def free(): + if ws is not None: + rt.free(ws) + + return run, free, pool + + +def _streamk_mode_label(mode: str, reduction: str) -> str: + """Short ranked-table label of a stream-K kernel, e.g. ``skdp-tree``.""" + m = {"dp_sk": "dp", "persistent": "pers"}.get(mode, mode) + r = {"linear": "lin", "workspace": "ws", "atomic": "atom"}.get(reduction, reduction) + return f"sk{m}-{r}" + + +def _cached_streamk_run( + args, + problem, + ident, + rt, + arch, + launcher, + ptrs, + sizes, + stage2, + KernelLauncher, + LaunchConfig, +): + """``(run, free, label)`` for a cached stream-K wgrad binary. + + The launch helpers take a spec; this one describes the cached binary on + the launch shape. Nothing is built from it -- the partition, workspace and + kernargs depend only on the tile, the warp layout, the stream-K mode and + reduction, the scratch replicas and the launch shape. + """ + from kernels.common._conv_implicit_gemm_common import ConvDataSpec + from kernels.common.conv_implicit_gemm_wgrad import WgradConvSpec + + spec = WgradConvSpec( + problem=problem, + data=ConvDataSpec( + dtype_a=ident.dtype_a, dtype_b=ident.dtype_b, dtype_d=ident.dtype_d + ), + tile_m=ident.tile_m, + tile_n=ident.tile_n, + tile_k=ident.tile_k, + warp_m=ident.warp_m, + warp_n=ident.warp_n, + warp_tile_m=ident.warp_tile_m, + warp_tile_n=ident.warp_tile_n, + warp_tile_k=ident.warp_tile_k, + wave_size=ident.wave_size, + pipeline=ident.pipeline, + epilogue=ident.epilogue, + streamk=ident.streamk, + streamk_reduction=ident.streamk_reduction, + streamk_ctas=args.streamk_ctas, + ws_replicas=_ws_replicas_of(ident), + ) + run, free, pool = _streamk_launch( + rt=rt, + launcher=launcher, + spec=spec, + problem=problem, + arch=arch, + ptrs=ptrs, + sizes=sizes, + stage2=stage2, + KernelLauncher=KernelLauncher, + LaunchConfig=LaunchConfig, + ) + return run, free, f"{ident.label()} @streamk_ctas={pool}" + + def _split_k_degrees(problem, tile_m: int, tile_n: int, tile_k: int) -> tuple: """Split-K degrees to launch a wgrad split-K kernel with. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/grouped_conv/benchmark_grouped_conv.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/grouped_conv/benchmark_grouped_conv.py index 66f4275670aa..13bd21657f1d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/grouped_conv/benchmark_grouped_conv.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/grouped_conv/benchmark_grouped_conv.py @@ -540,27 +540,6 @@ def main() -> int: choices=["fp16", "bf16", "fp32"], help="data type (default: fp16)", ) - parser.add_argument( - "--streamk", - default="auto", - choices=["auto", "off", "dp_sk", "persistent"], - help=( - "wgrad only: auto (default) lets the dispatcher pick stream-K where it " - "is expected to win; off never uses it; dp_sk (data-parallel + " - "stream-K CTAs) / persistent (a fixed CTA pool round-robining tiles) " - "force it" - ), - ) - parser.add_argument( - "--streamk-reduction", - default="auto", - dest="streamk_reduction", - choices=["auto", "linear", "tree", "workspace"], - help=( - "wgrad stream-K fixup: linear/tree (deterministic flag hand-off) or " - "workspace (f32 scratch atomics + reduce kernel); auto = workspace" - ), - ) parser.add_argument( "--warmup", type=int, default=3, help="warmup iterations (default: 3)" ) @@ -844,10 +823,6 @@ def _u8(t): dtype=dtype, direction=args.direction, ) - if args.direction == "wgrad": - req_base.update( - streamk=args.streamk, streamk_reduction=args.streamk_reduction - ) if problem.is_3d: req_base.update( Di=problem.Di, @@ -1135,120 +1110,6 @@ def _make(*shape): return 0 -def _streamk_launcher( - *, - rt, - spec, - kernel, - arch: str, - grid, - block, - stream, - ws_dev, - ws_nbytes: int, - buffers, - nbytes, - compile_kernel, - KernelLauncher, - LaunchConfig, -): - """``(launcher, kernel_name, launch(fence))`` for a stream-K wgrad spec. - - Each launch resets exactly what the reduction's contract requires before - the kernel runs: the flag table for linear/tree (dW is stored once by the - tile owner), the f32 scratch for the workspace reduction, which is then - cast into dW by the Stage 2 reduce kernel on the same stream. - """ - from kernels.common.conv_implicit_gemm_wgrad import ( - wgrad_streamk_launch_values, - wgrad_streamk_signature, - wgrad_streamk_workspace_layout, - ) - from kernels.common.conv_wgrad_workspace_reduce import ( - WgradReduceSpec, - build_conv_wgrad_workspace_reduce, - wgrad_reduce_grid, - wgrad_reduce_signature, - ) - - dY_dev, X_dev, dW_dev = buffers - dY_nb, X_nb, dW_nb = nbytes - artifact = compile_kernel(kernel, arch=arch) - launcher = KernelLauncher( - hsaco=artifact.hsaco, - kernel_name=artifact.kernel_name, - signature=wgrad_streamk_signature(spec), - ) - reducer = None - extra = {} - if spec.streamk_reduction == "workspace": - extra = dict(ws_ptr=ws_dev, ws_bytes=ws_nbytes) - reset_bytes = ws_nbytes - else: - flags_bytes, _ = wgrad_streamk_workspace_layout(spec, arch=arch) - extra = dict(sk_flags=ws_dev, sk_partials=ws_dev + flags_bytes) - reset_bytes = flags_bytes - values = wgrad_streamk_launch_values( - spec, - dY_dev, - X_dev, - dW_dev, - dY_nb, - X_nb, - dW_nb, - arch=arch, - **extra, - ) - if spec.streamk_reduction == "workspace": - p = spec.problem - rspec = WgradReduceSpec( - problem=p, - dtype_d=spec.data.dtype_d, - groups=p.groups, - ws_replicas=spec.ws_replicas, - ) - rart = compile_kernel( - build_conv_wgrad_workspace_reduce(rspec, arch=arch), arch=arch - ) - reducer = KernelLauncher( - hsaco=rart.hsaco, - kernel_name=rart.kernel_name, - signature=wgrad_reduce_signature(rspec), - ) - rvalues = { - "ws_ptr": ws_dev, - "dw_ptr": dW_dev, - "wg_M": spec.wg_M, - "wg_N": spec.wg_N, - "ws_bytes": ws_nbytes, - "dw_bytes": dW_nb, - "groups": p.groups, - } - - def _launch(fence: bool): - if reset_bytes: - rt.memset(ws_dev, 0, reset_bytes) - launcher( - values, - config=LaunchConfig( - grid=grid, block=block, stream=stream, fence=fence and reducer is None - ), - ) - if reducer is not None: - reducer( - rvalues, - config=LaunchConfig( - grid=wgrad_reduce_grid(rspec), - block=(rspec.block_size, 1, 1), - stream=stream, - fence=fence, - ), - ) - - name = artifact.kernel_name + ("+cast" if reducer is not None else "") - return launcher, name, _launch - - def _run_wgrad( *, args, @@ -1277,9 +1138,6 @@ def _run_wgrad( _DEFAULT_TILE_M as _WGRAD_REDUCE_TILE_M, _DEFAULT_TILE_N as _WGRAD_REDUCE_TILE_N, ) - from kernels.common.conv_implicit_gemm_wgrad import ( - wgrad_streamk_workspace_nbytes, - ) _u8 = u8 p = problem @@ -1368,9 +1226,6 @@ def _make(*shape): launcher, ws_nbytes = build_implicit_gemm_conv_wgrad_two_stage( instance_spec, arch=arch ) - elif instance_spec.streamk != "off": - kernel = build_implicit_gemm_conv_wgrad(instance_spec, arch=arch) - ws_nbytes = wgrad_streamk_workspace_nbytes(instance_spec, arch=arch) else: kernel = build_implicit_gemm_conv_wgrad(instance_spec, arch=arch) ws_nbytes = 0 @@ -1389,27 +1244,7 @@ def _make(*shape): # factor from z entirely, so every grouped launch covered group 0 only. grid = _wgrad_grid(dspec, req) - if instance_spec.streamk != "off": - if ws_dev is not None: - rt.free(ws_dev) - ws_dev = rt.alloc(max(ws_nbytes, 4)) - launcher, kernel_name, _launch = _streamk_launcher( - rt=rt, - spec=instance_spec, - kernel=kernel, - arch=arch, - grid=grid, - block=block, - stream=stream, - ws_dev=ws_dev, - ws_nbytes=ws_nbytes, - buffers=(dY_dev, X_dev, dW_dev), - nbytes=(dY_t.nbytes, X_t.nbytes, dW_t.nbytes), - compile_kernel=compile_kernel, - KernelLauncher=KernelLauncher, - LaunchConfig=LaunchConfig, - ) - elif instance_spec.two_stage: + if instance_spec.two_stage: if ws_dev is not None: rt.free(ws_dev) ws_dev = rt.alloc(ws_nbytes) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_cache.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_cache.py index ad6ecaf7cca9..f21025f7ccdd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_cache.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_cache.py @@ -200,6 +200,12 @@ class KernelIdentity: # binary's identity rather than a default the two could drift apart on. # 0 = no scratch (not two-stage). ws_replicas: int = 0 + # wgrad stream-K (WgradConvSpec.streamk / streamk_reduction): the launch + # mode ("dp_sk" or "persistent") and the partial-sum reduction are baked + # into the ISA; the CTA pool is a kernarg, so it is not here. "off" / "" + # for every other kernel. + streamk: str = "off" + streamk_reduction: str = "" group_merge: int = 1 num_load_waves: int = 0 cshuffle_no_alias: bool = False @@ -271,6 +277,9 @@ def stable_hash(self) -> str: # field keep their hash (and their cache entries). if not fields["ws_replicas"]: del fields["ws_replicas"] + # Same for the stream-K pair at "no stream-K". + if fields["streamk"] == "off" and not fields["streamk_reduction"]: + del fields["streamk"], fields["streamk_reduction"] payload = json.dumps(fields, sort_keys=True, separators=(",", ":")) return hashlib.sha1(payload.encode()).hexdigest() diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_label.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_label.py index b1fb7cfe491e..d4195736f94d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_label.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_label.py @@ -68,6 +68,8 @@ split_k=1, two_stage=False, ws_replicas=0, + streamk="off", + streamk_reduction="", group_merge=1, num_load_waves=0, cshuffle_no_alias=False, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py index 1f0c6a6ae5e1..fc1f17582874 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py @@ -51,6 +51,11 @@ ``ws_replicas`` = 8). Plus group-merged depthwise kernels (``group_merge`` in 2..64, two-stage, load widths derived by the builder), served only to depthwise problems whose merged GEMM fits the tile. + Opt-in (``streamk_modes``; ``--compile-all --streamk``): stream-K + kernels, default epilogue only (c = 1, b in {1, 2, 4, 8}), ``split_k`` = 1, + one per launch mode (``dp_sk``, ``persistent``) x reduction (``linear``, + ``tree``, ``workspace`` with ``ws_replicas`` = 8, and ``atomic`` where dW + is fp32). The CTA pool is a kernarg, so it is not swept here. * **dgrad** -- vector widths: a in {1, 2, 4, 8} (dY along kpg), b = c in {1, 2, 4, 8} (W and dX along cpg); with the default epilogue c = 1 and b in {1, 2, 4, 8}. stride and dilation are baked into the @@ -175,6 +180,22 @@ # separately if you need it; the wgrad benchmark's JIT sweep still covers it. CACHE_WGRAD_SPLIT_KS: Tuple[int, ...] = (2,) +# ---- wgrad stream-K axis (opt-in) ------------------------------------------- +# Stream-K (WgradConvSpec.streamk) bakes its launch mode and its partial-sum +# reduction into the kernel; the CTA pool -- and with it the whole iteration +# partition -- is a kernarg, so one binary per (mode, reduction) serves every +# shape and pool. It is off unless a caller asks for it (--streamk): it adds +# a default-epilogue kernel per (mode, reduction) to every wgrad geometry, +# which more than doubles the wgrad grid. "atomic" accumulates into dW with +# f32 atomics and the validator keeps it to fp32 dW only. +CACHE_WGRAD_STREAMK_MODES: Tuple[str, ...] = ("dp_sk", "persistent") +CACHE_WGRAD_STREAMK_REDUCTIONS: Tuple[str, ...] = ( + "linear", + "tree", + "workspace", + "atomic", +) + # Private aliases for internal use -- the generators below refer to these. _TILE_MN = CACHE_TILE_MN _TILE_K = CACHE_TILE_K @@ -414,6 +435,21 @@ def _async_chunks(cfg: dict, dtypes, caps: dict) -> dict: ) +def _wgrad_streamk_variants( + epilogue: str, + streamk_modes: Sequence[str], + streamk_reductions: Sequence[str], +) -> List[Tuple[str, str]]: + """``(mode, reduction)`` stream-K variants built per wgrad geometry. + + Stream-K runs its epilogue once per tile inside the stream-K tile loop and + only the default epilogue supports that, so cshuffle geometries get none. + """ + if epilogue != "default": + return [] + return list(itertools.product(streamk_modes, streamk_reductions)) + + def _wgrad_jobs( arch, dtype, @@ -424,6 +460,8 @@ def _wgrad_jobs( *, geometries: Optional[Sequence[tuple]] = None, llvm_flavor: Optional[str] = None, + streamk_modes: Sequence[str] = (), + streamk_reductions: Sequence[str] = CACHE_WGRAD_STREAMK_REDUCTIONS, ) -> Iterator[BuildJob]: from kernels.common.conv_implicit_gemm_wgrad import _DEFAULT_WS_REPLICAS @@ -492,6 +530,52 @@ def _wgrad_jobs( spec_kwargs=cfg, caps=caps, ) + # Stream-K kernels (opt-in). Default epilogue only, so the store is + # scalar and B takes its own width; split_k stays 1 (stream-K owns the + # K partition). The workspace reduction reuses the two-stage scratch, + # so its replica count is pinned exactly as for two-stage. + for (mode, reduction), vec_a, vec_b, caps in itertools.product( + _wgrad_streamk_variants(epilogue, streamk_modes, streamk_reductions), + _VECS, + _VECS, + _layout_caps(), + ): + cfg = dict( + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + warp_m=warp_m, + warp_n=warp_n, + warp_tile_m=wt, + warp_tile_n=wt, + warp_tile_k=atom.k, + pipeline=pipeline, + epilogue=epilogue, + wave_size=wave_size, + vector_size_a=vec_a, + vector_size_b=vec_b, + vector_size_c=1, + streamk=mode, + streamk_reduction=reduction, + ) + if reduction == "workspace": + cfg["ws_replicas"] = _DEFAULT_WS_REPLICAS + yield BuildJob( + identity=KernelIdentity( + arch=arch, + direction="wgrad", + algorithm="implicit_gemm", + llvm_flavor=flavor, + dtype_a=da, + dtype_b=db, + dtype_d=dd, + grouped=caps["grouped"], + **cfg, + ), + direction="wgrad", + spec_kwargs=cfg, + caps=caps, + ) # Group-merged depthwise kernels. Merged split-K must take the # two-stage path (the packed atomic cannot drop off-diagonal pairs). # The vector widths are left to the builder: it derives them from the @@ -673,6 +757,14 @@ def job_for_identity(identity: KernelIdentity) -> BuildJob: atom, ) want = identity.stable_hash() + streamk = ( + dict( + streamk_modes=(identity.streamk,), + streamk_reductions=(identity.streamk_reduction,), + ) + if identity.direction == "wgrad" and identity.streamk != "off" + else {} + ) for job in _direction_jobs( identity.direction, identity.arch, @@ -682,6 +774,7 @@ def job_for_identity(identity: KernelIdentity) -> BuildJob: identity.max_sub_gemms, geometries=[geometry], llvm_flavor=identity.llvm_flavor, + **streamk, ): if job.identity.stable_hash() == want: return job @@ -741,6 +834,8 @@ def count_jobs( target, directions: Sequence[str], split_ks: Sequence[int] = CACHE_WGRAD_SPLIT_KS, + streamk_modes: Sequence[str] = (), + streamk_reductions: Sequence[str] = CACHE_WGRAD_STREAMK_REDUCTIONS, ) -> int: """Raw candidate count :func:`enumerate_jobs` walks, before dedup/validation. @@ -762,6 +857,14 @@ def count_jobs( if "wgrad" in directions and _aot_pipeline("wgrad", pipeline): total += vecs * sum(len(_wgrad_two_stages(sk)) for sk in split_ks) * layouts total += merged + total += ( + len(_VECS) + * len(_VECS) + * layouts + * len( + _wgrad_streamk_variants(geo[7], streamk_modes, streamk_reductions) + ) + ) if "dgrad" in directions and _aot_pipeline("dgrad", pipeline): total += vecs * dgrad_caps return total @@ -774,6 +877,9 @@ def _direction_jobs( target, split_ks: Sequence[int], max_sub_gemms: int, + *, + streamk_modes: Sequence[str] = (), + streamk_reductions: Sequence[str] = CACHE_WGRAD_STREAMK_REDUCTIONS, **kw, ) -> Iterator[BuildJob]: wave_size = target.wave_size @@ -781,7 +887,17 @@ def _direction_jobs( if direction == "fwd": return _fwd_jobs(arch, dtype, wave_size, mma_family, target, **kw) if direction == "wgrad": - return _wgrad_jobs(arch, dtype, wave_size, mma_family, target, split_ks, **kw) + return _wgrad_jobs( + arch, + dtype, + wave_size, + mma_family, + target, + split_ks, + streamk_modes=streamk_modes, + streamk_reductions=streamk_reductions, + **kw, + ) if direction == "dgrad": return _dgrad_jobs( arch, dtype, wave_size, mma_family, target, max_sub_gemms, **kw @@ -805,6 +921,7 @@ def _enumerate_chunk( validate: bool, llvm_flavor: str, geometries: Sequence[tuple], + streamk: Tuple[Sequence[str], Sequence[str]] = ((), ()), ) -> Tuple[int, List[Tuple[str, BuildJob]]]: """Walk one direction over ``geometries``. @@ -825,6 +942,8 @@ def _enumerate_chunk( max_sub_gemms, geometries=geometries, llvm_flavor=llvm_flavor, + streamk_modes=streamk[0], + streamk_reductions=streamk[1], ): n_raw += 1 key = job.identity.stable_hash() @@ -845,9 +964,18 @@ def _enumerate_chunk( def _enumerate_worker(payload): - (direction, arch, dtype, split_ks, max_sub_gemms, validate, flavor, lo, hi) = ( - payload - ) + ( + direction, + arch, + dtype, + split_ks, + max_sub_gemms, + validate, + flavor, + streamk, + lo, + hi, + ) = payload from rocke.core.arch import ArchTarget target = ArchTarget.from_gfx(arch) @@ -864,6 +992,7 @@ def _enumerate_worker(payload): validate, flavor, geos[lo:hi], + streamk, ) @@ -879,9 +1008,14 @@ def enumerate_jobs( jobs: int = 1, log=None, log_every_s: float = 5.0, + streamk_modes: Sequence[str] = (), + streamk_reductions: Sequence[str] = CACHE_WGRAD_STREAMK_REDUCTIONS, ) -> List[BuildJob]: """All buildable jobs for the requested directions, deduped by identity. + ``streamk_modes`` (empty by default) adds the opt-in wgrad stream-K + kernels for those launch modes x ``streamk_reductions``. + With ``validate=True`` (the default) each candidate is run through its direction's spec validator first, so the returned list is what will actually compile rather than the raw cross product. @@ -908,8 +1042,18 @@ def enumerate_jobs( (len(geos) * i // n_chunks, len(geos) * (i + 1) // n_chunks) for i in range(n_chunks) ] + streamk = (tuple(streamk_modes), tuple(streamk_reductions)) payloads = [ - (direction, arch, dtype, tuple(split_ks), max_sub_gemms, validate, flavor) + ( + direction, + arch, + dtype, + tuple(split_ks), + max_sub_gemms, + validate, + flavor, + streamk, + ) + bound for direction in directions for bound in bounds @@ -918,7 +1062,12 @@ def enumerate_jobs( total = 0 if log is not None: total = count_jobs( - dtype=dtype, target=target, directions=directions, split_ks=split_ks + dtype=dtype, + target=target, + directions=directions, + split_ks=split_ks, + streamk_modes=streamk_modes, + streamk_reductions=streamk_reductions, ) log(f" {total} candidates to check with {jobs} process(es)") @@ -958,6 +1107,7 @@ def _done(idx: int, chunk_raw: int, chunk_out) -> None: validate, flavor, geos[lo:hi], + streamk, ), ) else: @@ -1241,9 +1391,14 @@ def compile_all( jobs: int = 1, limit: Optional[int] = None, log=print, + streamk_modes: Sequence[str] = (), + streamk_reductions: Sequence[str] = CACHE_WGRAD_STREAMK_REDUCTIONS, ) -> int: """Populate ``cache`` with every implicit-GEMM variant for ``arch``/``dtype``. + ``streamk_modes`` opts the wgrad stream-K kernels in (see + :data:`CACHE_WGRAD_STREAMK_MODES`). + Compilation is embarrassingly parallel and dominated by LLVM, so it fans out over processes; ``jobs`` defaults to one but the caller normally passes ``os.cpu_count()``. @@ -1262,6 +1417,8 @@ def compile_all( directions=directions, jobs=jobs, log=log, + streamk_modes=streamk_modes, + streamk_reductions=streamk_reductions, ), build=build_kernel, arch=arch, @@ -1289,6 +1446,10 @@ def compile_all( lambda i: ("two-stage" if i.two_stage else "atomic") if i.split_k > 1 else "no", ), ("group_merge", lambda i: i.group_merge), + ( + "stream-K", + lambda i: f"{i.streamk}/{i.streamk_reduction}" if i.streamk != "off" else "no", + ), ("stride x dilation", lambda i: f"{i.stride_h}x{i.dilation_h}"), ("filter", lambda i: f"{i.filter_h}x{i.filter_w}"), ("pad", lambda i: i.pad_h), @@ -1594,6 +1755,11 @@ def _launch_values_for(direction, problem, identity, ptrs, sizes, extras): split_k=extras.get("split_k", 1), ws_ptr=extras.get("ws_ptr"), ws_bytes=extras.get("ws_bytes"), + # Stream-K: the StreamKLaunch (pool, mode, reduction) that fills + # the partition block, plus the linear/tree flag/partial buffers. + streamk=extras.get("streamk"), + sk_flags=extras.get("sk_flags"), + sk_partials=extras.get("sk_partials"), ) return ConvArgs.from_problem( problem, direction="dgrad", tile_m=tm, tile_n=tn diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py index 401af0882e25..3bdb2d80c6eb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/grouped_convolution.py @@ -156,9 +156,6 @@ _DEFAULT_WS_REPLICAS as _WGRAD_WS_REPLICAS, is_valid_wgrad_spec as _wgrad_is_valid_spec, wgrad_atomic_epilogue_available as _wgrad_atomic_epilogue_available, - wgrad_streamk_default_pool as _wgrad_streamk_default_pool, - wgrad_streamk_grid as _wgrad_streamk_grid, - wgrad_streamk_launch_values as _wgrad_streamk_launch_values, ) from kernels.common.conv_implicit_gemm_dgrad import ( DgradConvSpec, @@ -278,13 +275,6 @@ class ConvGroupedRequest(OperatorRequest): dilation_d: Optional[int] = None # optional vec_size_c override; None = let the candidate decide vec_size_c: Optional[int] = None - # wgrad only: "auto" (default) picks stream-K for the shapes it is measured - # to win on (_wgrad_streamk_auto) and the grid-per-tile kernel with auto - # split-K otherwise; "off" never uses stream-K; "dp_sk" / "persistent" - # force it. streamk_reduction picks how the CTAs sharing an output tile - # combine their partial sums; "auto" is the f32 workspace reduction. - streamk: str = "auto" - streamk_reduction: str = "auto" op: str = "conv_grouped" algorithm: str = "auto" spec_id: str = "auto" @@ -379,19 +369,6 @@ def _request_errors(req: OperatorRequest) -> list[str]: errors.append(f"unsupported dtype {req.dtype!r}; fp16 or bf16 only") if req.layout.upper() != "NHWC": errors.append(f"unsupported layout {req.layout!r}; NHWC only") - if req.streamk not in ("off", "auto", "dp_sk", "persistent"): - errors.append( - "streamk must be 'off', 'auto', 'dp_sk' or 'persistent', got " - f"{req.streamk!r}" - ) - elif req.streamk in ("dp_sk", "persistent") and req.direction != "wgrad": - errors.append(f"streamk is wgrad-only, got direction={req.direction!r}") - if req.streamk_reduction not in ("auto", "linear", "tree", "workspace"): - # 'atomic' needs an fp32 dW, which a fp16/bf16 request never has. - errors.append( - "streamk_reduction must be 'auto', 'linear', 'tree' or 'workspace', " - f"got {req.streamk_reduction!r}" - ) try: ArchTarget.from_gfx(req.arch) except KeyError as e: @@ -459,9 +436,7 @@ def _epilogue_for(req: ConvGroupedRequest) -> str: return "cshuffle" if _vec_size_c(req) > 1 else "default" -def _wgrad_grouped_overrides( - req: ConvGroupedRequest, streamk: str = "off" -) -> Tuple[str, int]: +def _wgrad_grouped_overrides(req: ConvGroupedRequest) -> Tuple[str, int]: """(epilogue, split_k) for a wgrad spec. Grouping is orthogonal to the epilogue and split-K: every epilogue (direct, @@ -471,102 +446,10 @@ def _wgrad_grouped_overrides( same vec-derived epilogue and auto split-K formula (split_k=-1) as ungrouped. (WMMA grouped is handled by its own candidate, which forces default@split_k=1 -- WMMA has neither a cshuffle nor a split-K path.) - - ``streamk`` is the resolved mode from :func:`_wgrad_streamk_choice`. """ - if streamk != "off": - # Stream-K owns the K range (no split-K) and stores each tile through - # the direct epilogue once per tile inside its tile loop. - return "default", 1 return _epilogue_for(req), -1 -# Auto-selection thresholds for stream-K wgrad, in units of the candidate's -# own tiles. Stream-K pays off when a handful of output tiles sit on a long -# reduction: split-K can only cut that reduction into equal slices per tile, -# while stream-K balances iterations across the whole CTA pool. With many -# tiles, or a short reduction, the grid-per-tile kernel is already balanced and -# stream-K's per-tile fixup is pure overhead. Fitted on a benchmark corpus of -# production wgrad shapes on gfx950 and kept deliberately coarse (two -# thresholds) so it does not encode that corpus. -_WGRAD_STREAMK_MAX_TILES = 128 -_WGRAD_STREAMK_MIN_ITERS_PER_TILE = 1024 - - -def _wgrad_streamk_auto( - req: ConvGroupedRequest, - tile_m: int, - tile_n: int, - tile_k: int, - reduction: str = "workspace", -) -> bool: - """Whether ``streamk="auto"`` should pick stream-K for this request. - - Besides the two thresholds, the partition the launch will use must give - the stream-K CTAs work: a pool larger than ``tiles * iters_per_tile`` - (one tile on a bigger arch's pool) falls back to all-DP, which runs each - tile's whole reduction on one CTA, slower than the split-K path it would - replace. - """ - from rocke.helpers.streamk import streamk_iter_partition - - p = _problem(req) - spatial = (p.Z if p.is_3d else 1) * p.Y * p.X - wg_M = p.K // p.groups - wg_N = spatial * (p.C // p.groups) - wg_K = p.N * p.Ho * p.Wo * (p.Do if p.is_3d else 1) - tiles = -(-wg_M // tile_m) * -(-wg_N // tile_n) * p.groups - iters_per_tile = -(-wg_K // tile_k) - if ( - tiles > _WGRAD_STREAMK_MAX_TILES - or iters_per_tile < _WGRAD_STREAMK_MIN_ITERS_PER_TILE - ): - return False - part = streamk_iter_partition( - m_tiles=tiles, - n_tiles=1, - iters_per_tile=iters_per_tile, - max_active_wgs=_wgrad_streamk_default_pool(reduction, req.arch), - ) - return part.sk_ctas > 0 - - -def _wgrad_streamk_choice( - req: ConvGroupedRequest, tile_m: int, tile_n: int, tile_k: int -) -> Tuple[str, str]: - """Resolve the request's ``(streamk, streamk_reduction)`` to concrete values. - - The one resolver every wgrad candidate's instance spec, dispatch spec and - grid go through, so they cannot disagree on whether the kernel is stream-K. - """ - # The f32 workspace reduction never waits on another CTA, so its pool can be - # sized for occupancy; it is also the fastest of the reductions measured. - reduction = ( - "workspace" if req.streamk_reduction == "auto" else req.streamk_reduction - ) - mode = req.streamk - if mode == "auto": - mfma = req.arch in ("gfx942", "gfx950") - mode = ( - "dp_sk" - if mfma and _wgrad_streamk_auto(req, tile_m, tile_n, tile_k, reduction) - else "off" - ) - return mode, reduction - - -def _wgrad_streamk_fields(mode: str, reduction: str) -> dict: - """The WgradConvSpec stream-K fields for a resolved choice. - - ``streamk_ctas`` stays -1 (sized by the instance from the arch and the - reduction) so the spec, the kernel and the launch grid all derive the pool - from the same place. - """ - if mode == "off": - return {} - return {"streamk": mode, "streamk_reduction": reduction, "streamk_ctas": -1} - - def _data_spec(req: ConvGroupedRequest) -> ConvDataSpec: return ConvDataSpec( dtype_a=req.dtype.lower(), @@ -610,8 +493,6 @@ class ConvGroupedSpec: arch: str split_k: int = 1 # wgrad only lds_k_outer: bool = False # wgrad only - streamk: str = "off" # wgrad only - streamk_reduction: str = "linear" # wgrad only name: str = "rocke_conv_grouped" def kernel_name(self) -> str: @@ -635,9 +516,6 @@ def kernel_name(self) -> str: # fetch rather than a transpose-on-store. if self.direction in ("wgrad", "dgrad") and self.lds_k_outer: parts.append("kouter") - # streamk: a different grid, K ranges and fixup epilogue. - if self.direction == "wgrad" and self.streamk != "off": - parts.append(f"sk_{self.streamk}_{self.streamk_reduction}") return kernel_name_join(self.name, *parts) def to_fwd_spec(self, problem: "ConvProblem") -> "ImplicitGemmConvSpec": @@ -675,18 +553,7 @@ def to_wgrad_spec(self, problem: "ConvProblem") -> "WgradConvSpec": """ assert self.direction == "wgrad", "to_wgrad_spec is only valid for wgrad specs" target = ArchTarget.from_gfx(self.arch) - if self.streamk != "off": - # Stream-K owns the K range: no split-K, no two-stage, and the - # direct epilogue runs per tile. - resolved_split_k, two_stage = 1, False - streamk = { - "streamk": self.streamk, - "streamk_reduction": self.streamk_reduction, - "streamk_ctas": -1, - } - else: - resolved_split_k, two_stage, _ = _resolve_wgrad_split_k(self, problem) - streamk = {} + resolved_split_k, two_stage, _ = _resolve_wgrad_split_k(self, problem) return WgradConvSpec( problem=problem, name=self.name, @@ -708,8 +575,7 @@ def to_wgrad_spec(self, problem: "ConvProblem") -> "WgradConvSpec": pipeline=self.pipeline, # two_stage writes one f32 per element via workspace store; cshuffle # is not used and would produce an invalid spec (validator rejects it). - # Stream-K likewise stores through the direct epilogue. - epilogue="default" if (two_stage or streamk) else self.epilogue, + epilogue="default" if two_stage else self.epilogue, split_k=resolved_split_k, two_stage=two_stage, # Pin the replica count explicitly rather than inheriting the @@ -717,7 +583,6 @@ def to_wgrad_spec(self, problem: "ConvProblem") -> "WgradConvSpec": # the i32 ws_bytes ABI using this same constant, and a cap computed # over a different R than the spec carries bounds nothing. ws_replicas=_WGRAD_WS_REPLICAS, - **streamk, ) def to_dgrad_spec(self, problem: "ConvProblem") -> "DgradConvSpec": @@ -857,10 +722,6 @@ def _resolve_wgrad_split_k( def _wgrad_grid(spec: ConvGroupedSpec, req: OperatorRequest) -> Tuple[int, int, int]: assert isinstance(req, ConvGroupedRequest) p = _problem(req) - if spec.streamk != "off": - # One-dimensional stream-K launch, sized from the same spec the kernel - # is built from so the two cannot disagree on the CTA pool. - return _wgrad_streamk_grid(spec.to_wgrad_spec(p), arch=spec.arch, problem=p) spatial = (p.Z if p.is_3d else 1) * p.Y * p.X # Per-group GEMM dims. For the ungrouped path G==1 these reduce to wg_M=K, # wg_N=spatial*C. @@ -1248,8 +1109,7 @@ def _tile(req: ConvGroupedRequest): def _build_instance_spec(req: ConvGroupedRequest) -> WgradConvSpec: tm, tn, tk, wm, wn, wtmn, wtk = _tile(req) - _skm, _skr = _wgrad_streamk_choice(req, tm, tn, tk) - _ep, _sk = _wgrad_grouped_overrides(req, _skm) + _ep, _sk = _wgrad_grouped_overrides(req) return WgradConvSpec( problem=_problem(req), name=name, @@ -1266,7 +1126,6 @@ def _build_instance_spec(req: ConvGroupedRequest) -> WgradConvSpec: pipeline=_PIPELINE, epilogue=_ep, split_k=_sk, - **_wgrad_streamk_fields(_skm, _skr), ) def support(req: OperatorRequest) -> Tuple[bool, str]: @@ -1292,8 +1151,7 @@ def select(req: OperatorRequest) -> ConvGroupedSpec: raise ValueError(f"{name} does not support request: {why}") assert isinstance(req, ConvGroupedRequest) tm, tn, tk, wm, wn, wtmn, wtk = _tile(req) - _skm, _skr = _wgrad_streamk_choice(req, tm, tn, tk) - _ep, _sk = _wgrad_grouped_overrides(req, _skm) + _ep, _sk = _wgrad_grouped_overrides(req) return ConvGroupedSpec( direction="wgrad", tile_m=tm, @@ -1308,8 +1166,6 @@ def select(req: OperatorRequest) -> ConvGroupedSpec: dtype=req.dtype.lower(), arch=req.arch, split_k=_sk, - streamk=_skm, - streamk_reduction=_skr, name=name, ) @@ -1444,8 +1300,7 @@ def _pipeline(req: ConvGroupedRequest) -> str: def _build_instance_spec(req: ConvGroupedRequest) -> WgradConvSpec: tm, tn, tk, wm, wn, wtmn, wtk = _tile(req) - _skm, _skr = _wgrad_streamk_choice(req, tm, tn, tk) - _ep, _sk = _wgrad_grouped_overrides(req, _skm) + _ep, _sk = _wgrad_grouped_overrides(req) return WgradConvSpec( problem=_problem(req), name=name, @@ -1463,7 +1318,6 @@ def _build_instance_spec(req: ConvGroupedRequest) -> WgradConvSpec: pipeline=_pipeline(req), epilogue=_ep, split_k=_sk, - **_wgrad_streamk_fields(_skm, _skr), ) def support(req: OperatorRequest) -> Tuple[bool, str]: @@ -1489,8 +1343,7 @@ def select(req: OperatorRequest) -> ConvGroupedSpec: raise ValueError(f"{name} does not support request: {why}") assert isinstance(req, ConvGroupedRequest) tm, tn, tk, wm, wn, wtmn, wtk = _tile(req) - _skm, _skr = _wgrad_streamk_choice(req, tm, tn, tk) - _ep, _sk = _wgrad_grouped_overrides(req, _skm) + _ep, _sk = _wgrad_grouped_overrides(req) return ConvGroupedSpec( direction="wgrad", tile_m=tm, @@ -1506,8 +1359,6 @@ def select(req: OperatorRequest) -> ConvGroupedSpec: dtype=req.dtype.lower(), arch=req.arch, split_k=_sk, - streamk=_skm, - streamk_reduction=_skr, name=name, ) @@ -1680,9 +1531,7 @@ def _make_gfx1250_wgrad_candidate() -> KernelCandidate: test_gfx1250_grouped_wgrad_dual_engine). Group merging (``group_merge > 1``) is MFMA-only and additionally needs split_k > 1 on the two-stage path, so it is unreachable here on both counts: is_valid_wgrad_spec rejects it for - wave32 and for the split_k=1 this candidate forces. Stream-K is MFMA-only - too, so a stream-K request is turned away rather than silently served by - the grid-per-tile kernel. + wave32 and for the split_k=1 this candidate forces. """ name = "implicit_gemm_conv_wgrad_gfx1250" spec_id = "igemm_conv_wgrad_gfx1250_32x32" @@ -1729,8 +1578,6 @@ def support(req: OperatorRequest) -> Tuple[bool, str]: return False, f"gfx1250 candidate requires arch=gfx1250 (got {req.arch!r})" if req.direction != "wgrad": return False, f"candidate handles 'wgrad', got direction={req.direction!r}" - if req.streamk in ("dp_sk", "persistent"): - return False, "stream-K wgrad is MFMA-only (gfx942, gfx950)" ok, why = selector_matches(req, candidate) if not ok: return False, why @@ -1912,23 +1759,13 @@ def _signature_for(req: ConvGroupedRequest, spec: ConvGroupedSpec) -> tuple: The trailing args are not cosmetic: dgrad always carries the tilde record buffer, wgrad always carries ``ks``/``ks_count``, and a two-stage wgrad - also carries the workspace pair; a stream-K wgrad carries its reduction's - buffers and the partition block. Kernargs pack positionally, so omitting + also carries the workspace pair. Kernargs pack positionally, so omitting one shifts every argument after it. """ from kernels.common.conv_abi import conv_args_signature p = _problem(req) dtype = req.dtype - if req.direction == "wgrad" and spec.streamk != "off": - return tuple( - conv_args_signature( - dtype, - direction="wgrad", - is_3d=p.is_3d, - streamk=spec.streamk_reduction, - ) - ) if req.direction == "wgrad": # Two-stage is decided by the resolver, not carried on the dispatcher # spec: it is what to_wgrad_spec builds the kernel from, so asking @@ -1961,8 +1798,6 @@ def launch_values_for( ws_bytes: Optional[int] = None, sub_gemm_buf: Optional[int] = None, num_sub_gemms: Optional[int] = None, - sk_flags: Optional[int] = None, - sk_partials: Optional[int] = None, ) -> dict: """Compute the full AOT ``values`` dict for a dispatch request. @@ -1973,34 +1808,11 @@ def launch_values_for( The tile size comes from ``spec`` because the chiplet-swizzle path decodes the workgroup id against ``p_num_pid_m/n``, which are tile-count values. - - A stream-K wgrad launch takes ``ws_ptr``/``ws_bytes`` for the workspace - reduction or ``sk_flags``/``sk_partials`` for linear/tree (laid out by - ``wgrad_streamk_workspace_layout``); its partition block is computed here - from the request's shape and the arch's CTA pool. """ from kernels.common.conv_args import ConvArgs p = _problem(req) tile_m, tile_n = spec.tile_m, spec.tile_n - if req.direction == "wgrad" and spec.streamk != "off": - return _wgrad_streamk_launch_values( - spec.to_wgrad_spec(p), - A_ptr, - B_ptr, - D_ptr, - A_bytes, - B_bytes, - D_bytes, - arch=spec.arch, - problem=p, - ws_ptr=ws_ptr, - ws_bytes=ws_bytes, - sk_flags=sk_flags, - sk_partials=sk_partials, - ) - if sk_flags is not None or sk_partials is not None: - raise ValueError("sk_flags/sk_partials belong to a stream-K wgrad launch") if req.direction == "wgrad": # The resolved degree, not spec.split_k: the latter may be -1 (auto) # or exceed the clamp, while the kernel and _wgrad_grid both use the @@ -2070,18 +1882,7 @@ def dispatch_conv_grouped( f"spec_hash={kid.spec_hash}", f"request_hash={kid.request_hash}", ] - if req.direction == "wgrad" and spec.streamk != "off": - why = ( - f"auto: <= {_WGRAD_STREAMK_MAX_TILES} tiles over >= " - f"{_WGRAD_STREAMK_MIN_ITERS_PER_TILE} K iterations each" - if req.streamk == "auto" - else "requested" - ) - explanation.append( - f"streamk={spec.streamk} reduction={spec.streamk_reduction} ({why}; " - "split-K off)" - ) - elif req.direction == "wgrad": + if req.direction == "wgrad": split_k, _two_stage, requested = _resolve_wgrad_split_k(spec, _problem(req)) if split_k != requested: explanation.append( diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md index d261bbc98bcf..9b9fd921e795 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md @@ -323,16 +323,20 @@ exact for any `wg_M`, avoiding CK's `GemmM % MPerBlock` caveat. straddles the dispatch front (its owner for `linear`, about one per round for `tree`) can wait on a CTA that has not started yet. -Reachable from dispatch: `ConvGroupedRequest.streamk` defaults to `"auto"`, -which selects stream-K (with the `workspace` reduction) for few output tiles on -a long reduction -- at most `_WGRAD_STREAMK_MAX_TILES` tiles of at least -`_WGRAD_STREAMK_MIN_ITERS_PER_TILE` K iterations each -- and the grid-per-tile -split-K kernel otherwise. Elsewhere split-K's equal slices already balance the -grid and stream-K's per-tile fixup is pure overhead. `"off"` never selects it; -`"dp_sk"` / `"persistent"` force it. A stream-K choice pins `split_k=1` and the -direct epilogue and routes `_wgrad_grid` through the instance's own -`wgrad_streamk_grid`, so the spec, the kernargs and the launch grid cannot -disagree about the pool. +Not selected by the grouped-convolution dispatcher: `ConvGroupedRequest` has +no stream-K field and the dispatcher only builds the grid-per-tile kernels, so +stream-K is reached by building a `WgradConvSpec` with `streamk` set. To +measure it, use the implicit-GEMM sweep +(`library/benchmarks/common/benchmark_implicit_gemm_conv.py`, also driven by +`benchmark_conv_compare.py`) with `--direction wgrad --streamk {dp_sk,persistent,all}`, +optionally narrowed with `--streamk-reductions` and pinned to a pool with +`--streamk-ctas` (default `-1`, the per-reduction pool above). The JIT sweep then +adds a stream-K leg, one kernel per geometry x pipeline x mode x reduction +(`atomic` only for fp32 dW); `--compile-all --streamk ...` adds the same kernels +to the AOT cache (`kernel_sweep.py`, opt-in because it more than doubles the +wgrad grid) and `--run-from-cache --streamk ...` runs the cached ones. Every +launch follows the zeroing contract above, and `workspace` launches are timed +together with their Stage 2 reduce. ## Next steps diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py index b4dfa8b824cf..0af7dd2b35fb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py @@ -422,140 +422,6 @@ def test_ws_replicas_changes_the_instance_name(self): ) -class TestStreamKSelection(unittest.TestCase): - """Stream-K is opt-in on the request and reaches the instance spec intact.""" - - def test_streamk_reaches_the_instance_spec(self): - for arch in ("gfx942", "gfx950"): - for mode in ("dp_sk", "persistent"): - for reduction in ("linear", "tree", "workspace"): - r = dispatch_conv_grouped( - _wgrad(arch, G=4, streamk=mode, streamk_reduction=reduction) - ) - ws = r.spec.to_wgrad_spec(_problem(r.request)) - with self.subTest(arch=arch, mode=mode, reduction=reduction): - self.assertEqual(ws.streamk, mode) - self.assertEqual(ws.streamk_reduction, reduction) - self.assertEqual(ws.split_k, 1) - self.assertFalse(ws.two_stage) - self.assertEqual(ws.epilogue, "default") - - def test_auto_keeps_short_k_shapes_on_split_k(self): - # Default request: streamk="auto". A short reduction (wg_K = 2*14*14) - # never reaches the iteration threshold, so the split-K kernel stays. - r = dispatch_conv_grouped(_wgrad("gfx950", G=4)) - self.assertEqual(r.request.streamk, "auto") - self.assertEqual(r.spec.streamk, "off") - self.assertEqual(r.spec.to_wgrad_spec(_problem(r.request)).streamk, "off") - - def test_auto_picks_streamk_for_few_tiles_long_k(self): - # Few output tiles (K=96, C=24, 3x3) on a long reduction (N*Ho*Wo is - # 128*96*96): the shape stream-K is selected for. - for arch in ("gfx942", "gfx950"): - r = dispatch_conv_grouped(_wgrad(arch, N=128, C=24, K=96, Hi=96, Wi=96)) - with self.subTest(arch=arch): - self.assertEqual(r.spec.streamk, "dp_sk") - self.assertEqual(r.spec.streamk_reduction, "workspace") - self.assertEqual(r.spec.split_k, 1) - self.assertTrue(any("streamk=dp_sk" in e for e in r.explanation)) - - def test_auto_respects_thresholds(self): - from dispatch.grouped_convolution import ( - _WGRAD_STREAMK_MAX_TILES, - _WGRAD_STREAMK_MIN_ITERS_PER_TILE, - _wgrad_streamk_auto, - ) - - long_k = _wgrad("gfx950", N=128, C=24, K=96, Hi=96, Wi=96) - self.assertTrue(_wgrad_streamk_auto(long_k, 64, 64, 64)) - # Same shape with a tile_k that leaves fewer iterations than the floor. - big_tk = -(-(128 * 96 * 96) // (_WGRAD_STREAMK_MIN_ITERS_PER_TILE - 1)) - self.assertFalse(_wgrad_streamk_auto(long_k, 64, 64, big_tk)) - # Many tiles: 16 M tiles x 144 N tiles is far above the ceiling. - many = _wgrad("gfx950", N=128, C=1024, K=1024, Hi=96, Wi=96) - self.assertGreater(16 * 144, _WGRAD_STREAMK_MAX_TILES) - self.assertFalse(_wgrad_streamk_auto(many, 64, 64, 64)) - - def test_auto_skips_an_all_dp_partition(self): - # One output tile with 1024..1215 iterations clears both thresholds, - # but on gfx942 that is below the 1216-CTA workspace pool, so the - # partition would be all-DP: one CTA running the whole reduction. - # Auto keeps split-K there; gfx950's 1024-CTA pool still splits it. - for N, arch, mode in ( - (64, "gfx942", "off"), - (72, "gfx942", "off"), - (76, "gfx942", "dp_sk"), - (64, "gfx950", "dp_sk"), - ): - req = _wgrad( - arch, N=N, Hi=32, Wi=32, Y=1, X=1, pad_h=0, pad_w=0, dtype="bf16" - ) - with self.subTest(arch=arch, N=N): - self.assertEqual(dispatch_conv_grouped(req).spec.streamk, mode) - - def test_off_and_gfx1250_never_pick_streamk(self): - req = _wgrad("gfx950", N=128, C=24, K=96, Hi=96, Wi=96, streamk="off") - self.assertEqual(dispatch_conv_grouped(req).spec.streamk, "off") - r = dispatch_conv_grouped(_wgrad("gfx1250", N=128, C=24, K=96, Hi=96, Wi=96)) - self.assertEqual(r.spec.streamk, "off") - - def test_streamk_builds_from_dispatch(self): - from kernels.common.conv_implicit_gemm_wgrad import ( - build_implicit_gemm_conv_wgrad, - ) - - r = dispatch_conv_grouped(_wgrad("gfx950", G=2, streamk="dp_sk")) - k = build_implicit_gemm_conv_wgrad( - r.spec.to_wgrad_spec(_problem(r.request)), arch="gfx950" - ) - # streamk_reduction defaults to the workspace reduction. The CTA pool - # is a launch parameter, so it is not part of the kernel name. - self.assertTrue(k.name.endswith("_sk_skrworkspace"), k.name) - - def test_gfx1250_turns_streamk_away(self): - with self.assertRaises(ValueError) as cm: - dispatch_conv_grouped(_wgrad("gfx1250", streamk="dp_sk")) - self.assertIn("stream-K wgrad is MFMA-only", str(cm.exception)) - - def test_request_validation(self): - from dispatch.grouped_convolution import _request_errors - - self.assertTrue(_request_errors(_wgrad("gfx950", streamk="sideways"))) - self.assertTrue(_request_errors(_wgrad("gfx950", streamk_reduction="atomic"))) - self.assertFalse(_request_errors(_wgrad("gfx950", streamk="auto"))) - fwd = ConvGroupedRequest( - N=2, C=64, K=64, Hi=8, Wi=8, Y=3, X=3, arch="gfx950", streamk="dp_sk" - ) - self.assertTrue(_request_errors(fwd)) - - -class TestStreamKGridShape(unittest.TestCase): - """The dispatch grid is the instance's stream-K grid, one-dimensional.""" - - def test_grid_matches_instance_partition(self): - from kernels.common.conv_implicit_gemm_wgrad import wgrad_streamk_grid - - for arch in ("gfx942", "gfx950"): - for mode in ("dp_sk", "persistent"): - r = dispatch_conv_grouped(_wgrad(arch, G=4, streamk=mode)) - ws = r.spec.to_wgrad_spec(_problem(r.request)) - with self.subTest(arch=arch, mode=mode): - self.assertEqual(r.grid, wgrad_streamk_grid(ws, arch=arch)) - self.assertEqual(r.grid[1:], (1, 1)) - - def test_streamk_changes_the_dispatch_name(self): - off = dispatch_conv_grouped(_wgrad("gfx950", G=4)).spec.kernel_name() - names = { - dispatch_conv_grouped( - _wgrad("gfx950", G=4, streamk=m, streamk_reduction=red) - ).spec.kernel_name() - for m in ("dp_sk", "persistent") - for red in ("linear", "tree") - } - self.assertEqual(len(names), 4) - self.assertNotIn(off, names) - - class TestLaunchContractMatchesKernel(unittest.TestCase): """The dispatcher's signature and launch values must describe the kernel ``to_wgrad_spec`` builds and the grid ``_wgrad_grid`` launches. @@ -607,52 +473,6 @@ def test_odd_wg_N_two_stage(self): if ws.split_k > 1: self.assertTrue(ws.two_stage) - def test_streamk(self): - # The stream-K ABI: the reduction's buffers after ks_count, then the - # partition block, filled from the launch shape and the arch pool. - from dispatch.grouped_convolution import launch_values_for - from kernels.common.conv_args import ConvArgs - from kernels.common.conv_implicit_gemm_wgrad import wgrad_streamk_partition - - for mode in ("dp_sk", "persistent"): - for reduction in ("linear", "tree", "workspace"): - req = _wgrad("gfx950", G=4, streamk=mode, streamk_reduction=reduction) - r = dispatch_conv_grouped(req) - p = _problem(r.request) - ws = r.spec.to_wgrad_spec(p) - abi = ConvArgs.from_problem( - p, - direction="wgrad", - tile_m=ws.tile_m, - tile_n=ws.tile_n, - tile_k=ws.tile_k, - ).arg_names(streamk=reduction) - with self.subTest(mode=mode, reduction=reduction): - self.assertEqual( - [a["name"] for a in r.signature], [n for n, _ in abi] - ) - kw = ( - dict(ws_ptr=0x9000, ws_bytes=64) - if reduction == "workspace" - else dict(sk_flags=0x9000, sk_partials=0xA000) - ) - values = launch_values_for( - r.request, - r.spec, - A_ptr=0x1000, - B_ptr=0x2000, - D_ptr=0x3000, - A_bytes=1, - B_bytes=1, - D_bytes=1, - **kw, - ) - part = wgrad_streamk_partition(ws, arch="gfx950") - self.assertEqual(values["sk_dp_tiles"], part.dp_tiles) - self.assertEqual(values["sk_ctas"], part.sk_ctas) - self.assertEqual(values["sk_pool"], part.max_active_wgs) - self.assertEqual(r.grid, (part.grid_size, 1, 1)) - if __name__ == "__main__": unittest.main() diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md index e215fbf7bc96..15ae339a07d8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md @@ -1994,7 +1994,7 @@ Per-arch facts belong in the arch references — see §21. | `streamk` | `WgradConvSpec` | `"off"` | `"dp_sk"` launches one CTA per data-parallel tile plus `streamk_ctas` stream-K CTAs; `"persistent"` launches only the pool and round-robins the data-parallel tiles. Both balance MAC iterations across CTAs instead of tiles, so a shape with few tiles and a long `wg_K` still fills the device. Mutually exclusive with `split_k` | | `streamk_reduction` | `WgradConvSpec` | `"linear"` | How CTAs sharing a tile combine partials. `"linear"`: the owner waits on each later contributor in turn and stores once. `"tree"`: pairwise fan-in, `ceil(log2)` rounds. Both are deterministic and need no dW pre-zeroing. `"atomic"`: f32 atomics into dW (fp32 only). `"workspace"`: f32 scratch atomics plus the Stage 2 reduce kernel | | `streamk_ctas` | `WgradConvSpec` | `-1` | The CTA pool the stream-K remainder is spread over (CK `max_active_wgs`). A launch parameter, not part of the kernel: it sizes the host partition that arrives as the `sk_*` kernargs. `-1` resolves per reduction: `workspace`/`atomic` (no waiting) get four CTAs per CU (`_STREAMK_CTAS_PER_CU`, counted in CTAs, not waves); `linear`/`tree` get one CTA per CU, because their fixups get slower as more CTAs share a tile. The pool need not be resident: a CTA waits only on higher-index CTAs and publishes its own partial first, so ascending workgroup dispatch keeps the waits from starving their producers | -| `streamk` (dispatch) | `ConvGroupedRequest` | `"auto"` | Picks stream-K + `workspace` when the candidate's tiling yields at most `_WGRAD_STREAMK_MAX_TILES` output tiles of at least `_WGRAD_STREAMK_MIN_ITERS_PER_TILE` K iterations; split-K otherwise. `"off"` disables, `"dp_sk"`/`"persistent"` force | +| `--streamk` (sweep) | `benchmark_implicit_gemm_conv.py` / `kernel_sweep.py` | `off` | Opt-in wgrad stream-K candidates (`dp_sk`/`persistent`/`all` × `--streamk-reductions`); the dispatcher does not select stream-K | #### 12.1.J Quantization From 8b0f49b03d80aa99038da53ff5a2fde0c09bd8a5 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Thu, 8 Oct 2026 17:50:52 +0000 Subject: [PATCH 16/20] fix(rocke): type-check the flag ops and reject bad sync attrs in both engines - global_flag_store / global_flag_wait_eq now require a global i32 pointer in the builders and verifiers of both engines, so the LLVM and HIP lowerings cannot address different memory. - All four lowerers reject a missing or unknown scope/ordering with the same message instead of falling back to defaults or emitting invalid IR. - Document the per-(CTA, wave) flag protocol and its forward-progress assumption (waits only target higher-index CTAs). ISSUE ID : AICK-1943 Co-Authored-By: Claude Opus 5.5 --- .../parity/conv_implicit_gemm_wgrad_emit.c | 2 +- .../rocke/platform/cpp/core/ir/ir_mem.cpp | 23 +++ .../cpp/core/lower_hip/lower_hip_mem.cpp | 98 ++++++----- .../platform/cpp/core/lower_llvm/mem.cpp | 40 ++++- .../rocke/platform/cpp/core/verify.cpp | 6 +- .../primitives/intrinsics_and_primitives.md | 47 +++++- .../rocke/platform/python/rocke/core/ir.py | 33 ++++ .../platform/python/rocke/core/lower_hip.py | 24 ++- .../platform/python/rocke/core/lower_llvm.py | 17 +- .../platform/python/rocke/core/verify.py | 7 +- .../tests/core/future_intrinsic_lowering.cpp | 159 +++++++++++++++--- .../rocke/platform/tests/test_rocke.py | 91 ++++++++++ 12 files changed, 444 insertions(+), 103 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c index da6480bc7dc6..ce94560db390 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/conv_implicit_gemm_wgrad_emit.c @@ -54,7 +54,7 @@ #include "rocke/lower_llvm.h" #include "rocke/verify.h" -/* Stream-K fields of configs 18-27 (Python _STREAMK_CFGS). */ +/* Stream-K fields of configs 21-29 (Python _STREAMK_CFGS). */ static int streamk_cfg(int idx, rocke_implicit_gemm_conv_wgrad_spec_t* spec, const char** arch) { /* idx: dtype a/b, dtype d, atom edge, streamk, reduction, ctas, arch */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp index 245216d6497d..fa16a53af676 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp @@ -536,6 +536,13 @@ void rocke_b_fence(rocke_ir_builder_t* b, const char* scope, const char* orderin (void)rocke_i_op0(b, ROCKE_OP_MEMREF_FENCE, NULL, 0, &a); } +/* ir._is_flag_ptr: the ptr a cross-workgroup flag op addresses. */ +static bool rocke_i_flag_ptr_ok(const rocke_type_t* t) +{ + return t && t->kind == ROCKE_TYPE_PTR && t->space && !strcmp(t->space, "global") && t->pointee + && rocke_i_type_is(t->pointee, "i32"); +} + void rocke_b_global_flag_store(rocke_ir_builder_t* b, rocke_value_t* ptr, rocke_value_t* idx, @@ -556,6 +563,14 @@ void rocke_b_global_flag_store(rocke_ir_builder_t* b, scope = "agent"; if(ordering == NULL) ordering = "release"; + if(!rocke_i_flag_ptr_ok(ptr->type)) + { + (void)rocke_i_set_err(b, + ROCKE_ERR_VALUE, + "global_flag_store expects a global i32 pointer, got %s", + ptr->type ? ptr->type->name : "?"); + return; + } if(!rocke_i_type_is(value->type, "i32")) { (void)rocke_i_set_err(b, @@ -600,6 +615,14 @@ void rocke_b_global_flag_wait_eq(rocke_ir_builder_t* b, } if(scope == NULL) scope = "agent"; + if(!rocke_i_flag_ptr_ok(ptr->type)) + { + (void)rocke_i_set_err(b, + ROCKE_ERR_VALUE, + "global_flag_wait_eq expects a global i32 pointer, got %s", + ptr->type ? ptr->type->name : "?"); + return; + } if(!rocke_i_type_is(expect->type, "i32")) { (void)rocke_i_set_err(b, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp index 5b8641390ae3..8023901a6163 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp @@ -762,8 +762,37 @@ static rocke_status_t _op_memref_global_atomic_add(rocke_h_lowerer_t* lw, const return lw->status; } -/* Python _HIP_ATOMIC_ORDER / _HIP_FENCE_SCOPE / _HIP_MEMORY_SCOPE. Unknown - * spellings return NULL so the ops reject them, like Python's table lookup. */ +/* ir._FENCE_SCOPES / _FENCE_ORDERINGS / _FLAG_STORE_ORDERINGS. */ +static const char* const k_hip_sync_scopes[] = {"workgroup", "agent", "system", NULL}; +static const char* const k_hip_fence_orderings[] + = {"acquire", "release", "acq_rel", "seq_cst", NULL}; +static const char* const k_hip_flag_store_orderings[] = {"monotonic", "release", "seq_cst", NULL}; + +/* Python ir._sync_attr: the lowerer may run on unverified IR, so a missing or + * unknown scope / ordering is rejected rather than defaulted. */ +static const char* hip_sync_attr(rocke_h_lowerer_t* lw, + const rocke_op_t* op, + const char* key, + const char* const* allowed) +{ + const char* s = rocke_attr_get_str(&op->attrs, key); + if(!s) + { + rocke_h_fail( + lw, ROCKE_ERR_VALUE, "%s: %s attribute missing or not a string", op->name, key); + } + for(const char* const* a = allowed; *a; ++a) + { + if(!strcmp(s, *a)) + { + return s; + } + } + rocke_h_fail(lw, ROCKE_ERR_VALUE, "%s: unknown %s '%s'", op->name, key, s); +} + +/* Python _HIP_ATOMIC_ORDER / _HIP_FENCE_SCOPE / _HIP_MEMORY_SCOPE, applied to + * values hip_sync_attr already accepted. */ static const char* hip_atomic_order(const char* o) { if(!strcmp(o, "monotonic")) @@ -774,9 +803,7 @@ static const char* hip_atomic_order(const char* o) return "__ATOMIC_RELEASE"; if(!strcmp(o, "acq_rel")) return "__ATOMIC_ACQ_REL"; - if(!strcmp(o, "seq_cst")) - return "__ATOMIC_SEQ_CST"; - return NULL; + return "__ATOMIC_SEQ_CST"; } static const char* hip_fence_scope(const char* s) @@ -785,9 +812,7 @@ static const char* hip_fence_scope(const char* s) return "workgroup"; if(!strcmp(s, "agent")) return "agent"; - if(!strcmp(s, "system")) - return ""; - return NULL; + return ""; } static const char* hip_memory_scope(const char* s) @@ -796,39 +821,28 @@ static const char* hip_memory_scope(const char* s) return "__HIP_MEMORY_SCOPE_WORKGROUP"; if(!strcmp(s, "agent")) return "__HIP_MEMORY_SCOPE_AGENT"; - if(!strcmp(s, "system")) - return "__HIP_MEMORY_SCOPE_SYSTEM"; - return NULL; + return "__HIP_MEMORY_SCOPE_SYSTEM"; } /* Python _op_memref_fence */ static rocke_status_t _op_memref_fence(rocke_h_lowerer_t* lw, const rocke_op_t* op) { - const char *ord_s, *scp_s, *ord, *scp; + const char *scp, *ord; if(!rocke_h_live(lw)) { return lw->status; } - ord_s = mem_attr_str(op, "ordering", "acq_rel"); - scp_s = mem_attr_str(op, "scope", "agent"); - ord = hip_atomic_order(ord_s); - scp = hip_fence_scope(scp_s); - if(!ord) - { - return rocke_h_fail(lw, ROCKE_ERR_VALUE, "memref.fence: unknown ordering '%s'", ord_s); - } - if(!scp) - { - return rocke_h_fail(lw, ROCKE_ERR_VALUE, "memref.fence: unknown scope '%s'", scp_s); - } - rocke_h_emitf(lw, "__builtin_amdgcn_fence(%s, \"%s\");", ord, scp); + scp = hip_sync_attr(lw, op, "scope", k_hip_sync_scopes); + ord = hip_sync_attr(lw, op, "ordering", k_hip_fence_orderings); + rocke_h_emitf( + lw, "__builtin_amdgcn_fence(%s, \"%s\");", hip_atomic_order(ord), hip_fence_scope(scp)); return lw->status; } /* Python _op_memref_global_flag_store */ static rocke_status_t _op_memref_global_flag_store(rocke_h_lowerer_t* lw, const rocke_op_t* op) { - const char *ord_s, *scp_s, *ord, *scp; + const char *scp, *ord; if(!rocke_h_live(lw)) { return lw->status; @@ -837,34 +851,22 @@ static rocke_status_t _op_memref_global_flag_store(rocke_h_lowerer_t* lw, const { return rocke_h_fail(lw, ROCKE_ERR_VALUE, "memref.global_flag_store: too few operands"); } - ord_s = mem_attr_str(op, "ordering", "release"); - scp_s = mem_attr_str(op, "scope", "agent"); - ord = hip_atomic_order(ord_s); - scp = hip_memory_scope(scp_s); - if(!ord) - { - return rocke_h_fail( - lw, ROCKE_ERR_VALUE, "memref.global_flag_store: unknown ordering '%s'", ord_s); - } - if(!scp) - { - return rocke_h_fail( - lw, ROCKE_ERR_VALUE, "memref.global_flag_store: unknown scope '%s'", scp_s); - } + scp = hip_sync_attr(lw, op, "scope", k_hip_sync_scopes); + ord = hip_sync_attr(lw, op, "ordering", k_hip_flag_store_orderings); rocke_h_emitf(lw, "__hip_atomic_store(&%s[%s], %s, %s, %s);", rocke_h_name(lw, op->operands[0]), rocke_h_name(lw, op->operands[1]), rocke_h_name(lw, op->operands[2]), - ord, - scp); + hip_atomic_order(ord), + hip_memory_scope(scp)); return lw->status; } /* Python _op_memref_global_flag_wait_eq */ static rocke_status_t _op_memref_global_flag_wait_eq(rocke_h_lowerer_t* lw, const rocke_op_t* op) { - const char *scp_s, *scp; + const char* scp; if(!rocke_h_live(lw)) { return lw->status; @@ -873,18 +875,12 @@ static rocke_status_t _op_memref_global_flag_wait_eq(rocke_h_lowerer_t* lw, cons { return rocke_h_fail(lw, ROCKE_ERR_VALUE, "memref.global_flag_wait_eq: too few operands"); } - scp_s = mem_attr_str(op, "scope", "agent"); - scp = hip_memory_scope(scp_s); - if(!scp) - { - return rocke_h_fail( - lw, ROCKE_ERR_VALUE, "memref.global_flag_wait_eq: unknown scope '%s'", scp_s); - } + scp = hip_sync_attr(lw, op, "scope", k_hip_sync_scopes); rocke_h_emitf(lw, "while (__hip_atomic_load(&%s[%s], __ATOMIC_ACQUIRE, %s) != %s) {}", rocke_h_name(lw, op->operands[0]), rocke_h_name(lw, op->operands[1]), - scp, + hip_memory_scope(scp), rocke_h_name(lw, op->operands[2])); return lw->status; } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp index f55290176212..114e8ff853c6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp @@ -183,6 +183,30 @@ static void op_memref_global_store_typed(rocke_lower_t* L, const rocke_op_t* op) (long long)align); } +/* ir._FENCE_SCOPES / _FENCE_ORDERINGS / _FLAG_STORE_ORDERINGS. */ +static const char* const k_ll_sync_scopes[] = {"workgroup", "agent", "system", NULL}; +static const char* const k_ll_fence_orderings[] + = {"acquire", "release", "acq_rel", "seq_cst", NULL}; +static const char* const k_ll_flag_store_orderings[] = {"monotonic", "release", "seq_cst", NULL}; + +/* Python ir._sync_attr: the lowerer may run on unverified IR, so a missing or + * unknown scope / ordering is rejected rather than defaulted or rendered into + * invalid IR. */ +static const char* ll_sync_attr(rocke_lower_t* L, + const rocke_op_t* op, + const char* key, + const char* const* allowed) +{ + const char* s = rocke_attr_get_str(&op->attrs, key); + if(!s) + rocke_ll_fail( + L, ROCKE_ERR_VALUE, "%s: %s attribute missing or not a string", op->name, key); + for(const char* const* a = allowed; *a; ++a) + if(strcmp(s, *a) == 0) + return s; + rocke_ll_fail(L, ROCKE_ERR_VALUE, "%s: unknown %s '%s'", op->name, key, s); +} + /* Python _llvm_syncscope: "system" is LLVM's default scope (no clause). */ static const char* ll_syncscope(rocke_lower_t* L, const char* scope) { @@ -194,10 +218,9 @@ static const char* ll_syncscope(rocke_lower_t* L, const char* scope) /* Python _op_memref_fence. */ static void op_memref_fence(rocke_lower_t* L, const rocke_op_t* op) { - rocke_ll_emitf(L, - " fence%s %s", - ll_syncscope(L, ll_attr_str(op, "scope", "agent")), - ll_attr_str(op, "ordering", "acq_rel")); + const char* scope = ll_sync_attr(L, op, "scope", k_ll_sync_scopes); + const char* ordering = ll_sync_attr(L, op, "ordering", k_ll_fence_orderings); + rocke_ll_emitf(L, " fence%s %s", ll_syncscope(L, scope), ordering); } /* Python _op_memref_global_flag_store. */ @@ -208,6 +231,8 @@ static void op_memref_global_flag_store(rocke_lower_t* L, const rocke_op_t* op) const rocke_value_t* ptr = op->operands[0]; const rocke_value_t* idx = op->operands[1]; const rocke_value_t* val = op->operands[2]; + const char* scope = ll_sync_attr(L, op, "scope", k_ll_sync_scopes); + const char* ordering = ll_sync_attr(L, op, "ordering", k_ll_flag_store_orderings); const char* gep = rocke_ll_fresh(L, "gep"); rocke_ll_emitf(L, " %s = getelementptr inbounds i32, ptr addrspace(1) %s, i32 %s", @@ -218,8 +243,8 @@ static void op_memref_global_flag_store(rocke_lower_t* L, const rocke_op_t* op) " store atomic i32 %s, ptr addrspace(1) %s%s %s, align 4", rocke_ll_operand(L, val), gep, - ll_syncscope(L, ll_attr_str(op, "scope", "agent")), - ll_attr_str(op, "ordering", "release")); + ll_syncscope(L, scope), + ordering); } /* Python _op_memref_global_flag_wait_eq: a self-contained spin loop around an @@ -231,6 +256,7 @@ static void op_memref_global_flag_wait_eq(rocke_lower_t* L, const rocke_op_t* op const rocke_value_t* ptr = op->operands[0]; const rocke_value_t* idx = op->operands[1]; const rocke_value_t* expect = op->operands[2]; + const char* scope = ll_sync_attr(L, op, "scope", k_ll_sync_scopes); const char* gep = rocke_ll_fresh(L, "gep"); rocke_ll_block_t* cur = rocke_ll_current(L); rocke_ll_block_t* spin; @@ -255,7 +281,7 @@ static void op_memref_global_flag_wait_eq(rocke_lower_t* L, const rocke_op_t* op " %s = load atomic i32, ptr addrspace(1) %s%s acquire, align 4", val, gep, - ll_syncscope(L, ll_attr_str(op, "scope", "agent"))); + ll_syncscope(L, scope)); rocke_ll_block_emitf( L, spin, " %s = icmp eq i32 %s, %s", hit, val, rocke_ll_operand(L, expect)); done = rocke_ll_new_block(L, "flag.done"); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp index bb3ac250de7d..62afcbad52fe 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp @@ -529,9 +529,11 @@ static void check_contract(verifier_t* v, const rocke_op_t* op) return; } const rocke_type_t* pt = op->operands[0]->type; - if(!pt || pt->kind != ROCKE_TYPE_PTR || !pt->space || strcmp(pt->space, "global")) + /* ir._is_flag_ptr: both lowerers address the flag as an i32. */ + if(!pt || pt->kind != ROCKE_TYPE_PTR || !pt->space || strcmp(pt->space, "global") + || strcmp(tn(pt->pointee), "i32")) { - v_errf(v, op, "%s: expected a global pointer, got %s", name, tn(pt)); + v_errf(v, op, "%s: expected a global i32 pointer, got %s", name, tn(pt)); } const char* what[] = {"index", "value"}; for(int i = 0; i < 2; ++i) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/intrinsics_and_primitives.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/intrinsics_and_primitives.md index 11488b3917c6..2c8da4830649 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/intrinsics_and_primitives.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/intrinsics_and_primitives.md @@ -252,8 +252,7 @@ local_off = b.global_atomic_add(Counter, eid, b.const_i32(1)) ### Cross-workgroup flags and fences Stream-K fixups hand a partial tile from one workgroup to another through -global memory, guarded by a per-workgroup i32 flag. Three leaf ops cover -the protocol: +global memory, guarded by i32 flags. Three leaf ops cover the protocol: ```text fence(*, scope="agent", ordering="acq_rel") @@ -262,12 +261,44 @@ global_flag_wait_eq(ptr, idx, expect, *, scope="agent") ``` `scope` is one of `workgroup`, `agent`, `system` (LLVM `syncscope`; -`system` has no clause). The producer stores its partial tile, every wave -issues `fence(ordering="release")`, the workgroup barriers, and then the -flag is stored. The consumer calls `global_flag_wait_eq`, which lowers to a -self-contained loop around an `acquire` atomic load, so its later global -loads observe the released partial. At `agent` scope on multi-XCD CDNA -parts the release writes back and the acquire invalidates the per-XCD L2. +`system` has no clause). `ptr` must be a `ptr`: the builders and +the verifier reject any other element type, because the LLVM lowering +addresses the flag as an i32 while the HIP lowering indexes the pointee. +Both lowerers also reject a missing or unknown `scope` / `ordering` instead +of defaulting it, so IR that skipped the verifier cannot lower to an invalid +`fence` or atomic. + +The wgrad stream-K fixups (`linear` and `tree` in +`library/kernels/common/conv_implicit_gemm_wgrad.py`) keep one flag per +(stream-K CTA, wave), at index `sk_cta * n_waves + warp_id`. A consumer lane +only reads the partial written by the same lane of the same wave in the +producer CTA, so each wave hands off on its own, with no workgroup barrier +on either side: + +- Producer wave: store this wave's slice of the partial tile, issue + `fence(scope="agent", ordering="release")`, then + `global_flag_store(flags, sk_cta * n_waves + warp_id, 1)`. +- Consumer wave: `global_flag_wait_eq(flags, partner * n_waves + warp_id, 1)`, + which lowers to a self-contained loop around an `acquire` atomic load, + then load the partner's partial. The acquire orders those loads after the + observed flag, so they see what the producer wave released. + +At `agent` scope on multi-XCD CDNA parts the release writes back and the +acquire invalidates the per-XCD L2. The flag region must be zeroed before +each launch; the partials are always written before they are read. + +Forward progress: a CTA only waits on a higher-index stream-K CTA (the +`linear` owner folds `sk_cta + 1 ...`, a `tree` receiver folds +`sk_cta + stride`), never on a lower one. If the whole CTA pool is resident +at once, every wait terminates. When the pool is larger than what can be +resident, the protocol relies on workgroups being dispatched in ascending +index order, which the hardware does in practice but the programming model +does not guarantee. Under that order only the CTAs of the tile straddling +the dispatch front can wait on a CTA that has not started, and they hold a +bounded number of slots (one for `linear`, about one per round for +`tree`), so the other resident CTAs finish and free slots for the producers. +Do not add a wait on a lower-index CTA, and do not reorder the grid (for +example by remapping `block_id_x`) without revisiting this. | IR op | LLVM | HIP | |---|---|---| diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py index a839de396f4f..66cbee6d2cc6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py @@ -143,6 +143,7 @@ def _streaming(temporal_hint: TemporalHint) -> bool: raise TypeError(f"temporal_hint must be a TemporalHint, got {temporal_hint!r}") return temporal_hint is TemporalHint.STREAMING + # LLVM ``syncscope`` values accepted by the cross-workgroup fence / flag ops. # ``system`` renders as the empty (default) scope. _FENCE_SCOPES = ("workgroup", "agent", "system") @@ -151,6 +152,36 @@ def _streaming(temporal_hint: TemporalHint) -> bool: _FLAG_STORE_ORDERINGS = ("monotonic", "release", "seq_cst") +def _is_flag_ptr(t: "Type") -> bool: + """True for the ``ptr`` a cross-workgroup flag op addresses. + + Both lowerers hard-code an i32 element (LLVM GEPs ``i32``; HIP indexes + the pointer's own element type), so any other pointee would make the two + disagree on the flag's address and width. + """ + return isinstance(t, PtrType) and t.space == "global" and t.pointee.name == "i32" + + +def _check_flag_ptr(op_name: str, ptr: "Value") -> None: + if not _is_flag_ptr(ptr.type): + raise ValueError(f"{op_name} expects a global i32 pointer, got {ptr.type.name}") + + +def _sync_attr(op: "Op", key: str, allowed: Sequence[str]) -> str: + """Scope / ordering attr of a fence or flag op, checked for the lowerers. + + The lowerers can run on IR that never went through :func:`verify`, so a + missing or unknown value is rejected here instead of being defaulted or + rendered into invalid IR. The C++ engine raises the same messages. + """ + value = op.attrs.get(key) + if not isinstance(value, str): + raise ValueError(f"{op.name}: {key} attribute missing or not a string") + if value not in allowed: + raise ValueError(f"{op.name}: unknown {key} '{value}'") + return value + + # ----- target-neutral MMA metadata --------------------------------------- # # ``IRBuilder.mma`` emits a single ``tile.mma`` op keyed by ``op_id``; the ISA @@ -1505,6 +1536,7 @@ def global_flag_store( Lowers to ``store atomic i32 ... syncscope(scope) ordering``. Pair with :meth:`global_flag_wait_eq` on the consumer side. """ + _check_flag_ptr("global_flag_store", ptr) if value.type.name != "i32": raise ValueError( f"global_flag_store expects an i32 value, got {value.type.name}" @@ -1535,6 +1567,7 @@ def global_flag_wait_eq( data-dependent trip count in the IR; it is a leaf op so structured control flow stays ``scf.for`` / ``scf.if`` only. """ + _check_flag_ptr("global_flag_wait_eq", ptr) if expect.type.name != "i32": raise ValueError( f"global_flag_wait_eq expects an i32 value, got {expect.type.name}" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py index 93f39ff8b1ef..a18b3d9a6fcb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py @@ -20,6 +20,9 @@ from .arch.wmma_scale import gfx1250_scaled_wmma from .dtypes import dtype_info from .ir import ( + _FENCE_ORDERINGS, + _FENCE_SCOPES, + _FLAG_STORE_ORDERINGS, KernelDef, Op, PtrType, @@ -28,6 +31,7 @@ Type, Value, VectorType, + _sync_attr, require_streaming_arch, ) @@ -269,9 +273,6 @@ def _encode_waitcnt_gfx11(vmcnt: int, expcnt: int, lgkmcnt: int) -> int: return (ec_b & 0x7) | ((lk_b & 0x3F) << 4) | ((vm_b & 0x3F) << 10) -# Default arch for the HIP path. gfx950 is the historical CDNA target whose -# emitted source is the byte-identical baseline; callers that don't pass an -# arch (e.g. the in-tree coverage tests) keep getting exactly that output. # Fence / flag-op scope and ordering spellings for the HIP atomic builtins. _HIP_ATOMIC_ORDER = { "monotonic": "__ATOMIC_RELAXED", @@ -287,6 +288,9 @@ def _encode_waitcnt_gfx11(vmcnt: int, expcnt: int, lgkmcnt: int) -> int: "system": "__HIP_MEMORY_SCOPE_SYSTEM", } +# Default arch for the HIP path. gfx950 is the historical CDNA target whose +# emitted source is the byte-identical baseline; callers that don't pass an +# arch (e.g. the in-tree coverage tests) keep getting exactly that output. _DEFAULT_HIP_ARCH = "gfx950" @@ -1748,26 +1752,30 @@ def _op_memref_global_atomic_add(self, op: Op) -> None: def _op_memref_fence(self, op: Op) -> None: """Lower ``memref.fence`` via ``__builtin_amdgcn_fence``.""" + scope = _sync_attr(op, "scope", _FENCE_SCOPES) + ordering = _sync_attr(op, "ordering", _FENCE_ORDERINGS) self._emit( - f"__builtin_amdgcn_fence({_HIP_ATOMIC_ORDER[op.attrs['ordering']]}, " - f"\"{_HIP_FENCE_SCOPE[op.attrs['scope']]}\");" + f"__builtin_amdgcn_fence({_HIP_ATOMIC_ORDER[ordering]}, " + f'"{_HIP_FENCE_SCOPE[scope]}");' ) def _op_memref_global_flag_store(self, op: Op) -> None: """Lower ``memref.global_flag_store`` via ``__hip_atomic_store``.""" ptr, idx, val = op.operands + scope = _sync_attr(op, "scope", _FENCE_SCOPES) + ordering = _sync_attr(op, "ordering", _FLAG_STORE_ORDERINGS) self._emit( f"__hip_atomic_store(&{_name(ptr)}[{_name(idx)}], {_name(val)}, " - f"{_HIP_ATOMIC_ORDER[op.attrs['ordering']]}, " - f"{_HIP_MEMORY_SCOPE[op.attrs['scope']]});" + f"{_HIP_ATOMIC_ORDER[ordering]}, {_HIP_MEMORY_SCOPE[scope]});" ) def _op_memref_global_flag_wait_eq(self, op: Op) -> None: """Lower ``memref.global_flag_wait_eq`` to an acquire-load spin loop.""" ptr, idx, expect = op.operands + scope = _sync_attr(op, "scope", _FENCE_SCOPES) self._emit( f"while (__hip_atomic_load(&{_name(ptr)}[{_name(idx)}], __ATOMIC_ACQUIRE, " - f"{_HIP_MEMORY_SCOPE[op.attrs['scope']]}) != {_name(expect)}) {{}}" + f"{_HIP_MEMORY_SCOPE[scope]}) != {_name(expect)}) {{}}" ) def _op_memref_global_atomic_add_pk_bf16(self, op: Op) -> None: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py index 70a5606809e3..589ef503fab5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py @@ -39,6 +39,9 @@ from .codegen_policy import codegen_policy_for_kernel from .ir import ( + _FENCE_ORDERINGS, + _FENCE_SCOPES, + _FLAG_STORE_ORDERINGS, BF16, F32, F16, @@ -54,6 +57,7 @@ Type, Value, VectorType, + _sync_attr, require_streaming_arch, split_loc, ) @@ -2962,13 +2966,15 @@ def _op_memref_global_atomic_add(self, op: Op) -> None: def _op_memref_fence(self, op: Op) -> None: """Lower ``memref.fence`` to an LLVM ``fence`` at the op's scope.""" - self._current().emit( - f" fence{_llvm_syncscope(op.attrs['scope'])} {op.attrs['ordering']}" - ) + scope = _sync_attr(op, "scope", _FENCE_SCOPES) + ordering = _sync_attr(op, "ordering", _FENCE_ORDERINGS) + self._current().emit(f" fence{_llvm_syncscope(scope)} {ordering}") def _op_memref_global_flag_store(self, op: Op) -> None: """Lower ``memref.global_flag_store`` to a scoped ``store atomic i32``.""" ptr, idx, val = op.operands + scope = _sync_attr(op, "scope", _FENCE_SCOPES) + ordering = _sync_attr(op, "ordering", _FLAG_STORE_ORDERINGS) gep = self._fresh("gep") self._current().emit( f" {gep} = getelementptr inbounds i32, ptr addrspace(1) " @@ -2976,7 +2982,7 @@ def _op_memref_global_flag_store(self, op: Op) -> None: ) self._current().emit( f" store atomic i32 {self._operand(val)}, ptr addrspace(1) {gep}" - f"{_llvm_syncscope(op.attrs['scope'])} {op.attrs['ordering']}, align 4" + f"{_llvm_syncscope(scope)} {ordering}, align 4" ) def _op_memref_global_flag_wait_eq(self, op: Op) -> None: @@ -2987,6 +2993,7 @@ def _op_memref_global_flag_wait_eq(self, op: Op) -> None: orders every later global load after the observed flag value. """ ptr, idx, expect = op.operands + scope = _sync_attr(op, "scope", _FENCE_SCOPES) gep = self._fresh("gep") cur = self._current() cur.emit( @@ -3000,7 +3007,7 @@ def _op_memref_global_flag_wait_eq(self, op: Op) -> None: hit = self._fresh("flag.hit") spin.emit( f" {val} = load atomic i32, ptr addrspace(1) {gep}" - f"{_llvm_syncscope(op.attrs['scope'])} acquire, align 4" + f"{_llvm_syncscope(scope)} acquire, align 4" ) spin.emit(f" {hit} = icmp eq i32 {val}, {self._operand(expect)}") done = self._new_block("flag.done") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/verify.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/verify.py index df2e2091e549..8a2dbdca93c8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/verify.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/verify.py @@ -38,6 +38,7 @@ _FENCE_ORDERINGS, _FENCE_SCOPES, _FLAG_STORE_ORDERINGS, + _is_flag_ptr, KernelDef, Op, Region, @@ -400,8 +401,10 @@ def _check_contract(self, op: Op, scope: Dict[str, Value]) -> None: ) return ptr, idx, val = op.operands - if not isinstance(ptr.type, PtrType) or ptr.type.space != "global": - self.err(f"{name}: expected a global pointer, got {ptr.type.name}", op) + if not _is_flag_ptr(ptr.type): + self.err( + f"{name}: expected a global i32 pointer, got {ptr.type.name}", op + ) for what, v in (("index", idx), ("value", val)): if v.type.name != "i32": self.err(f"{name}: {what} must be i32, got {v.type.name}", op) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp index 98d03e9bc0ef..2e715e38b0e0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp @@ -17,6 +17,7 @@ #include #include #include +#include #include #include "rocke/arch_target.h" @@ -659,12 +660,15 @@ void case_flag_ops_hip() EXPECT_IR(hip, "while (__hip_atomic_load("); } -/* Raw op with a bad scope/ordering, as serialized IR could carry past the - * builder's checks: the HIP lowerer must reject it, not pick a default. */ +/* Raw op with a bad or missing scope/ordering, as serialized IR could carry + * past the builder's checks: both lowerers must reject it (ROCKE_ERR_VALUE), + * never pick a default or render it into invalid IR. Every other attr is + * valid, so the rejection is attributable to `key`; value == nullptr drops + * the key. Mirrors test_rocke.py test_flag_ops_lowerers_reject_bad_attrs. */ void expect_flag_attr_rejected(rocke_opcode_t kind, const char* key, const char* value) { rocke_ir_builder_t b; - rocke_ir_builder_init(&b, "flag_bad_attr_hip"); + rocke_ir_builder_init(&b, "flag_bad_attr"); rocke_value_t* operands[3]; int n = 0; if(kind != ROCKE_OP_MEMREF_FENCE) @@ -673,9 +677,20 @@ void expect_flag_attr_rejected(rocke_opcode_t kind, const char* key, const char* operands[n++] = rocke_b_const_i32(&b, 0); operands[n++] = rocke_b_const_i32(&b, 1); } + const char* scope = "agent"; + const char* ordering = kind == ROCKE_OP_MEMREF_FENCE ? "acq_rel" + : kind == ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE ? "release" + : nullptr; + if(!strcmp(key, "scope")) + scope = value; + else + ordering = value; rocke_attr_map_t attrs; rocke_attr_map_init(&attrs); - rocke_attr_set_str(&b, &attrs, key, value); + if(scope) + rocke_attr_set_str(&b, &attrs, "scope", scope); + if(ordering) + rocke_attr_set_str(&b, &attrs, "ordering", ordering); rocke_b_op(&b, kind, operands, n, nullptr, 0, &attrs, nullptr, 0, nullptr, nullptr); rocke_b_ret(&b); @@ -686,8 +701,102 @@ void expect_flag_attr_rejected(rocke_opcode_t kind, const char* key, const char* const rocke_status_t st = rocke_lower_kernel_to_hip(&b, rocke_ir_builder_kernel(&b), &opts, &out); if(st != ROCKE_ERR_VALUE) - fail("HIP lowering must reject an unknown flag/fence scope or ordering", __LINE__); + fail("HIP lowering must reject a missing/unknown flag/fence scope or ordering", __LINE__); rocke_strbuf_free(&out); + + char* ll = nullptr; + char err[ROCKE_ERR_MSG_CAP]; + err[0] = '\0'; + const rocke_status_t lst = rocke_lower_kernel_to_llvm_ex( + rocke_ir_builder_kernel(&b), ROCKE_LLVM_FLAVOR_AUTO, "gfx950", &ll, err, sizeof(err)); + if(lst != ROCKE_ERR_VALUE) + fail("LLVM lowering must reject a missing/unknown flag/fence scope or ordering", __LINE__); + char want[160]; + if(value) + snprintf(want, sizeof(want), "unknown %s '%s'", key, value); + else + snprintf(want, sizeof(want), "%s attribute missing or not a string", key); + if(!strstr(err, want)) + { + std::fprintf(stderr, "LLVM error '%s' lacks '%s'\n", err, want); + fail("LLVM lowering rejected a flag/fence attr with the wrong message", __LINE__); + } + std::free(ll); + rocke_ir_builder_free(&b); +} + +/* A flag op on a non-i32 pointer: the LLVM lowerer would GEP i32 while HIP + * indexes the pointee, so the builder and the verifier both refuse it. */ +void case_flag_ops_reject_non_i32_ptr() +{ + for(int wait = 0; wait < 2; ++wait) + { + rocke_ir_builder_t b; + rocke_ir_builder_init(&b, "flag_f16_ptr"); + const char* want = wait ? "global_flag_wait_eq expects a global i32 pointer, got " + "ptr" + : "global_flag_store expects a global i32 pointer, got " + "ptr"; + /* The builder raises (ckc::ValueError) or latches the status. */ + std::string got; + try + { + rocke_value_t* f = global_ptr_param(&b, "flags", rocke_f16()); + if(wait) + rocke_b_global_flag_wait_eq( + &b, f, rocke_b_const_i32(&b, 0), rocke_b_const_i32(&b, 1), NULL); + else + rocke_b_global_flag_store( + &b, f, rocke_b_const_i32(&b, 0), rocke_b_const_i32(&b, 1), NULL, NULL); + if(rocke_ir_builder_status(&b) == ROCKE_ERR_VALUE && rocke_ir_builder_error(&b)) + got = rocke_ir_builder_error(&b); + } + catch(const std::exception& e) + { + got = e.what(); + } + if(got != want) + { + std::fprintf(stderr, "builder error '%s', want '%s'\n", got.c_str(), want); + fail("flag-op builder must reject a non-i32 pointer", __LINE__); + } + rocke_ir_builder_free(&b); + } + + rocke_ir_builder_t b; + rocke_ir_builder_init(&b, "flag_f16_ptr_verify"); + rocke_value_t* operands[] = {global_ptr_param(&b, "flags", rocke_f16()), + rocke_b_const_i32(&b, 0), + rocke_b_const_i32(&b, 1)}; + const rocke_opcode_t kinds[] + = {ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ}; + for(rocke_opcode_t kind : kinds) + { + rocke_attr_map_t attrs; + rocke_attr_map_init(&attrs); + rocke_attr_set_str(&b, &attrs, "scope", "agent"); + if(kind == ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE) + rocke_attr_set_str(&b, &attrs, "ordering", "release"); + rocke_b_op(&b, kind, operands, 3, nullptr, 0, &attrs, nullptr, 0, nullptr, nullptr); + } + rocke_b_ret(&b); + rocke_diag_t* diags = nullptr; + size_t n = 0; + rocke_verify(rocke_ir_builder_kernel(&b), &diags, &n); + bool saw_store = false, saw_wait = false; + for(size_t i = 0; i < n; ++i) + { + const std::string m = diags[i].message ? diags[i].message : ""; + saw_store |= m.find("memref.global_flag_store: expected a global i32 pointer, got " + "ptr") + != std::string::npos; + saw_wait |= m.find("memref.global_flag_wait_eq: expected a global i32 pointer, got " + "ptr") + != std::string::npos; + } + rocke_diags_free(diags, n); + if(!saw_store || !saw_wait) + fail("verifier must reject a flag op on a non-i32 pointer", __LINE__); rocke_ir_builder_free(&b); } @@ -703,17 +812,17 @@ void case_flag_ops_verify_and_llvm_arity() rocke_attr_map_init(&fence_attrs); rocke_attr_set_str(&b, &fence_attrs, "scope", "agent"); rocke_attr_set_str(&b, &fence_attrs, "ordering", "monotonic"); - rocke_b_op(&b, - ROCKE_OP_MEMREF_FENCE, - nullptr, - 0, - nullptr, - 0, - &fence_attrs, - nullptr, - 0, - nullptr, - nullptr); + rocke_op_t* fence = rocke_b_op(&b, + ROCKE_OP_MEMREF_FENCE, + nullptr, + 0, + nullptr, + 0, + &fence_attrs, + nullptr, + 0, + nullptr, + nullptr); rocke_attr_map_t store_attrs; rocke_attr_map_init(&store_attrs); rocke_attr_set_str(&b, &store_attrs, "scope", "agent"); @@ -748,24 +857,35 @@ void case_flag_ops_verify_and_llvm_arity() if(!saw_arity) fail("verifier must reject a two-operand flag store", __LINE__); + /* Make the fence valid so the LLVM rejection is the flag store's arity. */ + rocke_attr_set_str(&b, &fence->attrs, "ordering", "acq_rel"); char* ll = nullptr; char err[ROCKE_ERR_MSG_CAP]; err[0] = '\0'; const rocke_status_t st = rocke_lower_kernel_to_llvm_ex( rocke_ir_builder_kernel(&b), ROCKE_LLVM_FLAVOR_AUTO, "gfx950", &ll, err, sizeof(err)); - if(st != ROCKE_ERR_VALUE) + if(st != ROCKE_ERR_VALUE || !strstr(err, "memref.global_flag_store expects three operands")) fail("LLVM lowering must reject a two-operand flag store", __LINE__); std::free(ll); rocke_ir_builder_free(&b); } -void case_flag_ops_hip_reject_bad_attrs() +void case_flag_ops_lowerers_reject_bad_attrs() { expect_flag_attr_rejected(ROCKE_OP_MEMREF_FENCE, "ordering", "bogus"); expect_flag_attr_rejected(ROCKE_OP_MEMREF_FENCE, "scope", "bogus"); expect_flag_attr_rejected(ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, "ordering", "bogus"); expect_flag_attr_rejected(ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, "scope", "bogus"); expect_flag_attr_rejected(ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ, "scope", "bogus"); + /* A load-only ordering is not a fence / store ordering. */ + expect_flag_attr_rejected(ROCKE_OP_MEMREF_FENCE, "ordering", "monotonic"); + expect_flag_attr_rejected(ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, "ordering", "acquire"); + /* Missing attrs are rejected, not defaulted. */ + expect_flag_attr_rejected(ROCKE_OP_MEMREF_FENCE, "ordering", nullptr); + expect_flag_attr_rejected(ROCKE_OP_MEMREF_FENCE, "scope", nullptr); + expect_flag_attr_rejected(ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, "ordering", nullptr); + expect_flag_attr_rejected(ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, "scope", nullptr); + expect_flag_attr_rejected(ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ, "scope", nullptr); } /* ---- av.load / av.store (agent-scope 128-bit vector mem) ---- */ @@ -1183,8 +1303,9 @@ const TestCase k_cases[] = { {"flag_ops", case_flag_ops}, {"fence_system_scope", case_fence_system_scope}, {"flag_ops_hip", case_flag_ops_hip}, - {"flag_ops_hip_reject_bad_attrs", case_flag_ops_hip_reject_bad_attrs}, + {"flag_ops_lowerers_reject_bad_attrs", case_flag_ops_lowerers_reject_bad_attrs}, {"flag_ops_verify_and_llvm_arity", case_flag_ops_verify_and_llvm_arity}, + {"flag_ops_reject_non_i32_ptr", case_flag_ops_reject_non_i32_ptr}, {"opcode_names_are_aligned", case_opcode_names_are_aligned}, }; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py index abb81c023320..3b8d8dd7380b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py @@ -21,6 +21,7 @@ from __future__ import annotations +import re import unittest from unittest.mock import patch, Mock @@ -3344,6 +3345,96 @@ def test_flag_ops_reject_bad_scope_ordering_and_type(self): with self.assertRaisesRegex(ValueError, "unknown flag scope"): b.global_flag_wait_eq(flags, b.const_i32(0), b.const_i32(1), scope="gpu") + def test_flag_ops_reject_non_i32_pointer(self): + # LLVM GEPs the flag as i32 while HIP indexes the pointee type, so any + # other element type would make the two engines address different + # bytes. Builder and verifier both refuse it (C++ mirror: + # case_flag_ops_reject_non_i32_ptr). + from rocke.core.ir import F16, PtrType + from rocke.core.verify import ERROR, verify + + b = self._builder("flagf16") + half = b.param("flags", PtrType(F16, "global")) + for name, emit in ( + ("global_flag_store", b.global_flag_store), + ("global_flag_wait_eq", b.global_flag_wait_eq), + ): + with self.assertRaisesRegex( + ValueError, + f"^{name} expects a global i32 pointer, got ptr$", + ): + emit(half, b.const_i32(0), b.const_i32(1)) + + b = self._builder("flagf16verify") + flags = self._flag_kernel(b) + half = b.param("half", PtrType(F16, "global")) + b.global_flag_store(flags, b.const_i32(0), b.const_i32(1)) + b.global_flag_wait_eq(flags, b.const_i32(0), b.const_i32(1)) + b.ret() + for op in b.kernel.body.ops: + if op.name.startswith("memref.global_flag_"): + op.operands[0] = half + msgs = [d.message for d in verify(b.kernel) if d.severity == ERROR] + for name in ("memref.global_flag_store", "memref.global_flag_wait_eq"): + needle = f"{name}: expected a global i32 pointer, got ptr" + self.assertTrue(any(needle in m for m in msgs), (needle, msgs)) + + def test_flag_ops_lowerers_reject_bad_attrs(self): + # Python counterpart of the C++ expect_flag_attr_rejected: an op with a + # missing or unknown scope / ordering (serialized IR that skipped the + # builder and the verifier) is rejected by BOTH lowerers instead of + # defaulting or rendering invalid IR such as ``fence ... bogus``. + from rocke.core.backend import BackendError, _lower_via_cpp_engine + from rocke.core.ir_serialize import serialize + from rocke.core.lower_hip import lower_kernel_to_hip + from rocke.core.lower_llvm import _lower_kernel_to_llvm_python + + valid = { + "memref.fence": {"scope": "agent", "ordering": "acq_rel"}, + "memref.global_flag_store": {"scope": "agent", "ordering": "release"}, + "memref.global_flag_wait_eq": {"scope": "agent"}, + } + cases = [ + ("memref.fence", "ordering", "bogus"), + ("memref.fence", "scope", "bogus"), + ("memref.fence", "ordering", "monotonic"), + ("memref.global_flag_store", "ordering", "bogus"), + ("memref.global_flag_store", "ordering", "acquire"), + ("memref.global_flag_store", "scope", "bogus"), + ("memref.global_flag_wait_eq", "scope", "bogus"), + ] + [(name, key, None) for name, attrs in valid.items() for key in attrs] + for name, key, value in cases: + with self.subTest(op=name, key=key, value=value): + b = self._builder("flag_bad_attr") + if name == "memref.fence": + b.fence() + else: + flags = self._flag_kernel(b) + emit = getattr(b, name.split(".", 1)[1]) + emit(flags, b.const_i32(0), b.const_i32(1)) + b.ret() + (op,) = [o for o in b.kernel.body.ops if o.name == name] + op.attrs.clear() + op.attrs.update(valid[name]) + if value is None: + del op.attrs[key] + want = f"{name}: {key} attribute missing or not a string" + else: + op.attrs[key] = value + want = f"{name}: unknown {key} '{value}'" + with self.assertRaisesRegex(ValueError, re.escape(want)): + lower_kernel_to_hip(b.kernel) + with self.assertRaisesRegex(ValueError, re.escape(want)): + _lower_kernel_to_llvm_python(b.kernel) + try: + _lower_via_cpp_engine(serialize(b.kernel), "gfx950", "llvm20") + except BackendError: + continue # engine not built: C++ is covered by its ctest + except Exception as e: # noqa: BLE001 -- engine error type + self.assertIn(want, str(e)) + else: + self.fail(f"C++ engine accepted {name} with {key}={value!r}") + def test_verifier_rejects_malformed_flag_ops(self): # Serialized IR can carry flag ops the builder would have refused; the # verifier has to catch them before a lowerer unpacks the operands. From 7c820fe3c96e683b84c4f3dfe0a5234732346fbe Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Thu, 8 Oct 2026 17:50:52 +0000 Subject: [PATCH 17/20] fix(rocke): align the stream-K wgrad host API and validator with Python - The C host helpers (partition, grid, workspace layout/bytes) take the target arch so the auto CTA pool matches the kernargs on gfx942. - The C validator follows Python's gate order and ports the vector_size_c and group_merge gates, so multi-cause rejects give the same reason. - Python rejects streamk_ctas of 0 or below -1, as C does. - Add a ctest with C-side host-plan and validator vectors so parity is checked without the Python binding, and test a shape that actually splits under the default pool. - The GPU repeat test now claims release detection only; codegen tests pin the acquire spin and release publish in IR and ISA. ISSUE ID : AICK-1943 Co-Authored-By: Claude Opus 5.5 --- ...ROCKE_ENGINE_test_categories_external.yaml | 4 + .../common/conv_implicit_gemm_wgrad.py | 6 + .../library/tests/test_conv_wgrad_streamk.py | 105 ++++- .../tests/test_conv_wgrad_streamk_gpu.py | 295 ++++++++++++- .../platform/cpp/bindings/rocke_engine.cpp | 34 +- .../rocke/instance_conv_implicit_gemm_wgrad.h | 46 +- .../common/conv_implicit_gemm_wgrad.cpp | 195 +++++++-- .../rocke/platform/tests/CMakeLists.txt | 21 + .../tests/instances/wgrad_streamk_host.cpp | 404 ++++++++++++++++++ 9 files changed, 1021 insertions(+), 89 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/wgrad_streamk_host.cpp diff --git a/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml b/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml index f2f3948a392c..f6b0c1a524d2 100644 --- a/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml +++ b/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml @@ -29,6 +29,8 @@ test_categories: - "rocke_ir_serialize_roundtrip" - "rocke_mma_frag_ssot" - "rocke_streamk_iter_partition" + - "rocke_wgrad_streamk_host" + - "rocke_future_intrinsic_lowering" - "rocke_tiled_attention_2d_reentrancy" - "rocke_python_import_smoke" - "rocke_golden_static" @@ -52,6 +54,8 @@ test_categories: - "rocke_ir_serialize_roundtrip" - "rocke_mma_frag_ssot" - "rocke_streamk_iter_partition" + - "rocke_wgrad_streamk_host" + - "rocke_future_intrinsic_lowering" - "rocke_tiled_attention_2d_reentrancy" - "rocke_python_import_smoke" - "rocke_golden_static" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py index dc0397223f6f..5833d6b46e8b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py @@ -1239,6 +1239,12 @@ def wgrad_streamk_launch(spec: "WgradConvSpec", *, arch: str = "gfx950"): if spec.streamk == "off": raise ValueError("wgrad_streamk_launch requires streamk != 'off'") + # Same contract as the validator and the C port (rocke_wgrad_conv_streamk_ctas): + # -1 is auto, anything else not > 0 is an error rather than the auto pool. + if spec.streamk_ctas == 0 or spec.streamk_ctas < -1: + raise ValueError( + f"streamk_ctas must be -1 (auto) or > 0 (got {spec.streamk_ctas})" + ) pool = ( spec.streamk_ctas if spec.streamk_ctas > 0 diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py index 970694d35e15..185f6ce20904 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py @@ -39,6 +39,10 @@ _GROUPED = ConvProblem(N=2, Hi=10, Wi=10, C=48, K=120, Y=3, X=3, pH=1, pW=1, groups=3) # Depthwise: the only shape group merging admits, so its stream-K gate is reached. _DW = ConvProblem(N=2, Hi=8, Wi=8, C=16, K=16, Y=3, X=3, pH=1, pW=1, groups=16) +# 17 x 27 = 459 output tiles at 64x64: more than the default pool of either +# stream-K arch (256 CTAs on gfx950, 304 on gfx942 for linear/tree), so an +# auto-pool launch really splits tiles and needs flags and partial slots. +_SPLITS = ConvProblem(N=2, Hi=12, Wi=12, C=192, K=1056, Y=3, X=3, pH=1, pW=1) def _spec(problem=_DENSE, *, dtype_d="fp16", edge=32, **kw) -> WgradConvSpec: @@ -218,8 +222,7 @@ def test_cpp_rejects_with_the_same_reason(self): self.assertIn(why, str(cm.exception)) # (label, problem, arch, spec overrides). Each rejected config trips one - # gate only, so the engines' differing check order cannot pick a different - # reason. + # gate only; _MULTI below covers configs that trip several. _AGREEMENT = [ ("ok_dp_sk_linear", _DENSE, "gfx950", {}), ( @@ -270,12 +273,55 @@ def test_cpp_validity_agrees(self): self.assertEqual(label.startswith("ok_"), want[0], want[1]) got = eng.conv_wgrad_is_valid_spec(_spec_dict(spec), arch) self.assertEqual(tuple(got), want) - # A wave64 spec on gfx1250 trips Python's wave-size check before the - # stream-K arch gate, which the C port has no counterpart of; the - # reasons differ but both engines reject it. - spec = _spec() - self.assertFalse(is_valid_wgrad_spec(spec, "gfx1250")[0]) - self.assertFalse(eng.conv_wgrad_is_valid_spec(_spec_dict(spec), "gfx1250")[0]) + + # Configs that fail more than one gate: the C port runs the gates in + # Python's order, so both report the first one Python reaches. + _MULTI = [ + ("ctas_zero+tile_m", _DENSE, "gfx950", dict(streamk_ctas=0, tile_m=96)), + ( + "block_4096+ctas_zero", + _DENSE, + "gfx950", + dict(warp_m=8, warp_n=8, tile_m=256, tile_n=256, streamk_ctas=0), + ), + ("unknown_arch", _DENSE, "gfx9999", {}), + ("unknown_arch+tile_m", _DENSE, "gfx9999", dict(tile_m=96)), + ("wave32_on_gfx950+ctas", _DENSE, "gfx950", dict(wave_size=32, streamk_ctas=-5)), + ("wave64_on_gfx1250", _DENSE, "gfx1250", {}), + ("async_dma", _DENSE, "gfx950", dict(async_dma=True)), + ("async_dma_off", _DENSE, "gfx950", dict(async_dma=True, streamk="off")), + ("vec_c+ctas", _DENSE, "gfx950", dict(vector_size_c=4, streamk_ctas=-5)), + ("vec_c_off", _DENSE, "gfx950", dict(vector_size_c=4, streamk="off")), + ("group_merge_degree", _DW, "gfx950", dict(group_merge=3)), + ("group_merge+tile_m", _DW, "gfx950", dict(group_merge=2, tile_m=96)), + ("gfx90a+split_k", _DENSE, "gfx90a", dict(split_k=4)), + ] + + def test_cpp_validity_agrees_multi_failure(self): + eng = _engine() + for label, problem, arch, kw in self._MULTI: + spec = _spec(problem, **kw) + want = is_valid_wgrad_spec(spec, arch) + with self.subTest(label): + self.assertFalse(want[0]) + got = eng.conv_wgrad_is_valid_spec(_spec_dict(spec), arch) + self.assertEqual(tuple(got), want) + + def test_cpp_rejects_grouped_non_streamk(self): + # Intentional divergence: grouped wgrad is only ported on the stream-K + # path, so the C port rejects a grouped spec Python accepts -- but only + # after every gate the two share. + eng = _engine() + spec = _spec(_GROUPED, streamk="off") + self.assertEqual(is_valid_wgrad_spec(spec, "gfx950"), (True, "ok")) + ok, why = eng.conv_wgrad_is_valid_spec(_spec_dict(spec), "gfx950") + self.assertFalse(ok) + self.assertIn("grouped convolution (groups=3 > 1)", why) + bad = _spec(_GROUPED, streamk="off", tile_m=96) + self.assertEqual( + tuple(eng.conv_wgrad_is_valid_spec(_spec_dict(bad), "gfx950")), + is_valid_wgrad_spec(bad, "gfx950"), + ) def test_cpp_rejects_group_merge_with_the_same_reason(self): eng = _engine() @@ -351,6 +397,27 @@ def test_workspace_layout(self): ws = _spec(streamk_reduction="workspace") self.assertEqual(wgrad_streamk_workspace_nbytes(ws), 8 * 96 * 288 * 4) + def test_launch_rejects_bad_pool(self): + from kernels.common.conv_implicit_gemm_wgrad import wgrad_streamk_launch + + for ctas in (0, -2, -3): + with self.subTest(ctas=ctas), self.assertRaisesRegex( + ValueError, "streamk_ctas must be -1" + ): + wgrad_streamk_launch(_spec(streamk_ctas=ctas)) + self.assertEqual(wgrad_streamk_launch(_spec(streamk_ctas=-1)).pool, 256) + self.assertEqual( + wgrad_streamk_launch(_spec(streamk_ctas=-1), arch="gfx942").pool, 304 + ) + + def test_cpp_host_rejects_bad_pool(self): + eng = _engine() + for ctas in (0, -3): + with self.subTest(ctas=ctas), self.assertRaises(ValueError): + eng.conv_wgrad_streamk_host(_spec_dict(_spec(streamk_ctas=ctas))) + with self.assertRaises(ValueError): + eng.conv_wgrad_streamk_host(_spec_dict(_spec(streamk="off"))) + def test_signature_tracks_reduction(self): names = lambda s: [p["name"] for p in wgrad_streamk_signature(s)] # noqa: E731 plain = [n for n, _ in conv_arg_names(direction="wgrad")] @@ -396,9 +463,15 @@ def test_cpp_default_pool_agrees(self): from kernels.common.conv_implicit_gemm_wgrad import wgrad_streamk_default_pool from rocke.helpers.split_k import _ARCH_NUM_CUS + from rocke.core.arch.target import known_arches + eng = _engine() - # Every arch of the Python CU table, plus one it does not list. - for arch in [*_ARCH_NUM_CUS, "gfx90a"]: + # Both directions: every arch of the Python CU table, and every arch + # rocke knows of plus a few it does not -- a key only the C table has + # would make the C pool differ from Python's 256-CU fallback there. + arches = {*_ARCH_NUM_CUS, *known_arches(), "gfx90a", "gfx908", "gfx1100"} + arches |= {"gfx1200", "gfx9999"} + for arch in sorted(arches): for reduction in ("linear", "tree", "atomic", "workspace"): with self.subTest(arch=arch, r=reduction): self.assertEqual( @@ -410,8 +483,9 @@ def test_cpp_host_plan_agrees_auto_pool(self): # streamk_ctas=-1 (the spec default) resolves to the arch's CU-based # pool in both engines. eng = _engine() + split = 0 for arch, edge in (("gfx950", 32), ("gfx942", 16)): - for problem in (_DENSE, _GROUPED): + for problem in (_DENSE, _GROUPED, _SPLITS): for mode in ("dp_sk", "persistent"): for reduction, dtype_d in ( ("linear", "fp16"), @@ -443,6 +517,15 @@ def test_cpp_host_plan_agrees_auto_pool(self): got["workspace_bytes"], wgrad_streamk_workspace_nbytes(spec, arch=arch), ) + if problem is _SPLITS and reduction in ("linear", "tree"): + # 459 tiles over the default pool: every CTA is a + # stream-K CTA, so both regions are non-empty. + pool = 256 if arch == "gfx950" else 304 + self.assertEqual(got["grid"][0], pool) + self.assertEqual(flags, pool * 4 * 4) + self.assertEqual(partials, pool * 64 * 64 * 4) + split += 1 + self.assertEqual(split, 2 * 2 * 2) # Without an arch both engines default to gfx950. spec = _spec(streamk="persistent", streamk_ctas=-1) got = eng.conv_wgrad_streamk_host(_spec_dict(spec)) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk_gpu.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk_gpu.py index bed52a41e7f2..25e70741d8d7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk_gpu.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk_gpu.py @@ -11,23 +11,37 @@ The shapes are chosen to exercise the partitioner, not the GEMM: several stream-K CTAs per tile (multi-round linear/tree fixups), a data-parallel plus stream-K mix, a grouped problem whose last M tile is partial, and a small CTA -pool so the stream-K path is taken at these sizes. The linear and tree fixups -are also run repeatedly and must be bit-identical across launches -- the only -check that catches a missing release/acquire on the partial hand-off. Before -every launch the partial slots and dW are poisoned with NaN bytes (only the -flags are zeroed, as the contract requires), so a consumer that reads a slot -before its producer's release lands sees NaN rather than the previous -launch's bit-identical partial. - -Requires a gfx942/gfx950 GPU and torch; skips cleanly otherwise. Host-side -coverage lives in ``test_conv_wgrad_streamk.py``. +pool so the stream-K path is taken at these sizes. A shape large enough to +split under the default (per-CU) pool runs every reduction too. + +The linear and tree fixups are also run repeatedly and must be bit-identical +across launches. Before every launch the partial slots and dW are poisoned +with NaN bytes (only the flags are zeroed, as the contract requires), so a +consumer that reads a slot before its producer's partial is visible sees NaN +rather than the previous launch's bit-identical partial. That catches a +missing *release* on the producer side. It does not catch a missing +*acquire*: the consumer's partial loads are data- and control-dependent on +the spin exit and no stale cache line exists to observe, so a relaxed spin +still reads the right bytes on this hardware. The acquire side is pinned by +codegen instead -- ``TestConvWgradStreamKFlagOrderingIR`` checks the lowered +``.ll`` (host-only), ``TestConvWgradStreamKFlagOrderingIsa`` the gfx942/gfx950 +ISA (COMGR + ``llvm-objdump``, no GPU; skips when either is missing). + +The numeric cases require a gfx942/gfx950 GPU and torch and skip cleanly +otherwise. Host-side coverage lives in ``test_conv_wgrad_streamk.py``. """ from __future__ import annotations import ctypes import importlib.util +import os +import re +import shutil +import subprocess +import tempfile import unittest +from pathlib import Path from rocke.runtime.hip_module import get_device_arch @@ -301,8 +315,20 @@ def test_modes_and_reductions(self): def test_default_pool(self): # streamk_ctas=-1: one CTA per CU for linear/tree, the occupancy-sized # pool (several CTAs per CU) for the reductions that never wait. - for name in ("dp_plus_sk", "multi_partner"): - problem, _ = _problems()[name] + # dp_plus_sk / multi_partner have far fewer tiles than that pool, so + # they run all-DP; "splits" (17 x 27 = 459 tiles) is the one that + # actually takes a linear/tree/workspace/atomic fixup at the default + # pool (on gfx950: 256 linear/tree CTAs, 1024 atomic/workspace CTAs). + from kernels.common._conv_implicit_gemm_common import ConvProblem + from kernels.common.conv_implicit_gemm_wgrad import wgrad_streamk_partition + + shapes = { + name: _problems()[name][0] for name in ("dp_plus_sk", "multi_partner") + } + shapes["splits"] = ConvProblem( + N=2, Hi=12, Wi=12, C=192, K=1056, Y=3, X=3, pH=1, pW=1 + ) + for name, problem in shapes.items(): for mode in ("dp_sk", "persistent"): for reduction, dtype_d in _REDUCTIONS: spec = _spec( @@ -313,6 +339,10 @@ def test_default_pool(self): ctas=-1, ) with self.subTest(shape=name, mode=mode, reduction=reduction): + if name == "splits": + part = wgrad_streamk_partition(spec, arch=GPU_ARCH) + self.assertGreater(part.sk_ctas, 0, part) + self.assertGreater(part.sk_tiles, 0, part) _check(self, spec) @@ -401,7 +431,12 @@ def test_one_binary_many_shapes(self): @unittest.skipUnless(not _SKIP_REASON, _SKIP_REASON or "no GPU") class TestConvWgradStreamKDeterminism(unittest.TestCase): - """Linear and tree fold partials in a fixed order: dW is bit-exact.""" + """Linear and tree fold partials in a fixed order: dW is bit-exact. + + With the partial slots NaN-poisoned this also catches a missing release on + the partial hand-off, but not a missing acquire -- see the module + docstring and ``TestConvWgradStreamKFlagOrderingIR``/``...Isa``. + """ def test_repeat_launches_are_bit_identical(self): import torch @@ -421,5 +456,239 @@ def test_repeat_launches_are_bit_identical(self): ) +# -------------------------------------------------------------------------- +# Flag hand-off ordering, pinned in codegen. +# +# The repeat-launch test above catches a missing release but cannot catch a +# missing acquire (see the module docstring), so the acquire -- and, for +# symmetry, the release -- is asserted on the code a real linear and tree +# kernel lowers to. Host-only: no GPU is needed for either level. + +_FLAG_ARCHES = ("gfx950", "gfx942") +_FLAG_KERNELS = ( + ("dp_sk", "linear"), + ("dp_sk", "tree"), + ("persistent", "linear"), + ("persistent", "tree"), +) + + +def _flag_comgr_input(arch: str, mode: str, reduction: str): + """The lowered ``.ll`` + COMGR target of a real stream-K fixup kernel. + + Always the Python engine (the C++ engine is byte-identical by contract and + its extension may not be importable), and lowered with ``arch`` as the + target rather than the detected GPU, so the gfx942 code is checked on a + gfx950 host and vice versa. + """ + from rocke.helpers.compile import lower_kernel_for_comgr + + from kernels.common._conv_implicit_gemm_common import ConvDataSpec, ConvProblem + from kernels.common.conv_implicit_gemm_wgrad import ( + WgradConvSpec, + build_implicit_gemm_conv_wgrad, + ) + + edge = 32 if arch == "gfx950" else 16 + spec = WgradConvSpec( + problem=ConvProblem(N=8, Hi=16, Wi=16, C=16, K=96, Y=3, X=3, pH=1, pW=1), + data=ConvDataSpec(dtype_a="fp16", dtype_b="fp16", dtype_d="fp16"), + tile_m=64, + tile_n=64, + tile_k=64, + warp_m=2, + warp_n=2, + warp_tile_m=edge, + warp_tile_n=edge, + warp_tile_k=16, + streamk=mode, + streamk_reduction=reduction, + streamk_ctas=16, + ) + kernel = build_implicit_gemm_conv_wgrad(spec, arch=arch) + return lower_kernel_for_comgr(kernel, arch=arch, backend="python") + + +def _ll_blocks(ll: str) -> dict: + """``{label: [instruction lines]}`` of every basic block in ``ll``.""" + blocks, cur = {}, None + for line in ll.splitlines(): + m = re.match(r"^([A-Za-z0-9_.$-]+):", line) + if m: + cur = blocks.setdefault(m.group(1), []) + elif cur is not None and line.startswith(" "): + cur.append(line.strip()) + return blocks + + +def _gep_base(ll: str, ptr: str) -> str: + """Base pointer of the ``getelementptr`` that defines ``ptr``.""" + m = re.search( + rf"^\s*{re.escape(ptr)} = getelementptr\b[^,]*, ptr addrspace\(1\) (%[\w.]+),", + ll, + re.MULTILINE, + ) + return m.group(1) if m else "" + + +def _assert_flag_ll(test, ll: str) -> None: + """Acquire on every flag spin, fence + release on every flag publish.""" + blocks = _ll_blocks(ll) + load_re = re.compile( + r"^(%[\w.]+) = load atomic i32, ptr addrspace\(1\) (%[\w.]+) (.*), align 4$" + ) + store_re = re.compile( + r"^store atomic i32 1, ptr addrspace\(1\) (%[\w.]+) (.*), align 4$" + ) + spins = publishes = 0 + for label, insts in blocks.items(): + for i, inst in enumerate(insts): + m = load_re.match(inst) + if m: + spins += 1 + test.assertEqual(_gep_base(ll, m.group(2)), "%sk_flags", inst) + test.assertEqual(m.group(3), 'syncscope("agent") acquire', inst) + # A spin: the block re-branches to itself until the flag hits. + test.assertRegex( + insts[-1], rf"^br i1 .*label %{re.escape(label)}\b", label + ) + continue + m = store_re.match(inst) + if m: + publishes += 1 + test.assertEqual(_gep_base(ll, m.group(1)), "%sk_flags", inst) + test.assertEqual(m.group(2), 'syncscope("agent") release', inst) + # The nearest preceding memory op is the agent release fence, + # and every partial store of the block is ordered before it. + mem = [ + j + for j in range(i) + if re.match( + r"^(fence|store|%[\w.]+ = (load|atomicrmw))\b", insts[j] + ) + ] + test.assertTrue(mem, f"{label}: flag store with no fence before it") + test.assertEqual( + insts[mem[-1]], 'fence syncscope("agent") release', label + ) + test.assertTrue( + any(re.match(r"^store float ", insts[j]) for j in mem[:-1]), + f"{label}: no partial store before the release fence", + ) + test.assertGreater(spins, 0, "no flag-wait spin loop in the kernel") + test.assertGreater(publishes, 0, "no flag publish in the kernel") + + +def _objdump() -> str | None: + for candidate in ( + os.environ.get("LLVM_OBJDUMP"), + os.path.join(os.environ.get("ROCM_PATH", "/opt/rocm"), "llvm/bin/llvm-objdump"), + shutil.which("llvm-objdump"), + ): + if candidate and os.path.isfile(candidate): + return candidate + return None + + +def _isa_lines(objdump: str, comgr_input, arch: str) -> list: + from rocke.runtime.comgr import build_hsaco_from_llvm_ir + + hsaco, _ = build_hsaco_from_llvm_ir( + comgr_input.llvm_text, + isa=comgr_input.isa, + options=list(comgr_input.options), + ) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "streamk.hsaco" + path.write_bytes(hsaco) + text = subprocess.run( + [objdump, "-d", f"--mcpu={arch}", str(path)], + check=True, + capture_output=True, + text=True, + ).stdout + lines = (ln.split("//")[0].strip() for ln in text.splitlines()) + return [ln for ln in lines if ln and not ln.endswith(">:")] + + +_SC1_LOAD = re.compile(r"^global_load_dword\b.*\bsc1\b") +_SC1_STORE = re.compile(r"^global_store_dword\b.*\bsc1\b") + + +def _assert_flag_isa(test, isa: list) -> None: + """gfx94x agent-scope acquire/release: ``buffer_inv sc1`` after each spin + load, ``buffer_wbl2 sc1`` + ``s_waitcnt vmcnt(0)`` before each flag store. + """ + spins = [i for i, ln in enumerate(isa) if _SC1_LOAD.match(ln)] + stores = [i for i, ln in enumerate(isa) if _SC1_STORE.match(ln)] + test.assertTrue(spins, "no sc1 flag load in the ISA") + test.assertTrue(stores, "no sc1 flag store in the ISA") + for i in spins: + # Up to the spin's back-branch: wait for the load, then invalidate so + # the partial loads after the exit cannot hit a stale line. + end = next( + (j for j in range(i + 1, len(isa)) if isa[j].startswith("s_cbranch")), + len(isa), + ) + body = isa[i + 1 : end] + test.assertIn("buffer_inv sc1", body, f"spin at {i}: no acquire invalidate") + test.assertLess( + body.index("s_waitcnt vmcnt(0)"), + body.index("buffer_inv sc1"), + f"spin at {i}: buffer_inv before the flag load completed", + ) + for i in stores: + # Back to the last partial store: write back L2 and drain every + # outstanding store before the flag goes out. + start = max( + ( + j + for j in range(i) + if isa[j].startswith(("global_store", "buffer_store", "global_atomic")) + ), + default=-1, + ) + window = isa[start + 1 : i] + test.assertIn("buffer_wbl2 sc1", window, f"flag store at {i}: no writeback") + test.assertIn( + "s_waitcnt vmcnt(0)", window, f"flag store at {i}: stores not drained" + ) + + +class TestConvWgradStreamKFlagOrderingIR(unittest.TestCase): + """The lowered ``.ll`` of real linear/tree kernels orders the hand-off.""" + + def test_acquire_spin_and_release_publish(self): + for arch in _FLAG_ARCHES: + for mode, reduction in _FLAG_KERNELS: + with self.subTest(arch=arch, mode=mode, reduction=reduction): + ll = _flag_comgr_input(arch, mode, reduction).llvm_text + _assert_flag_ll(self, ll) + + +class TestConvWgradStreamKFlagOrderingIsa(unittest.TestCase): + """The same kernels' gfx942/gfx950 ISA keeps the cache maintenance. + + Offline compile (COMGR) + ``llvm-objdump``; no GPU. Skips when either + tool is unavailable. + """ + + def test_acquire_invalidate_and_release_writeback(self): + from rocke.runtime import comgr + + objdump = _objdump() + if objdump is None: + self.skipTest("llvm-objdump not found (LLVM_OBJDUMP / ROCM_PATH / PATH)") + try: + comgr._resolve_lib() + except (comgr.ComgrError, OSError) as e: + self.skipTest(f"COMGR not loadable: {str(e)[:200]}") + for arch in _FLAG_ARCHES: + for mode, reduction in _FLAG_KERNELS: + with self.subTest(arch=arch, mode=mode, reduction=reduction): + c = _flag_comgr_input(arch, mode, reduction) + _assert_flag_isa(self, _isa_lines(objdump, c, arch)) + + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp index 4d91ddb8198b..46da5f6ec152 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp @@ -2273,6 +2273,20 @@ rocke_implicit_gemm_conv_wgrad_spec_t conv_wgrad_build_spec(const py::dict& d, s.chiplet_wgm = dict_int(d, "chiplet_wgm", s.chiplet_wgm); s.chiplet_num_xcds = dict_int(d, "chiplet_num_xcds", s.chiplet_num_xcds); s.chiplet_chunk_size = dict_int(d, "chiplet_chunk_size", s.chiplet_chunk_size); + /* Optional ints: absent or None keeps Python's None (has_* = false). + * Only the fields the C builder or validator honours are read: the C wgrad + * builder does not apply lds_k_pad or vector_size_a/b, so reading them + * would tag the kernel name (e.g. "pad") without changing its body. */ + if(d.contains("waves_per_eu") && !d["waves_per_eu"].is_none()) + { + s.has_waves_per_eu = true; + s.waves_per_eu = d["waves_per_eu"].cast(); + } + if(d.contains("vector_size_c") && !d["vector_size_c"].is_none()) + { + s.has_vector_size_c = true; + s.vector_size_c = d["vector_size_c"].cast(); + } { std::string v; if(dict_str(d, "name", v)) @@ -3902,24 +3916,26 @@ PYBIND11_MODULE(rocke_engine, m) [](const py::dict& d, const std::string& arch) -> py::dict { std::deque store; rocke_implicit_gemm_conv_wgrad_spec_t s = conv_wgrad_build_spec(d, store); - /* Resolve the auto pool (-1) for `arch` up front, as Python - * wgrad_streamk_launch does, so every helper below sees it. */ - const int pool = rocke_wgrad_conv_streamk_ctas(&s, arch_or_default(arch)); - if(pool > 0) - s.streamk_ctas = pool; + /* Every helper resolves an auto pool (-1) for `arch` itself, as the + * Python wgrad_streamk_*(spec, arch=arch) twins do. Errors are + * ValueError, matching Python. */ + const char* a = arch_or_default(arch); int grid[3] = {0, 0, 0}; - if(rocke_wgrad_conv_streamk_grid(&s, grid) != ROCKE_OK) - throw std::runtime_error( + if(rocke_wgrad_conv_streamk_grid(&s, a, grid) != ROCKE_OK) + throw py::value_error( "rocke_engine.conv_wgrad_streamk_host: needs streamk != 'off' and " "streamk_ctas == -1 or > 0"); size_t flags_bytes = 0; size_t partials_bytes = 0; - rocke_wgrad_conv_streamk_workspace_layout(&s, &flags_bytes, &partials_bytes); + if(rocke_wgrad_conv_streamk_workspace_layout(&s, a, &flags_bytes, &partials_bytes) + != ROCKE_OK) + throw py::value_error( + "rocke_engine.conv_wgrad_streamk_host: invalid stream-K workspace layout"); py::dict out; out["grid"] = py::make_tuple(grid[0], grid[1], grid[2]); out["flags_bytes"] = flags_bytes; out["partials_bytes"] = partials_bytes; - out["workspace_bytes"] = rocke_wgrad_conv_streamk_workspace_bytes(&s); + out["workspace_bytes"] = rocke_wgrad_conv_streamk_workspace_bytes(&s, a); out["is_deterministic"] = rocke_wgrad_conv_spec_is_deterministic(&s); return out; }, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h index cfabdc546498..f28e23a747a9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_wgrad.h @@ -185,7 +185,9 @@ typedef struct rocke_implicit_gemm_conv_wgrad_spec * streamk_ctas CTA pool the stream-K remainder is spread over; * a launch parameter, not part of the kernel, so the * builder accepts -1 (auto) (unlike split_k=-1); the - * host helpers need it resolved to > 0. + * rocke_wgrad_conv_streamk_* host helpers resolve -1 + * to the per-arch default pool of the arch they are + * given (rocke_wgrad_conv_streamk_default_pool). * "off" keeps every existing config byte-identical. Stream-K folds the * conv groups into GEMM-M, so it is the one path on which this port * accepts groups > 1. */ @@ -262,26 +264,38 @@ int rocke_wgrad_conv_streamk_default_pool(const char* reduction, const char* arc * count. A NULL arch is gfx950. */ int rocke_wgrad_conv_streamk_ctas(const rocke_implicit_gemm_conv_wgrad_spec_t* s, const char* arch); -/* Python wgrad_streamk_partition(spec): the iteration partition the kernel and - * its launch grid share. Requires streamk != "off" and streamk_ctas == -1 - * (auto, resolved for gfx950 like Python's default arch; resolve it with - * rocke_wgrad_conv_streamk_ctas for another arch) or > 0; otherwise +/* The helpers below take the target arch because an auto pool + * (streamk_ctas == -1) resolves to a per-arch CU count; every one of them + * resolves it the same way, so the grid, the workspace and the kernarg + * partition block of one launch always agree. A NULL arch is gfx950, the + * default of the Python twins (wgrad_streamk_*(spec, arch="gfx950")). */ + +/* Python wgrad_streamk_partition(spec, arch=arch): the iteration partition the + * kernel and its launch grid share. Requires a non-NULL spec with + * streamk != "off" and streamk_ctas == -1 (auto) or > 0; otherwise * *out_status (if non-NULL) is ROCKE_ERR_VALUE and the result is zeroed. */ -rocke_streamk_iter_partition_t - rocke_wgrad_conv_streamk_partition(const rocke_implicit_gemm_conv_wgrad_spec_t* s, - rocke_status_t* out_status); +rocke_streamk_iter_partition_t rocke_wgrad_conv_streamk_partition( + const rocke_implicit_gemm_conv_wgrad_spec_t* s, const char* arch, rocke_status_t* out_status); -/* Python wgrad_streamk_grid(spec): one-dimensional launch grid. */ +/* Python wgrad_streamk_grid(spec, arch=arch): one-dimensional launch grid. + * ROCKE_ERR_VALUE when the partition is invalid or out is NULL. */ rocke_status_t rocke_wgrad_conv_streamk_grid(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + const char* arch, int out[3]); -/* Python wgrad_streamk_workspace_layout(spec): (flags_bytes, partials_bytes) of - * a linear/tree kernel; both 0 for the atomic and workspace reductions. */ -rocke_status_t rocke_wgrad_conv_streamk_workspace_layout( - const rocke_implicit_gemm_conv_wgrad_spec_t* s, size_t* flags_bytes, size_t* partials_bytes); - -/* Python wgrad_streamk_workspace_nbytes(spec). */ -size_t rocke_wgrad_conv_streamk_workspace_bytes(const rocke_implicit_gemm_conv_wgrad_spec_t* s); +/* Python wgrad_streamk_workspace_layout(spec, arch=arch): (flags_bytes, + * partials_bytes) of a linear/tree kernel; both 0 for the atomic and workspace + * reductions. ROCKE_ERR_VALUE for a NULL argument or an invalid partition. */ +rocke_status_t + rocke_wgrad_conv_streamk_workspace_layout(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + const char* arch, + size_t* flags_bytes, + size_t* partials_bytes); + +/* Python wgrad_streamk_workspace_nbytes(spec, arch=arch); 0 for a NULL spec, + * streamk == "off", or an invalid partition. */ +size_t rocke_wgrad_conv_streamk_workspace_bytes(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + const char* arch); /* 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); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp index 2cdba4ee88a7..038841808564 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp @@ -36,6 +36,7 @@ #include /* strcmp, memset, memcpy */ #include "rocke/error_boundary.hpp" +#include "rocke/helper_rocke.core.arch.h" /* rocke_archtarget_* */ #include "rocke/helper_rocke.helpers.spec.h" #include "rocke/helper_rocke.helpers.transforms.h" #include "rocke/instance_conv_abi.h" @@ -209,14 +210,13 @@ int rocke_wgrad_conv_streamk_ctas(const rocke_implicit_gemm_conv_wgrad_spec_t* s return 0; } -/* Python wgrad_streamk_partition (arch defaults to gfx950, as in Python). */ -rocke_streamk_iter_partition_t - rocke_wgrad_conv_streamk_partition(const rocke_implicit_gemm_conv_wgrad_spec_t* s, - rocke_status_t* out_status) +/* Python wgrad_streamk_partition (a NULL arch is gfx950, Python's default). */ +rocke_streamk_iter_partition_t rocke_wgrad_conv_streamk_partition( + const rocke_implicit_gemm_conv_wgrad_spec_t* s, const char* arch, rocke_status_t* out_status) { rocke_streamk_iter_partition_t zero; memset(&zero, 0, sizeof(zero)); - const int pool = rocke_wgrad_conv_streamk_ctas(s, NULL); + const int pool = rocke_wgrad_conv_streamk_ctas(s, arch); if(pool <= 0) { if(out_status) @@ -236,10 +236,11 @@ rocke_streamk_iter_partition_t } rocke_status_t rocke_wgrad_conv_streamk_grid(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + const char* arch, int out[3]) { rocke_status_t st = ROCKE_OK; - rocke_streamk_iter_partition_t part = rocke_wgrad_conv_streamk_partition(s, &st); + rocke_streamk_iter_partition_t part = rocke_wgrad_conv_streamk_partition(s, arch, &st); if(st != ROCKE_OK || out == NULL) return ROCKE_ERR_VALUE; out[0] = rocke_streamk_iter_plan(&part).grid_size; @@ -248,10 +249,13 @@ rocke_status_t rocke_wgrad_conv_streamk_grid(const rocke_implicit_gemm_conv_wgra return ROCKE_OK; } -rocke_status_t rocke_wgrad_conv_streamk_workspace_layout( - const rocke_implicit_gemm_conv_wgrad_spec_t* s, size_t* flags_bytes, size_t* partials_bytes) +rocke_status_t + rocke_wgrad_conv_streamk_workspace_layout(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + const char* arch, + size_t* flags_bytes, + size_t* partials_bytes) { - if(flags_bytes == NULL || partials_bytes == NULL) + if(s == NULL || flags_bytes == NULL || partials_bytes == NULL) return ROCKE_ERR_VALUE; *flags_bytes = 0; *partials_bytes = 0; @@ -259,7 +263,7 @@ rocke_status_t rocke_wgrad_conv_streamk_workspace_layout( if(strcmp(red, "linear") != 0 && strcmp(red, "tree") != 0) return ROCKE_OK; rocke_status_t st = ROCKE_OK; - rocke_streamk_iter_partition_t part = rocke_wgrad_conv_streamk_partition(s, &st); + rocke_streamk_iter_partition_t part = rocke_wgrad_conv_streamk_partition(s, arch, &st); if(st != ROCKE_OK) return st; const rocke_streamk_iter_plan_t plan = rocke_streamk_iter_plan(&part); @@ -270,7 +274,8 @@ rocke_status_t rocke_wgrad_conv_streamk_workspace_layout( return ROCKE_OK; } -size_t rocke_wgrad_conv_streamk_workspace_bytes(const rocke_implicit_gemm_conv_wgrad_spec_t* s) +size_t rocke_wgrad_conv_streamk_workspace_bytes(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + const char* arch) { if(s == NULL || !wgrad_streamk_on(s)) return 0; @@ -285,7 +290,8 @@ size_t rocke_wgrad_conv_streamk_workspace_bytes(const rocke_implicit_gemm_conv_w } size_t flags_bytes = 0; size_t partials_bytes = 0; - if(rocke_wgrad_conv_streamk_workspace_layout(s, &flags_bytes, &partials_bytes) != ROCKE_OK) + if(rocke_wgrad_conv_streamk_workspace_layout(s, arch, &flags_bytes, &partials_bytes) + != ROCKE_OK) return 0; return flags_bytes + partials_bytes; } @@ -535,6 +541,56 @@ static bool wgrad_streamk_gates(const rocke_implicit_gemm_conv_wgrad_spec_t* s, #undef SK_REJECT } +/* Python wgrad_group_merge_available, with byte-identical reason strings. + * Python checks it right after the arch lookup, before any other gate. */ +static bool wgrad_group_merge_available(const rocke_implicit_gemm_conv_wgrad_spec_t* s, + char* reason, + size_t reason_cap) +{ +#define GM_REJECT(...) \ + do \ + { \ + if(reason && reason_cap) \ + snprintf(reason, reason_cap, __VA_ARGS__); \ + return false; \ + } while(0) + const int gm = s->group_merge; + if(gm == 1) + return true; + /* Python _GROUP_MERGE_DEGREES */ + if(gm != 2 && gm != 4 && gm != 8 && gm != 16 && gm != 32 && gm != 64) + GM_REJECT("group_merge must be 1 or one of (2, 4, 8, 16, 32, 64); got %d", gm); + const rocke_conv_problem_t* p = &s->problem; + if(p->groups % gm != 0) + GM_REJECT("group_merge %d does not divide groups %d: a workgroup would own a " + "partial group", + gm, + p->groups); + const int cpg = rocke_conv_problem_cpg(p); + const int kpg = rocke_conv_problem_kpg(p); + if(cpg != 1 || kpg != 1) + GM_REJECT("group_merge is implemented for depthwise only (cpg == kpg == 1); got " + "cpg=%d kpg=%d. The merged index math is correct for wider groups but is " + "not covered, so it is gated off.", + cpg, + kpg); + if(s->wave_size != 64) + GM_REJECT("group_merge is MFMA-only (wave_size 64); got wave_size=%d", s->wave_size); + if(s->split_k > 1 && !s->two_stage) + GM_REJECT("group_merge with split_k=%d needs the two-stage path (two_stage=True); " + "the packed-atomic split-K epilogue cannot drop off-diagonal group pairs", + s->split_k); + /* grid_M / grid_N: the merged GEMM extents, kpg*Gm and Z*Y*X*cpg*Gm. */ + const int grid_M = kpg * gm; + const int grid_N = (p->is_3d ? p->Z : 1) * p->Y * p->X * cpg * gm; + if(grid_M > s->tile_m) + GM_REJECT("merged GEMM-M %d (kpg*%d) exceeds tile_m %d", grid_M, gm, s->tile_m); + if(grid_N > s->tile_n) + GM_REJECT("merged GEMM-N %d (Z*Y*X*cpg*%d) exceeds tile_n %d", grid_N, gm, s->tile_n); + return true; +#undef GM_REJECT +} + bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv_wgrad_spec_t* s, const char* arch, char* reason, @@ -556,60 +612,109 @@ bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv if(arch == NULL) arch = "gfx950"; - /* groups > 1 is only ported on the stream-K path, which folds the group - * into its linear tile index; the grid-per-group path stays Python-only. */ - const bool sk_on = wgrad_streamk_on(s); - if(s->problem.groups != 1 && !sk_on) + /* The gates below run in Python is_valid_wgrad_spec's order -- arch lookup, + * group_merge, geometry, block size, vec_c, wave size, stream-K -- so a + * spec that fails several of them is rejected with the same reason by + * both engines. */ + + /* Python: ArchTarget.from_gfx(arch) -> KeyError. The forward validator + * checks the arch first and reproduces str(KeyError) verbatim, so an + * unknown arch is handed to it for the reason text. */ + const rocke_archtarget_t* target = rocke_archtarget_from_gfx(arch); + if(target == NULL) + { + rocke_implicit_gemm_conv_spec_t probe = rocke_implicit_gemm_conv_spec_default(); + return rocke_implicit_gemm_conv_is_valid_spec(&probe, arch, reason, reason_cap); + } + + if(!wgrad_group_merge_available(s, reason, reason_cap)) + return false; + + /* geometry */ + if(s->tile_m % (s->warp_m * s->warp_tile_m)) { if(reason && reason_cap) - snprintf(reason, - reason_cap, - "grouped convolution (groups=%d > 1) is not supported for wgrad", - s->problem.groups); + snprintf(reason, reason_cap, "tile_m not divisible by warp_m * warp_tile_m"); return false; } - if(!wgrad_streamk_gates(s, arch, reason, reason_cap)) - return false; - if(s->group_merge != 1 && !sk_on) + if(s->tile_n % (s->warp_n * s->warp_tile_n)) { if(reason && reason_cap) - snprintf(reason, - reason_cap, - "group_merge=%d is not supported for wgrad in the C port", - s->group_merge); + snprintf(reason, reason_cap, "tile_n not divisible by warp_n * warp_tile_n"); return false; } - if(s->problem.groups > 1 && rocke_conv_problem_is_pointwise(&s->problem)) + if(s->tile_k % s->warp_tile_k) { if(reason && reason_cap) - snprintf(reason, reason_cap, "grouped pointwise (1x1) wgrad is not yet supported"); + snprintf(reason, reason_cap, "tile_k not divisible by warp_tile_k"); return false; } - - /* geometry */ - if(s->tile_m % (s->warp_m * s->warp_tile_m)) + const int block_size = rocke_wgrad_conv_spec_block_size(s); + const int max_tpb = rocke_archtarget_max_threads_per_block(target); + if(block_size > max_tpb) { if(reason && reason_cap) - snprintf(reason, reason_cap, "tile_m not divisible by warp_m * warp_tile_m"); + snprintf(reason, + reason_cap, + "block_size %d > %d (hardware cap) on %s", + block_size, + max_tpb, + arch); return false; } - if(s->tile_n % (s->warp_n * s->warp_tile_n)) + + /* vector_size_c > 1 needs the cshuffle epilogue. Mirrors Python + * is_valid_wgrad_spec (and the dgrad port's gate). */ + { + const char* epi = s->epilogue ? s->epilogue : "default"; + if(s->has_vector_size_c && s->vector_size_c > 1 && strcmp(epi, "default") == 0) + { + if(reason && reason_cap) + snprintf(reason, + reason_cap, + "default epilogue is not supported with vector size c: %d", + s->vector_size_c); + return false; + } + } + + if(s->wave_size != target->wave_size) { if(reason && reason_cap) - snprintf(reason, reason_cap, "tile_n not divisible by warp_n * warp_tile_n"); + snprintf(reason, + reason_cap, + "spec wave_size %d != %s wave_size %d", + s->wave_size, + arch, + target->wave_size); return false; } - if(s->tile_k % s->warp_tile_k) + + if(!wgrad_streamk_gates(s, arch, reason, reason_cap)) + return false; + + /* C-port-only gates, after every gate Python shares: grouped wgrad + * (groups > 1) and group merging are only ported on the stream-K path, + * which folds the group into its linear tile index; the grid-per-group + * path stays Python-only. This is an intentional divergence -- Python + * accepts these specs, the C port rejects them. */ + const bool sk_on = wgrad_streamk_on(s); + if(s->problem.groups != 1 && !sk_on) { if(reason && reason_cap) - snprintf(reason, reason_cap, "tile_k not divisible by warp_tile_k"); + snprintf(reason, + reason_cap, + "grouped convolution (groups=%d > 1) is not supported for wgrad", + s->problem.groups); return false; } - int block_size = rocke_wgrad_conv_spec_block_size(s); - if(block_size > 1024) + if(s->group_merge != 1 && !sk_on) { if(reason && reason_cap) - snprintf(reason, reason_cap, "block_size %d > 1024", block_size); + snprintf(reason, + reason_cap, + "group_merge=%d is not supported for wgrad in the C port", + s->group_merge); return false; } @@ -921,6 +1026,16 @@ bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv } } + /* Python checks this between the MMA atom and the LDS budget, both of + * which the forward validator below owns; checking it just before is the + * closest order the delegation allows. */ + if(s->problem.groups > 1 && rocke_conv_problem_is_pointwise(&s->problem)) + { + if(reason && reason_cap) + snprintf(reason, reason_cap, "grouped pointwise (1x1) wgrad is not yet supported"); + return false; + } + /* Delegate the MMA-atom + LDS + WMMA gates to the forward validator via an * adapter spec. The forward validator only reads: tile_m/n/k, warp_m/n, * warp_tile_m/n/k, wave_size, pipeline, epilogue, async_dma, unroll_k, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt b/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt index b320a8f8ef4f..66d974f86fdd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt @@ -285,6 +285,27 @@ if(BUILD_TESTING) rocke_streamk_iter_partition) endif() + # Stream-K wgrad host helpers and validator gates (auto CTA pool per arch, + # grid, workspace layout, reject reasons): also invisible to the byte-identity + # gate, and the Python-vs-C pytest needs the pybind module CI does not build, + # so the C port is pinned to Python reference vectors here. + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/instances/wgrad_streamk_host.cpp") + add_executable(rocke_wgrad_streamk_host instances/wgrad_streamk_host.cpp) + target_link_libraries(rocke_wgrad_streamk_host PRIVATE rocke_core) + if(UNIX) + target_link_libraries(rocke_wgrad_streamk_host PRIVATE m) + endif() + add_test(NAME rocke_wgrad_streamk_host COMMAND rocke_wgrad_streamk_host) + set_tests_properties(rocke_wgrad_streamk_host PROPERTIES LABELS "ckc;unit;host") + rocke_cxx_test_installed_name(_rocke_installed_name rocke_wgrad_streamk_host) + set_target_properties(rocke_wgrad_streamk_host + PROPERTIES OUTPUT_NAME "${_rocke_installed_name}") + install(TARGETS rocke_wgrad_streamk_host + RUNTIME DESTINATION ${ROCKE_CXX_TEST_INSTALL_DIR}) + set_property(GLOBAL APPEND PROPERTY ROCKE_INSTALL_TEST_TARGETS + rocke_wgrad_streamk_host) + endif() + # Apply YAML-driven quick/standard/comprehensive/full tier labels to the # tests above. Must run in THIS directory scope: CMake (pre-3.28, # which this project's cmake_minimum_required predates) only allows diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/wgrad_streamk_host.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/wgrad_streamk_host.cpp new file mode 100644 index 000000000000..74a5fd16e02d --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/wgrad_streamk_host.cpp @@ -0,0 +1,404 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* + * tests/instances/wgrad_streamk_host.cpp -- C parity for the stream-K wgrad + * host helpers and validator gates of instance_conv_implicit_gemm_wgrad.h. + * + * The launch grid, workspace layout and auto CTA pool these helpers compute + * never reach the emitted IR, so the byte-identity gate cannot see a C/Python + * divergence in them; it would surface as skipped work or an undersized + * buffer. The Python-vs-C pytest (library/tests/test_conv_wgrad_streamk.py) + * needs the rocke_engine binding, which CI does not build, so the vectors are + * pinned here as well. + * + * Every expected value below was produced by the Python implementation + * (kernels.common.conv_implicit_gemm_wgrad: wgrad_streamk_default_pool, + * wgrad_streamk_launch(...).pool, wgrad_streamk_partition, wgrad_streamk_grid, + * wgrad_streamk_workspace_layout, wgrad_streamk_workspace_nbytes and + * is_valid_wgrad_spec) for the same specs; regenerate them from there if the + * Python reference changes. + */ +#include +#include + +#include "rocke/instance_conv_implicit_gemm_wgrad.h" + +namespace +{ + +int g_failures = 0; + +#define CHECK(cond) \ + do \ + { \ + if(!(cond)) \ + { \ + fprintf(stderr, "FAIL: %s (%s:%d)\n", #cond, __FILE__, __LINE__); \ + ++g_failures; \ + } \ + } while(0) + +/* ConvProblem(N=2, Hi=12, Wi=12, C=32, K=96, Y=3, X=3, pH=1, pW=1): 10 tiles. */ +rocke_conv_problem_t dense() +{ + rocke_conv_problem_t p = rocke_conv_problem_default(2, 12, 12, 32, 96, 3, 3); + p.pH = 1; + p.pW = 1; + return p; +} + +/* ConvProblem(N=2, Hi=12, Wi=12, C=192, K=1056, Y=3, X=3, pH=1, pW=1): + * 17 x 27 = 459 tiles, more than either arch's default pool, so an auto-pool + * launch splits. */ +rocke_conv_problem_t splits() +{ + rocke_conv_problem_t p = rocke_conv_problem_default(2, 12, 12, 192, 1056, 3, 3); + p.pH = 1; + p.pW = 1; + return p; +} + +/* The library test's _spec(): 64x64x32 tiles, 2x2 warps, 32x32x16 atoms, + * dp_sk, linear, streamk_ctas=6. */ +rocke_implicit_gemm_conv_wgrad_spec_t spec(const rocke_conv_problem_t& p) +{ + rocke_implicit_gemm_conv_wgrad_spec_t s = rocke_implicit_gemm_conv_wgrad_spec_default(); + s.problem = p; + s.tile_m = 64; + s.tile_n = 64; + s.tile_k = 32; + s.warp_m = 2; + s.warp_n = 2; + s.warp_tile_m = 32; + s.warp_tile_n = 32; + s.warp_tile_k = 16; + s.streamk = "dp_sk"; + s.streamk_ctas = 6; + return s; +} + +void check_default_pool() +{ + static const struct + { + const char* arch; + int linear, atomic; + } k_pools[] = { + {"gfx950", 256, 1024}, + {"gfx942", 304, 1216}, + {"gfx1151", 64, 256}, + {"gfx90a", 256, 1024}, /* not in the CU table: 256-CU fallback */ + {"gfx9999", 256, 1024}, + {NULL, 256, 1024}, /* NULL arch is gfx950 */ + }; + for(const auto& e : k_pools) + { + CHECK(rocke_wgrad_conv_streamk_default_pool("linear", e.arch) == e.linear); + CHECK(rocke_wgrad_conv_streamk_default_pool("tree", e.arch) == e.linear); + CHECK(rocke_wgrad_conv_streamk_default_pool("atomic", e.arch) == e.atomic); + CHECK(rocke_wgrad_conv_streamk_default_pool("workspace", e.arch) == e.atomic); + } + CHECK(rocke_wgrad_conv_streamk_default_pool(NULL, "gfx942") == 304); /* linear */ +} + +void check_ctas() +{ + rocke_implicit_gemm_conv_wgrad_spec_t s = spec(dense()); + CHECK(rocke_wgrad_conv_streamk_ctas(&s, "gfx942") == 6); + s.streamk_ctas = -1; + CHECK(rocke_wgrad_conv_streamk_ctas(&s, NULL) == 256); + CHECK(rocke_wgrad_conv_streamk_ctas(&s, "gfx950") == 256); + CHECK(rocke_wgrad_conv_streamk_ctas(&s, "gfx942") == 304); + s.streamk_reduction = "workspace"; + CHECK(rocke_wgrad_conv_streamk_ctas(&s, "gfx942") == 1216); + /* Python wgrad_streamk_launch raises for these; C reports no pool. */ + s.streamk_ctas = 0; + CHECK(rocke_wgrad_conv_streamk_ctas(&s, "gfx950") == 0); + s.streamk_ctas = -3; + CHECK(rocke_wgrad_conv_streamk_ctas(&s, "gfx950") == 0); + s.streamk_ctas = 6; + s.streamk = "off"; + CHECK(rocke_wgrad_conv_streamk_ctas(&s, "gfx950") == 0); + CHECK(rocke_wgrad_conv_streamk_ctas(NULL, "gfx950") == 0); +} + +struct PlanVector +{ + const char* arch; + const char* mode; + int ctas; + int sk_tiles, sk_ctas, dp_tiles, grid, ipsc, extra; + size_t flags_bytes, partials_bytes; +}; + +void check_plans() +{ + /* linear reduction on the 459-tile problem. */ + static const PlanVector k_plans[] = { + {"gfx950", "dp_sk", -1, 459, 256, 0, 256, 16, 35, 4096, 4194304}, + {"gfx950", "persistent", -1, 459, 256, 0, 256, 16, 35, 4096, 4194304}, + {"gfx942", "dp_sk", -1, 459, 304, 0, 304, 13, 179, 4864, 4980736}, + {"gfx942", "persistent", -1, 459, 304, 0, 304, 13, 179, 4864, 4980736}, + {"gfx950", "dp_sk", 7, 11, 7, 448, 455, 14, 1, 256, 114688}, + {"gfx942", "persistent", 7, 11, 7, 448, 7, 14, 1, 256, 114688}, + {NULL, "dp_sk", -1, 459, 256, 0, 256, 16, 35, 4096, 4194304}, + }; + for(const PlanVector& v : k_plans) + { + rocke_implicit_gemm_conv_wgrad_spec_t s = spec(splits()); + s.streamk = v.mode; + s.streamk_ctas = v.ctas; + rocke_status_t st = ROCKE_OK; + const rocke_streamk_iter_partition_t part + = rocke_wgrad_conv_streamk_partition(&s, v.arch, &st); + CHECK(st == ROCKE_OK); + CHECK(part.m_tiles == 17); + CHECK(part.n_tiles == 27); + CHECK(part.iters_per_tile == 9); + const rocke_streamk_iter_plan_t r = rocke_streamk_iter_plan(&part); + CHECK(r.sk_tiles == v.sk_tiles); + CHECK(r.sk_ctas == v.sk_ctas); + CHECK(r.dp_tiles == v.dp_tiles); + CHECK(r.grid_size == v.grid); + CHECK(r.iters_per_sk_cta == v.ipsc); + CHECK(r.extra_iters == v.extra); + + int grid[3] = {0, 0, 0}; + CHECK(rocke_wgrad_conv_streamk_grid(&s, v.arch, grid) == ROCKE_OK); + CHECK(grid[0] == v.grid && grid[1] == 1 && grid[2] == 1); + + size_t flags = 1, partials = 1; + CHECK(rocke_wgrad_conv_streamk_workspace_layout(&s, v.arch, &flags, &partials) == ROCKE_OK); + CHECK(flags == v.flags_bytes); + CHECK(partials == v.partials_bytes); + CHECK(rocke_wgrad_conv_streamk_workspace_bytes(&s, v.arch) + == v.flags_bytes + v.partials_bytes); + } + + /* The other reductions on the same shape, auto pool. */ + static const struct + { + const char* arch; + const char* reduction; + const char* dtype_d; + int grid; + size_t workspace_bytes; + } k_reductions[] = { + {"gfx950", "tree", "fp16", 256, 4198400}, + {"gfx942", "tree", "fp16", 304, 4985600}, + {"gfx950", "atomic", "fp32", 1024, 0}, + {"gfx942", "atomic", "fp32", 1216, 0}, + {"gfx950", "workspace", "fp16", 1024, 58392576}, + {"gfx942", "workspace", "fp16", 1216, 58392576}, + }; + for(const auto& e : k_reductions) + { + rocke_implicit_gemm_conv_wgrad_spec_t s = spec(splits()); + s.streamk_ctas = -1; + s.streamk_reduction = e.reduction; + s.dtype_d = e.dtype_d; + int grid[3] = {0, 0, 0}; + CHECK(rocke_wgrad_conv_streamk_grid(&s, e.arch, grid) == ROCKE_OK); + CHECK(grid[0] == e.grid); + size_t flags = 1, partials = 1; + CHECK(rocke_wgrad_conv_streamk_workspace_layout(&s, e.arch, &flags, &partials) == ROCKE_OK); + const bool fixup = strcmp(e.reduction, "tree") == 0; + CHECK(fixup ? flags + partials == e.workspace_bytes : flags == 0 && partials == 0); + CHECK(rocke_wgrad_conv_streamk_workspace_bytes(&s, e.arch) == e.workspace_bytes); + CHECK(rocke_wgrad_conv_spec_is_deterministic(&s) == fixup); + } + + /* 10 tiles over a pool of 5: all data-parallel, no flags or slots. */ + rocke_implicit_gemm_conv_wgrad_spec_t s = spec(dense()); + s.streamk_ctas = 5; + size_t flags = 1, partials = 1; + CHECK(rocke_wgrad_conv_streamk_workspace_layout(&s, "gfx950", &flags, &partials) == ROCKE_OK); + CHECK(flags == 0 && partials == 0); + /* 10 tiles over 7: 256 flag bytes, 7 partial tiles. */ + s.streamk_ctas = 7; + CHECK(rocke_wgrad_conv_streamk_workspace_layout(&s, "gfx950", &flags, &partials) == ROCKE_OK); + CHECK(flags == 256 && partials == 7 * 64 * 64 * 4); +} + +void check_invalid_inputs() +{ + rocke_status_t st = ROCKE_OK; + int grid[3] = {0, 0, 0}; + size_t flags = 0, partials = 0; + + rocke_streamk_iter_partition_t part = rocke_wgrad_conv_streamk_partition(NULL, NULL, &st); + CHECK(st == ROCKE_ERR_VALUE); + CHECK(part.m_tiles == 0 && part.max_active_wgs == 0); + CHECK(rocke_wgrad_conv_streamk_grid(NULL, NULL, grid) == ROCKE_ERR_VALUE); + CHECK(rocke_wgrad_conv_streamk_workspace_layout(NULL, NULL, &flags, &partials) + == ROCKE_ERR_VALUE); + CHECK(rocke_wgrad_conv_streamk_workspace_bytes(NULL, NULL) == 0); + + rocke_implicit_gemm_conv_wgrad_spec_t s = spec(dense()); + CHECK(rocke_wgrad_conv_streamk_grid(&s, NULL, NULL) == ROCKE_ERR_VALUE); + CHECK(rocke_wgrad_conv_streamk_workspace_layout(&s, NULL, NULL, &partials) == ROCKE_ERR_VALUE); + CHECK(rocke_wgrad_conv_streamk_workspace_layout(&s, NULL, &flags, NULL) == ROCKE_ERR_VALUE); + + const int bad_ctas[] = {0, -2, -3}; + for(int c : bad_ctas) + { + s.streamk_ctas = c; + st = ROCKE_OK; + (void)rocke_wgrad_conv_streamk_partition(&s, "gfx950", &st); + CHECK(st == ROCKE_ERR_VALUE); + CHECK(rocke_wgrad_conv_streamk_grid(&s, "gfx950", grid) == ROCKE_ERR_VALUE); + CHECK(rocke_wgrad_conv_streamk_workspace_layout(&s, "gfx950", &flags, &partials) + == ROCKE_ERR_VALUE); + CHECK(rocke_wgrad_conv_streamk_workspace_bytes(&s, "gfx950") == 0); + } + s.streamk_ctas = 6; + s.streamk = "off"; + st = ROCKE_OK; + (void)rocke_wgrad_conv_streamk_partition(&s, "gfx950", &st); + CHECK(st == ROCKE_ERR_VALUE); + CHECK(rocke_wgrad_conv_streamk_workspace_bytes(&s, "gfx950") == 0); +} + +/* is_valid_wgrad_spec accept/reject reasons, verbatim from Python. */ +void check_validator() +{ + char why[1024]; + auto valid = [&](const rocke_implicit_gemm_conv_wgrad_spec_t& s, const char* arch) { + why[0] = '\0'; + return rocke_implicit_gemm_conv_wgrad_is_valid_spec(&s, arch, why, sizeof(why)); + }; + auto rejects = + [&](const rocke_implicit_gemm_conv_wgrad_spec_t& s, const char* arch, const char* reason) { + const bool ok = valid(s, arch); + if(ok || strcmp(why, reason) != 0) + { + fprintf(stderr, " want: %s\n got: %s\n", reason, ok ? "(accepted)" : why); + return false; + } + return true; + }; + + rocke_implicit_gemm_conv_wgrad_spec_t s = spec(dense()); + CHECK(valid(s, "gfx950")); + CHECK(valid(s, NULL)); + s.streamk_ctas = -1; + CHECK(valid(s, "gfx950")); + s.streamk = "persistent"; + s.streamk_reduction = "tree"; + CHECK(valid(s, "gfx950")); + s = spec(dense()); + s.warp_tile_m = 16; + s.warp_tile_n = 16; + CHECK(valid(s, "gfx942")); + s = spec(dense()); + s.dtype_d = "fp32"; + s.streamk_reduction = "atomic"; + CHECK(valid(s, "gfx950")); + + s = spec(dense()); + s.streamk_ctas = 0; + CHECK(rejects(s, "gfx950", "streamk_ctas must be -1 (auto) or > 0 (got 0)")); + s.streamk_ctas = -3; + CHECK(rejects(s, "gfx950", "streamk_ctas must be -1 (auto) or > 0 (got -3)")); + + s = spec(dense()); + s.streamk = "sideways"; + CHECK(rejects(s, "gfx950", "streamk must be 'off', 'dp_sk' or 'persistent' (got 'sideways')")); + s = spec(dense()); + s.streamk_reduction = "ring"; + CHECK(rejects(s, + "gfx950", + "streamk_reduction must be 'linear', 'tree', 'atomic' or 'workspace' " + "(got 'ring')")); + s = spec(dense()); + s.split_k = 4; + CHECK(rejects(s, + "gfx950", + "streamk and split_k both partition K_wg; use split_k=1 with streamk " + "(got split_k=4)")); + s = spec(dense()); + s.async_dma = true; + CHECK(rejects(s, + "gfx950", + "streamk runs one plain K loop per work item; async_dma/unroll_k " + "pipeline the K loop across items")); + s = spec(dense()); + s.streamk_reduction = "atomic"; + CHECK(rejects(s, + "gfx950", + "streamk_reduction='atomic' needs dtype_d='fp32' (got 'fp16'); use " + "'linear', 'tree' or 'workspace'")); + s = spec(dense()); + CHECK(rejects(s, "gfx90a", "stream-K wgrad supports only gfx942, gfx950 (got gfx90a)")); + + /* Multi-cause rejects: the first gate in Python's order wins. */ + s = spec(dense()); + s.streamk_ctas = 0; + s.tile_m = 96; + CHECK(rejects(s, "gfx950", "tile_m not divisible by warp_m * warp_tile_m")); + s = spec(dense()); + s.warp_m = 8; + s.warp_n = 8; + s.tile_m = 256; + s.tile_n = 256; + s.streamk_ctas = 0; + CHECK(rejects(s, "gfx950", "block_size 4096 > 1024 (hardware cap) on gfx950")); + s = spec(dense()); + s.wave_size = 32; + s.streamk_ctas = 0; + CHECK(rejects(s, "gfx950", "spec wave_size 32 != gfx950 wave_size 64")); + s = spec(dense()); + CHECK(rejects(s, "gfx1250", "spec wave_size 64 != gfx1250 wave_size 32")); + s = spec(dense()); + s.has_vector_size_c = true; + s.vector_size_c = 4; + s.streamk_ctas = -5; + CHECK(rejects(s, "gfx950", "default epilogue is not supported with vector size c: 4")); + s = spec(dense()); + s.async_dma = true; + s.streamk = "off"; + CHECK(rejects(s, + "gfx950", + "wgrad async_dma requires lds_k_outer=True: the direct global->LDS load " + "needs a stride-1 reduction axis, which wgrad only has once the tile is " + "stored K-outer")); + s = spec(dense()); + s.tile_m = 96; + CHECK(rejects(s, + "gfx9999", + "\"unknown gfx target 'gfx9999'; known: ['gfx11-generic', 'gfx1151', " + "'gfx1201', 'gfx1250', 'gfx90a', 'gfx942', 'gfx950']. Add a row to " + "arch_specs.json.\"")); + + /* Depthwise ConvProblem(N=2, Hi=8, Wi=8, C=16, K=16, Y=3, X=3, pH=1, pW=1, + * groups=16): the shape group merging admits. */ + rocke_conv_problem_t dw = rocke_conv_problem_default(2, 8, 8, 16, 16, 3, 3); + dw.pH = 1; + dw.pW = 1; + dw.groups = 16; + s = spec(dw); + s.group_merge = 3; + CHECK(rejects(s, "gfx950", "group_merge must be 1 or one of (2, 4, 8, 16, 32, 64); got 3")); + s.group_merge = 2; + CHECK(rejects(s, "gfx950", "streamk folds conv groups into GEMM-M; use group_merge=1 (got 2)")); + s.tile_m = 96; + CHECK(rejects(s, "gfx950", "tile_m not divisible by warp_m * warp_tile_m")); +} + +} // namespace + +int main() +{ + check_default_pool(); + check_ctas(); + check_plans(); + check_invalid_inputs(); + check_validator(); + if(g_failures) + { + fprintf(stderr, "rocke_wgrad_streamk_host: %d failure(s)\n", g_failures); + return 1; + } + printf("rocke_wgrad_streamk_host: OK\n"); + return 0; +} From 4cf1e5c46849607591d3a1d5fcf0ab9c64939379 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Thu, 8 Oct 2026 17:50:53 +0000 Subject: [PATCH 18/20] fix(rocke): handle all-DP stream-K partitions in the sweep and share launch code - The sweep's stream-K launcher always allocates a workspace, so a shape whose partition has no stream-K CTAs no longer aborts the sweep. - Stream-K builds reuse the regular wgrad build path, Stage 2 launches share one helper, and the run-from-cache stream-K branch is additive. - Drop the unused stream-K extras from kernel_sweep's launch values and warn when stream-K options are given with --streamk off. ISSUE ID : AICK-1943 Co-Authored-By: Claude Opus 5.5 --- .../common/benchmark_conv_compare.py | 10 +- .../common/benchmark_implicit_gemm_conv.py | 505 ++++++++---------- .../library/benchmarks/common/kernel_sweep.py | 13 +- .../tests/test_streamk_sweep_launch.py | 277 ++++++++++ .../rocke/platform/CMakeLists.txt | 1 + 5 files changed, 528 insertions(+), 278 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_streamk_sweep_launch.py diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_conv_compare.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_conv_compare.py index ab7c62d3d5b0..ddd6cabc8349 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_conv_compare.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_conv_compare.py @@ -230,6 +230,14 @@ def _build_implicit_only_args(args) -> list[str]: a += ["--split-k", str(args.split_k)] if args.split_k_prune is not None: a += ["--split-k-prune", str(args.split_k_prune)] + if args.streamk == "off" and ( + args.streamk_reductions is not None or args.streamk_ctas != -1 + ): + print( + "warning: --streamk-reductions / --streamk-ctas are ignored " + "without --streamk", + file=sys.stderr, + ) if args.streamk != "off": a += ["--streamk", args.streamk] if args.streamk_reductions is not None: @@ -458,7 +466,7 @@ def main() -> int: dest="streamk_ctas", metavar="N", help=( - "stream-K CTA pool (-1 = arch default). Forwarded to implicit-GEMM " "only." + "stream-K CTA pool (-1 = arch default). Forwarded to implicit-GEMM only." ), ) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py index f95008aebae8..e2938cab6290 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/benchmark_implicit_gemm_conv.py @@ -77,6 +77,7 @@ EarlyStop, add_early_stop_arg, ) +from benchmarks.common.kernel_sweep import _AOT_ALIAS_PIPELINES # noqa: E402 from benchmarks.common.kernel_sweep import ( # noqa: E402 CACHE_EPILOGUES as _EPILOGUES, CACHE_PIPELINES as _PIPELINES, @@ -86,7 +87,6 @@ CACHE_WARP_TILE_MN as _WARP_TILE_MN, CACHE_WGRAD_STREAMK_MODES as _STREAMK_MODES, CACHE_WGRAD_STREAMK_REDUCTIONS as _STREAMK_REDUCTIONS, - _AOT_ALIAS_PIPELINES, ) # gfx1250-only: wider CUs allow larger tiles and more warp concurrency. @@ -1191,6 +1191,15 @@ def main() -> int: print(f"--csv-top must be >= 1 (got {args.csv_top})", file=sys.stderr) return 2 + if args.streamk == "off" and ( + args.streamk_reductions != ",".join(_STREAMK_REDUCTIONS) + or args.streamk_ctas != -1 + ): + print( + "warning: --streamk-reductions / --streamk-ctas are ignored " + "without --streamk", + file=sys.stderr, + ) args.streamk_modes = { "off": (), "all": _STREAMK_MODES, @@ -1708,6 +1717,16 @@ def _build_wgrad_one(args_tuple): group_merge=group_merge, lds_k_outer=lds_k_outer, async_dma=async_dma, + # Stream-K leg: the launch mode and reduction pick the kernel; the CTA + # pool stays at -1 here (it is a kernarg, resolved at launch). + **( + {} + if combo.streamk == "off" + else { + "streamk": combo.streamk, + "streamk_reduction": combo.streamk_reduction, + } + ), ) ok, _ = is_valid_wgrad_spec(spec, arch) if not ok: @@ -1878,87 +1897,6 @@ def _build_wgrad_two_stage_one(args_tuple): return combo, spec, resolved_split_k, s1_kernel, s2_kernel -def _build_wgrad_streamk_one(args_tuple): - """Top-level picklable worker: validate + build IR for one stream-K combo. - - The combo's ``streamk`` / ``streamk_reduction`` pick the kernel; the CTA - pool is left at -1 (it is a kernarg, resolved at launch). Returns - ``(combo, spec, 1, kernel)`` on success, or ``None``. - """ - combo, problem, dtype, arch = args_tuple - combo = WgradCombo(*combo) - - from rocke.core.arch import ArchTarget - from kernels.common.conv_implicit_gemm import ConvDataSpec - from kernels.common.conv_implicit_gemm_wgrad import ( - WgradConvSpec, - build_implicit_gemm_conv_wgrad, - is_valid_wgrad_spec, - ) - - target = ArchTarget.from_gfx(arch) - _mma_family = "wmma" if target.wave_size == 32 else "mma" - atom = target.mma.select_largest_k( - family=_mma_family, - a_dtype=dtype, - b_dtype=dtype, - c_dtype="fp32", - m=combo.warp_tile_mn, - n=combo.warp_tile_mn, - k_max=combo.tile_k, - ) - if atom is None: - return None - spec = WgradConvSpec( - problem=problem, - name="rocke_bench_igemm_wgrad", - data=ConvDataSpec(dtype_a=dtype, dtype_b=dtype, dtype_d=dtype), - tile_m=combo.tile_m, - tile_n=combo.tile_n, - tile_k=combo.tile_k, - warp_m=combo.warp_m, - warp_n=combo.warp_n, - warp_tile_m=combo.warp_tile_mn, - warp_tile_n=combo.warp_tile_mn, - warp_tile_k=atom.k, - wave_size=target.wave_size, - pipeline=combo.pipeline, - epilogue=combo.epilogue, - split_k=1, - # Same per-combo K-outer gate as the other legs, so the timings stay - # comparable. - lds_k_outer=WgradConvSpec.default_lds_k_outer( - arch=arch, - dtype_a=dtype, - dtype_b=dtype, - warp_tile_m=combo.warp_tile_mn, - warp_tile_n=combo.warp_tile_mn, - wave_size=target.wave_size, - ), - streamk=combo.streamk, - streamk_reduction=combo.streamk_reduction, - ) - ok, _ = is_valid_wgrad_spec(spec, arch) - if not ok: - return None - try: - kernel = build_implicit_gemm_conv_wgrad(spec, arch=arch) - except ValueError: - return None - return combo, spec, 1, kernel - - -def _build_and_compile_wgrad_streamk_one(args_tuple): - """Merged worker: :func:`_build_wgrad_streamk_one` + compile.""" - result = _build_wgrad_streamk_one(args_tuple) - if result is None: - return None - combo, spec, split_k, kernel = result - from rocke import compile_kernel as _compile_kernel - - return combo, spec, split_k, _compile_kernel(kernel, arch=args_tuple[3]) - - def _build_dgrad_one(args_tuple): """Top-level picklable worker: validate + build IR for one dgrad combo. @@ -2908,7 +2846,7 @@ def _degrees(pipeline: str, async_dma: bool) -> tuple: ) pending_sk = _build_ir_parallel( [(c, problem, dtype, arch) for c in sk_combos], - _build_and_compile_wgrad_streamk_one, + _build_and_compile_wgrad_one, jobs, ) n_skipped += len(sk_combos) - len(pending_sk) @@ -3187,7 +3125,6 @@ def _launch_spk(_v=_vals_snap, _c=_cfg_snap): ) from kernels.common.conv_wgrad_workspace_reduce import ( WgradReduceSpec, - wgrad_reduce_grid, wgrad_reduce_signature, ) @@ -3232,9 +3169,6 @@ def _launch_spk(_v=_vals_snap, _c=_cfg_snap): # Must match Stage 1 or the fold covers the wrong slab count. ws_replicas=spec.ws_replicas, ) - s2_grid = wgrad_reduce_grid(s2_spec) - s2_block = (s2_spec.tile_m * s2_spec.tile_n, 1, 1) - try: s1_launcher = KernelLauncher( hsaco=s1_art.hsaco, @@ -3279,19 +3213,16 @@ def _launch_spk(_v=_vals_snap, _c=_cfg_snap): ws_ptr=int(ws_dev), ws_bytes=ws_nbytes, ) - s2_values = { - "ws_ptr": ws_dev, - "dw_ptr": dW_dev2, - "wg_M": spec.wg_M, - "wg_N": spec.wg_N, - "ws_bytes": ws_nbytes, - "dw_bytes": dW_t.nbytes, - "groups": spec.problem.groups, - } + s2_values, s2_cfg = _stage2_launch( + s2_spec, + ws_ptr=ws_dev, + ws_bytes=ws_nbytes, + dw_ptr=dW_dev2, + dw_bytes=dW_t.nbytes, + ) s1_cfg = LaunchConfig( grid=s1_grid, block=(spec.block_size, 1, 1), stream=0 ) - s2_cfg = LaunchConfig(grid=s2_grid, block=s2_block, stream=0) def _launch_two_stage( _s1=s1_launcher, @@ -3979,7 +3910,6 @@ def _run_from_cache(args, arch, target, cases, cache, directions) -> int: from rocke.runtime import synchronize_and_release from rocke.runtime.hip_module import HipError, Runtime from rocke.runtime.launcher import KernelLauncher, LaunchConfig - from benchmarks.common.kernel_sweep import _launch_values_for _TORCH_DT = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32} @@ -4148,6 +4078,31 @@ def _u8(t): print(f" [skip] {ident.label()}: {e}", flush=True) continue + # A stream-K binary partitions K itself over a CTA pool (a + # kernarg) and has its own ABI, workspace and zeroing + # contract per reduction, so it gets its own launch setup and + # one launch; its workspace is released after the loop. + sk_bound = sk_free = None + if direction == "wgrad" and ident.streamk != "off": + try: + sk_run, sk_free, sk_label = _cached_streamk_run( + args, + problem, + ident, + rt, + arch, + launcher, + (A_dev, B_dev, D_dev), + (A_t.nbytes, B_t.nbytes, D_t.nbytes), + stage2_kernels, + KernelLauncher, + LaunchConfig, + ) + except ValueError as e: + print(f" [skip] {ident.label()}: {e}", flush=True) + continue + sk_bound = (sk_run, sk_label) + # A split-K wgrad binary is one kernel for every degree > 1 # (the degree is a kernarg), so it is timed at each degree it # can use. @@ -4164,156 +4119,74 @@ def _u8(t): degrees = (None,) for split_k in degrees: - free_ws = None - try: - if direction == "wgrad" and ident.streamk != "off": - # Stream-K: the kernel partitions K itself over a CTA pool - # (a kernarg) and has its own ABI, workspace and zeroing - # contract per reduction. - try: - run, free_ws, label = _cached_streamk_run( - args, - problem, - ident, - rt, - arch, - launcher, - (A_dev, B_dev, D_dev), - (A_t.nbytes, B_t.nbytes, D_t.nbytes), - stage2_kernels, - KernelLauncher, - LaunchConfig, - ) - except ValueError as e: - print(f" [skip] {ident.label()}: {e}", flush=True) - continue - else: - try: - extras, grid = _extras_and_grid( - direction, - problem, - ident, - rt, - case_dtype, - split_k=split_k, - ) - except ValueError as e: - print(f" [skip] {ident.label()}: {e}", flush=True) - continue - values = _launch_values_for( - direction, - problem, - ident, - (A_dev, B_dev, D_dev), - (A_t.nbytes, B_t.nbytes, D_t.nbytes), - extras, - ) - cfg = LaunchConfig( - grid=grid, - block=( - meta.get("block_size") - or (ident.warp_m * ident.warp_n * ident.wave_size), - 1, - 1, - ), - stream=0, - ) - run = launcher.bind(values, config=cfg) - if direction == "wgrad" and ident.two_stage: - # Two-stage is three steps and all of them are the - # algorithm's cost: zero the scratch (Stage 1 - # accumulates), Stage 1, then the fold/cast into dW. - from kernels.common.conv_wgrad_workspace_reduce import ( - wgrad_reduce_grid, - ) - - s2_launcher, s2_spec = _stage2_reduce( - problem, - ident.dtype_d, - arch, - stage2_kernels, - KernelLauncher, - _ws_replicas_of(ident), - ) - s2_values = { - "ws_ptr": extras["ws_ptr"], - "dw_ptr": D_dev, - "wg_M": s2_spec.wg_M, - "wg_N": s2_spec.wg_N, - "ws_bytes": extras["ws_bytes"], - "dw_bytes": D_t.nbytes, - "groups": problem.groups, - } - s2_cfg = LaunchConfig( - grid=wgrad_reduce_grid(s2_spec), - block=(s2_spec.block_size, 1, 1), - stream=0, - ) - - def run( - values=values, - cfg=cfg, - extras=extras, - s2_launcher=s2_launcher, - s2_values=s2_values, - s2_cfg=s2_cfg, - ): - rt.memset(extras["ws_ptr"], 0, extras["ws_bytes"]) - launcher(values, config=cfg) - s2_launcher(s2_values, config=s2_cfg) - - label = ident.label() - if split_k is not None: - label += f" @split_k={split_k}" - if ref_out is not None: - # The output is zeroed first: the atomic split-K - # epilogues accumulate into it. A kernel that fails is - # reported and not timed, so it can never rank. - stopped, passed = _verify_kernel( - rt=rt, - launch_fn=run, - out_dev=D_dev, - out_t=D_t, - zero_init_out=True, - ref_out=ref_out, - kernel_name=label, - dump_fail=args.dump_fail, - u8=_u8, - arch=arch, - compute_dtype=case_dtype, - ) - if stopped: - return 1 - if not passed: - n_failed += 1 - overall_rc = 1 - continue - try: - ms = _stop.measure( - run, - warmup=args.warmup, - iters=args.iters, - stream=0, - ) - except (HipError, RuntimeError) as e: - print(f" [skip] {ident.label()}: {e}", flush=True) - continue - if ms is None: - _stop.report(label) - continue - results.append( - (flop / (ms * 1e9), bytes_xfer / (ms * 1e6), ms, label) + bound = sk_bound or _cached_bind( + direction, + problem, + ident, + meta, + rt, + launcher, + case_dtype, + split_k, + (A_dev, B_dev, D_dev), + (A_t.nbytes, B_t.nbytes, D_t.nbytes), + arch, + stage2_kernels, + KernelLauncher, + LaunchConfig, + ) + if bound is None: + continue + run, label = bound + if ref_out is not None: + # The output is zeroed first: the atomic split-K + # epilogues accumulate into it. A kernel that fails is + # reported and not timed, so it can never rank. + stopped, passed = _verify_kernel( + rt=rt, + launch_fn=run, + out_dev=D_dev, + out_t=D_t, + zero_init_out=True, + ref_out=ref_out, + kernel_name=label, + dump_fail=args.dump_fail, + u8=_u8, + arch=arch, + compute_dtype=case_dtype, ) - # One line per measured launch, as the JIT sweep prints, - # so a long run shows its progress and partial results. - print( - f" [{cand_idx:>{len(str(n_cand))}}/{n_cand}] " - f"{results[-1][0]:7.1f} TFLOPS {ms:8.3f} ms {label}", - flush=True, + if stopped: + return 1 + if not passed: + n_failed += 1 + overall_rc = 1 + continue + try: + ms = _stop.measure( + run, + warmup=args.warmup, + iters=args.iters, + stream=0, ) - finally: - if free_ws is not None: - free_ws() + except (HipError, RuntimeError) as e: + print(f" [skip] {ident.label()}: {e}", flush=True) + continue + if ms is None: + _stop.report(label) + continue + results.append( + (flop / (ms * 1e9), bytes_xfer / (ms * 1e6), ms, label) + ) + # One line per measured launch, as the JIT sweep prints, + # so a long run shows its progress and partial results. + print( + f" [{cand_idx:>{len(str(n_cand))}}/{n_cand}] " + f"{results[-1][0]:7.1f} TFLOPS {ms:8.3f} ms {label}", + flush=True, + ) + + if sk_free is not None: + sk_free() synchronize_and_release() if ref_out is not None: @@ -4507,6 +4380,31 @@ def _stage2_reduce(problem, dtype_d, arch, compiled, KernelLauncher, ws_replicas return compiled[key], spec +def _stage2_launch(s2_spec, *, ws_ptr, ws_bytes, dw_ptr, dw_bytes): + """``(values, config)`` of the Stage 2 fold of ``s2_spec`` into dW. + + Stage 2 (:class:`WgradReduceSpec`) folds the f32 scratch at ``ws_ptr`` + into dW for both the two-stage split-K path and the stream-K + ``workspace`` reduction; every leg launches it through here. + """ + from rocke.runtime.launcher import LaunchConfig + from kernels.common.conv_wgrad_workspace_reduce import wgrad_reduce_grid + + values = { + "ws_ptr": ws_ptr, + "dw_ptr": dw_ptr, + "wg_M": s2_spec.wg_M, + "wg_N": s2_spec.wg_N, + "ws_bytes": ws_bytes, + "dw_bytes": dw_bytes, + "groups": s2_spec.groups, + } + config = LaunchConfig( + grid=wgrad_reduce_grid(s2_spec), block=(s2_spec.block_size, 1, 1), stream=0 + ) + return values, config + + def _streamk_launch( *, rt, @@ -4544,21 +4442,22 @@ def _streamk_launch( flags_bytes, _ = wgrad_streamk_workspace_layout(spec, arch=arch, problem=problem) grid = wgrad_streamk_grid(spec, arch=arch, problem=problem) pool = wgrad_streamk_launch(spec, arch=arch).pool - ws = rt.alloc(ws_bytes) if ws_bytes else None + # An all-data-parallel partition (no stream-K CTAs) needs no workspace, + # but the kernargs still carry the pointers, so a 4-byte stub stands in + # for the empty allocation. + ws = rt.alloc(max(ws_bytes, 4)) try: extra = {} if reduction == "workspace": - extra = dict(ws_ptr=ws, ws_bytes=ws_bytes) + extra = {"ws_ptr": ws, "ws_bytes": ws_bytes} elif reduction in ("linear", "tree"): - extra = dict(sk_flags=ws, sk_partials=int(ws) + flags_bytes) + extra = {"sk_flags": ws, "sk_partials": int(ws) + flags_bytes} values = wgrad_streamk_launch_values( spec, *ptrs, *sizes, arch=arch, problem=problem, **extra ) cfg = LaunchConfig(grid=grid, block=(spec.block_size, 1, 1), stream=0) s2 = None if reduction == "workspace": - from kernels.common.conv_wgrad_workspace_reduce import wgrad_reduce_grid - s2_launcher, s2_spec = _stage2_reduce( problem, spec.data.dtype_d, @@ -4569,24 +4468,16 @@ def _streamk_launch( ) s2 = ( s2_launcher, - { - "ws_ptr": ws, - "dw_ptr": ptrs[2], - "wg_M": s2_spec.wg_M, - "wg_N": s2_spec.wg_N, - "ws_bytes": ws_bytes, - "dw_bytes": sizes[2], - "groups": problem.groups, - }, - LaunchConfig( - grid=wgrad_reduce_grid(s2_spec), - block=(s2_spec.block_size, 1, 1), - stream=0, + *_stage2_launch( + s2_spec, + ws_ptr=ws, + ws_bytes=ws_bytes, + dw_ptr=ptrs[2], + dw_bytes=sizes[2], ), ) except Exception: - if ws is not None: - rt.free(ws) + rt.free(ws) raise def run(): @@ -4601,8 +4492,7 @@ def run(): s2[0](s2[1], config=s2[2]) def free(): - if ws is not None: - rt.free(ws) + rt.free(ws) return run, free, pool @@ -4704,6 +4594,85 @@ def _runtime_split_ks(problem, ident) -> tuple: return _split_k_degrees(problem, ident.tile_m, ident.tile_n, ident.tile_k) +def _cached_bind( + direction, + problem, + ident, + meta, + rt, + launcher, + case_dtype, + split_k, + ptrs, + sizes, + arch, + stage2_kernels, + KernelLauncher, + LaunchConfig, +): + """``(run, label)`` of a cached (non-stream-K) binary at ``split_k``. + + ``ptrs`` / ``sizes`` are the A, B, D buffers. Returns None, after saying + why, when the binary cannot launch ``problem``. + """ + from benchmarks.common.kernel_sweep import _launch_values_for + + try: + extras, grid = _extras_and_grid( + direction, problem, ident, rt, case_dtype, split_k=split_k + ) + except ValueError as e: + print(f" [skip] {ident.label()}: {e}", flush=True) + return None + values = _launch_values_for(direction, problem, ident, ptrs, sizes, extras) + cfg = LaunchConfig( + grid=grid, + block=( + meta.get("block_size") or (ident.warp_m * ident.warp_n * ident.wave_size), + 1, + 1, + ), + stream=0, + ) + run = launcher.bind(values, config=cfg) + if direction == "wgrad" and ident.two_stage: + # Two-stage is three steps and all of them are the algorithm's cost: + # zero the scratch (Stage 1 accumulates), Stage 1, then the fold/cast + # into dW. + s2_launcher, s2_spec = _stage2_reduce( + problem, + ident.dtype_d, + arch, + stage2_kernels, + KernelLauncher, + _ws_replicas_of(ident), + ) + s2_values, s2_cfg = _stage2_launch( + s2_spec, + ws_ptr=extras["ws_ptr"], + ws_bytes=extras["ws_bytes"], + dw_ptr=ptrs[2], + dw_bytes=sizes[2], + ) + + def run( + values=values, + cfg=cfg, + extras=extras, + s2_launcher=s2_launcher, + s2_values=s2_values, + s2_cfg=s2_cfg, + ): + rt.memset(extras["ws_ptr"], 0, extras["ws_bytes"]) + launcher(values, config=cfg) + s2_launcher(s2_values, config=s2_cfg) + + label = ident.label() + if split_k is not None: + label += f" @split_k={split_k}" + return run, label + + def _extras_and_grid(direction, problem, ident, rt, dtype, split_k=None): """Trailing kernargs and the launch grid for one cached kernel. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py index fc1f17582874..301d86a0622a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py @@ -758,10 +758,10 @@ def job_for_identity(identity: KernelIdentity) -> BuildJob: ) want = identity.stable_hash() streamk = ( - dict( - streamk_modes=(identity.streamk,), - streamk_reductions=(identity.streamk_reduction,), - ) + { + "streamk_modes": (identity.streamk,), + "streamk_reductions": (identity.streamk_reduction,), + } if identity.direction == "wgrad" and identity.streamk != "off" else {} ) @@ -1755,11 +1755,6 @@ def _launch_values_for(direction, problem, identity, ptrs, sizes, extras): split_k=extras.get("split_k", 1), ws_ptr=extras.get("ws_ptr"), ws_bytes=extras.get("ws_bytes"), - # Stream-K: the StreamKLaunch (pool, mode, reduction) that fills - # the partition block, plus the linear/tree flag/partial buffers. - streamk=extras.get("streamk"), - sk_flags=extras.get("sk_flags"), - sk_partials=extras.get("sk_partials"), ) return ConvArgs.from_problem( problem, direction="dgrad", tile_m=tm, tile_n=tn diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_streamk_sweep_launch.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_streamk_sweep_launch.py new file mode 100644 index 000000000000..15bc917bd8c6 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_streamk_sweep_launch.py @@ -0,0 +1,277 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""The implicit-GEMM sweep's stream-K wgrad launch helper. + +``benchmark_implicit_gemm_conv._streamk_launch`` sets up one stream-K launch +for the JIT sweep and for ``--run-from-cache``: it sizes and allocates the +reduction's workspace, fills the kernargs and wraps the per-reduction zeroing +contract (and the Stage 2 fold for ``workspace``) in a ``run()`` closure. + +The case pinned here is a partition with no stream-K CTAs at all: when every +output tile is data-parallel the workspace is zero bytes, and the helper must +still hand the kernel valid flag / partial pointers instead of failing (an +uncaught error there aborts the whole sweep). The host tests drive the helper +with a recording runtime; the GPU test launches the real kernel through it +and checks dW, and skips without a gfx942/gfx950 GPU and torch. +""" + +from __future__ import annotations + +import ctypes +import importlib.util +import unittest + +from rocke.runtime.hip_module import get_device_arch + +from benchmarks.common.benchmark_implicit_gemm_conv import _streamk_launch + +_HAS_TORCH = importlib.util.find_spec("torch") is not None +GPU_ARCH = get_device_arch(0) + +# Every output tile is data-parallel under the default gfx950 pool: 1 x 9 +# tiles of 64x64 over K=64, Y*X*C=576, so the stream-K CTA count is 0. +_ALL_DP = dict(N=2, Hi=14, Wi=14, C=64, K=64, Y=3, X=3, pH=1, pW=1) +# 459 tiles: more than the default pool, so stream-K CTAs share some of them. +_SPLITS = dict(N=2, Hi=12, Wi=12, C=192, K=1056, Y=3, X=3, pH=1, pW=1) +# (reduction, dW dtype): "atomic" needs fp32 dW. +_REDUCTIONS = ( + ("linear", "fp16"), + ("tree", "fp16"), + ("atomic", "fp32"), + ("workspace", "bf16"), +) + + +def _spec(shape, *, arch, reduction, dtype_d, mode="dp_sk"): + from kernels.common._conv_implicit_gemm_common import ConvDataSpec, ConvProblem + from kernels.common.conv_implicit_gemm_wgrad import WgradConvSpec + + ab = "bf16" if dtype_d == "bf16" else "fp16" + edge = 32 if arch == "gfx950" else 16 + return WgradConvSpec( + problem=ConvProblem(**shape), + data=ConvDataSpec(dtype_a=ab, dtype_b=ab, dtype_d=dtype_d), + tile_m=64, + tile_n=64, + tile_k=32, + warp_m=2, + warp_n=2, + warp_tile_m=edge, + warp_tile_n=edge, + warp_tile_k=16, + streamk=mode, + streamk_reduction=reduction, + ) + + +class _RecordingRuntime: + """The subset of ``Runtime`` the helper uses, recording every call.""" + + def __init__(self): + self.allocs: dict = {} + self.freed: list = [] + self.memsets: list = [] + self._next = 0x10000 + + def alloc(self, nbytes): + ptr = self._next + self._next += max(nbytes, 1) + 0x1000 + self.allocs[ptr] = nbytes + return ptr + + def free(self, ptr): + self.freed.append(ptr) + + def memset(self, ptr, value, nbytes): + self.memsets.append((ptr, value, nbytes)) + + +class _RecordingLauncher: + def __init__(self): + self.calls: list = [] + + def __call__(self, values, *, config): + self.calls.append((values, config)) + + +def _launch_host(spec, arch): + from rocke.runtime.launcher import LaunchConfig + + rt = _RecordingRuntime() + launcher = _RecordingLauncher() + stage2 = _RecordingLauncher() + # Pre-seed the Stage 2 memo so the workspace reduction needs no compiler. + stage2_memo = {(spec.data.dtype_d, spec.ws_replicas): stage2} + run, free, pool = _streamk_launch( + rt=rt, + launcher=launcher, + spec=spec, + problem=spec.problem, + arch=arch, + ptrs=(0x100, 0x200, 0x300), + sizes=(64, 64, 64), + stage2=stage2_memo, + KernelLauncher=None, + LaunchConfig=LaunchConfig, + ) + return rt, launcher, stage2, run, free, pool + + +class TestStreamKSweepLaunchHost(unittest.TestCase): + """Host-only: the helper's workspace handling, with a recording runtime.""" + + ARCH = "gfx950" + + def test_all_dp_partition_has_no_workspace(self): + # The premise of the other cases: this shape takes no stream-K CTAs. + from kernels.common.conv_implicit_gemm_wgrad import ( + wgrad_streamk_partition, + wgrad_streamk_workspace_nbytes, + ) + + spec = _spec(_ALL_DP, arch=self.ARCH, reduction="linear", dtype_d="fp16") + part = wgrad_streamk_partition(spec, arch=self.ARCH) + self.assertEqual(part.sk_ctas, 0) + self.assertEqual(wgrad_streamk_workspace_nbytes(spec, arch=self.ARCH), 0) + + def test_all_dp_partition_launches(self): + for reduction, dtype_d in _REDUCTIONS: + with self.subTest(reduction=reduction): + spec = _spec( + _ALL_DP, arch=self.ARCH, reduction=reduction, dtype_d=dtype_d + ) + rt, launcher, stage2, run, free, _ = _launch_host(spec, self.ARCH) + # One (stub) allocation, never a zero-byte one. + self.assertEqual(len(rt.allocs), 1) + (ws,) = rt.allocs + self.assertGreaterEqual(rt.allocs[ws], 4) + run() + self.assertEqual(len(launcher.calls), 1) + values, _ = launcher.calls[0] + if reduction in ("linear", "tree"): + self.assertEqual(values["sk_flags"], ws) + self.assertEqual(values["sk_partials"], ws) + self.assertEqual(len(stage2.calls), reduction == "workspace") + free() + self.assertEqual(rt.freed, [ws]) + + def test_split_partition_lays_out_flags_then_partials(self): + from kernels.common.conv_implicit_gemm_wgrad import ( + wgrad_streamk_partition, + wgrad_streamk_workspace_layout, + ) + + spec = _spec(_SPLITS, arch=self.ARCH, reduction="linear", dtype_d="fp16") + self.assertGreater(wgrad_streamk_partition(spec, arch=self.ARCH).sk_ctas, 0) + flags_bytes, partials_bytes = wgrad_streamk_workspace_layout( + spec, arch=self.ARCH + ) + self.assertGreater(flags_bytes, 0) + rt, launcher, _, run, free, _ = _launch_host(spec, self.ARCH) + (ws,) = rt.allocs + self.assertEqual(rt.allocs[ws], flags_bytes + partials_bytes) + run() + values, _ = launcher.calls[0] + self.assertEqual(values["sk_flags"], ws) + self.assertEqual(values["sk_partials"], ws + flags_bytes) + # linear/tree reset only the flag table before each launch. + self.assertEqual(rt.memsets, [(ws, 0, flags_bytes)]) + free() + self.assertEqual(rt.freed, [ws]) + + +def _gpu_skip_reason() -> str: + if not GPU_ARCH: + return "no ROCm GPU detected" + if not _HAS_TORCH: + return "torch not importable" + if GPU_ARCH not in ("gfx942", "gfx950"): + return f"stream-K wgrad is gfx942/gfx950 only; got {GPU_ARCH}" + return "" + + +_GPU_SKIP_REASON = _gpu_skip_reason() + + +@unittest.skipUnless(not _GPU_SKIP_REASON, _GPU_SKIP_REASON or "no GPU") +class TestStreamKSweepLaunchGpu(unittest.TestCase): + """The helper's ``run()`` on the GPU for the all-data-parallel shape.""" + + def _check(self, reduction, dtype_d): + import torch + + from rocke import compile_kernel + from rocke.runtime import synchronize_and_release + from rocke.runtime.hip_module import Runtime + from rocke.runtime.launcher import KernelLauncher, LaunchConfig + from kernels.common.conv_implicit_gemm_wgrad import ( + build_implicit_gemm_conv_wgrad, + wgrad_streamk_signature, + ) + + arch = GPU_ARCH + spec = _spec(_ALL_DP, arch=arch, reduction=reduction, dtype_d=dtype_d) + p = spec.problem + art = compile_kernel(build_implicit_gemm_conv_wgrad(spec, arch=arch), arch=arch) + launcher = KernelLauncher( + hsaco=art.hsaco, + kernel_name=art.kernel_name, + signature=wgrad_streamk_signature(spec), + ) + dt = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32} + torch.manual_seed(0) + X = torch.empty(p.N, p.Hi, p.Wi, p.C).uniform_(-1, 1).to(dt[spec.data.dtype_b]) + dY = torch.empty(p.N, p.Ho, p.Wo, p.K).uniform_(-1, 1).to(dt[spec.data.dtype_a]) + dW = torch.empty(p.K, p.Y, p.X, p.C, dtype=dt[dtype_d]) + ref = torch.nn.grad.conv2d_weight( + X.float().permute(0, 3, 1, 2), + weight_size=(p.K, p.C, p.Y, p.X), + grad_output=dY.float().permute(0, 3, 1, 2), + padding=(p.pH, p.pW), + ).permute(0, 2, 3, 1) + + def u8(t): + return (ctypes.c_uint8 * t.nbytes).from_address(t.data_ptr()) + + rt = Runtime() + bufs = [rt.alloc(t.nbytes) for t in (dY, X, dW)] + try: + rt.memcpy_h2d(bufs[0], u8(dY), dY.nbytes) + rt.memcpy_h2d(bufs[1], u8(X), X.nbytes) + # Poison dW: a reduction that leaves it alone must overwrite it. + rt.memset(bufs[2], 0xFF if reduction != "atomic" else 0, dW.nbytes) + run, free, _ = _streamk_launch( + rt=rt, + launcher=launcher, + spec=spec, + problem=p, + arch=arch, + ptrs=tuple(int(b) for b in bufs), + sizes=(dY.nbytes, X.nbytes, dW.nbytes), + stage2={}, + KernelLauncher=KernelLauncher, + LaunchConfig=LaunchConfig, + ) + try: + run() + synchronize_and_release(0) + finally: + free() + rt.memcpy_d2h(u8(dW), bufs[2], dW.nbytes) + finally: + for b in bufs: + rt.free(b) + out = dW.float() + rel = float((out - ref).abs().max() / ref.abs().max().clamp(min=1.0)) + tol = {"fp32": 1e-4, "fp16": 2e-3, "bf16": 2e-2}[dtype_d] + self.assertLess(rel, tol, f"{reduction}: rel_err={rel:.3e}") + + def test_all_dp_linear(self): + self._check("linear", "fp16") + + def test_all_dp_workspace(self): + self._check("workspace", "bf16") + + +if __name__ == "__main__": + unittest.main() diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt index f72d5a311ebc..a7db23e478bf 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt @@ -438,6 +438,7 @@ if(ROCKE_INSTALL_TESTS) set(_ROCKE_LIBRARY_CONV_TEST_RELPATHS "tests/library/tests/dispatch/test_grouped_conv_wgrad_dispatch.py" "tests/library/tests/test_conv_wgrad_streamk.py" + "tests/library/tests/test_streamk_sweep_launch.py" "tests/library/tests/test_conv_build_signatures.py" "tests/library/tests/test_builder_signature_contract.py" "tests/library/tests/test_library_layering.py") From 33e50deeb7fe051e8dbf1c4c5c203da5e20a77bc Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Fri, 9 Oct 2026 11:44:20 +0000 Subject: [PATCH 19/20] style(rocke): black-format the stream-K wgrad validator test ISSUE ID : AICK-1943 Co-Authored-By: Claude Opus 5.5 --- .../rocke/library/tests/test_conv_wgrad_streamk.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py index 185f6ce20904..5c7127ab7ab0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_conv_wgrad_streamk.py @@ -286,7 +286,12 @@ def test_cpp_validity_agrees(self): ), ("unknown_arch", _DENSE, "gfx9999", {}), ("unknown_arch+tile_m", _DENSE, "gfx9999", dict(tile_m=96)), - ("wave32_on_gfx950+ctas", _DENSE, "gfx950", dict(wave_size=32, streamk_ctas=-5)), + ( + "wave32_on_gfx950+ctas", + _DENSE, + "gfx950", + dict(wave_size=32, streamk_ctas=-5), + ), ("wave64_on_gfx1250", _DENSE, "gfx1250", {}), ("async_dma", _DENSE, "gfx950", dict(async_dma=True)), ("async_dma_off", _DENSE, "gfx950", dict(async_dma=True, streamk="off")), From 2845319a3af4fe12213d33711546e3ea2f953445 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Mon, 12 Oct 2026 13:05:28 +0000 Subject: [PATCH 20/20] perf(rocke): cache only a slice of the stream-K wgrad kernels Stream-K wgrad AOT kernels were built for every default-epilogue geometry x launch mode x reduction x all 16 (vector_size_a, vector_size_b) pairs. Cache them only where they are used: - at the widest vector width W of the dtype (16-byte loads) for (W, W) and (W, 1), with every requested reduction except "tree"; - fully scalar, dp_sk x {tree, workspace}, on tiles with tile_n >= 128, for large-filter depthwise layers. Shapes at other vector widths are served by the split-K kernels, which are still cached at every width. The selection lives in _wgrad_streamk_builds and count_jobs mirrors it. Stream-K instances per dtype on gfx950: bf16/fp16 178560 -> 17076, fp32 171720 -> 31374; gfx942: bf16/fp16 223200 -> 21234, fp32 133704 -> 24192. Co-Authored-By: Claude Opus 5.5 --- .../library/benchmarks/common/kernel_sweep.py | 81 +++++++++++++++---- .../common/README_conv_implicit_gemm_wgrad.md | 9 ++- 2 files changed, 73 insertions(+), 17 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py index ace675ce76f0..73e2d931cd5e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/kernel_sweep.py @@ -52,10 +52,13 @@ (``group_merge`` in 2..64, two-stage, load widths derived by the builder), served only to depthwise problems whose merged GEMM fits the tile. Opt-in (``streamk_modes``; ``--compile-all --streamk``): stream-K - kernels, default epilogue only (c = 1, b in {1, 2, 4, 8}), ``split_k`` = 1, - one per launch mode (``dp_sk``, ``persistent``) x reduction (``linear``, - ``tree``, ``workspace`` with ``ws_replicas`` = 8, and ``atomic`` where dW - is fp32). The CTA pool is a kernarg, so it is not swept here. + kernels, default epilogue only, ``split_k`` = 1, c = 1. At the widest + vector width W of the dtype (a = W; b in {W, 1}): one per launch mode + (``dp_sk``, ``persistent``) x reduction (``linear``, ``workspace`` with + ``ws_replicas`` = 8, and ``atomic`` where dW is fp32). Fully scalar + (a = b = 1), on tiles with tile_n >= 128 only: ``dp_sk`` x (``tree``, + ``workspace``), for large-filter depthwise layers. The CTA pool is a + kernarg, so it is not swept here. * **dgrad** -- vector widths: a in {1, 2, 4, 8} (dY along kpg), b = c in {1, 2, 4, 8} (W and dX along cpg); with the default epilogue c = 1 and b in {1, 2, 4, 8}. stride and dilation are baked into the @@ -188,6 +191,15 @@ # a default-epilogue kernel per (mode, reduction) to every wgrad geometry, # which more than doubles the wgrad grid. "atomic" accumulates into dW with # f32 atomics and the validator keeps it to fp32 dW only. +# +# Only a slice of the (mode, reduction) x vector-width grid is cached (see +# _wgrad_streamk_builds). Stream-K pays off where the shape loads at the +# widest vector width, so the requested variants are built there, minus +# "tree" ("linear" is the other deterministic fixup). Narrower widths are left +# to the split-K kernels, which are cached at every width and serve those +# shapes as well -- except for one case: large-filter depthwise layers, whose +# per-group reduction is long and whose GEMM is wide. They load scalar, and +# stream-K wins there, so a scalar dp_sk slice is kept on the wide tiles. CACHE_WGRAD_STREAMK_MODES: Tuple[str, ...] = ("dp_sk", "persistent") CACHE_WGRAD_STREAMK_REDUCTIONS: Tuple[str, ...] = ( "linear", @@ -195,6 +207,15 @@ "workspace", "atomic", ) +# Bytes per stream-K operand load at the widest vector width: 8 elements for +# fp16/bf16, 4 for fp32. +CACHE_WGRAD_STREAMK_LOAD_BYTES = 16 +# Reductions not cached at the widest vector widths. +CACHE_WGRAD_STREAMK_WIDE_EXCLUDED_REDUCTIONS: Tuple[str, ...] = ("tree",) +# The scalar (vector width 1) slice for large-filter depthwise layers. +CACHE_WGRAD_STREAMK_SCALAR_MODES: Tuple[str, ...] = ("dp_sk",) +CACHE_WGRAD_STREAMK_SCALAR_REDUCTIONS: Tuple[str, ...] = ("tree", "workspace") +CACHE_WGRAD_STREAMK_SCALAR_MIN_TILE_N = 128 # Private aliases for internal use -- the generators below refer to these. _TILE_MN = CACHE_TILE_MN @@ -454,6 +475,41 @@ def _wgrad_streamk_variants( return list(itertools.product(streamk_modes, streamk_reductions)) +def _wgrad_streamk_builds( + epilogue: str, + tile_n: int, + dtype: str, + streamk_modes: Sequence[str], + streamk_reductions: Sequence[str], +) -> List[Tuple[str, str, int, int]]: + """``(mode, reduction, vector_size_a, vector_size_b)`` built per geometry. + + Of the requested variants: every one but the + ``CACHE_WGRAD_STREAMK_WIDE_EXCLUDED_REDUCTIONS`` at the widest width W of + the dtype on both operands, and at W on dY with a scalar X for the + narrow-channel layers (C = 3 stems) that cannot load X wide; plus the + ``CACHE_WGRAD_STREAMK_SCALAR_*`` slice, fully scalar, on tiles at least + ``CACHE_WGRAD_STREAMK_SCALAR_MIN_TILE_N`` wide. Every other combination + falls to the split-K kernels. + """ + variants = _wgrad_streamk_variants(epilogue, streamk_modes, streamk_reductions) + w = CACHE_WGRAD_STREAMK_LOAD_BYTES // {"fp16": 2, "bf16": 2, "fp32": 4}[dtype] + builds = [ + (mode, reduction, vec_a, vec_b) + for mode, reduction in variants + if reduction not in CACHE_WGRAD_STREAMK_WIDE_EXCLUDED_REDUCTIONS + for vec_a, vec_b in ((w, w), (w, 1)) + ] + if tile_n >= CACHE_WGRAD_STREAMK_SCALAR_MIN_TILE_N: + builds += [ + (mode, reduction, 1, 1) + for mode, reduction in variants + if mode in CACHE_WGRAD_STREAMK_SCALAR_MODES + and reduction in CACHE_WGRAD_STREAMK_SCALAR_REDUCTIONS + ] + return builds + + def _wgrad_jobs( arch, dtype, @@ -538,10 +594,10 @@ def _wgrad_jobs( # scalar and B takes its own width; split_k stays 1 (stream-K owns the # K partition). The workspace reduction reuses the two-stage scratch, # so its replica count is pinned exactly as for two-stage. - for (mode, reduction), vec_a, vec_b, caps in itertools.product( - _wgrad_streamk_variants(epilogue, streamk_modes, streamk_reductions), - _VECS, - _VECS, + for (mode, reduction, vec_a, vec_b), caps in itertools.product( + _wgrad_streamk_builds( + epilogue, tile_n, dtype, streamk_modes, streamk_reductions + ), _layout_caps(), ): cfg = dict( @@ -861,12 +917,9 @@ def count_jobs( if "wgrad" in directions and _aot_pipeline("wgrad", pipeline): total += vecs * sum(len(_wgrad_two_stages(sk)) for sk in split_ks) * layouts total += merged - total += ( - len(_VECS) - * len(_VECS) - * layouts - * len( - _wgrad_streamk_variants(geo[7], streamk_modes, streamk_reductions) + total += layouts * len( + _wgrad_streamk_builds( + geo[7], geo[1], dtype, streamk_modes, streamk_reductions ) ) if "dgrad" in directions and _aot_pipeline("dgrad", pipeline): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md index 9b9fd921e795..a05d54d52ff9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/README_conv_implicit_gemm_wgrad.md @@ -332,9 +332,12 @@ measure it, use the implicit-GEMM sweep optionally narrowed with `--streamk-reductions` and pinned to a pool with `--streamk-ctas` (default `-1`, the per-reduction pool above). The JIT sweep then adds a stream-K leg, one kernel per geometry x pipeline x mode x reduction -(`atomic` only for fp32 dW); `--compile-all --streamk ...` adds the same kernels -to the AOT cache (`kernel_sweep.py`, opt-in because it more than doubles the -wgrad grid) and `--run-from-cache --streamk ...` runs the cached ones. Every +(`atomic` only for fp32 dW); `--compile-all --streamk ...` adds stream-K kernels +to the AOT cache (`kernel_sweep.py`, opt-in) and `--run-from-cache --streamk ...` +runs the cached ones. The cache keeps a subset: the widest vector width of the +dtype without `tree`, plus a scalar `dp_sk` `tree`/`workspace` slice on wide +tiles for large-filter depthwise layers; other narrow-channel shapes fall to the +split-K kernels. Every launch follows the zeroing contract above, and `workspace` launches are timed together with their Stage 2 reduce.