diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/data.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/data.cpp index 1c950b67ff6a..3644634e3bc6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/data.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/data.cpp @@ -1017,6 +1017,8 @@ static const rocke_ati_mma_frag_row_t rocke_ati_mma_frag[] = { {"wmma_gfx1250_f32_16x16x64_fp8_bf8", 8}, {"wmma_gfx1250_f32_16x16x64_bf8_fp8", 8}, {"wmma_gfx1250_f32_16x16x64_bf8_bf8", 8}, + {"wmma_scale_f32_16x16x128_fp8_fp8", 8}, + {"wmma_scale16_f32_16x16x128_fp8_fp8", 8}, }; int rocke_arch_mma_c_frag_len(const char* op_id) @@ -1799,6 +1801,36 @@ static const rocke_mma_op_t k_mma_gfx1250[] = { NULL, NULL, &lm_wmma_gfx12_c}, + {"wmma_scale", + "fp8", + "fp8", + "fp32", + 16, + 16, + 128, + "wmma_scale_f32_16x16x128_fp8_fp8", + 16, + 16, + 8, + 32, + NULL, + NULL, + &lm_wmma_gfx12_c}, + {"wmma_scale16", + "fp8", + "fp8", + "fp32", + 16, + 16, + 128, + "wmma_scale16_f32_16x16x128_fp8_fp8", + 16, + 16, + 8, + 32, + NULL, + NULL, + &lm_wmma_gfx12_c}, }; /* ========================================================================= diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_tile.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_tile.cpp index 74a0801cb32e..358e870d3ef7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_tile.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_tile.cpp @@ -49,6 +49,8 @@ static const rocke_mma_hint_row_t ROCKE_MMA_RESULT_HINT[] = { {"mfma_f32_16x16x96_fp6", "acc6"}, {"mfma_f32_16x16x128_fp8", "acc128"}, {"mfma_scale_f32_16x16x128_f8f6f4", "mxacc"}, + {"wmma_scale_f32_16x16x128_fp8_fp8", "mxacc"}, + {"wmma_scale16_f32_16x16x128_fp8_fp8", "mxacc"}, }; /* Accumulator fragment length for op_id, from the arch SSOT @@ -651,6 +653,32 @@ rocke_value_t* rocke_b_mfma_scale_f32_16x16x128_f8f6f4(rocke_ir_builder_t* b, return rocke_b_mma(b, "mfma_scale_f32_16x16x128_f8f6f4", a, bb, c, extra, 2); } +rocke_value_t* rocke_b_wmma_scale_f32_16x16x128_fp8_fp8(rocke_ir_builder_t* b, + rocke_value_t* a, + rocke_value_t* bb, + rocke_value_t* c, + rocke_value_t* a_scale, + rocke_value_t* b_scale) +{ + rocke_value_t* extra[2]; + extra[0] = a_scale; + extra[1] = b_scale; + return rocke_b_mma(b, "wmma_scale_f32_16x16x128_fp8_fp8", a, bb, c, extra, 2); +} + +rocke_value_t* rocke_b_wmma_scale16_f32_16x16x128_fp8_fp8(rocke_ir_builder_t* b, + rocke_value_t* a, + rocke_value_t* bb, + rocke_value_t* c, + rocke_value_t* a_scale, + rocke_value_t* b_scale) +{ + rocke_value_t* extra[2]; + extra[0] = a_scale; + extra[1] = b_scale; + return rocke_b_mma(b, "wmma_scale16_f32_16x16x128_fp8_fp8", a, bb, c, extra, 2); +} + /* ===================================================================== */ /* register-fragment reshape (P13) */ /* ===================================================================== */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp index 54ab0451d157..02cb5a9176cd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp @@ -76,11 +76,12 @@ const char* const ROCKE_HIP_PROLOGUE "_ROCKE_VEC(float, f32x, 1); _ROCKE_VEC(float, f32x, 2); _ROCKE_VEC(float, f32x, 4);\n" "_ROCKE_VEC(float, f32x, 8); _ROCKE_VEC(float, f32x, 16);\n" "_ROCKE_VEC(int, i32x, 1); _ROCKE_VEC(int, i32x, 2); _ROCKE_VEC(int, i32x, 3);\n" - "_ROCKE_VEC(int, i32x, 4); _ROCKE_VEC(int, i32x, 8);\n" + "_ROCKE_VEC(int, i32x, 4); _ROCKE_VEC(int, i32x, 8); _ROCKE_VEC(int, i32x, 16);\n" "_ROCKE_VEC(int16_t, i16x, 1); _ROCKE_VEC(int16_t, i16x, 2);\n" "_ROCKE_VEC(int16_t, i16x, 4); _ROCKE_VEC(int16_t, i16x, 8);\n" "_ROCKE_VEC(int8_t, i8x, 1); _ROCKE_VEC(int8_t, i8x, 2);\n" "_ROCKE_VEC(int8_t, i8x, 4); _ROCKE_VEC(int8_t, i8x, 8); _ROCKE_VEC(int8_t, i8x, 16);\n" + "_ROCKE_VEC(int8_t, i8x, 32); _ROCKE_VEC(int8_t, i8x, 48); _ROCKE_VEC(int8_t, i8x, 64);\n" "_ROCKE_VEC(bool, boolx, 2); _ROCKE_VEC(bool, boolx, 4); _ROCKE_VEC(bool, boolx, 8);\n" "_ROCKE_VEC(bool, boolx, 16);\n" "#undef _ROCKE_VEC\n" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_cast.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_cast.cpp index 33b684bc8f36..67b6cf66bbf8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_cast.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_cast.cpp @@ -32,16 +32,41 @@ static const rocke_value_t* h_res(const rocke_op_t* op) /* ----------------------------- arith: int width casts ---------------------- */ -/* def _op_arith_zext(self, op): - * (v,) = op.operands - * self._emit(f"{_type_to_hip(op.result.type)} {_name(op.result)} = " - * f"({_type_to_hip(op.result.type)}){_name(v)};") */ +static const char* h_unsigned_int_type(const rocke_type_t* t) +{ + const char* name = t ? t->name : NULL; + if(!name) + return NULL; + if(__builtin_strcmp(name, "i1") == 0) + return "bool"; + if(__builtin_strcmp(name, "i8") == 0) + return "uint8_t"; + if(__builtin_strcmp(name, "i16") == 0) + return "uint16_t"; + if(__builtin_strcmp(name, "i32") == 0) + return "uint32_t"; + if(__builtin_strcmp(name, "i64") == 0) + return "uint64_t"; + return NULL; +} + +/* Zero extension must cast through the unsigned spelling of the source type. + * A direct int8_t -> int/int64_t cast sign-extends bytes 0x80..0xff. */ static rocke_status_t rocke_h_op_arith_zext(rocke_h_lowerer_t* lw, const rocke_op_t* op) { const rocke_value_t* v = op->operands[0]; const rocke_value_t* r = h_res(op); const char* t = rocke_h_type_to_hip(lw, r->type); - rocke_h_emitf(lw, "%s %s = (%s)%s;", t, rocke_h_name(lw, r), t, rocke_h_name(lw, v)); + const char* source_unsigned = h_unsigned_int_type(v->type); + if(!source_unsigned) + { + return rocke_h_fail(lw, + ROCKE_ERR_NOTIMPL, + "HIP zext requires an integer scalar source, got '%s'", + v->type ? v->type->name : "(null)"); + } + rocke_h_emitf( + lw, "%s %s = (%s)(%s)%s;", t, rocke_h_name(lw, r), t, source_unsigned, rocke_h_name(lw, v)); return lw->status; } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mma.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mma.cpp index 3d9b98248ebf..959d67dd1bf9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mma.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mma.cpp @@ -75,6 +75,38 @@ static const char* h_elem_scalar(const rocke_type_t* t) * op.result (and the WMMA gate keys off the op_id *string* it is passed, not * op.attrs), so a synthetic op aliasing the same operands/results/regions and * reusing the original attrs map reproduces the Python emission exactly. */ +static rocke_status_t + h_emit_gfx1250_scaled_wmma(rocke_h_lowerer_t* lw, const rocke_op_t* op, bool scale16) +{ + const char* op_id + = scale16 ? "wmma_scale16_f32_16x16x128_fp8_fp8" : "wmma_scale_f32_16x16x128_fp8_fp8"; + const char* builtin = scale16 ? "__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4" + : "__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4"; + if(!lw->arch.gfx || __builtin_strcmp(lw->arch.gfx, "gfx1250") != 0) + { + return rocke_h_fail(lw, + ROCKE_ERR_NOTIMPL, + "WMMA op '%s' is not available on %s", + op_id, + lw->arch.gfx ? lw->arch.gfx : "(unknown)"); + } + if(op->num_operands != 5 || op->num_results != 1) + { + return rocke_h_fail(lw, ROCKE_ERR_VALUE, "%s expects 5 operands and 1 result", op_id); + } + rocke_h_emitf(lw, + "f32x8 %s = %s(0, %s, 0, %s, (int16_t)0, %s, " + "0, 0, %s, 0, 0, %s, false, false);", + rocke_h_name(lw, op->results[0]), + builtin, + rocke_h_name(lw, op->operands[0]), + rocke_h_name(lw, op->operands[1]), + rocke_h_name(lw, op->operands[2]), + rocke_h_name(lw, op->operands[3]), + rocke_h_name(lw, op->operands[4])); + return lw->status; +} + static rocke_status_t rocke_h_op_tile_mma(rocke_h_lowerer_t* lw, const rocke_op_t* op) { const char* op_id; @@ -91,6 +123,14 @@ static rocke_status_t rocke_h_op_tile_mma(rocke_h_lowerer_t* lw, const rocke_op_ { return rocke_h_fail(lw, ROCKE_ERR_KEY, "tile.mma: missing 'op_id' attr"); } + if(__builtin_strcmp(op_id, "wmma_scale_f32_16x16x128_fp8_fp8") == 0) + { + return h_emit_gfx1250_scaled_wmma(lw, op, false); + } + if(__builtin_strcmp(op_id, "wmma_scale16_f32_16x16x128_fp8_fp8") == 0) + { + return h_emit_gfx1250_scaled_wmma(lw, op, true); + } snprintf(dotted, sizeof(dotted), "tile.%s", op_id); legacy_opcode = rocke_opcode_from_name(dotted); fn = rocke_h_dispatch(legacy_opcode); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/data.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/data.cpp index d57c49f577b5..1ee6248e7f8d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/data.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/data.cpp @@ -213,6 +213,14 @@ const rocke_ll_decl_t ROCKE_LL_INTRINSIC_DECLS[] = { {"wmma.gfx1250.f32.16x16x64.bf8.bf8", "declare <8 x float> @llvm.amdgcn.wmma.f32.16x16x64.bf8.bf8.v8f32.v8i32(<8 x i32>, <8 x i32>, " "i16 immarg, <8 x float>, i1 immarg, i1 immarg)"}, + {"wmma.scale.gfx1250.f32.16x16x128.fp8.fp8", + "declare <8 x float> @llvm.amdgcn.wmma.scale.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32(i32 " + "immarg, <16 x i32>, i32 immarg, <16 x i32>, i16 immarg, <8 x float>, i32 immarg, i32 immarg, " + "i32, i32 immarg, i32 immarg, i32, i1 immarg, i1 immarg)"}, + {"wmma.scale16.gfx1250.f32.16x16x128.fp8.fp8", + "declare <8 x float> @llvm.amdgcn.wmma.scale16.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32(i32 " + "immarg, <16 x i32>, i32 immarg, <16 x i32>, i16 immarg, <8 x float>, i32 immarg, i32 immarg, " + "i64, i32 immarg, i32 immarg, i64, i1 immarg, i1 immarg)"}, {"mfma.f32.16x16x16f16", "declare <4 x float> @llvm.amdgcn.mfma.f32.16x16x16f16(<4 x half>, <4 x half>, <4 x float>, " "i32 immarg, i32 immarg, i32 immarg)"}, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mma.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mma.cpp index 636b1bd6daa1..0c9c79e1c3d5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mma.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mma.cpp @@ -35,6 +35,7 @@ namespace ckc static void _op_tile_wmma_f32_16x16x16_f16(rocke_lower_t* L, const rocke_op_t* op); static void _op_tile_wmma_f32_16x16x16_bf16(rocke_lower_t* L, const rocke_op_t* op); static void _emit_wmma(rocke_lower_t* L, const rocke_op_t* op, const char* op_id); +static void _emit_wmma_scale(rocke_lower_t* L, const rocke_op_t* op, bool scale16); static void _op_tile_mma(rocke_lower_t* L, const rocke_op_t* op); static void _op_tile_mfma_f32_16x16x32_fp8(rocke_lower_t* L, const rocke_op_t* op); static void _op_tile_mfma_f32_16x16x32_bf8(rocke_lower_t* L, const rocke_op_t* op); @@ -133,6 +134,14 @@ static void _op_tile_mma(rocke_lower_t* L, const rocke_op_t* op) * reject (Python ISABackend.emit_wmma raises NotImplementedError). The * gfx12-specific op_ids ("wmma_gfx12_*") can only occur on RDNA4. */ } + else if(strcmp(op_id, "wmma_scale_f32_16x16x128_fp8_fp8") == 0) + { + _emit_wmma_scale(L, op, false); + } + else if(strcmp(op_id, "wmma_scale16_f32_16x16x128_fp8_fp8") == 0) + { + _emit_wmma_scale(L, op, true); + } else if(strncmp(op_id, "wmma_", 5) == 0) { if(L->backend && L->backend->kind == ROCKE_LL_ISA_RDNA) @@ -164,6 +173,84 @@ static void _op_tile_mma(rocke_lower_t* L, const rocke_op_t* op) } } +/* gfx1250 native MX FP8 WMMA. ROCm 7.13 introduced the LLVM 23 ABI: + * matrix operands are <16 x i32>; SCALE carries four packed E8M0 bytes in + * each i32 scale operand and SCALE16 carries eight in i64. */ +static void _emit_wmma_scale(rocke_lower_t* L, const rocke_op_t* op, bool scale16) +{ + const char* intrinsic; + const char* decl_key; + const char* scale_ty; + const char* op_name; + + if(!rocke_ll_live(L)) + { + return; + } + op_name = scale16 ? "tile.wmma_scale16_f32_16x16x128_fp8_fp8" + : "tile.wmma_scale_f32_16x16x128_fp8_fp8"; + if(!L->backend || strcmp(L->backend->gfx, "gfx1250") != 0) + { + rocke_ll_fail(L, + ROCKE_ERR_NOTIMPL, + "%s is only available on gfx1250 (got %s)", + op_name, + L->backend ? L->backend->gfx : "(unknown)"); + } + if(L->flavor != ROCKE_LLVM_FLAVOR_LLVM23) + { + rocke_ll_fail(L, + ROCKE_ERR_NOTIMPL, + "%s requires llvm23 (ROCm 7.13+), got %s", + op_name, + rocke_llvm_flavor_name(L->flavor)); + } + if(op->num_operands != 5) + { + rocke_ll_fail( + L, ROCKE_ERR_VALUE, "%s expects 5 operands, got %d", op_name, op->num_operands); + } + + scale_ty = scale16 ? "i64" : "i32"; + if(strcmp(op->operands[3]->type->name, scale_ty) != 0 + || strcmp(op->operands[4]->type->name, scale_ty) != 0) + { + rocke_ll_fail(L, + ROCKE_ERR_VALUE, + "%s expects %s scale operands, got %s/%s", + op_name, + scale_ty, + op->operands[3]->type->name, + op->operands[4]->type->name); + } + + if(scale16) + { + decl_key = "wmma.scale16.gfx1250.f32.16x16x128.fp8.fp8"; + intrinsic = "llvm.amdgcn.wmma.scale16.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32"; + } + else + { + decl_key = "wmma.scale.gfx1250.f32.16x16x128.fp8.fp8"; + intrinsic = "llvm.amdgcn.wmma.scale.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32"; + } + rocke_ll_need(L, decl_key); + rocke_ll_emitf(L, + " %s = call <8 x float> @%s(" + "i32 0, <16 x i32> %s, i32 0, <16 x i32> %s, " + "i16 0, <8 x float> %s, i32 0, i32 0, %s %s, " + "i32 0, i32 0, %s %s, i1 false, i1 false)", + mma_result_name(L, op), + intrinsic, + rocke_ll_operand(L, op->operands[0]), + rocke_ll_operand(L, op->operands[1]), + rocke_ll_operand(L, op->operands[2]), + scale_ty, + rocke_ll_operand(L, op->operands[3]), + scale_ty, + rocke_ll_operand(L, op->operands[4])); +} + /* ====================================================================== */ /* WMMA on CDNA: faithful rejection (Python ISABackend.emit_wmma raises) */ /* ====================================================================== */ 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 10f731c13e0b..5e36e56c1b51 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 @@ -1052,6 +1052,18 @@ rocke_value_t* rocke_b_mfma_scale_f32_16x16x128_f8f6f4(rocke_ir_builder_t* b, rocke_value_t* c, rocke_value_t* a_scale, rocke_value_t* b_scale); +rocke_value_t* rocke_b_wmma_scale_f32_16x16x128_fp8_fp8(rocke_ir_builder_t* b, + rocke_value_t* a, + rocke_value_t* bb, + rocke_value_t* c, + rocke_value_t* a_scale, + rocke_value_t* b_scale); +rocke_value_t* rocke_b_wmma_scale16_f32_16x16x128_fp8_fp8(rocke_ir_builder_t* b, + rocke_value_t* a, + rocke_value_t* bb, + rocke_value_t* c, + rocke_value_t* a_scale, + rocke_value_t* b_scale); rocke_value_t* rocke_b_wmma_f32_16x16x16_f16(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_t* bb, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/arch_specs.json b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/arch_specs.json index bb838829d485..0607dc02932c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/arch_specs.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/arch_specs.json @@ -142,7 +142,7 @@ ] }, "gfx1250": { - "_comment": "gfx1250-class CDNA multi-chip device on the GFX12 programming model (wave32, WMMA, no MFMA). Primary fp16/bf16 WMMA atom is 16x16x32 (NOT gfx1201's 16x16x16). ABI empirically confirmed on ROCm 7.13 clang 23 (gfx1250 device target): __builtin_amdgcn_wmma_f32_16x16x32_f16 -> llvm.amdgcn.wmma.f32.16x16x32.f16.v8f32.v16f16(i1,<16 x half>,i1,<16 x half>,i16,<8 x float>,i1,i1); bf16 uses v16bf16 (<16 x bfloat>) directly, not the gfx11 <16 x i16> bitcast. Datalayout/triple byte-identical to gfx950/gfx1201. This catalog ships only the validated fp16/bf16 and fp8/bf8 WMMA atoms; block-scaled F4/block16 forms are intentionally omitted. Lane maps are a HYPOTHESIS verified by examples/gfx1250/wmma_probe.py. MEMORY: gfx1250 HAS async global->LDS DMA (its own GFX12 family: global_load_async_to_lds_b128 + s_wait_asynccnt -- NOT the gfx9 buffer_load_lds/global_load_lds, which do not select), so has_async_global_lds=true (GPU-validated via the DTLA path). It also HAS a hardware transpose-LDS read, ds_load_tr16_b128 (GPU-validated via compute_pv_dstr). The has_async_lds and has_ds_read_tr flags stay false on purpose: in this catalog they denote the gfx9/gfx950 ds_read_*_tr_* / async-LDS *ABI* used to admit the gfx950 MFMA tiled builder (attention_arch.py), which gfx1250 does NOT have -- its async/transpose use distinct GFX12 opcodes gated directly in the ISA backend, not via these MFMA-pipeline bits. has_tdm=false: TDM not verified on this part.", + "_comment": "gfx1250-class CDNA multi-chip device on the GFX12 programming model (wave32, WMMA, no MFMA). Primary fp16/bf16 WMMA atom is 16x16x32 (NOT gfx1201's 16x16x16). ABI empirically confirmed on ROCm 7.13 clang 23 (gfx1250 device target): __builtin_amdgcn_wmma_f32_16x16x32_f16 -> llvm.amdgcn.wmma.f32.16x16x32.f16.v8f32.v16f16(i1,<16 x half>,i1,<16 x half>,i16,<8 x float>,i1,i1); bf16 uses v16bf16 (<16 x bfloat>) directly, not the gfx11 <16 x i16> bitcast. Datalayout/triple byte-identical to gfx950/gfx1201. This catalog ships validated fp16/bf16 and fp8/bf8 WMMA atoms plus the FP8xFP8 SCALE/SCALE16 prerequisite slice; FP4/FP6/BF6 block-scaled forms remain intentionally omitted. Lane maps are a HYPOTHESIS verified by examples/gfx1250/wmma_probe.py. MEMORY: gfx1250 HAS async global->LDS DMA (its own GFX12 family: global_load_async_to_lds_b128 + s_wait_asynccnt -- NOT the gfx9 buffer_load_lds/global_load_lds, which do not select), so has_async_global_lds=true (GPU-validated via the DTLA path). It also HAS a hardware transpose-LDS read, ds_load_tr16_b128 (GPU-validated via compute_pv_dstr). The has_async_lds and has_ds_read_tr flags stay false on purpose: in this catalog they denote the gfx9/gfx950 ds_read_*_tr_* / async-LDS *ABI* used to admit the gfx950 MFMA tiled builder (attention_arch.py), which gfx1250 does NOT have -- its async/transpose use distinct GFX12 opcodes gated directly in the ISA backend, not via these MFMA-pipeline bits. has_tdm=false: TDM not verified on this part.", "family": "cdna", "target_family": "gfx12_cdna", "matrix_path": "wmma", @@ -176,7 +176,9 @@ {"family": "wmma", "a": "fp8", "b": "fp8", "c": "fp32", "m": 16, "n": 16, "k": 64, "op_id": "wmma_gfx1250_f32_16x16x64_fp8_fp8"}, {"family": "wmma", "a": "fp8", "b": "bf8", "c": "fp32", "m": 16, "n": 16, "k": 64, "op_id": "wmma_gfx1250_f32_16x16x64_fp8_bf8"}, {"family": "wmma", "a": "bf8", "b": "fp8", "c": "fp32", "m": 16, "n": 16, "k": 64, "op_id": "wmma_gfx1250_f32_16x16x64_bf8_fp8"}, - {"family": "wmma", "a": "bf8", "b": "bf8", "c": "fp32", "m": 16, "n": 16, "k": 64, "op_id": "wmma_gfx1250_f32_16x16x64_bf8_bf8"} + {"family": "wmma", "a": "bf8", "b": "bf8", "c": "fp32", "m": 16, "n": 16, "k": 64, "op_id": "wmma_gfx1250_f32_16x16x64_bf8_bf8"}, + {"family": "wmma_scale", "a": "fp8", "b": "fp8", "c": "fp32", "m": 16, "n": 16, "k": 128, "op_id": "wmma_scale_f32_16x16x128_fp8_fp8"}, + {"family": "wmma_scale16", "a": "fp8", "b": "fp8", "c": "fp32", "m": 16, "n": 16, "k": 128, "op_id": "wmma_scale16_f32_16x16x128_fp8_fp8"} ] }, "gfx11-generic": { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/target.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/target.py index 5254a30e1a8b..fd8f09e839f1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/target.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/target.py @@ -749,6 +749,27 @@ class _FragInfo: None, _wmma_gfx12_acc_16x16, ), + # Native gfx1250 MX FP8 WMMA. A/B each carry 64 bytes as <16 x i32>; + # SCALE packs four K=32 E8M0 bytes in i32 and SCALE16 packs eight K=16 + # bytes in i64. Both share the gfx12 column-distributed accumulator. + "wmma_scale_f32_16x16x128_fp8_fp8": _FragInfo( + 16, + 16, + 8, + 32, + None, + None, + _wmma_gfx12_acc_16x16, + ), + "wmma_scale16_f32_16x16x128_fp8_fp8": _FragInfo( + 16, + 16, + 8, + 32, + None, + None, + _wmma_gfx12_acc_16x16, + ), "wmma_gfx1250_f32_16x16x32_bf16": _FragInfo( 16, 16, 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 074d55a80023..f77f6463f7f2 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 @@ -29,7 +29,6 @@ from dataclasses import dataclass, field from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple - # ----------------------------- Types -------------------------------------- @@ -86,6 +85,8 @@ def __repr__(self) -> str: "mfma_f32_16x16x96_fp6": "acc6", "mfma_f32_16x16x128_fp8": "acc128", "mfma_scale_f32_16x16x128_f8f6f4": "mxacc", + "wmma_scale_f32_16x16x128_fp8_fp8": "mxacc", + "wmma_scale16_f32_16x16x128_fp8_fp8": "mxacc", } @@ -1962,6 +1963,38 @@ def wmma_gfx1250_f32_16x16x64_bf8_bf8(self, a: Value, b: Value, c: Value) -> Val """gfx1250 (gfx1250) BF8 K=64 WMMA. Thin wrapper over :meth:`mma`.""" return self.mma("wmma_gfx1250_f32_16x16x64_bf8_bf8", a, b, c) + def wmma_scale_f32_16x16x128_fp8_fp8( + self, + a: Value, + b: Value, + c: Value, + a_scale: Value, + b_scale: Value, + ) -> Value: + """gfx1250 native SCALE FP8 WMMA with packed E8M0 scale operands. + + A and B are ``<16 x i32>`` fragments (64 FP8 bytes per lane), C is + ``<8 x f32>``, and each scale operand is one i32 packing four E8M0 + bytes for the instruction's four K=32 scale blocks. + """ + return self.mma("wmma_scale_f32_16x16x128_fp8_fp8", a, b, c, a_scale, b_scale) + + def wmma_scale16_f32_16x16x128_fp8_fp8( + self, + a: Value, + b: Value, + c: Value, + a_scale: Value, + b_scale: Value, + ) -> Value: + """gfx1250 native SCALE16 FP8 WMMA with eight packed E8M0 scales. + + The matrix and accumulator fragments match + :meth:`wmma_scale_f32_16x16x128_fp8_fp8`; each scale operand is i64 + because SCALE16 carries eight K=16 E8M0 scale bytes. + """ + return self.mma("wmma_scale16_f32_16x16x128_fp8_fp8", a, b, c, a_scale, b_scale) + def mfma_f32_16x16x16_f16(self, a: Value, b: Value, c: Value) -> Value: return self.mma("mfma_f32_16x16x16_f16", a, b, c) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/isa/backend.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/isa/backend.py index e82d62354dde..8970d949570e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/isa/backend.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/isa/backend.py @@ -332,6 +332,23 @@ class Gfx9MfmaBackend(ISABackend): } +# gfx1250 native MX FP8 WMMA, available with the LLVM 23 toolchain used by +# ROCm 7.13+. Both operations consume <16 x i32> matrix fragments; SCALE packs +# four E8M0 scale bytes in i32 while SCALE16 packs eight in i64. +_GFX1250_WMMA_FP8_SCALE = { + "tile.wmma_scale_f32_16x16x128_fp8_fp8": ( + "wmma.scale.gfx1250.f32.16x16x128.fp8.fp8", + "llvm.amdgcn.wmma.scale.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32", + "i32", + ), + "tile.wmma_scale16_f32_16x16x128_fp8_fp8": ( + "wmma.scale16.gfx1250.f32.16x16x128.fp8.fp8", + "llvm.amdgcn.wmma.scale16.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32", + "i64", + ), +} + + class Gfx11RdnaBackend(ISABackend): """RDNA3 / RDNA3.5 (gfx11, e.g. gfx1151 Strix Halo). **wave32**, **WMMA** (no MFMA), and a distinct ``s_waitcnt`` layout from gfx9/10. Datalayout + @@ -558,6 +575,10 @@ def emit_lds_barrier_drain(self, lowerer, *, drain_vmem: bool) -> None: lowerer._current().emit(" call void @llvm.amdgcn.s.wait.dscnt(i16 0)") def emit_wmma(self, lowerer, op) -> None: + scale_spec = _GFX1250_WMMA_FP8_SCALE.get(op.name) + if scale_spec is not None: + self._emit_wmma_scale(lowerer, op, scale_spec) + return fp8_spec = _GFX1250_WMMA_FP8.get(op.name) if fp8_spec is not None: self._emit_wmma_fp8(lowerer, op, fp8_spec) @@ -566,7 +587,7 @@ def emit_wmma(self, lowerer, op) -> None: if spec is None: raise NotImplementedError( f"WMMA op {op.name!r} not yet wired for {self.arch.gfx}; " - f"known: {sorted(_GFX1250_WMMA) + sorted(_GFX1250_WMMA_FP8)}" + f"known: {sorted(_GFX1250_WMMA) + sorted(_GFX1250_WMMA_FP8) + sorted(_GFX1250_WMMA_FP8_SCALE)}" ) decl_key, intrinsic, elt = spec a, b, c = op.operands @@ -601,6 +622,32 @@ def _emit_wmma_fp8(self, lowerer, op, spec) -> None: f"i1 false, i1 false)" ) + def _emit_wmma_scale(self, lowerer, op, spec) -> None: + """Emit the ROCm 7.13+ gfx1250 SCALE/SCALE16 FP8 call.""" + if lowerer._flavor != "llvm23": + raise NotImplementedError( + f"{op.name} requires llvm23 (ROCm 7.13+), got {lowerer._flavor}" + ) + if len(op.operands) != 5: + raise ValueError(f"{op.name} expects 5 operands, got {len(op.operands)}") + decl_key, intrinsic, scale_ty = spec + a, b, c, a_scale, b_scale = op.operands + if a_scale.type.name != scale_ty or b_scale.type.name != scale_ty: + raise ValueError( + f"{op.name} expects {scale_ty} scale operands, got " + f"{a_scale.type.name}/{b_scale.type.name}" + ) + lowerer._need(decl_key) + lowerer._current().emit( + f" {op.result.name} = call <8 x float> @{intrinsic}(" + f"i32 0, <16 x i32> {lowerer._operand(a)}, " + f"i32 0, <16 x i32> {lowerer._operand(b)}, " + f"i16 0, <8 x float> {lowerer._operand(c)}, " + f"i32 0, i32 0, {scale_ty} {lowerer._operand(a_scale)}, " + f"i32 0, i32 0, {scale_ty} {lowerer._operand(b_scale)}, " + f"i1 false, i1 false)" + ) + # gfx -> backend class. Adding a CDNA gfx is one row here plus, when its codegen # actually diverges, a new subclass. 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 ed8b8a704a7a..bbd4e482c8de 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 @@ -41,6 +41,14 @@ "bf8e5m2": "bf8e5m2", } +_HIP_UNSIGNED_INT_TYPE = { + "i1": "bool", + "i8": "uint8_t", + "i16": "uint16_t", + "i32": "uint32_t", + "i64": "uint64_t", +} + def _type_to_hip(t) -> str: if isinstance(t, PtrType): @@ -112,11 +120,12 @@ def _type_to_hip(t) -> str: _ROCKE_VEC(float, f32x, 1); _ROCKE_VEC(float, f32x, 2); _ROCKE_VEC(float, f32x, 4); _ROCKE_VEC(float, f32x, 8); _ROCKE_VEC(float, f32x, 16); _ROCKE_VEC(int, i32x, 1); _ROCKE_VEC(int, i32x, 2); _ROCKE_VEC(int, i32x, 3); -_ROCKE_VEC(int, i32x, 4); _ROCKE_VEC(int, i32x, 8); +_ROCKE_VEC(int, i32x, 4); _ROCKE_VEC(int, i32x, 8); _ROCKE_VEC(int, i32x, 16); _ROCKE_VEC(int16_t, i16x, 1); _ROCKE_VEC(int16_t, i16x, 2); _ROCKE_VEC(int16_t, i16x, 4); _ROCKE_VEC(int16_t, i16x, 8); _ROCKE_VEC(int8_t, i8x, 1); _ROCKE_VEC(int8_t, i8x, 2); _ROCKE_VEC(int8_t, i8x, 4); _ROCKE_VEC(int8_t, i8x, 8); _ROCKE_VEC(int8_t, i8x, 16); +_ROCKE_VEC(int8_t, i8x, 32); _ROCKE_VEC(int8_t, i8x, 48); _ROCKE_VEC(int8_t, i8x, 64); _ROCKE_VEC(bool, boolx, 2); _ROCKE_VEC(bool, boolx, 4); _ROCKE_VEC(bool, boolx, 8); _ROCKE_VEC(bool, boolx, 16); #undef _ROCKE_VEC @@ -403,8 +412,15 @@ def _op_arith_smin(self, op: Op) -> None: def _op_arith_zext(self, op: Op) -> None: (v,) = op.operands + source_unsigned = _HIP_UNSIGNED_INT_TYPE.get(v.type.name) + if source_unsigned is None: + raise NotImplementedError( + f"HIP zext requires an integer scalar source, got {v.type.name!r}" + ) + target = _type_to_hip(op.result.type) self._emit( - f"{_type_to_hip(op.result.type)} {_name(op.result)} = ({_type_to_hip(op.result.type)}){_name(v)};" + f"{target} {_name(op.result)} = " + f"({target})({source_unsigned}){_name(v)};" ) def _op_arith_sext(self, op: Op) -> None: @@ -636,6 +652,12 @@ def _op_tile_wmma_gfx1250_f32_16x16x64_bf8_fp8(self, op: Op) -> None: def _op_tile_wmma_gfx1250_f32_16x16x64_bf8_bf8(self, op: Op) -> None: self._emit_wmma_gfx1250_fp8(op, "bf8_bf8") + def _op_tile_wmma_scale_f32_16x16x128_fp8_fp8(self, op: Op) -> None: + self._emit_wmma_gfx1250_scaled(op, scale16=False) + + def _op_tile_wmma_scale16_f32_16x16x128_fp8_fp8(self, op: Op) -> None: + self._emit_wmma_gfx1250_scaled(op, scale16=True) + def _emit_wmma_gfx1250_fp8(self, op: Op, ab: str) -> None: # gfx1250 K=64 FP8/BF8 builtin: A/B are <8 x i32> (32 low-bit # bytes per lane), 6-operand form (A, B, fmt, C, reuseA, reuseB). @@ -647,6 +669,25 @@ def _emit_wmma_gfx1250_fp8(self, op: Op, ab: str) -> None: f"{_name(a)}, {_name(b)}, (int16_t)0, {_name(c)}, false, false);" ) + def _emit_wmma_gfx1250_scaled(self, op: Op, *, scale16: bool) -> None: + op_id = ( + "wmma_scale16_f32_16x16x128_fp8_fp8" + if scale16 + else "wmma_scale_f32_16x16x128_fp8_fp8" + ) + self._require_wmma_arch(op_id) + a, b, c, a_scale, b_scale = op.operands + builtin = ( + "__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4" + if scale16 + else "__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4" + ) + self._emit( + f"f32x8 {_name(op.result)} = {builtin}(" + f"0, {_name(a)}, 0, {_name(b)}, (int16_t)0, {_name(c)}, " + f"0, 0, {_name(a_scale)}, 0, 0, {_name(b_scale)}, false, false);" + ) + def _op_tile_wmma_gfx1250_f32_16x16x32_bf16(self, op: Op) -> None: # Same gfx1250 K=32 ABI as the f16 form, but operands are true bf16 vectors. self._require_wmma_arch("wmma_gfx1250_f32_16x16x32_bf16") 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 5c6c4ba32887..e63366632a7e 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 @@ -608,6 +608,18 @@ def _resolve_llvm_flavor() -> str: "declare <8 x float> @llvm.amdgcn.wmma.f32.16x16x64.bf8.bf8.v8f32.v8i32(" "<8 x i32>, <8 x i32>, i16 immarg, <8 x float>, i1 immarg, i1 immarg)" ), + "wmma.scale.gfx1250.f32.16x16x128.fp8.fp8": ( + "declare <8 x float> @llvm.amdgcn.wmma.scale.f32.16x16x128.f8f6f4." + "v8f32.v16i32.v16i32(i32 immarg, <16 x i32>, i32 immarg, " + "<16 x i32>, i16 immarg, <8 x float>, i32 immarg, i32 immarg, i32, " + "i32 immarg, i32 immarg, i32, i1 immarg, i1 immarg)" + ), + "wmma.scale16.gfx1250.f32.16x16x128.fp8.fp8": ( + "declare <8 x float> @llvm.amdgcn.wmma.scale16.f32.16x16x128.f8f6f4." + "v8f32.v16i32.v16i32(i32 immarg, <16 x i32>, i32 immarg, " + "<16 x i32>, i16 immarg, <8 x float>, i32 immarg, i32 immarg, i64, " + "i32 immarg, i32 immarg, i64, i1 immarg, i1 immarg)" + ), "mfma.f32.16x16x16f16": ( "declare <4 x float> @llvm.amdgcn.mfma.f32.16x16x16f16(" "<4 x half>, <4 x half>, <4 x float>, " diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx1250/gemm/block_scaled_gemm_verify.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx1250/gemm/block_scaled_gemm_verify.py index 567277edc711..f703106d894e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx1250/gemm/block_scaled_gemm_verify.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx1250/gemm/block_scaled_gemm_verify.py @@ -1,20 +1,22 @@ # Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -"""Build + numeric-verify the gfx1250 K=64 FP8/BF8 block-scaled WMMA GEMM. +"""Launch gfx1250 block-scaled GEMM and compare with an independent reference. -Computes ``C[m, n] = sum_kg a_scale[m, kg] * b_scale[kg, n] - * sum_{k in group} A[m, k] * B[n, k]`` (RCR, ``C = A @ B^T``) -on a gfx1250 device and compares against a numpy fp32 reference. A/B carry -fp8e4m3 / bf8e5m2 mantissas (encoded via ml_dtypes); per-block f32 scales use -``group_size_k = block_k``. +The default invocation keeps the K=64 FP8/BF8 WMMA + FP32-scale verifier. +Native ``--matrix-path wmma_scale`` / ``wmma_scale16`` use FP8 and E8M0 +scales with K=32 / K=16 groups. Native fixtures cover K=128 or 256 and use +bounded dyadic values, permitting exact comparison after BF16 rounding. -A wrong K=64 fragment lane map produces a transposed / permuted result and -fails verify, so a PASS at multiple shapes confirms the mapping. +Run on visible HIP device 0 (must be gfx1250), for example:: -Must run on a gfx1250 device: + python -m rocke.examples.gfx1250.gemm.block_scaled_gemm_verify \ + --matrix-path wmma_scale16 --case all + python -m rocke.examples.gfx1250.gemm.block_scaled_gemm_verify \ + --matrix-path wmma_scale --m 32 --n 48 --k 256 --compile-route hip - PYTHONPATH=python python3 -m rocke.examples.gfx1250.gemm.block_scaled_gemm_verify \ - --m 16 --n 16 --k 128 --dtype fp8e4m3 +Native COMGR lowering requires ROCKE_LLVM_FLAVOR=llvm23 and a matching +compiler. HIP verification requires a hipcc supporting the scaled builtins. +NumPy and ml_dtypes are required; torch is not used. """ from __future__ import annotations @@ -23,137 +25,254 @@ import ctypes import struct -from rocke.helpers import compile_kernel -from rocke.instances.gfx1250.block_scaled_gemm import ( +import numpy as np + +from ....helpers import compile_kernel +from ....helpers.compile import compile_kernel_via_hipcc +from ....instances.gfx1250.block_scaled_gemm import ( BlockScaledGemmSpec, block_scaled_gemm_grid, build_block_scaled_gemm, + is_valid_spec, ) -from rocke.runtime.hip_module import Runtime +from ....runtime.hip_module import Runtime, get_device_arch + + +def decode_e8m0(encoded: np.ndarray) -> np.ndarray: + """Decode finite unsigned E8M0 bytes without using kernel packing logic.""" + if encoded.dtype != np.uint8 or np.any(encoded == 0xFF): + raise ValueError("expected finite uint8 E8M0 scales (0xff encodes NaN)") + return np.ldexp(np.ones(encoded.shape), encoded.astype(np.int32) - 127) -def _ml_fp8_dtype(np_mod, dtype: str): +def reference_result( + a: np.ndarray, + b: np.ndarray, + a_scale: np.ndarray, + b_scale: np.ndarray, + block_k: int, + *, + native: bool, +) -> np.ndarray: + """Expand scales onto logical A/B elements, multiply, then round to BF16. + + The native fixtures use quarter-integer inputs of magnitude <=1 and + scales 2**[-2,3]. At K<=256, even the sum of absolute products fits in + 2**22 units of 2**-8, so every FP32 partial sum is exact. Float64 host + arithmetic and a single BF16 rounding provide an independent oracle. + """ import ml_dtypes - if dtype == "fp8e4m3": - return ml_dtypes.float8_e4m3fn - if dtype == "bf8e5m2": - return ml_dtypes.float8_e5m2 - raise SystemExit(f"unsupported low-bit dtype {dtype!r}") + sa = decode_e8m0(a_scale) if native else a_scale.astype(np.float64) + sb = decode_e8m0(b_scale) if native else b_scale.astype(np.float64) + scaled_a = a.astype(np.float64) * np.repeat(sa, block_k, axis=1) + scaled_b = b.astype(np.float64) * np.repeat(sb.T, block_k, axis=1) + ref = scaled_a @ scaled_b.T + if native: + ref = ref.astype(ml_dtypes.bfloat16) + return ref.astype(np.float32) -def main() -> int: - from rocke.runtime.comgr import prefer_bundled_lib +def make_case_inputs( + spec: BlockScaledGemmSpec, case: str +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Create asymmetric mantissas and scales, optionally isolating a K group.""" + import ml_dtypes - prefer_bundled_lib() # pin newest comgr/LLVM flavor before lowering (gfx1250 needs ROCm>=7.2) - p = argparse.ArgumentParser(description=__doc__) - p.add_argument("--arch", default="gfx1250") - p.add_argument("--m", type=int, default=16) - p.add_argument("--n", type=int, default=16) - p.add_argument("--k", type=int, default=128) - p.add_argument("--block-k", type=int, default=128) - p.add_argument("--dtype", default="fp8e4m3", choices=("fp8e4m3", "bf8e5m2")) - p.add_argument("--tol", type=float, default=2e-2) - args = p.parse_args() - - import numpy as np - - for d, name, mult in ( - (args.m, "M", 16), - (args.n, "N", 16), - (args.k, "K", args.block_k), - ): - if d % mult: - raise SystemExit(f"{name}={d} must be a multiple of {mult}") - if args.block_k % 64: - raise SystemExit("block_k must be a multiple of 64") - - M, N, K, BK = args.m, args.n, args.k, args.block_k - groups = K // BK - fp8 = _ml_fp8_dtype(np, args.dtype) + native = spec.resolved_matrix_path() != "wmma" + if native and spec.K not in (128, 256): + raise ValueError("exact native fixtures require K=128 or K=256") + groups = spec.K // spec.block_k + lowbit_types = { + "fp8": ml_dtypes.float8_e4m3fn, + "fp8e4m3": ml_dtypes.float8_e4m3fn, + "bf8": ml_dtypes.float8_e5m2, + "bf8e5m2": ml_dtypes.float8_e5m2, + } rng = np.random.default_rng(0xB10C) + magnitude = 0.25 if native else 0.5 + a = (rng.integers(-4, 5, size=(spec.M, spec.K)) * magnitude).astype( + lowbit_types[spec.dtype_a] + ) + b = (rng.integers(-4, 5, size=(spec.N, spec.K)) * magnitude).astype( + lowbit_types[spec.dtype_b] + ) + if native: + sa = rng.integers(125, 131, size=(spec.M, groups), dtype=np.uint8) + sb = rng.integers(125, 131, size=(groups, spec.N), dtype=np.uint8) + neutral = 127 + else: + sa = rng.uniform(0.5, 1.5, size=(spec.M, groups)).astype(np.float32) + sb = rng.uniform(0.5, 1.5, size=(groups, spec.N)).astype(np.float32) + neutral = 1.0 + if case in ("neutral", "b-only"): + sa.fill(neutral) + if case in ("neutral", "a-only"): + sb.fill(neutral) + if case.startswith("group-"): + group = int(case.removeprefix("group-")) + if not 0 <= group < groups: + raise ValueError(f"scale group {group} outside [0, {groups})") + active = (np.arange(spec.K) // spec.block_k) == group + a[:, ~active] = 0 + b[:, ~active] = 0 + elif case not in ("neutral", "a-only", "b-only", "mixed"): + raise ValueError(f"unknown verification case {case!r}") + return a, b, sa, sb - # Small, mostly fp8-representable magnitudes so quantisation noise is tiny. - A_f = (rng.integers(-4, 5, size=(M, K)).astype(np.float32)) * 0.5 - B_f = (rng.integers(-4, 5, size=(N, K)).astype(np.float32)) * 0.5 - A_q = A_f.astype(fp8) - B_q = B_f.astype(fp8) - A_dq = A_q.astype(np.float32) - B_dq = B_q.astype(np.float32) - - A_scale = (rng.uniform(0.5, 1.5, size=(M, groups))).astype(np.float32) - B_scale = (rng.uniform(0.5, 1.5, size=(groups, N))).astype(np.float32) - # Reference: per-group scaled accumulation. - ref = np.zeros((M, N), dtype=np.float32) - for kg in range(groups): - sl = slice(kg * BK, (kg + 1) * BK) - partial = A_dq[:, sl] @ B_dq[:, sl].T # (M, N) - ref += partial * np.outer(A_scale[:, kg], B_scale[kg, :]) +def check_result( + got: np.ndarray, expected: np.ndarray, *, exact: bool, tol: float = 2e-2 +) -> None: + """Fail on incomplete/non-finite output or a numerical mismatch.""" + if got.shape != expected.shape: + raise AssertionError(f"output shape {got.shape} != {expected.shape}") + if not np.isfinite(expected).all() or not np.isfinite(got).all(): + raise AssertionError("unexpected non-finite output or reference") + diff = np.abs(got.astype(np.float64) - expected.astype(np.float64)) + bad = got != expected if exact else diff > tol * np.maximum(np.abs(expected), 1) + if np.any(bad): + worst = np.unravel_index(int(np.argmax(diff)), diff.shape) + raise AssertionError( + f"bad={np.count_nonzero(bad)}/{diff.size}; worst={worst}: " + f"got={got[worst]}, expected={expected[worst]}, diff={diff[worst]}" + ) - spec = BlockScaledGemmSpec( - name=f"verify_{args.arch}", - M=M, - N=N, - K=K, - dtype_a=args.dtype, - dtype_b=args.dtype, - dtype_c="bf16", - scale_dtype="fp32", - block_k=BK, - ) - art = compile_kernel(build_block_scaled_gemm(spec, arch=args.arch), arch=args.arch) - print(f"[{args.arch}] built {art.kernel_name} ({art.hsaco_bytes} B, isa={art.isa})") - rt = Runtime() - module = rt.load_module(art.hsaco) - fn = module.get_function(art.kernel_name) +def _u8_buffer(array: np.ndarray): + array = np.ascontiguousarray(array) + return (ctypes.c_uint8 * int(array.nbytes)).from_buffer_copy(array) - def u8(a): - a = np.ascontiguousarray(a) - return (ctypes.c_uint8 * int(a.nbytes)).from_buffer_copy(a) +def _launch(rt, fn, spec, inputs): import ml_dtypes - A_bytes = A_q.view(np.uint8) - B_bytes = B_q.view(np.uint8) - out = np.zeros((M, N), dtype=ml_dtypes.bfloat16) - - ad = rt.alloc(A_bytes.nbytes) - bd = rt.alloc(B_bytes.nbytes) - asd = rt.alloc(A_scale.nbytes) - bsd = rt.alloc(B_scale.nbytes) - cd = rt.alloc(out.nbytes) - rt.memcpy_h2d(ad, u8(A_bytes), A_bytes.nbytes) - rt.memcpy_h2d(bd, u8(B_bytes), B_bytes.nbytes) - rt.memcpy_h2d(asd, u8(A_scale), A_scale.nbytes) - rt.memcpy_h2d(bsd, u8(B_scale), B_scale.nbytes) - rt.memset(cd, 0, out.nbytes) - - grid = block_scaled_gemm_grid(spec) - block = (spec.block_size, 1, 1) - packed = struct.pack(" int: + """Compile once and launch each case; return the verified case count.""" + ok, reason = is_valid_spec(spec, arch="gfx1250") + if not ok: + raise ValueError(reason) + if spec.dtype_c != "bf16": + raise ValueError("this verifier requires BF16 output") + native = spec.resolved_matrix_path() != "wmma" + if native and spec.K not in (128, 256): + raise ValueError("exact native fixtures require K=128 or K=256") + if not cases: + raise ValueError("at least one verification case is required") + if compile_route not in ("comgr", "hip"): + raise ValueError(f"unknown compile route {compile_route!r}") + if not np.isfinite(tol) or tol < 0: + raise ValueError("tol must be finite and nonnegative") + arch = get_device_arch(0) + if arch != "gfx1250": + raise RuntimeError(f"visible HIP device 0 must be gfx1250, got {arch!r}") + kernel = build_block_scaled_gemm(spec, arch=arch) + compile_fn = ( + compile_kernel if compile_route == "comgr" else compile_kernel_via_hipcc + ) + art = compile_fn(kernel, arch=arch) print( - f"[{args.arch}] block-scaled {args.dtype} {M}x{N}x{K} bk{BK} groups={groups}: " - f"max_abs={max_abs:.3e} max_rel={max_rel:.3e} tol={args.tol:.0e} -> {tag}" + f"[{arch}/{compile_route}] compiled {art.kernel_name} isa={art.isa}", flush=True + ) + rt = Runtime() + module = rt.load_module(art.hsaco) + try: + fn = module.get_function(art.kernel_name) + for case in cases: + inputs = make_case_inputs(spec, case) + expected = reference_result(*inputs, spec.block_k, native=native) + label = ( + f"{spec.resolved_matrix_path()}/{compile_route}/{case} " + f"{spec.M}x{spec.N}x{spec.K} bk{spec.block_k}" + ) + got = _launch(rt, fn, spec, inputs) + try: + check_result(got, expected, exact=native, tol=tol) + except AssertionError as exc: + raise AssertionError(f"{label}: {exc}") from exc + print(f"PASS: {label} bad=0", flush=True) + finally: + module.unload() + return len(cases) + + +def main(argv: list[str] | None = None) -> int: + """Run the legacy verifier or a bounded native SCALE/SCALE16 case set.""" + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--arch", default="gfx1250", choices=("gfx1250",)) + p.add_argument("--m", type=int, default=16) + p.add_argument("--n", type=int, default=16) + p.add_argument("--k", type=int, default=128) + p.add_argument("--block-k", type=int, default=None) + p.add_argument("--dtype", default="fp8e4m3", choices=("fp8e4m3", "bf8e5m2")) + p.add_argument("--tol", type=float, default=2e-2, help="legacy WMMA tolerance only") + p.add_argument( + "--matrix-path", default="wmma", choices=("wmma", "wmma_scale", "wmma_scale16") + ) + p.add_argument("--compile-route", default="comgr", choices=("comgr", "hip")) + p.add_argument( + "--case", + default="mixed", + help="mixed, neutral, a-only, b-only, group-N, or all (including every K group)", + ) + args = p.parse_args(argv) + native = args.matrix_path != "wmma" + block_k = args.block_k + if block_k is None: + block_k = {"wmma": 128, "wmma_scale": 32, "wmma_scale16": 16}[args.matrix_path] + spec = BlockScaledGemmSpec( + name="verify_gfx1250", + M=args.m, + N=args.n, + K=args.k, + dtype_a=args.dtype, + dtype_b=args.dtype, + dtype_c="bf16", + scale_dtype="e8m0" if native else "fp32", + block_k=block_k, + matrix_path=args.matrix_path, + ) + ok, reason = is_valid_spec(spec, arch=args.arch) + if not ok: + p.error(reason) + cases = ( + ("neutral", "a-only", "b-only", "mixed") + + tuple(f"group-{g}" for g in range(spec.K // block_k)) + if args.case == "all" + else (args.case,) ) - return 0 if ok else 1 + count = run_cases(spec, cases, compile_route=args.compile_route, tol=args.tol) + print(f"PASS: verified {count} cases", flush=True) + return 0 if __name__ == "__main__": diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/gfx1250/block_scaled_gemm.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/gfx1250/block_scaled_gemm.py index 91764786b5c1..4a8b128d4169 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/gfx1250/block_scaled_gemm.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/gfx1250/block_scaled_gemm.py @@ -2,11 +2,10 @@ # SPDX-License-Identifier: MIT """gfx1250 block-scaled low-bit dense GEMM. -The production target is a WMMA-backed FP8/BF8 expert GEMM. ROCm 7.13 currently -does not select gfx1250 FP8 WMMA builtins, so this day-0 instance ships a real -scalar correctness kernel behind the same gfx1250-owned ABI. It validates -Qwen expert shapes, scale indexing, graph/dispatch plumbing, and GPU numerics -while the low-bit WMMA atom plumbing is brought up. +The default path preserves the existing K=64 FP8/BF8 WMMA plus software +post-scaling contract. ``matrix_path="wmma_scale"`` and ``"wmma_scale16"`` +select the native gfx1250 K=128 instructions and consume packed E8M0 scale +bytes directly. """ from __future__ import annotations @@ -18,7 +17,9 @@ BF16, F16, F32, + I8, I32, + I64, IRBuilder, KernelDef, PtrType, @@ -31,10 +32,18 @@ _LOWBIT_DTYPES = {"fp8", "fp8e4m3", "bf8", "bf8e5m2"} _OUTPUT_DTYPES = {"fp16", "f16", "bf16"} _SCALE_DTYPES = {"fp16", "f16", "fp32", "f32"} -_SUPPORTED_MATRIX_PATHS = {"auto", "wmma", "wmma_scaffold", "mfma"} +_SUPPORTED_MATRIX_PATHS = { + "auto", + "wmma", + "wmma_scaffold", + "wmma_scale", + "wmma_scale16", + "mfma", +} _BLOCK_M = 16 _BLOCK_N = 16 _WMMA_K = 64 +_WMMA_SCALE_K = 128 _WAVE = 32 _HALF_K = 32 # K-elements per lane-half for the K=64 atom _ACC = 8 # accumulator slots per lane (<8 x f32>) @@ -94,7 +103,7 @@ def kernel_name(self) -> str: f"M{self.M}N{self.N}K{self.K}", f"bk{self.block_k}", f"t{self.tile_m}x{self.tile_n}x{self.tile_k}", - flags={"wmma": self.resolved_matrix_path() == "wmma"}, + flags={self.resolved_matrix_path(): True}, ) def resolved_matrix_path(self) -> str: @@ -106,7 +115,7 @@ def resolved_matrix_path(self) -> str: def is_valid_spec(spec: BlockScaledGemmSpec, arch: str = "gfx1250") -> Tuple[bool, str]: - """Return ``(ok, reason)`` for the build-only gfx1250 scaffold.""" + """Return ``(ok, reason)`` for the gfx1250 block-scaled GEMM contract.""" from ...core.arch import ArchTarget try: @@ -137,46 +146,73 @@ def is_valid_spec(spec: BlockScaledGemmSpec, arch: str = "gfx1250") -> Tuple[boo return False, ( f"A/B must be fp8 or bf8 (got A={spec.dtype_a!r}, B={spec.dtype_b!r})" ) + matrix_path = spec.resolved_matrix_path() + native_scale = matrix_path in ("wmma_scale", "wmma_scale16") + family = matrix_path if native_scale else "wmma" + atom_k = _WMMA_SCALE_K if native_scale else _WMMA_K + if native_scale and ( + _canon_lowbit(spec.dtype_a) != "fp8" or _canon_lowbit(spec.dtype_b) != "fp8" + ): + return False, "native gfx1250 SCALE/SCALE16 currently supports fp8 x fp8 only" if not target.mma.has_shape( - family="wmma", + family=family, a_dtype=_canon_lowbit(spec.dtype_a), b_dtype=_canon_lowbit(spec.dtype_b), c_dtype="fp32", m=_BLOCK_M, n=_BLOCK_N, - k=_WMMA_K, + k=atom_k, ): return False, ( - f"no gfx1250 16x16x64 WMMA atom for " + f"no gfx1250 16x16x{atom_k} {family} atom for " f"{_canon_lowbit(spec.dtype_a)}/{_canon_lowbit(spec.dtype_b)}" ) if spec.dtype_c not in ("bf16", "fp16", "f16"): - return False, f"day-0 gfx1250 output must be bf16/fp16 (got {spec.dtype_c!r})" + return False, f"gfx1250 output must be bf16/fp16 (got {spec.dtype_c!r})" if spec.dtype_acc not in ("fp32", "f32"): return False, f"accumulator dtype must be fp32 (got {spec.dtype_acc!r})" if spec.layout != "RCR": return False, f"block_scaled_gemm supports RCR only (got {spec.layout!r})" - try: - _wire_scale_dtype(spec.scale_dtype) - except ValueError as e: - return False, str(e) + if native_scale: + if spec.scale_dtype not in ("e8m0", "i8"): + return False, ( + "native gfx1250 SCALE/SCALE16 requires packed E8M0 scale bytes " + f"(got {spec.scale_dtype!r})" + ) + required_block_k = 16 if matrix_path == "wmma_scale16" else 32 + if spec.block_k != required_block_k: + return False, ( + f"{matrix_path} requires block_k={required_block_k} E8M0 groups " + f"(got {spec.block_k})" + ) + else: + try: + _wire_scale_dtype(spec.scale_dtype) + except ValueError as e: + return False, str(e) if spec.block_k <= 0 or spec.K % spec.block_k: return False, f"K ({spec.K}) must be divisible by block_k ({spec.block_k})" - if spec.block_k % _WMMA_K: + if not native_scale and spec.block_k % _WMMA_K: return False, f"block_k ({spec.block_k}) must be a multiple of {_WMMA_K}" - if spec.K % _WMMA_K: - return False, f"K ({spec.K}) must be a multiple of the WMMA K={_WMMA_K}" + if spec.K % atom_k: + return False, f"K ({spec.K}) must be a multiple of the WMMA K={atom_k}" if (spec.tile_m, spec.tile_n) != (_BLOCK_M, _BLOCK_N): return False, "gfx1250 block_scaled_gemm uses fixed 16x16 output tiles" if spec.M % _BLOCK_M or spec.N % _BLOCK_N: return False, "M and N must be multiples of 16" + if native_scale: + return True, f"ok: gfx1250 K=128 native {matrix_path} FP8 GEMM" return True, "ok: gfx1250 K=64 FP8/BF8 WMMA block-scaled GEMM" def block_scaled_gemm_signature(spec: BlockScaledGemmSpec) -> List[dict]: - """Manifest-style signature for the scaffold contract.""" - scale = _wire_scale_dtype(spec.scale_dtype) + """Return the manifest signature for the selected scaling path.""" + scale = ( + "i8" + if spec.resolved_matrix_path() in ("wmma_scale", "wmma_scale16") + else _wire_scale_dtype(spec.scale_dtype) + ) return ( SignatureBuilder() .ptr("A", _canon_lowbit(spec.dtype_a)) @@ -214,18 +250,20 @@ def _as_f32(b: IRBuilder, v): def build_block_scaled_gemm( spec: BlockScaledGemmSpec, arch: str = "gfx1250" ) -> KernelDef: - """Build the gfx1250 K=64 FP8/BF8 block-scaled GEMM (RCR, ``C = A @ B^T``). - - One wave (32 lanes) computes one 16x16 output tile, no LDS. The K loop runs - in ``block_k``-sized groups; each group accumulates ``block_k / 64`` WMMA - K=64 steps into a fresh ``<8 x f32>`` accumulator, then applies the per-block - scales ``a_scale[m, kg] * b_scale[kg, n]`` to each of the lane's 8 output - slots (group scales are constant within a group, so post-scaling the group - accumulator is exact). The fragment layout mirrors the verified gfx1250 f16 - K=32 atom, doubled to K=64: lane ``l`` carries 32 low-bit bytes - (``<8 x i32>``) starting at K = ``k0 + (l // 16) * 32``; the accumulator is - the column-distributed gfx12 ``<8 x f32>`` (slot ``i`` of lane ``l`` -> - output ``(row = m0 + (l // 16) * 8 + i, col = n0 + l % 16)``). + """Build a gfx1250 FP8/BF8 block-scaled GEMM (RCR, ``C = A @ B^T``). + + One wave (32 lanes) computes one 16x16 output tile without LDS. The legacy + ``wmma`` path uses K=64 FP8/BF8 atoms, accumulates each ``block_k`` group, + and applies FP16/FP32 A/B scales in software. The native ``wmma_scale`` and + ``wmma_scale16`` paths use K=128 FP8 atoms and pass packed E8M0 scale bytes + directly to the instruction, with K=32 and K=16 scale groups respectively. + + Lane ``l`` owns output column ``l % 16`` and rows + ``(l // 16) * 8 : (l // 16 + 1) * 8``. Legacy matrix fragments carry 32 + low-bit bytes per lane as ``<8 x i32>``. Native fragments carry 64 bytes as + ``<16 x i32>`` as four 16-byte K chunks, alternating chunks between lane + halves. Both paths use the gfx12 column-distributed ``<8 x f32>`` + accumulator layout. """ ok, reason = is_valid_spec(spec, arch=arch) if not ok: @@ -234,9 +272,17 @@ def build_block_scaled_gemm( a_ty = _storage_type(spec.dtype_a) b_ty = _storage_type(spec.dtype_b) c_ty = _storage_type(spec.dtype_c) - scale_ty = _scale_type(spec.scale_dtype) - op_id = _wmma_op_id(spec.dtype_a, spec.dtype_b) - a_frag_ty = VectorType(I32, _ACC) + matrix_path = spec.resolved_matrix_path() + native_scale = matrix_path in ("wmma_scale", "wmma_scale16") + scale_ty = I8 if native_scale else _scale_type(spec.scale_dtype) + op_id = ( + f"{matrix_path}_f32_16x16x128_fp8_fp8" + if native_scale + else _wmma_op_id(spec.dtype_a, spec.dtype_b) + ) + atom_k = _WMMA_SCALE_K if native_scale else _WMMA_K + frag_words = 16 if native_scale else _ACC + a_frag_ty = VectorType(I32, frag_words) groups = spec.K // spec.block_k steps_per_group = spec.block_k // _WMMA_K @@ -265,7 +311,8 @@ def build_block_scaled_gemm( lane = ir.mod(ir.thread_id_x(), c32) frag = ir.mod(lane, c16) # lane%16: A row / B col / output col half = ir.div(lane, c16) # lane//16: K-half (operands) + row-block (acc) - half_k = ir.mul(half, ir.const_i32(_HALF_K)) # (l//16)*32 K offset within step + if not native_scale: + half_k = ir.mul(half, ir.const_i32(_HALF_K)) # (l//16)*32 K offset within step m0 = ir.mul(ir.block_id_y(), c16) n0 = ir.mul(ir.block_id_x(), c16) @@ -275,13 +322,77 @@ def build_block_scaled_gemm( b_base = ir.mul(b_row, cK) def _load_frag(ptr, base, storage_ty, k0): - # 32 low-bit bytes per lane. fp8/bf8 vector loads cap at 16 elements, - # so issue two 16-wide loads and concat into the <8 x i32> fragment. - off0 = ir.add(ir.add(base, ir.const_i32(k0)), half_k) - off1 = ir.add(off0, ir.const_i32(16)) - lo = ir.global_load_vN(ptr, off0, storage_ty, 16, align=16) - hi = ir.global_load_vN(ptr, off1, storage_ty, 16, align=16) - return ir.bitcast(ir.vec_concat(lo, hi), a_frag_ty) + if not native_scale: + off0 = ir.add(ir.add(base, ir.const_i32(k0)), half_k) + off1 = ir.add(off0, ir.const_i32(16)) + lo = ir.global_load_vN(ptr, off0, storage_ty, 16, align=16) + hi = ir.global_load_vN(ptr, off1, storage_ty, 16, align=16) + return ir.bitcast(ir.vec_concat(lo, hi), a_frag_ty) + + # CK's gfx1250 SCALE and SCALE16 wrappers use the same four 16-byte + # matrix chunks. The instruction variants differ in scale grouping, + # not in their <16 x i32> A/B fragment layout. + chunk_width = 16 + lane_chunk = ir.mul(half, ir.const_i32(chunk_width)) + step_base = ir.add(base, ir.const_i32(k0)) + chunks = [ + ir.global_load_vN( + ptr, + ir.add( + ir.add(step_base, ir.const_i32(offset)), + lane_chunk, + ), + storage_ty, + chunk_width, + align=chunk_width, + ) + for offset in range(0, atom_k, 2 * chunk_width) + ] + packed = chunks[0] + for chunk in chunks[1:]: + packed = ir.vec_concat(packed, chunk) + return ir.bitcast(packed, a_frag_ty) + + def _pack_strided_scales(ptr, row_or_col, call_idx, *, for_b): + count = 8 if matrix_path == "wmma_scale16" else 4 + word_ty = I64 if count == 8 else I32 + packed = ir.const_i64(0) if count == 8 else ir.const_i32(0) + scale_groups = spec.K // spec.block_k + for j in range(count): + scale_group = call_idx * count + j + if for_b: + idx = ir.add(ir.mul(ir.const_i32(scale_group), cN), row_or_col) + else: + idx = ir.add( + ir.mul(row_or_col, ir.const_i32(scale_groups)), + ir.const_i32(scale_group), + ) + byte = ir.global_load(ptr, idx, I8, align=1) + widened = ir.zext(byte, word_ty) + shift = ir.const_i64(j * 8) if count == 8 else ir.const_i32(j * 8) + shifted = ir.shl(widened, shift) + packed = ir.lor(packed, shifted) + return packed + + if native_scale: + acc = ir.zero_vec_f32(_ACC) + for step in range(spec.K // _WMMA_SCALE_K): + k0 = step * _WMMA_SCALE_K + a_frag = _load_frag(A, a_base, a_ty, k0) + b_frag = _load_frag(B, b_base, b_ty, k0) + a_scale = _pack_strided_scales(AScale, a_row, step, for_b=False) + b_scale = _pack_strided_scales(BScale, b_row, step, for_b=True) + acc = ir.mma(op_id, a_frag, b_frag, acc, a_scale, b_scale) + + out_col = ir.add(n0, frag) + row_base = ir.add(m0, ir.mul(half, ir.const_i32(_ACC))) + for i in range(_ACC): + out_row = ir.add(row_base, ir.const_i32(i)) + idx = ir.add(ir.mul(out_row, cN), out_col) + ir.global_store( + C, idx, ir.cast_f32_to(ir.vec_extract(acc, i), c_ty), align=2 + ) + return ir.kernel # Per-lane f32 output accumulators (8 column-distributed slots). outer = [ir.const_f32(0.0) for _ in range(_ACC)] 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 d2b763e3b651..09fa98506079 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 @@ -20,6 +20,7 @@ #include #include "rocke/ir.h" +#include "rocke/lower_hip.h" #include "rocke/lower_llvm.h" namespace @@ -98,6 +99,38 @@ std::string lower_one(const char* name, return ir; } +/* Build the same kernel through the public HIP-source lowerer. */ +template +std::string lower_one_hip(const char* name, BuildFn build, const char* arch) +{ + rocke_ir_builder_t b; + rocke_strbuf_t out; + rocke_lower_hip_opts_t opts{}; + if(rocke_ir_builder_init(&b, name) != ROCKE_OK) + { + fail("rocke_ir_builder_init", __LINE__); + return std::string(); + } + build(&b); + rocke_b_ret(&b); + rocke_strbuf_init(&out, 0); + opts.arch = arch; + const rocke_status_t st + = rocke_lower_kernel_to_hip(&b, rocke_ir_builder_kernel(&b), &opts, &out); + std::string hip; + if(st != ROCKE_OK) + { + char msg[128]; + snprintf(msg, sizeof(msg), "HIP lower failed (status %d)", (int)st); + fail(msg, __LINE__); + } + else + hip.assign(rocke_strbuf_cstr(&out)); + rocke_strbuf_free(&out); + rocke_ir_builder_free(&b); + return hip; +} + rocke_value_t* global_ptr_param(rocke_ir_builder_t* b, const char* name, const rocke_type_t* elem); void case_gfx1250_standalone_bridge() @@ -512,6 +545,78 @@ void case_global_load_async_to_lds_b8() EXPECT_IR(ir, "getelementptr inbounds"); } +/* ---- HIP zero-extension source signedness ---- */ +void case_hip_zext_uses_unsigned_source_cast() +{ + const std::string hip = lower_one_hip( + "zext_i8", + [](rocke_ir_builder_t* b) { + rocke_value_t* byte = rocke_b_param(b, "byte", rocke_i8(), nullptr); + rocke_b_zext(b, byte, rocke_i32()); + rocke_b_zext(b, byte, rocke_i64()); + }, + "gfx950"); + EXPECT_IR_COUNT(hip, "(int)(uint8_t)byte", 1); + EXPECT_IR_COUNT(hip, "(int64_t)(uint8_t)byte", 1); + EXPECT_NO_IR(hip, "(int)byte"); + EXPECT_NO_IR(hip, "(int64_t)byte"); +} + +/* ---- gfx1250 native SCALE / SCALE16 FP8 WMMA ---- */ +void case_gfx1250_scaled_wmma() +{ + struct Variant + { + bool scale16; + const rocke_type_t* (*scale_type)(void); + const char* intrinsic; + const char* builtin; + const char* scale_llvm_type; + }; + static const Variant variants[] = { + {false, + rocke_i32, + "llvm.amdgcn.wmma.scale.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32", + "__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4", + "i32"}, + {true, + rocke_i64, + "llvm.amdgcn.wmma.scale16.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32", + "__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4", + "i64"}, + }; + + for(const Variant& v : variants) + { + const auto build = [&v](rocke_ir_builder_t* b) { + rocke_value_t* matrix = global_ptr_param(b, "matrix", rocke_i32()); + rocke_value_t* accum = global_ptr_param(b, "accum", rocke_f32()); + rocke_value_t* scales = global_ptr_param(b, "scales", v.scale_type()); + rocke_value_t* lane = rocke_b_thread_id_x(b); + rocke_value_t* lo = rocke_b_global_load_vN(b, matrix, lane, rocke_i32(), 8, 0); + rocke_value_t* eight = rocke_b_const_i32(b, 8); + rocke_value_t* hi_idx = rocke_b_add(b, lane, eight); + rocke_value_t* hi = rocke_b_global_load_vN(b, matrix, hi_idx, rocke_i32(), 8, 0); + rocke_value_t* fragment = rocke_b_vec_concat(b, lo, hi); + rocke_value_t* c = rocke_b_global_load_vN(b, accum, lane, rocke_f32(), 8, 0); + rocke_value_t* scale = rocke_b_global_load(b, scales, lane, v.scale_type(), 1); + rocke_value_t* d = v.scale16 ? rocke_b_wmma_scale16_f32_16x16x128_fp8_fp8( + b, fragment, fragment, c, scale, scale) + : rocke_b_wmma_scale_f32_16x16x128_fp8_fp8( + b, fragment, fragment, c, scale, scale); + rocke_b_global_store(b, accum, lane, rocke_b_vec_extract(b, d, 0), 1); + }; + const char* name = v.scale16 ? "wmma_scale16" : "wmma_scale"; + const std::string ir = lower_one(name, build, "gfx1250", ROCKE_LLVM_FLAVOR_LLVM23); + EXPECT_IR(ir, v.intrinsic); + const std::string scale_arg = std::string(", ") + v.scale_llvm_type + " %"; + EXPECT_IR(ir, scale_arg.c_str()); + + const std::string hip = lower_one_hip(name, build, "gfx1250"); + EXPECT_IR(hip, v.builtin); + } +} + /* ---- opcode table alignment ---- */ /* These opcodes were spliced into rocke_opcode_t's family groups rather than @@ -614,6 +719,8 @@ const TestCase k_cases[] = { {"s_prefetch_inst", case_s_prefetch_inst}, {"buffer_load_lds_async", case_buffer_load_lds_async}, {"global_load_async_to_lds_b8", case_global_load_async_to_lds_b8}, + {"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}, {"opcode_names_are_aligned", case_opcode_names_are_aligned}, }; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/mma_frag_ssot.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/mma_frag_ssot.cpp index 4b77f1136edb..ceb797c3d5cc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/mma_frag_ssot.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/mma_frag_ssot.cpp @@ -17,6 +17,7 @@ #include #include +#include "rocke/arch_target.h" #include "rocke/ir.h" static int g_failures = 0; @@ -59,6 +60,19 @@ static void check_atom(rocke_ir_builder_t* b, const char* op_id, int expect_frag } } +static void check_catalog_fragments(const rocke_mma_catalog_t* catalog, const char* op_id) +{ + const rocke_mma_op_t* op = rocke_mma_catalog_by_op_id(catalog, op_id); + CHECK(op != NULL, op_id); + if(!op) + { + return; + } + CHECK(op->a_frag_len == 16, op_id); + CHECK(op->b_frag_len == 16, op_id); + CHECK(op->c_frag_len == 8, op_id); +} + int main(void) { rocke_ir_builder_t b; @@ -78,6 +92,17 @@ int main(void) check_atom(&b, "wmma_i32_16x16x16_iu8", 8, true); check_atom(&b, "wmma_i32_16x16x16_iu4", 8, true); + /* Keep the C catalog in parity with Python _MMA_FRAGMENT_INFO. The scaled + * matrix operands are <16 x i32>, so fragment length is 16 elements, not + * the equivalent 64-byte storage width. */ + const rocke_arch_target_t* gfx1250 = rocke_arch_target_from_gfx("gfx1250"); + CHECK(gfx1250 != NULL, "gfx1250 target"); + if(gfx1250) + { + check_catalog_fragments(&gfx1250->mma, "wmma_scale_f32_16x16x128_fp8_fp8"); + check_catalog_fragments(&gfx1250->mma, "wmma_scale16_f32_16x16x128_fp8_fp8"); + } + /* Unknown op_id must be rejected. The engine's error path either returns * NULL with a sticky builder error or raises (ckc::ValueError) depending on * build config, so accept either form of rejection. */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_gfx1250_scaled_wmma.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_gfx1250_scaled_wmma.py new file mode 100644 index 000000000000..71028eed346a --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_gfx1250_scaled_wmma.py @@ -0,0 +1,201 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Tests for the gfx1250 native FP8 SCALE and SCALE16 WMMA slice.""" + +from __future__ import annotations + +import unittest + +from rocke.core.arch import ArchTarget +from rocke.core.ir import F32, I32, I64, IRBuilder, PtrType +from rocke.core.lower_hip import lower_kernel_to_hip +from rocke.core.lower_llvm import lower_kernel_to_llvm +from rocke.instances.gfx1250.block_scaled_gemm import ( + BlockScaledGemmSpec, + block_scaled_gemm_signature, + build_block_scaled_gemm, + is_valid_spec, +) + + +def _build_scaled_atom(*, scale16: bool): + b = IRBuilder("gfx1250_scaled_wmma") + matrix = b.param("matrix", PtrType(I32, "global"), readonly=True) + accum = b.param("accum", PtrType(F32, "global")) + scale_ty = I64 if scale16 else I32 + scales = b.param("scales", PtrType(scale_ty, "global"), readonly=True) + lane = b.thread_id_x() + lo = b.global_load_vN(matrix, lane, I32, 8) + hi = b.global_load_vN(matrix, b.add(lane, b.const_i32(8)), I32, 8) + fragment = b.vec_concat(lo, hi) + c = b.global_load_vN(accum, lane, F32, 8) + scale = b.global_load(scales, lane, scale_ty) + if scale16: + d = b.wmma_scale16_f32_16x16x128_fp8_fp8(fragment, fragment, c, scale, scale) + else: + d = b.wmma_scale_f32_16x16x128_fp8_fp8(fragment, fragment, c, scale, scale) + b.global_store(accum, lane, b.vec_extract(d, 0)) + return b.kernel + + +class TestGfx1250ScaledWmma(unittest.TestCase): + def test_catalog_fragment_lengths(self): + target = ArchTarget.from_gfx("gfx1250") + for op_id in ( + "wmma_scale_f32_16x16x128_fp8_fp8", + "wmma_scale16_f32_16x16x128_fp8_fp8", + ): + with self.subTest(op_id=op_id): + op = target.mma.by_op_id(op_id) + self.assertIsNotNone(op) + self.assertEqual( + (op.a_frag_len, op.b_frag_len, op.c_frag_len), (16, 16, 8) + ) + + def test_scale_and_scale16_exact_llvm23_abi(self): + cases = ( + ( + False, + ("llvm.amdgcn.wmma.scale.f32.16x16x128.f8f6f4." "v8f32.v16i32.v16i32"), + "i32", + ), + ( + True, + ( + "llvm.amdgcn.wmma.scale16.f32.16x16x128.f8f6f4." + "v8f32.v16i32.v16i32" + ), + "i64", + ), + ) + for scale16, intrinsic, scale_ty in cases: + with self.subTest(scale16=scale16): + ll = lower_kernel_to_llvm( + _build_scaled_atom(scale16=scale16), + llvm_flavor="llvm23", + arch="gfx1250", + ) + self.assertIn(f"declare <8 x float> @{intrinsic}", ll) + self.assertIn( + "i32 0, <16 x i32>", + next( + line + for line in ll.splitlines() + if f"call <8 x float> @{intrinsic}" in line + ), + ) + self.assertIn(f", {scale_ty} %", ll) + + def test_scaled_wmma_rejects_pre_llvm23_flavors(self): + for flavor in ("llvm20", "llvm22"): + with self.subTest(flavor=flavor), self.assertRaisesRegex( + NotImplementedError, "requires llvm23" + ): + lower_kernel_to_llvm( + _build_scaled_atom(scale16=False), + llvm_flavor=flavor, + arch="gfx1250", + ) + + def test_scale_and_scale16_hip_builtin_lowering(self): + for scale16, builtin in ( + (False, "__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4"), + (True, "__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4"), + ): + with self.subTest(scale16=scale16): + hip = lower_kernel_to_hip( + _build_scaled_atom(scale16=scale16), arch="gfx1250" + ) + self.assertIn(builtin, hip) + + def test_native_block_scaled_gemm_uses_packed_e8m0_in_instruction(self): + for matrix_path, block_k, scale_ty, fragment_load, load_count in ( + ("wmma_scale", 32, "i32", "load <16 x i8>", 8), + ("wmma_scale16", 16, "i64", "load <16 x i8>", 8), + ): + with self.subTest(matrix_path=matrix_path): + spec = BlockScaledGemmSpec( + name="native_mx", + M=16, + N=16, + K=128, + scale_dtype="e8m0", + block_k=block_k, + matrix_path=matrix_path, + ) + ok, why = is_valid_spec(spec) + self.assertTrue(ok, why) + signature = block_scaled_gemm_signature(spec) + self.assertEqual(signature[2]["type"], "ptr") + ll = lower_kernel_to_llvm( + build_block_scaled_gemm(spec), + llvm_flavor="llvm23", + arch="gfx1250", + ) + intrinsic_family = "wmma.scale16" if block_k == 16 else "wmma.scale" + self.assertIn(f"llvm.amdgcn.{intrinsic_family}.", ll) + self.assertIn(f", {scale_ty} %", ll) + self.assertEqual(ll.count(fragment_load), load_count) + self.assertNotIn("fmul float", ll) + + def test_native_block_scaled_gemm_hip_zero_extends_e8m0_bytes(self): + cases = ( + ("wmma_scale", 32, "(int)(uint8_t)", 8, r" int zx\d+ = \(int\)\w+;"), + ( + "wmma_scale16", + 16, + "(int64_t)(uint8_t)", + 16, + r" int64_t zx\d+ = \(int64_t\)\w+;", + ), + ) + for matrix_path, block_k, unsigned_cast, count, direct_signed in cases: + with self.subTest(matrix_path=matrix_path): + spec = BlockScaledGemmSpec( + name="native_mx_hip", + M=16, + N=16, + K=128, + scale_dtype="e8m0", + block_k=block_k, + matrix_path=matrix_path, + ) + hip = lower_kernel_to_hip(build_block_scaled_gemm(spec), arch="gfx1250") + self.assertEqual(hip.count(unsigned_cast), count) + self.assertNotRegex(hip, direct_signed) + + scale_bytes = [0x80, 0xFF, 0x00, 0x01] + packed = sum(byte << (8 * index) for index, byte in enumerate(scale_bytes)) + self.assertEqual(packed, 0x0100FF80) + + def test_native_slice_rejects_unsupported_contracts(self): + bad_dtype = BlockScaledGemmSpec( + name="bad_dtype", + M=16, + N=16, + K=128, + dtype_b="bf8", + scale_dtype="e8m0", + block_k=32, + matrix_path="wmma_scale", + ) + ok, why = is_valid_spec(bad_dtype) + self.assertFalse(ok) + self.assertIn("fp8 x fp8 only", why) + + bad_block = BlockScaledGemmSpec( + name="bad_block", + M=16, + N=16, + K=128, + scale_dtype="e8m0", + block_k=32, + matrix_path="wmma_scale16", + ) + ok, why = is_valid_spec(bad_block) + self.assertFalse(ok) + self.assertIn("requires block_k=16", why) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_gfx1250_scaled_wmma_reference.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_gfx1250_scaled_wmma_reference.py new file mode 100644 index 000000000000..e93bd0514f33 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_gfx1250_scaled_wmma_reference.py @@ -0,0 +1,119 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""CPU checks for the scaled-WMMA numerical oracle and diagnostic fixtures.""" + +from __future__ import annotations + +import numpy as np +import pytest + +pytest.importorskip("ml_dtypes") + +from rocke.examples.gfx1250.gemm.block_scaled_gemm_verify import ( + check_result, + decode_e8m0, + make_case_inputs, + reference_result, +) +from rocke.instances.gfx1250.block_scaled_gemm import BlockScaledGemmSpec + + +def test_e8m0_golden_values_include_high_bit(): + encoded = np.array([125, 126, 127, 128, 129, 130], dtype=np.uint8) + np.testing.assert_array_equal(decode_e8m0(encoded), [0.25, 0.5, 1, 2, 4, 8]) + with pytest.raises(ValueError, match="0xff"): + decode_e8m0(np.array([255], dtype=np.uint8)) + + +def test_reference_matches_hand_computed_group_scales(): + a = np.array([[1, 2, 3, 4], [2, 0, 1, -1]], dtype=np.float32) + b = np.array([[1, 0, 1, 0], [0, 1, 0, 1], [1, 1, 1, 1]], dtype=np.float32) + sa = np.array([[127, 128], [126, 127]], dtype=np.uint8) + sb = np.array([[127, 128, 126], [128, 127, 129]], dtype=np.uint8) + expected = np.array([[13, 12, 57.5], [3, -1, 0.5]], dtype=np.float32) + got = reference_result(a, b, sa, sb, 2, native=True) + np.testing.assert_array_equal(got, expected) + # Either operand's scale-group permutation must be observable. + for wrong_sa, wrong_sb in ((sa[:, ::-1], sb), (sa, sb[::-1, :])): + wrong = reference_result(a, b, wrong_sa, wrong_sb, 2, native=True) + with pytest.raises(AssertionError, match="bad="): + check_result(wrong, expected, exact=True) + + +@pytest.mark.parametrize("path,bk", [("wmma_scale", 32), ("wmma_scale16", 16)]) +@pytest.mark.parametrize("k", [128, 256]) +def test_fixtures_exercise_each_scale_group_and_high_bytes(path, bk, k): + spec = BlockScaledGemmSpec( + name="reference", + M=32, + N=48, + K=k, + matrix_path=path, + block_k=bk, + scale_dtype="e8m0", + ) + a, b, sa, sb = make_case_inputs(spec, "mixed") + assert set(np.unique(sa)) == set(range(125, 131)) + assert set(np.unique(sb)) == set(range(125, 131)) + expected = reference_result(a, b, sa, sb, bk, native=True) + for wrong_sa, wrong_sb in ((sa[:, ::-1], sb), (sa, sb[::-1, :])): + wrong = reference_result(a, b, wrong_sa, wrong_sb, bk, native=True) + with pytest.raises(AssertionError, match="bad="): + check_result(wrong, expected, exact=True) + # At K=256, reusing the first call's scales must also be detectable. + if k == 256: + first_call = 128 // bk + reused_sa = np.tile(sa[:, :first_call], (1, 2)) + reused_sb = np.tile(sb[:first_call, :], (2, 1)) + wrong = reference_result(a, b, reused_sa, reused_sb, bk, native=True) + with pytest.raises(AssertionError, match="bad="): + check_result(wrong, expected, exact=True) + for group in range(k // bk): + ga, gb, gsa, gsb = make_case_inputs(spec, f"group-{group}") + sl = slice(group * bk, (group + 1) * bk) + np.testing.assert_array_equal(ga[:, sl], a[:, sl]) + np.testing.assert_array_equal(gb[:, sl], b[:, sl]) + outside = (np.arange(k) // bk) != group + assert not np.any(ga[:, outside].astype(np.float32)) + assert not np.any(gb[:, outside].astype(np.float32)) + isolated = reference_result(ga, gb, gsa, gsb, bk, native=True) + assert np.isfinite(isolated).all() and np.any(isolated) + # Wrong scales for this one active group cannot hide in accumulation. + wrong = reference_result(ga, gb, gsa[:, ::-1], gsb, bk, native=True) + with pytest.raises(AssertionError, match="bad="): + check_result(wrong, isolated, exact=True) + + +@pytest.mark.parametrize("case", ["neutral", "a-only", "b-only"]) +def test_neutral_and_one_operand_scale_fixtures(case): + spec = BlockScaledGemmSpec( + name="reference", + M=16, + N=16, + K=128, + matrix_path="wmma_scale", + block_k=32, + scale_dtype="e8m0", + ) + _, _, sa, sb = make_case_inputs(spec, case) + assert bool(np.all(sa == 127)) == (case in ("neutral", "b-only")) + assert bool(np.all(sb == 127)) == (case in ("neutral", "a-only")) + + +@pytest.mark.parametrize("value", [np.nan, np.inf, -np.inf]) +@pytest.mark.parametrize("exact", [False, True]) +def test_comparison_rejects_nonfinite_output_and_reference(value, exact): + finite = np.array([[1.0, 0.0]], dtype=np.float32) + nonfinite = np.array([[1.0, value]], dtype=np.float32) + for got, ref in ((nonfinite, finite), (finite, nonfinite), (nonfinite, nonfinite)): + with pytest.raises(AssertionError, match="non-finite"): + check_result(got, ref, exact=exact) + + +def test_exact_comparison_detects_one_bf16_step_and_shape_mismatch(): + expected = np.array([[1.0]], dtype=np.float32) + check_result(expected.copy(), expected, exact=True) + with pytest.raises(AssertionError, match="bad=1/1"): + check_result(expected + 2**-7, expected, exact=True) + with pytest.raises(AssertionError, match="shape"): + check_result(expected.ravel(), expected, exact=True) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx1250_lowering_emit.c b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx1250_lowering_emit.c index 3d1042bbdf73..bc5f0b71ea66 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx1250_lowering_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx1250_lowering_emit.c @@ -117,6 +117,43 @@ static void build_wmma_k64_bf8_bf8(rocke_ir_builder_t* b) wmma_k64(b, "wmma_gfx1250_f32_16x16x64_bf8_bf8"); } +/* K=128 FP8 SCALE/SCALE16 WMMA. Matrix fragments are <16 x i32>; packed + * E8M0 scale operands are i32 for SCALE and i64 for SCALE16. */ +static void wmma_scaled(rocke_ir_builder_t* b, bool scale16) +{ + const rocke_type_t* scale_ty = scale16 ? rocke_i64() : rocke_i32(); + rocke_value_t* a_ptr = frag_param(b, "A", rocke_i32(), true); + rocke_value_t* b_ptr = frag_param(b, "B", rocke_i32(), true); + rocke_value_t* c_ptr = frag_param(b, "C", rocke_f32(), false); + rocke_value_t* scale_ptr = frag_param(b, "scale", scale_ty, true); + rocke_value_t* tid = rocke_b_thread_id_x(b); + rocke_value_t* a_lo = rocke_b_global_load_vN(b, a_ptr, tid, rocke_i32(), 8, /*align=*/0); + rocke_value_t* eight = rocke_b_const_i32(b, 8); + rocke_value_t* hi_idx = rocke_b_add(b, tid, eight); + rocke_value_t* a_hi = rocke_b_global_load_vN(b, a_ptr, hi_idx, rocke_i32(), 8, /*align=*/0); + rocke_value_t* a = rocke_b_vec_concat(b, a_lo, a_hi); + rocke_value_t* b_lo = rocke_b_global_load_vN(b, b_ptr, tid, rocke_i32(), 8, /*align=*/0); + rocke_value_t* b_hi = rocke_b_global_load_vN(b, b_ptr, hi_idx, rocke_i32(), 8, /*align=*/0); + rocke_value_t* bb = rocke_b_vec_concat(b, b_lo, b_hi); + rocke_value_t* c = rocke_b_global_load_vN(b, c_ptr, tid, rocke_f32(), 8, /*align=*/0); + rocke_value_t* scale = rocke_b_global_load(b, scale_ptr, tid, scale_ty, /*align=*/1); + rocke_value_t* d = scale16 + ? rocke_b_wmma_scale16_f32_16x16x128_fp8_fp8(b, a, bb, c, scale, scale) + : rocke_b_wmma_scale_f32_16x16x128_fp8_fp8(b, a, bb, c, scale, scale); + rocke_b_global_store(b, c_ptr, tid, d, /*align=*/1); + rocke_b_ret(b); +} + +static void build_wmma_scale(rocke_ir_builder_t* b) +{ + wmma_scaled(b, false); +} + +static void build_wmma_scale16(rocke_ir_builder_t* b) +{ + wmma_scaled(b, true); +} + /* ds_read_b128_tr_b16. gfx950 has one type-agnostic opcode returning * <8 x i16> that the handler reinterprets; gfx1250 has per-element-type * opcodes (.v8f16 / .v8bf16) that land in the right type with no reinterpret. */ @@ -309,6 +346,8 @@ static const config_t CONFIGS[] = { {build_wmma_k64_fp8_bf8, "gfx1250"}, {build_wmma_k64_bf8_fp8, "gfx1250"}, {build_wmma_k64_bf8_bf8, "gfx1250"}, + {build_wmma_scale, "gfx1250"}, + {build_wmma_scale16, "gfx1250"}, {build_tr16_f16, "gfx1250"}, {build_tr16_f16, "gfx950"}, {build_tr16_bf16, "gfx1250"}, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx1250_lowering_emit.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx1250_lowering_emit.py index 1ce4e5b94ed7..80ca5f0f3075 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx1250_lowering_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx1250_lowering_emit.py @@ -7,13 +7,13 @@ # # gfx1250 kernels are authored in Python (there is no cpp/instances/gfx1250/), # so no instance-builder family in this gate exercises the gfx1250 *lowerer*. -# That left the six places Gfx1250Backend diverges from its Gfx12RdnaBackend +# That left the target-specific places Gfx1250Backend diverges from its Gfx12RdnaBackend # parent covered only by the ROCKE_BACKEND=both pytest lane, and invisible to # check_byte_identity.py -- the tree's stated definition-of-done. This family # closes that: one config per divergence, built from the smallest kernel that # reaches it, byte-compared against gfx1250_lowering_emit.c. # -# The gfx950 twins are deliberately included. Four of the six divergences are a +# The gfx950 twins are deliberately included. Several divergences are a # *choice between two encodings*, so a lowering that ignored the backend and # always picked the gfx1250 form would still pass a gfx1250-only family. Pairing # each with its gfx950 counterpart pins both sides of the choice. @@ -74,6 +74,44 @@ def build(b: IRBuilder) -> None: return build +def _wmma_scaled(a_kind, b_kind, scale_mode): + """K=128 scaled WMMA, parameterized by operand dtypes and scale mode.""" + scale_ty = {"scale": I32, "scale16": I64}[scale_mode] + op_id = f"wmma_{scale_mode}_f32_16x16x128_{a_kind}_{b_kind}" + + def build(b: IRBuilder) -> None: + a_ptr = b.param( + "A", PtrType(I32, "global"), noalias=True, readonly=True, align=16 + ) + b_ptr = b.param( + "B", PtrType(I32, "global"), noalias=True, readonly=True, align=16 + ) + c_ptr = b.param("C", PtrType(F32, "global"), noalias=True, align=16) + scale_ptr = b.param( + "scale", + PtrType(scale_ty, "global"), + noalias=True, + readonly=True, + align=16, + ) + tid = b.thread_id_x() + a_lo = b.global_load_vN(a_ptr, tid, dtype=I32, n=8) + eight = b.const_i32(8) + hi_idx = b.add(tid, eight) + a_hi = b.global_load_vN(a_ptr, hi_idx, dtype=I32, n=8) + a = b.vec_concat(a_lo, a_hi) + b_lo = b.global_load_vN(b_ptr, tid, dtype=I32, n=8) + b_hi = b.global_load_vN(b_ptr, hi_idx, dtype=I32, n=8) + bb = b.vec_concat(b_lo, b_hi) + c = b.global_load_vN(c_ptr, tid, dtype=F32, n=8) + scale = b.global_load(scale_ptr, tid, dtype=scale_ty) + d = b.mma(op_id, a, bb, c, scale, scale) + b.global_store(c_ptr, tid, d) + b.ret() + + return build + + def _tr16_b128(elem): """ds_read_b128_tr_b16. @@ -197,6 +235,8 @@ def build_tensor_transfers(b: IRBuilder) -> None: (_wmma_k64("fp8", "bf8"), "gfx1250"), (_wmma_k64("bf8", "fp8"), "gfx1250"), (_wmma_k64("bf8", "bf8"), "gfx1250"), + (_wmma_scaled("fp8", "fp8", "scale"), "gfx1250"), + (_wmma_scaled("fp8", "fp8", "scale16"), "gfx1250"), (_tr16_b128(F16), "gfx1250"), (_tr16_b128(F16), "gfx950"), (_tr16_b128(BF16), "gfx1250"), diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx1250_scaled_wmma_numeric.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx1250_scaled_wmma_numeric.py new file mode 100644 index 000000000000..9ceb92484625 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx1250_scaled_wmma_numeric.py @@ -0,0 +1,96 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Real gfx1250 SCALE/SCALE16 launches; no torch dependency. + +Ordinary CPU/other-GPU runs skip. Set ROCKE_REQUIRE_GFX1250=1 on the validation +node to fail instead of silently accepting an all-skipped run. Select the GPU +with HIP_VISIBLE_DEVICES, and use ROCKE_LLVM_FLAVOR=llvm23 with matching COMGR. +""" + +from __future__ import annotations + +import importlib.util +import os +from pathlib import Path +import subprocess +import sys + +import pytest + + +_PYROOT = Path(__file__).resolve().parents[2] / "python" +_MODULE = "rocke.examples.gfx1250.gemm.block_scaled_gemm_verify" + + +@pytest.fixture(scope="module") +def gpu_env(): + env = dict(os.environ) + env["PYTHONPATH"] = os.pathsep.join([str(_PYROOT), env.get("PYTHONPATH", "")]) + # Keep HIP initialization out of pytest collection and bound a hung probe. + probe = subprocess.run( + [ + sys.executable, + "-c", + "from rocke.runtime.hip_module import get_device_arch; " + "print(get_device_arch(0))", + ], + env=env, + capture_output=True, + text=True, + timeout=30, + ) + assert probe.returncode == 0, probe.stdout + probe.stderr + arch = probe.stdout.strip().splitlines()[-1] + if arch != "gfx1250": + reason = f"requires visible HIP device 0 = gfx1250; detected {arch}" + if env.get("ROCKE_REQUIRE_GFX1250") == "1": + pytest.fail(reason) + pytest.skip(reason) + for dependency in ("numpy", "ml_dtypes"): + assert importlib.util.find_spec( + dependency + ), f"install {dependency} for numerics" + return env + + +@pytest.mark.parametrize( + "matrix_path,route,m,n,k,case,count", + [ + ("wmma_scale", "comgr", 16, 16, 128, "all", 8), + ("wmma_scale16", "comgr", 16, 16, 128, "all", 12), + ("wmma_scale", "comgr", 32, 48, 256, "mixed", 1), + ("wmma_scale16", "comgr", 32, 48, 256, "mixed", 1), + ("wmma_scale", "hip", 32, 48, 256, "mixed", 1), + ("wmma_scale16", "hip", 32, 48, 256, "mixed", 1), + ("wmma", "comgr", 16, 16, 128, "mixed", 1), + ], +) +def test_scaled_wmma_numeric(gpu_env, matrix_path, route, m, n, k, case, count): + result = subprocess.run( + [ + sys.executable, + "-m", + _MODULE, + "--matrix-path", + matrix_path, + "--compile-route", + route, + "--m", + str(m), + "--n", + str(n), + "--k", + str(k), + "--case", + case, + ], + env=gpu_env, + capture_output=True, + text=True, + timeout=300, + ) + output = result.stdout + result.stderr + assert result.returncode == 0, output + assert f"PASS: verified {count} cases" in output, output + assert output.count("bad=0") == count, output + print(output, end="")