feat(rocke): add logical TF32 and native gfx942 atoms

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

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},
@@ -1040,6 +1103,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",

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},
 };
 

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)

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)
 {

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;
     }

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 <stdio.h>
 #include <string.h>
 
@@ -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);