feat(rocke): add FP64 support for elementwise and reduce

dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/dtypes.cpp

@@ -29,7 +29,8 @@ 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"},         {"tf32", "tf32"},
-    {"xf32", "tf32"},       {"iu8", "iu8"},           {"iu4", "iu4"},
+    {"xf32", "tf32"},       {"f64", "fp64"},          {"double", "fp64"},
+    {"fp64", "fp64"},       {"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"},
@@ -123,16 +124,16 @@ 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},   {"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},
+    {"fp64", ROCKE_DTYPE_FLOAT, 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},
 };
 
 const rocke_dtype_info_t* rocke_dtype_info(const char* name)

dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp

@@ -54,6 +54,7 @@ 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_f64, ROCKE_SCALAR_F64, "f64")
 ROCKE_SCALAR_SINGLETON(rocke_fp8e4m3, ROCKE_SCALAR_FP8E4M3, "fp8e4m3")
 ROCKE_SCALAR_SINGLETON(rocke_bf8e5m2, ROCKE_SCALAR_BF8E5M2, "bf8e5m2")
 
@@ -87,6 +88,8 @@ const rocke_type_t* rocke_scalar_by_name(const char* name)
         return rocke_tf32();
     if(strcmp(name, "f32") == 0)
         return rocke_f32();
+    if(strcmp(name, "f64") == 0)
+        return rocke_f64();
     if(strcmp(name, "fp8e4m3") == 0 || strcmp(name, "e4m3") == 0)
         return rocke_fp8e4m3();
     if(strcmp(name, "bf8e5m2") == 0)
@@ -113,6 +116,8 @@ const rocke_type_t* rocke_dtype_to_ir_type(const char* dtype)
         return rocke_f16();
     if(strcmp(info->name, "fp32") == 0)
         return rocke_f32();
+    if(strcmp(info->name, "fp64") == 0)
+        return rocke_f64();
     return rocke_scalar_by_name(info->name);
 }
 

dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_arith.cpp

@@ -51,6 +51,17 @@ rocke_value_t* rocke_b_const_f32(rocke_ir_builder_t* b, double value)
     return rocke_i_op1(b, ROCKE_OP_ARITH_CONSTANT, NULL, 0, rocke_f32(), &a, "c");
 }
 
