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 770808827ede..720da7886a36 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_mfma_gemm.h" #include "rocke/instance_mx_gemm.h" #include "rocke/instance_streamk_gemm.h" +#include "rocke/instance_tf32_mma_probe.h" #include "rocke/ir.h" #include "rocke/ir_serialize.h" #include "rocke/lower_llvm.h" @@ -3517,6 +3518,25 @@ PYBIND11_MODULE(rocke_engine, m) []() { return std::string(rocke_engine_version()); }, "Human-readable engine version of this module."); + m.def( + "tf32_mma_probe_serialize_ir", + [](int shape, const std::string& preparation) { + rocke_ir_builder_t b; + rocke_kernel_def_t* kernel = rocke_build_tf32_mma_probe(&b, shape, preparation.c_str()); + char* text = nullptr; + rocke_status_t status + = kernel ? rocke_ir_serialize(kernel, &text) : rocke_ir_builder_status(&b); + std::string error = rocke_ir_builder_error(&b); + std::string result = text ? text : ""; + free(text); + rocke_ir_builder_free(&b); + if(status != ROCKE_OK) + throw std::runtime_error(error); + return result; + }, + py::arg("m") = 16, + py::arg("preparation") = "raw"); + /* ---- family-agnostic lower-from-serialized-IR ---- * The keystone of the ROCKE_BACKEND=cpp default for Python-authored * kernels: a Python front end serializes any KernelDef to ck.dsl.ir/v1 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 756aae8e0a03..3e5b73948ebd 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 @@ -708,6 +708,67 @@ static const rocke_layout_map_t lm_mfma_32x32x16_a = {ROCKE_MMA_ROLE_A, 8, 64, _ static const rocke_layout_map_t lm_mfma_32x32x16_b = {ROCKE_MMA_ROLE_B, 8, 64, _mfma_b_32x32x16}; static const rocke_layout_map_t lm_mfma_32x32x16_c = {ROCKE_MMA_ROLE_ACC, 16, 64, _mfma_acc_32x32}; +static void _mfma_a_16x16x8_xf32( + rocke_ir_builder_t* b, rocke_value_t* lane, int slot, rocke_value_t** c0, rocke_value_t** c1) +{ + rocke_value_t* c = rocke_b_const_i32(b, 16); + rocke_value_t* axis = rocke_b_mod(b, lane, c); + rocke_value_t* group = rocke_b_div(b, lane, c); + rocke_value_t* two = rocke_b_const_i32(b, 2); + rocke_value_t* base = rocke_b_mul(b, group, two); + rocke_value_t* index = rocke_b_const_i32(b, slot); + rocke_value_t* k = rocke_b_add(b, base, index); + *c0 = axis; + *c1 = k; +} + +static void _mfma_b_16x16x8_xf32( + rocke_ir_builder_t* b, rocke_value_t* lane, int slot, rocke_value_t** c0, rocke_value_t** c1) +{ + rocke_value_t* c = rocke_b_const_i32(b, 16); + rocke_value_t* axis = rocke_b_mod(b, lane, c); + rocke_value_t* group = rocke_b_div(b, lane, c); + rocke_value_t* two = rocke_b_const_i32(b, 2); + rocke_value_t* base = rocke_b_mul(b, group, two); + rocke_value_t* index = rocke_b_const_i32(b, slot); + rocke_value_t* k = rocke_b_add(b, base, index); + *c0 = k; + *c1 = axis; +} + +static void _mfma_a_32x32x4_xf32( + rocke_ir_builder_t* b, rocke_value_t* lane, int slot, rocke_value_t** c0, rocke_value_t** c1) +{ + rocke_value_t* c = rocke_b_const_i32(b, 32); + rocke_value_t* axis = rocke_b_mod(b, lane, c); + rocke_value_t* group = rocke_b_div(b, lane, c); + rocke_value_t* two = rocke_b_const_i32(b, 2); + rocke_value_t* base = rocke_b_mul(b, group, two); + rocke_value_t* index = rocke_b_const_i32(b, slot); + rocke_value_t* k = rocke_b_add(b, base, index); + *c0 = axis; + *c1 = k; +} + +static void _mfma_b_32x32x4_xf32( + rocke_ir_builder_t* b, rocke_value_t* lane, int slot, rocke_value_t** c0, rocke_value_t** c1) +{ + rocke_value_t* c = rocke_b_const_i32(b, 32); + rocke_value_t* axis = rocke_b_mod(b, lane, c); + rocke_value_t* group = rocke_b_div(b, lane, c); + rocke_value_t* two = rocke_b_const_i32(b, 2); + rocke_value_t* base = rocke_b_mul(b, group, two); + rocke_value_t* index = rocke_b_const_i32(b, slot); + rocke_value_t* k = rocke_b_add(b, base, index); + *c0 = k; + *c1 = axis; +} + +static const rocke_layout_map_t lm_xf32_16_a = {ROCKE_MMA_ROLE_A, 2, 64, _mfma_a_16x16x8_xf32}; +static const rocke_layout_map_t lm_xf32_16_b = {ROCKE_MMA_ROLE_B, 2, 64, _mfma_b_16x16x8_xf32}; +static const rocke_layout_map_t lm_xf32_32_a = {ROCKE_MMA_ROLE_A, 2, 64, _mfma_a_32x32x4_xf32}; +static const rocke_layout_map_t lm_xf32_32_b = {ROCKE_MMA_ROLE_B, 2, 64, _mfma_b_32x32x4_xf32}; + /* --- mfma_f32_16x16x4_f32: a/b/c present (frag 1/1/4, wave64) (#8348) --- */ static const rocke_layout_map_t lm_mfma_16x16x4_f32_a = {ROCKE_MMA_ROLE_A, 1, 64, _mfma_a_16x16x4_f32}; @@ -890,6 +951,8 @@ static_assert(offsetof(rocke_ati_mma_frag_row_t, c_frag_len) == sizeof(const cha "c_frag_len must immediately follow op_id (no leading padding)"); static const rocke_ati_mma_frag_row_t rocke_ati_mma_frag[] = { + {"mfma_f32_32x32x4_xf32", 16}, + {"mfma_f32_16x16x8_xf32", 4}, /* --- MFMA fp32 (wave64) --- */ {"mfma_f32_16x16x4_f32", 4}, {"mfma_f32_32x32x2_f32", 16}, @@ -1042,6 +1105,37 @@ static const rocke_mma_op_t k_mma_gfx90a[] = { /* ----------------------------- gfx942 (CDNA) ----------------------------- */ static const rocke_mma_op_t k_mma_gfx942[] = { + {"mma", + "tf32", + "tf32", + "fp32", + 16, + 16, + 8, + "mfma_f32_16x16x8_xf32", + 2, + 2, + 4, + 64, + &lm_xf32_16_a, + &lm_xf32_16_b, + &lm_mfma_16x16x4_f32_c}, + {"mma", + "tf32", + "tf32", + "fp32", + 32, + 32, + 4, + "mfma_f32_32x32x4_xf32", + 2, + 2, + 16, + 64, + &lm_xf32_32_a, + &lm_xf32_32_b, + &lm_mfma_32x32x2_f32_c}, + {"mma", "fp32", "fp32", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/dtypes.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/dtypes.cpp index dcc32275b741..362d722b9e26 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/dtypes.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/dtypes.cpp @@ -26,33 +26,16 @@ typedef struct rocke_dtype_alias /* Byte-for-byte the _DTYPE_ALIASES map (insertion order is irrelevant: lookup is * by exact lowercased key). */ static const rocke_dtype_alias_t k_dtype_aliases[] = { - {"f16", "fp16"}, - {"half", "fp16"}, - {"fp16", "fp16"}, - {"bf16", "bf16"}, - {"bfloat16", "bf16"}, - {"f32", "fp32"}, - {"float", "fp32"}, - {"fp32", "fp32"}, - {"iu8", "iu8"}, - {"iu4", "iu4"}, - {"i8", "i8"}, - {"int8", "i8"}, - {"i4", "i4"}, - {"int4", "i4"}, - {"i32", ROCKE_DTYPE_I32}, - {"int32", ROCKE_DTYPE_I32}, - {"fp8", "fp8e4m3"}, - {"fp8e4m3", "fp8e4m3"}, - {"e4m3", "fp8e4m3"}, - {"bf8", "bf8e5m2"}, - {"bf8e5m2", "bf8e5m2"}, - {"fp6", "fp6e2m3"}, - {"fp6e2m3", "fp6e2m3"}, - {"bf6", "fp6e3m2"}, - {"fp6e3m2", "fp6e3m2"}, - {"fp4", "fp4e2m1"}, - {"fp4e2m1", "fp4e2m1"}, + {"f16", "fp16"}, {"half", "fp16"}, {"fp16", "fp16"}, + {"bf16", "bf16"}, {"bfloat16", "bf16"}, {"f32", "fp32"}, + {"float", "fp32"}, {"fp32", "fp32"}, {"tf32", "tf32"}, + {"xf32", "tf32"}, {"iu8", "iu8"}, {"iu4", "iu4"}, + {"i8", "i8"}, {"int8", "i8"}, {"i4", "i4"}, + {"int4", "i4"}, {"i32", ROCKE_DTYPE_I32}, {"int32", ROCKE_DTYPE_I32}, + {"fp8", "fp8e4m3"}, {"fp8e4m3", "fp8e4m3"}, {"e4m3", "fp8e4m3"}, + {"bf8", "bf8e5m2"}, {"bf8e5m2", "bf8e5m2"}, {"fp6", "fp6e2m3"}, + {"fp6e2m3", "fp6e2m3"}, {"bf6", "fp6e3m2"}, {"fp6e3m2", "fp6e3m2"}, + {"fp4", "fp4e2m1"}, {"fp4e2m1", "fp4e2m1"}, }; #define K_NUM_DTYPE_ALIASES ((int)(sizeof(k_dtype_aliases) / sizeof(k_dtype_aliases[0]))) @@ -140,23 +123,15 @@ const char* rocke_normalize_dtype(const char* name, char* lowered, size_t cap) } static const rocke_dtype_info_t k_dtype_info[] = { - {"fp32", ROCKE_DTYPE_FLOAT, 32}, - {"fp16", ROCKE_DTYPE_FLOAT, 16}, - {"bf16", ROCKE_DTYPE_FLOAT, 16}, - {"fp8e4m3", ROCKE_DTYPE_FLOAT, 8}, - {"bf8e5m2", ROCKE_DTYPE_FLOAT, 8}, - {"fp6e2m3", ROCKE_DTYPE_FLOAT, 6}, - {"fp6e3m2", ROCKE_DTYPE_FLOAT, 6}, - {"fp4e2m1", ROCKE_DTYPE_FLOAT, 4}, - {"e8m0", ROCKE_DTYPE_FLOAT, 8}, - {"e5m3", ROCKE_DTYPE_FLOAT, 8}, - {"i4", ROCKE_DTYPE_INTEGER, 4}, - {"iu4", ROCKE_DTYPE_INTEGER, 4}, - {"i8", ROCKE_DTYPE_INTEGER, 8}, - {"iu8", ROCKE_DTYPE_INTEGER, 8}, - {"i16", ROCKE_DTYPE_INTEGER, 16}, - {"i32", ROCKE_DTYPE_INTEGER, 32}, - {"i64", ROCKE_DTYPE_INTEGER, 64}, + {"fp32", ROCKE_DTYPE_FLOAT, 32}, {"tf32", ROCKE_DTYPE_FLOAT, 32}, + {"fp16", ROCKE_DTYPE_FLOAT, 16}, {"bf16", ROCKE_DTYPE_FLOAT, 16}, + {"fp8e4m3", ROCKE_DTYPE_FLOAT, 8}, {"bf8e5m2", ROCKE_DTYPE_FLOAT, 8}, + {"fp6e2m3", ROCKE_DTYPE_FLOAT, 6}, {"fp6e3m2", ROCKE_DTYPE_FLOAT, 6}, + {"fp4e2m1", ROCKE_DTYPE_FLOAT, 4}, {"e8m0", ROCKE_DTYPE_FLOAT, 8}, + {"e5m3", ROCKE_DTYPE_FLOAT, 8}, {"i4", ROCKE_DTYPE_INTEGER, 4}, + {"iu4", ROCKE_DTYPE_INTEGER, 4}, {"i8", ROCKE_DTYPE_INTEGER, 8}, + {"iu8", ROCKE_DTYPE_INTEGER, 8}, {"i16", ROCKE_DTYPE_INTEGER, 16}, + {"i32", ROCKE_DTYPE_INTEGER, 32}, {"i64", ROCKE_DTYPE_INTEGER, 64}, {"i1", ROCKE_DTYPE_PREDICATE, 1}, }; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp index 6062cadb30c0..cd34e3b984c0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp @@ -53,6 +53,7 @@ ROCKE_SCALAR_SINGLETON(rocke_i64, ROCKE_SCALAR_I64, "i64") ROCKE_SCALAR_SINGLETON(rocke_bf16, ROCKE_SCALAR_BF16, "bf16") ROCKE_SCALAR_SINGLETON(rocke_f16, ROCKE_SCALAR_F16, "f16") ROCKE_SCALAR_SINGLETON(rocke_f32, ROCKE_SCALAR_F32, "f32") +ROCKE_SCALAR_SINGLETON(rocke_tf32, ROCKE_SCALAR_TF32, "tf32") ROCKE_SCALAR_SINGLETON(rocke_fp8e4m3, ROCKE_SCALAR_FP8E4M3, "fp8e4m3") ROCKE_SCALAR_SINGLETON(rocke_bf8e5m2, ROCKE_SCALAR_BF8E5M2, "bf8e5m2") @@ -82,6 +83,8 @@ const rocke_type_t* rocke_scalar_by_name(const char* name) return rocke_bf16(); if(strcmp(name, "f16") == 0) return rocke_f16(); + if(strcmp(name, "tf32") == 0) + return rocke_tf32(); if(strcmp(name, "f32") == 0) return rocke_f32(); if(strcmp(name, "fp8e4m3") == 0 || strcmp(name, "e4m3") == 0) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_arith.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_arith.cpp index 650ab185dffb..7d0b95e6cf4b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_arith.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_arith.cpp @@ -609,6 +609,48 @@ rocke_value_t* rocke_b_cast_to_f32(rocke_ir_builder_t* b, rocke_value_t* v) return rocke_i_op1(b, ROCKE_OP_ARITH_CAST_TO_F32, operands, 1, rocke_f32(), NULL, "f32"); } +rocke_value_t* rocke_b_cvt_f32_to_tf32(rocke_ir_builder_t* b, rocke_value_t* v) +{ + if(!rocke_i_live(b)) + return NULL; + if(!v || !rocke_type_eq(v->type, rocke_f32())) + return (rocke_value_t*)rocke_i_set_err( + b, ROCKE_ERR_VALUE, "cvt_f32_to_tf32 expects f32 input"); + rocke_value_t* bits = rocke_b_bitcast(b, v, rocke_i32()); + rocke_value_t* magnitude_mask = rocke_b_const_i32(b, 2147483647); + rocke_value_t* mag = rocke_b_land(b, bits, magnitude_mask); + rocke_value_t* sign_mask = rocke_b_const_i32(b, -2147483648); + rocke_value_t* sign = rocke_b_land(b, bits, sign_mask); + rocke_value_t* exponent_mask = rocke_b_const_i32(b, 2139095040); + rocke_value_t* exp = rocke_b_land(b, mag, exponent_mask); + rocke_value_t* exponent_all_ones = rocke_b_const_i32(b, 2139095040); + rocke_value_t* special = rocke_b_cmp_eq(b, exp, exponent_all_ones); + rocke_value_t* zero = rocke_b_const_i32(b, 0); + rocke_value_t* safe = rocke_b_select(b, special, zero, mag); + rocke_value_t* discarded_bits = rocke_b_const_i32(b, 13); + rocke_value_t* shift = rocke_b_lshr(b, safe, discarded_bits); + rocke_value_t* low_bit = rocke_b_const_i32(b, 1); + rocke_value_t* odd = rocke_b_land(b, shift, low_bit); + rocke_value_t* rounding_bias = rocke_b_const_i32(b, 4095); + rocke_value_t* bias = rocke_b_add(b, odd, rounding_bias); + rocke_value_t* rounded = rocke_b_add(b, safe, bias); + rocke_value_t* precision_mask = rocke_b_const_i32(b, -8192); + rounded = rocke_b_land(b, rounded, precision_mask); + rounded = rocke_b_lor(b, rounded, sign); + rocke_value_t* fraction_mask = rocke_b_const_i32(b, 8388607); + rocke_value_t* frac = rocke_b_land(b, mag, fraction_mask); + rocke_value_t* zero_fraction = rocke_b_const_i32(b, 0); + rocke_value_t* is_nan = rocke_b_cmp_ne(b, frac, zero_fraction); + rocke_value_t* quiet_bit = rocke_b_const_i32(b, 4194304); + rocke_value_t* quiet = rocke_b_lor(b, bits, quiet_bit); + rocke_value_t* nan_precision_mask = rocke_b_const_i32(b, -8192); + quiet = rocke_b_land(b, quiet, nan_precision_mask); + rocke_value_t* nonfinite = rocke_b_select(b, is_nan, quiet, bits); + rocke_value_t* result = rocke_b_select(b, special, nonfinite, rounded); + result = rocke_b_bitcast(b, result, rocke_tf32()); + return result; +} + rocke_value_t* rocke_b_cast_f32_to(rocke_ir_builder_t* b, rocke_value_t* v, const rocke_type_t* target) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp index b07e81cb6b16..05608fbafc5d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp @@ -231,7 +231,8 @@ rocke_value_t* rocke_b_global_load_vN(rocke_ir_builder_t* b, return (rocke_value_t*)rocke_i_set_err( b, ROCKE_ERR_VALUE, "unsupported vector width for global_load_vN: %d", n); } - else if(rocke_i_type_is(dtype, "f32") || rocke_i_type_is(dtype, "i32")) + else if(rocke_i_type_is(dtype, "f32") || rocke_i_type_is(dtype, "i32") + || rocke_i_type_is(dtype, "tf32")) { elem_bytes = 4; if(n != 2 && n != 3 && n != 4 && n != 8) @@ -251,7 +252,7 @@ rocke_value_t* rocke_b_global_load_vN(rocke_ir_builder_t* b, return (rocke_value_t*)rocke_i_set_err( b, ROCKE_ERR_VALUE, - "global_load_vN supports f16/bf16/i16/f32/i32/fp8e4m3/bf8e5m2/i8, got %s", + "global_load_vN supports f16/bf16/i16/f32/i32/tf32/fp8e4m3/bf8e5m2/i8, got %s", en); } vt = rocke_vector_type(b, dtype, n); @@ -339,7 +340,7 @@ void rocke_b_global_store_vN(rocke_ir_builder_t* b, return; } } - else if(rocke_i_type_is(et, "f32") || rocke_i_type_is(et, "i32")) + else if(rocke_i_type_is(et, "f32") || rocke_i_type_is(et, "i32") || rocke_i_type_is(et, "tf32")) { elem_bytes = 4; if(n == 16) @@ -359,7 +360,7 @@ void rocke_b_global_store_vN(rocke_ir_builder_t* b, (void)rocke_i_set_err( b, ROCKE_ERR_VALUE, - "global_store_vN supports f16/bf16/i16/f32/i32/i8/fp8e4m3/bf8e5m2, got %s", + "global_store_vN supports f16/bf16/i16/f32/i32/tf32/i8/fp8e4m3/bf8e5m2, got %s", en); return; } 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 0585df2a5aab..3311d8e603ba 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/tf32_internal.h" #include #include @@ -98,7 +99,8 @@ static int rocke_elem_bytes_name(const char* elem_name) { return 1; } - if(strcmp(elem_name, "f32") == 0 || strcmp(elem_name, "i32") == 0) + if(strcmp(elem_name, "f32") == 0 || strcmp(elem_name, "i32") == 0 + || strcmp(elem_name, "tf32") == 0) { return 4; } @@ -363,13 +365,13 @@ rocke_value_t* rocke_b_smem_load_vN(rocke_ir_builder_t* b, } dn = dtype->name; if(!(strcmp(dn, "f16") == 0 || strcmp(dn, "bf16") == 0 || strcmp(dn, "f32") == 0 - || strcmp(dn, "i32") == 0 || strcmp(dn, "fp8e4m3") == 0 || strcmp(dn, "bf8e5m2") == 0 - || strcmp(dn, "i8") == 0)) + || strcmp(dn, "tf32") == 0 || strcmp(dn, "i32") == 0 || strcmp(dn, "fp8e4m3") == 0 + || strcmp(dn, "bf8e5m2") == 0 || strcmp(dn, "i8") == 0)) { return (rocke_value_t*)rocke_i_set_err( b, ROCKE_ERR_VALUE, - "smem_load_vN supports f16 / bf16 / f32 / i32 / fp8e4m3 / " + "smem_load_vN supports f16 / bf16 / f32 / i32 / tf32 / fp8e4m3 / " "bf8e5m2 / i8, got %s", dn); } @@ -607,6 +609,9 @@ rocke_value_t* rocke_b_mma(rocke_ir_builder_t* b, { ops[3 + i] = extra[i]; } + 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); 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/core.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp index 035c013c9e53..2df181aaa4e6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp @@ -28,6 +28,7 @@ * arena-backed table keyed by the producing Value pointer, since the frozen IR * attrs must not be mutated. */ +#include "rocke/tf32_internal.h" #include #include #include @@ -295,7 +296,7 @@ const char* rocke_h_hip_scalar(const char* ir_scalar_name) { return "int16_t"; } - if(strcmp(ir_scalar_name, "i32") == 0) + if(strcmp(ir_scalar_name, "i32") == 0 || strcmp(ir_scalar_name, "tf32") == 0) { return "int"; } @@ -345,7 +346,7 @@ const char* rocke_h_vec_prefix(const char* ir_scalar_name, bool full_map) { return "f32x"; } - if(strcmp(ir_scalar_name, "i32") == 0) + if(strcmp(ir_scalar_name, "i32") == 0 || strcmp(ir_scalar_name, "tf32") == 0) { return "i32x"; } @@ -383,8 +384,8 @@ static bool rocke_h_scalar_in_vec_map(const char* name, bool full_map) { return false; } - return strcmp(name, "f32") == 0 || strcmp(name, "i32") == 0 || strcmp(name, "i16") == 0 - || strcmp(name, "i8") == 0 || strcmp(name, "fp8e4m3") == 0 + return strcmp(name, "tf32") == 0 || strcmp(name, "f32") == 0 || strcmp(name, "i32") == 0 + || strcmp(name, "i16") == 0 || strcmp(name, "i8") == 0 || strcmp(name, "fp8e4m3") == 0 || strcmp(name, "bf8e5m2") == 0; } @@ -450,8 +451,9 @@ const char* rocke_h_type_to_hip(rocke_h_lowerer_t* lw, const rocke_type_t* t) * through to the KeyError. Detect the listed set explicitly so an * unknown vector elem is an error rather than silently "f16x". */ if(strcmp(elem, "f16") != 0 && strcmp(elem, "bf16") != 0 && strcmp(elem, "f32") != 0 - && strcmp(elem, "i32") != 0 && strcmp(elem, "i16") != 0 && strcmp(elem, "i8") != 0 - && strcmp(elem, "fp8e4m3") != 0 && strcmp(elem, "bf8e5m2") != 0) + && strcmp(elem, "tf32") != 0 && strcmp(elem, "i32") != 0 && strcmp(elem, "i16") != 0 + && strcmp(elem, "i8") != 0 && strcmp(elem, "fp8e4m3") != 0 + && strcmp(elem, "bf8e5m2") != 0) { rocke_h_fail(lw, ROCKE_ERR_KEY, "type_to_hip: unmappable vector elem '%s'", elem); return ""; @@ -796,6 +798,9 @@ rocke_status_t rocke_h_lower_op(rocke_h_lowerer_t* lw, const rocke_op_t* op) { return rocke_h_fail(lw, ROCKE_ERR_VALUE, "lower_op: NULL op"); } + const char* tf32_error = rocke_tf32_op_error(op); + if(tf32_error) + return rocke_h_fail(lw, ROCKE_ERR_VALUE, "%s", tf32_error); fn = rocke_h_dispatch(op->opcode); if(!fn) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp index 60129168ffcd..ced6af4ad0a9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp @@ -641,6 +641,22 @@ static rocke_status_t _op_memref_global_store_vN(rocke_h_lowerer_t* lw, const ro n = mem_attr_int(op, "vec", 0); elem_name = mem_attr_str(op, "elem_type", "f16"); prefix = rocke_h_vec_prefix_checked(lw, elem_name, /*full_map=*/true, "global_store_vN"); + const int64_t byte_count = n * (rocke_dtype_info(elem_name)->encoded_bits / 8); + const int64_t align = mem_attr_int(op, "align", byte_count); + if(align <= 0 || (align & (align - 1))) + return rocke_h_fail( + lw, ROCKE_ERR_VALUE, "global_store_vN: alignment must be a positive power of two"); + if(align < byte_count || (byte_count & (byte_count - 1))) + { + rocke_h_emitf(lw, + "__builtin_memcpy(__builtin_assume_aligned(%s + %s, %lld), &%s, %lld);", + rocke_h_name(lw, ptr), + rocke_h_name(lw, idx), + (long long)align, + rocke_h_name(lw, val), + (long long)byte_count); + return lw->status; + } rocke_h_emitf(lw, "*reinterpret_cast<%s%lld*>(%s + %s) = %s;", prefix, 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 3abfd360a3c5..6abf236d5547 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/tf32_internal.h" #include "rocke/wmma_scale_internal.h" #include /* snprintf */ @@ -134,6 +135,36 @@ 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"); } + int tf32_count = rocke_tf32_mma_count(op_id); + if(tf32_count) + { + if(strcmp(lw->arch.gfx, "gfx942") != 0) + return rocke_h_fail(lw, ROCKE_ERR_VALUE, "XF32 MMA requires gfx942"); + rocke_h_emitf(lw, + "f32x%d %s = __builtin_amdgcn_%s(__builtin_bit_cast(f32x2, %s), " + "__builtin_bit_cast(f32x2, %s), %s, 0, 0, 0);", + tf32_count, + rocke_h_name(lw, op->results[0]), + op_id, + rocke_h_name(lw, op->operands[0]), + rocke_h_name(lw, op->operands[1]), + rocke_h_name(lw, op->operands[2])); + return ROCKE_OK; + } + /* Ordinary FP32 controls use catalog IDs that have no legacy opcode. */ + if(strcmp(op_id, "mfma_f32_16x16x4_f32") == 0 || strcmp(op_id, "mfma_f32_32x32x2_f32") == 0) + { + bool small = strcmp(op_id, "mfma_f32_16x16x4_f32") == 0; + rocke_h_emitf(lw, + "f32x%d %s = __builtin_amdgcn_%s(%s, %s, %s, 0, 0, 0);", + small ? 4 : 16, + rocke_h_name(lw, op->results[0]), + small ? "mfma_f32_16x16x4f32" : "mfma_f32_32x32x2f32", + rocke_h_name(lw, op->operands[0]), + rocke_h_name(lw, op->operands[1]), + rocke_h_name(lw, op->operands[2])); + return ROCKE_OK; + } if(rocke_gfx1250_scaled_wmma(op_id)) { return h_emit_gfx1250_scaled_wmma(lw, op); 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 f0e1c6153c65..6076f3d8727d 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 @@ -20,6 +20,7 @@ * stub bodies remain in this file. */ #include "rocke/lower_llvm_internal.h" +#include "rocke/tf32_internal.h" #include #include @@ -978,7 +979,7 @@ const char* rocke_ll_llvm_type(rocke_lower_t* L, const rocke_type_t* t) return "i8"; if(strcmp(n, "i16") == 0) return "i16"; - if(strcmp(n, "i32") == 0) + if(strcmp(n, "i32") == 0 || strcmp(n, "tf32") == 0) return "i32"; if(strcmp(n, "i64") == 0) return "i64"; @@ -1573,7 +1574,7 @@ static int ll_smem_seg_size(const rocke_type_t* stype) eb = 1; else if(strcmp(n, "f16") == 0 || strcmp(n, "bf16") == 0) eb = 2; - else if(strcmp(n, "i32") == 0 || strcmp(n, "f32") == 0) + else if(strcmp(n, "i32") == 0 || strcmp(n, "tf32") == 0 || strcmp(n, "f32") == 0) eb = 4; else if(strcmp(n, "i64") == 0) eb = 8; @@ -1876,6 +1877,9 @@ void rocke_ll_lower_op(rocke_lower_t* L, const rocke_op_t* op) { return; } + const char* tf32_error = rocke_tf32_op_error(op); + if(tf32_error) + rocke_ll_fail(L, ROCKE_ERR_VALUE, "%s", tf32_error); rocke_opcode_t oc = op->opcode; rocke_ll_op_fn fn = NULL; if(oc > ROCKE_OP_INVALID && oc < ROCKE_OP__COUNT) 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 327475dc4836..007887954a5f 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 @@ -243,6 +243,12 @@ const rocke_ll_decl_t ROCKE_LL_INTRINSIC_DECLS[] = { {"mfma.f32.32x32x16.f16", "declare <16 x float> @llvm.amdgcn.mfma.f32.32x32x16.f16(<8 x half>, <8 x half>, <16 x " "float>, i32 immarg, i32 immarg, i32 immarg)"}, + {"mfma.f32.16x16x8.xf32", + "declare <4 x float> @llvm.amdgcn.mfma.f32.16x16x8.xf32(<2 x float>, <2 x float>, <4 x " + "float>, i32 immarg, i32 immarg, i32 immarg)"}, + {"mfma.f32.32x32x4.xf32", + "declare <16 x float> @llvm.amdgcn.mfma.f32.32x32x4.xf32(<2 x float>, <2 x float>, <16 x " + "float>, i32 immarg, i32 immarg, i32 immarg)"}, {"mfma.f32.16x16x4f32", "declare <4 x float> @llvm.amdgcn.mfma.f32.16x16x4f32(float, float, <4 x float>, i32 immarg, " "i32 immarg, i32 immarg)"}, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp index 5cf265459213..ab59c3061643 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp @@ -45,7 +45,7 @@ static int ll_elem_bytes(const char* name) { return 2; } - if(strcmp(name, "i32") == 0 || strcmp(name, "f32") == 0) + if(strcmp(name, "i32") == 0 || strcmp(name, "tf32") == 0 || strcmp(name, "f32") == 0) { return 4; } @@ -338,7 +338,10 @@ static void op_memref_global_store_vN(rocke_lower_t* L, const rocke_op_t* op) : rocke_ll_llvm_type(L, val->type); const char* elem_name = ll_is_vec(val->type) ? val->type->elem->name : val->type->name; int elem_bytes = ll_elem_bytes(elem_name); - int64_t align = vec * elem_bytes; + int64_t align = ll_attr_int(op, "align", vec * elem_bytes); + if(align <= 0 || (align & (align - 1))) + rocke_ll_fail( + L, ROCKE_ERR_VALUE, "global_store_vN: alignment must be a positive power of two"); const char* ty = rocke_ll_llvm_type(L, val->type); rocke_ll_emitf(L, " %s = getelementptr inbounds %s, ptr addrspace(1) %s, i32 %s", @@ -561,7 +564,7 @@ static void op_tile_smem_load_vN(rocke_lower_t* L, const rocke_op_t* op) { elem_bytes = 2; } - else if(strcmp(en, "i32") == 0 || strcmp(en, "f32") == 0) + else if(strcmp(en, "i32") == 0 || strcmp(en, "tf32") == 0 || strcmp(en, "f32") == 0) { elem_bytes = 4; } 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 e801c63d4544..da86d56c51f2 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/tf32_internal.h" #include #include @@ -90,6 +91,9 @@ 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(rocke_tf32_mma_count(op_id) && strcmp(L->backend->gfx, "gfx942") != 0) + rocke_ll_fail(L, ROCKE_ERR_VALUE, "XF32 MMA requires gfx942"); + /* f16 / bf16 / f32 dense atoms resolve from the table; the scaled / fp4 / * fp6 / fp8-bf8 / hero atoms keep their dedicated bodies below. */ if(_try_emit_mfma_table(L, op, op_id)) @@ -623,6 +627,18 @@ typedef struct _mfma_spec } _mfma_spec_t; static const _mfma_spec_t MFMA_SPECS[] = { + {"mfma_f32_32x32x4_xf32", + "mfma.f32.32x32x4.xf32", + "llvm.amdgcn.mfma.f32.32x32x4.xf32", + "<2 x i32>", + "<16 x float>", + "<2 x float>"}, + {"mfma_f32_16x16x8_xf32", + "mfma.f32.16x16x8.xf32", + "llvm.amdgcn.mfma.f32.16x16x8.xf32", + "<2 x i32>", + "<4 x float>", + "<2 x float>"}, {"mfma_f32_16x16x16_f16", "mfma.f32.16x16x16f16", "llvm.amdgcn.mfma.f32.16x16x16f16", 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 f9991640c7e1..d30fcbb53c6a 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/tf32_internal.h" #include #include @@ -321,6 +322,9 @@ static void check_op(verifier_t* v, const rocke_op_t* op); static void check_contract(verifier_t* v, const rocke_op_t* op) { + const char* tf32_error = rocke_tf32_op_error(op); + if(tf32_error) + v_errf(v, op, "%s", tf32_error); rocke_opcode_t o = op->opcode; const char* name = op->name; if(is_binary_same_type(o)) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp index 9d8230329102..fa0e634d0aab 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp @@ -13,6 +13,8 @@ const rocke_type_t* rocke_storage_ir_type(const char* dtype) const rocke_dtype_info_t* info = rocke_dtype_info(dtype); if(!info) return NULL; + if(strcmp(info->name, "tf32") == 0) + return rocke_i32(); if(info->encoded_bits % 8 || strcmp(info->name, "e8m0") == 0 || strcmp(info->name, "e5m3") == 0) return rocke_i8(); return rocke_dtype_to_ir_type(info->name); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_tf32_mma_probe.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_tf32_mma_probe.h new file mode 100644 index 000000000000..9431f548329d --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_tf32_mma_probe.h @@ -0,0 +1,15 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +#ifndef ROCKE_INSTANCE_TF32_MMA_PROBE_H +#define ROCKE_INSTANCE_TF32_MMA_PROBE_H +#include "rocke/ir.h" +#ifdef __cplusplus +extern "C" { +#endif +/* Initializes b; returned kernel is owned by b. Mirrors Tf32MmaProbeSpec. */ +rocke_kernel_def_t* + rocke_build_tf32_mma_probe(rocke_ir_builder_t* b, int m, const char* preparation); +#ifdef __cplusplus +} +#endif +#endif diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h index a6fcf74b8ab5..8c6b09ed045a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h @@ -110,6 +110,7 @@ typedef enum rocke_scalar_kind ROCKE_SCALAR_FP6E3M2, ROCKE_SCALAR_E8M0, ROCKE_SCALAR_E5M3, + ROCKE_SCALAR_TF32, ROCKE_SCALAR__COUNT } rocke_scalar_kind_t; @@ -577,6 +578,7 @@ const rocke_type_t* rocke_i64(void); const rocke_type_t* rocke_bf16(void); const rocke_type_t* rocke_f16(void); const rocke_type_t* rocke_f32(void); +const rocke_type_t* rocke_tf32(void); const rocke_type_t* rocke_fp8e4m3(void); const rocke_type_t* rocke_bf8e5m2(void); const rocke_type_t* rocke_fp4e2m1(void); @@ -785,6 +787,8 @@ rocke_value_t* rocke_b_cvt_scalef32_pk_f32_fp8x4(rocke_ir_builder_t* b, rocke_value_t* rocke_b_cvt_scalef32_pk_f32_bf8x4(rocke_ir_builder_t* b, rocke_value_t* v, rocke_value_t* scale); +/* Explicit RNE conversion; bitcast preserves raw FP32 payloads instead. */ +rocke_value_t* rocke_b_cvt_f32_to_tf32(rocke_ir_builder_t* b, rocke_value_t* v); rocke_value_t* rocke_b_cvt_f32_to_fp8(rocke_ir_builder_t* b, rocke_value_t* v); rocke_value_t* rocke_b_cvt_f32_to_bf8(rocke_ir_builder_t* b, rocke_value_t* v); rocke_value_t* rocke_b_cvt_f32_to_i8_sat(rocke_ir_builder_t* b, rocke_value_t* v); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/tf32_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/tf32_internal.h new file mode 100644 index 000000000000..bbbc8ca68721 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/tf32_internal.h @@ -0,0 +1,75 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* Mirrors core/tf32.py. Validation sees logical types, never carrier equality. */ +#ifndef ROCKE_TF32_INTERNAL_H +#define ROCKE_TF32_INTERNAL_H +#include "rocke/ir.h" +#include +static inline int rocke_tf32_mma_count(const char* id) +{ + if(!id) + return 0; + if(strcmp(id, "mfma_f32_16x16x8_xf32") == 0) + return 4; + if(strcmp(id, "mfma_f32_32x32x4_xf32") == 0) + return 16; + return 0; +} +static inline const char* rocke_tf32_mma_error(const char* id, rocke_value_t* const* args, int n) +{ + int count = rocke_tf32_mma_count(id); + if(!count) + { + for(int i = 0; i < n; ++i) + if(args[i] + && (strcmp(args[i]->type->name, "tf32") == 0 + || strncmp(args[i]->type->name, "vec" : "vec"; + if(n != 3 || !args[0] || !args[1] || !args[2] || strcmp(args[0]->type->name, "vec") != 0 + || strcmp(args[1]->type->name, "vec") != 0 || strcmp(args[2]->type->name, acc) != 0) + return "XF32 MMA requires two vec operands and its FP32 accumulator"; + return NULL; +} +static inline const char* rocke_tf32_op_error(const rocke_op_t* op) +{ + const char* id = rocke_attr_get_str(&op->attrs, "op_id"); + if(!id && strncmp(op->name, "tile.", 5) == 0) + id = op->name + 5; + const int is_mma = strcmp(op->name, "tile.mma") == 0 || strncmp(op->name, "tile.mfma", 9) == 0 + || strncmp(op->name, "tile.wmma", 9) == 0; + const char* error = is_mma ? rocke_tf32_mma_error(id, op->operands, op->num_operands) : NULL; + if(error) + return error; + int count = rocke_tf32_mma_count(id); + if(is_mma && !count) + for(int i = 0; i < op->num_results; ++i) + if(strcmp(op->results[i]->type->name, "tf32") == 0 + || strncmp(op->results[i]->type->name, "vecnum_results != 1 + || strcmp(op->results[0]->type->name, count == 4 ? "vec" : "vec") != 0)) + return "XF32 MMA result must match its FP32 accumulator"; + /* Only payload transport, reinterpretation, and selection preserve TF32. */ + if((strncmp(op->name, "arith.", 6) == 0 || strncmp(op->name, "math.", 5) == 0 + || strncmp(op->name, "vector.", 7) == 0) + && strcmp(op->name, "arith.bitcast") != 0 && strcmp(op->name, "arith.select") != 0 + && strcmp(op->name, "vector.bitcast") != 0 && strcmp(op->name, "vector.extract") != 0 + && strcmp(op->name, "vector.pack") != 0 && strcmp(op->name, "vector.concat") != 0 + && strcmp(op->name, "vector.insert") != 0 && strcmp(op->name, "vector.splat") != 0 + && strcmp(op->name, "vector.select") != 0) + { + for(int i = 0; i < op->num_operands + op->num_results; ++i) + { + const rocke_value_t* v + = i < op->num_operands ? op->operands[i] : op->results[i - op->num_operands]; + if(strcmp(v->type->name, "tf32") == 0 || strncmp(v->type->name, "vec +#include + +rocke_kernel_def_t* rocke_build_tf32_mma_probe(rocke_ir_builder_t* b, int m, const char* mode) +{ + char name[96]; + snprintf(name, sizeof(name), "tf32_probe_%d_%s", m, mode ? mode : "invalid"); + if(rocke_ir_builder_init(b, name) != ROCKE_OK) + return NULL; + return ckc::guard_builder(b, [&]() -> rocke_kernel_def_t* { + if((m != 16 && m != 32) || !mode + || (strcmp(mode, "raw") && strcmp(mode, "carrier") && strcmp(mode, "rne") + && strcmp(mode, "prepacked") && strcmp(mode, "fp32"))) + return (rocke_kernel_def_t*)rocke_i_set_err( + b, ROCKE_ERR_VALUE, "TF32 probe requires m=16/32 and a known preparation"); + int k = 128 / m; + const rocke_arch_target_t* target = rocke_arch_target_from_gfx("gfx942"); + const rocke_mma_op_t* atom = rocke_mma_catalog_op_for_shape( + &target->mma, "mma", "tf32", "tf32", "fp32", m, m, k, NULL); + const rocke_type_t* ty = strcmp(mode, "prepacked") == 0 ? rocke_i32() : rocke_f32(); + rocke_param_opts_t opts = {}; + opts.align = 4; + opts.align_set = true; + const rocke_type_t* input_ptr = rocke_ptr_type(b, ty, "global"); + rocke_value_t* a = rocke_b_param(b, "A", input_ptr, &opts); + rocke_value_t* bb = rocke_b_param(b, "B", input_ptr, &opts); + const rocke_type_t* float_ptr = rocke_ptr_type(b, rocke_f32(), "global"); + rocke_value_t* cc = rocke_b_param(b, "C", float_ptr, &opts); + rocke_value_t* d = rocke_b_param(b, "D", float_ptr, &opts); + const rocke_type_t* bits_ptr = rocke_ptr_type(b, rocke_i32(), "global"); + rocke_value_t* pa = rocke_b_param(b, "PA", bits_ptr, &opts); + rocke_value_t* pb = rocke_b_param(b, "PB", bits_ptr, &opts); + rocke_value_t* lane = rocke_b_thread_id_x(b); + rocke_value_t* batch = rocke_b_block_id_x(b); + rocke_value_t* cm = rocke_b_const_i32(b, m); + rocke_value_t* axis = rocke_b_mod(b, lane, cm); + rocke_value_t* group = rocke_b_div(b, lane, cm); + rocke_value_t* ck = rocke_b_const_i32(b, k); + rocke_value_t* row = rocke_b_mul(b, axis, ck); + rocke_value_t* size = rocke_b_const_i32(b, m * k); + rocke_value_t* batch_input = rocke_b_mul(b, batch, size); + rocke_value_t* row_base = rocke_b_add(b, batch_input, row); + size = rocke_b_const_i32(b, m * m); + rocke_value_t* batch_output = rocke_b_mul(b, batch, size); + rocke_value_t* indices[16]; + rocke_value_t* c_values[16]; + for(int slot = 0; slot < atom->c_frag_len; ++slot) + { + rocke_value_t *r, *col; + rocke_layout_map_coord(atom->c_layout, b, lane, slot, &r, &col); + rocke_value_t* offset = rocke_b_mul(b, r, cm); + offset = rocke_b_add(b, offset, col); + indices[slot] = rocke_b_add(b, batch_output, offset); + c_values[slot] = rocke_b_global_load(b, cc, indices[slot], rocke_f32(), 4); + } + rocke_value_t* acc = rocke_b_vec_pack(b, c_values, atom->c_frag_len, rocke_f32()); + if(strcmp(mode, "fp32") == 0) + { + const rocke_mma_op_t* full = rocke_mma_catalog_op_for_shape( + &target->mma, "mma", "fp32", "fp32", "fp32", m, m, k / 2, NULL); + for(int step = 0; step < 2; ++step) + { + rocke_value_t* delta = rocke_b_const_i32(b, step * (64 / m)); + delta = rocke_b_add(b, group, delta); + rocke_value_t* index = rocke_b_add(b, row_base, delta); + rocke_value_t* av = rocke_b_global_load(b, a, index, rocke_f32(), 4); + rocke_value_t* bv = rocke_b_global_load(b, bb, index, rocke_f32(), 4); + rocke_value_t* bits = rocke_b_bitcast(b, av, rocke_i32()); + rocke_b_global_store(b, pa, index, bits, 4); + bits = rocke_b_bitcast(b, bv, rocke_i32()); + rocke_b_global_store(b, pb, index, bits, 4); + acc = rocke_b_mma(b, full->op_id, av, bv, acc, NULL, 0); + } + } + else + { + rocke_matrix_fragment_layout_t layout = {{{32, 32}, 2, 32, 2}, 2, 64 / m, m}; + rocke_value_t* inputs[] = {a, bb}; + rocke_value_t* prepared[] = {pa, pb}; + rocke_value_t* fragments[2]; + for(int role = 0; role < 2; ++role) + { + bool packed = strcmp(mode, "prepacked") == 0; + const rocke_type_t* carrier + = packed || strcmp(mode, "carrier") == 0 ? rocke_i32() : rocke_f32(); + rocke_value_t* values = rocke_h_load_matrix_fragment(b, + inputs[role], + row_base, + group, + 0, + packed ? "tf32" : "fp32", + &layout, + carrier, + 4); + if(strcmp(mode, "rne") == 0) + { + rocke_value_t* elems[2]; + for(int i = 0; i < 2; ++i) + { + rocke_value_t* e = rocke_b_vec_extract(b, values, i); + elems[i] = rocke_b_cvt_f32_to_tf32(b, e); + } + values = rocke_b_vec_pack(b, elems, 2, rocke_tf32()); + } + else + values = rocke_b_bitcast(b, values, rocke_vector_type(b, rocke_tf32(), 2)); + for(int slot = 0; slot < 2; ++slot) + { + rocke_value_t* two = rocke_b_const_i32(b, 2); + rocke_value_t* offset = rocke_b_mul(b, group, two); + rocke_value_t* index = rocke_b_const_i32(b, slot); + offset = rocke_b_add(b, offset, index); + index = rocke_b_add(b, row_base, offset); + rocke_value_t* e = rocke_b_vec_extract(b, values, slot); + rocke_value_t* bits = rocke_b_bitcast(b, e, rocke_i32()); + rocke_b_global_store(b, prepared[role], index, bits, 4); + } + fragments[role] = values; + } + acc = rocke_b_mma(b, atom->op_id, fragments[0], fragments[1], acc, NULL, 0); + } + for(int slot = 0; slot < atom->c_frag_len; ++slot) + { + rocke_value_t* value = rocke_b_vec_extract(b, acc, slot); + rocke_b_global_store(b, d, indices[slot], value, 4); + } + rocke_b_ret(b); + return rocke_ir_builder_kernel(b); + }); +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md index ab5a8349a6f6..e88fbe70f20b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md @@ -40,6 +40,12 @@ from raw byte patterns with the common bit packer. Scale-format validation remains separate from matrix alias normalization: E5M3 is not BF8 E5M2. Registering a scale dtype does not enable a backend selector. +Logical `tf32` (`xf32` alias) is a distinct 32-bit type carried as I32. +Fragment loading can reinterpret FP32 storage into I32 without conversion; +`bitcast` then wraps those bits as logical TF32. Explicit `cvt_f32_to_tf32` +performs RNE preparation. Packing itself never rounds. See the +[TF32 numerical example](../../python/rocke/examples/gfx942/tf32_numerics/README.md). + ## Common packing `BitPacking` / `rocke_bit_packing_t` describes unsigned bit patterns with an @@ -49,6 +55,7 @@ bit first in a little-endian byte stream. | Format | Dense group occupying whole 32-bit words | |---|---| +| TF32 | 1 element in 1 word | | FP16/BF16 | 2 elements in 1 word | | FP8/BF8 | 4 elements in 1 word | | FP6/BF6 | 16 elements across 3 words | diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/ir_serialization_format.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/ir_serialization_format.md index 6e4cc3199de5..a96210026896 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/ir_serialization_format.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/ir_serialization_format.md @@ -386,3 +386,9 @@ kernel @add_one { (`parse` of that text rebuilds a `KernelDef` whose `lower_kernel_to_llvm` output is identical to the original's.) + +The scalar name `tf32` denotes logical TF32 with a 32-bit I32 carrier. It may +appear in vectors, pointers, shared-memory types and loop carries, using the +existing type grammar (for example, `vec`). Serialization preserves its +logical identity independently of `f32` and `i32`; no format-version change is +needed. The public dtype alias `xf32` normalizes to `tf32` before IR construction. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/dtype_support.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/dtype_support.md index 8e8b4921351a..079a77a13115 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/dtype_support.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/dtype_support.md @@ -23,12 +23,18 @@ estimates, shipping order, and acceptance checklists are maintained separately. | **Out of scope** | Runtime/dispatch dtype validation, the provider's C-api surface, and any measurement. | | **Keeping it current** | When a dtype's support changes at any layer, update the affected row(s) and the §2.6 anchors in the same change. The criteria (§2.2) and rubric (§2.3) are the stable part; the scores are not. **Line anchors and the §4B atom tables are hand-maintained mirrors of the basis tree and go stale on rebase** — an arch entry gaining or losing atoms upstream silently invalidates §4B without touching this file. Re-check both against the new base before merging any rebase. | +**TF32 coverage update.** The TF32 entries in §1, §4B, and §4D reflect the +logical dtype and native gfx942 atoms added after the original basis snapshot. +Other scores and source line numbers retain their stated historical basis. + --- ## 1. The canonical dtype set -rocKE's compiler knows exactly **10 scalar IR types**, defined identically in -both engines: +The original basis snapshot listed ten scalar IR types. The table below also +includes the subsequently added logical TF32 type. For the newer logical +low-bit and scale types, see [dtype storage and packing](../architecture/dtype_storage_packing.md). +The definitions in both engines are: - Python: `platform/python/rocke/core/ir.py:43-52` - C++: `platform/cpp/include/rocke/ir.h:93-108`, @@ -44,6 +50,7 @@ both engines: | `f16` | `half` | 16 | float | | `bf16` | `bfloat` | 16 | float | | `f32` | `float` | 32 | float | +| `tf32` | `i32` (carrier) | 32 | logical float; native gfx942 XF32 input | | `fp8e4m3` | `i8` (storage) | 8 | float (OCP E4M3) | | `bf8e5m2` | `i8` (storage) | 8 | float (OCP E5M2) | @@ -52,9 +59,10 @@ LLVM mapping: `lower_llvm.py:1093-1129` (Python, `_llvm_type`) ↔ storage and materialized through `llvm.amdgcn.cvt.*` intrinsics — there is no native `<8 x fp8>` LLVM type in the emitter. -**Sub-byte / block formats are NOT scalar IR types.** `i4`, `fp4`, `fp6`, and -the `e8m0` MX scale exist only as *packed encodings* consumed by dedicated -helpers/atoms, never as first-class `Type` objects (see §4). +**Sub-byte / block formats in the original snapshot.** The basis tree treated +`i4`, `fp4`, `fp6`, and the `e8m0` MX scale as packed encodings consumed by +dedicated helpers/atoms. Those historical scores do not describe the newer +logical-type and storage APIs linked above. > CK-Tile gives these formats named logical types backed by a packed storage > representation: `pk_fp4_t = pk_float4_e2m1_t` @@ -67,15 +75,27 @@ helpers/atoms, never as first-class `Type` objects (see §4). > rocKE's dedicated packed-format paths can provide support without a first-class > IR `Type`; C5 separately measures whether kernel families accept the format. -**Reduced-precision compute (the "tf32 request") is a mode, not a storage type.** -It never appears as a tensor dtype: storage stays `f32` and the *precision* is -reduced before the MMA. What is missing is the **mode itself**: no spec knob selects -reduced-precision compute, so a caller's TF32 request cannot be expressed at all. - -Native XF32 is available on gfx942, a supported rocKE target, but absent on -gfx950 and gfx1250 (§4D). The catalog has no XF32 atom, including on gfx942. -Native instruction absence does not rule out software implementations of TF32 -semantics. +**TF32 is a logical dtype with a 32-bit I32 carrier.** `tf32` and its `xf32` +alias select two native gfx942 XF32 atoms independently of ordinary FP32 +selection. The logical type survives serialization; only A/B are bitcast to +the intrinsic's float-vector ABI, while C/D remain FP32. + +`bitcast(value, TF32)` wraps an FP32 or I32 payload without rounding. +`cvt_f32_to_tf32(value)` explicitly prepares FP32 inputs with round-to-nearest, +ties-to-even; bitcasting back to F32 reinterprets the prepared bits. Generic +TF32 scalar and vector arithmetic is rejected. Bitcasts, vector transport, +selection, and the native atoms are supported. Direct global vector loads +accept 2, 3, 4, or 8 TF32 elements; global vector stores accept 1, 2, 4, or 8. +Both use the same alignment rules as I32. Payload width and address alignment +are independent: underaligned HIP stores copy exactly the payload bytes, and +target/alignment determine the machine transfer instructions. +The fragment loader uses `storage_ir_type("tf32") == I32` and returns carriers +that the caller wraps as logical TF32. + +Native XF32 is available on gfx942, but absent on gfx950 and gfx1250 (§4D). +Production GEMM/convolution selectors and cross-target TF32 emulation remain +outside this support. The [numerical probe](../../python/rocke/examples/gfx942/tf32_numerics/README.md) +covers raw payloads, explicit RNE preparation, and ordinary FP32 controls. A **BF16 compute mode over `f32` storage** could convert operands to `bf16` and use existing bf16 MFMA/WMMA atoms. Its precision differs from TF32: the reference @@ -85,14 +105,17 @@ of an `f32`, `GpuRefTypes.h:92-103`, gated by `USE_TF32`), while `bf16` keeps `1.0` in bf16. Widening validation tolerances does not make these contracts equivalent; `USE_TF32` output is not a bit-accuracy oracle for BF16 compute. -The matrix below therefore labels this possible mode `bf16-mode`. No API selector -or TF32-request mapping currently exists. +The matrix below therefore labels this possible mode `bf16-mode`. It is separate +from logical TF32; no production selector maps TF32 requests to BF16 compute. **Types that are entirely absent from the vocabulary:** `f64`, unsigned `u8/u16/u32/u64`, and a distinct `bool` (predicates reuse `i1`). No engine references them. -Both engines define the same set of **10 scalar IR types**. +TF32 source anchors: [Python type and builders](../../python/rocke/core/ir.py), +[native type declarations](../../cpp/include/rocke/ir.h), +[Python contracts](../../python/rocke/core/tf32.py), and +[native contracts](../../cpp/include/rocke/tf32_internal.h). --- @@ -162,6 +185,8 @@ planning questions and deliberately live outside this document (see Scores 0–3 per criterion; `—` = N/A (excluded from the denominator). Rows are grouped by type family, not ranked — see "What the scores are not" in §2.3. +TF32 was added after this scored snapshot; its current layer boundaries are +described in §1 and §4D rather than assigned a retrospective percentage here. Gap categories: **[C]** conversion gap · **[P]** partial/plumbing gap · **[M]** missing entirely. @@ -406,13 +431,13 @@ to gfx1251, so this is not a missing gfx1250 atom. `core/arch/data/arch_specs.json` is rocKE's architecture SSOT. An atom that is not declared there is unusable by any kernel regardless of what the silicon provides. -Verbatim at the basis commit: +The original basis snapshot, with the gfx942 XF32 addition noted below: | Arch | Declared matrix atoms | dtypes reachable | |---|---|---| | `gfx90a` | `mfma_f32_{16x16x16,32x32x8}_{f16,bf16}` | f16, bf16 | -| `gfx942` | the gfx90a set + `mfma_f32_{16x16x4,32x32x2}_f32` + `mfma_f32_{16x16x32,32x32x16}_{fp8,bf8}` | f32, f16, bf16, fp8, bf8 | -| `gfx950` | the gfx942 set + K-packed `{16x16x32,32x32x16}_{f16,bf16}` + `mfma_f32_16x16x128_fp4` + `mfma_f32_16x16x96_fp6` | f32, f16, bf16, fp8, bf8, **fp4, fp6** | +| `gfx942` | the gfx90a set + `mfma_f32_{16x16x4,32x32x2}_f32` + `mfma_f32_{16x16x32,32x32x16}_{fp8,bf8}` + `mfma_f32_{16x16x8,32x32x4}_xf32` | f32, f16, bf16, fp8, bf8, **tf32** | +| `gfx950` | the gfx942 set **excluding XF32** + K-packed `{16x16x32,32x32x16}_{f16,bf16}` + `mfma_f32_16x16x128_fp4` + `mfma_f32_16x16x96_fp6` | f32, f16, bf16, fp8, bf8, **fp4, fp6** | | `gfx1151` | `wmma_f32_16x16x16_{f16,bf16}`, `wmma_i32_16x16x16_{iu8,iu4}` | f16, bf16, **iu8, iu4** | | `gfx1201` | `wmma_gfx12_f32_16x16x16_{f16,bf16}` | f16, bf16 | | `gfx1250` | `wmma_gfx1250_f32_16x16x4_f32`, `..._16x16x32_{f16,bf16}`, `..._16x16x64_{fp8_fp8,fp8_bf8,bf8_fp8,bf8_bf8}`, `wmma_scale_f32_16x16x128_fp8_fp8`, `wmma_scale16_f32_16x16x128_fp8_fp8` | f32, f16, bf16, fp8, bf8 (+ fp8 block-scaled) | @@ -428,8 +453,8 @@ Verbatim at the basis commit: ### 4C. Deltas — silicon capability vs. rocKE's declared atoms The following gaps compare §4A with §4B. The first five concern atom declarations -only; f64 (4C.6) also lacks an IR type. The gfx942 XF32 declaration gap is covered -separately in §4D. +only; f64 (4C.6) also lacks an IR type. The gfx942 XF32 declaration gap is now +closed; §4D describes its logical dtype and remaining integration boundaries. #### 4C.1 gfx1250 — fp4/fp6/bf6 block-scaled atoms are omitted (deliberately) The part carries `V_WMMA_F32_16X16X128_F8F6F4`, a dedicated @@ -478,11 +503,13 @@ gfx1250 has no native F64 matrix instruction and is excluded from this gap. | gfx1151 / gfx1201 (RDNA3.5 / RDNA4) | ❌ never present | | gfx1250 (CDNA5) | ❌ absent from the public XML and LLVM target features | -**Consequence for rocKE.** gfx942 is a supported target: it has a catalog entry -and block-scale GEMM accepts it for FP8/BF8. It provides native XF32, but rocKE -does not declare XF32 atoms. gfx950/gfx1250 have no native XF32 instruction. -See [§1](#1-the-canonical-dtype-set) for the distinction between native XF32, -software TF32 implementations, and a BF16 compute mode. +**Consequence for rocKE.** The gfx942 catalog declares +`mfma_f32_16x16x8_xf32` and `mfma_f32_32x32x4_xf32`, with logical TF32 A/B and +FP32 C/D. Both Python and C++ lowerers support them and reject these atoms on +other targets. gfx950/gfx1250 have no native XF32 instruction. This atom support +does not add production GEMM/convolution dispatch or cross-target emulation. +See [§1](#1-the-canonical-dtype-set) for the distinction between logical TF32, +raw payload wrapping, explicit RNE preparation, and a BF16 compute mode. ### 4E. The FP8 encoding dialect (FNUZ vs OCP) 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 24ccdf131042..c2d53b3c61db 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 @@ -137,3 +137,10 @@ A `tile_m=128, tile_n=128, warp_m=2, warp_n=2, warp_tile_m=32, warp_tile_n=32` G - Using `f16_16x16x32` or `f16_32x32x16` on gfx940/gfx942: link-time failure (intrinsic absent). - Using bf16 atoms without going through `mfma_16x16x16_for_dtype`: declaration mismatch (plain `bf16.16x16x16` doesn't exist; must use the `_1k` variant). - `4x4x4` epilogue stores not compositing `batch_idx` into the per-group offset: 16 batches collide on the same outputs. + +Logical `tf32` on gfx942 has two native atoms: `mfma_f32_16x16x8_xf32` and +`mfma_f32_32x32x4_xf32`. Both use two TF32 operands per lane for each input, +stored as two I32 words; accumulator sizes are four and sixteen FP32 values. +The [numerical example](../../python/rocke/examples/gfx942/tf32_numerics/README.md) +demonstrates raw input transport and explicit RNE preparation. These atoms do +not alter ordinary FP32 selection or enable TF32 production GEMM dispatch. 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 41642b1170ce..3f9597229488 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 @@ -128,3 +128,15 @@ Existing `wmma_scale*_f32_*` and dotted `wmma.scaled.*` IDs are retired, as are the dedicated scaled builder wrappers; serialized IR using those IDs must be regenerated. Use `tile.mma` with a resolved catalog atom. Other MMA operation IDs retain their existing spelling. + +## Logical TF32 operands + +On gfx942, `tf32` (`xf32` alias) selects `mfma_f32_16x16x8_xf32` or +`mfma_f32_32x32x4_xf32`. Both take two logical TF32 A/B elements per lane; +C/D are FP32 with four or sixteen elements. FP32 queries retain the ordinary +FP32 atoms. The logical type uses one I32 storage word per element and is +bitcast to the intrinsic's float-vector ABI during lowering. + +See the [TF32 numerical example](../../python/rocke/examples/gfx942/tf32_numerics/README.md) +for the distinction between raw FP32 payloads and explicit RNE preparation, +and for a reproducible comparison through both engines. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/op_vocabulary.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/op_vocabulary.md index 231496fc4ac1..da95fd84bb0d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/op_vocabulary.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/op_vocabulary.md @@ -179,3 +179,10 @@ static_if(host_bool, body) # raises TypeError on SSA Value - `ir_lowering/ir_model.md` for the IRBuilder API by category. - `ir_lowering/lowering_pipeline.md` for the op-to-LLVM mapping. - `ir_lowering/backend_details.md` for purity classification and waitcnt encoding. + +`IRBuilder.cvt_f32_to_tf32` / `rocke_b_cvt_f32_to_tf32` is a scalar F32-to-TF32 +RNE recipe built from existing integer arithmetic, select, and bitcast ops; it +introduces no serialized opcode. It quiets NaNs and retains their upper payload +bits. `bitcast` to logical TF32 only reinterprets the payload. Logical TF32 uses +I32 storage; general arithmetic requires explicitly reinterpreting it as F32. +See [TF32 numerics](../../python/rocke/examples/gfx942/tf32_numerics/README.md). diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/__init__.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/__init__.py index ffbea41535ac..cf129d2b4945 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/__init__.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/__init__.py @@ -68,6 +68,7 @@ I16, I32, I64, + TF32, IRBuilder, KernelDef, Op, @@ -174,6 +175,7 @@ "I16", "I32", "I64", + "TF32", "IRBuilder", "KernelDef", "Op", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/__init__.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/__init__.py index c6b12bbdb49e..7abcf1acbf25 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/__init__.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/__init__.py @@ -61,6 +61,7 @@ I8, I32, I64, + TF32, IRBuilder, KernelDef, Op, @@ -123,6 +124,7 @@ "I8", "I32", "I64", + "TF32", "IRBuilder", "KernelDef", "Op", 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 50aae7319e14..757292d9c853 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 @@ -45,6 +45,8 @@ "sgprs": 102 }, "mma": [ + {"family": "mma", "a": "tf32", "b": "tf32", "c": "fp32", "m": 16, "n": 16, "k": 8, "op_id": "mfma_f32_16x16x8_xf32"}, + {"family": "mma", "a": "tf32", "b": "tf32", "c": "fp32", "m": 32, "n": 32, "k": 4, "op_id": "mfma_f32_32x32x4_xf32"}, {"family": "mma", "a": "fp32", "b": "fp32", "c": "fp32", "m": 16, "n": 16, "k": 4, "op_id": "mfma_f32_16x16x4_f32"}, {"family": "mma", "a": "fp32", "b": "fp32", "c": "fp32", "m": 32, "n": 32, "k": 2, "op_id": "mfma_f32_32x32x2_f32"}, {"family": "mma", "a": "fp16", "b": "fp16", "c": "fp32", "m": 16, "n": 16, "k": 16, "op_id": "mfma_f32_16x16x16_f16"}, 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 1400110ccc78..244985e0762a 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 @@ -330,6 +330,38 @@ def _mfma_b_16x16(builder, lane, slot): return k, n_in_atom +def _mfma_a_16x16x8_xf32(builder, lane, slot): + c = builder.const_i32(16) + axis = builder.mod(lane, c) + group = builder.div(lane, c) + k = builder.add(builder.mul(group, builder.const_i32(2)), builder.const_i32(slot)) + return axis, k + + +def _mfma_b_16x16x8_xf32(builder, lane, slot): + c = builder.const_i32(16) + axis = builder.mod(lane, c) + group = builder.div(lane, c) + k = builder.add(builder.mul(group, builder.const_i32(2)), builder.const_i32(slot)) + return k, axis + + +def _mfma_a_32x32x4_xf32(builder, lane, slot): + c = builder.const_i32(32) + axis = builder.mod(lane, c) + group = builder.div(lane, c) + k = builder.add(builder.mul(group, builder.const_i32(2)), builder.const_i32(slot)) + return axis, k + + +def _mfma_b_32x32x4_xf32(builder, lane, slot): + c = builder.const_i32(32) + axis = builder.mod(lane, c) + group = builder.div(lane, c) + k = builder.add(builder.mul(group, builder.const_i32(2)), builder.const_i32(slot)) + return k, axis + + def _mfma_a_16x16x4_f32(builder, lane, slot): """MFMA 16x16x4 fp32 A operand: each lane holds one fp32 scalar. @@ -659,6 +691,12 @@ class _FragInfo: # layout-map functions are populated for the atoms whose lane math is verified. # Adding a new atom is one row here. _MMA_FRAGMENT_INFO: Dict[str, _FragInfo] = { + "mfma_f32_32x32x4_xf32": _FragInfo( + 2, 2, 16, 64, _mfma_a_32x32x4_xf32, _mfma_b_32x32x4_xf32, _mfma_acc_32x32 + ), + "mfma_f32_16x16x8_xf32": _FragInfo( + 2, 2, 4, 64, _mfma_a_16x16x8_xf32, _mfma_b_16x16x8_xf32, _mfma_acc_16x16 + ), # --- MFMA fp32 (wave64) ----------------------------------------------- # A/B are scalar float per lane (a_frag_len=b_frag_len=1); accumulator # shares the standard 16x16 / 32x32 layout (c_frag_len=4 / 16). diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/dtypes.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/dtypes.py index b0745c226610..731eafe5efa6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/dtypes.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/dtypes.py @@ -23,6 +23,8 @@ "f32": "fp32", "float": "fp32", "fp32": "fp32", + "tf32": "tf32", + "xf32": "tf32", "fp8": "fp8e4m3", "fp8e4m3": "fp8e4m3", "e4m3": "fp8e4m3", @@ -76,6 +78,7 @@ class DTypeInfo: DTypeCategory.FLOAT, ( ("fp32", 32), + ("tf32", 32), ("fp16", 16), ("bf16", 16), ("fp8e4m3", 8), 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 b93f34fef188..d08e3a6b54e8 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 @@ -51,6 +51,7 @@ def __repr__(self) -> str: BF16 = Type("bf16") F16 = Type("f16") F32 = Type("f32") +TF32 = Type("tf32") FP8E4M3 = Type("fp8e4m3") BF8E5M2 = Type("bf8e5m2") FP4E2M1 = Type("fp4e2m1") @@ -75,6 +76,7 @@ def dtype_to_ir_type(dtype: str) -> Type: "fp16": F16, "bf16": BF16, "fp32": F32, + "tf32": TF32, "fp8e4m3": FP8E4M3, "bf8e5m2": BF8E5M2, "fp4e2m1": FP4E2M1, @@ -919,6 +921,49 @@ def cast_to_f32(self, v: Value) -> Value: raise ValueError(f"cast_to_f32 unsupported from {v.type.name}") return self._op("arith.cast_to_f32", [v], [F32], result_name_hint="f32").result + def cvt_f32_to_tf32(self, v: Value) -> Value: + """Round f32 to TF32 precision (RNE), carried as i32 bits. + + Finite encodings have 13 zero low bits. Infinities and signed zero are + preserved; NaNs are quieted and retain their upper payload bits. + Use bitcast(v, TF32) to request native XF32 truncation without rounding. + """ + if v.type != F32: + raise ValueError("cvt_f32_to_tf32 expects f32 input") + bits = self.bitcast(v, I32) + magnitude_mask = self.const_i32(2147483647) + mag = self.land(bits, magnitude_mask) + sign_mask = self.const_i32(-2147483648) + sign = self.land(bits, sign_mask) + exponent_mask = self.const_i32(2139095040) + exp = self.land(mag, exponent_mask) + exponent_all_ones = self.const_i32(2139095040) + special = self.cmp_eq(exp, exponent_all_ones) + zero = self.const_i32(0) + safe = self.select(special, zero, mag) + discarded_bits = self.const_i32(13) + shift = self.lshr(safe, discarded_bits) + low_bit = self.const_i32(1) + odd = self.land(shift, low_bit) + rounding_bias = self.const_i32(4095) + bias = self.add(odd, rounding_bias) + rounded = self.add(safe, bias) + precision_mask = self.const_i32(-8192) + rounded = self.land(rounded, precision_mask) + rounded = self.lor(rounded, sign) + fraction_mask = self.const_i32(8388607) + frac = self.land(mag, fraction_mask) + zero_fraction = self.const_i32(0) + is_nan = self.cmp_ne(frac, zero_fraction) + quiet_bit = self.const_i32(4194304) + quiet = self.lor(bits, quiet_bit) + nan_precision_mask = self.const_i32(-8192) + quiet = self.land(quiet, nan_precision_mask) + nonfinite = self.select(is_nan, quiet, bits) + result = self.select(special, nonfinite, rounded) + result = self.bitcast(result, TF32) + return result + def cast_f32_to(self, v: Value, target: Type) -> Value: if v.type.name != "f32": raise ValueError("cast_f32_to expects f32 input") @@ -1572,7 +1617,7 @@ def global_load_vN( ) -> Value: """Vectorised global load of N consecutive values. - Supports f16/bf16/i16 (N in {2, 4, 6, 8, 16}), f32/i32 + Supports f16/bf16/i16 (N in {2, 4, 6, 8, 16}), f32/i32/tf32 (N in {2, 3, 4, 8}), and fp8e4m3/bf8e5m2/i8 (N in {2, 4, 8, 12, 16}). Loads exactly N elements. Instruction selection depends on target and alignment; 96-bit payloads do not require a 96-bit scalar type. @@ -1587,7 +1632,7 @@ def global_load_vN( # <16 x half> operand fragment; AMDGPU coalesces it when aligned. if n not in (2, 4, 6, 8, 16): raise ValueError(f"unsupported vector width for global_load_vN: {n}") - elif dtype.name in ("f32", "i32"): + elif dtype.name in ("f32", "i32", "tf32"): elem_bytes = 4 if n not in (2, 3, 4, 8): raise ValueError( @@ -1601,7 +1646,7 @@ def global_load_vN( ) else: raise ValueError( - "global_load_vN supports f16/bf16/i16/f32/i32/fp8e4m3/bf8e5m2/i8, " + "global_load_vN supports f16/bf16/i16/f32/i32/tf32/fp8e4m3/bf8e5m2/i8, " f"got {dtype.name}" ) return self._op( @@ -1786,7 +1831,7 @@ def smem_store_vN( elem_bytes = ( 1 if elem_name in ("i8", "fp8e4m3", "bf8e5m2") - else 4 if elem_name in ("f32", "i32") else 2 + else 4 if elem_name in ("f32", "i32", "tf32") else 2 ) self._op( "tile.smem_store_vN", @@ -1824,9 +1869,18 @@ def smem_load_vN(self, smem: Value, *indices, dtype: Type, n: int = 0) -> Value: 96-bit payloads (12 bytes, six halfwords, or three words), using element alignment. The target and alignment determine instruction selection. """ - if dtype.name not in ("f16", "bf16", "f32", "i32", "fp8e4m3", "bf8e5m2", "i8"): + if dtype.name not in ( + "f16", + "bf16", + "f32", + "i32", + "tf32", + "fp8e4m3", + "bf8e5m2", + "i8", + ): raise ValueError( - "smem_load_vN supports f16 / bf16 / f32 / i32 / fp8e4m3 / " + "smem_load_vN supports f16 / bf16 / f32 / i32 / tf32 / fp8e4m3 / " f"bf8e5m2 / i8, got {dtype.name}" ) allowed_n = ( @@ -1882,6 +1936,11 @@ def mma( (``a_scale``, ``b_scale``); ordinary atoms take exactly ``a, b, c``. """ 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]) + if error: + raise ValueError(error) c_frag_len = ( op.c_frag_len if hasattr(op, "c_frag_len") and op.c_frag_len @@ -4179,10 +4238,11 @@ def global_store_vN( """Vectorised global store of N consecutive elements. Supports the full element-type catalog the LLVM lowering already - emits: ``f16`` / ``bf16`` / ``i16`` (2-byte), ``f32`` / ``i32`` + emits: ``f16`` / ``bf16`` / ``i16`` (2-byte), ``f32`` / ``i32`` / ``tf32`` (4-byte), ``i8`` / ``fp8e4m3`` / ``bf8e5m2`` (1-byte). Lowers to - a single ``store `` and AMDGPU coalesces into one - ``global_store_dwordxN`` transaction. + a single ``store `` with the supplied address alignment. + Payload width and address alignment are independent; target and + alignment determine whether the transfer uses one machine instruction. """ if n not in (1, 2, 4, 8, 16): raise ValueError(f"global_store_vN n must be 1, 2, 4, 8, or 16 (got {n})") @@ -4195,7 +4255,7 @@ def global_store_vN( elem_bytes = 2 if n == 16: raise ValueError(f"global_store_vN n=16 not supported for {elem_name}") - elif elem_name in ("f32", "i32"): + elif elem_name in ("f32", "i32", "tf32"): elem_bytes = 4 if n == 16: raise ValueError(f"global_store_vN n=16 not supported for {elem_name}") @@ -4203,7 +4263,7 @@ def global_store_vN( elem_bytes = 1 else: raise ValueError( - "global_store_vN supports f16/bf16/i16/f32/i32/i8/fp8e4m3/bf8e5m2, " + "global_store_vN supports f16/bf16/i16/f32/i32/tf32/i8/fp8e4m3/bf8e5m2, " f"got {elem_name}" ) self._op( 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 ab29bc41586b..32249f619793 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 @@ -36,6 +36,7 @@ "i8": "int8_t", "i16": "int16_t", "i32": "int", + "tf32": "int", "i64": "int64_t", "f16": "fp16", "bf16": "bf16", @@ -67,7 +68,7 @@ def _type_to_hip(t) -> str: return f"bf16x{t.count}" if elem == "f32": return f"f32x{t.count}" - if elem == "i32": + if elem in ("i32", "tf32"): return f"i32x{t.count}" if elem == "i16": return f"i16x{t.count}" @@ -182,6 +183,7 @@ def _name(v: Value) -> str: "bf16": "bf16x", "f32": "f32x", "i32": "i32x", + "tf32": "i32x", "i16": "i16x", "i8": "i8x", "fp8e4m3": "i8x", @@ -312,6 +314,11 @@ def _pop_indent(self) -> None: self._indent -= 1 def lower_op(self, op: Op) -> None: + from .tf32 import tf32_op_error + + error = tf32_op_error(op) + if error: + raise ValueError(error) if op.name.startswith("tile.") and gfx1250_scaled_wmma(op.name) is not None: self._emit_wmma_gfx1250_scaled(op) return @@ -552,6 +559,19 @@ def _op_tile_mma(self, op: Op) -> None: This keeps the HIP emission identical to the legacy ISA-named path while the IRBuilder helpers route through :meth:`IRBuilder.mma`. """ + from .tf32 import TF32_MMA + + op_id = op.attrs["op_id"] + if op_id in TF32_MMA: + if self.arch.gfx != "gfx942": + raise ValueError("XF32 MMA requires gfx942") + a, b, c = op.operands + self._emit( + f"f32x{TF32_MMA[op_id]} {_name(op.result)} = __builtin_amdgcn_{op_id}(" + f"__builtin_bit_cast(f32x2, {_name(a)}), " + f"__builtin_bit_cast(f32x2, {_name(b)}), {_name(c)}, 0, 0, 0);" + ) + return op_id = op.attrs["op_id"] legacy = Op( name=f"tile.{op_id}", @@ -1817,6 +1837,19 @@ def _op_memref_global_store_vN(self, op: Op) -> None: n = int(op.attrs["vec"]) elem_name = op.attrs.get("elem_type", "f16") prefix = _vec_prefix(elem_name, "global_store_vN") + byte_count = n * (dtype_info(elem_name).encoded_bits // 8) + align = int(op.attrs.get("align", byte_count)) + if align <= 0 or align & (align - 1): + raise ValueError( + "global_store_vN: alignment must be a positive power of two" + ) + if align < byte_count or byte_count & (byte_count - 1): + self._emit( + f"__builtin_memcpy(" + f"__builtin_assume_aligned({_name(ptr)} + {_name(idx)}, {align}), " + f"&{_name(val)}, {byte_count});" + ) + return self._emit( f"*reinterpret_cast<{prefix}{n}*>({_name(ptr)} + {_name(idx)}) = " f"{_name(val)};" 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 f75a4f438f30..ea3b0b47c9ee 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 @@ -662,6 +662,8 @@ def _resolve_llvm_flavor() -> str: "<8 x half>, <8 x half>, <16 x float>, " "i32 immarg, i32 immarg, i32 immarg)" ), + "mfma.f32.16x16x8.xf32": "declare <4 x float> @llvm.amdgcn.mfma.f32.16x16x8.xf32(<2 x float>, <2 x float>, <4 x float>, i32 immarg, i32 immarg, i32 immarg)", + "mfma.f32.32x32x4.xf32": "declare <16 x float> @llvm.amdgcn.mfma.f32.32x32x4.xf32(<2 x float>, <2 x float>, <16 x float>, i32 immarg, i32 immarg, i32 immarg)", "mfma.f32.16x16x4f32": ( "declare <4 x float> @llvm.amdgcn.mfma.f32.16x16x4f32(" "float, float, <4 x float>, " @@ -1118,7 +1120,7 @@ def _llvm_type(t: Type) -> str: return "i8" if t.name == "i16": return "i16" - if t.name == "i32": + if t.name in ("i32", "tf32"): return "i32" if t.name == "i64": return "i64" @@ -1921,6 +1923,7 @@ def _compute_smem_layout(self) -> None: "f16": 2, "bf16": 2, "i32": 4, + "tf32": 4, "f32": 4, "i64": 8, } @@ -2102,6 +2105,11 @@ def _emit_smem_base_ptr(self, gname: str, stype: SmemType) -> str: # ----- per-op lowerings ----- def lower_op(self, op: Op) -> None: + from .tf32 import tf32_op_error + + error = tf32_op_error(op) + if error: + raise ValueError(error) method = getattr(self, f"_op_{op.name.replace('.', '_')}", None) if method is None: raise NotImplementedError(f"no LLVM lowering for op {op.name!r}") @@ -3123,6 +3131,7 @@ def _op_tile_smem_store(self, op: Op) -> None: "f16": 2, "bf16": 2, "i32": 4, + "tf32": 4, "f32": 4, "i64": 8, }.get(value.type.name, 2) @@ -3187,6 +3196,7 @@ def _op_tile_smem_store_vN(self, op: Op) -> None: "f16": 2, "bf16": 2, "i32": 4, + "tf32": 4, "f32": 4, "i64": 8, }.get( @@ -3264,7 +3274,15 @@ def _op_tile_smem_load_vN(self, op: Op) -> None: elem_ty = _llvm_type(op.result.type.elem) # type: ignore[attr-defined] # Element byte size drives the vector alignment. 16-bit # (f16 / bf16): 2 bytes; 32-bit (f32 / i32): 4 bytes. - elem_bytes = {"i8": 1, "f16": 2, "bf16": 2, "i32": 4, "f32": 4, "i64": 8}.get( + elem_bytes = { + "i8": 1, + "f16": 2, + "bf16": 2, + "i32": 4, + "tf32": 4, + "f32": 4, + "i64": 8, + }.get( op.result.type.elem.name, 2, # type: ignore[attr-defined] ) @@ -3307,6 +3325,10 @@ def _op_tile_wmma_f32_16x16x16_bf16(self, op: Op) -> None: self._backend.emit_wmma(self, op) def _op_tile_mma(self, op: Op) -> None: + from .tf32 import TF32_MMA + + if op.attrs.get("op_id") in TF32_MMA and self._backend.arch.gfx != "gfx942": + raise ValueError("XF32 MMA requires gfx942") # Target-neutral MMA: the ISA backend maps ``op.attrs["op_id"]`` to the # matching MFMA (CDNA) or WMMA (RDNA) emission. CDNA backends reuse the # existing ``_op_tile_`` handler verbatim, so the output is @@ -3366,6 +3388,31 @@ def _op_tile_mfma_f32_16x16x32_bf16(self, op: Op) -> None: f"i32 0, i32 0, i32 0)" ) + def _op_tile_mfma_f32_16x16x8_xf32(self, op: Op) -> None: + self._emit_xf32(op, "mfma.f32.16x16x8.xf32", 4) + + def _op_tile_mfma_f32_32x32x4_xf32(self, op: Op) -> None: + self._emit_xf32(op, "mfma.f32.32x32x4.xf32", 16) + + def _emit_xf32(self, op: Op, intrinsic: str, count: int) -> None: + if self._backend.arch.gfx != "gfx942": + raise ValueError("XF32 MMA requires gfx942") + a, b, c = op.operands + self._need(intrinsic) + a_cast = self._fresh("mfma_a_i16") + b_cast = self._fresh("mfma_b_i16") + self._current().emit( + f" {a_cast} = bitcast <2 x i32> {self._operand(a)} to <2 x float>" + ) + self._current().emit( + f" {b_cast} = bitcast <2 x i32> {self._operand(b)} to <2 x float>" + ) + self._current().emit( + f" {op.result.name} = call <{count} x float> @llvm.amdgcn.{intrinsic}(" + f"<2 x float> {a_cast}, <2 x float> {b_cast}, " + f"<{count} x float> {self._operand(c)}, i32 0, i32 0, i32 0)" + ) + def _op_tile_mfma_f32_16x16x4_f32(self, op: Op) -> None: a, b, c = op.operands self._need("mfma.f32.16x16x4f32") @@ -5467,10 +5514,15 @@ def _op_memref_global_store_vN(self, op: Op) -> None: "f16": 2, "bf16": 2, "i32": 4, + "tf32": 4, "f32": 4, "i64": 8, }.get(elem_name, 2) - align = vec * elem_bytes + align = int(op.attrs.get("align", vec * elem_bytes)) + if align <= 0 or align & (align - 1): + raise ValueError( + "global_store_vN: alignment must be a positive power of two" + ) ty = _llvm_type(val.type) self._current().emit( f" store {ty} {self._operand(val)}, ptr addrspace(1) {gep}, align {align}" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/tf32.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/tf32.py new file mode 100644 index 000000000000..6aea0c65f4ee --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/tf32.py @@ -0,0 +1,57 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Logical TF32 operand contracts, independent of the physical i32 carrier.""" + +TF32_MMA = {"mfma_f32_16x16x8_xf32": 4, "mfma_f32_32x32x4_xf32": 16} + + +def tf32_mma_error(op_id, operands, results=None): + count = TF32_MMA.get(op_id) + if count is None: + if any( + v.type.name == "tf32" or v.type.name.startswith("vec", "vec", f"vec"] + if [v.type.name for v in operands] != expected: + return "XF32 MMA requires two vec operands and its FP32 accumulator" + if results is not None and [v.type.name for v in results] != [expected[2]]: + return "XF32 MMA result must match its FP32 accumulator" + return None + + +def tf32_op_error(op): + op_id = op.attrs.get("op_id", op.name.removeprefix("tile.")) + error = ( + tf32_mma_error(op_id, op.operands, op.results) + if op.name == "tile.mma" or op.name.startswith(("tile.mfma", "tile.wmma")) + else None + ) + if error: + return error + # Vector arithmetic must not operate on the integer carrier bits either. + # Keep the operations that only move, reinterpret, or select payloads. + if op.name.startswith(("arith.", "math.", "vector.")) and op.name not in ( + "arith.bitcast", + "arith.select", + "vector.bitcast", + "vector.extract", + "vector.pack", + "vector.concat", + "vector.insert", + "vector.splat", + "vector.select", + ): + if any( + v.type.name == "tf32" or v.type.name.startswith("vec None: # ---- contracts ---- def _check_contract(self, op: Op, scope: Dict[str, Value]) -> None: + from .tf32 import tf32_op_error + + error = tf32_op_error(op) + if error: + self.err(error, op) name = op.name if name in _BINARY_SAME_TYPE: if len(op.operands) != 2 or len(op.results) != 1: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics.py new file mode 100644 index 000000000000..ae21ad10248a --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics.py @@ -0,0 +1,356 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Compare native XF32 input preparation with independent NumPy references.""" + +import argparse +import ctypes +import json +import struct +from pathlib import Path + +import numpy as np + +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, _resolve_llvm_flavor +from rocke.instances.gfx942.tf32_mma_probe import ( + PREPARATIONS, + Tf32MmaProbeSpec, + build_tf32_mma_probe, +) + + +def rne_bits(bits): + """Independent integer quotient/remainder reference; quiet NaNs.""" + u = np.asarray(bits, dtype=np.uint32) + sign = u & np.uint32(0x80000000) + magnitude = u & np.uint32(0x7FFFFFFF) + quotient = magnitude.astype(np.uint64) // 8192 + remainder = magnitude % 8192 + rounded = ( + quotient + ((remainder > 4096) | ((remainder == 4096) & (quotient % 2 == 1))) + ) * 8192 + finite = rounded.astype(np.uint32) | sign + special = u & np.uint32(0xFFFFE000) + special = np.where(magnitude > 0x7F800000, special | np.uint32(0x400000), special) + return np.where(magnitude < 0x7F800000, finite, special).astype(np.uint32) + + +def corpus(m): + k = 128 // m + # Exact values, ties (even/odd), either side, carry, normals, subnormals, + # infinities, signed zeros, signaling and quiet NaNs, both signs. + bits = np.array( + [ + 0, + 1, + 0x1000, + 0x7FFFFF, + 0x800000, + 0x801001, + 0x3F800000, + 0x3F800001, + 0x3F800FFF, + 0x3F801000, + 0x3F801001, + 0x3F802FFF, + 0x3F803000, + 0x3F803001, + 0x3FFFFFFF, + 0x7F7FFFFF, + 0x7F800000, + 0x7F800001, + 0x7FC00001, + ], + dtype=np.uint32, + ) + bits = np.concatenate((bits, bits | np.uint32(0x80000000))) + rng = np.random.default_rng(1714) + finite = rng.uniform(-8, 8, 1024).astype(np.float32).view(np.uint32) + bits = np.concatenate((bits, finite)) + single = len(bits) * 2 + a = np.zeros((single + 32, m, k), dtype=np.float32) + b = np.zeros_like(a) + c = np.zeros((len(a), m, m), dtype=np.float32) + a[: len(bits), :, 0] = bits.view(np.float32)[:, None] + b[: len(bits), :, 0] = 1 + a[len(bits) : single, :, 0] = 1 + b[len(bits) : single, :, 0] = bits.view(np.float32)[:, None] + a[single:] = rng.uniform(-2, 2, a[single:].shape) + b[single:] = rng.uniform(-2, 2, b[single:].shape) + c[single:] = rng.uniform(-1, 1, c[single:].shape) + # Asymmetric dense inputs and cancellation with a nonzero accumulator. + a[-1, :, 1::2] = -a[-1, :, ::2] + b[-1, :, 1::2] = b[-1, :, ::2] + return a, b, c, bits, single + + +def same_values(a, b): + # Compare every finite bit (including zero sign); NaN payloads are not an + # arithmetic-output contract. Prepared input payloads are checked separately. + return (a.view(np.uint32) == b.view(np.uint32)) | (np.isnan(a) & np.isnan(b)) + + +def launch(blob, name, a, b, c): + from rocke.runtime.hip_module import Runtime + + runtime = Runtime() + module = runtime.load_module(blob) + arrays = [ + a.copy(), + b.copy(), + c.copy(), + np.full_like(c, np.nan), + np.full(a.shape, 0xDEADBEEF, np.uint32), + np.full(b.shape, 0xDEADBEEF, np.uint32), + ] + ptrs = [] + try: + for array in arrays: + ptr = runtime.alloc(array.nbytes) + ptrs.append(ptr) + runtime.memcpy_h2d( + ptr, (ctypes.c_ubyte * array.nbytes).from_buffer(array), array.nbytes + ) + runtime.launch_blocking( + module.get_function(name), + (len(a), 1, 1), + (64, 1, 1), + struct.pack("6Q", *ptrs), + ) + for array, ptr in zip(arrays[3:], ptrs[3:]): + runtime.memcpy_d2h( + (ctypes.c_ubyte * array.nbytes).from_buffer(array), ptr, array.nbytes + ) + return arrays[3:] + finally: + for ptr in ptrs: + runtime.free(ptr) + module.unload() + + +def run(output_dir, backend="both", shapes=(16, 32)): + from rocke.runtime.comgr import build_hsaco_from_llvm_ir, resolved_lib_path + from rocke.runtime.hip_module import get_device_arch + + arch = get_device_arch() + if not arch or arch.split(":")[0] != "gfx942": + raise RuntimeError(f"TF32 numerical validation requires gfx942, got {arch}") + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + engines = ("python", "cpp") if backend == "both" else (backend,) + native = None + if "cpp" in engines: + import rocke_engine as native # Required: no fallback. + flavor = _resolve_llvm_flavor() + summary = { + "arch": arch, + "numpy": np.__version__, + "flavor": flavor, + "results": [], + } + for m in shapes: + a, b, c, bits, single = corpus(m) + np.savez(output_dir / f"inputs_{m}.npz", a=a, b=b, c=c) + all_results = {} + for engine in engines: + results = {} + for variant in (*PREPARATIONS, "masked"): + mode = "prepacked" if variant == "masked" else variant + spec = Tf32MmaProbeSpec(m, mode) + kernel = build_tf32_mma_probe(spec) + ir = serialize(kernel) + if engine == "cpp": + ir_native = native.tf32_mma_probe_serialize_ir(m, mode) + if ir != ir_native: + raise RuntimeError( + f"m={m} mode={mode}: builder parity mismatch" + ) + ll = native.lower_serialized_ir( + ir_native, + arch="gfx942", + flavor=flavor, + ) + if ll != _lower_kernel_to_llvm_python( + kernel, arch="gfx942", llvm_flavor=flavor + ): + raise RuntimeError( + f"m={m} mode={mode}: lowerer parity mismatch" + ) + else: + ll = _lower_kernel_to_llvm_python( + kernel, arch="gfx942", llvm_flavor=flavor + ) + stem = output_dir / f"{engine}_{m}_{variant}" + stem.with_suffix(".ll").write_text(ll) + stem.with_suffix(".hip").write_text( + lower_kernel_to_hip(kernel, arch="gfx942") + ) + blob, _ = build_hsaco_from_llvm_ir(ll, isa="amdgcn-amd-amdhsa--gfx942") + stem.with_suffix(".hsaco").write_bytes(blob) + aa, bb = a, b + if variant == "prepacked": + aa, bb = rne_bits(a.view(np.uint32)), rne_bits(b.view(np.uint32)) + elif variant == "masked": + aa, bb = a.view(np.uint32) & np.uint32(0xFFFFE000), b.view( + np.uint32 + ) & np.uint32(0xFFFFE000) + d, pa, pb = launch(blob, spec.name, aa, bb, c) + expected_a = ( + rne_bits(a.view(np.uint32)) + if variant in ("rne", "prepacked") + else aa.view(np.uint32) + ) + expected_b = ( + rne_bits(b.view(np.uint32)) + if variant in ("rne", "prepacked") + else bb.view(np.uint32) + ) + context = f"{engine} m={m} variant={variant}" + if not np.array_equal(pa, expected_a): + raise RuntimeError(f"{context}: A preparation mismatch") + if not np.array_equal(pb, expected_b): + raise RuntimeError(f"{context}: B preparation mismatch") + np.savez(stem.with_suffix(".npz"), d=d, pa=pa, pb=pb) + results[variant] = d + # Dense reference uses float64 scalar products and a conservative + # FP32 accumulation bound, proportional to sum(abs(products))+abs(C). + ra, rb = expected_a, expected_b + if variant in ("raw", "carrier", "masked"): + ra, rb = ra & np.uint32(0xFFFFE000), rb & np.uint32(0xFFFFE000) + da, db = ra.view(np.float32)[single:].astype(np.float64), rb.view( + np.float32 + )[single:].astype(np.float64) + ref = np.einsum("bik,bjk->bij", da, db) + c[single:].astype(np.float64) + scale = np.einsum("bik,bjk->bij", abs(da), abs(db)) + abs(c[single:]) + error = abs(d[single:].astype(np.float64) - ref) + if not np.all( + error <= 4 * (128 // m) * np.finfo(np.float32).eps * scale + 1e-7 + ): + raise RuntimeError( + f"{context}: dense reference mismatch (max error {float(error.max())})" + ) + # Exact single-product normal results avoid reduction-order noise. + input_bits = np.concatenate((bits, bits)) + expected = ( + input_bits + if variant == "fp32" + else ( + rne_bits(input_bits) + if variant in ("rne", "prepacked") + else input_bits & np.uint32(0xFFFFE000) + ) + ) + # +0 accumulation turns a negative-zero product into +0. + expected = np.where( + (expected & 0x7FFFFFFF) == 0, np.uint32(0), expected + ).astype(np.uint32) + if not np.all( + same_values( + d[:single], + np.broadcast_to( + expected.view(np.float32)[:, None, None], d[:single].shape + ), + ) + ): + raise RuntimeError(f"{context}: single product mismatch") + summary["results"].append( + { + "engine": engine, + "m": m, + "variant": variant, + "batches": len(a), + "prepared_words": int(pa.size + pb.size), + "dense_max_abs_error": float(error.max()), + "status": "pass", + } + ) + for left, right in ( + ("raw", "carrier"), + ("raw", "masked"), + ("rne", "prepacked"), + ): + if not np.all(same_values(results[left], results[right])): + raise RuntimeError( + f"{engine} m={m}: variant parity mismatch ({left}, {right})" + ) + # Preserve all boundary observations, including denormal and NaN behavior. + observations = [ + { + "input_bits": f"{int(u):08x}", + "rne_bits": f"{int(rne_bits(u)):08x}", + **{ + v: f"{int(d[i,0,0].view(np.uint32)):08x}" + for v, d in results.items() + }, + } + for i, u in enumerate(bits) + ] + for row in observations: + original = np.array(int(row["input_bits"], 16), dtype=np.uint32).view( + np.float32 + ) + row["input_value"] = str(original) + for variant in ("raw", "rne"): + actual_bits = int(row[variant], 16) + actual = np.array(actual_bits, dtype=np.uint32).view(np.float32) + if np.isfinite(original) and np.isfinite(actual): + row[f"{variant}_abs_error_vs_fp32"] = abs( + float(actual) - float(original) + ) + # Within one sign, IEEE encoding order gives ULP distance. + if (actual_bits ^ int(row["input_bits"], 16)) < 0x80000000: + row[f"{variant}_ulp_error_vs_fp32"] = abs( + actual_bits - int(row["input_bits"], 16) + ) + (output_dir / f"{engine}_{m}_boundaries.json").write_text( + json.dumps(observations, indent=2) + "\n" + ) + counterexample = next( + row + for row in observations + if row["raw"] != row["rne"] and row["input_bits"] == "3f801001" + ) + print(engine, m, json.dumps(counterexample), flush=True) + all_results[engine] = results + if backend == "both": + for variant in all_results["python"]: + if not np.all( + same_values( + all_results["python"][variant], all_results["cpp"][variant] + ) + ): + raise RuntimeError( + f"m={m} variant={variant}: backend parity mismatch (python, cpp)" + ) + summary["comgr"] = resolved_lib_path() + summary["status"] = "pass" + (output_dir / "results.json").write_text(json.dumps(summary, indent=2) + "\n") + return summary + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--backend", choices=("python", "cpp", "both"), default="both") + parser.add_argument("--shape", choices=("16", "32", "all"), default="all") + args = parser.parse_args() + args.output_dir.mkdir(parents=True, exist_ok=True) + result_path = args.output_dir / "results.json" + result_path.write_text(json.dumps({"status": "running"}) + "\n") + try: + result = run( + args.output_dir, + args.backend, + (16, 32) if args.shape == "all" else (int(args.shape),), + ) + except Exception as error: + result_path.write_text( + json.dumps({"status": "fail", "error": str(error)}) + "\n" + ) + raise + print(json.dumps({"status": result["status"], "variants": len(result["results"])})) + + +if __name__ == "__main__": + main() diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics/README.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics/README.md new file mode 100644 index 000000000000..389971de7609 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics/README.md @@ -0,0 +1,74 @@ +# Native TF32 numerics on gfx942 + +Run the [numerical example](../tf32_numerics.py) with NumPy and a built +`rocke_engine` on `PYTHONPATH`, using a ROCm runtime/compiler pair that supports +gfx942: + +```console +python -m rocke.examples.gfx942.tf32_numerics --backend both --shape all --output-dir results +``` + +The example resolves one LLVM flavor from the installed COMGR release for both +engines and records it in `results.json`. Override with `ROCKE_LLVM_FLAVOR` when +needed (`llvm20` for ROCm <7.2, `llvm22` for 7.2–7.12, `llvm23` for 7.13+). +The default runs both engines and both native shapes. Missing gfx942 hardware +or the native extension is an error. +The [spec-driven probe](../../../instances/gfx942/tf32_mma_probe.py) is also +available through the native builder and the paired test emitters. + +The example loads dense A and transposed B tiles, accumulates into FP32 C, and +saves D plus the actual prepared A/B bits. Each shape runs raw FP32, the same bits +loaded through I32, explicit GPU RNE preparation, independent host RNE preparation, +host masking, and ordinary FP32 MFMA over the same total K. The C++ lane builds +its own IR and lowers it natively; it also checks exact Python/C++ byte identity. + +The corpus includes both signs, even/odd ties and their neighbors, exponent +carry, subnormals, signed zero, infinities, signaling/quiet NaNs, deterministic +dense data, cancellation, asymmetric matrices, and nonzero C. Single products +are checked exactly (NaNs by classification, zero after positive-zero +accumulation). Dense results use float64 products with an explicit FP32 +accumulation bound based on the sum of absolute products. No BLAS TF32 policy +can influence the reference. + +The output directory contains inputs, prepared words and outputs (`.npz`), LLVM, +HIP source, HSACO, full boundary observations, and `results.json`. The program +exits nonzero on compilation, launch, preparation, parity, or numerical failure. +Boundary JSON preserves hexadecimal results for every variant; the console +prints a representative RNE counterexample. This is a correctness example, not +a performance benchmark. + +## Input preparation and logical type + +`TF32` is a distinct logical IR type with an **I32 carrier**, one 32-bit word per +element. The dtype names `tf32` and `xf32` select the same catalog format. +`bitcast(fp32_value, TF32)` preserves the raw payload; it does not round. +`cvt_f32_to_tf32(fp32_value)` explicitly rounds to nearest, ties to even, setting +the low 13 fraction bits to zero. It preserves infinities and signed zero and +quiets NaNs while retaining the upper payload bits. A bitcast back to F32 +reinterprets the stored bits exactly; it does not truncate a raw payload. + +On gfx942 with ROCm 7.14, both native XF32 shapes accepted raw FP32 payloads. +For the tested corpus their outputs matched masking the low 13 bits; I32 +transport was bit-preserving. RNE preparation changed rounding-boundary results. +For a single product with B=1 and C=0: + +| Path | Input bits | Output bits | +|---|---|---| +| Ordinary FP32 MFMA | `3f801001` | `3f801001` | +| Raw XF32 / I32 carrier | `3f801001` | `3f800000` | +| GPU / host RNE → XF32 | `3f802000` | `3f802000` | + +Thus preparation is unnecessary for **instruction acceptance**, but explicit +preparation is necessary when the caller requires **RNE inputs**. Raw payloads +retain native XF32 behavior, including special encodings: in the tested corpus +a signaling NaN whose payload occupied only discarded bits became infinity. +The RNE conversion prevents that by quieting NaNs before clearing low bits. +These observations describe this corpus and toolchain; they are not a guarantee +that raw payloads have general FP32 arithmetic semantics. + +The catalog exposes only gfx942 `16x16x8` and `32x32x4` TF32 atoms, with two +logical TF32 elements per lane in A/B and FP32 accumulation. The LLVM intrinsic +ABI takes two floats, so lowering bitcasts the I32 carrier at the intrinsic +boundary. Ordinary FP32 catalog selection is unchanged. Generic TF32 arithmetic +is rejected; bitcasts, vector transport, selection, and the native atoms are +supported. Production GEMM dispatch and cross-target emulation are separate work. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py index ebcfc0e21a42..dcb5898824d8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py @@ -25,6 +25,8 @@ def storage_ir_type(dtype: str) -> Type: """Addressable storage unit, distinct from the logical dtype resolver.""" info = dtype_info(dtype) + if info.name == "tf32": + return I32 if info.encoded_bits % 8 or info.name in ("e8m0", "e5m3"): return I8 return dtype_to_ir_type(info.name) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/gfx942/tf32_mma_probe.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/gfx942/tf32_mma_probe.py new file mode 100644 index 000000000000..3bce6406cc0c --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/gfx942/tf32_mma_probe.py @@ -0,0 +1,120 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""One wave per dense atom, with explicit FP32 input-preparation variants.""" + +from dataclasses import dataclass + +from rocke.core.arch import ArchTarget +from rocke.core.ir import F32, I32, TF32, IRBuilder, PtrType, VectorType +from rocke.core.storage import BitPacking, FragmentPacking, MatrixFragmentLayout +from rocke.helpers.mma_io import load_matrix_fragment + +PREPARATIONS = ("raw", "carrier", "rne", "prepacked", "fp32") + + +@dataclass(frozen=True) +class Tf32MmaProbeSpec: + m: int = 16 + preparation: str = "raw" + + def __post_init__(self): + if self.m not in (16, 32) or self.preparation not in PREPARATIONS: + raise ValueError("TF32 probe requires m=16/32 and a known preparation") + + @property + def name(self): + return f"tf32_probe_{self.m}_{self.preparation}" + + +def build_tf32_mma_probe(spec: Tf32MmaProbeSpec): + m, mode = spec.m, spec.preparation + k = 128 // m + target = ArchTarget.from_gfx("gfx942") + atom = target.mma.op_for_shape( + family="mma", + a_dtype="tf32", + b_dtype="tf32", + c_dtype="fp32", + m=m, + n=m, + k=k, + ) + b = IRBuilder(spec.name) + ty = I32 if mode == "prepacked" else F32 + a = b.param("A", PtrType(ty, "global"), align=4) + bb = b.param("B", PtrType(ty, "global"), align=4) + cc = b.param("C", PtrType(F32, "global"), align=4) + d = b.param("D", PtrType(F32, "global"), align=4) + pa = b.param("PA", PtrType(I32, "global"), align=4) + pb = b.param("PB", PtrType(I32, "global"), align=4) + lane = b.thread_id_x() + batch = b.block_id_x() + cm = b.const_i32(m) + axis = b.mod(lane, cm) + group = b.div(lane, cm) + ck = b.const_i32(k) + row = b.mul(axis, ck) + batch_input = b.mul(batch, b.const_i32(m * k)) + row_base = b.add(batch_input, row) + batch_output = b.mul(batch, b.const_i32(m * m)) + c_values, out_indices = [], [] + for slot in range(atom.c_frag_len): + r, col = atom.c_layout().coord(b, lane, slot) + index = b.add(batch_output, b.add(b.mul(r, cm), col)) + out_indices.append(index) + c_values.append(b.global_load(cc, index, F32, align=4)) + acc = b.vec_pack(c_values, F32) + if mode == "fp32": + full = target.mma.op_for_shape( + family="mma", + a_dtype="fp32", + b_dtype="fp32", + c_dtype="fp32", + m=m, + n=m, + k=k // 2, + ) + for step in range(2): + index = b.add(row_base, b.add(group, b.const_i32(step * (64 // m)))) + av = b.global_load(a, index, F32, align=4) + bv = b.global_load(bb, index, F32, align=4) + b.global_store(pa, index, b.bitcast(av, I32), align=4) + b.global_store(pb, index, b.bitcast(bv, I32), align=4) + acc = b.mma(full, av, bv, acc) + else: + layout = MatrixFragmentLayout( + FragmentPacking(BitPacking(32), 2, 32, 2), 2, 64 // m, m + ) + fragments = [] + for ptr, prepared in ((a, pa), (bb, pb)): + carrier = I32 if mode in ("carrier", "prepacked") else F32 + values = load_matrix_fragment( + b, + ptr, + row_base, + group, + 0, + dtype="tf32" if mode == "prepacked" else "fp32", + layout=layout, + carrier_type=carrier, + alignment_bytes=4, + ) + if mode == "rne": + values = b.vec_pack( + [b.cvt_f32_to_tf32(b.vec_extract(values, i)) for i in range(2)], + TF32, + ) + else: + values = b.bitcast(values, VectorType(TF32, 2)) + for slot in range(2): + index = b.add( + row_base, b.add(b.mul(group, b.const_i32(2)), b.const_i32(slot)) + ) + bits = b.bitcast(b.vec_extract(values, slot), I32) + b.global_store(prepared, index, bits, align=4) + fragments.append(values) + acc = b.mma(atom, *fragments, acc) + for slot, index in enumerate(out_indices): + b.global_store(d, index, b.vec_extract(acc, slot), align=4) + b.ret() + return b.kernel diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt b/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt index cedfb57338f6..87d32004f5d2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt @@ -14,6 +14,15 @@ if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/core/smoke.cpp") endif() if(BUILD_TESTING) + add_executable(rocke_tf32 core/test_tf32.cpp) + target_link_libraries(rocke_tf32 PRIVATE rocke_core) + add_test(NAME rocke_tf32 COMMAND rocke_tf32) + set_tests_properties(rocke_tf32 PROPERTIES LABELS "ckc;unit;host") + rocke_cxx_test_installed_name(_rocke_installed_name rocke_tf32) + set_target_properties(rocke_tf32 PROPERTIES OUTPUT_NAME "${_rocke_installed_name}") + install(TARGETS rocke_tf32 RUNTIME DESTINATION ${ROCKE_CXX_TEST_INSTALL_DIR}) + set_property(GLOBAL APPEND PROPERTY ROCKE_INSTALL_TEST_TARGETS rocke_tf32) + add_executable(rocke_storage core/test_storage.cpp) target_link_libraries(rocke_storage PRIVATE rocke_core) add_test(NAME rocke_storage COMMAND rocke_storage) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.cpp new file mode 100644 index 000000000000..33b64dd7b735 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.cpp @@ -0,0 +1,316 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +#include "rocke/error.hpp" +#include "rocke/instance_tf32_mma_probe.h" +#include "rocke/ir_serialize.h" +#include "rocke/lower_hip.h" +#include "rocke/lower_llvm.h" +#include "rocke/verify.h" +#include +#include +#include +#include + +#define CHECK(expr) \ + do \ + { \ + if(!(expr)) \ + { \ + std::fprintf(stderr, "line %d: %s\n", __LINE__, #expr); \ + return 1; \ + } \ + } while(0) + +static int test_invalid_tf32_ops_rejected() +{ + for(int variant = 0; variant < 8; ++variant) + { + rocke_ir_builder_t b, parsed; + CHECK(rocke_ir_builder_init(&b, "invalid_tf32_vector") == ROCKE_OK); + CHECK(rocke_ir_builder_init(&parsed, "parsed") == ROCKE_OK); + auto* elem = variant == 5 ? rocke_i32() : rocke_tf32(); + auto* v = rocke_b_param(&b, "v", rocke_vector_type(&b, elem, 2), nullptr); + switch(variant) + { + case 0: + rocke_b_vector_add(&b, v, v); + break; + case 1: + rocke_b_vector_mul(&b, v, v); + break; + case 2: + rocke_b_vector_sum(&b, v); + break; + case 3: + rocke_b_vector_reduce_max(&b, v); + break; + case 4: + rocke_b_vector_cmp(&b, "lt", v, v); + break; + case 5: + rocke_b_vector_trunc(&b, v, rocke_tf32()); + break; + } + if(variant >= 6) + { + rocke_ir_builder_free(&b); + auto* probe = rocke_build_tf32_mma_probe(&b, variant == 6 ? 16 : 32, "fp32"); + CHECK(probe); + for(int i = 0; i < probe->body->num_ops; ++i) + { + auto* op = probe->body->ops[i]; + if(std::strcmp(op->name, "tile.mma") == 0) + op->results[0]->type + = rocke_vector_type(&b, rocke_tf32(), variant == 6 ? 4 : 16); + } + } + else + rocke_b_ret(&b); + const char* message = variant >= 6 ? "MMA results must not use TF32" : "TF32 arithmetic"; + char* text = nullptr; + CHECK(rocke_ir_serialize(b.kernel, &text) == ROCKE_OK); + rocke_kernel_def_t* kernel = nullptr; + CHECK(rocke_ir_parse(text, &parsed, &kernel) == ROCKE_OK); + std::free(text); + rocke_diag_t* diagnostics = nullptr; + size_t count = 0; + CHECK(rocke_verify(kernel, &diagnostics, &count) == ROCKE_OK); + bool rejected = false; + for(size_t i = 0; i < count; ++i) + rejected |= std::strstr(diagnostics[i].message, message) != nullptr; + CHECK(rejected); + rocke_diags_free(diagnostics, count); + for(int f = 0; f < rocke_llvm_flavor_count(); ++f) + { + char* ll = nullptr; + auto flavor = rocke_llvm_flavor_from_name(rocke_llvm_flavor_at(f)); + CHECK(rocke_lower_kernel_to_llvm(kernel, flavor, "gfx942", &ll) == ROCKE_ERR_VALUE); + std::free(ll); + } + rocke_strbuf_t hip; + CHECK(rocke_strbuf_init(&hip, 0) == 0); + rocke_lower_hip_opts_t opts = {}; + opts.arch = "gfx942"; + CHECK(rocke_lower_kernel_to_hip(&parsed, kernel, &opts, &hip) == ROCKE_ERR_VALUE); + rocke_strbuf_free(&hip); + rocke_ir_builder_free(&parsed); + rocke_ir_builder_free(&b); + } + return 0; +} + +static int test_vector_load() +{ + for(int n : {0, 1, 2, 3, 4, 6, 8, 16}) + for(int alignment : {0, 4}) + { + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "tf32_vector_load") == ROCKE_OK); + auto* p = rocke_b_param(&b, "p", rocke_ptr_type(&b, rocke_tf32(), "global"), nullptr); + auto* index = rocke_b_const_i32(&b, 0); + if(n != 2 && n != 3 && n != 4 && n != 8) + { + try + { + rocke_b_global_load_vN(&b, p, index, rocke_tf32(), n, alignment); + CHECK(false); + } + catch(const ckc::Error& error) + { + CHECK(error.code() == ROCKE_ERR_VALUE); + CHECK(std::strstr(error.what(), "unsupported vector width for tf32")); + } + rocke_ir_builder_free(&b); + continue; + } + auto* values = rocke_b_global_load_vN(&b, p, index, rocke_tf32(), n, alignment); + CHECK(values); + rocke_b_global_store(&b, p, index, rocke_b_vec_extract(&b, values, n - 1), 4); + rocke_b_ret(&b); + int expected_align = alignment ? alignment : (n == 3 ? 4 : n * 4); + for(int f = 0; f < rocke_llvm_flavor_count(); ++f) + { + char* ll = nullptr; + auto flavor = rocke_llvm_flavor_from_name(rocke_llvm_flavor_at(f)); + CHECK(rocke_lower_kernel_to_llvm(b.kernel, flavor, "gfx942", &ll) == ROCKE_OK); + char load[64], align[32]; + std::snprintf(load, sizeof(load), "load <%d x i32>", n); + std::snprintf(align, sizeof(align), "align %d", expected_align); + CHECK(std::strstr(ll, load)); + CHECK(std::strstr(ll, align)); + std::free(ll); + } + rocke_strbuf_t hip; + CHECK(rocke_strbuf_init(&hip, 0) == 0); + rocke_lower_hip_opts_t opts = {}; + opts.arch = "gfx942"; + CHECK(rocke_lower_kernel_to_hip(&b, b.kernel, &opts, &hip) == ROCKE_OK); + if(expected_align < n * 4 || n == 3) + { + char size[32]; + std::snprintf(size, sizeof(size), "%d);", n * 4); + CHECK(std::strstr(hip.data, "__builtin_memcpy")); + CHECK(std::strstr(hip.data, size)); + } + rocke_strbuf_free(&hip); + rocke_ir_builder_free(&b); + } + return 0; +} + +static int test_invalid_store_alignment() +{ + for(int align : {0, -4, 3, 12}) + { + rocke_ir_builder_t b, parsed; + CHECK(rocke_ir_builder_init(&b, "invalid_store_alignment") == ROCKE_OK); + CHECK(rocke_ir_builder_init(&parsed, "parsed") == ROCKE_OK); + auto* p = rocke_b_param(&b, "p", rocke_ptr_type(&b, rocke_tf32(), "global"), nullptr); + auto* zero = rocke_b_const_i32(&b, 0); + auto* value = rocke_b_bitcast(&b, zero, rocke_tf32()); + rocke_value_t* lanes[] = {value, value, value, value}; + auto* values = rocke_b_vec_pack(&b, lanes, 4, rocke_tf32()); + rocke_b_global_store_vN(&b, p, zero, values, 4, align); + auto* body = b.kernel->body; + rocke_attr_set_int(&b, &body->ops[body->num_ops - 1]->attrs, "align", align); + rocke_b_ret(&b); + char* text = nullptr; + CHECK(rocke_ir_serialize(b.kernel, &text) == ROCKE_OK); + rocke_kernel_def_t* kernel = nullptr; + CHECK(rocke_ir_parse(text, &parsed, &kernel) == ROCKE_OK); + std::free(text); + for(int f = 0; f < rocke_llvm_flavor_count(); ++f) + { + char* ll = nullptr; + auto flavor = rocke_llvm_flavor_from_name(rocke_llvm_flavor_at(f)); + CHECK(rocke_lower_kernel_to_llvm(kernel, flavor, "gfx942", &ll) == ROCKE_ERR_VALUE); + std::free(ll); + } + rocke_strbuf_t hip; + CHECK(rocke_strbuf_init(&hip, 0) == 0); + rocke_lower_hip_opts_t opts = {}; + opts.arch = "gfx942"; + CHECK(rocke_lower_kernel_to_hip(&parsed, kernel, &opts, &hip) == ROCKE_ERR_VALUE); + rocke_strbuf_free(&hip); + rocke_ir_builder_free(&parsed); + rocke_ir_builder_free(&b); + } + return 0; +} + +int main() +{ + CHECK(test_invalid_store_alignment() == 0); + CHECK(test_invalid_tf32_ops_rejected() == 0); + CHECK(test_vector_load() == 0); + CHECK(!rocke_type_eq(rocke_tf32(), rocke_i32())); + CHECK(!rocke_type_eq(rocke_tf32(), rocke_f32())); + CHECK(rocke_scalar_by_name("tf32") == rocke_tf32()); + const char* modes[] = {"raw", "carrier", "rne", "prepacked", "fp32"}; + for(int m : {16, 32}) + { + for(const char* mode : modes) + { + rocke_ir_builder_t b; + auto* kernel = rocke_build_tf32_mma_probe(&b, m, mode); + CHECK(kernel); + rocke_diag_t* diagnostics = nullptr; + size_t count = 0; + rocke_verify(kernel, &diagnostics, &count); + CHECK(count == 0); + rocke_diags_free(diagnostics, count); + rocke_lower_hip_opts_t opts = {}; + opts.arch = "gfx942"; + rocke_strbuf_t hip; + CHECK(rocke_strbuf_init(&hip, 0) == 0); + CHECK(rocke_lower_kernel_to_hip(&b, kernel, &opts, &hip) == ROCKE_OK); + CHECK(std::strstr(rocke_strbuf_cstr(&hip), "__builtin_amdgcn_mfma_f32_")); + rocke_strbuf_free(&hip); + if(std::strcmp(mode, "fp32") != 0) + { + char* ll = nullptr; + CHECK(rocke_lower_kernel_to_llvm(kernel, ROCKE_LLVM_FLAVOR_AUTO, "gfx950", &ll) + != ROCKE_OK); + std::free(ll); + } + rocke_ir_builder_free(&b); + } + } + // Mirror test_global_vector_store: native authoring must accept the same widths. + for(int n : {1, 2, 4, 8, 16}) + for(int alignment : {0, 4, 64}) + { + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "tf32_vector_store") == ROCKE_OK); + auto* p = rocke_b_param(&b, "p", rocke_ptr_type(&b, rocke_tf32(), "global"), nullptr); + auto* index = rocke_b_const_i32(&b, 0); + auto* value = rocke_b_bitcast(&b, index, rocke_tf32()); + rocke_value_t* components[16]; + for(int i = 0; i < n; ++i) + components[i] = value; + auto* values = rocke_b_vec_pack(&b, components, n, rocke_tf32()); + if(n == 16) + { + try + { + rocke_b_global_store_vN(&b, p, index, values, n, alignment); + CHECK(false); + } + catch(const ckc::Error& error) + { + CHECK(error.code() == ROCKE_ERR_VALUE); + CHECK(std::strstr(error.what(), "n=16 not supported for tf32")); + } + } + else + { + rocke_b_global_store_vN(&b, p, index, values, n, alignment); + rocke_b_ret(&b); + for(int f = 0; f < rocke_llvm_flavor_count(); ++f) + { + char* ll = nullptr; + auto flavor = rocke_llvm_flavor_from_name(rocke_llvm_flavor_at(f)); + CHECK(rocke_lower_kernel_to_llvm( + rocke_ir_builder_kernel(&b), flavor, "gfx950", &ll) + == ROCKE_OK); + char store[64], align[32]; + std::snprintf(store, sizeof(store), "store <%d x i32>", n); + std::snprintf(align, sizeof(align), "align %d", alignment ? alignment : n * 4); + CHECK(std::strstr(ll, store)); + CHECK(std::strstr(ll, align)); + std::free(ll); + } + rocke_strbuf_t hip; + CHECK(rocke_strbuf_init(&hip, 0) == 0); + rocke_lower_hip_opts_t opts = {}; + opts.arch = "gfx942"; + CHECK(rocke_lower_kernel_to_hip(&b, b.kernel, &opts, &hip) == ROCKE_OK); + if(alignment && alignment < n * 4) + { + char size[32]; + std::snprintf(size, sizeof(size), ", %d);", n * 4); + CHECK(std::strstr(hip.data, "__builtin_memcpy(__builtin_assume_aligned(")); + CHECK(std::strstr(hip.data, size)); + CHECK(!std::strstr(hip.data, "*reinterpret_cast<")); + } + else + CHECK(std::strstr(hip.data, "*reinterpret_cast<")); + rocke_strbuf_free(&hip); + } + rocke_ir_builder_free(&b); + } + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "invalid") == ROCKE_OK); + auto* integer = rocke_b_const_i32(&b, 1); + try + { + rocke_b_cvt_f32_to_tf32(&b, integer); + CHECK(false); + } + catch(const ckc::Error& error) + { + CHECK(error.code() == ROCKE_ERR_VALUE); + } + rocke_ir_builder_free(&b); + return 0; +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py new file mode 100644 index 000000000000..07cdf2eb1bf4 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py @@ -0,0 +1,530 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Logical TF32 contracts and strict native parity, independent of GPU access.""" +import numpy as np +import pytest + +from rocke.core.arch import ArchTarget +from rocke.core.dtypes import dtype_info, normalize_dtype +from rocke.core.ir import ( + F32, + I1, + I32, + TF32, + IRBuilder, + PtrType, + VectorType, + dtype_to_ir_type, +) +from rocke.core.ir_serialize import serialize, parse +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.examples.gfx942.tf32_numerics import rne_bits +from rocke.instances.gfx942.tf32_mma_probe import ( + PREPARATIONS, + Tf32MmaProbeSpec, + build_tf32_mma_probe, +) + + +def test_public_tf32_exports(): + import rocke + import rocke.core + from rocke import TF32 as public_tf32 + from rocke.core import TF32 as core_tf32 + + assert public_tf32 is core_tf32 is TF32 + assert "TF32" in rocke.__all__ + assert "TF32" in rocke.core.__all__ + + +def test_dtype_identity(): + assert TF32 != F32 and TF32 != I32 + for alias in ("tf32", "xf32", " XF32 "): + assert normalize_dtype(alias) == "tf32" + assert dtype_to_ir_type(alias) == TF32 + assert dtype_info(alias).encoded_bits == 32 + + +@pytest.mark.parametrize("m,k,acc", [(16, 8, 4), (32, 4, 16)]) +def test_catalog_isolation(m, k, acc): + catalog = ArchTarget.from_gfx("gfx942").mma + atom = catalog.op_for_shape( + family="mma", a_dtype="xf32", b_dtype="tf32", c_dtype="fp32", m=m, n=m, k=k + ) + assert (atom.a_frag_len, atom.b_frag_len, atom.c_frag_len) == (2, 2, acc) + assert atom.op_id == f"mfma_f32_{m}x{m}x{k}_xf32" + full = catalog.op_for_shape( + family="mma", a_dtype="fp32", b_dtype="fp32", c_dtype="fp32", m=m, n=m, k=k // 2 + ) + assert full.op_id == f"mfma_f32_{m}x{m}x{k//2}_f32" + assert not catalog.has_shape( + family="mma", a_dtype="fp32", b_dtype="fp32", c_dtype="fp32", m=m, n=m, k=k + ) + for gfx in ("gfx950", "gfx1151", "gfx1250"): + assert not any(op.a_dtype == "tf32" for op in ArchTarget.from_gfx(gfx).mma.ops) + + class Integers: + const_i32 = staticmethod(int) + mod = staticmethod(lambda x, y: x % y) + div = staticmethod(lambda x, y: x // y) + mul = staticmethod(lambda x, y: x * y) + add = staticmethod(lambda x, y: x + y) + + calc = Integers() + for layout, shape in ( + (atom.a_layout(), (m, k)), + (atom.b_layout(), (k, m)), + (atom.c_layout(), (m, m)), + ): + count = acc if shape == (m, m) else 2 + coords = [ + layout.coord(calc, lane, slot) + for lane in range(64) + for slot in range(count) + ] + assert len(set(coords)) == shape[0] * shape[1] == len(coords) + assert set(coords) == {(i, j) for i in range(shape[0]) for j in range(shape[1])} + + +@pytest.mark.parametrize("m", [16, 32]) +@pytest.mark.parametrize("mode", PREPARATIONS) +def test_both_builders_and_lowerers(m, mode): + native = pytest.importorskip("rocke_engine") + kernel = build_tf32_mma_probe(Tf32MmaProbeSpec(m, mode)) + ir = serialize(kernel) + assert serialize(parse(ir)) == ir + assert native.tf32_mma_probe_serialize_ir(m, mode) == ir + for flavor in ("llvm20", "llvm22", "llvm23"): + ll = _lower_kernel_to_llvm_python(kernel, arch="gfx942", llvm_flavor=flavor) + assert native.lower_serialized_ir(ir, arch="gfx942", flavor=flavor) == ll + if mode != "fp32": + assert f"llvm.amdgcn.mfma.f32.{m}x{m}x{128//m}.xf32" in ll + assert "sitofp" not in ll and "uitofp" not in ll + if mode != "fp32": + assert "__builtin_bit_cast(f32x2," in lower_kernel_to_hip(kernel, arch="gfx942") + + +@pytest.mark.parametrize("dtype,n", [(F32, 2), (I32, 2), (TF32, 1), (TF32, 4)]) +def test_mma_requires_logical_tf32(dtype, n): + b = IRBuilder("bad") + a = b.param("a", VectorType(dtype, n)) + c = b.param("c", VectorType(F32, 4)) + with pytest.raises(ValueError, match="XF32 MMA requires"): + b.mma("mfma_f32_16x16x8_xf32", a, a, c) + + +@pytest.mark.parametrize("gfx", ["gfx950", "gfx1151", "gfx1250"]) +def test_rejects_unsupported_target(gfx): + kernel = build_tf32_mma_probe(Tf32MmaProbeSpec()) + with pytest.raises(ValueError, match="gfx942"): + _lower_kernel_to_llvm_python(kernel, arch=gfx) + with pytest.raises(ValueError, match="gfx942"): + lower_kernel_to_hip(kernel, arch=gfx) + native = pytest.importorskip("rocke_engine") + with pytest.raises(Exception): + native.lower_serialized_ir(serialize(kernel), arch=gfx) + + +def test_arithmetic_rejected_after_serialization(): + b = IRBuilder("bad_arithmetic") + v = b.param("x", TF32) + b.add(v, v) + b.ret() + with pytest.raises(ValueError, match="TF32 arithmetic"): + _lower_kernel_to_llvm_python(b.kernel) + native = pytest.importorskip("rocke_engine") + with pytest.raises(Exception): + native.lower_serialized_ir(serialize(b.kernel)) + + +@pytest.mark.parametrize( + "operation", + [ + "add", + "sub", + "mul", + "and", + "or", + "shl", + "lshr", + "smax", + "smin", + "max", + "fma", + "sum", + "reduce_max", + "cmp", + "trunc", + "sext", + "trunc_result", + ], +) +def test_vector_arithmetic_rejected_after_serialization(operation): + b = IRBuilder("invalid_tf32_vector_arithmetic") + elem = I32 if operation == "trunc_result" else TF32 + v = b.param("v", VectorType(elem, 2)) + if operation == "fma": + b.vector_fma(v, v, v) + elif operation in ("sum", "reduce_max"): + getattr(b, f"vector_{operation}")(v) + elif operation == "cmp": + b.vector_cmp("lt", v, v) + elif operation in ("trunc", "sext", "trunc_result"): + if operation == "trunc_result": + b.vector_trunc(v, TF32) + else: + getattr(b, f"vector_{operation}")(v, I32) + else: + getattr(b, f"vector_{operation}")(v, v) + b.ret() + ir = serialize(b.kernel) + kernel = parse(ir) + message = "TF32 arithmetic requires an explicit conversion to f32" + assert any(message in str(d) for d in verify(kernel)) + with pytest.raises(ValueError, match=message): + _lower_kernel_to_llvm_python(kernel) + with pytest.raises(ValueError, match=message): + lower_kernel_to_hip(kernel) + native = pytest.importorskip("rocke_engine") + with pytest.raises(Exception, match=message): + native.lower_serialized_ir(ir) + + +def test_vector_payload_operations_remain_supported(): + b = IRBuilder("tf32_vector_payload") + p = b.param("p", PtrType(TF32, "global")) + mask = b.param("mask", I1) + zero = b.const_i32(0) + v = b.global_load(p, zero, TF32, align=4) + pair = b.vec_pack([v, v], TF32) + pair = b.vec_insert(pair, b.vec_extract(pair, 0), 1) + pair = b.vec_bitcast(b.vec_bitcast(pair, VectorType(I32, 2)), VectorType(TF32, 2)) + pair = b.vector_select(mask, pair, b.vector_splat(v, 2)) + pair = b.select(mask, pair, pair) + b.global_store_vN(p, zero, b.vec_concat(pair, pair), 4) + b.ret() + ir = serialize(b.kernel) + kernel = parse(ir) + assert verify(kernel) == [] + assert "i32x4" in lower_kernel_to_hip(kernel) + native = pytest.importorskip("rocke_engine") + for flavor in ("llvm20", "llvm22", "llvm23"): + ll = _lower_kernel_to_llvm_python(kernel, llvm_flavor=flavor) + assert native.lower_serialized_ir(ir, flavor=flavor) == ll + + +@pytest.mark.parametrize("n", [2, 3, 4, 8]) +@pytest.mark.parametrize("align", [None, 4]) +def test_global_vector_load(n, align): + b = IRBuilder("tf32_vector_load") + p = b.param("p", PtrType(TF32, "global")) + zero = b.const_i32(0) + values = b.global_load_vN(p, zero, TF32, n, align=align) + b.global_store(p, zero, b.vec_extract(values, n - 1), align=4) + b.ret() + ir = serialize(b.kernel) + kernel = parse(ir) + assert verify(kernel) == [] + expected_align = align or (4 if n == 3 else n * 4) + hip = lower_kernel_to_hip(kernel) + if expected_align < n * 4 or n == 3: + assert f"{n * 4});" in hip + assert "__builtin_memcpy" in hip + native = pytest.importorskip("rocke_engine") + for flavor in ("llvm20", "llvm22", "llvm23"): + ll = _lower_kernel_to_llvm_python(kernel, llvm_flavor=flavor) + assert f"load <{n} x i32>" in ll + assert f"align {expected_align}" in ll + assert native.lower_serialized_ir(ir, flavor=flavor) == ll + + +@pytest.mark.parametrize("n", [0, 1, 6, 16]) +def test_global_vector_load_rejects_unsupported_width(n): + b = IRBuilder("invalid_tf32_vector_load") + p = b.param("p", PtrType(TF32, "global")) + with pytest.raises(ValueError, match="unsupported vector width for tf32"): + b.global_load_vN(p, b.const_i32(0), TF32, n) + + +def test_transport(): + b = IRBuilder("transport") + p = b.param("p", PtrType(TF32, "global")) + i = b.const_i32(0) + v = b.global_load(p, i, TF32, align=4) + pair = b.vec_pack([v, v], TF32) + v = b.vec_extract(pair, 1) + b.global_store(p, i, v, align=4) + b.ret() + ir = serialize(b.kernel) + assert serialize(parse(ir)) == ir + ll = _lower_kernel_to_llvm_python(b.kernel) + assert "load i32" in ll and "store i32" in ll + native = pytest.importorskip("rocke_engine") + assert native.lower_serialized_ir(ir) == ll + + +def test_rne_reference_edges(): + inputs = np.array( + [ + 0, + 0x80000000, + 0x3F801000, + 0x3F803000, + 0x3F801001, + 0xBF801001, + 0x7F7FFFFF, + 0xFF7FFFFF, + 0x7F800000, + 0xFF800000, + 0x7F800001, + 0xFF800001, + 0x1000, + 0x1001, + 0x7FFFFF, + ], + np.uint32, + ) + expected = np.array( + [ + 0, + 0x80000000, + 0x3F800000, + 0x3F804000, + 0x3F802000, + 0xBF802000, + 0x7F800000, + 0xFF800000, + 0x7F800000, + 0xFF800000, + 0x7FC00000, + 0xFFC00000, + 0, + 0x2000, + 0x800000, + ], + np.uint32, + ) + assert np.array_equal(rne_bits(inputs), expected) + assert np.array_equal(rne_bits(expected), expected) + + +@pytest.mark.parametrize("n", [1, 2, 4]) +def test_lds_and_loop_transport(n): + b = IRBuilder("tf32_lds_loop") + p = b.param("p", PtrType(TF32, "global")) + zero = b.const_i32(0) + one = b.const_i32(1) + v = b.global_load(p, zero, TF32, align=4) + pair = b.vec_pack([v] * n, TF32) + shared = b.smem_alloc(TF32, (64,)) + b.smem_store_vN(shared, [zero], v if n == 1 else pair, n) + loaded = b.smem_load_vN(shared, zero, dtype=TF32, n=n) + loop = b.scf_for_iter(zero, one, one, [("payload", loaded)]) + with loop as (_, [payload]): + b.scf_yield(payload) + b.global_store(p, zero, b.vec_extract(loop.results[0], 0), align=4) + b.ret() + ir = serialize(b.kernel) + assert serialize(parse(ir)) == ir + ll = _lower_kernel_to_llvm_python(b.kernel) + assert "[256 x i8]" in ll + assert f"phi <{n} x i32>" in ll + assert "int smem" in lower_kernel_to_hip( + b.kernel + ) or "int shared" in lower_kernel_to_hip(b.kernel) + native = pytest.importorskip("rocke_engine") + assert native.lower_serialized_ir(ir) == ll + + +def test_fp32_atom_rejects_tf32_carrier(): + b = IRBuilder("wrong_atom") + a = b.param("a", TF32) + c = b.param("c", VectorType(F32, 4)) + with pytest.raises(ValueError, match="TF32 operands require"): + b.mma("mfma_f32_16x16x4_f32", a, a, c) + + +def test_parsed_mma_rejects_wrong_operands(): + native = pytest.importorskip("rocke_engine") + for bad in ("vec", "vec", "vec"): + ir = serialize(build_tf32_mma_probe(Tf32MmaProbeSpec())).replace( + "vec", bad + ) + with pytest.raises(ValueError, match="XF32 MMA requires"): + _lower_kernel_to_llvm_python(parse(ir), arch="gfx942") + with pytest.raises(Exception): + native.lower_serialized_ir(ir, arch="gfx942") + + +@pytest.mark.parametrize("m", [16, 32]) +@pytest.mark.parametrize("mode", PREPARATIONS) +def test_recipe_replay_preserves_ir(m, mode, monkeypatch): + from rocke.portable_ir.src.recording_builder import record_kernel + from rocke.portable_ir.src import online, recipe_bundle + + spec = Tf32MmaProbeSpec(m, mode) + recorded, recipe = record_kernel(lambda: build_tf32_mma_probe(spec)) + assert serialize(recorded) == serialize(build_tf32_mma_probe(spec)) + for flavor in ("llvm20", "llvm22", "llvm23"): + monkeypatch.setenv("ROCKE_LLVM_FLAVOR", flavor) + replayed, _ = online.recipe_cbor_to_llvm( + recipe_bundle.cbor_encode(recipe), arch="gfx942" + ) + assert replayed == _lower_kernel_to_llvm_python( + recorded, arch="gfx942", llvm_flavor=flavor + ) + + +@pytest.mark.parametrize("n", [1, 2, 4, 8, 16]) +@pytest.mark.parametrize("align", [None, 4, 64, "legacy"]) +@pytest.mark.parametrize("dtype", [TF32, I32, F32]) +def test_global_vector_store(n, align, dtype): + b = IRBuilder("vector_store_alignment") + p = b.param("p", PtrType(dtype, "global")) + i = b.const_i32(0) + v = b.global_load(p, i, dtype, align=4) + values = b.vec_pack([v] * n, dtype) + if n == 16: + with pytest.raises(ValueError, match="n=16 not supported"): + b.global_store_vN(p, i, values, n, align=align) + return + # Legacy serialized stores without an align attribute retain their default. + legacy = align == "legacy" + align = None if legacy else align + # The guarantee applies to the accessed address, including the index. + index = b.const_i32(1 if align == 4 else 0) + b.global_store_vN(p, index, values, n, align=align) + if legacy: + b.kernel.body.ops[-1].attrs.pop("align") + b.ret() + ir = serialize(b.kernel) + kernel = parse(ir) + assert serialize(kernel) == ir + assert verify(kernel) == [] + expected_align = align or n * 4 + hip = lower_kernel_to_hip(kernel) + if expected_align < n * 4: + assert ( + f"__builtin_assume_aligned(p + {index.name.lstrip('%')}, {expected_align})" + in hip + ) + assert f"&{values.name.lstrip('%')}, {n * 4});" in hip + else: + assert "*reinterpret_cast<" in hip + native = pytest.importorskip("rocke_engine") + for flavor in ("llvm20", "llvm22", "llvm23"): + ll = _lower_kernel_to_llvm_python(kernel, llvm_flavor=flavor) + store = next(line for line in ll.splitlines() if "store <" in line) + assert f"store <{n} x {'float' if dtype == F32 else 'i32'}>" in store + assert store.endswith(f"align {expected_align}") + assert native.lower_serialized_ir(ir, flavor=flavor) == ll + + +@pytest.mark.parametrize("align", [0, -4, 3, 12]) +def test_global_vector_store_rejects_invalid_alignment(align): + b = IRBuilder("invalid_store_alignment") + p = b.param("p", PtrType(TF32, "global")) + zero = b.const_i32(0) + value = b.bitcast(zero, TF32) + b.global_store_vN(p, zero, b.vec_pack([value] * 4, TF32), 4, align=align) + # The builder treats zero as its default; serialized IR can contain it literally. + b.kernel.body.ops[-1].attrs["align"] = align + b.ret() + ir = serialize(b.kernel) + kernel = parse(ir) + message = "global_store_vN: alignment must be a positive power of two" + with pytest.raises(ValueError, match=message): + lower_kernel_to_hip(kernel) + native = pytest.importorskip("rocke_engine") + for flavor in ("llvm20", "llvm22", "llvm23"): + with pytest.raises(ValueError, match=message): + _lower_kernel_to_llvm_python(kernel, llvm_flavor=flavor) + with pytest.raises(Exception, match=message): + native.lower_serialized_ir(ir, flavor=flavor) + + +@pytest.mark.parametrize("m", [16, 32]) +@pytest.mark.parametrize("scalar_result", [False, True]) +def test_non_xf32_mma_rejects_tf32_result(m, scalar_result): + b = IRBuilder("non_xf32_tf32_result") + out = b.param("out", PtrType(TF32, "global")) + a = b.param("a", F32) + c = b.param("c", VectorType(F32, 4 if m == 16 else 16)) + atom = ArchTarget.from_gfx("gfx942").mma.op_for_shape( + family="mma", + a_dtype="fp32", + b_dtype="fp32", + c_dtype="fp32", + m=m, + n=m, + k=64 // m, + ) + result = b.mma(atom, a, a, c) + # Model malformed serialized IR while keeping all downstream types consistent. + result.type = TF32 if scalar_result else VectorType(TF32, 4 if m == 16 else 16) + value = result if scalar_result else b.vec_extract(result, 0) + b.global_store(out, b.const_i32(0), value, align=4) + b.ret() + ir = serialize(b.kernel) + kernel = parse(ir) + message = "MMA results must not use TF32" + assert any(message in str(d) for d in verify(kernel)) + with pytest.raises(ValueError, match=message): + lower_kernel_to_hip(kernel, arch="gfx942") + native = pytest.importorskip("rocke_engine") + for flavor in ("llvm20", "llvm22", "llvm23"): + with pytest.raises(ValueError, match=message): + _lower_kernel_to_llvm_python(kernel, arch="gfx942", llvm_flavor=flavor) + with pytest.raises(Exception, match=message): + native.lower_serialized_ir(ir, arch="gfx942", flavor=flavor) + + +@pytest.mark.parametrize("m", [16, 32]) +@pytest.mark.parametrize("role", ["accumulator", "result"]) +@pytest.mark.parametrize("bad", ["integer", "tf32", "width"]) +def test_parsed_mma_rejects_wrong_accumulator_and_result(m, role, bad): + native = pytest.importorskip("rocke_engine") + kernel = build_tf32_mma_probe(Tf32MmaProbeSpec(m)) + mma = next(op for op in kernel.body.ops if op.name == "tile.mma") + value = mma.operands[2] if role == "accumulator" else mma.result + count = 4 if m == 16 else 16 + value.type = VectorType( + {"integer": I32, "tf32": TF32, "width": F32}[bad], + count + 1 if bad == "width" else count, + ) + ir = serialize(kernel) + message = "XF32 MMA requires" if role == "accumulator" else "XF32 MMA result" + with pytest.raises(ValueError, match=message): + _lower_kernel_to_llvm_python(parse(ir), arch="gfx942") + with pytest.raises(Exception, match=message): + native.lower_serialized_ir(ir, arch="gfx942") + + +@pytest.mark.parametrize("flavor", ["llvm20", "llvm22", "llvm23"]) +def test_numeric_harness_auto_flavor(flavor, monkeypatch, tmp_path): + pytest.importorskip("rocke_engine") + from rocke.examples.gfx942 import tf32_numerics + from rocke.runtime import comgr, hip_module + + monkeypatch.delenv("ROCKE_LLVM_FLAVOR", raising=False) + monkeypatch.setattr(tf32_numerics, "_resolve_llvm_flavor", lambda: flavor) + monkeypatch.setattr(hip_module, "get_device_arch", lambda: "gfx942") + expected = _lower_kernel_to_llvm_python( + build_tf32_mma_probe(Tf32MmaProbeSpec()), + arch="gfx942", + llvm_flavor=flavor, + ) + + class ReachedCompiler(Exception): + pass + + def check_compiler_input(ll, *, isa): + assert ll == expected + raise ReachedCompiler + + # Exercise the actual builders/lowerers, then stop before GPU compilation. + monkeypatch.setattr(comgr, "build_hsaco_from_llvm_ir", check_compiler_input) + with pytest.raises(ReachedCompiler): + tf32_numerics.run(tmp_path, backend="cpp", shapes=(16,)) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32_numerics_validation.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32_numerics_validation.py new file mode 100644 index 000000000000..aac9de10985b --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32_numerics_validation.py @@ -0,0 +1,132 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Host-only CLI checks with injected outputs; no GPU correctness evidence.""" + +import json +import os +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import numpy as np +import pytest +from rocke.examples.gfx942 import tf32_numerics as numerics +from rocke.runtime import comgr, hip_module + + +def _run_mocked_cli(fault: str, output_dir: str) -> None: + launches = 0 + + def corpus(m): + bits = np.array([0x3F801001], dtype=np.uint32) + a = np.zeros((3, m, 128 // m), dtype=np.float32) + b = np.zeros_like(a) + a[:, :, 0] = bits.view(np.float32)[0] + b[:, :, 0] = 1 + c = np.zeros((3, m, m), dtype=np.float32) + return a, b, c, bits, 2 + + def launch(blob, name, a, b, c): + nonlocal launches + launches += 1 + mode = name.rsplit("_", 1)[1] + pa, pb = a.view(np.uint32).copy(), b.view(np.uint32).copy() + if mode == "rne": + pa, pb = numerics.rne_bits(pa), numerics.rne_bits(pb) + ra, rb = pa, pb + if mode != "fp32": + ra, rb = pa & np.uint32(0xFFFFE000), pb & np.uint32(0xFFFFE000) + d = np.einsum("bik,bjk->bij", ra.view(np.float32), rb.view(np.float32)) + c + if fault == "A preparation": + pa.flat[0] ^= np.uint32(1) + elif fault == "B preparation": + pb.flat[0] ^= np.uint32(1) + elif fault == "dense reference": + d[-1, 0, 0] += 1 + elif fault == "single product": + d[0, 0, 0] += 1 + elif (fault == "variant parity" and mode == "carrier") or ( + fault == "backend parity" and launches > 6 + ): + # Stay within the dense-reference tolerance while breaking exact parity. + d[-1, 0, 0] = np.nextafter(d[-1, 0, 0], np.float32(np.inf)) + return d, pa, pb + + def native_ir(m, mode): + kernel = numerics.build_tf32_mma_probe(numerics.Tf32MmaProbeSpec(m, mode)) + return "wrong IR" if fault == "builder parity" else numerics.serialize(kernel) + + native = SimpleNamespace( + tf32_mma_probe_serialize_ir=native_ir, + lower_serialized_ir=lambda *a, **k: ( + "wrong LLVM" if fault == "lowerer parity" else "mock LLVM" + ), + ) + argv = ["tf32_numerics", "--output-dir", output_dir, "--shape", "16"] + with ( + patch.object(numerics, "corpus", corpus), + patch.object(numerics, "launch", launch), + patch.object(numerics, "_resolve_llvm_flavor", return_value="llvm23"), + patch.object( + numerics, "_lower_kernel_to_llvm_python", return_value="mock LLVM" + ), + patch.object(numerics, "lower_kernel_to_hip", return_value="mock HIP"), + patch.object(hip_module, "get_device_arch", return_value="gfx942"), + patch.object(comgr, "build_hsaco_from_llvm_ir", return_value=(b"mock", None)), + patch.object(comgr, "resolved_lib_path", return_value="mock COMGR"), + patch.dict(sys.modules, {"rocke_engine": native}), + patch.object(sys, "argv", argv), + ): + numerics.main() + + +@pytest.mark.parametrize("optimization", ["normal", "flag", "environment"]) +@pytest.mark.parametrize( + "fault", + [ + "none", + "A preparation", + "B preparation", + "dense reference", + "single product", + "builder parity", + "lowerer parity", + "variant parity", + "backend parity", + ], +) +def test_cli_validation_survives_optimization(tmp_path, optimization, fault): + env = os.environ.copy() + env.pop("PYTHONOPTIMIZE", None) + env["PYTHONPATH"] = os.pathsep.join(sys.path) + env["ROCKE_BACKEND"] = "python" + if optimization == "environment": + env["PYTHONOPTIMIZE"] = "1" + command = [sys.executable] + if optimization == "flag": + command.append("-O") + command += [ + "-c", + ( + "import runpy, sys; " + "runpy.run_path(sys.argv[1])['_run_mocked_cli'](sys.argv[2], sys.argv[3])" + ), + str(Path(__file__).resolve()), + fault, + str(tmp_path), + ] + result = subprocess.run( + command, env=env, capture_output=True, text=True, timeout=60, check=False + ) + report = json.loads((tmp_path / "results.json").read_text()) + if fault == "none": + assert result.returncode == 0, result.stderr + assert report["status"] == "pass" + assert len(report["results"]) == 12 + else: + assert result.returncode != 0, result.stdout + assert report["status"] == "fail" + assert fault in report["error"] + assert fault in result.stderr diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json index 265655ff4be6..454e8df24747 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json @@ -1583,6 +1583,18 @@ "14": "9997bc5e87821bf254534c0b7e2c55b928c0d2420828e52ca09b46a149c0dc93", "15": "17efbd8d97270032cb45822ab296ee7b546323e0158f32bde502c9592b2a7259", "16": "eed1a46189fb8296576cd4e1f840c8f8c9c137f4aadf2193331bc74dd92d2d51" + }, + "tf32_mma": { + "0": "5c55f4ced229ea0f5e87fa2fd95d79ca163cab8e74a0e561e416f067ec7b5c12", + "1": "a04ceff3d0aa833d139f5552b6dcd3348e8ba57f81b1d6903bb19f4388024822", + "2": "27f2533b40d0a64fc63ecf2338c9c9a1d7aa8952386558986f88db12ce92c573", + "3": "a56f37a803b2506e7223c2f41cf10bad35cc67dd2f5b2f534f2bc7d74b9a7404", + "4": "fac881d8ca2ca0cac250de093201c31d271976a10a958aa66e6e4698f45ecd2b", + "5": "71ef4ae6c1e94e84d20801fbbcf4519c0a5578cacb7223d8230dab7fc349c976", + "6": "6eb563ec23871f49e4858d7bf7c1a97218093c70d2a758551754934fbb02a887", + "7": "aeb835b4a246d15f8e55c705a5f109d4422eb27bb4fd4320e6efecd7dbdc8526", + "8": "18d9204ef8738449e215fe36bdb58e68c055a9b947611f02dba05204b9aa2821", + "9": "913879a7f860a5c9c0eb61a967a16efc0732e957d674833b8b532457173ec834" } } } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/tf32_mma_emit.c b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/tf32_mma_emit.c new file mode 100644 index 000000000000..73308b4d8b43 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/tf32_mma_emit.c @@ -0,0 +1,74 @@ +/* Copyright (c) Advanced Micro Devices, Inc., or its affiliates. + * SPDX-License-Identifier: MIT */ +#include "rocke/instance_tf32_mma_probe.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 >= 10) + { + fprintf(stderr, "unknown config index %d\n", idx); + return 2; + } + const char* modes[] = {"raw", "carrier", "rne", "prepacked", "fp32"}; + const char* mode = argc > 2 ? argv[2] : "ll"; + rocke_ir_builder_t b; + rocke_kernel_def_t* kernel = rocke_build_tf32_mma_probe(&b, idx < 5 ? 16 : 32, modes[idx % 5]); + 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, "gfx942", &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 = "gfx942"; + 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/tf32_mma_emit.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/tf32_mma_emit.py new file mode 100644 index 000000000000..f58bcda3e7a8 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/tf32_mma_emit.py @@ -0,0 +1,20 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +from rocke.instances.gfx942.tf32_mma_probe import ( + PREPARATIONS, + Tf32MmaProbeSpec, + build_tf32_mma_probe, +) +from _emit_common import run_emit + + +def _spec(idx): + if not 0 <= idx < 10: + raise SystemExit(f"unknown config index {idx}") + return Tf32MmaProbeSpec(16 if idx < 5 else 32, PREPARATIONS[idx % 5]) + + +if __name__ == "__main__": + raise SystemExit( + run_emit(_spec, lambda spec, arch: build_tf32_mma_probe(spec), arch="gfx942") + ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/numeric/test_gfx942_tf32_numeric.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/numeric/test_gfx942_tf32_numeric.py new file mode 100644 index 000000000000..ae7d2589485a --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/numeric/test_gfx942_tf32_numeric.py @@ -0,0 +1,17 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +import os +import pytest +from rocke.runtime.hip_module import get_device_arch +from rocke.examples.gfx942.tf32_numerics import run + + +def test_gfx942_tf32_numeric(tmp_path): + arch = get_device_arch() + if not arch or arch.split(":")[0] != "gfx942": + if os.getenv("ROCKE_REQUIRE_GFX942") == "1": + pytest.fail(f"Required gfx942 device absent: {arch}") + pytest.skip("requires gfx942") + result = run(tmp_path, backend="both") + assert result["status"] == "pass" + assert len(result["results"]) == 24