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 c2cfc9422769..58554acad8e4 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 @@ -52,6 +52,7 @@ extern "C" { #include "rocke/instance_gemm_multi_abd.h" #include "rocke/instance_gemm_multi_d.h" #include "rocke/instance_gemm_universal.h" +#include "rocke/instance_gfx950_block_scaled_gemm.h" #include "rocke/instance_grouped_gemm.h" #include "rocke/instance_matmul_nbits.h" #include "rocke/instance_mfma_gemm.h" @@ -1031,6 +1032,63 @@ std::vector mx_gemm_verify(const py::dict& d, const std::string& ar rocke_build_mx_gemm_new(&b, &s, arch_or_default(arch))); } +rocke_gfx950_block_scaled_gemm_spec_t gfx950_scaled_build_spec(const py::dict& d, + std::deque& store) +{ + auto s = rocke_gfx950_block_scaled_gemm_spec_default(); + s.M = dict_int(d, "M", s.M); + s.N = dict_int(d, "N", s.N); + s.K = dict_int(d, "K", s.K); + s.block_k = dict_int(d, "block_k", s.block_k); + const char* fields[] = {"name", "dtype_a", "dtype_b", "dtype_c"}; + const char** outputs[] = {&s.name, &s.dtype_a, &s.dtype_b, &s.dtype_c}; + for(int i = 0; i < 4; ++i) + { + std::string v; + if(dict_str(d, fields[i], v)) + { + store.push_back(v); + *outputs[i] = store.back().c_str(); + } + } + return s; +} +std::string gfx950_scaled_serialize_ir(const py::dict& d, const std::string& arch) +{ + ROCKE_FAMILY_SERIALIZE_BODY( + "rocke_engine.gfx950_block_scaled_gemm_serialize_ir", + rocke_gfx950_block_scaled_gemm_spec_t, + gfx950_scaled_build_spec, + rocke_build_gfx950_block_scaled_gemm_new(&b, &s, arch_or_default(arch))); +} +std::vector gfx950_scaled_verify(const py::dict& d, const std::string& arch) +{ + ROCKE_FAMILY_VERIFY_BODY( + "rocke_engine.gfx950_block_scaled_gemm_verify", + rocke_gfx950_block_scaled_gemm_spec_t, + gfx950_scaled_build_spec, + rocke_build_gfx950_block_scaled_gemm_new(&b, &s, arch_or_default(arch))); +} +std::string gfx950_scaled_lower_llvm(const py::dict& d, const std::string& arch) +{ + std::deque store; + auto s = gfx950_scaled_build_spec(d, store); + rocke_ir_builder_t b; + auto* k = rocke_build_gfx950_block_scaled_gemm_new(&b, &s, arch_or_default(arch)); + if(!k) + { + std::string err = rocke_ir_builder_error(&b); + rocke_ir_builder_free(&b); + throw std::runtime_error(err); + } + char* ll = nullptr; + char err[ROCKE_ERR_MSG_CAP] = {}; + auto status = rocke_lower_kernel_to_llvm_ex( + k, ROCKE_LLVM_FLAVOR_AUTO, arch_or_default(arch), &ll, err, sizeof(err)); + rocke_ir_builder_free(&b); + return take_lowered(status, ll, err, "gfx950_block_scaled_gemm"); +} + /* ============================== mfma GEMM ============================= */ rocke_mfma_gemm_spec_t mfma_build_spec(const py::dict& d, std::deque& store) @@ -3787,6 +3845,10 @@ PYBIND11_MODULE(rocke_engine, m) &block_scale_gemm_lower_llvm, &block_scale_gemm_serialize_ir, &block_scale_gemm_verify); + reg3("gfx950_block_scaled_gemm", + &gfx950_scaled_lower_llvm, + &gfx950_scaled_serialize_ir, + &gfx950_scaled_verify); reg3("mx_gemm", &mx_gemm_lower_llvm, &mx_gemm_serialize_ir, &mx_gemm_verify); reg3("mfma_gemm", &mfma_gemm_lower_llvm, &mfma_gemm_serialize_ir, &mfma_gemm_verify); m.def("mfma_gemm_is_valid", &mfma_gemm_is_valid, py::arg("spec"), py::arg("arch") = "gfx950"); 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 a0cf89927941..9e31b93084d2 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 @@ -1060,6 +1060,10 @@ int rocke_arch_mma_c_frag_len(const char* op_id) return rocke_ati_mma_frag[i].c_frag_len; } } + const rocke_arch_target_t* target = rocke_arch_target_from_gfx("gfx950"); + const rocke_mma_op_t* atom = rocke_mma_catalog_by_op_id(&target->mma, op_id); + if(atom) + return atom->c_frag_len; return 0; /* unknown atom: the zero-length _frag_info fallback */ } @@ -1333,7 +1337,607 @@ static const rocke_mma_op_t k_mma_gfx942[] = { }; /* ----------------------------- gfx950 (CDNA) ----------------------------- */ +static void mfma_scaled_a( + rocke_ir_builder_t* b, rocke_value_t* lane, int slot, rocke_value_t** row, rocke_value_t** k) +{ + *row = rocke_b_mod(b, lane, rocke_b_const_i32(b, 32)); + *k = rocke_b_add( + b, + rocke_b_mul(b, rocke_b_div(b, lane, rocke_b_const_i32(b, 32)), rocke_b_const_i32(b, 32)), + rocke_b_const_i32(b, slot)); +} +static void mfma_scaled_f8_a( + rocke_ir_builder_t* b, rocke_value_t* lane, int slot, rocke_value_t** row, rocke_value_t** k) +{ + *row = rocke_b_mod(b, lane, rocke_b_const_i32(b, 32)); + *k = rocke_b_add( + b, + rocke_b_mul(b, rocke_b_div(b, lane, rocke_b_const_i32(b, 32)), rocke_b_const_i32(b, 16)), + rocke_b_const_i32(b, (slot / 16) * 32 + slot % 16)); +} +static void mfma_scaled_f8_b( + rocke_ir_builder_t* b, rocke_value_t* lane, int slot, rocke_value_t** k, rocke_value_t** col) +{ + mfma_scaled_f8_a(b, lane, slot, col, k); +} +static void mfma_scaled_b( + rocke_ir_builder_t* b, rocke_value_t* lane, int slot, rocke_value_t** k, rocke_value_t** col) +{ + mfma_scaled_a(b, lane, slot, col, k); +} +static void mfma_scaled_scale_a( + rocke_ir_builder_t* b, rocke_value_t* lane, int slot, rocke_value_t** row, rocke_value_t** k) +{ + (void)slot; + *row = rocke_b_mod(b, lane, rocke_b_const_i32(b, 32)); + *k = rocke_b_div(b, lane, rocke_b_const_i32(b, 32)); +} +static void mfma_scaled_scale_b( + rocke_ir_builder_t* b, rocke_value_t* lane, int slot, rocke_value_t** k, rocke_value_t** col) +{ + mfma_scaled_scale_a(b, lane, slot, col, k); +} +static const rocke_layout_map_t lm_mfma_scaled_f8_a = {ROCKE_MMA_ROLE_A, 32, 64, mfma_scaled_f8_a}; +static const rocke_layout_map_t lm_mfma_scaled_f8_b = {ROCKE_MMA_ROLE_B, 32, 64, mfma_scaled_f8_b}; +static const rocke_layout_map_t lm_mfma_scaled_a = {ROCKE_MMA_ROLE_A, 32, 64, mfma_scaled_a}; +static const rocke_layout_map_t lm_mfma_scaled_b = {ROCKE_MMA_ROLE_B, 32, 64, mfma_scaled_b}; +static const rocke_layout_map_t lm_mfma_scaled_c = {ROCKE_MMA_ROLE_ACC, 16, 64, _mfma_acc_32x32}; +static const rocke_layout_map_t lm_mfma_scaled_sa + = {ROCKE_MMA_ROLE_A_SCALE, 1, 64, mfma_scaled_scale_a}; +static const rocke_layout_map_t lm_mfma_scaled_sb + = {ROCKE_MMA_ROLE_B_SCALE, 1, 64, mfma_scaled_scale_b}; + static const rocke_mma_op_t k_mma_gfx950[] = { + {"mfma_scaled", + "fp8e4m3", + "fp8e4m3", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_fp8_fp8_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_f8_a, + &lm_mfma_scaled_f8_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "fp8e4m3", + "bf8e5m2", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_fp8_bf8_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_f8_a, + &lm_mfma_scaled_f8_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "fp8e4m3", + "fp6e2m3", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_fp8_fp6_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_f8_a, + &lm_mfma_scaled_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "fp8e4m3", + "fp6e3m2", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_fp8_bf6_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_f8_a, + &lm_mfma_scaled_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "fp8e4m3", + "fp4e2m1", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_fp8_fp4_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_f8_a, + &lm_mfma_scaled_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "bf8e5m2", + "fp8e4m3", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_bf8_fp8_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_f8_a, + &lm_mfma_scaled_f8_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "bf8e5m2", + "bf8e5m2", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_bf8_bf8_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_f8_a, + &lm_mfma_scaled_f8_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "bf8e5m2", + "fp6e2m3", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_bf8_fp6_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_f8_a, + &lm_mfma_scaled_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "bf8e5m2", + "fp6e3m2", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_bf8_bf6_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_f8_a, + &lm_mfma_scaled_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "bf8e5m2", + "fp4e2m1", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_bf8_fp4_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_f8_a, + &lm_mfma_scaled_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "fp6e2m3", + "fp8e4m3", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_fp6_fp8_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_a, + &lm_mfma_scaled_f8_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "fp6e2m3", + "bf8e5m2", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_fp6_bf8_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_a, + &lm_mfma_scaled_f8_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "fp6e2m3", + "fp6e2m3", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_fp6_fp6_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_a, + &lm_mfma_scaled_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "fp6e2m3", + "fp6e3m2", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_fp6_bf6_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_a, + &lm_mfma_scaled_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "fp6e2m3", + "fp4e2m1", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_fp6_fp4_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_a, + &lm_mfma_scaled_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "fp6e3m2", + "fp8e4m3", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_bf6_fp8_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_a, + &lm_mfma_scaled_f8_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "fp6e3m2", + "bf8e5m2", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_bf6_bf8_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_a, + &lm_mfma_scaled_f8_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "fp6e3m2", + "fp6e2m3", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_bf6_fp6_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_a, + &lm_mfma_scaled_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "fp6e3m2", + "fp6e3m2", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_bf6_bf6_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_a, + &lm_mfma_scaled_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "fp6e3m2", + "fp4e2m1", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_bf6_fp4_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_a, + &lm_mfma_scaled_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "fp4e2m1", + "fp8e4m3", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_fp4_fp8_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_a, + &lm_mfma_scaled_f8_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "fp4e2m1", + "bf8e5m2", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_fp4_bf8_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_a, + &lm_mfma_scaled_f8_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "fp4e2m1", + "fp6e2m3", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_fp4_fp6_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_a, + &lm_mfma_scaled_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "fp4e2m1", + "fp6e3m2", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_fp4_bf6_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_a, + &lm_mfma_scaled_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, + {"mfma_scaled", + "fp4e2m1", + "fp4e2m1", + "fp32", + 32, + 32, + 64, + "mfma_gfx950_f32_32x32x64_fp4_fp4_scale_e8m0_e8m0_k32", + 32, + 32, + 16, + 64, + &lm_mfma_scaled_a, + &lm_mfma_scaled_b, + &lm_mfma_scaled_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 1, + 1, + &lm_mfma_scaled_sa, + &lm_mfma_scaled_sb}, {"mma", "fp32", "fp32", 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 540a27a2f754..f18114e2c112 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 @@ -13,6 +13,7 @@ * rocke_i_attrs / type helpers) lives in bucket 0 (ir_core.c) and is declared * in rocke/ir_internal.h. */ +#include "rocke/mfma_scale_internal.h" #include "rocke/tf32_internal.h" #include #include @@ -71,7 +72,7 @@ static bool rocke_mma_is_int_acc(const char* op_id) static const char* rocke_mma_result_hint(const char* op_id) { const char* family = rocke_arch_mma_op_id_family(op_id); - if(family && strcmp(family, "wmma_scaled") == 0) + if(family && (strcmp(family, "wmma_scaled") == 0 || strcmp(family, "mfma_scaled") == 0)) return "mxacc"; size_t i; if(op_id) @@ -639,6 +640,9 @@ rocke_value_t* rocke_b_mma(rocke_ir_builder_t* b, const char* tf32_error = rocke_tf32_mma_error(op_id, ops, nops); if(tf32_error) return (rocke_value_t*)rocke_i_set_err(b, ROCKE_ERR_VALUE, "%s", tf32_error); + const char* mx_error = rocke_scaled_mfma_error(op_id, ops, nops); + if(mx_error) + return (rocke_value_t*)rocke_i_set_err(b, ROCKE_ERR_VALUE, "%s", mx_error); attrs = rocke_i_attrs(b); rocke_attr_set_str(b, &attrs, "op_id", op_id); return rocke_i_op1(b, ROCKE_OP_TILE_MMA, ops, nops, vt, &attrs, hint); 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 5cc77c73348c..0004ee9968ad 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 @@ -30,6 +30,7 @@ #include "rocke/ir.h" #include "rocke/lower_hip.h" #include "rocke/lower_hip_internal.h" +#include "rocke/mfma_scale_internal.h" #include "rocke/tf32_internal.h" #include "rocke/wmma_scale_internal.h" @@ -135,6 +136,27 @@ 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(const auto* atom = rocke_gfx950_scaled_mfma(op_id)) + { + if(strcmp(lw->arch.gfx, "gfx950") != 0) + return rocke_h_fail(lw, ROCKE_ERR_NOTIMPL, "scaled MFMA requires gfx950"); + if(const char* error = rocke_scaled_mfma_error(op_id, op->operands, op->num_operands)) + return rocke_h_fail(lw, ROCKE_ERR_VALUE, "%s", error); + if(op->num_results != 1 || strcmp(op->results[0]->type->name, "vec") != 0) + return rocke_h_fail(lw, ROCKE_ERR_VALUE, "scaled MFMA result must be vec"); + rocke_h_emitf(lw, + "f32x16 %s = __builtin_amdgcn_mfma_scale_f32_32x32x64_f8f6f4(%s, %s, %s, %d, " + "%d, 0, %s, 0, %s);", + rocke_h_name(lw, op->results[0]), + rocke_h_name(lw, op->operands[0]), + rocke_h_name(lw, op->operands[1]), + rocke_h_name(lw, op->operands[2]), + rocke_mfma_matrix_format(atom->a_dtype), + rocke_mfma_matrix_format(atom->b_dtype), + rocke_h_name(lw, op->operands[3]), + rocke_h_name(lw, op->operands[4])); + return ROCKE_OK; + } int tf32_count = rocke_tf32_mma_count(op_id); if(tf32_count) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/core.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/core.cpp index 94c2faea4d6b..510ca02cedf4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/core.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/core.cpp @@ -1876,6 +1876,18 @@ void rocke_ll_finalize(rocke_lower_t* L, rocke_strbuf_t* out) any_need = true; } } + /* dict.update appends override-only keys after the base keys. */ + for(int i = 0; i < novr; ++i) + { + bool in_base = false; + for(int j = 0; j < ROCKE_LL_INTRINSIC_DECLS_COUNT; ++j) + in_base |= strcmp(ovr[i].key, ROCKE_LL_INTRINSIC_DECLS[j].key) == 0; + if(!in_base && ll_need_has(L, ovr[i].key)) + { + rocke_strbuf_appendf(out, "%s\n", ovr[i].decl); + any_need = true; + } + } /* Dynamic decls not in either static table. */ for(size_t i = 0; i < L->dyn_decls.len; i++) { 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 08b6535eb06f..cbc45ccf05bc 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 @@ -417,6 +417,9 @@ const rocke_ll_decl_t ROCKE_LL_INTRINSIC_DECLS[] = { {"s.prefetch.inst.p0", "declare void @llvm.amdgcn.s.prefetch.inst.p0(ptr, i32)"}, {"s.prefetch.inst.p1", "declare void @llvm.amdgcn.s.prefetch.inst.p1(ptr addrspace(1), i32)"}, {"s.prefetch.inst.p4", "declare void @llvm.amdgcn.s.prefetch.inst.p4(ptr addrspace(4), i32)"}, + {"mfma.scale.f32.32x32x64.f8f6f4.v8i32.v8i32", + "declare <16 x float> @llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4.v8i32.v8i32(<8 x i32>, <8 x " + "i32>, <16 x float>, i32 immarg, i32 immarg, i32 immarg, i32, i32 immarg, i32)"}, /* gfx1250 data prefetch (llvm23). s.prefetch.data is llvm_anyptr_ty like * s.prefetch.inst above; see ROCKE_LL_S_PREFETCH_DATA_PTR_TYPES. */ {"s.setreg", "declare void @llvm.amdgcn.s.setreg(i32 immarg, i32)"}, 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 da86d56c51f2..518ab5738433 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 @@ -19,6 +19,7 @@ * reached from _op_tile_mma's op_id routing, which mirrors the Python CDNA * ISABackend.emit_mma rebuilding ``tile.`` and re-dispatching. */ +#include "rocke/mfma_scale_internal.h" #include "rocke/tf32_internal.h" #include #include @@ -91,6 +92,30 @@ static void _op_tile_mma(rocke_lower_t* L, const rocke_op_t* op) rocke_ll_fail(L, ROCKE_ERR_KEY, "tile.mma: missing op_id attribute"); } + if(const auto* atom = rocke_gfx950_scaled_mfma(op_id)) + { + if(!L->backend || strcmp(L->backend->gfx, "gfx950") != 0) + rocke_ll_fail(L, ROCKE_ERR_NOTIMPL, "scaled MFMA requires gfx950"); + if(const char* error = rocke_scaled_mfma_error(op_id, op->operands, op->num_operands)) + rocke_ll_fail(L, ROCKE_ERR_VALUE, "%s", error); + if(op->num_results != 1 || strcmp(op->results[0]->type->name, "vec") != 0) + rocke_ll_fail(L, ROCKE_ERR_VALUE, "scaled MFMA result must be vec"); + rocke_ll_need(L, "mfma.scale.f32.32x32x64.f8f6f4.v8i32.v8i32"); + rocke_ll_emitf( + L, + " %s = call <16 x float> @llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4.v8i32.v8i32(" + "<8 x i32> %s, <8 x i32> %s, <16 x float> %s, i32 %d, i32 %d, i32 0, i32 %s, i32 0, " + "i32 %s)", + mma_result_name(L, op), + rocke_ll_operand(L, op->operands[0]), + rocke_ll_operand(L, op->operands[1]), + rocke_ll_operand(L, op->operands[2]), + rocke_mfma_matrix_format(atom->a_dtype), + rocke_mfma_matrix_format(atom->b_dtype), + rocke_ll_operand(L, op->operands[3]), + rocke_ll_operand(L, op->operands[4])); + return; + } if(rocke_tf32_mma_count(op_id) && strcmp(L->backend->gfx, "gfx942") != 0) rocke_ll_fail(L, ROCKE_ERR_VALUE, "XF32 MMA requires gfx942"); 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 62afcbad52fe..13494047fdf4 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 @@ -23,6 +23,7 @@ */ #include "rocke/verify.h" +#include "rocke/mfma_scale_internal.h" #include "rocke/tf32_internal.h" #include @@ -356,6 +357,15 @@ static void check_sync_attrs(verifier_t* v, const rocke_op_t* op, const char* co static void check_contract(verifier_t* v, const rocke_op_t* op) { + if(op->opcode == ROCKE_OP_TILE_MMA) + { + const char* id = rocke_attr_get_str(&op->attrs, "op_id"); + if(const char* error = rocke_scaled_mfma_error(id, op->operands, op->num_operands)) + v_errf(v, op, "%s", error); + if(rocke_gfx950_scaled_mfma(id) + && (op->num_results != 1 || strcmp(op->results[0]->type->name, "vec") != 0)) + v_errf(v, op, "scaled MFMA result must be vec"); + } const char* tf32_error = rocke_tf32_op_error(op); if(tf32_error) v_errf(v, op, "%s", tf32_error); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_block_scaled_gemm.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_block_scaled_gemm.h new file mode 100644 index 000000000000..938a0ef3ff77 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_block_scaled_gemm.h @@ -0,0 +1,26 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +#ifndef ROCKE_INSTANCE_GFX950_BLOCK_SCALED_GEMM_H +#define ROCKE_INSTANCE_GFX950_BLOCK_SCALED_GEMM_H +#include "rocke/ir.h" +#ifdef __cplusplus +extern "C" { +#endif +typedef struct rocke_gfx950_block_scaled_gemm_spec +{ + const char* name; + int M, N, K; + const char* dtype_a; + const char* dtype_b; + const char* dtype_c; + int block_k; +} rocke_gfx950_block_scaled_gemm_spec_t; +rocke_gfx950_block_scaled_gemm_spec_t rocke_gfx950_block_scaled_gemm_spec_default(void); +bool rocke_gfx950_block_scaled_gemm_is_valid_spec(const rocke_gfx950_block_scaled_gemm_spec_t* spec, + const char* arch); +rocke_kernel_def_t* rocke_build_gfx950_block_scaled_gemm_new( + rocke_ir_builder_t* b, const rocke_gfx950_block_scaled_gemm_spec_t* spec, const char* arch); +#ifdef __cplusplus +} +#endif +#endif diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/mfma_scale_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/mfma_scale_internal.h new file mode 100644 index 000000000000..9cf5ef0f2940 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/mfma_scale_internal.h @@ -0,0 +1,51 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +#ifndef ROCKE_MFMA_SCALE_INTERNAL_H +#define ROCKE_MFMA_SCALE_INTERNAL_H +#include "rocke/arch_target.h" +#include "rocke/error.hpp" +#include "rocke/ir.h" +#include "rocke/storage.h" +#include +static inline const rocke_mma_op_t* rocke_gfx950_scaled_mfma(const char* id) +{ + const auto* target = rocke_arch_target_from_gfx("gfx950"); + const auto* atom = rocke_mma_catalog_by_op_id(&target->mma, id); + return atom && strcmp(atom->family, "mfma_scaled") == 0 ? atom : NULL; +} +static inline int rocke_mfma_matrix_format(const char* dtype) +{ + const char* formats[] = {"fp8e4m3", "bf8e5m2", "fp6e2m3", "fp6e3m2", "fp4e2m1"}; + for(int i = 0; i < 5; ++i) + if(strcmp(dtype, formats[i]) == 0) + return i; + return -1; +} +static inline rocke_matrix_fragment_layout_t rocke_mfma_matrix_layout(const char* dtype) +{ + int fmt = rocke_mfma_matrix_format(dtype); + int bits = fmt < 2 ? 8 : fmt < 4 ? 6 : 4; + rocke_bit_packing_t packing; + rocke_fragment_packing_t fragment; + rocke_matrix_fragment_layout_t layout; + if(fmt < 0 || !rocke_bit_packing_init(&packing, bits, 0) + || !rocke_fragment_packing_init(&fragment, &packing, 32, 32, 8) + || !rocke_matrix_fragment_layout_init(&layout, &fragment, fmt < 2 ? 16 : 32, 2, 32)) + ckc::raise_status(ROCKE_ERR_VALUE, "unsupported scaled MFMA matrix dtype"); + return layout; +} +static inline const char* rocke_scaled_mfma_error(const char* id, rocke_value_t* const* args, int n) +{ + if(!id || !rocke_arch_mma_op_id_family(id) + || strcmp(rocke_arch_mma_op_id_family(id), "mfma_scaled") != 0) + return NULL; + const char* types[] = {"vec", "vec", "vec", "i32", "i32"}; + if(n != 5) + return "scaled MFMA requires five operands"; + for(int i = 0; i < 5; ++i) + if(!args[i] || strcmp(args[i]->type->name, types[i]) != 0) + return "scaled MFMA requires two vec payloads, vec " + "accumulator and two i32 E8M0 scales"; + return NULL; +} +#endif diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/block_scaled_gemm.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/block_scaled_gemm.cpp new file mode 100644 index 000000000000..eb7659e3ee08 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/block_scaled_gemm.cpp @@ -0,0 +1,141 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +#include "rocke/dtypes.h" +#include "rocke/error_boundary.hpp" +#include "rocke/helper_rocke.helpers.mma_io.h" +#include "rocke/instance_gfx950_block_scaled_gemm.h" +#include "rocke/mfma_scale_internal.h" +#include +#include +#include + +rocke_gfx950_block_scaled_gemm_spec_t rocke_gfx950_block_scaled_gemm_spec_default(void) +{ + return {"block_scaled_gemm", 0, 0, 0, "fp8e4m3", "fp8e4m3", "bf16", 32}; +} + +bool rocke_gfx950_block_scaled_gemm_is_valid_spec(const rocke_gfx950_block_scaled_gemm_spec_t* s, + const char* arch) +{ + if(!s || !arch || strcmp(arch, "gfx950") != 0 || !s->name || !s->dtype_a || !s->dtype_b + || !s->dtype_c) + return false; + char a_scratch[64], b_scratch[64], c_scratch[64]; + const char* a = rocke_normalize_dtype(s->dtype_a, a_scratch, sizeof(a_scratch)); + const char* bb = rocke_normalize_dtype(s->dtype_b, b_scratch, sizeof(b_scratch)); + const char* c = rocke_normalize_dtype(s->dtype_c, c_scratch, sizeof(c_scratch)); + return a && bb && c && rocke_mfma_matrix_format(a) >= 0 && rocke_mfma_matrix_format(bb) >= 0 + && (strcmp(c, "bf16") == 0 || strcmp(c, "fp16") == 0) && s->block_k == 32 && s->M > 0 + && s->N > 0 && s->K > 0 && s->M % 32 == 0 && s->N % 32 == 0 && s->K % 64 == 0 + && (int64_t)s->M * s->K <= INT32_MAX && (int64_t)s->N * s->K <= INT32_MAX + && (int64_t)s->M * s->N <= INT32_MAX; +} + +rocke_kernel_def_t* rocke_build_gfx950_block_scaled_gemm_new( + rocke_ir_builder_t* b, const rocke_gfx950_block_scaled_gemm_spec_t* s, const char* arch) +{ + return ckc::guard_builder(b, [&]() -> rocke_kernel_def_t* { + if(!b || !s) + return NULL; + if(!rocke_gfx950_block_scaled_gemm_is_valid_spec(s, arch)) + { + rocke_ir_builder_init(b, "invalid_scaled_gemm"); + ckc::raise_status(ROCKE_ERR_VALUE, "invalid gfx950 block-scaled GEMM spec"); + } + char a_scratch[64], b_scratch[64], c_scratch[64]; + const char* da = rocke_normalize_dtype(s->dtype_a, a_scratch, sizeof(a_scratch)); + const char* db = rocke_normalize_dtype(s->dtype_b, b_scratch, sizeof(b_scratch)); + const char* dc = rocke_normalize_dtype(s->dtype_c, c_scratch, sizeof(c_scratch)); + if(rocke_ir_builder_init(b, "") != ROCKE_OK) + return NULL; + char* name = rocke_arena_printf( + &b->arena, "%s_mfma_scaled_%s_%s_%s_M%dN%dK%d", s->name, da, db, dc, s->M, s->N, s->K); + if(!name) + ckc::raise_status(ROCKE_ERR_OOM, "cannot allocate scaled GEMM kernel name"); + b->kernel->name = name; + char id[160]; + const char* tokens[] = {"fp8", "bf8", "fp6", "bf6", "fp4"}; + snprintf(id, + sizeof(id), + "mfma_gfx950_f32_32x32x64_%s_%s_scale_e8m0_e8m0_k32", + tokens[rocke_mfma_matrix_format(da)], + tokens[rocke_mfma_matrix_format(db)]); + const auto* atom = rocke_gfx950_scaled_mfma(id); + rocke_attr_set_int(b, &b->kernel->attrs, "max_workgroup_size", 64); + rocke_value_t* params[5]; + const char* names[] = {"A", "B", "A_scale", "B_scale", "C"}; + const char* dtypes[] = {da, db, "e8m0", "e8m0", dc}; + for(int i = 0; i < 5; ++i) + { + rocke_param_opts_t opts = {}; + opts.noalias = true; + opts.noalias_set = true; + if(i == 4) + { + opts.writeonly = true; + opts.writeonly_set = true; + } + else + { + opts.readonly = true; + opts.readonly_set = true; + } + opts.align = (i == 2 || i == 3) ? 1 : 16; + opts.align_set = true; + params[i] = rocke_b_param( + b, names[i], rocke_ptr_type(b, rocke_storage_ir_type(dtypes[i]), "global"), &opts); + } + for(const char* n : {"M", "N", "K"}) + rocke_b_param(b, n, rocke_i32(), NULL); + auto* c32 = rocke_b_const_i32(b, 32); + auto* cN = rocke_b_const_i32(b, s->N); + auto* tid = rocke_b_thread_id_x(b); + auto* c64 = rocke_b_const_i32(b, 64); + auto* lane = rocke_b_mod(b, tid, c64); + auto* row = rocke_b_mod(b, lane, c32); + auto* half = rocke_b_div(b, lane, c32); + auto* m0 = rocke_b_mul(b, rocke_b_block_id_y(b), c32); + auto* n0 = rocke_b_mul(b, rocke_b_block_id_x(b), c32); + auto* ar = rocke_b_add(b, m0, row); + auto* br = rocke_b_add(b, n0, row); + auto* abase = rocke_b_mul( + b, ar, rocke_b_const_i32(b, (int64_t)s->K * rocke_dtype_info(da)->encoded_bits / 8)); + auto* bbase = rocke_b_mul( + b, br, rocke_b_const_i32(b, (int64_t)s->K * rocke_dtype_info(db)->encoded_bits / 8)); + auto* acc = rocke_b_zero_vec_f32(b, 16); + auto al = rocke_mfma_matrix_layout(da), bl = rocke_mfma_matrix_layout(db); + for(int step = 0; step < s->K / 64; ++step) + { + auto* a = rocke_h_load_matrix_fragment( + b, params[0], abase, half, step * 64, da, &al, rocke_i32(), 16); + auto* bb = rocke_h_load_matrix_fragment( + b, params[1], bbase, half, step * 64, db, &bl, rocke_i32(), 16); + auto* group = rocke_b_add(b, rocke_b_const_i32(b, step * 2), half); + auto* ai = rocke_b_add(b, rocke_b_mul(b, ar, rocke_b_const_i32(b, s->K / 32)), group); + auto* bi = rocke_b_add(b, rocke_b_mul(b, group, cN), br); + auto* sa = rocke_b_zext( + b, rocke_b_global_load(b, params[2], ai, rocke_i8(), 1), rocke_i32()); + auto* sb = rocke_b_zext( + b, rocke_b_global_load(b, params[3], bi, rocke_i8(), 1), rocke_i32()); + rocke_value_t* scales[] = {sa, sb}; + acc = rocke_b_mma(b, id, a, bb, acc, scales, 2); + } + for(int slot = 0; slot < 16; ++slot) + { + rocke_value_t *r, *c; + rocke_layout_map_coord(atom->c_layout, b, lane, slot, &r, &c); + auto* rr = rocke_b_add(b, m0, r); + auto* offset = rocke_b_mul(b, rr, cN); + auto* cc = rocke_b_add(b, n0, c); + auto* index = rocke_b_add(b, offset, cc); + rocke_b_global_store(b, + params[4], + index, + rocke_b_cast_f32_to(b, + rocke_b_vec_extract(b, acc, slot), + rocke_storage_ir_type(dc)), + 2); + } + return b->status == ROCKE_OK ? b->kernel : NULL; + }); +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/kernel_taxonomy.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/kernel_taxonomy.md index 1d2ad0dc9b13..82681b6dd5d1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/kernel_taxonomy.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/kernel_taxonomy.md @@ -94,6 +94,7 @@ and exact gfx targets. | `build_wmma_gemm` | [P/gfx1151/wmma_gemm.py](../../python/rocke/instances/gfx1151/wmma_gemm.py), [P/gfx1201/wmma_gemm.py](../../python/rocke/instances/gfx1201/wmma_gemm.py), [P/gfx1250/wmma_gemm.py](../../python/rocke/instances/gfx1250/wmma_gemm.py) | exact-target GEMM | target-catalog WMMA | module's exact gfx target | | `build_wmma_gemm_int8`, `build_wmma_gemm_iu8`, `build_wmma_gemm_iu8_dequant` | [P/gfx1151/wmma_gemm_int8.py](../../python/rocke/instances/gfx1151/wmma_gemm_int8.py), [P/gfx1151/wmma_gemm_iu8.py](../../python/rocke/instances/gfx1151/wmma_gemm_iu8.py), [P/gfx1151/wmma_gemm_iu8_dequant.py](../../python/rocke/instances/gfx1151/wmma_gemm_iu8_dequant.py) | integer GEMM, optionally dequantized | target-catalog integer WMMA plus optional conversion | `gfx1151` | | `build_block_scaled_gemm` | [P/gfx1250/block_scaled_gemm.py](../../python/rocke/instances/gfx1250/block_scaled_gemm.py) | block-scaled low-precision GEMM | target-catalog WMMA plus scale handling | `gfx1250` | +| `build_block_scaled_gemm` | [P/gfx950/block_scaled_gemm.py](../../python/rocke/instances/gfx950/block_scaled_gemm.py) | packed RCR block-scaled low-precision GEMM | target-catalog 32x32x64 scaled MFMA with independent A/B K32 E8M0 scales | `gfx950`; direct spec-driven builder | ## Reduction, Pointwise, And Data-Movement Kernels diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/gfx950_block_scaled_gemm.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/gfx950_block_scaled_gemm.md new file mode 100644 index 000000000000..7f85194e4548 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/gfx950_block_scaled_gemm.md @@ -0,0 +1,45 @@ +# Native gfx950 block-scaled GEMM + +The spec-driven Python and native C++ builders implement densely packed RCR +GEMM: A[M,K], B[N,K], A_scale[M,K/32], B_scale[K/32,N] and C[M,N]. M/N are +positive multiples of 32 and K is a positive multiple of 64. All tensor indices +fit signed i32. M/N/K scalar parameters mirror the specialization and launch +grid; callers must supply the spec's dimensions. Each wave64 owns one 32x32 +output tile, with grid `(N/32, M/32, 1)` and block `(64, 1, 1)`. + +A/B may independently use FP8 E4M3, BF8 E5M2, FP6 E2M3, BF6 E3M2 or FP4 E2M1. +Scales are raw E8M0 bytes, one K32 group per half-wave, with independent A/B +VGPR sources. The kernel uses FP32 accumulation and BF16 or FP16 output. Base +matrix/output pointers must meet their declared 16-byte alignment; scale pointers +need byte alignment. All five buffers are disjoint (`noalias`). + +Use `build_block_scaled_gemm(spec, arch="gfx950")` directly. The four focused +examples construct a spec and invoke the shared verifier to prepare inputs, +compile, launch and compare with an independent decoded reference. The bounded +fixtures accept K64/K128/K256. These are untuned correctness examples +demonstrating native scaled MFMA support. They make no performance claim. + +The native ABI is `rocke_gfx950_block_scaled_gemm_spec_t` and +`rocke_build_gfx950_block_scaled_gemm_new`. The pybind engine exposes native +spec-to-IR, LLVM lowering and verification. Full modules match Python byte for +byte. LLVM 20, 22 and 23 use the same nine-argument overloaded intrinsic ABI. +Other targets reject this atom family. + +[Example commands, storage, ISA references and numerical limits](../../python/rocke/examples/gfx950/gemm/README.md) +explain the four focused entrypoints and their mandatory non-unity scale cases. + +## Relationship to Existing gfx950 MFMA Support + +This builder consumes the 32x32x64 `mfma_scaled` target-catalog family. It +selects the complete shape/dtype/scale contract with `op_for_shape`, loads and +packs the matrices and scales, then emits `IRBuilder.mma(atom, a, b, c, sa, sb)`. +The Python/C++ builder modules are siblings of the gfx1250 block-scaled GEMM +modules, and both instruction shapes share the existing MFMA lowering files. + +The existing `MfmaAtom.emit_scaled()` helper always calls the 16x16x128 +scaled wrapper. Separately, `MfmaAtom.fp8_16x16x128().emit()` is used by the +gfx950 FP8 fused-MoE dispatcher, with constant in-instruction scales and +FP32 block-scale factors applied outside the MFMA. This 32x32x64 builder +passes independent E8M0 scales to the instruction and currently runs directly +from a spec, without a GEMM dispatcher registration. See the +[API and consumer comparison](../reference/mma_metadata.md#gfx950-scaled-mfma-paths). diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/index.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/index.md index deb72217ec72..5c42e2031871 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/index.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/index.md @@ -210,3 +210,11 @@ art = compile_kernel(kernel) `art.hsaco` is ready for `KernelLauncher`; `art.llvm_text` is ready for `analyze_llvm_ir`; `art.timings` is the per-stage codegen budget. For end-to-end execution, see `runtime/compile_launch_and_manifest.md` and `runtime/manifest_schema.md`. + +## Native gfx950 block-scaled GEMM + +[`gfx950/block_scaled_gemm.py`](../../python/rocke/instances/gfx950/block_scaled_gemm.py) +provides `BlockScaledGemmSpec`, `build_block_scaled_gemm`, signature and grid helpers. +It consumes all 25 F8/F6/F4 format pairs with K32 E8M0 scales and a wave64 +32x32x64 native scaled MFMA. See [the family contract](gfx950_block_scaled_gemm.md) +and [F8, F6, F4 and mixed examples](../../python/rocke/examples/gfx950/gemm/README.md). diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/arch/gfx950.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/arch/gfx950.md index e4370d8ad392..36e65d4de758 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/arch/gfx950.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/arch/gfx950.md @@ -30,6 +30,12 @@ historical `21.x` numbering so existing cross-references resolve.) ### 21.1 MFMA Atom Catalog (gfx950) +Native 32x32x64 scaled MFMA is exposed through the target catalog's +`mfma_scaled` family, separately from the `MfmaAtom` factories below. It +supports independent FP8/BF8/FP6/BF6/FP4 inputs and K32 E8M0 scales. See +[MMA entrypoints and dispatcher coverage](../../reference/mma_metadata.md#gfx950-scaled-mfma-paths) +and the [untuned scaled-GEMM examples](../../../python/rocke/examples/gfx950/gemm/README.md). + The full DSL atom catalog lives in `helpers/atoms.py::MFMA_*_ATOMS`. On gfx950 the following are available: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/mfma_atom_catalog.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/mfma_atom_catalog.md index c2d53b3c61db..4ef8e8938f60 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/mfma_atom_catalog.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/mfma_atom_catalog.md @@ -26,6 +26,24 @@ The bf16 16x16x16 path uses the `_1k` variant (`mfma.f32.16x16x16bf16.1k` with ` The wide-K / MX atoms (`fp8_16x16x128`, `fp4_16x16x128`, `fp6_16x16x96`) are gfx950-only — they map onto the `f8f6f4` / MX-scaled MFMA instructions that only the CDNA4 target exposes. +## Native gfx950 32x32x64 Scaled MFMA + +The `helpers.atoms.MfmaAtom` table above is separate from the target catalog's +`mfma_scaled` family. On gfx950, that family exposes 32x32x64 operations for +all 25 ordered FP8/BF8/FP6/BF6/FP4 pairs, with independent A/B E8M0 scales +over K32 and 16 FP32 accumulator values per lane. Select the complete contract +with `ArchTarget.mma.op_for_shape(...)`, then emit it with `IRBuilder.mma(...)`. + +`MfmaAtom.emit_scaled()` remains a fixed 16x16x128 wrapper; it does not select +the 32x32x64 operation. The existing `fp8_16x16x128` atom uses `emit()` and +the scaled intrinsic with constant in-instruction scales. The gfx950 FP8 +fused-MoE dispatcher uses that atom and applies its FP32 block scales outside +the MFMA. The 32x32x64 family currently has direct GEMM examples. + +See [MMA selection, emission and dispatcher paths](mma_metadata.md#selection-ir-emission-and-dispatch), +the [gfx950 GEMM contract](../instances/gfx950_block_scaled_gemm.md), and +[runnable scaled-GEMM examples](../../python/rocke/examples/gfx950/gemm/README.md). + ## Lane Layouts (Output) The kernel epilogue must agree with `MfmaAtom.lane_to_output(b, lane, i)`. The layouts below come from `helpers/atoms.py`. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/mma_metadata.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/mma_metadata.md index 3f9597229488..71a98de09154 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/mma_metadata.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/mma_metadata.md @@ -4,6 +4,52 @@ This reference describes the `MmaOp` metadata and `MmaCatalog` queries used to select matrix operations. For the overall kernel-building workflow, see the [authoring model](../architecture/authoring_model.md). +## Selection, IR Emission, and Dispatch + +These APIs operate at different levels: + +| API | Responsibility | Result | +|---|---|---| +| `target.mma.op_for_shape(...)` | Match one instruction's shape, A/B/C dtypes, family and scale contract in the target catalog | An `MmaOp`; no match returns `None`, ambiguity raises `ValueError`; no IR is emitted | +| `builder.mma(atom, a, b, c, *extra)` | Record the selected operation and its operands | One `tile.mma` IR operation; lowering later emits the LLVM intrinsic selected by `op_id` | +| `MfmaAtom.emit_scaled(...)` | Call the existing named 16x16x128 scaled-MFMA wrapper | Delegates to `builder.mma` with a fixed operation ID, regardless of the helper atom's shape or dtype | +| Operator dispatcher | Select a kernel implementation and specialization for a workload | A kernel builder and spec; that builder selects and emits instructions | + +`mma` also accepts a raw operation-ID string, as the existing named wrappers +do. The selected `MmaOp` supplies accumulator size and dtype; the operation ID +identifies the concrete lowering. Matrix loads, packing and scale-tensor loads +belong to the surrounding kernel builder. A dispatcher is above this authoring +layer: request -> kernel/spec selection -> IR construction -> compilation/launch. +Direct examples may start from a spec and builder without dispatch. + +### gfx950 Scaled-MFMA Paths + +Both instruction shapes have rocKE lowering support, through different +authoring entrypoints: + +| Path | Authoring entrypoint | Per-lane FP32 accumulator | Current consumer | +|---|---|---|---| +| Existing 16x16x128 scaled wrapper | `IRBuilder.mfma_scale_f32_16x16x128_f8f6f4(...)`, also called by `MfmaAtom.emit_scaled(...)` | 4 values | Direct operation authoring; A/B format selectors are fixed to zero in lowering | +| Existing 16x16x128 FP8 atom | `MfmaAtom.fp8_16x16x128().emit(...)` | 4 values | gfx950 FP8 fused-MoE dispatcher; the same scaled intrinsic is used with constant in-instruction scales | +| Native 32x32x64 scaled family | `op_for_shape(family="mfma_scaled", ...)`, then `IRBuilder.mma(atom, a, b, c, a_scale, b_scale)` | 16 values | Direct packed GEMM builder and untuned examples, with independent A/B formats and K32 E8M0 scales | + +The [fused-MoE dispatcher](../../python/rocke/dispatch/families/moe.py) selects +the 16x16x128 FP8 atom on gfx950 and the 16x16x32 FP8 atom on gfx942. Its +[FP8 builder](../../python/rocke/instances/common/moe_fused_mega_fp8.py) uses +`atom.emit()` and applies FP32 block-scale factors outside the MFMA. It does +not call `emit_scaled()` with variable E8M0 operands. The native 32x32x64 +GEMM instead loads raw E8M0 bytes and passes separate A/B scales to the +instruction. It currently has no GEMM dispatcher registration. + +The new 32x32x64 lowering lives beside the existing 16x16x128 lowering in +[`core/lower_llvm.py`](../../python/rocke/core/lower_llvm.py) and +[`cpp/core/lower_llvm/mma.cpp`](../../cpp/core/lower_llvm/mma.cpp). Its catalog +and lane maps are in the existing architecture layer. The Python/C++ +`instances/gfx950/block_scaled_gemm` builders and `examples/gfx950/gemm` +entrypoints mirror the gfx1250 scaled-GEMM structure. See the +[gfx950 family contract](../instances/gfx950_block_scaled_gemm.md) and +[example guide](../../python/rocke/examples/gfx950/gemm/README.md). + ## Operand Metadata MMA metadata uses `a_dtype`, `b_dtype`, and `c_dtype`, with C describing both @@ -62,6 +108,26 @@ assert atom is not None # result = builder.mma(atom, a, b, c, scale_a, scale_b) ``` +For gfx950, select the 32x32x64 scaled-MFMA family in the same way. This +example selects FP8 A and FP4 B, with independent E8M0 scales over K32: + +```python +atom = ArchTarget.from_gfx("gfx950").mma.op_for_shape( + family="mfma_scaled", + a_dtype="fp8", b_dtype="fp4", c_dtype="fp32", + scales=(MmaScaleDType.E8M0, MmaScaleDType.E8M0, MmaScaleBlockK.K32), + m=32, n=32, k=64, +) +assert atom is not None +# result = builder.mma(atom, a, b, c, scale_a, scale_b) +``` + +This operation takes two `vec` packed payloads, a `vec` +accumulator and two i32 scale operands. The catalog describes all 25 ordered +FP8/BF8/FP6/BF6/FP4 pairs. The scale-group K=32 is separate from instruction +K=64. The existing 16x16x128 scaled wrapper retains its fixed operation ID; +`MfmaAtom.emit_scaled()` does not perform this catalog lookup. + Omitting `scales` leaves scales unconstrained. Passing `(a_type, b_type, block_k)` matches the scale contract exactly; `(None, None, None)` selects unscaled atoms. Partially specified scale contracts are invalid. Enumeration 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 f87a7c7387d6..b34dfcb6c92e 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 @@ -96,7 +96,32 @@ {"family": "mma", "a": "bf8", "b": "bf8", "c": "fp32", "m": 16, "n": 16, "k": 32, "op_id": "mfma_f32_16x16x32_bf8"}, {"family": "mma", "a": "bf8", "b": "bf8", "c": "fp32", "m": 32, "n": 32, "k": 16, "op_id": "mfma_f32_32x32x16_bf8"}, {"family": "mma", "a": "fp4", "b": "fp4", "c": "fp32", "m": 16, "n": 16, "k": 128, "op_id": "mfma_f32_16x16x128_fp4"}, - {"family": "mma", "a": "fp6", "b": "fp6", "c": "fp32", "m": 16, "n": 16, "k": 96, "op_id": "mfma_f32_16x16x96_fp6"} + {"family": "mma", "a": "fp6", "b": "fp6", "c": "fp32", "m": 16, "n": 16, "k": 96, "op_id": "mfma_f32_16x16x96_fp6"}, + {"family": "mfma_scaled", "a": "fp8e4m3", "b": "fp8e4m3", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_fp8_fp8_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "fp8e4m3", "b": "bf8e5m2", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_fp8_bf8_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "fp8e4m3", "b": "fp6e2m3", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_fp8_fp6_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "fp8e4m3", "b": "fp6e3m2", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_fp8_bf6_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "fp8e4m3", "b": "fp4e2m1", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_fp8_fp4_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "bf8e5m2", "b": "fp8e4m3", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_bf8_fp8_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "bf8e5m2", "b": "bf8e5m2", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_bf8_bf8_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "bf8e5m2", "b": "fp6e2m3", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_bf8_fp6_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "bf8e5m2", "b": "fp6e3m2", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_bf8_bf6_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "bf8e5m2", "b": "fp4e2m1", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_bf8_fp4_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "fp6e2m3", "b": "fp8e4m3", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_fp6_fp8_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "fp6e2m3", "b": "bf8e5m2", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_fp6_bf8_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "fp6e2m3", "b": "fp6e2m3", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_fp6_fp6_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "fp6e2m3", "b": "fp6e3m2", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_fp6_bf6_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "fp6e2m3", "b": "fp4e2m1", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_fp6_fp4_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "fp6e3m2", "b": "fp8e4m3", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_bf6_fp8_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "fp6e3m2", "b": "bf8e5m2", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_bf6_bf8_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "fp6e3m2", "b": "fp6e2m3", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_bf6_fp6_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "fp6e3m2", "b": "fp6e3m2", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_bf6_bf6_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "fp6e3m2", "b": "fp4e2m1", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_bf6_fp4_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "fp4e2m1", "b": "fp8e4m3", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_fp4_fp8_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "fp4e2m1", "b": "bf8e5m2", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_fp4_bf8_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "fp4e2m1", "b": "fp6e2m3", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_fp4_fp6_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "fp4e2m1", "b": "fp6e3m2", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_fp4_bf6_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32}, + {"family": "mfma_scaled", "a": "fp4e2m1", "b": "fp4e2m1", "c": "fp32", "m": 32, "n": 32, "k": 64, "op_id": "mfma_gfx950_f32_32x32x64_fp4_fp4_scale_e8m0_e8m0_k32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32} ] }, "gfx1151": { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm20.json b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm20.json index 15e68c02ef53..66caa049132d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm20.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm20.json @@ -72,6 +72,7 @@ "mfma.f32.32x32x8f16": "llvm.amdgcn.mfma.f32.32x32x8f16", "mfma.f32.4x4x4f16": "llvm.amdgcn.mfma.f32.4x4x4f16", "mfma.scale.f32.16x16x128.f8f6f4": "llvm.amdgcn.mfma.scale.f32.16x16x128.f8f6f4", + "mfma.scale.f32.32x32x64.f8f6f4.v8i32.v8i32": "llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4.v8i32.v8i32", "minnum.bf16": "llvm.minnum.bf16", "minnum.f16": "llvm.minnum.f16", "minnum.f32": "llvm.minnum.f32", @@ -3157,6 +3158,42 @@ "verified_on": "llvm20" } }, + "mfma.scale.f32.32x32x64.f8f6f4.v8i32.v8i32": { + "gfx11-generic": { + "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4", + "status": "arch_absent", + "verified_on": "llvm20" + }, + "gfx1151": { + "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4", + "status": "arch_absent", + "verified_on": "llvm20" + }, + "gfx1201": { + "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4", + "status": "arch_absent", + "verified_on": "llvm20" + }, + "gfx1250": { + "evidence": "invalid target ID", + "status": "target_unsupported", + "verified_on": "llvm20" + }, + "gfx90a": { + "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4", + "status": "arch_absent", + "verified_on": "llvm20" + }, + "gfx942": { + "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4", + "status": "arch_absent", + "verified_on": "llvm20" + }, + "gfx950": { + "status": "ok", + "verified_on": "llvm20" + } + }, "minnum.bf16": { "gfx11-generic": { "status": "ok", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm22.json b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm22.json index f56744efb1ca..33babccf53d8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm22.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm22.json @@ -81,6 +81,7 @@ "mfma.f32.32x32x8bf16.1k": "llvm.amdgcn.mfma.f32.32x32x8bf16.1k", "mfma.f32.32x32x8f16": "llvm.amdgcn.mfma.f32.32x32x8f16", "mfma.f32.4x4x4f16": "llvm.amdgcn.mfma.f32.4x4x4f16", + "mfma.scale.f32.32x32x64.f8f6f4.v8i32.v8i32": "llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4.v8i32.v8i32", "minnum.bf16": "llvm.minnum.bf16", "minnum.f16": "llvm.minnum.f16", "minnum.f32": "llvm.minnum.f32", @@ -3104,6 +3105,42 @@ "verified_on": "llvm22" } }, + "mfma.scale.f32.32x32x64.f8f6f4.v8i32.v8i32": { + "gfx11-generic": { + "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4", + "status": "arch_absent", + "verified_on": "llvm22" + }, + "gfx1151": { + "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4", + "status": "arch_absent", + "verified_on": "llvm22" + }, + "gfx1201": { + "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4", + "status": "arch_absent", + "verified_on": "llvm22" + }, + "gfx1250": { + "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4", + "status": "arch_absent", + "verified_on": "llvm22" + }, + "gfx90a": { + "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4", + "status": "arch_absent", + "verified_on": "llvm22" + }, + "gfx942": { + "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4", + "status": "arch_absent", + "verified_on": "llvm22" + }, + "gfx950": { + "status": "ok", + "verified_on": "llvm22" + } + }, "minnum.bf16": { "gfx11-generic": { "status": "ok", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm23.json b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm23.json index 3277414f7d66..99e32cbd37d5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm23.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm23.json @@ -87,6 +87,7 @@ "mfma.f32.32x32x8f16": "llvm.amdgcn.mfma.f32.32x32x8f16", "mfma.f32.4x4x4f16": "llvm.amdgcn.mfma.f32.4x4x4f16", "mfma.scale.f32.16x16x128.f8f6f4": "llvm.amdgcn.mfma.scale.f32.16x16x128.f8f6f4.v8i32.v8i32", + "mfma.scale.f32.32x32x64.f8f6f4.v8i32.v8i32": "llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4.v8i32.v8i32", "minnum.bf16": "llvm.minnum.bf16", "minnum.f16": "llvm.minnum.f16", "minnum.f32": "llvm.minnum.f32", @@ -3086,6 +3087,42 @@ "verified_on": "llvm23" } }, + "mfma.scale.f32.32x32x64.f8f6f4.v8i32.v8i32": { + "gfx11-generic": { + "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4", + "status": "arch_absent", + "verified_on": "llvm23" + }, + "gfx1151": { + "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4", + "status": "arch_absent", + "verified_on": "llvm23" + }, + "gfx1201": { + "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4", + "status": "arch_absent", + "verified_on": "llvm23" + }, + "gfx1250": { + "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4", + "status": "arch_absent", + "verified_on": "llvm23" + }, + "gfx90a": { + "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4", + "status": "arch_absent", + "verified_on": "llvm23" + }, + "gfx942": { + "evidence": "fatal error: error in backend: Cannot select: intrinsic %llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4", + "status": "arch_absent", + "verified_on": "llvm23" + }, + "gfx950": { + "status": "ok", + "verified_on": "llvm23" + } + }, "minnum.bf16": { "gfx11-generic": { "status": "ok", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/mfma_scale.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/mfma_scale.py new file mode 100644 index 000000000000..0c42f914e2a3 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/mfma_scale.py @@ -0,0 +1,44 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""gfx950 scaled MFMA: 32 logical values and one K32 E8M0 scale per lane.""" + +from ..dtypes import dtype_info +from ..storage import BitPacking, FragmentPacking, MatrixFragmentLayout +from .target import ArchTarget, _op_id_family + +MATRIX_FORMATS = {"fp8e4m3": 0, "bf8e5m2": 1, "fp6e2m3": 2, "fp6e3m2": 3, "fp4e2m1": 4} +INTRINSIC_KEY = "mfma.scale.f32.32x32x64.f8f6f4.v8i32.v8i32" + + +def gfx950_scaled_mfma(op_id): + if _op_id_family().get(op_id.removeprefix("tile.")) != "mfma_scaled": + return None + atom = ArchTarget.from_gfx("gfx950").mma.by_op_id(op_id.removeprefix("tile.")) + return atom if atom is not None and atom.family == "mfma_scaled" else None + + +def matrix_layout(dtype): + if dtype not in MATRIX_FORMATS: + raise ValueError("unsupported scaled MFMA matrix dtype") + return MatrixFragmentLayout( + FragmentPacking(BitPacking(dtype_info(dtype).encoded_bits), 32, 32, 8), + 16 if MATRIX_FORMATS[dtype] < 2 else 32, + 2, + 32, + ) + + +def scaled_mfma_error(op_id, operands, results=None): + if gfx950_scaled_mfma(op_id) is None: + return None + if [v.type.name for v in operands] != [ + "vec", + "vec", + "vec", + "i32", + "i32", + ]: + return "scaled MFMA requires two vec payloads, vec accumulator and two i32 E8M0 scales" + if results is not None and [v.type.name for v in results] != ["vec"]: + return "scaled MFMA result must be vec" + return None 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 11c1a2a9427a..6f9290e5f393 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 @@ -1037,6 +1037,66 @@ class _FragInfo: ) +def _mfma_scaled_a(builder, lane, slot): + row = builder.mod(lane, builder.const_i32(32)) + k = builder.add( + builder.mul(builder.div(lane, builder.const_i32(32)), builder.const_i32(32)), + builder.const_i32(slot), + ) + return row, k + + +def _mfma_scaled_f8_a(builder, lane, slot): + row = builder.mod(lane, builder.const_i32(32)) + k = builder.add( + builder.mul(builder.div(lane, builder.const_i32(32)), builder.const_i32(16)), + builder.const_i32((slot // 16) * 32 + slot % 16), + ) + return row, k + + +def _mfma_scaled_f8_b(builder, lane, slot): + row, k = _mfma_scaled_f8_a(builder, lane, slot) + return k, row + + +def _mfma_scaled_b(builder, lane, slot): + row, k = _mfma_scaled_a(builder, lane, slot) + return k, row + + +def _mfma_scaled_scale_a(builder, lane, slot): + return builder.mod(lane, builder.const_i32(32)), builder.div( + lane, builder.const_i32(32) + ) + + +def _mfma_scaled_scale_b(builder, lane, slot): + row, k = _mfma_scaled_scale_a(builder, lane, slot) + return k, row + + +_MMA_FRAGMENT_INFO.update( + { + f"mfma_gfx950_f32_32x32x64_{a}_{b}_scale_e8m0_e8m0_k32": _FragInfo( + 32, + 32, + 16, + 64, + _mfma_scaled_f8_a if a in ("fp8", "bf8") else _mfma_scaled_a, + _mfma_scaled_f8_b if b in ("fp8", "bf8") else _mfma_scaled_b, + _mfma_acc_32x32, + 1, + 1, + _mfma_scaled_scale_a, + _mfma_scaled_scale_b, + ) + for a in ("fp8", "bf8", "fp6", "bf6", "fp4") + for b in ("fp8", "bf8", "fp6", "bf6", "fp4") + } +) + + def _frag_info(op_id: str) -> _FragInfo: info = _MMA_FRAGMENT_INFO.get(op_id) if info is None: 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 575d34f99582..45c63cfa8bbf 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 @@ -2283,7 +2283,11 @@ def mma( op_id = op.op_id if hasattr(op, "op_id") else str(op) from .tf32 import tf32_mma_error - error = tf32_mma_error(op_id, [a, b, c, *extra]) + from .arch.mfma_scale import scaled_mfma_error + + error = scaled_mfma_error(op_id, [a, b, c, *extra]) or tf32_mma_error( + op_id, [a, b, c, *extra] + ) if error: raise ValueError(error) c_frag_len = ( @@ -2299,7 +2303,7 @@ def mma( c_elem = I32 if is_int_acc else F32 hint = ( "mxacc" - if _arch._op_id_family().get(op_id) == "wmma_scaled" + if _arch._op_id_family().get(op_id) in ("wmma_scaled", "mfma_scaled") else _MMA_RESULT_HINT.get(op_id, "acc") ) return self._op( @@ -2662,17 +2666,16 @@ def mfma_scale_f32_16x16x128_f8f6f4( a_scale: Value, b_scale: Value, ) -> Value: - """MX-MFMA with in-instruction scales (P15). + """Existing gfx950 16x16x128 MFMA with in-instruction scales (P15). - ``mfma_scale_f32_16x16x128_f8f6f4`` is the gfx950 MX MFMA - instruction family. Operands are 16-byte packed mantissa - vectors (fp8/fp6/fp4) plus per-warp E8M0 scales applied - in-instruction. The scale broadcast inside the instruction is - per-output-row (vs per-input-row), which is what fixes the - B_MX1 row-aware correctness gap that the post-hoc - scale-apply chain has. + This fixed-operation wrapper delegates to ``mma``. LLVM lowering + uses two ``<8 x i32>`` payloads, a ``<4 x float>`` accumulator and + separate i32 scale operands; A/B format selectors are fixed to zero. + It lowers to ``llvm.amdgcn.mfma.scale.f32.16x16x128.f8f6f4``. - Lowers to ``llvm.amdgcn.mfma.scale.f32.16x16x128.f8f6f4``. + For 32x32x64 with independently selected A/B formats and K32 E8M0 + scales, resolve a ``mfma_scaled`` catalog atom and call ``mma`` + with that atom and both scales. See ``reference/mma_metadata.md``. """ return self.mma("mfma_scale_f32_16x16x128_f8f6f4", a, b, c, a_scale, b_scale) 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 76723b53bf53..64c21a678e91 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 @@ -136,6 +136,11 @@ def emit_mma(self, lowerer, op) -> None: from ..ir import Op op_id = op.attrs["op_id"] + from ..arch.mfma_scale import gfx950_scaled_mfma + + if gfx950_scaled_mfma(op_id) is not None: + lowerer._emit_gfx950_scaled_mfma(op) + return legacy = Op( name=f"tile.{op_id}", operands=list(op.operands), @@ -362,6 +367,11 @@ def emit_mma(self, lowerer, op) -> None: from ..ir import Op op_id = op.attrs["op_id"] + from ..arch.mfma_scale import gfx950_scaled_mfma + + if gfx950_scaled_mfma(op_id) is not None: + lowerer._emit_gfx950_scaled_mfma(op) + return legacy = Op( name=f"tile.{op_id}", operands=list(op.operands), 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 3b2dc303e470..2288b209d071 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 @@ -631,6 +631,26 @@ def _op_tile_mma(self, op: Op) -> None: from .tf32 import TF32_MMA op_id = op.attrs["op_id"] + from .arch.mfma_scale import ( + gfx950_scaled_mfma, + scaled_mfma_error, + MATRIX_FORMATS, + ) + + atom = gfx950_scaled_mfma(op_id) + if atom is not None: + if self.arch.gfx != "gfx950": + raise NotImplementedError("scaled MFMA requires gfx950") + error = scaled_mfma_error(op_id, op.operands, op.results) + if error: + raise ValueError(error) + a, b, c, sa, sb = op.operands + self._emit( + f"f32x16 {_name(op.result)} = __builtin_amdgcn_mfma_scale_f32_32x32x64_f8f6f4(" + f"{_name(a)}, {_name(b)}, {_name(c)}, {MATRIX_FORMATS[atom.a_dtype]}, {MATRIX_FORMATS[atom.b_dtype]}, " + f"0, {_name(sa)}, 0, {_name(sb)});" + ) + return if op_id in TF32_MMA: if self.arch.gfx != "gfx942": raise ValueError("XF32 MMA requires gfx942") 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 0d104e53c528..5b0bc2dafd65 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 @@ -955,6 +955,7 @@ def _resolve_llvm_flavor() -> str: "s.prefetch.inst.p4": ( "declare void @llvm.amdgcn.s.prefetch.inst.p4(ptr addrspace(4), i32)" ), + "mfma.scale.f32.32x32x64.f8f6f4.v8i32.v8i32": "declare <16 x float> @llvm.amdgcn.mfma.scale.f32.32x32x64.f8f6f4.v8i32.v8i32(<8 x i32>, <8 x i32>, <16 x float>, i32 immarg, i32 immarg, i32 immarg, i32, i32 immarg, i32)", # gfx1250 data prefetch (llvm23). ``s.prefetch.data`` is overloaded by # address space exactly like ``s.prefetch.inst``; see # _S_PREFETCH_DATA_PTR_TYPES. @@ -3654,6 +3655,29 @@ def _op_tile_mfma_f32_32x32x16_bf8(self, op: Op) -> None: op, dtype="bf8", out_vec=16, intrinsic="32x32x16.bf8.bf8" ) + def _emit_gfx950_scaled_mfma(self, op: Op) -> None: + from .arch.mfma_scale import ( + gfx950_scaled_mfma, + scaled_mfma_error, + MATRIX_FORMATS, + INTRINSIC_KEY, + ) + + atom = gfx950_scaled_mfma(op.attrs["op_id"]) + if self._backend.arch.gfx != "gfx950": + raise NotImplementedError("scaled MFMA requires gfx950") + error = scaled_mfma_error(atom.op_id, op.operands, op.results) + if error: + raise ValueError(error) + a, b, c, sa, sb = op.operands + self._need(INTRINSIC_KEY) + self._current().emit( + f" {op.result.name} = call <16 x float> @llvm.amdgcn.{INTRINSIC_KEY}(" + f"<8 x i32> {self._operand(a)}, <8 x i32> {self._operand(b)}, <16 x float> {self._operand(c)}, " + f"i32 {MATRIX_FORMATS[atom.a_dtype]}, i32 {MATRIX_FORMATS[atom.b_dtype]}, " + f"i32 0, i32 {self._operand(sa)}, i32 0, i32 {self._operand(sb)})" + ) + def _op_tile_mfma_scale_f32_16x16x128_f8f6f4(self, op: Op) -> None: """Lower P15 MX MFMA scaled intrinsic. 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 8a2dbdca93c8..2cb447d78f97 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 @@ -285,7 +285,14 @@ def _check_op_regions(self, op: Op, scope: Dict[str, Value]) -> None: def _check_contract(self, op: Op, scope: Dict[str, Value]) -> None: from .tf32 import tf32_op_error - error = tf32_op_error(op) + from .arch.mfma_scale import scaled_mfma_error + + mx_error = ( + scaled_mfma_error(op.attrs.get("op_id", ""), op.operands, op.results) + if op.name == "tile.mma" + else None + ) + error = mx_error or tf32_op_error(op) if error: self.err(error, op) name = op.name diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/common/scaled_gemm_reference.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/common/scaled_gemm_reference.py new file mode 100644 index 000000000000..c0733e2ce271 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/common/scaled_gemm_reference.py @@ -0,0 +1,263 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Independent finite low-bit GEMM fixtures and references shared by examples.""" + +from __future__ import annotations + +import numpy as np + +from ...core.dtypes import normalize_dtype +from ...core.storage import BitPacking + + +def _canon_lowbit(dtype): + return { + "fp8e4m3": "fp8", + "bf8e5m2": "bf8", + "fp6e2m3": "fp6", + "fp6e3m2": "bf6", + "fp4e2m1": "fp4", + }[normalize_dtype(dtype)] + + +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 pack_fp6_codes(codes: np.ndarray) -> np.ndarray: + """Pack four consecutive six-bit codes into three little-endian bytes.""" + if ( + codes.dtype != np.uint8 + or codes.ndim != 2 + or codes.shape[1] % 4 + or np.any(codes > 63) + ): + raise ValueError( + "expected rank-2 uint8 FP6 codes in [0, 63] and K divisible by 4" + ) + packed = b"".join(BitPacking(6).pack(row.tolist()) for row in codes) + return np.frombuffer(packed, dtype=np.uint8).reshape(codes.shape[0], -1) + + +def decode_fp6(packed: np.ndarray, dtype: str = "fp6") -> np.ndarray: + """Decode packed E2M3 or E3M2, including finite extrema and signed zero.""" + if packed.dtype != np.uint8 or packed.ndim != 2 or packed.shape[1] % 3: + raise ValueError("expected rank-2 uint8 FP6 bytes with row size divisible by 3") + if dtype not in ("fp6", "fp6e2m3", "bf6", "fp6e3m2"): + raise ValueError("FP6 dtype must be fp6/E2M3 or bf6/E3M2") + # Bit addressing is independent of the packer's four-code groups. + bits = np.unpackbits(packed, axis=1, bitorder="little").reshape( + packed.shape[0], -1, 6 + ) + codes = (bits * (1 << np.arange(6))).sum(axis=-1) + mantissa_bits, bias = (3, 1) if _canon_lowbit(dtype) == "fp6" else (2, 3) + fraction = (codes & ((1 << mantissa_bits) - 1)) / (1 << mantissa_bits) + exponent = (codes & 31) >> mantissa_bits + magnitude = np.where( + exponent == 0, + np.ldexp(fraction, 1 - bias), + np.ldexp(1 + fraction, exponent - bias), + ) + return np.copysign(magnitude, np.where(codes & 32, -1.0, 1.0)) + + +def _output_dtype(dtype: str): + import ml_dtypes + + if dtype == "bf16": + return ml_dtypes.bfloat16 + if dtype in ("fp16", "f16"): + return np.float16 + raise ValueError(f"verifier output must be bf16/fp16 (got {dtype!r})") + + +def decode_fp4(packed: np.ndarray) -> np.ndarray: + """Decode low-nibble-first E2M1 pairs by the format's exponent formula.""" + if packed.dtype != np.uint8 or packed.ndim != 2: + raise ValueError("expected a rank-2 uint8 packed FP4 matrix") + codes = np.empty((packed.shape[0], packed.shape[1] * 2), dtype=np.uint8) + codes[:, 0::2] = packed & 15 + codes[:, 1::2] = packed >> 4 + exponent = ((codes >> 1) & 3).astype(np.int32) + mantissa = (codes & 1).astype(np.float64) + magnitude = np.where( + exponent == 0, mantissa * 0.5, np.ldexp(1.0 + mantissa * 0.5, exponent - 1) + ) + return np.copysign(magnitude, np.where(codes & 8, -1.0, 1.0)) + + +def pack_fp4_codes(codes: np.ndarray) -> np.ndarray: + """Pack E2M1 codes, not arbitrary floating values; no quantization policy.""" + if ( + codes.dtype != np.uint8 + or codes.ndim != 2 + or codes.shape[1] % 2 + or np.any(codes > 15) + ): + raise ValueError("expected rank-2 uint8 FP4 codes in [0, 15] and even K") + return codes[:, 0::2] | (codes[:, 1::2] << 4) + + +def reference_result( + a: np.ndarray, + b: np.ndarray, + a_scale: np.ndarray, + b_scale: np.ndarray, + block_k: int, + *, + native: bool, + dtype_a: str | None = None, + dtype_b: str | None = None, + dtype_c: str = "bf16", +) -> np.ndarray: + """Expand scales onto logical A/B elements, multiply, then round the output. + + 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 output-type rounding provide an independent oracle. + FP6 all-code fixtures isolate one K element, avoiding accumulation error. + Input decoding preserves signed zero, but this reference and its numerical + comparison do not establish output zero-sign or accumulation sign semantics. + FP4 fixtures cover all E2M1 values (magnitude <=6), scales 2**[-2,1], + and K<=256: absolute partial sums are below 2**22 units of 2**-6, + so FP32 accumulation is also exact before the final output-type rounding. + """ + output_dtype = _output_dtype(dtype_c) + + 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) + + def matrix_values(data, dtype): + if dtype in ("fp6", "fp6e2m3", "bf6", "fp6e3m2"): + return decode_fp6(data, dtype) + if dtype in ("fp4", "fp4e2m1") or (dtype is None and data.dtype == np.uint8): + return decode_fp4(data) + return data.astype(np.float64) + + a_values = matrix_values(a, dtype_a) + b_values = matrix_values(b, dtype_b) + scaled_a = a_values * np.repeat(sa, block_k, axis=1) + scaled_b = b_values * np.repeat(sb.T, block_k, axis=1) + ref = scaled_a @ scaled_b.T + if native: + ref = ref.astype(output_dtype) + return ref.astype(np.float32) + + +def make_case_inputs( + spec, case: str +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Create asymmetric mantissas and scales, optionally isolating a K group.""" + import ml_dtypes + + native = spec.resolved_matrix_path() != "wmma" + if native and spec.K not in (64, 128, 256): + raise ValueError("exact native fixtures require K=64, 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) + + def operand(dtype, rows): + kind = _canon_lowbit(dtype) + if kind == "fp4": + # Keep E2M1 codes unpacked until group masking is finished. + return rng.integers(0, 16, size=(rows, spec.K), dtype=np.uint8) + if kind in ("fp6", "bf6"): + # Small dyadic values ensure exact FP32 partial sums in these tests. + codes = rng.integers( + 0, 9 if kind == "fp6" else 13, size=(rows, spec.K), dtype=np.uint8 + ) + return codes | (rng.integers(0, 2, size=codes.shape, dtype=np.uint8) << 5) + magnitude = 0.25 if native else 0.5 + return (rng.integers(-4, 5, size=(rows, spec.K)) * magnitude).astype( + lowbit_types[dtype] + ) + + a, b = operand(spec.dtype_a, spec.M), operand(spec.dtype_b, spec.N) + if native: + small = any( + _canon_lowbit(d) in ("fp6", "bf6", "fp4") + for d in (spec.dtype_a, spec.dtype_b) + ) + + sa = rng.integers( + 125, 129 if small else 131, size=(spec.M, groups), dtype=np.uint8 + ) + sb = rng.integers( + 125, 129 if small else 131, size=(groups, spec.N), dtype=np.uint8 + ) + neutral_a = neutral_b = 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_a = neutral_b = 1.0 + if case in ("neutral", "b-only"): + sa.fill(neutral_a) + if case in ("neutral", "a-only"): + sb.fill(neutral_b) + if case.startswith("codes-"): + if ( + _canon_lowbit(spec.dtype_a) not in ("fp6", "bf6") + or _canon_lowbit(spec.dtype_b) not in ("fp6", "bf6") + or min(spec.M, spec.N) < 64 + ): + raise ValueError("codes-N requires FP6/BF6 operands and M/N >= 64") + k = int(case.removeprefix("codes-")) + if not 0 <= k < spec.K: + raise ValueError("codes-N K index outside the matrix") + a.fill(0) + b.fill(0) + a[:, k] = np.arange(spec.M, dtype=np.uint8) % 64 + b[:, k] = np.arange(spec.N, dtype=np.uint8) % 64 + sa.fill(neutral_a) + sb.fill(neutral_b) + elif 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}") + + def pack(data, dtype): + kind = _canon_lowbit(dtype) + if kind == "fp4": + return pack_fp4_codes(data) + return pack_fp6_codes(data) if kind in ("fp6", "bf6") else data + + a, b = pack(a, spec.dtype_a), pack(b, spec.dtype_b) + return a, b, sa, sb + + +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. + + Exact mode requires numerical equality after output-type rounding, with + positive and negative zero treated as equal. Tolerance applies only to the + legacy software-scaled WMMA path. + """ + 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]}" + ) 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 a6a5013d7a5f..0075e1a04dc8 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 @@ -29,262 +29,43 @@ import numpy as np +from ....core.dtypes import normalize_dtype from ....helpers import compile_kernel from ....helpers.compile import compile_kernel_via_hipcc -from ....core.dtypes import normalize_dtype -from ....core.storage import BitPacking from ....instances.gfx1250.block_scaled_gemm import ( _LOWBIT_DTYPES, BlockScaledGemmSpec, - _canon_lowbit, block_scaled_gemm_grid, build_block_scaled_gemm, is_valid_spec, ) from ....runtime.hip_module import Runtime, get_device_arch +# Keep the reference helpers importable from this established example module. +from ...common.scaled_gemm_reference import ( + _output_dtype, + check_result, + decode_e8m0, + decode_fp4, + decode_fp6, + make_case_inputs, + pack_fp4_codes, + pack_fp6_codes, + reference_result, +) -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 pack_fp6_codes(codes: np.ndarray) -> np.ndarray: - """Pack four consecutive six-bit codes into three little-endian bytes.""" - if ( - codes.dtype != np.uint8 - or codes.ndim != 2 - or codes.shape[1] % 4 - or np.any(codes > 63) - ): - raise ValueError( - "expected rank-2 uint8 FP6 codes in [0, 63] and K divisible by 4" - ) - packed = b"".join(BitPacking(6).pack(row.tolist()) for row in codes) - return np.frombuffer(packed, dtype=np.uint8).reshape(codes.shape[0], -1) - - -def decode_fp6(packed: np.ndarray, dtype: str = "fp6") -> np.ndarray: - """Decode packed E2M3 or E3M2, including finite extrema and signed zero.""" - if packed.dtype != np.uint8 or packed.ndim != 2 or packed.shape[1] % 3: - raise ValueError("expected rank-2 uint8 FP6 bytes with row size divisible by 3") - if dtype not in ("fp6", "fp6e2m3", "bf6", "fp6e3m2"): - raise ValueError("FP6 dtype must be fp6/E2M3 or bf6/E3M2") - # Bit addressing is independent of the packer's four-code groups. - bits = np.unpackbits(packed, axis=1, bitorder="little").reshape( - packed.shape[0], -1, 6 - ) - codes = (bits * (1 << np.arange(6))).sum(axis=-1) - mantissa_bits, bias = (3, 1) if _canon_lowbit(dtype) == "fp6" else (2, 3) - fraction = (codes & ((1 << mantissa_bits) - 1)) / (1 << mantissa_bits) - exponent = (codes & 31) >> mantissa_bits - magnitude = np.where( - exponent == 0, - np.ldexp(fraction, 1 - bias), - np.ldexp(1 + fraction, exponent - bias), - ) - return np.copysign(magnitude, np.where(codes & 32, -1.0, 1.0)) - - -def _output_dtype(dtype: str): - import ml_dtypes - - if dtype == "bf16": - return ml_dtypes.bfloat16 - if dtype in ("fp16", "f16"): - return np.float16 - raise ValueError(f"verifier output must be bf16/fp16 (got {dtype!r})") - - -def decode_fp4(packed: np.ndarray) -> np.ndarray: - """Decode low-nibble-first E2M1 pairs by the format's exponent formula.""" - if packed.dtype != np.uint8 or packed.ndim != 2: - raise ValueError("expected a rank-2 uint8 packed FP4 matrix") - codes = np.empty((packed.shape[0], packed.shape[1] * 2), dtype=np.uint8) - codes[:, 0::2] = packed & 15 - codes[:, 1::2] = packed >> 4 - exponent = ((codes >> 1) & 3).astype(np.int32) - mantissa = (codes & 1).astype(np.float64) - magnitude = np.where( - exponent == 0, mantissa * 0.5, np.ldexp(1.0 + mantissa * 0.5, exponent - 1) - ) - return np.copysign(magnitude, np.where(codes & 8, -1.0, 1.0)) - - -def pack_fp4_codes(codes: np.ndarray) -> np.ndarray: - """Pack E2M1 codes, not arbitrary floating values; no quantization policy.""" - if ( - codes.dtype != np.uint8 - or codes.ndim != 2 - or codes.shape[1] % 2 - or np.any(codes > 15) - ): - raise ValueError("expected rank-2 uint8 FP4 codes in [0, 15] and even K") - return codes[:, 0::2] | (codes[:, 1::2] << 4) - - -def reference_result( - a: np.ndarray, - b: np.ndarray, - a_scale: np.ndarray, - b_scale: np.ndarray, - block_k: int, - *, - native: bool, - dtype_a: str | None = None, - dtype_b: str | None = None, - dtype_c: str = "bf16", -) -> np.ndarray: - """Expand scales onto logical A/B elements, multiply, then round the output. - - 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 output-type rounding provide an independent oracle. - FP6 all-code fixtures isolate one K element, avoiding accumulation error. - Input decoding preserves signed zero, but this reference and its numerical - comparison do not establish output zero-sign or accumulation sign semantics. - FP4 fixtures cover all E2M1 values (magnitude <=6), scales 2**[-2,1], - and K<=256: absolute partial sums are below 2**22 units of 2**-6, - so FP32 accumulation is also exact before the final output-type rounding. - """ - output_dtype = _output_dtype(dtype_c) - - 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) - - def matrix_values(data, dtype): - if dtype in ("fp6", "fp6e2m3", "bf6", "fp6e3m2"): - return decode_fp6(data, dtype) - if dtype in ("fp4", "fp4e2m1") or (dtype is None and data.dtype == np.uint8): - return decode_fp4(data) - return data.astype(np.float64) - - a_values = matrix_values(a, dtype_a) - b_values = matrix_values(b, dtype_b) - scaled_a = a_values * np.repeat(sa, block_k, axis=1) - scaled_b = b_values * np.repeat(sb.T, block_k, axis=1) - ref = scaled_a @ scaled_b.T - if native: - ref = ref.astype(output_dtype) - return ref.astype(np.float32) - - -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 - - 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) - - def operand(dtype, rows): - kind = _canon_lowbit(dtype) - if kind == "fp4": - # Keep E2M1 codes unpacked until group masking is finished. - return rng.integers(0, 16, size=(rows, spec.K), dtype=np.uint8) - if kind in ("fp6", "bf6"): - # Small dyadic values ensure exact FP32 partial sums in these tests. - codes = rng.integers( - 0, 9 if kind == "fp6" else 13, size=(rows, spec.K), dtype=np.uint8 - ) - return codes | (rng.integers(0, 2, size=codes.shape, dtype=np.uint8) << 5) - magnitude = 0.25 if native else 0.5 - return (rng.integers(-4, 5, size=(rows, spec.K)) * magnitude).astype( - lowbit_types[dtype] - ) - - a, b = operand(spec.dtype_a, spec.M), operand(spec.dtype_b, spec.N) - if native: - small = any( - _canon_lowbit(d) in ("fp6", "bf6", "fp4") - for d in (spec.dtype_a, spec.dtype_b) - ) - - sa = rng.integers( - 125, 129 if small else 131, size=(spec.M, groups), dtype=np.uint8 - ) - sb = rng.integers( - 125, 129 if small else 131, size=(groups, spec.N), dtype=np.uint8 - ) - neutral_a = neutral_b = 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_a = neutral_b = 1.0 - if case in ("neutral", "b-only"): - sa.fill(neutral_a) - if case in ("neutral", "a-only"): - sb.fill(neutral_b) - if case.startswith("codes-"): - if ( - _canon_lowbit(spec.dtype_a) not in ("fp6", "bf6") - or _canon_lowbit(spec.dtype_b) not in ("fp6", "bf6") - or min(spec.M, spec.N) < 64 - ): - raise ValueError("codes-N requires FP6/BF6 operands and M/N >= 64") - k = int(case.removeprefix("codes-")) - if not 0 <= k < spec.K: - raise ValueError("codes-N K index outside the matrix") - a.fill(0) - b.fill(0) - a[:, k] = np.arange(spec.M, dtype=np.uint8) % 64 - b[:, k] = np.arange(spec.N, dtype=np.uint8) % 64 - sa.fill(neutral_a) - sb.fill(neutral_b) - elif 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}") - - def pack(data, dtype): - kind = _canon_lowbit(dtype) - if kind == "fp4": - return pack_fp4_codes(data) - return pack_fp6_codes(data) if kind in ("fp6", "bf6") else data - - a, b = pack(a, spec.dtype_a), pack(b, spec.dtype_b) - return a, b, sa, sb - - -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. - - Exact mode requires numerical equality after output-type rounding, with - positive and negative zero treated as equal. Tolerance applies only to the - legacy software-scaled WMMA path. - """ - 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]}" - ) +__all__ = [ + "check_result", + "decode_e8m0", + "decode_fp4", + "decode_fp6", + "main", + "make_case_inputs", + "pack_fp4_codes", + "pack_fp6_codes", + "reference_result", + "run_cases", +] def _u8_buffer(array: np.ndarray): diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/FP4_SCALE.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/FP4_SCALE.md new file mode 100644 index 000000000000..b6498ba49f86 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/FP4_SCALE.md @@ -0,0 +1,13 @@ +# gfx950 FP4 and E8M0 storage + +FP4 E2M1 (`fp4`) stores two consecutive logical K values per byte, first code +in the low nibble. Each row occupies K/2 bytes. The 16 codes cover sign, two +exponent bits and one mantissa bit; finite magnitudes extend through 6. +Packing accepts codes rather than imposing a quantization policy. + +The native atom consumes 32 values per lane in four live i32 words, padded +with four zero words for the LLVM ABI. E8M0 scale tensors are separate byte +arrays, A_scale[M,K/32] and B_scale[K/32,N]. Every raw byte scales 32 logical +values; it is zero-extended into a separate A/B i32 VGPR operand. The examples +use 1/4, 1/2, 2 and 4 and an independent decoded oracle. See [the common +contract and commands](README.md). diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/FP6.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/FP6.md new file mode 100644 index 000000000000..82dfa3903981 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/FP6.md @@ -0,0 +1,13 @@ +# gfx950 FP6 storage + +FP6 E2M3 (`fp6`) and E3M2 (`bf6`) use six-bit codes. Four consecutive logical +K values occupy three little-endian bytes, with the first code in the low six +bits. Each row occupies K*6/8 bytes. The shared packer writes codes; it does +not define floating-point quantization. The independent reference decodes bits +using the exponent/mantissa formula, including all 64 finite codes and signed +zeros. Focused GPU tests isolate each end of every K32 group and exercise all +64x64 code products for both homogeneous and mixed FP6/BF6 pairs. + +The native atom consumes 32 values per lane in six live i32 words, padded with +two zero words for the LLVM ABI. One raw E8M0 byte belongs to each row/column +and K32 group. See [the common contract and commands](README.md). diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/README.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/README.md new file mode 100644 index 000000000000..0b9c9d6f3154 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/README.md @@ -0,0 +1,73 @@ +# gfx950 scaled GEMM examples + +These untuned correctness examples use the native 32x32x64 scaled MFMA and the +reusable `instances/gfx950/block_scaled_gemm.py` builder. They make no +performance claim. + +These entrypoints construct `BlockScaledGemmSpec` and call the builder directly. +The builder selects a `mfma_scaled` catalog atom and emits `IRBuilder.mma`; +there is currently no GEMM dispatcher registration for this family. The +existing gfx950 FP8 fused-MoE dispatcher uses the separate 16x16x128 FP8 +atom, with its FP32 block scales applied outside the MFMA. See the +[family contract](../../../../../dsl_docs/instances/gfx950_block_scaled_gemm.md) +and [selection/emission comparison](../../../../../dsl_docs/reference/mma_metadata.md#gfx950-scaled-mfma-paths). + +Run from an installed rocKE package or with the platform `python` directory on +`PYTHONPATH`. A gfx950 device, matching LLVM 20, 22 or 23 ROCm compiler/runtime +libraries, NumPy and +`ml_dtypes` are required. The commands below use LLVM 23; select the flavor +matching your compiler. + +| Example | A/B matrix formats | +| --- | --- | +| [mxfp8_gemm](mxfp8_gemm.py) | Homogeneous FP8 E4M3 or BF8 E5M2; both by default | +| [mxfp6_gemm](mxfp6_gemm.py) | Homogeneous FP6 E2M3 or BF6 E3M2; both by default | +| [mxfp4_gemm](mxfp4_gemm.py) | Homogeneous FP4 E2M1 | +| [mixed_scaled_gemm](mixed_scaled_gemm.py) | Any of the 20 ordered unequal pairs; defaults to FP8 x FP4 | + +```sh +ROCKE_LLVM_FLAVOR=llvm23 python -m rocke.examples.gfx950.gemm.mxfp8_gemm --case all +ROCKE_LLVM_FLAVOR=llvm23 python -m rocke.examples.gfx950.gemm.mxfp6_gemm --case all +ROCKE_LLVM_FLAVOR=llvm23 python -m rocke.examples.gfx950.gemm.mxfp4_gemm --case all +ROCKE_LLVM_FLAVOR=llvm23 python -m rocke.examples.gfx950.gemm.mixed_scaled_gemm --dtype-a fp6 --dtype-b fp4 --case all --compile-route hip --output-dtype fp16 +``` + +Short aliases `fp8`, `bf8`, `fp6`, `bf6`, `fp4` normalize to `fp8e4m3`, +`bf8e5m2`, `fp6e2m3`, `fp6e3m2`, `fp4e2m1`. Both matrices are densely packed +row-major, A[M,K] and B[N,K], yielding C[M,N]. [FP6 packing](FP6.md) and +[FP4 packing](FP4_SCALE.md) specify byte storage. A and B may use different +formats and strides. C uses BF16 (default) or FP16, with FP32 accumulation. + +Scale tensors contain raw finite E8M0 bytes: A_scale[M,K/32] and +B_scale[K/32,N]. Each scale covers 32 logical K values. The consumer loads +one byte per lane, zero-extends it into an i32 VGPR, and supplies separate A/B +scale operands. Code 127 is unity; codes 125, 126, 128 and 129 represent +1/4, 1/2, 2 and 4. The default `mixed` case uses independent non-unity patterns +on both operands, varying rows, columns, K32 groups, K64 steps and output tiles. +`--case all` adds unity and one-sided controls, signed and positive mantissas, +and every isolated K32 group. A NaN output sentinel detects missing stores. + +M and N must be positive multiples of 32; the builder accepts positive K +multiples of 64. The bounded example fixtures accept K64, K128 and K256 and +default to M32/N64/K256. Native comparisons require exact numerical equality +after output-type rounding, with signed zeros treated as equal. Fixture bounds +ensure exact FP32 partial sums. These checks do not establish arbitrary-input +rounding, zero-sign semantics, extreme E8M0 exponents or NaN behavior. + +The catalog exposes all 25 ordered pairs as `mfma_scaled` atoms. Each lane +owns 32 logical A/B values; the LLVM ABI uses padded eight-i32 carriers, with +8/6/4 physical VGPR words for F8/F6/F4. F8/BF8 lane payloads interleave 16-value +chunks between K32 halves; F6/BF6/F4 payloads use contiguous 32-value chunks. +C owns 16 FP32 values per lane in a wave64. Scale ownership is one K32 group +per half-wave. The Python and C++ builders and lowerers mirror these contracts. +LLVM 20, 22 and 23 support this family on gfx950 with the same overloaded +nine-argument intrinsic ABI. Select the flavor matching the compiler when +emitting offline IR; other targets reject this family explicitly. + +The public [CDNA4 machine-readable ISA](https://gpuopen.com/machine-readable-isa/) +provides the `V_MFMA_SCALE_F32_32X32X64_F8F6F4` encoding (ENC_VOP3PX2, opcode +46) and separate SCALE_SRC0/SCALE_SRC1 operands. LLVM's +[AMDGPU intrinsic definitions](https://github.com/llvm/llvm-project/blob/0bace1908348b840e6aa1b4b6e12151dae208158/llvm/include/llvm/IR/IntrinsicsAMDGPU.td) +provide the overloaded nine-argument ABI. Compiler/disassembly tests check the +instruction, format selectors, physical tuples and VGPR scales; GPU numeric +tests check the loaded-scale interpretation and layout independently. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/__init__.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/__init__.py new file mode 100644 index 000000000000..1df485718439 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/_scaled_gemm_example.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/_scaled_gemm_example.py new file mode 100644 index 000000000000..1a41eabfc210 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/_scaled_gemm_example.py @@ -0,0 +1,46 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Shared arguments and execution for focused scaled GEMM examples.""" + +from __future__ import annotations + +import argparse + +from .block_scaled_gemm_verify import run_cases + + +def argument_parser(description: str) -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=description) + parser.add_argument("--m", type=int, default=32) + parser.add_argument("--n", type=int, default=64) + parser.add_argument("--k", type=int, choices=(64, 128, 256), default=256) + parser.add_argument( + "--output-dtype", + choices=("bf16", "fp16"), + default="bf16", + help="output storage and reference rounding (default: bf16)", + ) + parser.add_argument("--compile-route", choices=("comgr", "hip"), default="comgr") + parser.add_argument( + "--case", + default="mixed", + help="neutral, a-only, b-only, mixed, positive, group-N, or all", + ) + return parser + + +def verify(spec, args: argparse.Namespace) -> int: + """Prepare bounded inputs, compile, launch, and compare with the reference. + + The shared verifier owns packing, independent decoding, HIP allocations, + and result checks. Family examples select the public input contract. + """ + cases = ( + ("neutral", "a-only", "b-only", "mixed", "positive") + + tuple(f"group-{g}" for g in range(spec.K // spec.block_k)) + if args.case == "all" + else (args.case,) + ) + count = run_cases(spec, cases, compile_route=args.compile_route) + print(f"PASS: verified {count} cases", flush=True) + return 0 diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/block_scaled_gemm_verify.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/block_scaled_gemm_verify.py new file mode 100644 index 000000000000..a9a24b21879a --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/block_scaled_gemm_verify.py @@ -0,0 +1,177 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""gfx950 scaled MFMA numerical verification with independent non-unity scales.""" + +from __future__ import annotations + +import ctypes +import struct + +import numpy as np + +from ....helpers import compile_kernel +from ....helpers.compile import compile_kernel_via_hipcc +from ....instances.gfx950.block_scaled_gemm import ( + BlockScaledGemmSpec, + block_scaled_gemm_grid, + build_block_scaled_gemm, + is_valid_spec, +) +from ....runtime.hip_module import Runtime, get_device_arch +from ...common.scaled_gemm_reference import ( + _output_dtype, + check_result, + make_case_inputs, + pack_fp4_codes, + pack_fp6_codes, + reference_result, +) + + +def make_inputs(spec, case): + """Use asymmetric K32 scales from 1/4, 1/2, 2, 4. + + Non-FP4 mantissas have magnitude <=1 and quantum >=1/16; FP4 + mantissas have magnitude <=6 and quantum 1/2. At K<=256, scaled + product quantum counts are bounded by 9*2**20 for FP4 x FP4, + 3*2**22 for mixed FP4 pairs, and 2**24 otherwise. Every integer + count through 2**24 is exactly representable in FP32, so each + partial sum is exact before output rounding. + """ + a, b, sa, sb = make_case_inputs(spec, "mixed" if case == "positive" else case) + groups = spec.K // 32 + values = np.array([125, 126, 128, 129], dtype=np.uint8) + g = np.arange(groups)[None, :] + rows = np.arange(spec.M)[:, None] + cols = np.arange(spec.N)[None, :] + # Include tile index and step index without repeating the same K64 halves. + sa[:] = values[(rows + rows // 32 + g + g // 2) % 4] + sb[:] = values[(2 * cols + cols // 32 + 3 * g.T + g.T // 2 + 1) % 4] + if case in ("neutral", "b-only") or case.startswith("codes-"): + sa.fill(127) + if case in ("neutral", "a-only") or case.startswith("codes-"): + sb.fill(127) + if case == "positive": + # Positive dyadic mantissas expose lost scaling without cancellation. + def positive(dtype, count): + codes = (np.arange(count)[:, None] + np.arange(spec.K)[None, :]) % 3 + 1 + if dtype == "fp4e2m1": + return pack_fp4_codes(codes.astype(np.uint8)) + if dtype in ("fp6e2m3", "fp6e3m2"): + return pack_fp6_codes(codes.astype(np.uint8)) + import ml_dtypes + + return (codes * 0.25).astype( + ml_dtypes.float8_e4m3fn if dtype == "fp8e4m3" else ml_dtypes.float8_e5m2 + ) + + a = positive(spec.dtype_a, spec.M) + b = positive(spec.dtype_b, spec.N) + return a, b, sa, sb + + +def assert_scale_coverage(spec, case, sa, sb): + if sa.shape != (spec.M, spec.K // 32) or sb.shape != (spec.K // 32, spec.N): + raise AssertionError("incorrect scale fixture shape") + if case == "neutral" or case.startswith("codes-"): + return + for active, scales in ((case != "b-only", sa), (case != "a-only", sb)): + if active and not ({125, 126, 128, 129} <= set(np.unique(scales))): + raise AssertionError( + "non-trivial scale fixture must include scales below/above unity" + ) + + +def _u8_buffer(array: np.ndarray): + array = np.ascontiguousarray(array) + return (ctypes.c_uint8 * int(array.nbytes)).from_buffer_copy(array) + + +def _launch(rt, fn, spec, inputs): + output_dtype = _output_dtype(spec.dtype_c) + + # A NaN sentinel makes missing output stores fail, including expected zeros. + out = np.full((spec.M, spec.N), np.nan, dtype=output_dtype) + allocations = [] + try: + for array in (*inputs, out): + ptr = rt.alloc(array.nbytes) + allocations.append(ptr) + rt.memcpy_h2d(ptr, _u8_buffer(array), array.nbytes) + packed = struct.pack(" int: + """Compile once and launch each case; return the verified case count.""" + ok, reason = is_valid_spec(spec, arch="gfx950") + if not ok: + raise ValueError(reason) + native = spec.resolved_matrix_path() != "wmma" + if native and spec.K not in (64, 128, 256): + raise ValueError("exact native fixtures require K=64, 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 != "gfx950": + raise RuntimeError(f"visible HIP device 0 must be gfx950, 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"[{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_inputs(spec, case) + assert_scale_coverage(spec, case, inputs[2], inputs[3]) + expected = reference_result( + *inputs, + spec.block_k, + native=native, + dtype_a=spec.dtype_a, + dtype_b=spec.dtype_b, + dtype_c=spec.dtype_c, + ) + label = ( + f"{spec.resolved_matrix_path()}/{spec.dtype_a}x{spec.dtype_b}/{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 nonunity_a={np.count_nonzero(inputs[2] != 127)} nonunity_b={np.count_nonzero(inputs[3] != 127)}", + flush=True, + ) + finally: + module.unload() + return len(cases) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/mixed_scaled_gemm.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/mixed_scaled_gemm.py new file mode 100644 index 000000000000..ca1cd88a7c42 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/mixed_scaled_gemm.py @@ -0,0 +1,48 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""GEMM with different A/B matrix formats and E8M0 scales on gfx950.""" + +from __future__ import annotations + +from ....core.dtypes import normalize_dtype +from ....instances.gfx950.block_scaled_gemm import LOWBIT_DTYPES, BlockScaledGemmSpec +from ._scaled_gemm_example import argument_parser, verify + +FORMATS = tuple(sorted(LOWBIT_DTYPES)) + + +def make_spec(args) -> BlockScaledGemmSpec: + if normalize_dtype(args.dtype_a) == normalize_dtype(args.dtype_b): + raise ValueError( + "choose different A/B formats or use a homogeneous mxfp example" + ) + return BlockScaledGemmSpec( + name="mixed_scaled_gemm", + M=args.m, + N=args.n, + K=args.k, + dtype_a=args.dtype_a, + dtype_b=args.dtype_b, + dtype_c=args.output_dtype, + block_k=32, + ) + + +def main(argv: list[str] | None = None) -> int: + parser = argument_parser(__doc__) + parser.add_argument( + "--dtype-a", type=normalize_dtype, choices=FORMATS, default="fp8e4m3" + ) + parser.add_argument( + "--dtype-b", type=normalize_dtype, choices=FORMATS, default="fp4e2m1" + ) + args = parser.parse_args(argv) + try: + spec = make_spec(args) + except ValueError as exc: + parser.error(str(exc)) + return verify(spec, args) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/mxfp4_gemm.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/mxfp4_gemm.py new file mode 100644 index 000000000000..0f12e11755c1 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/mxfp4_gemm.py @@ -0,0 +1,30 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Packed FP4 E2M1 x E2M1 GEMM with E8M0 block scales on gfx950.""" + +from __future__ import annotations + +from ....instances.gfx950.block_scaled_gemm import BlockScaledGemmSpec +from ._scaled_gemm_example import argument_parser, verify + + +def make_spec(args) -> BlockScaledGemmSpec: + return BlockScaledGemmSpec( + name="mxfp4_gemm", + M=args.m, + N=args.n, + K=args.k, + dtype_a="fp4", + dtype_b="fp4", + dtype_c=args.output_dtype, + block_k=32, + ) + + +def main(argv: list[str] | None = None) -> int: + args = argument_parser(__doc__).parse_args(argv) + return verify(make_spec(args), args) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/mxfp6_gemm.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/mxfp6_gemm.py new file mode 100644 index 000000000000..a889f2656e9e --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/mxfp6_gemm.py @@ -0,0 +1,42 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Homogeneous packed FP6 E2M3 or E3M2 GEMM with packed E8M0 block scales on gfx950.""" + +from __future__ import annotations + +from ....core.arch.target import normalize_dtype +from ....instances.gfx950.block_scaled_gemm import BlockScaledGemmSpec +from ._scaled_gemm_example import argument_parser, verify + + +def make_spec(args) -> BlockScaledGemmSpec: + return BlockScaledGemmSpec( + name="mxfp6_gemm", + M=args.m, + N=args.n, + K=args.k, + dtype_a=args.dtype, + dtype_b=args.dtype, + dtype_c=args.output_dtype, + block_k=32, + ) + + +def main(argv: list[str] | None = None) -> int: + parser = argument_parser(__doc__) + parser.add_argument( + "--dtype", + type=normalize_dtype, + choices=("fp6e2m3", "fp6e3m2", "both"), + default="both", + ) + args = parser.parse_args(argv) + dtypes = ("fp6e2m3", "fp6e3m2") if args.dtype == "both" else (args.dtype,) + for dtype in dtypes: + args.dtype = dtype + verify(make_spec(args), args) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/mxfp8_gemm.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/mxfp8_gemm.py new file mode 100644 index 000000000000..b775793c81dd --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx950/gemm/mxfp8_gemm.py @@ -0,0 +1,42 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Homogeneous FP8 E4M3 or BF8 E5M2 GEMM with E8M0 block scales on gfx950.""" + +from __future__ import annotations + +from ....core.dtypes import normalize_dtype +from ....instances.gfx950.block_scaled_gemm import BlockScaledGemmSpec +from ._scaled_gemm_example import argument_parser, verify + + +def make_spec(args) -> BlockScaledGemmSpec: + return BlockScaledGemmSpec( + name="mxfp8_gemm", + M=args.m, + N=args.n, + K=args.k, + dtype_a=args.dtype, + dtype_b=args.dtype, + dtype_c=args.output_dtype, + block_k=32, + ) + + +def main(argv: list[str] | None = None) -> int: + parser = argument_parser(__doc__) + parser.add_argument( + "--dtype", + type=normalize_dtype, + choices=("fp8e4m3", "bf8e5m2", "both"), + default="both", + ) + args = parser.parse_args(argv) + dtypes = ("fp8e4m3", "bf8e5m2") if args.dtype == "both" else (args.dtype,) + for dtype in dtypes: + args.dtype = dtype + verify(make_spec(args), args) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/atoms.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/atoms.py index c3acdaeeb014..6d830f387791 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/atoms.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/atoms.py @@ -582,13 +582,15 @@ def emit_scaled( a_scale: Value, b_scale: Value, ) -> Value: - """Issue a scaled MX MFMA (P15). - - Currently routes every shape through the - ``mfma_scale_f32_16x16x128_f8f6f4`` intrinsic; the - per-output-row scale broadcast happens inside the - instruction. Future expansion: dedicated scaled intrinsics - for the 32x32 hero shape if AMDGPU exposes them. + """Issue the existing 16x16x128 scaled MFMA wrapper (P15). + + This always calls ``mfma_scale_f32_16x16x128_f8f6f4``; + it does not select an instruction from this atom's shape or dtype. + For gfx950 32x32x64 with independent A/B formats and K32 E8M0 + scales, resolve ``family="mfma_scaled"`` through + ``ArchTarget.mma.op_for_shape`` and pass that ``MmaOp`` and both + scale operands to ``IRBuilder.mma``. See the field manual's + ``reference/mma_metadata.md`` for the two paths. """ return b.mfma_scale_f32_16x16x128_f8f6f4(a, bb, c, a_scale, b_scale) 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 27cd346b2216..e4df6de9f920 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 @@ -8,7 +8,7 @@ matrix atom / unsupported path). | Arch | µarch | Matrix engine | Wave | f16 fp32-acc atoms | |---|---|---|---|---| | gfx942 | CDNA3 | MFMA | 64 | 16x16x16, 32x32x8 | -| gfx950 | CDNA4 | MFMA | 64 | 16x16x16, 16x16x32, 32x32x8, 32x32x16 | +| gfx950 | CDNA4 | MFMA | 64 | 16x16x16, 16x16x32, 32x32x8, 32x32x16, scaled 32x32x64 | | gfx1151 | RDNA3.5 | WMMA | 32 | 16x16x16 (`wmma_f32_16x16x16_{f16,bf16}`) | Matrix-core instances below were exercised with a portable f16 **16x16x16** @@ -46,6 +46,20 @@ compute, `vec=2`). f64 elementwise is limited to the native f64 ops `copy`, --- +## gfx950 native scaled GEMM + +`gfx950/block_scaled_gemm` supports all 25 ordered homogeneous/mixed FP8 E4M3, +BF8 E5M2, FP6 E2M3, BF6 E3M2 and FP4 E2M1 pairs with native 32x32x64 scaled +MFMA. E8M0 scales cover K32 groups; accumulation is FP32 and output BF16/FP16. +M/N must be multiples of 32 and K a multiple of 64, within signed i32 indexing. +LLVM 20/22/23 lowering is supported on gfx950. Use the spec-driven builder +directly. The [four focused examples](../examples/gfx950/gemm/README.md) +cover F8, F6, F4 and mixed inputs with independent non-unity scales. These are +untuned correctness examples and make no performance claim. Python/C++ +IR and LLVM identity, compiler ISA checks and bounded hardware numeric suites +cover this consumer; arbitrary-input rounding and special scales remain outside +the finite verification fixtures. + ## gfx1250 native scaled GEMM `block_scaled_gemm` supports homogeneous and mixed FP8 E4M3, BF8 E5M2, FP6 E2M3, FP6 E3M2, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/gfx950/block_scaled_gemm.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/gfx950/block_scaled_gemm.py new file mode 100644 index 000000000000..446269a0768f --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/gfx950/block_scaled_gemm.py @@ -0,0 +1,189 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Native gfx950 RCR GEMM with independent packed matrices and K32 E8M0 scales.""" + +from dataclasses import dataclass + +from ...core.arch import ArchTarget +from ...core.arch.mfma_scale import MATRIX_FORMATS, matrix_layout +from ...core.dtypes import dtype_info, normalize_dtype +from ...core.ir import I8, I32, IRBuilder, PtrType +from ...helpers.mma_io import load_matrix_fragment, storage_ir_type +from ...helpers.spec import SignatureBuilder + +LOWBIT_DTYPES = frozenset(MATRIX_FORMATS) + + +@dataclass(frozen=True) +class BlockScaledGemmSpec: + name: str + M: int + N: int + K: int + dtype_a: str = "fp8e4m3" + dtype_b: str = "fp8e4m3" + dtype_c: str = "bf16" + block_k: int = 32 + + def __post_init__(self): + for field in ("dtype_a", "dtype_b", "dtype_c"): + object.__setattr__(self, field, normalize_dtype(getattr(self, field))) + + @property + def block_size(self): + return 64 + + def kernel_name(self): + return f"{self.name}_mfma_scaled_{self.dtype_a}_{self.dtype_b}_{self.dtype_c}_M{self.M}N{self.N}K{self.K}" + + def resolved_matrix_path(self): + return "mfma_scale" + + +def is_valid_spec(spec, arch="gfx950"): + if arch != "gfx950": + return False, "scaled MFMA GEMM requires gfx950" + if spec.dtype_a not in MATRIX_FORMATS or spec.dtype_b not in MATRIX_FORMATS: + return False, "A/B must be FP8, BF8, FP6, BF6 or FP4" + if spec.dtype_c not in ("bf16", "fp16"): + return False, "output must be bf16 or fp16" + if spec.block_k != 32: + return False, "E8M0 scale groups must have block_k=32" + if ( + spec.M <= 0 + or spec.N <= 0 + or spec.K <= 0 + or spec.M % 32 + or spec.N % 32 + or spec.K % 64 + ): + return False, "positive M/N multiples of 32 and K multiples of 64 required" + # All tensor offsets use signed i32, including packed row strides. + if max(spec.M * spec.K, spec.N * spec.K, spec.M * spec.N) > 0x7FFFFFFF: + return False, "tensor indexing exceeds i32 range" + return True, "ok" + + +def block_scaled_gemm_grid(spec): + return spec.N // 32, spec.M // 32, 1 + + +def block_scaled_gemm_signature(spec): + signature = SignatureBuilder() + for name, dtype in ( + ("A", spec.dtype_a), + ("B", spec.dtype_b), + ("A_scale", "e8m0"), + ("B_scale", "e8m0"), + ("C", spec.dtype_c), + ): + signature.ptr(name, storage_ir_type(dtype).name) + for name in ("M", "N", "K"): + signature.scalar(name, "i32") + return signature.build() + + +def build_block_scaled_gemm(spec, *, arch="gfx950"): + """One wave computes a 32x32 tile. A[M,K], B[N,K] are densely packed. + + A_scale[M,K/32] and B_scale[K/32,N] contain raw E8M0 bytes. Each + lane's selected scale is zero-extended into the low byte of an i32 VGPR. + Matrix payloads use eight i32 ABI words, with zero padding for FP6/FP4. + """ + ok, reason = is_valid_spec(spec, arch) + if not ok: + raise ValueError(reason) + atom = ArchTarget.from_gfx(arch).mma.op_for_shape( + family="mfma_scaled", + a_dtype=spec.dtype_a, + b_dtype=spec.dtype_b, + c_dtype="fp32", + scales=("e8m0", "e8m0", 32), + m=32, + n=32, + k=64, + ) + ir = IRBuilder(spec.kernel_name()) + ir.kernel.attrs["max_workgroup_size"] = spec.block_size + A = ir.param( + "A", + PtrType(storage_ir_type(spec.dtype_a), "global"), + noalias=True, + readonly=True, + align=16, + ) + B = ir.param( + "B", + PtrType(storage_ir_type(spec.dtype_b), "global"), + noalias=True, + readonly=True, + align=16, + ) + AS = ir.param( + "A_scale", PtrType(I8, "global"), noalias=True, readonly=True, align=1 + ) + BS = ir.param( + "B_scale", PtrType(I8, "global"), noalias=True, readonly=True, align=1 + ) + C = ir.param( + "C", + PtrType(storage_ir_type(spec.dtype_c), "global"), + noalias=True, + writeonly=True, + align=16, + ) + for name in ("M", "N", "K"): + ir.param(name, I32) + c32 = ir.const_i32(32) + cN = ir.const_i32(spec.N) + lane = ir.mod(ir.thread_id_x(), ir.const_i32(64)) + row = ir.mod(lane, c32) + half = ir.div(lane, c32) + m0 = ir.mul(ir.block_id_y(), c32) + n0 = ir.mul(ir.block_id_x(), c32) + ar = ir.add(m0, row) + br = ir.add(n0, row) + abase = ir.mul( + ar, ir.const_i32(spec.K * dtype_info(spec.dtype_a).encoded_bits // 8) + ) + bbase = ir.mul( + br, ir.const_i32(spec.K * dtype_info(spec.dtype_b).encoded_bits // 8) + ) + acc = ir.zero_vec_f32(16) + for step in range(spec.K // 64): + a = load_matrix_fragment( + ir, + A, + abase, + half, + step * 64, + dtype=spec.dtype_a, + layout=matrix_layout(spec.dtype_a), + alignment_bytes=16, + ) + b = load_matrix_fragment( + ir, + B, + bbase, + half, + step * 64, + dtype=spec.dtype_b, + layout=matrix_layout(spec.dtype_b), + alignment_bytes=16, + ) + group = ir.add(ir.const_i32(step * 2), half) + ai = ir.add(ir.mul(ar, ir.const_i32(spec.K // 32)), group) + bi = ir.add(ir.mul(group, cN), br) + sa = ir.zext(ir.global_load(AS, ai, I8, align=1), I32) + sb = ir.zext(ir.global_load(BS, bi, I8, align=1), I32) + acc = ir.mma(atom, a, b, acc, sa, sb) + for slot in range(16): + r, c = atom.c_layout().coord(ir, lane, slot) + index = ir.add(ir.mul(ir.add(m0, r), cN), ir.add(n0, c)) + ir.global_store( + C, + index, + ir.cast_f32_to(ir.vec_extract(acc, slot), storage_ir_type(spec.dtype_c)), + align=2, + ) + return ir.kernel diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/gfx950_scaled_mfma_llvm23.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/gfx950_scaled_mfma_llvm23.json new file mode 100644 index 000000000000..1763623ef9dd --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/gfx950_scaled_mfma_llvm23.json @@ -0,0 +1,52 @@ +{ + "0": "a550fa8769391a2bf246ee73ad70769826d1f902db0a49bb60a80e7c0efda41e", + "1": "e5bd9254d6c97778113568e4ca7862ab0c25db1a1ee902586ce222bc61cd9339", + "2": "87efa34085c39923773b9021e293a027814ed0a329e38f0ab95d4b92b209f405", + "3": "49d8d97eb1cf0fe21dec104f571742cffab61adda6296b71bd8b326663fdf71a", + "4": "f64530bd7204524f13837c2ae2d243e5dfe8a316908cb9d352380de6de0ba877", + "5": "862f81d0416d03ef409474cbf78e26ba39b6ec9897925c6487d8efc4952d3904", + "6": "ba39d068b0d7b26d1f2dc119993017623d5c88d8a5bc28ec7231986484677bff", + "7": "00891085050c4fd17245108dc7589f1ee7a20235034bd9eacfa08e11888e6fca", + "8": "b399ff71877f53dc55ca31979bb00134ffa66e310419760f2402987121c58961", + "9": "9009f97e6414919acb77196a33bba4833a3dfbc7693df56682b545a2dbe3ee9b", + "10": "0a76bd80f669fc9edeec6a312d774f2c227a05e79a964cb5b03d01bb85473c96", + "11": "09c7fa2112758c7fb247396e3f8513714bc5a7efe11914e4f5ca263bf2da4130", + "12": "a9bf9980ea4d2e780a2a01d026457da88b229bb4a273c125e4c40c2189567d5e", + "13": "a1fe69a540bcb72690d58219e2cf580556a76fd3849d5da5ad5336dc67dc3285", + "14": "eee027dbe6453f3ee37689e40605ac66715a8350a366011b60a1a1f48e7cb67e", + "15": "b62d2d18201c42db4265c4d60a0ec9d6afb4c14156c7ee81220616509d9088cd", + "16": "9f39d1528b62af90490b6301a0ea9874863cc2298810bbec3b164353369e7fce", + "17": "9b0403d12dd32eed42f03467b85c0218b499115871f6df362aade7ab17c3eb82", + "18": "db7ac1683fe70a788fc4cac844f4f83c022276d6e0fbf4c6e4c26c5ce54ffe1f", + "19": "54bbb67595b52c2dbc7494ab69b292b07687b41a80a6d4d1383e0b193bf92667", + "20": "f9079a766f6f4cebcd253aa69cdda68941f263dea58b4fac6ece56338f479963", + "21": "a46e1ce819dfe7fd6f75ffb948f8693440e49c789090dc6f426bbff21155c641", + "22": "1e516434694f058a0084198e1cfc7b09883bb75f2bebc90ba51e0df04f7ad741", + "23": "22d8fa7ec449a566f0e1ee1ff86f36a450507aca0dd6726de2de59497a4087dc", + "24": "9750e5f55dfb5f6379258d041c7d0b9798a90cf825de480913609f5f2d0bd0ba", + "25": "4c6258a0f05b36b5b0b3dab83a531cec670c763504548955e80067453a76490f", + "26": "1640cce297fe7531d51e86830279b3e83355c8fef01951e8fcc4ea062d0b089b", + "27": "9d7cc5ad3cf800551092b4ae057aab09f7c4fb30f8b61ece449c13ed564bc7fb", + "28": "e00bb24ab015c10572164ea0931918396c56052d196a1e85b14c0fea92c67a06", + "29": "5abf9231d9881761f1d0019dc511145fb567831f266a35f1ddcb28830f5ef9c6", + "30": "8959cb7fad595e68aedeeecd60f132f1b81b5fde1dffd70f9f11bd1e5635919d", + "31": "7c872e16ef8e6d86a04bc02ae39202e8cd3a6bd504e92d43fa9f5c422b219e38", + "32": "e7893eecfb88fea46838b9bd644e88c848efe8055de064f41df263b4310d93c3", + "33": "1b22ae154a54a235ffc6a744f4cb1aea9b4e7fb940a44625f0f59464956387a2", + "34": "2021a124ecae7fd1d42f9d7526812ebc27f9c49b50fd5a960f2d56d3f6b1e028", + "35": "670d3e99f6755c2ee991b169b136e1562c5a36b2ce71314a01a3c5f356449e12", + "36": "108c62a2ad7cbad34dcd2a3c1d7a8138711cff1e274608356e28d86fdb7f1677", + "37": "ad250509429eac708941d2db36d6a75f4ec743662e35d6acf0d5f442f83c1e20", + "38": "360ea615c89e02e7b3e104935c75334016f1eec41d1677772c174499d505f8d5", + "39": "cd8da0d57389a538974fa795d7d01399a047bad7456f45713f45eaddc48db2ee", + "40": "505c93a00f3d3541300025f4a5a24c9af25fcf228ce47afe40e8c05a4077401b", + "41": "c9224321a599dfbd3bcbc453ebf3253055322536cccdb66fef3c6c9298eb029e", + "42": "c9a04c5ec2845e7217359a1ba885069e75626bac595d993dd0387ef443085769", + "43": "005eddd91d63e95d9acbc05dd88e3988eeb70ae0a5fb62a8e62ab5c5edabdbe4", + "44": "6b14bcca4ebcb2f41655af64bbb824dc7849915a4cc581776dce012673976bcf", + "45": "a386a1cd825be92bb7fe14230bc12c6f236809639778b96ce5fc7887f62a2809", + "46": "4644db1d674af4914eef62cdc0ae78db078337a1e6023e80302e9f9da84bb7b9", + "47": "fdbc9d9b0e588fa338f4b5b893531e7ce7841f4fdf29324fd2515869060f2268", + "48": "f505df5f18ea8c71232d65ed6872e262cd7afc0527f5cd8332322d9e372786c8", + "49": "2e13655cdc0997cb78da783af029aeba34bf53d9526ecd783f7ec8c8a7306744" +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_gfx950_scaled_mfma.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_gfx950_scaled_mfma.py new file mode 100644 index 000000000000..3013a1156a9e --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_gfx950_scaled_mfma.py @@ -0,0 +1,263 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Complete scaled MFMA contracts, logical layouts and dual-engine emission.""" + +import hashlib +import json +import os +import re +import shutil +import subprocess +import sys +from dataclasses import asdict +from itertools import product +from pathlib import Path + +import pytest +from rocke.core.arch import ArchTarget +from rocke.core.arch.mfma_scale import matrix_layout +from rocke.core.ir import IRBuilder +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 +from rocke.core.verify import verify +from rocke.instances.gfx950.block_scaled_gemm import ( + BlockScaledGemmSpec, + build_block_scaled_gemm, + is_valid_spec, +) + +FORMATS = ("fp8", "bf8", "fp6", "bf6", "fp4") + + +class Arithmetic: + def const_i32(self, x): + return x + + def mod(self, a, b): + return a % b + + def div(self, a, b): + return a // b + + def mul(self, a, b): + return a * b + + def add(self, a, b): + return a + b + + +@pytest.mark.parametrize("a,b", list(product(FORMATS, repeat=2))) +def test_catalog_and_logical_layouts(a, b): + atom = ArchTarget.from_gfx("gfx950").mma.op_for_shape( + family="mfma_scaled", + a_dtype=a, + b_dtype=b, + c_dtype="fp32", + scales=("e8m0", "e8m0", 32), + m=32, + n=32, + k=64, + ) + assert (atom.a_frag_len, atom.b_frag_len, atom.c_frag_len, atom.wave_size) == ( + 32, + 32, + 16, + 64, + ) + cpu = Arithmetic() + for layout, shape in ( + (atom.a_layout(), (32, 64)), + (atom.b_layout(), (64, 32)), + (atom.c_layout(), (32, 32)), + (atom.a_scale_layout(), (32, 2)), + (atom.b_scale_layout(), (2, 32)), + ): + points = [ + layout.coord(cpu, l, s) for l in range(64) for s in range(layout.frag_len) + ] + assert len(points) == len(set(points)) == shape[0] * shape[1] + assert set(points) == set(product(range(shape[0]), range(shape[1]))) + for dtype in (atom.a_dtype, atom.b_dtype): + packed = matrix_layout(dtype) + assert packed.fragment.count == 32 and packed.fragment.carrier_count == 8 + assert packed.fragment.live_carriers in (4, 6, 8) + for target in ("gfx942", "gfx1250", "gfx1201"): + assert ArchTarget.from_gfx(target).mma.by_op_id(atom.op_id) is None + + +@pytest.mark.parametrize("idx", range(50)) +@pytest.mark.parametrize("flavor", ["llvm20", "llvm22", "llvm23"]) +def test_complete_kernel_golden_and_native_identity(idx, flavor, monkeypatch): + + spec = BlockScaledGemmSpec( + "scaled_parity", + 32 if idx < 25 else 64, + 64 if idx < 25 else 96, + 128 if idx < 25 else 256, + FORMATS[(idx % 25) // 5], + FORMATS[idx % 5], + "bf16" if idx < 25 else "fp16", + ) + kernel = build_block_scaled_gemm(spec) + assert not verify(kernel) + llvm = _lower_kernel_to_llvm_python(kernel, arch="gfx950", llvm_flavor=flavor) + calls = [ + line + for line in llvm.splitlines() + if "call <16 x float> @llvm.amdgcn.mfma.scale" in line + ] + assert len(calls) == spec.K // 64 + for line in calls: + assert re.search(rf", i32 {(idx % 25) // 5}, i32 {idx % 5}, i32 0, i32 %", line) + if flavor == "llvm23": + golden = json.loads( + Path(__file__).with_name("gfx950_scaled_mfma_llvm23.json").read_text() + ) + assert hashlib.sha256(llvm.encode()).hexdigest() == golden[str(idx)] + engine = pytest.importorskip("rocke_engine") + assert engine.gfx950_block_scaled_gemm_serialize_ir( + asdict(spec), "gfx950" + ) == serialize(kernel) + assert engine.lower_serialized_ir(serialize(kernel), "gfx950", flavor) == llvm + monkeypatch.setenv("ROCKE_LLVM_FLAVOR", flavor) + assert engine.gfx950_block_scaled_gemm_lower_llvm(asdict(spec), "gfx950") == llvm + assert engine.gfx950_block_scaled_gemm_verify(asdict(spec), "gfx950") == [] + + +@pytest.mark.parametrize("arch", ["gfx942", "gfx1201", "gfx1250"]) +def test_wrong_target_rejection(arch): + k = build_block_scaled_gemm(BlockScaledGemmSpec("reject", 32, 32, 64)) + with pytest.raises(NotImplementedError, match="requires gfx950"): + _lower_kernel_to_llvm_python(k, arch=arch, llvm_flavor="llvm23") + with pytest.raises(NotImplementedError, match="requires gfx950"): + lower_kernel_to_hip(k, arch=arch) + engine = pytest.importorskip("rocke_engine") + with pytest.raises(RuntimeError, match="requires gfx950"): + engine.lower_serialized_ir(serialize(k), arch, "llvm23") + + +def test_native_lowering_preserves_diagnostic(tmp_path): + pytest.importorskip("rocke_engine") + compiler = shutil.which("cc") + if sys.platform != "linux" or not compiler: + pytest.skip("ELF loader fixture requires Linux and a C compiler") + source = tmp_path / "unqueryable.c" + library = tmp_path / "unqueryable.so" + source.write_text("void unqueryable_compiler(void) {}\n") + subprocess.run( + [compiler, "-shared", "-fPIC", str(source), "-o", str(library)], + check=True, + capture_output=True, + text=True, + ) + env = dict(os.environ) + env.pop("ROCKE_LLVM_FLAVOR", None) + env["ROCKE_COMGR_LIB"] = str(library) + env["PYTHONPATH"] = os.pathsep.join(sys.path) + # A loadable, unqueryable fixture prevents fallback to an installed COMGR. + # A fresh process also avoids a previously cached compiler masking failure. + child = subprocess.run( + [ + sys.executable, + "-c", + "import rocke_engine; " + "rocke_engine.gfx950_block_scaled_gemm_lower_llvm(" + "{'name': 'diagnostic', 'M': 32, 'N': 32, 'K': 64}, 'gfx950')", + ], + env=env, + capture_output=True, + text=True, + timeout=30, + ) + assert child.returncode != 0 + assert "Cannot determine LLVM flavor" in child.stderr + assert "unknown error" not in child.stderr + + +@pytest.mark.parametrize("name_bytes", [255, 256, 512]) +def test_long_native_kernel_names(name_bytes): + engine = pytest.importorskip("rocke_engine") + empty = BlockScaledGemmSpec("", 32, 32, 64) + spec = BlockScaledGemmSpec( + "x" * (name_bytes - len(empty.kernel_name())), 32, 32, 64 + ) + assert len(spec.kernel_name().encode("ascii")) == name_bytes + assert is_valid_spec(spec)[0] + kernel = build_block_scaled_gemm(spec) + native_ir = engine.gfx950_block_scaled_gemm_serialize_ir(asdict(spec), "gfx950") + assert native_ir == serialize(kernel) + assert engine.gfx950_block_scaled_gemm_verify(asdict(spec), "gfx950") == [] + for flavor in ("llvm20", "llvm22", "llvm23"): + assert engine.lower_serialized_ir(native_ir, "gfx950", flavor) == ( + _lower_kernel_to_llvm_python(kernel, arch="gfx950", llvm_flavor=flavor) + ) + + +def test_builder_and_serialized_type_rejection(): + atom = next( + x for x in ArchTarget.from_gfx("gfx950").mma.ops if x.family == "mfma_scaled" + ) + ir = IRBuilder("bad") + a = ir.zero_vec_f32(8) + c = ir.zero_vec_f32(16) + scale = ir.const_i32(127) + with pytest.raises(ValueError, match="scaled MFMA requires"): + ir.mma(atom, a, a, c, scale, scale) + k = build_block_scaled_gemm(BlockScaledGemmSpec("bad", 32, 32, 64)) + op = next(x for x in k.body.ops if x.name == "tile.mma") + op.operands[3] = c + assert any("scaled MFMA requires" in str(d) for d in verify(k)) + with pytest.raises(ValueError, match="scaled MFMA requires"): + _lower_kernel_to_llvm_python(k, arch="gfx950", llvm_flavor="llvm23") + + +@pytest.mark.parametrize( + "M,N,K,block_k", + [ + (16, 32, 64, 32), + (32, 48, 64, 32), + (32, 32, 96, 32), + (32, 32, 64, 16), + (2**30, 32, 64, 32), + ], +) +def test_invalid_shape(M, N, K, block_k): + s = BlockScaledGemmSpec("bad", M, N, K, block_k=block_k) + assert not is_valid_spec(s)[0] + with pytest.raises(ValueError): + build_block_scaled_gemm(s) + + +@pytest.mark.parametrize("mutation", ["payload", "acc", "scale", "arity", "result"]) +def test_serialized_native_contract_rejection(mutation): + from rocke.core.ir import F32, VectorType + + k = build_block_scaled_gemm(BlockScaledGemmSpec("malformed", 32, 32, 64)) + op = next(x for x in k.body.ops if x.name == "tile.mma") + if mutation == "arity": + op.operands.pop() + elif mutation == "result": + op.results[0].type = VectorType(F32, 8) + else: + index = {"payload": 0, "acc": 2, "scale": 3}[mutation] + op.operands[index] = op.operands[3] if mutation != "scale" else op.operands[2] + assert any("scaled MFMA" in str(d) for d in verify(k)) + with pytest.raises(ValueError, match="scaled MFMA"): + _lower_kernel_to_llvm_python(k, arch="gfx950", llvm_flavor="llvm23") + engine = pytest.importorskip("rocke_engine") + with pytest.raises(RuntimeError, match="scaled MFMA"): + engine.lower_serialized_ir(serialize(k), "gfx950", "llvm23") + + +@pytest.mark.parametrize("dtype", ["unknown", "fp32", "FP6 ", " BF8", "fp4"]) +def test_native_spec_format_validation(dtype): + engine = pytest.importorskip("rocke_engine") + spec = BlockScaledGemmSpec("formats", 32, 32, 64, dtype_a=dtype) + if dtype.strip().lower() in ("unknown", "fp32"): + with pytest.raises(RuntimeError, match="invalid gfx950"): + engine.gfx950_block_scaled_gemm_serialize_ir(asdict(spec), "gfx950") + else: + assert engine.gfx950_block_scaled_gemm_serialize_ir( + asdict(spec), "gfx950" + ) == serialize(build_block_scaled_gemm(spec)) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_catalog_lookup.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_catalog_lookup.cpp index 6804d32756bb..81f7593038c5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_catalog_lookup.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_catalog_lookup.cpp @@ -345,6 +345,33 @@ static int test_scale_layouts_and_families() continue; } ++scaled; + if(strcmp(arches[i], "gfx950") == 0) + { + CHECK(strcmp(op->family, "mfma_scaled") == 0); + CHECK(op->m == 32 && op->n == 32 && op->k == 64); + CHECK(op->a_frag_len == 32 && op->b_frag_len == 32 && op->c_frag_len == 16); + CHECK(op->scale_block_k == ROCKE_MMA_SCALE_K32); + for(int source = 0; source < 2; ++source) + { + const auto* map = source ? rocke_mma_op_b_scale_layout(op, NULL) + : rocke_mma_op_a_scale_layout(op, NULL); + CHECK(map && map->frag_len == 1 && map->wave_size == 64); + CHECK(map->role == (source ? ROCKE_MMA_ROLE_B_SCALE : ROCKE_MMA_ROLE_A_SCALE)); + for(int slot : {-1, 0, 1}) + { + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "mfma_scale_layout") == ROCKE_OK); + auto* lane = rocke_b_thread_id_x(&b); + rocke_value_t *x = NULL, *y = NULL; + const bool ok = rocke_layout_map_coord(map, &b, lane, slot, &x, &y); + CHECK(ok == (slot == 0)); + CHECK(slot == 0 ? (x && y && b.status == ROCKE_OK) + : (!x && !y && b.status == ROCKE_ERR_VALUE)); + rocke_ir_builder_free(&b); + } + } + continue; + } CHECK(strcmp(arches[i], "gfx1250") == 0); const int count = op->scale_block_k == ROCKE_MMA_SCALE_K16 ? 8 : 4; CHECK(op->a_scale_frag_len == count && op->b_scale_frag_len == count); @@ -376,7 +403,7 @@ static int test_scale_layouts_and_families() } } } - CHECK(scaled == 50); + CHECK(scaled == 75); return 0; } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/conftest.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/conftest.py index c6cd4d321b8f..ab0f63d407fa 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/conftest.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/conftest.py @@ -11,13 +11,12 @@ import importlib.util import os -from pathlib import Path import subprocess import sys +from pathlib import Path import pytest - _PYROOT = Path(__file__).resolve().parents[2] / "python" _MODULE = "rocke.examples.gfx1250.gemm.block_scaled_gemm_verify" @@ -31,9 +30,12 @@ def gpu_env(): [ sys.executable, "-c", - "from rocke.runtime.hip_module import get_device_arch; " - "print(get_device_arch(0))", + ( + "from rocke.runtime.hip_module import get_device_arch; " + "print(get_device_arch(0))" + ), ], + check=False, env=env, capture_output=True, text=True, @@ -74,6 +76,7 @@ def _run_numeric(gpu_env, dtype, matrix_path, route, m, n, k, case, count): "--case", case, ], + check=False, env=gpu_env, capture_output=True, text=True, @@ -92,3 +95,16 @@ def numeric_case(gpu_env): from functools import partial return partial(_run_numeric, gpu_env) + + +@pytest.fixture +def gfx950_numeric(monkeypatch): + """Load the gfx950 numeric helper independently of pytest import mode.""" + from functools import partial + + spec = importlib.util.spec_from_file_location( + "gfx950_scaled_numeric", Path(__file__).with_name("gfx950_scaled_numeric.py") + ) + helper = importlib.util.module_from_spec(spec) + spec.loader.exec_module(helper) + return partial(helper.run_numeric, monkeypatch) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/test_gfx950_scaled_mfma_emit.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/test_gfx950_scaled_mfma_emit.py new file mode 100644 index 000000000000..d3c313398534 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/test_gfx950_scaled_mfma_emit.py @@ -0,0 +1,44 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""The gfx950 scaled-MFMA emitter follows the differential batch protocol.""" + +import hashlib +from pathlib import Path +import subprocess +import sys + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import run_diff + + +@pytest.mark.parametrize("flavor", ["llvm20", "llvm22", "llvm23"]) +def test_scaled_mfma_batch_reenters_and_terminates(monkeypatch, flavor): + monkeypatch.setenv("ROCKE_LLVM_FLAVOR", flavor) + batch = run_diff._run_python_batch("gfx950_block_scaled_gemm", [0, 49, 0, 50]) + assert batch.error is None + assert len(batch.emissions) == 4 + first, last, repeated, ended = batch.emissions + assert all( + emission.returncode == 0 and emission.has_output + for emission in (first, last, repeated) + ) + assert first.output_sha == repeated.output_sha + assert first.output_sha != last.output_sha + assert run_diff._is_end_emission(ended) + + # Compare the in-process batch with the original single-config CLI path. + process = subprocess.run( + [ + sys.executable, + str(run_diff.PARITY / "gfx950_block_scaled_gemm_emit.py"), + "0", + "ll", + ], + env=run_diff._python_emitter_env(run_diff.PARITY), + capture_output=True, + check=True, + timeout=run_diff.TIMEOUT, + ) + assert hashlib.sha256(process.stdout).hexdigest() == first.output_sha diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/gfx950_scaled_numeric.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/gfx950_scaled_numeric.py new file mode 100644 index 000000000000..b655bccd82c3 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/gfx950_scaled_numeric.py @@ -0,0 +1,124 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Bounded gfx950 compiler and hardware numeric validation helpers.""" + +import os +import re +import shutil +import subprocess +import tempfile +from pathlib import Path + +import pytest + +FORMATS = ("fp8", "bf8", "fp6", "bf6", "fp4") + + +def require_gpu(): + if not Path("/dev/kfd").exists(): + if os.environ.get("ROCKE_REQUIRE_GFX950") == "1": + pytest.fail("gfx950 numeric validation requires /dev/kfd") + pytest.skip("requires an AMD GPU") + from rocke.runtime.hip_module import get_device_arch + + arch = get_device_arch(0) + if arch != "gfx950": + if os.environ.get("ROCKE_REQUIRE_GFX950") == "1": + pytest.fail(f"requires gfx950; found {arch}") + pytest.skip(f"requires gfx950; found {arch}") + + +def check_isa(artifact, spec): + objdump = shutil.which("llvm-objdump") + assert objdump, "llvm-objdump is required for the compiler evidence" + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "kernel.hsaco" + path.write_bytes(artifact.hsaco) + if artifact.hsaco.startswith(b"__CLANG_OFFLOAD_BUNDLE__"): + bundler = shutil.which("clang-offload-bundler") + assert bundler, "clang-offload-bundler is required for HIP ISA evidence" + listing = subprocess.run( + [bundler, "--list", "--type=o", f"--input={path}"], + check=True, + capture_output=True, + text=True, + timeout=60, + ) + targets = [x for x in listing.stdout.splitlines() if "gfx950" in x] + assert len(targets) == 1, listing.stdout + device = Path(tmp) / "device.hsaco" + subprocess.run( + [ + bundler, + "--unbundle", + "--type=o", + f"--targets={targets[0]}", + f"--input={path}", + f"--output={device}", + ], + check=True, + capture_output=True, + text=True, + timeout=60, + ) + path = device + proc = subprocess.run( + [objdump, "-d", str(path)], + check=False, + capture_output=True, + text=True, + timeout=60, + ) + assert proc.returncode == 0, proc.stderr + instructions = [ + line + for line in proc.stdout.splitlines() + if "v_mfma_scale_f32_32x32x64_f8f6f4" in line + ] + assert len(instructions) == spec.K // 64, proc.stdout + from rocke.core.arch.mfma_scale import MATRIX_FORMATS + + fa, fb = MATRIX_FORMATS[spec.dtype_a], MATRIX_FORMATS[spec.dtype_b] + for line in instructions: + # Omitted selectors encode format zero. Matrix VGPR tuples are 8/6/4 words; + # the final two operands must be independent VGPR scale sources. + assert ( + int(re.search(r"cbsz:(\d+)", line).group(1)) == fa + if "cbsz:" in line + else fa == 0 + ) + assert ( + int(re.search(r"blgp:(\d+)", line).group(1)) == fb + if "blgp:" in line + else fb == 0 + ) + tuples = re.findall(r"(?:a|v)\[(\d+):(\d+)\]", line) + assert [int(hi) - int(lo) + 1 for lo, hi in tuples[:3]] == [ + 16, + (8, 8, 6, 6, 4)[fa], + (8, 8, 6, 6, 4)[fb], + ], line + assert re.search(r", v\d+, v\d+(?:\s|$)", line), line + return proc.stdout + + +def run_numeric(monkeypatch, a, b, output, route, M, N, K, cases=None): + require_gpu() + from rocke.examples.gfx950.gemm import block_scaled_gemm_verify as verifier + from rocke.instances.gfx950.block_scaled_gemm import BlockScaledGemmSpec + + spec = BlockScaledGemmSpec("numeric_scaled", M, N, K, a, b, output) + cases = cases or ( + ("neutral", "a-only", "b-only", "mixed", "positive") + + tuple(f"group-{i}" for i in range(K // 32)) + ) + key = "compile_kernel" if route == "comgr" else "compile_kernel_via_hipcc" + original = getattr(verifier, key) + + def compile_checked(kernel, **kwargs): + art = original(kernel, **kwargs) + check_isa(art, spec) + return art + + monkeypatch.setattr(verifier, key, compile_checked) + assert verifier.run_cases(spec, cases, compile_route=route) == len(cases) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx950_block_scaled_gemm_emit.c b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx950_block_scaled_gemm_emit.c new file mode 100644 index 000000000000..f5e08736b482 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx950_block_scaled_gemm_emit.c @@ -0,0 +1,82 @@ +/* Copyright (c) Advanced Micro Devices, Inc., or its affiliates. + * SPDX-License-Identifier: MIT */ +#include "rocke/instance_gfx950_block_scaled_gemm.h" +#include "rocke/ir_serialize.h" +#include "rocke/lower_hip.h" +#include "rocke/lower_llvm.h" +#include "rocke/verify.h" +#include +#include +#include + +int main(int argc, char** argv) +{ + int idx = argc > 1 ? atoi(argv[1]) : -1; + if(idx < 0 || idx >= 50) + { + fprintf(stderr, "unknown config index %d\n", idx); + return 2; + } + const char* formats[] = {"fp8e4m3", "bf8e5m2", "fp6e2m3", "fp6e3m2", "fp4e2m1"}; + const char* mode = argc > 2 ? argv[2] : "ll"; + auto spec = rocke_gfx950_block_scaled_gemm_spec_default(); + spec.name = "scaled_parity"; + spec.M = idx < 25 ? 32 : 64; + spec.N = idx < 25 ? 64 : 96; + spec.K = idx < 25 ? 128 : 256; + spec.dtype_a = formats[(idx % 25) / 5]; + spec.dtype_b = formats[idx % 5]; + spec.dtype_c = idx < 25 ? "bf16" : "fp16"; + rocke_ir_builder_t b; + auto* kernel = rocke_build_gfx950_block_scaled_gemm_new(&b, &spec, "gfx950"); + if(!kernel) + { + fprintf(stderr, "%s\n", b.err); + rocke_ir_builder_free(&b); + return 1; + } + char* text = NULL; + rocke_status_t status = ROCKE_OK; + if(strcmp(mode, "ll") == 0) + status = rocke_lower_kernel_to_llvm(kernel, ROCKE_LLVM_FLAVOR_AUTO, "gfx950", &text); + else if(strcmp(mode, "ir") == 0) + status = rocke_ir_serialize(kernel, &text); + else if(strcmp(mode, "hip") == 0) + { + rocke_strbuf_t buf; + rocke_strbuf_init(&buf, 0); + rocke_lower_hip_opts_t opts = {0}; + opts.arch = "gfx950"; + status = rocke_lower_kernel_to_hip(&b, kernel, &opts, &buf); + if(status == ROCKE_OK) + fputs(buf.data, stdout); + rocke_strbuf_free(&buf); + } + else if(strcmp(mode, "verify") == 0) + { + rocke_diag_t* diagnostics = NULL; + size_t count = 0; + rocke_verify(kernel, &diagnostics, &count); + for(size_t i = 0; i < count; ++i) + { + char* line = rocke_diag_to_string(&diagnostics[i]); + if(line) + { + puts(line); + free(line); + } + } + rocke_diags_free(diagnostics, count); + } + else + status = ROCKE_ERR_VALUE; + if(text) + { + fputs(text, stdout); + free(text); + } + if(status != ROCKE_OK) + fprintf(stderr, "emit failed: %d %s\n", status, b.err); + rocke_ir_builder_free(&b); + return status != ROCKE_OK; +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx950_block_scaled_gemm_emit.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx950_block_scaled_gemm_emit.py new file mode 100644 index 000000000000..def2885822b1 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx950_block_scaled_gemm_emit.py @@ -0,0 +1,31 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +from _emit_common import run_emit +from rocke.instances.gfx950.block_scaled_gemm import ( + BlockScaledGemmSpec, + build_block_scaled_gemm, +) + +FORMATS = ("fp8e4m3", "bf8e5m2", "fp6e2m3", "fp6e3m2", "fp4e2m1") + + +def _spec(idx): + if not 0 <= idx < 50: + raise SystemExit(f"unknown config index {idx}") + return BlockScaledGemmSpec( + "scaled_parity", + 32 if idx < 25 else 64, + 64 if idx < 25 else 96, + 128 if idx < 25 else 256, + FORMATS[(idx % 25) // 5], + FORMATS[idx % 5], + "bf16" if idx < 25 else "fp16", + ) + + +def main(): + return run_emit(_spec, build_block_scaled_gemm) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx950_mixed_scaled_gemm_numeric.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx950_mixed_scaled_gemm_numeric.py new file mode 100644 index 000000000000..57938134cec9 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx950_mixed_scaled_gemm_numeric.py @@ -0,0 +1,28 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""All ordered gfx950 mixed-format pairs with independent K32 scales.""" + +from itertools import product + +import pytest + +FORMATS = ("fp8", "bf8", "fp6", "bf6", "fp4") + + +@pytest.mark.parametrize( + "a,b", [(a, b) for a, b in product(FORMATS, repeat=2) if a != b] +) +@pytest.mark.parametrize("output", ["bf16", "fp16"]) +@pytest.mark.parametrize("route", ["comgr", "hip"]) +@pytest.mark.parametrize("M,N,K", [(32, 64, 128), (64, 96, 256)]) +def test_numeric(gfx950_numeric, a, b, output, route, M, N, K): + gfx950_numeric(a, b, output, route, M, N, K) + + +@pytest.mark.parametrize( + "a,b", [("fp8", "fp8"), ("bf8", "fp4"), ("fp6", "bf6"), ("fp4", "fp4")] +) +@pytest.mark.parametrize("route", ["comgr", "hip"]) +@pytest.mark.parametrize("output", ["bf16", "fp16"]) +def test_single_k64_atom(gfx950_numeric, a, b, output, route): + gfx950_numeric(a, b, output, route, 64, 96, 64) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx950_mxfp4_numeric.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx950_mxfp4_numeric.py new file mode 100644 index 000000000000..dc1cc89da633 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx950_mxfp4_numeric.py @@ -0,0 +1,13 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""gfx950 MXFP4 non-unity scaled MFMA numerical cases.""" + +import pytest + + +@pytest.mark.parametrize("dtype", ["fp4"]) +@pytest.mark.parametrize("output", ["bf16", "fp16"]) +@pytest.mark.parametrize("route", ["comgr", "hip"]) +@pytest.mark.parametrize("M,N,K", [(32, 64, 128), (64, 96, 256)]) +def test_numeric(gfx950_numeric, dtype, output, route, M, N, K): + gfx950_numeric(dtype, dtype, output, route, M, N, K) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx950_mxfp6_numeric.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx950_mxfp6_numeric.py new file mode 100644 index 000000000000..2073f5cd85fc --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx950_mxfp6_numeric.py @@ -0,0 +1,31 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""gfx950 MXFP6 non-unity scaled MFMA numerical cases.""" + +import pytest + + +@pytest.mark.parametrize("dtype", ["fp6", "bf6"]) +@pytest.mark.parametrize("output", ["bf16", "fp16"]) +@pytest.mark.parametrize("route", ["comgr", "hip"]) +@pytest.mark.parametrize("M,N,K", [(32, 64, 128), (64, 96, 256)]) +def test_numeric(gfx950_numeric, dtype, output, route, M, N, K): + gfx950_numeric(dtype, dtype, output, route, M, N, K) + + +@pytest.mark.parametrize( + "a,b", [("fp6", "fp6"), ("bf6", "bf6"), ("fp6", "bf6"), ("bf6", "fp6")] +) +@pytest.mark.parametrize("route", ["comgr", "hip"]) +def test_all_fp6_codes(gfx950_numeric, a, b, route): + positions = tuple(k for start in range(0, 128, 32) for k in (start, start + 31)) + gfx950_numeric( + a, + b, + "bf16", + route, + 64, + 64, + 128, + tuple(f"codes-{k}" for k in positions), + ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx950_mxfp8_numeric.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx950_mxfp8_numeric.py new file mode 100644 index 000000000000..28b61e2a3e55 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx950_mxfp8_numeric.py @@ -0,0 +1,13 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""gfx950 MXFP8 non-unity scaled MFMA numerical cases.""" + +import pytest + + +@pytest.mark.parametrize("dtype", ["fp8", "bf8"]) +@pytest.mark.parametrize("output", ["bf16", "fp16"]) +@pytest.mark.parametrize("route", ["comgr", "hip"]) +@pytest.mark.parametrize("M,N,K", [(32, 64, 128), (64, 96, 256)]) +def test_numeric(gfx950_numeric, dtype, output, route, M, N, K): + gfx950_numeric(dtype, dtype, output, route, M, N, K) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx950_mxfp_examples.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx950_mxfp_examples.py new file mode 100644 index 000000000000..ed0956224eb6 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx950_mxfp_examples.py @@ -0,0 +1,123 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +import argparse +import importlib +from itertools import product + +import numpy as np +import pytest +from rocke.core.lower_llvm import _lower_kernel_to_llvm_python +from rocke.examples.common.scaled_gemm_reference import reference_result +from rocke.examples.gfx950.gemm.block_scaled_gemm_verify import ( + assert_scale_coverage, + make_inputs, +) +from rocke.instances.gfx950.block_scaled_gemm import ( + BlockScaledGemmSpec, + build_block_scaled_gemm, +) + +FORMATS = ("fp8", "bf8", "fp6", "bf6", "fp4") + + +@pytest.mark.parametrize( + "family,dtype", + [ + ("mxfp8", "fp8"), + ("mxfp8", "bf8"), + ("mxfp6", "fp6"), + ("mxfp6", "bf6"), + ("mxfp4", "fp4"), + ], +) +@pytest.mark.parametrize("output", ["bf16", "fp16"]) +def test_example_builds_real_kernel(family, dtype, output): + ex = importlib.import_module(f"rocke.examples.gfx950.gemm.{family}_gemm") + spec = ex.make_spec( + argparse.Namespace(m=32, n=64, k=128, dtype=dtype, output_dtype=output) + ) + assert spec.block_k == 32 and spec.block_size == 64 + assert "call <16 x float> @llvm.amdgcn.mfma.scale" in _lower_kernel_to_llvm_python( + build_block_scaled_gemm(spec), arch="gfx950", llvm_flavor="llvm23" + ) + + +@pytest.mark.parametrize( + "family,count", [("mxfp8", 2), ("mxfp6", 2), ("mxfp4", 1), ("mixed_scaled", 1)] +) +def test_default_and_all_case_cli(monkeypatch, family, count): + from rocke.examples.gfx950.gemm import _scaled_gemm_example as common + + ex = importlib.import_module( + f"rocke.examples.gfx950.gemm.{family}_gemm" + if family != "mixed_scaled" + else "rocke.examples.gfx950.gemm.mixed_scaled_gemm" + ) + seen = [] + monkeypatch.setattr( + common, + "run_cases", + lambda spec, cases, **kw: seen.append((spec, cases, kw)) or len(cases), + ) + assert ( + ex.main(["--case", "all", "--output-dtype", "fp16", "--compile-route", "hip"]) + == 0 + ) + assert len(seen) == count + for spec, cases, kw in seen: + assert (spec.N, spec.dtype_c) == (64, "fp16") + assert set(cases) >= { + "a-only", + "b-only", + "mixed", + "positive", + "group-0", + "group-7", + } + assert kw == {"compile_route": "hip"} + + +@pytest.mark.parametrize("a,b", list(product(FORMATS, repeat=2))) +@pytest.mark.parametrize("output", ["bf16", "fp16"]) +def test_nontrivial_scale_coverage_and_oracle_sensitivity(a, b, output): + spec = BlockScaledGemmSpec("fixture", 64, 96, 256, a, b, output) + + def reference(inputs): + return reference_result( + *inputs, + 32, + native=True, + dtype_a=spec.dtype_a, + dtype_b=spec.dtype_b, + dtype_c=spec.dtype_c, + ) + + x = make_inputs(spec, "positive") + assert_scale_coverage(spec, "positive", *x[2:]) + sa, sb = x[2:] + assert np.all(sa[:, 0::2] != sa[:, 1::2]) and np.all(sb[0::2] != sb[1::2]) + expected = reference(x) + faults = [ + (np.full_like(sa, 127), sb), + (sa, np.full_like(sb, 127)), + (np.full_like(sa, 127), np.full_like(sb, 127)), + (np.repeat(sa[:, 0::2], 2, axis=1), sb), + (sa, np.repeat(sb[0::2], 2, axis=0)), + (np.roll(sa, 1, axis=1), sb), + (sa, np.roll(sb, 1, axis=0)), + (np.roll(sa, 1, axis=0), sb), + (sa, np.roll(sb, 1, axis=1)), + ] + for bad_a, bad_b in faults: + assert np.any(reference((x[0], x[1], bad_a, bad_b)) != expected) + for case in ("a-only", "b-only", "mixed", "positive", "group-0", "group-7"): + y = make_inputs(spec, case) + assert_scale_coverage(spec, case, *y[2:]) + assert np.isfinite(reference(y)).all() + + +def test_mixed_rejects_equal_aliases(): + from rocke.examples.gfx950.gemm.mixed_scaled_gemm import make_spec + + with pytest.raises(ValueError, match="different A/B"): + make_spec(argparse.Namespace(dtype_a="fp8", dtype_b="fp8e4m3"))