+rocke_value_t* rocke_b_const_f64(rocke_ir_builder_t* b, double value)
+{
+    rocke_attr_map_t a;
+    if(!rocke_i_live(b))
+        return NULL;
+    a = rocke_i_attrs(b);
+    rocke_attr_set_float(b, &a, "value", value);
+    rocke_attr_set_str(b, &a, "ity", "f64");
+    return rocke_i_op1(b, ROCKE_OP_ARITH_CONSTANT, NULL, 0, rocke_f64(), &a, "c");
+}
+
 rocke_value_t* rocke_b_fp16_zero(rocke_ir_builder_t* b)
 {
     rocke_attr_map_t a;

dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_flow.cpp

@@ -533,6 +533,16 @@ rocke_value_t* rocke_b_warp_shuffle_xor(rocke_ir_builder_t* b, rocke_value_t* v,
         return (rocke_value_t*)rocke_i_set_err(
             b, ROCKE_ERR_VALUE, "warp_shuffle_xor: NULL operand");
 
+    /* f64: split into two i32 halves and shuffle each on the 32-bit path. */
+    if(rocke_i_type_is(v->type, "f64"))
+    {
+        rocke_value_t* halves = rocke_b_bitcast(b, v, rocke_vector_type(b, rocke_i32(), 2));
+        rocke_value_t* parts[2];
+        parts[0] = rocke_b_warp_shuffle_xor(b, rocke_b_vec_extract(b, halves, 0), lane_xor);
+        parts[1] = rocke_b_warp_shuffle_xor(b, rocke_b_vec_extract(b, halves, 1), lane_xor);
+        return rocke_b_bitcast(b, rocke_b_vec_pack(b, parts, 2, rocke_i32()), rocke_f64());
+    }
+
     if(lane_xor >= 1 && lane_xor <= 31)
     {
         /* Intra-32-lane XOR -> ds_swizzle (1 LDS op, no addr-compute). */

dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp

@@ -290,6 +290,13 @@ rocke_value_t* rocke_b_global_load_vN_ex(rocke_ir_builder_t* b,
             return (rocke_value_t*)rocke_i_set_err(
                 b, ROCKE_ERR_VALUE, "unsupported vector width for %s global_load_vN: %d", en, n);
     }
+    else if(rocke_i_type_is(dtype, "f64"))
+    {
+        elem_bytes = 8;
+        if(n != 2)
+            return (rocke_value_t*)rocke_i_set_err(
+                b, ROCKE_ERR_VALUE, "unsupported vector width for %s global_load_vN: %d", en, n);
+    }
     else if(rocke_i_type_is(dtype, "fp8e4m3") || rocke_i_type_is(dtype, "bf8e5m2")
             || rocke_i_type_is(dtype, "i8"))
     {
@@ -303,7 +310,7 @@ rocke_value_t* rocke_b_global_load_vN_ex(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/tf32/fp8e4m3/bf8e5m2/i8, got %s",
+            "global_load_vN supports f16/bf16/i16/f32/i32/tf32/f64/fp8e4m3/bf8e5m2/i8, got %s",
             en);
     }
     streaming = mem_opts_streaming(b, opts, "global_load_vN");
@@ -419,6 +426,16 @@ void rocke_b_global_store_vN_ex(rocke_ir_builder_t* b,
             return;
         }
     }
+    else if(rocke_i_type_is(et, "f64"))
+    {
+        elem_bytes = 8;
+        if(n > 2)
+        {
+            (void)rocke_i_set_err(
+                b, ROCKE_ERR_VALUE, "global_store_vN n=%d not supported for %s", n, en);
+            return;
+        }
+    }
     else if(rocke_i_type_is(et, "i8") || rocke_i_type_is(et, "fp8e4m3")
             || rocke_i_type_is(et, "bf8e5m2"))
     {
@@ -429,7 +446,7 @@ void rocke_b_global_store_vN_ex(rocke_ir_builder_t* b,
         (void)rocke_i_set_err(
             b,
             ROCKE_ERR_VALUE,
-            "global_store_vN supports f16/bf16/i16/f32/i32/tf32/i8/fp8e4m3/bf8e5m2, got %s",
+            "global_store_vN supports f16/bf16/i16/f32/i32/tf32/f64/i8/fp8e4m3/bf8e5m2, got %s",
             en);
         return;
     }

dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_tile.cpp

@@ -87,7 +87,7 @@ static const char* rocke_mma_result_hint(const char* op_id)
     return "acc";
 }
 
-/* element-byte width helper (mirrors smem_store_vN's inline ternary). */
+/* element-byte width helper (mirrors smem_store_vN's elem_bytes). */
 static int rocke_elem_bytes_name(const char* elem_name)
 {
     if(!elem_name)
@@ -99,6 +99,10 @@ static int rocke_elem_bytes_name(const char* elem_name)
     {
         return 1;
     }
+    if(strcmp(elem_name, "f64") == 0)
+    {
+        return 8;
+    }
     if(strcmp(elem_name, "f32") == 0 || strcmp(elem_name, "i32") == 0
        || strcmp(elem_name, "tf32") == 0)
     {
@@ -240,6 +244,7 @@ void rocke_b_smem_store_vN(rocke_ir_builder_t* b,
     const char* elem_name;
     static const int allowed_8bit[] = {2, 4, 8, 16};
     static const int allowed_other[] = {2, 4, 8};
+    static const int allowed_64bit[] = {2};
     int elem_bytes;
     if(!rocke_i_live(b))
     {
@@ -271,10 +276,19 @@ void rocke_b_smem_store_vN(rocke_ir_builder_t* b,
     }
     elem_name = value->type->elem->name;
     {
-        bool eight = (strcmp(elem_name, "i8") == 0 || strcmp(elem_name, "fp8e4m3") == 0
-                      || strcmp(elem_name, "bf8e5m2") == 0);
-        const int* allowed = eight ? allowed_8bit : allowed_other;
-        int acount = eight ? 4 : 3;
+        const int* allowed = allowed_other;
+        int acount = 3;
+        if(strcmp(elem_name, "i8") == 0 || strcmp(elem_name, "fp8e4m3") == 0
+           || strcmp(elem_name, "bf8e5m2") == 0)
+        {
+            allowed = allowed_8bit;
+            acount = 4;
+        }
+        else if(strcmp(elem_name, "f64") == 0)
+        {
+            allowed = allowed_64bit;
+            acount = 1;
+        }
         if(!rocke_n_in(n, allowed, acount))
         {
             rocke_i_set_err(b,
@@ -354,6 +368,7 @@ rocke_value_t* rocke_b_smem_load_vN(rocke_ir_builder_t* b,
     static const int allowed_8bit[] = {1, 2, 4, 8, 12, 16};
     static const int allowed_16bit[] = {1, 2, 4, 6, 8};
     static const int allowed_32bit[] = {1, 2, 3, 4, 8};
+    static const int allowed_64bit[] = {1, 2};
     char hint[16];
     if(!rocke_i_live(b))
     {
@@ -365,22 +380,34 @@ 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, "tf32") == 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, "f64") == 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 / tf32 / fp8e4m3 / "
-            "bf8e5m2 / i8, got %s",
+            "smem_load_vN supports f16 / bf16 / f32 / i32 / tf32 / f64 / "
+            "fp8e4m3 / bf8e5m2 / i8, got %s",
             dn);
     }
     {
-        bool eight
-            = (strcmp(dn, "fp8e4m3") == 0 || strcmp(dn, "bf8e5m2") == 0 || strcmp(dn, "i8") == 0);
-        bool half = strcmp(dn, "f16") == 0 || strcmp(dn, "bf16") == 0;
-        const int* allowed = eight ? allowed_8bit : half ? allowed_16bit : allowed_32bit;
-        int acount = eight ? 6 : 5;
+        const int* allowed = allowed_32bit;
+        int acount = 5;
+        if(strcmp(dn, "fp8e4m3") == 0 || strcmp(dn, "bf8e5m2") == 0 || strcmp(dn, "i8") == 0)
+        {
+            allowed = allowed_8bit;
+            acount = 6;
+        }
+        else if(strcmp(dn, "f16") == 0 || strcmp(dn, "bf16") == 0)
+        {
+            allowed = allowed_16bit;
+            acount = 5;
+        }
+        else if(strcmp(dn, "f64") == 0)
+        {
+            allowed = allowed_64bit;
+            acount = 2;
+        }
         if(!rocke_n_in(n, allowed, acount))
         {
             return (rocke_value_t*)rocke_i_set_err(

dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp

@@ -12,7 +12,7 @@
  *   - exception-based errors and the NULL guard (rocke_h_fail / rocke_h_live),
  *   - naming / type mapping (rocke_h_name / rocke_h_type_to_hip / rocke_h_hip_scalar /
  *     rocke_h_vec_prefix),
- *   - float literal formatting (rocke_h_f32_literal),
+ *   - float literal formatting (rocke_h_f32_literal / rocke_h_f64_literal),
  *   - waitcnt encoders (rocke_h_encode_waitcnt + the two raw encoders),
  *   - arch gates (rocke_h_require_ds_read_tr / rocke_h_require_wmma_arch),
  *   - the smem storage side table (rocke_h_smem_storage / rocke_h_smem_set_storage),
@@ -76,6 +76,7 @@ const char* const ROCKE_HIP_PROLOGUE
       "_ROCKE_VEC(bf16, bf16x, 8); _ROCKE_VEC(bf16, bf16x, 16);\n"
       "_ROCKE_VEC(float, f32x, 1); _ROCKE_VEC(float, f32x, 2); _ROCKE_VEC(float, f32x, 4);\n"
       "_ROCKE_VEC(float, f32x, 8); _ROCKE_VEC(float, f32x, 16);\n"
+      "_ROCKE_VEC(double, f64x, 1); _ROCKE_VEC(double, f64x, 2);\n"
       "_ROCKE_VEC(int, i32x, 1); _ROCKE_VEC(int, i32x, 2); _ROCKE_VEC(int, i32x, 3);\n"
       "_ROCKE_VEC(int, i32x, 4); _ROCKE_VEC(int, i32x, 8); _ROCKE_VEC(int, i32x, 16);\n"
       "_ROCKE_VEC(int16_t, i16x, 1); _ROCKE_VEC(int16_t, i16x, 2);\n"
@@ -316,6 +317,10 @@ const char* rocke_h_hip_scalar(const char* ir_scalar_name)
     {
         return "float";
     }
+    if(strcmp(ir_scalar_name, "f64") == 0)
+    {
+        return "double";
+    }
     if(strcmp(ir_scalar_name, "fp8e4m3") == 0)
     {
         return "fp8e4m3";
@@ -346,6 +351,10 @@ const char* rocke_h_vec_prefix(const char* ir_scalar_name, bool full_map)
             {
                 return "f32x";
             }
+            if(strcmp(ir_scalar_name, "f64") == 0)
+            {
+                return "f64x";
+            }
             if(strcmp(ir_scalar_name, "i32") == 0 || strcmp(ir_scalar_name, "tf32") == 0)
             {
                 return "i32x";
@@ -384,9 +393,9 @@ static bool rocke_h_scalar_in_vec_map(const char* name, bool full_map)
     {
         return false;
     }
-    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;
+    return strcmp(name, "tf32") == 0 || strcmp(name, "f32") == 0 || strcmp(name, "f64") == 0
+           || strcmp(name, "i32") == 0 || strcmp(name, "i16") == 0 || strcmp(name, "i8") == 0
+           || strcmp(name, "fp8e4m3") == 0 || strcmp(name, "bf8e5m2") == 0;
 }
 
 const char* rocke_h_vec_prefix_checked(rocke_h_lowerer_t* lw,
@@ -407,6 +416,20 @@ const char* rocke_h_vec_prefix_checked(rocke_h_lowerer_t* lw,
     return rocke_h_vec_prefix(ir_scalar_name, full_map);
 }
 
+/* Python _reject_f64(op, tname):
+ *   if tname == "f64":
+ *       raise NotImplementedError(f"{op.name} does not support f64") */
+void rocke_h_reject_f64(rocke_h_lowerer_t* lw, const rocke_op_t* op, const char* tname)
+{
+    if(tname && strcmp(tname, "f64") == 0)
+    {
+        rocke_h_fail(lw,
+                     ROCKE_ERR_NOTIMPL,
+                     "%s does not support f64",
+                     op->name ? op->name : rocke_opcode_name(op->opcode));
+    }
+}
+
 /* Python _type_to_hip(t). Returns an arena-owned string; throws on an unmappable
  * HIP type (KeyError parity), including logical types without direct HIP lowering. */
 const char* rocke_h_type_to_hip(rocke_h_lowerer_t* lw, const rocke_type_t* t)
@@ -451,9 +474,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, "tf32") != 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, "f64") != 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 "";
@@ -516,6 +539,35 @@ const char* rocke_h_f32_literal(rocke_h_lowerer_t* lw, double val)
     return out ? out : "";
 }
 
+const char* rocke_h_f64_literal(rocke_h_lowerer_t* lw, double val)
+{
+    char* out;
+    if(!lw || !lw->b)
+    {
+        return "";
+    }
+    if(isnan(val))
+    {
+        out = rocke_arena_strdup(&lw->b->arena, "((double)NAN)");
+        return out ? out : "";
+    }
+    if(isinf(val))
+    {
+        out = rocke_arena_strdup(&lw->b->arena,
+                                 val < 0 ? "((double)-INFINITY)" : "((double)INFINITY)");
+        return out ? out : "";
+    }
+    if(val == 0.0 && signbit(val))
+    {
+        out = rocke_arena_strdup(&lw->b->arena, "-0.0");
+        return out ? out : "";
+    }
+    /* NOTE(port): Python emits repr(float); %.17g is a close approximation but
+     * NOT guaranteed byte-identical to CPython repr. Known port hazard. */
+    out = rocke_arena_printf(&lw->b->arena, "%.17g", val);
+    return out ? out : "";
+}
+
 /* ============================== waitcnt encoders ===================== */
 
 static int h_clamp(int v, int lo, int hi)