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 3644634e3bc6..3774425da6ac 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 @@ -3,11 +3,10 @@ /* * arch_target_arch_target_data.c -- bucket 0 of the C99 port of * rocke.core.arch.target: the frozen SSOT data, the lane-coord IR emitters, - * the dtype-normalisation core, and the shared lookup tables/helpers declared in + * and the shared lookup tables/helpers declared in * rocke/arch_target_internal.h. * * This file is a faithful, byte-identical translation of: - * - normalize_dtype() / _DTYPE_ALIASES -> rocke_normalize_dtype * - the _mfma / _wmma lane-coord closures -> static rocke_lane_coord_fn * - LayoutMap.coord -> rocke_layout_map_coord * - _MMA_FRAGMENT_INFO + _build_mma_op (precomputed) -> per-op_id static @@ -29,135 +28,6 @@ #include #include -/* ========================================================================= - * dtype normalisation core (shared by both buckets) - * ========================================================================= - * - * Mirrors target.py::normalize_dtype: strip + lower(name), then - * _DTYPE_ALIASES.get(key, key). The canonical RHS values are interned static - * strings; unknown spellings pass through as the lowercased text in `lowered`. - */ - -typedef struct rocke_ati_dtype_alias -{ - const char* alias; /* lowercased key */ - const char* canonical; /* interned canonical value */ -} rocke_ati_dtype_alias_t; - -/* Byte-for-byte the _DTYPE_ALIASES map (insertion order is irrelevant: lookup is - * by exact lowercased key). */ -static const rocke_ati_dtype_alias_t k_dtype_aliases[] = { - {"f16", "fp16"}, - {"half", "fp16"}, - {"fp16", "fp16"}, - {"bf16", "bf16"}, - {"bfloat16", "bf16"}, - {"f32", "fp32"}, - {"float", "fp32"}, - {"fp32", "fp32"}, - {"fp8", "fp8e4m3"}, - {"fp8e4m3", "fp8e4m3"}, - {"bf8", "bf8e5m2"}, - {"bf8e5m2", "bf8e5m2"}, - {"iu8", "iu8"}, - {"iu4", "iu4"}, - {"i8", "i8"}, - {"int8", "i8"}, - {"i4", "i4"}, - {"int4", "i4"}, - {"i32", ROCKE_DTYPE_I32}, - {"int32", ROCKE_DTYPE_I32}, -}; -#define K_NUM_DTYPE_ALIASES ((int)(sizeof(k_dtype_aliases) / sizeof(k_dtype_aliases[0]))) - -/* str.strip(): Python strips ASCII whitespace from both ends. */ -static int rocke_ati_is_ws(char c) -{ - return c == ' ' || c == '\t' || c == '\n' || c == '\r' || c == '\f' || c == '\v'; -} - -const char* rocke_ati_normalize_dtype(const char* name, char* lowered, size_t cap) -{ - int i; - size_t n; - const char* start; - const char* end; - - if(name == NULL) - { - /* Python would raise on None.strip(); the C contract requires a result. - * Treat as empty string (no alias match, passes through as ""). */ - if(lowered != NULL && cap > 0) - { - lowered[0] = '\0'; - } - return (lowered != NULL && cap > 0) ? lowered : ""; - } - - /* strip(): advance over leading/trailing whitespace */ - start = name; - while(*start != '\0' && rocke_ati_is_ws(*start)) - { - start++; - } - end = start + strlen(start); - while(end > start && rocke_ati_is_ws(end[-1])) - { - end--; - } - - /* lower() into the caller buffer (the pass-through result) */ - n = (size_t)(end - start); - if(lowered == NULL || cap == 0) - { - /* No scratch: caller guarantees a known spelling. Build a small inline - * copy on a fixed buffer to look up; if unknown we have nowhere to - * return it, so fall back to "". */ - static char tmp[64]; - size_t m = n < sizeof(tmp) - 1 ? n : sizeof(tmp) - 1; - for(i = 0; (size_t)i < m; i++) - { - char c = start[i]; - tmp[i] = (c >= 'A' && c <= 'Z') ? (char)(c - 'A' + 'a') : c; - } - tmp[m] = '\0'; - for(i = 0; i < K_NUM_DTYPE_ALIASES; i++) - { - if(strcmp(tmp, k_dtype_aliases[i].alias) == 0) - { - return k_dtype_aliases[i].canonical; - } - } - return ""; - } - - if(n > cap - 1) - { - n = cap - 1; - } - for(i = 0; (size_t)i < n; i++) - { - char c = start[i]; - lowered[i] = (c >= 'A' && c <= 'Z') ? (char)(c - 'A' + 'a') : c; - } - lowered[n] = '\0'; - - /* _DTYPE_ALIASES.get(key, key) */ - for(i = 0; i < K_NUM_DTYPE_ALIASES; i++) - { - if(strcmp(lowered, k_dtype_aliases[i].alias) == 0) - { - return k_dtype_aliases[i].canonical; - } - } - return lowered; -} - -const char* rocke_normalize_dtype(const char* name, char* scratch, size_t scratch_cap) -{ - return rocke_ati_normalize_dtype(name, scratch, scratch_cap); -} - /* ========================================================================= * lane-coord emitters (static rocke_lane_coord_fn) * ========================================================================= @@ -736,9 +606,11 @@ bool rocke_layout_map_coord(const rocke_layout_map_t* m, { if(b != NULL) { - const char* role_txt = (m->role == ROCKE_MMA_ROLE_ACC) ? "acc" - : (m->role == ROCKE_MMA_ROLE_A) ? "a" - : "b"; + const char* role_txt = (m->role == ROCKE_MMA_ROLE_ACC) ? "acc" + : (m->role == ROCKE_MMA_ROLE_A) ? "a" + : (m->role == ROCKE_MMA_ROLE_B) ? "b" + : (m->role == ROCKE_MMA_ROLE_A_SCALE) ? "a_scale" + : "b_scale"; b->status = ROCKE_ERR_VALUE; snprintf(b->err, ROCKE_ERR_MSG_CAP, @@ -757,6 +629,39 @@ bool rocke_layout_map_coord(const rocke_layout_map_t* m, return true; } +/* Scale slots describe logical K groups; both half-waves duplicate them. */ +static void _wmma_gfx1250_a_scale(rocke_ir_builder_t* b, + rocke_value_t* lane, + int slot, + rocke_value_t** out0, + rocke_value_t** out1) +{ + ROCKE_ATI_COORD_GUARD(b, out0, out1); + rocke_value_t* c16 = rocke_b_const_i32(b, 16); + rocke_value_t* row = rocke_b_mod(b, lane, c16); + rocke_value_t* group = rocke_b_const_i32(b, slot); + if(out0) + *out0 = row; + if(out1) + *out1 = group; +} + +static void _wmma_gfx1250_b_scale(rocke_ir_builder_t* b, + rocke_value_t* lane, + int slot, + rocke_value_t** out0, + rocke_value_t** out1) +{ + ROCKE_ATI_COORD_GUARD(b, out0, out1); + rocke_value_t* c16 = rocke_b_const_i32(b, 16); + rocke_value_t* col = rocke_b_mod(b, lane, c16); + rocke_value_t* group = rocke_b_const_i32(b, slot); + if(out0) + *out0 = group; + if(out1) + *out1 = col; +} + /* ========================================================================= * Per-op_id LayoutMap statics (the precomputed _build_mma_op output) * ========================================================================= @@ -772,6 +677,15 @@ bool rocke_layout_map_coord(const rocke_layout_map_t* m, * struct; the rest carry NULL layout pointers in the catalog directly. */ +static const rocke_layout_map_t lm_wmma_scale_k32_a + = {ROCKE_MMA_ROLE_A_SCALE, 4, 32, _wmma_gfx1250_a_scale}; +static const rocke_layout_map_t lm_wmma_scale_k32_b + = {ROCKE_MMA_ROLE_B_SCALE, 4, 32, _wmma_gfx1250_b_scale}; +static const rocke_layout_map_t lm_wmma_scale_k16_a + = {ROCKE_MMA_ROLE_A_SCALE, 8, 32, _wmma_gfx1250_a_scale}; +static const rocke_layout_map_t lm_wmma_scale_k16_b + = {ROCKE_MMA_ROLE_B_SCALE, 8, 32, _wmma_gfx1250_b_scale}; + /* --- mfma_f32_16x16x16_f16 / _bf16: a/b/c all present (frag 4/4/4, wave64) --- */ static const rocke_layout_map_t lm_mfma_16x16x16_a = {ROCKE_MMA_ROLE_A, 4, 64, _mfma_a_16x16}; static const rocke_layout_map_t lm_mfma_16x16x16_b = {ROCKE_MMA_ROLE_B, 4, 64, _mfma_b_16x16}; @@ -1017,8 +931,10 @@ static const rocke_ati_mma_frag_row_t rocke_ati_mma_frag[] = { {"wmma_gfx1250_f32_16x16x64_fp8_bf8", 8}, {"wmma_gfx1250_f32_16x16x64_bf8_fp8", 8}, {"wmma_gfx1250_f32_16x16x64_bf8_bf8", 8}, - {"wmma_scale_f32_16x16x128_fp8_fp8", 8}, - {"wmma_scale16_f32_16x16x128_fp8_fp8", 8}, + {"wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k32", 8}, + {"wmma_gfx1250_f32_16x16x128_bf8_bf8_scale_e8m0_e8m0_k32", 8}, + {"wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k16", 8}, + {"wmma_gfx1250_f32_16x16x128_bf8_bf8_scale_e8m0_e8m0_k16", 8}, }; int rocke_arch_mma_c_frag_len(const char* op_id) @@ -1045,15 +961,14 @@ int rocke_arch_mma_c_frag_len(const char* op_id) * * Frag lengths come from _MMA_FRAGMENT_INFO[op_id]; op_ids absent from that * table carry (0,0,0,64) and NULL maps. dtype fields hold the normalised - * (canonical) catalog keys (interned strings from k_dtype_aliases), exactly as + * (canonical) catalog keys from rocke_normalize_dtype, exactly as * normalize_dtype(o["a"/"b"/"c"]) would produce. * * Field order of rocke_mma_op_t: * family, a_dtype, b_dtype, c_dtype, m, n, k, op_id, * a_frag_len, b_frag_len, c_frag_len, wave_size, a_layout, b_layout, c_layout * - * NOTE: fp4/fp6 dtypes are not in _DTYPE_ALIASES, so normalize_dtype passes them - * through as the lowercased spelling "fp4"/"fp6" (Python identity fallthrough). + * Low-bit spellings normalize to explicit exponent/mantissa catalog keys. */ /* ----------------------------- gfx90a (CDNA2) ---------------------------- */ @@ -1489,11 +1404,11 @@ static const rocke_mma_op_t k_mma_gfx950[] = { &lm_mfma_32x32x16_a, &lm_mfma_32x32x16_b, &lm_mfma_32x32x16_c}, - /* fp4/fp6 pass through normalize_dtype unchanged; op_ids carry frag lengths + /* Low-bit catalog keys use exponent/mantissa names; op_ids carry frag lengths * only (no verified maps), per _MMA_FRAGMENT_INFO. */ {"mma", - "fp4", - "fp4", + "fp4e2m1", + "fp4e2m1", "fp32", 16, 16, @@ -1507,8 +1422,8 @@ static const rocke_mma_op_t k_mma_gfx950[] = { NULL, NULL}, {"mma", - "fp6", - "fp6", + "fp6e2m3", + "fp6e2m3", "fp32", 16, 16, @@ -1742,95 +1657,153 @@ static const rocke_mma_op_t k_mma_gfx1250[] = { &lm_wmma_gfx1250_32_b, &lm_wmma_gfx12_c}, {"wmma", - "fp8", - "fp8", + "fp8e4m3", + "fp8e4m3", "fp32", 16, 16, 64, "wmma_gfx1250_f32_16x16x64_fp8_fp8", - 32, - 32, + 8, + 8, 8, 32, NULL, NULL, &lm_wmma_gfx12_c}, {"wmma", - "fp8", - "bf8", + "fp8e4m3", + "bf8e5m2", "fp32", 16, 16, 64, "wmma_gfx1250_f32_16x16x64_fp8_bf8", - 32, - 32, + 8, + 8, 8, 32, NULL, NULL, &lm_wmma_gfx12_c}, {"wmma", - "bf8", - "fp8", + "bf8e5m2", + "fp8e4m3", "fp32", 16, 16, 64, "wmma_gfx1250_f32_16x16x64_bf8_fp8", - 32, - 32, + 8, + 8, 8, 32, NULL, NULL, &lm_wmma_gfx12_c}, {"wmma", - "bf8", - "bf8", + "bf8e5m2", + "bf8e5m2", "fp32", 16, 16, 64, "wmma_gfx1250_f32_16x16x64_bf8_bf8", + 8, + 8, + 8, 32, + NULL, + NULL, + &lm_wmma_gfx12_c}, + {"wmma_scaled", + "fp8e4m3", + "fp8e4m3", + "fp32", + 16, + 16, + 128, + "wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k32", + 16, + 16, + 8, 32, + NULL, + NULL, + &lm_wmma_gfx12_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 4, + 4, + &lm_wmma_scale_k32_a, + &lm_wmma_scale_k32_b}, + {"wmma_scaled", + "bf8e5m2", + "bf8e5m2", + "fp32", + 16, + 16, + 128, + "wmma_gfx1250_f32_16x16x128_bf8_bf8_scale_e8m0_e8m0_k32", + 16, + 16, 8, 32, NULL, NULL, - &lm_wmma_gfx12_c}, - {"wmma_scale", - "fp8", - "fp8", + &lm_wmma_gfx12_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K32, + 4, + 4, + &lm_wmma_scale_k32_a, + &lm_wmma_scale_k32_b}, + {"wmma_scaled", + "fp8e4m3", + "fp8e4m3", "fp32", 16, 16, 128, - "wmma_scale_f32_16x16x128_fp8_fp8", + "wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k16", 16, 16, 8, 32, NULL, NULL, - &lm_wmma_gfx12_c}, - {"wmma_scale16", - "fp8", - "fp8", + &lm_wmma_gfx12_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K16, + 8, + 8, + &lm_wmma_scale_k16_a, + &lm_wmma_scale_k16_b}, + {"wmma_scaled", + "bf8e5m2", + "bf8e5m2", "fp32", 16, 16, 128, - "wmma_scale16_f32_16x16x128_fp8_fp8", + "wmma_gfx1250_f32_16x16x128_bf8_bf8_scale_e8m0_e8m0_k16", 16, 16, 8, 32, NULL, NULL, - &lm_wmma_gfx12_c}, + &lm_wmma_gfx12_c, + "e8m0", + "e8m0", + ROCKE_MMA_SCALE_K16, + 8, + 8, + &lm_wmma_scale_k16_a, + &lm_wmma_scale_k16_b}, }; /* ========================================================================= diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/helpers.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/helpers.cpp index 5f74ed7ddc7d..00abd5e5b4a3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/helpers.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/helpers.cpp @@ -87,14 +87,16 @@ const rocke_mmaop_t* rocke_archtarget_op_for_shape(const rocke_archtarget_t* t, const char* c_dtype, int m, int n, - int k) + int k, + const rocke_mma_scale_filter_t* scales) { /* target.mma.op_for_shape(family=..., a/b/c=..., m, n, k). */ if(t == NULL) { return NULL; } - return rocke_mma_catalog_op_for_shape(&t->mma, family, a_dtype, b_dtype, c_dtype, m, n, k); + return rocke_mma_catalog_op_for_shape( + &t->mma, family, a_dtype, b_dtype, c_dtype, m, n, k, scales); } const rocke_mmaop_t* rocke_archtarget_by_op_id(const rocke_archtarget_t* t, const char* op_id) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/mma_family_index.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/mma_family_index.h new file mode 100644 index 000000000000..118690ade150 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/mma_family_index.h @@ -0,0 +1,95 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* Private immutable counterpart of Python target._op_id_family(). Owns the + * sorted index; borrows IDs and families from the static architecture catalogs. */ +#ifndef ROCKE_MMA_FAMILY_INDEX_H +#define ROCKE_MMA_FAMILY_INDEX_H + +#include +#include + +#include "rocke/arch_target_internal.h" +#include "rocke/error.hpp" + +struct rocke_ati_mma_family_index +{ + rocke_ati_mma_family_index(const rocke_ati_arch_row_t* registry, int count) + { + size_t capacity = 0; + for(int i = 0; i < count; ++i) + if(registry[i].target) + capacity += registry[i].target->mma.num_ops; + if(!capacity) + return; + rows = static_cast(malloc(capacity * sizeof(row))); + if(!rows) + ckc::raise_status(ROCKE_ERR_OOM, "MMA family index allocation failed"); + for(int i = 0; i < count; ++i) + { + const auto* target = registry[i].target; + if(!target) + continue; + for(int j = 0; j < target->mma.num_ops; ++j) + { + const auto& op = target->mma.ops[j]; + rows[size++] = {op.op_id, op.family, false}; + } + } + qsort(rows, size, sizeof(row), [](const void* a, const void* b) { + return strcmp(static_cast(a)->op_id, static_cast(b)->op_id); + }); + size_t unique = 0; + for(size_t i = 0; i < size; ++i) + { + if(unique && strcmp(rows[unique - 1].op_id, rows[i].op_id) == 0) + rows[unique - 1].conflict |= strcmp(rows[unique - 1].family, rows[i].family) != 0; + else + rows[unique++] = rows[i]; + } + size = unique; + } + + ~rocke_ati_mma_family_index() + { + free(rows); + } + rocke_ati_mma_family_index(const rocke_ati_mma_family_index&) = delete; + rocke_ati_mma_family_index& operator=(const rocke_ati_mma_family_index&) = delete; + + const char* lookup(const char* op_id) const + { + if(!op_id) + return NULL; + size_t first = 0, last = size; + while(first < last) + { + const size_t mid = first + (last - first) / 2; + const int cmp = strcmp(op_id, rows[mid].op_id); + if(cmp < 0) + last = mid; + else if(cmp > 0) + first = mid + 1; + else + { + if(rows[mid].conflict) + ckc::raise_status( + ROCKE_ERR_VALUE, + "arch SSOT drift: op_id has inconsistent family across arches"); + return rows[mid].family; + } + } + return NULL; + } + +private: + struct row + { + const char* op_id; + const char* family; + bool conflict; + }; + row* rows = NULL; + size_t size = 0; +}; + +#endif /* ROCKE_MMA_FAMILY_INDEX_H */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/query.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/query.cpp index 009d15329ef3..27365b859cfb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/query.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/query.cpp @@ -14,6 +14,8 @@ #include #include +#include "mma_family_index.h" + #include "rocke/arch_target.h" #include "rocke/arch_target_internal.h" #include "rocke/error.hpp" @@ -48,6 +50,49 @@ static const char* rocke_ati_family_or_default(const char* family) return family ? family : "mma"; } +static const char* rocke_ati_scale_dtype(const char* dtype) +{ + return dtype && strcmp(dtype, "fp8e4m3") == 0 ? "e4m3" : dtype; +} + +static void rocke_ati_validate_scales(const rocke_mma_scale_filter_t* scales) +{ + if(!scales) + return; + if(!scales->a_scale_dtype && !scales->b_scale_dtype && scales->scale_block_k == 0) + return; + const char* dtypes[2] = {scales->a_scale_dtype, scales->b_scale_dtype}; + for(int i = 0; i < 2; ++i) + { + const char* dtype = rocke_ati_scale_dtype(dtypes[i]); + if(!dtype + || (strcmp(dtype, "e8m0") != 0 && strcmp(dtype, "e4m3") != 0 + && strcmp(dtype, "e5m3") != 0)) + ckc::raise_status(ROCKE_ERR_VALUE, "MMA scale dtype must be e8m0, e4m3, or e5m3"); + } + if(scales->scale_block_k != 16 && scales->scale_block_k != 32) + ckc::raise_status(ROCKE_ERR_VALUE, + "MMA scale_block_k must be an integer equal to 16 or 32"); +} + +static bool rocke_ati_scales_match(const rocke_mma_op_t* op, const rocke_mma_scale_filter_t* scales) +{ + if(!scales) + return true; + if(scales->scale_block_k != op->scale_block_k) + return false; + const char* actual[2] = {op->a_scale_dtype, op->b_scale_dtype}; + const char* queried[2] = {scales->a_scale_dtype, scales->b_scale_dtype}; + for(int i = 0; i < 2; ++i) + { + const char* expected = rocke_ati_scale_dtype(queried[i]); + const char* dtype = rocke_ati_scale_dtype(actual[i]); + if(bool(expected) != bool(dtype) || (dtype && strcmp(expected, dtype) != 0)) + return false; + } + return true; +} + /* ============================== MMA atom getters ====================== */ void rocke_mma_op_shape(const rocke_mma_op_t* op, int* m, int* n, int* k) @@ -115,6 +160,18 @@ const rocke_layout_map_t* rocke_mma_op_acc_layout(const rocke_mma_op_t* op, rock return rocke_mma_op_c_layout(op, b); } +const rocke_layout_map_t* rocke_mma_op_a_scale_layout(const rocke_mma_op_t* op, + rocke_ir_builder_t* b) +{ + return rocke_ati_require_layout(op, op ? op->a_scale_layout : NULL, "a_scale", b); +} + +const rocke_layout_map_t* rocke_mma_op_b_scale_layout(const rocke_mma_op_t* op, + rocke_ir_builder_t* b) +{ + return rocke_ati_require_layout(op, op ? op->b_scale_layout : NULL, "b_scale", b); +} + /* ============================== MMA catalog =========================== */ const rocke_mma_op_t* rocke_mma_catalog_ops(const rocke_mma_catalog_t* cat, int* num_out) @@ -139,7 +196,8 @@ static bool rocke_ati_op_matches(const rocke_mma_op_t* op, const char* b, const char* c, int m, - int n) + int n, + const rocke_mma_scale_filter_t* scales) { if(strcmp(op->family, family) != 0) return false; @@ -153,7 +211,7 @@ static bool rocke_ati_op_matches(const rocke_mma_op_t* op, return false; if(n >= 0 && op->n != n) return false; - return true; + return rocke_ati_scales_match(op, scales); } int rocke_mma_catalog_enumerate(const rocke_mma_catalog_t* cat, @@ -164,7 +222,8 @@ int rocke_mma_catalog_enumerate(const rocke_mma_catalog_t* cat, int m, int n, const rocke_mma_op_t** out, - int cap) + int cap, + const rocke_mma_scale_filter_t* scales) { char abuf[64], bbuf[64], cbuf[64]; const char *a, *bd, *c, *fam; @@ -172,15 +231,16 @@ int rocke_mma_catalog_enumerate(const rocke_mma_catalog_t* cat, if(!cat) return 0; + rocke_ati_validate_scales(scales); fam = rocke_ati_family_or_default(family); - a = rocke_ati_normalize_dtype(a_dtype, abuf, sizeof abuf); - bd = rocke_ati_normalize_dtype(b_dtype, bbuf, sizeof bbuf); - c = rocke_ati_normalize_dtype(c_dtype, cbuf, sizeof cbuf); + a = rocke_normalize_dtype(a_dtype, abuf, sizeof abuf); + bd = rocke_normalize_dtype(b_dtype, bbuf, sizeof bbuf); + c = rocke_normalize_dtype(c_dtype, cbuf, sizeof cbuf); for(i = 0; i < cat->num_ops; ++i) { const rocke_mma_op_t* op = &cat->ops[i]; - if(!rocke_ati_op_matches(op, fam, a, bd, c, m, n)) + if(!rocke_ati_op_matches(op, fam, a, bd, c, m, n, scales)) continue; if(out && total < cap) out[total] = op; @@ -196,7 +256,8 @@ bool rocke_mma_catalog_has_shape(const rocke_mma_catalog_t* cat, const char* c_dtype, int m, int n, - int k) + int k, + const rocke_mma_scale_filter_t* scales) { char abuf[64], bbuf[64], cbuf[64]; const char *a, *bd, *c, *fam; @@ -204,16 +265,17 @@ bool rocke_mma_catalog_has_shape(const rocke_mma_catalog_t* cat, if(!cat) return false; + rocke_ati_validate_scales(scales); fam = rocke_ati_family_or_default(family); - a = rocke_ati_normalize_dtype(a_dtype, abuf, sizeof abuf); - bd = rocke_ati_normalize_dtype(b_dtype, bbuf, sizeof bbuf); - c = rocke_ati_normalize_dtype(c_dtype, cbuf, sizeof cbuf); + a = rocke_normalize_dtype(a_dtype, abuf, sizeof abuf); + bd = rocke_normalize_dtype(b_dtype, bbuf, sizeof bbuf); + c = rocke_normalize_dtype(c_dtype, cbuf, sizeof cbuf); /* Python enumerates with m=m, n=n then checks op.shape == (m, n, k). */ for(i = 0; i < cat->num_ops; ++i) { const rocke_mma_op_t* op = &cat->ops[i]; - if(!rocke_ati_op_matches(op, fam, a, bd, c, m, n)) + if(!rocke_ati_op_matches(op, fam, a, bd, c, m, n, scales)) continue; if(op->m == m && op->n == n && op->k == k) return true; @@ -228,32 +290,41 @@ const rocke_mma_op_t* rocke_mma_catalog_select_largest_k(const rocke_mma_catalog const char* c_dtype, int m, int n, - int k_max) + int k_max, + const rocke_mma_scale_filter_t* scales) { char abuf[64], bbuf[64], cbuf[64]; const char *a, *bd, *c, *fam; const rocke_mma_op_t* best = NULL; int i; + bool ambiguous = false; if(!cat) return NULL; + rocke_ati_validate_scales(scales); fam = rocke_ati_family_or_default(family); - a = rocke_ati_normalize_dtype(a_dtype, abuf, sizeof abuf); - bd = rocke_ati_normalize_dtype(b_dtype, bbuf, sizeof bbuf); - c = rocke_ati_normalize_dtype(c_dtype, cbuf, sizeof cbuf); + a = rocke_normalize_dtype(a_dtype, abuf, sizeof abuf); + bd = rocke_normalize_dtype(b_dtype, bbuf, sizeof bbuf); + c = rocke_normalize_dtype(c_dtype, cbuf, sizeof cbuf); for(i = 0; i < cat->num_ops; ++i) { const rocke_mma_op_t* op = &cat->ops[i]; - if(!rocke_ati_op_matches(op, fam, a, bd, c, m, n)) + if(!rocke_ati_op_matches(op, fam, a, bd, c, m, n, scales)) continue; if(k_max >= 0 && op->k > k_max) continue; /* Python: k_max is None || op.k <= k_max */ - /* max(cands, key=op.k): first op wins ties (Python max keeps the first - * maximal element when iterating in catalog order). */ if(best == NULL || op->k > best->k) + { best = op; + ambiguous = false; + } + else if(op->k == best->k) + ambiguous = true; } + if(ambiguous) + ckc::raise_status(ROCKE_ERR_VALUE, + "ambiguous MMA query; specify the full operand contract"); return best; } @@ -278,33 +349,52 @@ const rocke_mma_op_t* rocke_mma_catalog_op_for_shape(const rocke_mma_catalog_t* const char* c_dtype, int m, int n, - int k) + int k, + const rocke_mma_scale_filter_t* scales) { char abuf[64], bbuf[64], cbuf[64]; const char *a, *bd, *c, *fam; int i; + const rocke_mma_op_t* match = NULL; if(!cat) return NULL; + rocke_ati_validate_scales(scales); fam = rocke_ati_family_or_default(family); - a = rocke_ati_normalize_dtype(a_dtype, abuf, sizeof abuf); - bd = rocke_ati_normalize_dtype(b_dtype, bbuf, sizeof bbuf); - c = rocke_ati_normalize_dtype(c_dtype, cbuf, sizeof cbuf); + a = rocke_normalize_dtype(a_dtype, abuf, sizeof abuf); + bd = rocke_normalize_dtype(b_dtype, bbuf, sizeof bbuf); + c = rocke_normalize_dtype(c_dtype, cbuf, sizeof cbuf); - /* Python enumerates with m=m, n=n, then returns the first op whose k == k. */ + /* Exact selection requires a unique contract. */ for(i = 0; i < cat->num_ops; ++i) { const rocke_mma_op_t* op = &cat->ops[i]; - if(!rocke_ati_op_matches(op, fam, a, bd, c, m, n)) + if(!rocke_ati_op_matches(op, fam, a, bd, c, m, n, scales)) continue; if(op->k == k) - return op; + { + if(match) + ckc::raise_status(ROCKE_ERR_VALUE, + "ambiguous MMA query; specify the full operand contract"); + match = op; + } } - return NULL; + return match; } /* ===================== bare-op_id SSOT lookups ======================== */ +const char* rocke_arch_mma_op_id_family(const char* op_id) +{ + if(!op_id) + return NULL; + // Function-local static initialization publishes the immutable index once, + // including when independent builders first query it concurrently. + static const rocke_ati_mma_family_index index(rocke_ati_arch_registry, + rocke_ati_arch_registry_len); + return index.lookup(op_id); +} + const char* rocke_arch_mma_op_id_c_dtype(const char* op_id) { int i; 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 new file mode 100644 index 000000000000..0f717a76e35f --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/dtypes.cpp @@ -0,0 +1,139 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* Target-independent dtype names. Mirrors rocke.core.dtypes: + * normalize_dtype -> rocke_normalize_dtype. + * No architecture, IR builder, or lowering dependencies. + */ +#include "rocke/dtypes.h" + +#include + +/* ========================================================================= + * target-independent dtype normalization + * ========================================================================= + * + * Mirrors core/dtypes.py::normalize_dtype: strip + lower(name), then + * _DTYPE_ALIASES.get(key, key). The canonical RHS values are interned static + * strings; unknown spellings pass through as the lowercased text in `lowered`. + */ + +typedef struct rocke_dtype_alias +{ + const char* alias; /* lowercased key */ + const char* canonical; /* interned canonical value */ +} rocke_dtype_alias_t; + +/* 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"}, + {"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]))) + +/* str.strip(): Python strips ASCII whitespace from both ends. */ +static int rocke_dtype_is_ws(char c) +{ + return c == ' ' || c == '\t' || c == '\n' || c == '\r' || c == '\f' || c == '\v'; +} + +const char* rocke_normalize_dtype(const char* name, char* lowered, size_t cap) +{ + int i; + size_t n; + const char* start; + const char* end; + + if(name == NULL) + { + /* Python would raise on None.strip(); the C contract requires a result. + * Treat as empty string (no alias match, passes through as ""). */ + if(lowered != NULL && cap > 0) + { + lowered[0] = '\0'; + } + return (lowered != NULL && cap > 0) ? lowered : ""; + } + + /* strip(): advance over leading/trailing whitespace */ + start = name; + while(*start != '\0' && rocke_dtype_is_ws(*start)) + { + start++; + } + end = start + strlen(start); + while(end > start && rocke_dtype_is_ws(end[-1])) + { + end--; + } + + /* lower() into the caller buffer (the pass-through result) */ + n = (size_t)(end - start); + if(lowered == NULL || cap == 0) + { + /* No scratch: caller guarantees a known spelling. Build a small inline + * copy on a fixed buffer to look up; if unknown we have nowhere to + * return it, so fall back to "". */ + static char tmp[64]; + size_t m = n < sizeof(tmp) - 1 ? n : sizeof(tmp) - 1; + for(i = 0; (size_t)i < m; i++) + { + char c = start[i]; + tmp[i] = (c >= 'A' && c <= 'Z') ? (char)(c - 'A' + 'a') : c; + } + tmp[m] = '\0'; + for(i = 0; i < K_NUM_DTYPE_ALIASES; i++) + { + if(strcmp(tmp, k_dtype_aliases[i].alias) == 0) + { + return k_dtype_aliases[i].canonical; + } + } + return ""; + } + + if(n > cap - 1) + { + n = cap - 1; + } + for(i = 0; (size_t)i < n; i++) + { + char c = start[i]; + lowered[i] = (c >= 'A' && c <= 'Z') ? (char)(c - 'A' + 'a') : c; + } + lowered[n] = '\0'; + + /* _DTYPE_ALIASES.get(key, key) */ + for(i = 0; i < K_NUM_DTYPE_ALIASES; i++) + { + if(strcmp(lowered, k_dtype_aliases[i].alias) == 0) + { + return k_dtype_aliases[i].canonical; + } + } + return lowered; +} 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 358e870d3ef7..8ebb8bb6191e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_tile.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_tile.cpp @@ -49,8 +49,6 @@ static const rocke_mma_hint_row_t ROCKE_MMA_RESULT_HINT[] = { {"mfma_f32_16x16x96_fp6", "acc6"}, {"mfma_f32_16x16x128_fp8", "acc128"}, {"mfma_scale_f32_16x16x128_f8f6f4", "mxacc"}, - {"wmma_scale_f32_16x16x128_fp8_fp8", "mxacc"}, - {"wmma_scale16_f32_16x16x128_fp8_fp8", "mxacc"}, }; /* Accumulator fragment length for op_id, from the arch SSOT @@ -71,6 +69,9 @@ static bool rocke_mma_is_int_acc(const char* op_id) static const char* rocke_mma_result_hint(const char* op_id) { + const char* family = rocke_arch_mma_op_id_family(op_id); + if(family && strcmp(family, "wmma_scaled") == 0) + return "mxacc"; size_t i; if(op_id) { @@ -653,32 +654,6 @@ rocke_value_t* rocke_b_mfma_scale_f32_16x16x128_f8f6f4(rocke_ir_builder_t* b, return rocke_b_mma(b, "mfma_scale_f32_16x16x128_f8f6f4", a, bb, c, extra, 2); } -rocke_value_t* rocke_b_wmma_scale_f32_16x16x128_fp8_fp8(rocke_ir_builder_t* b, - rocke_value_t* a, - rocke_value_t* bb, - rocke_value_t* c, - rocke_value_t* a_scale, - rocke_value_t* b_scale) -{ - rocke_value_t* extra[2]; - extra[0] = a_scale; - extra[1] = b_scale; - return rocke_b_mma(b, "wmma_scale_f32_16x16x128_fp8_fp8", a, bb, c, extra, 2); -} - -rocke_value_t* rocke_b_wmma_scale16_f32_16x16x128_fp8_fp8(rocke_ir_builder_t* b, - rocke_value_t* a, - rocke_value_t* bb, - rocke_value_t* c, - rocke_value_t* a_scale, - rocke_value_t* b_scale) -{ - rocke_value_t* extra[2]; - extra[0] = a_scale; - extra[1] = b_scale; - return rocke_b_mma(b, "wmma_scale16_f32_16x16x128_fp8_fp8", a, bb, c, extra, 2); -} - /* ===================================================================== */ /* register-fragment reshape (P13) */ /* ===================================================================== */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mma.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mma.cpp index 60e7ab17c364..3abfd360a3c5 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/wmma_scale_internal.h" #include /* snprintf */ #include /* atoi */ @@ -75,11 +76,17 @@ static const char* h_elem_scalar(const rocke_type_t* t) * op.result (and the WMMA gate keys off the op_id *string* it is passed, not * op.attrs), so a synthetic op aliasing the same operands/results/regions and * reusing the original attrs map reproduces the Python emission exactly. */ -static rocke_status_t - h_emit_gfx1250_scaled_wmma(rocke_h_lowerer_t* lw, const rocke_op_t* op, bool scale16) +static rocke_status_t h_emit_gfx1250_scaled_wmma(rocke_h_lowerer_t* lw, const rocke_op_t* op) { - const char* op_id - = scale16 ? "wmma_scale16_f32_16x16x128_fp8_fp8" : "wmma_scale_f32_16x16x128_fp8_fp8"; + const rocke_mma_op_t* atom = rocke_gfx1250_scaled_wmma_from_op(op); + if(!atom) + { + return rocke_h_fail(lw, ROCKE_ERR_NOTIMPL, "unsupported scaled WMMA op '%s'", op->name); + } + const rocke_scaled_wmma_op_t contract = rocke_scaled_wmma_contract(atom); + const rocke_scaled_wmma_op_t* spec = &contract; + const char* op_id = spec->op_id; + const bool scale16 = spec->scales.block_k == 16; const char* builtin = scale16 ? "__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4" : "__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4"; if(!lw->arch.gfx || __builtin_strcmp(lw->arch.gfx, "gfx1250") != 0) @@ -95,14 +102,18 @@ static rocke_status_t return rocke_h_fail(lw, ROCKE_ERR_VALUE, "%s expects 5 operands and 1 result", op_id); } rocke_h_emitf(lw, - "f32x8 %s = %s(0, %s, 0, %s, (int16_t)0, %s, " - "0, 0, %s, 0, 0, %s, false, false);", + "f32x8 %s = %s(%d, %s, %d, %s, (int16_t)0, %s, " + "0, %d, %s, 0, %d, %s, false, false);", rocke_h_name(lw, op->results[0]), builtin, + spec->matrix_formats[0], rocke_h_name(lw, op->operands[0]), + spec->matrix_formats[1], rocke_h_name(lw, op->operands[1]), rocke_h_name(lw, op->operands[2]), + spec->scale_formats[0], rocke_h_name(lw, op->operands[3]), + spec->scale_formats[1], rocke_h_name(lw, op->operands[4])); return lw->status; } @@ -123,13 +134,9 @@ static rocke_status_t rocke_h_op_tile_mma(rocke_h_lowerer_t* lw, const rocke_op_ { return rocke_h_fail(lw, ROCKE_ERR_KEY, "tile.mma: missing 'op_id' attr"); } - if(__builtin_strcmp(op_id, "wmma_scale_f32_16x16x128_fp8_fp8") == 0) - { - return h_emit_gfx1250_scaled_wmma(lw, op, false); - } - if(__builtin_strcmp(op_id, "wmma_scale16_f32_16x16x128_fp8_fp8") == 0) + if(rocke_gfx1250_scaled_wmma(op_id)) { - return h_emit_gfx1250_scaled_wmma(lw, op, true); + return h_emit_gfx1250_scaled_wmma(lw, op); } snprintf(dotted, sizeof(dotted), "tile.%s", op_id); legacy_opcode = rocke_opcode_from_name(dotted); 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 1ee6248e7f8d..e6fd436551cb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/data.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/data.cpp @@ -213,11 +213,11 @@ const rocke_ll_decl_t ROCKE_LL_INTRINSIC_DECLS[] = { {"wmma.gfx1250.f32.16x16x64.bf8.bf8", "declare <8 x float> @llvm.amdgcn.wmma.f32.16x16x64.bf8.bf8.v8f32.v8i32(<8 x i32>, <8 x i32>, " "i16 immarg, <8 x float>, i1 immarg, i1 immarg)"}, - {"wmma.scale.gfx1250.f32.16x16x128.fp8.fp8", + {"wmma.scale.block32.gfx1250.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32", "declare <8 x float> @llvm.amdgcn.wmma.scale.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32(i32 " "immarg, <16 x i32>, i32 immarg, <16 x i32>, i16 immarg, <8 x float>, i32 immarg, i32 immarg, " "i32, i32 immarg, i32 immarg, i32, i1 immarg, i1 immarg)"}, - {"wmma.scale16.gfx1250.f32.16x16x128.fp8.fp8", + {"wmma.scale.block16.gfx1250.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32", "declare <8 x float> @llvm.amdgcn.wmma.scale16.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32(i32 " "immarg, <16 x i32>, i32 immarg, <16 x i32>, i16 immarg, <8 x float>, i32 immarg, i32 immarg, " "i64, i32 immarg, i32 immarg, i64, i1 immarg, i1 immarg)"}, 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 0c9c79e1c3d5..b998e9263138 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 @@ -25,6 +25,7 @@ #include "rocke/ir.h" #include "rocke/lower_llvm.h" #include "rocke/lower_llvm_internal.h" +#include "rocke/wmma_scale_internal.h" namespace ckc { @@ -35,7 +36,7 @@ namespace ckc static void _op_tile_wmma_f32_16x16x16_f16(rocke_lower_t* L, const rocke_op_t* op); static void _op_tile_wmma_f32_16x16x16_bf16(rocke_lower_t* L, const rocke_op_t* op); static void _emit_wmma(rocke_lower_t* L, const rocke_op_t* op, const char* op_id); -static void _emit_wmma_scale(rocke_lower_t* L, const rocke_op_t* op, bool scale16); +static void _emit_wmma_scale(rocke_lower_t* L, const rocke_op_t* op); static void _op_tile_mma(rocke_lower_t* L, const rocke_op_t* op); static void _op_tile_mfma_f32_16x16x32_fp8(rocke_lower_t* L, const rocke_op_t* op); static void _op_tile_mfma_f32_16x16x32_bf8(rocke_lower_t* L, const rocke_op_t* op); @@ -134,13 +135,9 @@ static void _op_tile_mma(rocke_lower_t* L, const rocke_op_t* op) * reject (Python ISABackend.emit_wmma raises NotImplementedError). The * gfx12-specific op_ids ("wmma_gfx12_*") can only occur on RDNA4. */ } - else if(strcmp(op_id, "wmma_scale_f32_16x16x128_fp8_fp8") == 0) + else if(rocke_gfx1250_scaled_wmma(op_id)) { - _emit_wmma_scale(L, op, false); - } - else if(strcmp(op_id, "wmma_scale16_f32_16x16x128_fp8_fp8") == 0) - { - _emit_wmma_scale(L, op, true); + _emit_wmma_scale(L, op); } else if(strncmp(op_id, "wmma_", 5) == 0) { @@ -173,10 +170,10 @@ static void _op_tile_mma(rocke_lower_t* L, const rocke_op_t* op) } } -/* gfx1250 native MX FP8 WMMA. ROCm 7.13 introduced the LLVM 23 ABI: - * matrix operands are <16 x i32>; SCALE carries four packed E8M0 bytes in - * each i32 scale operand and SCALE16 carries eight in i64. */ -static void _emit_wmma_scale(rocke_lower_t* L, const rocke_op_t* op, bool scale16) +/* Emit the gfx1250 scaled-WMMA intrinsic using the resolved operand contract. + * The contract supplies packed carrier types, matrix and scale format selectors, + * and the declaration identity. This instruction family requires LLVM 23. */ +static void _emit_wmma_scale(rocke_lower_t* L, const rocke_op_t* op) { const char* intrinsic; const char* decl_key; @@ -187,8 +184,16 @@ static void _emit_wmma_scale(rocke_lower_t* L, const rocke_op_t* op, bool scale1 { return; } - op_name = scale16 ? "tile.wmma_scale16_f32_16x16x128_fp8_fp8" - : "tile.wmma_scale_f32_16x16x128_fp8_fp8"; + const rocke_mma_op_t* atom = rocke_gfx1250_scaled_wmma_from_op(op); + if(!atom) + { + rocke_ll_fail(L, ROCKE_ERR_NOTIMPL, "unsupported scaled WMMA op '%s'", op->name); + } + const rocke_scaled_wmma_op_t contract = rocke_scaled_wmma_contract(atom); + const rocke_scaled_wmma_op_t* spec = &contract; + char concrete_name[160]; + snprintf(concrete_name, sizeof(concrete_name), "tile.%s", spec->op_id); + op_name = concrete_name; if(!L->backend || strcmp(L->backend->gfx, "gfx1250") != 0) { rocke_ll_fail(L, @@ -211,7 +216,9 @@ static void _emit_wmma_scale(rocke_lower_t* L, const rocke_op_t* op, bool scale1 L, ROCKE_ERR_VALUE, "%s expects 5 operands, got %d", op_name, op->num_operands); } - scale_ty = scale16 ? "i64" : "i32"; + char packed_type[8]; + snprintf(packed_type, sizeof(packed_type), "i%d", rocke_scale_word_bits(&spec->scales)); + scale_ty = packed_type; if(strcmp(op->operands[3]->type->name, scale_ty) != 0 || strcmp(op->operands[4]->type->name, scale_ty) != 0) { @@ -224,29 +231,28 @@ static void _emit_wmma_scale(rocke_lower_t* L, const rocke_op_t* op, bool scale1 op->operands[4]->type->name); } - if(scale16) - { - decl_key = "wmma.scale16.gfx1250.f32.16x16x128.fp8.fp8"; - intrinsic = "llvm.amdgcn.wmma.scale16.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32"; - } - else - { - decl_key = "wmma.scale.gfx1250.f32.16x16x128.fp8.fp8"; - intrinsic = "llvm.amdgcn.wmma.scale.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32"; - } + // Physical declaration identity comes from the resolved backend contract. + decl_key = spec->declaration_key; + intrinsic = spec->intrinsic; rocke_ll_need(L, decl_key); rocke_ll_emitf(L, " %s = call <8 x float> @%s(" - "i32 0, <16 x i32> %s, i32 0, <16 x i32> %s, " - "i16 0, <8 x float> %s, i32 0, i32 0, %s %s, " - "i32 0, i32 0, %s %s, i1 false, i1 false)", + "i32 %d, <%d x i32> %s, i32 %d, <%d x i32> %s, " + "i16 0, <8 x float> %s, i32 0, i32 %d, %s %s, " + "i32 0, i32 %d, %s %s, i1 false, i1 false)", mma_result_name(L, op), intrinsic, + spec->matrix_formats[0], + spec->matrix_words[0], rocke_ll_operand(L, op->operands[0]), + spec->matrix_formats[1], + spec->matrix_words[1], rocke_ll_operand(L, op->operands[1]), rocke_ll_operand(L, op->operands[2]), + spec->scale_formats[0], scale_ty, rocke_ll_operand(L, op->operands[3]), + spec->scale_formats[1], scale_ty, rocke_ll_operand(L, op->operands[4])); } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/arch_target.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/arch_target.h index f7c9e238d4f5..12c52d012c2b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/arch_target.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/arch_target.h @@ -22,7 +22,6 @@ * class ResourceLimits rocke_resource_limits_t * class MmaCatalog rocke_mma_catalog_t + rocke_mma_catalog_* () * class ArchTarget (frozen) rocke_arch_target_t + rocke_arch_* () getters - * normalize_dtype() rocke_normalize_dtype() * ArchTarget.from_gfx(gfx) rocke_arch_target_from_gfx() * known_arches() rocke_known_arches() * arch_from_isa(isa) rocke_arch_from_isa() @@ -50,34 +49,24 @@ #include #include +#include "rocke/dtypes.h" #include "rocke/ir.h" #ifdef __cplusplus extern "C" { #endif -/* ============================== dtype keys ============================== */ - -/* Map a dtype spelling ("f16"/"half"/"fp16" -> "fp16", ...) to its canonical - * catalog key. Returns a pointer to a static, interned canonical string when the - * spelling is recognised; otherwise returns the *lowercased* spelling stored in - * `scratch` (caller-provided buffer of >= scratch_cap bytes), so unknown - * spellings pass through Python-identically. `scratch` may be NULL only if the - * caller guarantees a known spelling; pass a buffer to be safe. - * - * Mirrors target.py::normalize_dtype (strip + lower + _DTYPE_ALIASES.get). */ -const char* rocke_normalize_dtype(const char* name, char* scratch, size_t scratch_cap); - /* ============================== layout map ============================== */ -/* MMA fragment role. Mirrors LayoutMap.role ("acc"/"a"/"b"). The accumulator is - * spelled "c" in the C op_id maps (MmaOp.c_layout) but role text is "acc"/"a"/ - * "b" exactly as Python LayoutMap.role. */ +/* MMA fragment role. Matrix roles mirror LayoutMap.role ("acc"/"a"/"b"); + * scale roles are "a_scale"/"b_scale" and use logical K-group coordinates. */ typedef enum rocke_mma_role { ROCKE_MMA_ROLE_ACC = 0, /* accumulator C/D: coords (row, col) */ ROCKE_MMA_ROLE_A, /* A operand: coords (row, k) */ - ROCKE_MMA_ROLE_B /* B operand: coords (k, col) */ + ROCKE_MMA_ROLE_B, /* B operand: coords (k, col) */ + ROCKE_MMA_ROLE_A_SCALE, /* A scale: coords (row, K-group) */ + ROCKE_MMA_ROLE_B_SCALE /* B scale: coords (K-group, col) */ } rocke_mma_role_t; /* The lane/slot -> tile-coordinate emitter. Given the builder, a runtime i32 @@ -115,12 +104,28 @@ bool rocke_layout_map_coord(const rocke_layout_map_t* m, /* ============================== MMA atom =============================== */ -/* A single supported matrix-multiply-accumulate atom on a target. Frozen; lives - * in static storage. Mirrors MmaOp. `family` is "mma"|"wmma"; dtype fields hold - * canonical (normalised) catalog keys; op_id is the opaque backend handle (the - * ISA-named IRBuilder method, e.g. "mfma_f32_16x16x16_f16"). The *_layout - * pointers are NULL when no verified lane map is registered for the op_id (the - * Python `_*_layout is None` / NotImplementedError case). */ +/* Shared K-group size. NONE is only the unscaled sentinel. */ +typedef enum rocke_mma_scale_block_k +{ + ROCKE_MMA_SCALE_NONE = 0, + ROCKE_MMA_SCALE_K16 = 16, + ROCKE_MMA_SCALE_K32 = 32 +} rocke_mma_scale_block_k_t; + +/* Complete scale query: independent A/B types, shared K-group size. + * {NULL, NULL, 0} selects unscaled atoms. */ +typedef struct rocke_mma_scale_filter +{ + const char* a_scale_dtype; + const char* b_scale_dtype; + rocke_mma_scale_block_k_t scale_block_k; +} rocke_mma_scale_filter_t; + +/* A supported matrix-multiply-accumulate atom. Frozen, in static storage. + * Mirrors MmaOp: family is "mma", "wmma", or "wmma_scaled"; matrix dtypes + * are canonical catalog keys. C describes both accumulator input and result. + * op_id is an opaque backend handle; lowering reads metadata, not its spelling. + * Layout pointers are NULL when no verified lane map is registered. */ typedef struct rocke_mma_op { const char* family; @@ -138,21 +143,33 @@ typedef struct rocke_mma_op const rocke_layout_map_t* a_layout; /* may be NULL */ const rocke_layout_map_t* b_layout; /* may be NULL */ const rocke_layout_map_t* c_layout; /* may be NULL */ + const char* a_scale_dtype; /* NULL when unscaled */ + const char* b_scale_dtype; /* NULL when unscaled */ + rocke_mma_scale_block_k_t scale_block_k; /* shared by A/B; NONE when unscaled */ + /* Logical scale elements per lane, independent of the packed carrier. + * Appended ABI fields: native consumers must rebuild. */ + int a_scale_frag_len; + int b_scale_frag_len; + const rocke_layout_map_t* a_scale_layout; /* may be NULL */ + const rocke_layout_map_t* b_scale_layout; /* may be NULL */ } rocke_mma_op_t; /* MmaOp.shape -> (m, n, k) via out params. */ void rocke_mma_op_shape(const rocke_mma_op_t* op, int* m, int* n, int* k); -/* Physical-layout accessors. Each returns the verified map for the role, or NULL - * when none is registered. Unlike Python (which raises NotImplementedError), the - * C getters return NULL and -- when `b` is non-NULL -- set the builder's sticky - * error (ROCKE_ERR_NOTIMPL) with the same message text, so callers can either - * check NULL or rely on the sticky-fail builder. Pass b=NULL for a pure lookup. - * Mirrors MmaOp.a_layout / b_layout / c_layout / acc_layout. */ +/* Physical-layout accessors. Missing maps raise ckc::Error with + * ROCKE_ERR_NOTIMPL, like Python's NotImplementedError. The enclosing public + * builder/lowering boundary translates that exception into status + message; + * direct C++ callers must catch it, including when b is NULL. A NULL op returns + * NULL. Mirrors MmaOp's matrix and scale layout accessors. */ const rocke_layout_map_t* rocke_mma_op_a_layout(const rocke_mma_op_t* op, rocke_ir_builder_t* b); const rocke_layout_map_t* rocke_mma_op_b_layout(const rocke_mma_op_t* op, rocke_ir_builder_t* b); const rocke_layout_map_t* rocke_mma_op_c_layout(const rocke_mma_op_t* op, rocke_ir_builder_t* b); const rocke_layout_map_t* rocke_mma_op_acc_layout(const rocke_mma_op_t* op, rocke_ir_builder_t* b); +const rocke_layout_map_t* rocke_mma_op_a_scale_layout(const rocke_mma_op_t* op, + rocke_ir_builder_t* b); +const rocke_layout_map_t* rocke_mma_op_b_scale_layout(const rocke_mma_op_t* op, + rocke_ir_builder_t* b); /* ====================== memory caps / resource limits ================== */ @@ -184,6 +201,16 @@ typedef struct rocke_mma_catalog /* MmaCatalog.ops accessor (count + pointer). */ const rocke_mma_op_t* rocke_mma_catalog_ops(const rocke_mma_catalog_t* cat, int* num_out); +/* Optional scale filter: NULL leaves scales unconstrained; {NULL, NULL, 0} + * selects unscaled atoms. C callers supply the filter pointer explicitly. + * Exact selection and largest-K ties raise a query error on ambiguity. + * As with layout queries, C++ callers catch ckc::Error at their entry boundary. */ +#ifdef __cplusplus +#define ROCKE_MMA_DEFAULT_SCALES = nullptr +#else +#define ROCKE_MMA_DEFAULT_SCALES +#endif + /* MmaCatalog.enumerate: filter atoms by family + (normalised) dtype combo and * optional m/n. Writes up to `cap` matching op pointers into `out` and returns * the total number of matches (which may exceed `cap` -- the caller can size a @@ -197,7 +224,8 @@ int rocke_mma_catalog_enumerate(const rocke_mma_catalog_t* cat, int m, int n, const rocke_mma_op_t** out, - int cap); + int cap, + const rocke_mma_scale_filter_t* scales ROCKE_MMA_DEFAULT_SCALES); /* MmaCatalog.has_shape. family NULL => "mma". */ bool rocke_mma_catalog_has_shape(const rocke_mma_catalog_t* cat, @@ -207,7 +235,8 @@ bool rocke_mma_catalog_has_shape(const rocke_mma_catalog_t* cat, const char* c_dtype, int m, int n, - int k); + int k, + const rocke_mma_scale_filter_t* scales ROCKE_MMA_DEFAULT_SCALES); /* MmaCatalog.select_largest_k: the matching atom with the largest k (k <= k_max * if k_max >= 0; pass k_max<0 for Python None). Returns NULL if none match. */ @@ -218,21 +247,25 @@ const rocke_mma_op_t* rocke_mma_catalog_select_largest_k(const rocke_mma_catalog const char* c_dtype, int m, int n, - int k_max); + int k_max, + const rocke_mma_scale_filter_t* scales + ROCKE_MMA_DEFAULT_SCALES); /* MmaCatalog.by_op_id. Returns NULL if absent. */ const rocke_mma_op_t* rocke_mma_catalog_by_op_id(const rocke_mma_catalog_t* cat, const char* op_id); /* MmaCatalog.op_for_shape: the atom with exactly (m, n, k). NULL if absent. * family NULL => "mma". */ -const rocke_mma_op_t* rocke_mma_catalog_op_for_shape(const rocke_mma_catalog_t* cat, - const char* family, - const char* a_dtype, - const char* b_dtype, - const char* c_dtype, - int m, - int n, - int k); +const rocke_mma_op_t* + rocke_mma_catalog_op_for_shape(const rocke_mma_catalog_t* cat, + const char* family, + const char* a_dtype, + const char* b_dtype, + const char* c_dtype, + int m, + int n, + int k, + const rocke_mma_scale_filter_t* scales ROCKE_MMA_DEFAULT_SCALES); /* ===================== bare-op_id SSOT lookups ======================== */ /* These resolve a bare op_id string (no MmaOp / ArchTarget in hand) against the @@ -240,19 +273,14 @@ const rocke_mma_op_t* rocke_mma_catalog_op_for_shape(const rocke_mma_catalog_t* * any atom metadata. They mirror the target.py module-level helpers that * rocke.core.ir.IRBuilder.mma consults. */ +#undef ROCKE_MMA_DEFAULT_SCALES + /* Accumulator fragment length for op_id -- the c_frag_len projection of * target.py::_MMA_FRAGMENT_INFO (rocke_arch_mma_c_frag_len(op_id) == * target._frag_info(op_id).c_frag_len). Returns 0 for an unknown atom (the * zero-length _FragInfo fallback). */ int rocke_arch_mma_c_frag_len(const char* op_id); -/* Canonical (normalised) accumulator-dtype key, as produced by - * rocke_normalize_dtype / the _DTYPE_ALIASES table. Compare an op_id's - * accumulator dtype against this instead of a raw string literal so the spelling - * has a single definition site shared by producer (the normalise table) and - * consumers (e.g. the integer-accumulator predicate in ir_tile). */ -#define ROCKE_DTYPE_I32 "i32" - /* Normalised accumulator dtype for op_id, aggregated across every arch's * catalog (mirrors target._op_id_c_dtype()[op_id]). The dtype is invariant * across the arches that list an op_id, so the first catalog hit wins. Returns @@ -260,6 +288,10 @@ int rocke_arch_mma_c_frag_len(const char* op_id); * interned canonical key (e.g. ROCKE_DTYPE_I32 / "fp32"); do not free. */ const char* rocke_arch_mma_op_id_c_dtype(const char* op_id); +/* Target-independent family projection. Unknown/NULL IDs return NULL; + * conflicting families for one ID raise ROCKE_ERR_VALUE. */ +const char* rocke_arch_mma_op_id_family(const char* op_id); + /* ============================== arch target =========================== */ /* Hardware-facts surface for one gfx target. Frozen; cheap to pass around. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/arch_target_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/arch_target_internal.h index 1a272f900f9b..2c492d5021db 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/arch_target_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/arch_target_internal.h @@ -56,15 +56,6 @@ extern const int rocke_ati_arch_registry_len; * in sync. */ extern const char* const rocke_ati_known_arches[]; -/* ------------------------------------------------------- dtype normalisation */ - -/* The canonical-key core of rocke_normalize_dtype, shared so both buckets resolve - * spellings identically. `lowered` is a caller buffer that receives the - * strip+lower of `name` (used as the pass-through result for unknown spellings). - * Returns a static interned canonical string for a known alias, else `lowered`. - * Defined in bucket 0. */ -const char* rocke_ati_normalize_dtype(const char* name, char* lowered, size_t cap); - #ifdef __cplusplus } /* extern "C" */ #endif diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/dtypes.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/dtypes.h new file mode 100644 index 000000000000..f7da21efad31 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/dtypes.h @@ -0,0 +1,35 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* Target-independent dtype naming API. Mirrors rocke.core.dtypes: + * normalize_dtype -> rocke_normalize_dtype. + * Recognized names do not imply IR scalar or target instruction support. + */ +#ifndef ROCKE_DTYPES_H +#define ROCKE_DTYPES_H + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/* Canonical integer dtype key shared by normalization and its consumers. */ +#define ROCKE_DTYPE_I32 "i32" + +/* ============================== dtype keys ============================== */ + +/* Map a dtype spelling ("f16"/"half"/"fp16" -> "fp16", ...) to its canonical + * catalog key. Returns a pointer to a static, interned canonical string when the + * spelling is recognised; otherwise returns the *lowercased* spelling stored in + * `scratch` (caller-provided buffer of >= scratch_cap bytes), so unknown + * spellings pass through Python-identically. `scratch` may be NULL only if the + * caller guarantees a known spelling; pass a buffer to be safe. + * + * Mirrors core/dtypes.py::normalize_dtype (strip + lower + _DTYPE_ALIASES.get). */ +const char* rocke_normalize_dtype(const char* name, char* scratch, size_t scratch_cap); + +#ifdef __cplusplus +} /* extern "C" */ +#endif + +#endif /* ROCKE_DTYPES_H */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.core.arch.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.core.arch.h index 1fc0f9462403..ec228fb445aa 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.core.arch.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.core.arch.h @@ -95,6 +95,9 @@ const rocke_arch_mma_catalog_t* rocke_archtarget_mma(const rocke_archtarget_t* t /* target.mma.op_for_shape(...): the resolved MmaOp for an exact (m, n, k) atom * shape and (normalised) dtype combo, or NULL if absent. `family` NULL => "mma". + * The trailing scales argument is required: NULL leaves scales unconstrained; + * {NULL, NULL, ROCKE_MMA_SCALE_NONE} selects unscaled atoms. Invalid filters or + * ambiguous matches raise ckc::Error; callers handle errors at a C++ boundary. * This is the entry point build_universal_gemm uses to resolve the atom from the * target catalog. Forwards to rocke_mma_catalog_op_for_shape on t->mma. */ const rocke_mmaop_t* rocke_archtarget_op_for_shape(const rocke_archtarget_t* t, @@ -104,7 +107,8 @@ const rocke_mmaop_t* rocke_archtarget_op_for_shape(const rocke_archtarget_t* t, const char* c_dtype, int m, int n, - int k); + int k, + const rocke_mma_scale_filter_t* scales); /* target.mma.by_op_id(op_id): the catalog atom whose op_id handle matches * `op_id` (the backend's MMA key, e.g. "mfma_f32_16x16x16_f16"), or NULL if the 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 092330b16a11..59ae31feced3 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 @@ -1053,18 +1053,6 @@ rocke_value_t* rocke_b_mfma_scale_f32_16x16x128_f8f6f4(rocke_ir_builder_t* b, rocke_value_t* c, rocke_value_t* a_scale, rocke_value_t* b_scale); -rocke_value_t* rocke_b_wmma_scale_f32_16x16x128_fp8_fp8(rocke_ir_builder_t* b, - rocke_value_t* a, - rocke_value_t* bb, - rocke_value_t* c, - rocke_value_t* a_scale, - rocke_value_t* b_scale); -rocke_value_t* rocke_b_wmma_scale16_f32_16x16x128_fp8_fp8(rocke_ir_builder_t* b, - rocke_value_t* a, - rocke_value_t* bb, - rocke_value_t* c, - rocke_value_t* a_scale, - rocke_value_t* b_scale); rocke_value_t* rocke_b_wmma_f32_16x16x16_f16(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_t* bb, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/wmma_scale_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/wmma_scale_internal.h new file mode 100644 index 000000000000..c85df088e922 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/wmma_scale_internal.h @@ -0,0 +1,110 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* Shared gfx1250 operand contracts. Mirrors core/arch/wmma_scale.py: + * ScalePacking -> rocke_scale_packing_t + * ScaledWmmaOp -> rocke_scaled_wmma_op_t + * gfx1250_scaled_wmma -> rocke_gfx1250_scaled_wmma + */ +#ifndef ROCKE_WMMA_SCALE_INTERNAL_H +#define ROCKE_WMMA_SCALE_INTERNAL_H + +#include +#include + +#include "rocke/arch_target.h" +#include "rocke/error.hpp" +#include "rocke/ir.h" + +typedef struct rocke_scale_packing +{ + /* Eight-bit encoded scales for consecutive K groups, first group in the + * low byte, independently for A/B. */ + int count; + int block_k; +} rocke_scale_packing_t; + +typedef struct rocke_scaled_wmma_op +{ + const char* op_id; + int matrix_formats[2]; + int scale_formats[2]; + int matrix_words[2]; + char declaration_key[160]; + char intrinsic[160]; + rocke_scale_packing_t scales; +} rocke_scaled_wmma_op_t; + +static inline const rocke_mma_op_t* rocke_gfx1250_scaled_wmma(const char* op_id) +{ + if(!op_id) + return NULL; + if(strncmp(op_id, "tile.", 5) == 0) + op_id += 5; + const rocke_arch_target_t* target = rocke_arch_target_from_gfx("gfx1250"); + const rocke_mma_op_t* atom = rocke_mma_catalog_by_op_id(&target->mma, op_id); + return atom && strcmp(atom->family, "wmma_scaled") == 0 ? atom : NULL; +} + +static inline rocke_scaled_wmma_op_t rocke_scaled_wmma_contract(const rocke_mma_op_t* atom) +{ + rocke_scaled_wmma_op_t spec = {}; + spec.op_id = atom->op_id; + const char* dtypes[2] = {atom->a_dtype, atom->b_dtype}; + const int words[2] = {atom->a_frag_len, atom->b_frag_len}; + const char* scale_dtypes[2] = {atom->a_scale_dtype, atom->b_scale_dtype}; + for(int i = 0; i < 2; ++i) + { + if(strcmp(dtypes[i], "fp8e4m3") == 0) + spec.matrix_formats[i] = 0; + else if(strcmp(dtypes[i], "bf8e5m2") == 0) + spec.matrix_formats[i] = 1; + else + ckc::raise_status(ROCKE_ERR_VALUE, "unsupported scaled WMMA matrix format"); + if(!scale_dtypes[i] || strcmp(scale_dtypes[i], "e8m0") != 0) + ckc::raise_status(ROCKE_ERR_VALUE, "unsupported scaled WMMA scale format"); + spec.scale_formats[i] = 0; // E8M0. + spec.matrix_words[i] = words[i]; + } + if((atom->scale_block_k != 16 && atom->scale_block_k != 32) + || strcmp(atom->c_dtype, "fp32") != 0 || atom->m != 16 || atom->n != 16 || atom->k != 128) + ckc::raise_status(ROCKE_ERR_VALUE, "unsupported scaled WMMA backend contract"); + spec.scales.block_k = atom->scale_block_k; + const int count = atom->k / spec.scales.block_k; + if(atom->a_scale_frag_len != count || atom->b_scale_frag_len != count) + ckc::raise_status(ROCKE_ERR_VALUE, "unsupported scaled WMMA scale fragment lengths"); + spec.scales.count = atom->a_scale_frag_len; + char suffix[96]; + snprintf(suffix, + sizeof(suffix), + "f32.%dx%dx%d.f8f6f4.v%df32.v%di32.v%di32", + atom->m, + atom->n, + atom->k, + atom->c_frag_len, + spec.matrix_words[0], + spec.matrix_words[1]); + snprintf(spec.intrinsic, + sizeof(spec.intrinsic), + "llvm.amdgcn.wmma.%s.%s", + spec.scales.block_k == 16 ? "scale16" : "scale", + suffix); + snprintf(spec.declaration_key, + sizeof(spec.declaration_key), + "wmma.scale.block%d.gfx1250.%s", + spec.scales.block_k, + suffix); + return spec; +} + +static inline const rocke_mma_op_t* rocke_gfx1250_scaled_wmma_from_op(const rocke_op_t* op) +{ + const char* op_id = rocke_attr_get_str(&op->attrs, "op_id"); + return rocke_gfx1250_scaled_wmma(op_id ? op_id : op->name); +} + +static inline int rocke_scale_word_bits(const rocke_scale_packing_t* packing) +{ + return packing->count * 8; +} + +#endif /* ROCKE_WMMA_SCALE_INTERNAL_H */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/_matmul_nbits_common.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/_matmul_nbits_common.cpp index df41776315e3..e8233bf31b79 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/_matmul_nbits_common.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/_matmul_nbits_common.cpp @@ -306,7 +306,8 @@ bool rocke_matmul_nbits_validate_common_spec(const rocke_matmul_nbits_spec_t* sp ROCKE__WMMA_C_DTYPE, t->warp_tile_m, t->warp_tile_n, - t->warp_tile_k) + t->warp_tile_k, + nullptr) == NULL) { ROCKE_NBITS_REJECT("unsupported fp16 warp_tile (%d, %d, %d) on %s", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_spec_descriptors.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_spec_descriptors.cpp index cf3cded44af4..9b72cb723485 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_spec_descriptors.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_spec_descriptors.cpp @@ -959,7 +959,8 @@ const rocke_mmaop_t* rocke_conv_resolve_op(rocke_ir_builder_t* b, "fp32", spec->warp_tile_m, spec->warp_tile_n, - spec->warp_tile_k); + spec->warp_tile_k, + nullptr); } if(op == NULL) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_dgrad.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_dgrad.cpp index 52f30b6e97e9..70436929f1b8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_dgrad.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_dgrad.cpp @@ -857,7 +857,8 @@ static const rocke_mmaop_t* "fp32", spec->warp_tile_m, spec->warp_tile_n, - spec->warp_tile_k); + spec->warp_tile_k, + nullptr); if(!op) { rocke_i_set_err(b, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/fmha_arch.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/fmha_arch.cpp index fa15c2ec20be..b3356cbc1eef 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/fmha_arch.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/fmha_arch.cpp @@ -135,7 +135,8 @@ bool rocke_validate_fmha_mfma_atom(const char* dtype, const char* arch, char* ou * has_shape(m,n,k) is true iff a matching op exists; since enumerate() * filters on m,n and has_shape compares the full (m,n,k) shape, this is * exactly op_for_shape(...,k) != None. family defaults to "mma". */ - if(rocke_archtarget_op_for_shape(target, "mma", a_name, a_name, "fp32", blk, blk, blk) == NULL) + if(rocke_archtarget_op_for_shape(target, "mma", a_name, a_name, "fp32", blk, blk, blk, nullptr) + == NULL) { if(out != NULL && out_cap != 0) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/gemm_spec.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/gemm_spec.cpp index 2ba203aabc7d..cb28bb0feaa7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/gemm_spec.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/gemm_spec.cpp @@ -415,7 +415,8 @@ static const rocke_type_t* "fp32", t->warp_tile_m, t->warp_tile_n, - t->warp_tile_k); + t->warp_tile_k, + nullptr); } /* _ab_lds_plan(spec, arch) -> (ab_single, db, two_buf). Pure ints/bools. */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/gemm_universal.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/gemm_universal.cpp index a07976eef143..8dc0c91ded87 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/gemm_universal.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/gemm_universal.cpp @@ -81,7 +81,8 @@ static const rocke_mmaop_t* rocke_gemm_resolve_mma_op(const rocke_gemm_universal "fp32", t->warp_tile_m, t->warp_tile_n, - t->warp_tile_k); + t->warp_tile_k, + nullptr); } /* _storage_dtype(spec): validate homogeneous A/B/C + fp32 acc + RCR layout, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/mfma_gemm.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/mfma_gemm.cpp index 2de604ddb873..5fae31871d91 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/mfma_gemm.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/mfma_gemm.cpp @@ -254,7 +254,7 @@ bool rocke_mfma_gemm_is_valid_spec(const rocke_mfma_gemm_spec_t* spec, /* if not target.mma.has_shape(a=cat, b=cat, c="fp32", m, n, k): ... * op_for_shape returns NULL when the shape/dtype combo is absent. */ if(rocke_archtarget_op_for_shape( - target, "mma", cat_dtype, cat_dtype, "fp32", atom->m, atom->n, atom->k) + target, "mma", cat_dtype, cat_dtype, "fp32", atom->m, atom->n, atom->k, nullptr) == NULL) { snprintf(buf, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx1151/deep_fused_conv_pool_build_context_ctx_init_quant_primitives.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx1151/deep_fused_conv_pool_build_context_ctx_init_quant_primitives.cpp index 08a245652c5e..851bb2f993e8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx1151/deep_fused_conv_pool_build_context_ctx_init_quant_primitives.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx1151/deep_fused_conv_pool_build_context_ctx_init_quant_primitives.cpp @@ -576,7 +576,8 @@ rocke_status_t "fp32", ROCKE_GFX1151_DFCP_WMMA, ROCKE_GFX1151_DFCP_WMMA, - ROCKE_GFX1151_DFCP_WMMA); + ROCKE_GFX1151_DFCP_WMMA, + nullptr); /* op0 = target.mma.by_op_id(_OP_ID_IU8) if spec.native_int else op * op1 = target.mma.by_op_id(_OP_ID_IU4) if spec.native_int else op */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx1151/wmma_gemm.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx1151/wmma_gemm.cpp index c7bf55eae136..4e456a9b5b9c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx1151/wmma_gemm.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx1151/wmma_gemm.cpp @@ -145,7 +145,8 @@ bool rocke_wmma_gemm_gfx1151_is_valid_spec(const rocke_wmma_gemm_gfx1151_spec_t* "fp32", ROCKE_WMMA_M, ROCKE_WMMA_N, - ROCKE_WMMA_K) + ROCKE_WMMA_K, + nullptr) == NULL) { snprintf(buf, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx1151/wmma_gemm_int8.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx1151/wmma_gemm_int8.cpp index cad79e6eb429..fb6b0c138189 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx1151/wmma_gemm_int8.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx1151/wmma_gemm_int8.cpp @@ -133,8 +133,15 @@ bool rocke_wmma_gemm_int8_is_valid_spec(const rocke_wmma_gemm_int8_spec_t* spec, * Operands are int8 in memory but dequantized to f16 before the MMA, so the * *compute* atom we gate on is the fp16 WMMA 16x16x16 one. * op_for_shape returns NULL when the shape/dtype combo is absent. */ - if(rocke_archtarget_op_for_shape( - target, "wmma", "fp16", "fp16", "fp32", ROCKE_WMMA_M, ROCKE_WMMA_N, ROCKE_WMMA_K) + if(rocke_archtarget_op_for_shape(target, + "wmma", + "fp16", + "fp16", + "fp32", + ROCKE_WMMA_M, + ROCKE_WMMA_N, + ROCKE_WMMA_K, + nullptr) == NULL) { snprintf(buf, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx1201/wmma_gemm.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx1201/wmma_gemm.cpp index 21748c6fc3bf..d7fa0fdb44bc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx1201/wmma_gemm.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx1201/wmma_gemm.cpp @@ -136,7 +136,8 @@ bool rocke_wmma_gemm_gfx1201_is_valid_spec(const rocke_wmma_gemm_gfx1201_spec_t* "fp32", ROCKE_WMMA_M, ROCKE_WMMA_N, - ROCKE_WMMA_K) + ROCKE_WMMA_K, + nullptr) == NULL) { snprintf(buf, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp index fa786496833c..4b9fb4e6cf24 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp @@ -101,9 +101,9 @@ static void rocke_attn3d_set_err_buf(char* err, size_t err_cap, const char* msg) static bool rocke_attn3d_narrow_k_available(const rocke_archtarget_t* t) { const rocke_mmaop_t* f16 - = rocke_archtarget_op_for_shape(t, "mma", "f16", "f16", "fp32", 16, 16, 16); + = rocke_archtarget_op_for_shape(t, "mma", "f16", "f16", "fp32", 16, 16, 16, nullptr); const rocke_mmaop_t* bf16 - = rocke_archtarget_op_for_shape(t, "mma", "bf16", "bf16", "fp32", 16, 16, 16); + = rocke_archtarget_op_for_shape(t, "mma", "bf16", "bf16", "fp32", 16, 16, 16, nullptr); return f16 != NULL && bf16 != NULL; } @@ -112,7 +112,8 @@ static bool rocke_attn3d_narrow_k_available(const rocke_archtarget_t* t) * exposed by the C arch surface -- see the ARCH GATE NOTE above). */ static bool rocke_attn3d_wide_k_available(const rocke_archtarget_t* t) { - return rocke_archtarget_op_for_shape(t, "mma", "f16", "f16", "fp32", 16, 16, 32) != NULL; + return rocke_archtarget_op_for_shape(t, "mma", "f16", "f16", "fp32", 16, 16, 32, nullptr) + != NULL; } /* ===================================================================== * diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp index 66da4b755d44..26dd78a5f54b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp @@ -120,9 +120,9 @@ static void rocke_attn3d950_set_err_buf(char* err, size_t err_cap, const char* m static bool rocke_attn3d950_narrow_k_available(const rocke_archtarget_t* t) { const rocke_mmaop_t* f16 - = rocke_archtarget_op_for_shape(t, "mma", "f16", "f16", "fp32", 16, 16, 16); + = rocke_archtarget_op_for_shape(t, "mma", "f16", "f16", "fp32", 16, 16, 16, nullptr); const rocke_mmaop_t* bf16 - = rocke_archtarget_op_for_shape(t, "mma", "bf16", "bf16", "fp32", 16, 16, 16); + = rocke_archtarget_op_for_shape(t, "mma", "bf16", "bf16", "fp32", 16, 16, 16, nullptr); return f16 != NULL && bf16 != NULL; } @@ -131,7 +131,8 @@ static bool rocke_attn3d950_narrow_k_available(const rocke_archtarget_t* t) * exposed by the C arch surface -- see the ARCH GATE NOTE above). */ static bool rocke_attn3d950_wide_k_available(const rocke_archtarget_t* t) { - return rocke_archtarget_op_for_shape(t, "mma", "f16", "f16", "fp32", 16, 16, 32) != NULL; + return rocke_archtarget_op_for_shape(t, "mma", "f16", "f16", "fp32", 16, 16, 32, nullptr) + != NULL; } /* validate_tiled_attention_arch(arch) gate (shared by supports + config). On diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/authoring_model.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/authoring_model.md index 197f7782e90b..0fd7996f116b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/authoring_model.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/authoring_model.md @@ -29,7 +29,12 @@ out separately where they differ. Matrix-operation choices start from the exact gfx target. Resolve `ArchTarget.from_gfx(...)`, select an `MmaOp` from that target's `MmaCatalog`, and validate the operation's wave and layout contract. The exact catalog is -listed in [`kernel_taxonomy.md`](./kernel_taxonomy.md). +listed in [`kernel_taxonomy.md`](./kernel_taxonomy.md). For operand metadata, +scale layouts, query semantics, and migration details, see +[MMA metadata and queries](../reference/mma_metadata.md). +Native builders can use `rocke_archtarget_op_for_shape`; see +[native C/C++ atom lookup](../reference/mma_metadata.md#native-cc-atom-lookup) +for the helper's contract and a complete example. Kernel definitions in [`/python/rocke/instances/`](../../python/rocke/instances/) and diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md index 284a7c75448b..2856aad6821c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md @@ -2,6 +2,9 @@ Top-level re-exports from `rocke` and `rocke.helpers`. Use this as a quick lookup when reading other docs. +For `MmaOp` fields, scale layouts, and `MmaCatalog` selection rules, see +[MMA metadata and queries](mma_metadata.md). + ## `from rocke import ...` ```text 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 new file mode 100644 index 000000000000..41642b1170ce --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/mma_metadata.md @@ -0,0 +1,130 @@ +# MMA Metadata and Queries + +This reference describes the `MmaOp` metadata and `MmaCatalog` queries used to +select matrix operations. For the overall kernel-building workflow, see the +[authoring model](../architecture/authoring_model.md). + +## Operand Metadata + +MMA metadata uses `a_dtype`, `b_dtype`, and `c_dtype`, with C describing both +the accumulator input and result. Optional `a_scale_dtype`, `b_scale_dtype`, +and `scale_block_k` fields describe independent scale types and a shared K-group +size. `MmaScaleBlockK.K16` and `MmaScaleBlockK.K32` are the only scaled sizes; +all three fields are `None` for unscaled atoms. Integer 16/32 values from JSON +are normalized to the enum. Scale types are independent of matrix dtypes and +packed register types. Matrix fragment and layout accessors retain their A/B/C roles. + +Scale dtypes normalize to the string-backed `MmaScaleDType` enum: `E8M0`, +`E4M3`, or `E5M3`. Existing string inputs remain accepted, and string formatting +and JSON retain the canonical `e8m0`, `e4m3`, and `e5m3` spellings. +`fp8e4m3` is an alias for `E4M3`. E5M3 is an unsigned scale format, distinct from +the signed E5M2 matrix format: `e5m2`, `fp8e5m2`, `bf8e5m2`, and `bf8` are +rejected as scale dtypes. LLVM defines separate matrix and scale format names +in [WMMAMods](https://github.com/ROCm/llvm-project/blob/6bd443a039454b7d2d9f34740ce0a6e3c10ac8d4/llvm/lib/Target/AMDGPU/Utils/AMDGPUAsmUtils.h#L146). +Enum membership describes a format; supported operations still come from the +target catalog and backend. + +## Scale Layouts + +`a_scale_frag_len` and `b_scale_frag_len` count logical scale elements per lane. +`a_scale_layout()` maps `(lane, slot)` to `(row, K-group)`; +`b_scale_layout()` maps it to `(K-group, column)`. These instruction-local maps +are separate from register byte packing and global tensor strides. Unscaled +atoms default to zero scale fragments and absent maps. An unavailable map raises +`NotImplementedError`, including for a scaled atom whose layout is not verified. +The native record appends corresponding counts and map pointers, so native +consumers must rebuild; existing role enum values remain unchanged. + +### gfx1250 Scale Mapping and Packing + +The four gfx1250 FP8/BF8 scaled atoms provide these scale maps: lane `l` and +slot `j` map to `(l % 16, j)` for A and `(j, l % 16)` for B. Both half-waves +duplicate the scales. K32 uses four E8M0 elements per lane packed into i32; +K16 uses eight packed into i64, with slot `j` at bit offset `8*j`. Matrix A/B +maps for these atoms remain unavailable. The gfx1250 loader applies tile and +instruction-step offsets to the scale coordinates before applying tensor strides. + +## Catalog Queries + +For scaled operations, query the complete contract and pass the selected atom +to `IRBuilder.mma`: + +```python +from rocke.core.arch import ArchTarget, MmaScaleBlockK, MmaScaleDType + +atom = ArchTarget.from_gfx("gfx1250").mma.op_for_shape( + family="wmma_scaled", + a_dtype="fp8", b_dtype="fp8", c_dtype="fp32", + scales=(MmaScaleDType.E8M0, MmaScaleDType.E8M0, MmaScaleBlockK.K32), + m=16, n=16, k=128, +) +assert atom is not None +# result = builder.mma(atom, a, b, c, scale_a, scale_b) +``` + +Omitting `scales` leaves scales unconstrained. Passing `(a_type, b_type, block_k)` +matches the scale contract exactly; `(None, None, None)` selects unscaled atoms. +Partially specified scale contracts are invalid. Enumeration +and existence queries may match several records, while `op_for_shape` and +`select_largest_k` reject ambiguous exact matches or largest-K ties. The C queries +use the same rules: a NULL `rocke_mma_scale_filter_t` pointer is unconstrained, +and `{NULL, NULL, ROCKE_MMA_SCALE_NONE}` requests an unscaled atom. The C block +enum has `ROCKE_MMA_SCALE_K16` and `ROCKE_MMA_SCALE_K32`, plus the unscaled +sentinel. Scale value formats are `e8m0`, `e4m3`, and `e5m3`, with +`fp8e4m3` accepted as an alias for `e4m3`; actual target support comes from the +catalog. + +### Native C/C++ Atom Lookup + +When holding an `ArchTarget`, use `rocke_archtarget_op_for_shape` from +`rocke/helper_rocke.core.arch.h`. It forwards to +`rocke_mma_catalog_op_for_shape(&target->mma, ...)`, which is also available +directly from `rocke/arch_target.h`. Both select the same catalog record. +Kernel builders use that record's operation ID, fragment sizes, and layouts; +support checks use the lookup to establish whether the requested atom exists. +The `(m, n, k)` arguments describe one instruction atom, not the full GEMM. + +This C++ example selects the gfx1250 FP8 atom with E8M0 scales shared over K32: + +```cpp +#include "rocke/error.hpp" +#include "rocke/helper_rocke.core.arch.h" + +int main() +{ + try + { + const auto* target = rocke_archtarget_from_gfx("gfx1250"); + const rocke_mma_scale_filter_t scales = {"e8m0", "e8m0", ROCKE_MMA_SCALE_K32}; + const auto* atom = rocke_archtarget_op_for_shape( + target, "wmma_scaled", "fp8", "fp8", "fp32", 16, 16, 128, &scales); + return atom ? 0 : 1; + } + catch(const ckc::Error&) + { + return 2; + } +} +``` + +The helper requires the trailing filter argument in both C and C++. +Pass `NULL` (`nullptr` in C++) for unconstrained scales, or a pointer to +`{NULL, NULL, ROCKE_MMA_SCALE_NONE}` to select only unscaled atoms. +The scale-contract rules above apply. No match, or a null target, returns +`NULL`; invalid filters and ambiguous matches raise `ckc::Error`. Handle those +errors at a C++ boundary before returning to C. In particular, the scaled +FP8/BF8 `16x16x128` shapes have both K16 and K32 records, so an unconstrained +lookup is ambiguous. + +## Operation IDs and Migration + +Scaled-WMMA IDs have the form +`wmma______scale___k`. +For example, `wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k32` +has atom K=128 and one scale per 32 K elements for both inputs. Both scale +types are written even when equal. Lowering reads the catalog fields, never +parses the ID, and selects LLVM intrinsic names and packed carriers separately. +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. 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 efc5cd42b926..bd0f96fc8493 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 @@ -10,6 +10,7 @@ Modules: + - ``dtypes`` : target-independent dtype names and alias normalization. - ``ir`` : `IRBuilder`, `KernelDef`, `Value`, `Op`, `Region`, plus type system (`F16`, `F32`, `I32`, `I64`, `VectorType`, `PtrType`, `SmemType`). @@ -44,6 +45,7 @@ apply_codegen_policy, codegen_policy_for_kernel, ) +from .dtypes import normalize_dtype from .ir import ( BF16, F16, @@ -93,6 +95,7 @@ from .verify import Diagnostic, verify, verify_or_raise __all__ = [ + "normalize_dtype", "CodegenPolicy", "SchedulerStrategy", "apply_codegen_policy", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/__init__.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/__init__.py index 96b44b9b5859..f97e32b8c7b5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/__init__.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/__init__.py @@ -7,18 +7,20 @@ querying a GPU or compiler. """ +from ..dtypes import normalize_dtype from .target import ( # noqa: F401 ArchTarget, LayoutMap, MemoryCapabilities, MmaCatalog, MmaOp, + MmaScaleBlockK, + MmaScaleDType, ResourceLimits, arch_from_isa, base_arch_from_target_id, compiler_target_from_target_id, known_arches, - normalize_dtype, target_id_from_isa, validate_arch, ) @@ -29,6 +31,8 @@ "MemoryCapabilities", "MmaCatalog", "MmaOp", + "MmaScaleBlockK", + "MmaScaleDType", "ResourceLimits", "arch_from_isa", "base_arch_from_target_id", 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 0607dc02932c..3d12a94fae34 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/arch_specs.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/arch_specs.json @@ -142,7 +142,7 @@ ] }, "gfx1250": { - "_comment": "gfx1250-class CDNA multi-chip device on the GFX12 programming model (wave32, WMMA, no MFMA). Primary fp16/bf16 WMMA atom is 16x16x32 (NOT gfx1201's 16x16x16). ABI empirically confirmed on ROCm 7.13 clang 23 (gfx1250 device target): __builtin_amdgcn_wmma_f32_16x16x32_f16 -> llvm.amdgcn.wmma.f32.16x16x32.f16.v8f32.v16f16(i1,<16 x half>,i1,<16 x half>,i16,<8 x float>,i1,i1); bf16 uses v16bf16 (<16 x bfloat>) directly, not the gfx11 <16 x i16> bitcast. Datalayout/triple byte-identical to gfx950/gfx1201. This catalog ships validated fp16/bf16 and fp8/bf8 WMMA atoms plus the FP8xFP8 SCALE/SCALE16 prerequisite slice; FP4/FP6/BF6 block-scaled forms remain intentionally omitted. Lane maps are a HYPOTHESIS verified by examples/gfx1250/wmma_probe.py. MEMORY: gfx1250 HAS async global->LDS DMA (its own GFX12 family: global_load_async_to_lds_b128 + s_wait_asynccnt -- NOT the gfx9 buffer_load_lds/global_load_lds, which do not select), so has_async_global_lds=true (GPU-validated via the DTLA path). It also HAS a hardware transpose-LDS read, ds_load_tr16_b128 (GPU-validated via compute_pv_dstr). The has_async_lds and has_ds_read_tr flags stay false on purpose: in this catalog they denote the gfx9/gfx950 ds_read_*_tr_* / async-LDS *ABI* used to admit the gfx950 MFMA tiled builder (attention_arch.py), which gfx1250 does NOT have -- its async/transpose use distinct GFX12 opcodes gated directly in the ISA backend, not via these MFMA-pipeline bits. has_tdm=false: TDM not verified on this part.", + "_comment": "gfx1250-class CDNA multi-chip device on the GFX12 programming model (wave32, WMMA, no MFMA). Primary fp16/bf16 WMMA atom is 16x16x32 (NOT gfx1201's 16x16x16). ABI empirically confirmed on ROCm 7.13 clang 23 (gfx1250 device target): __builtin_amdgcn_wmma_f32_16x16x32_f16 -> llvm.amdgcn.wmma.f32.16x16x32.f16.v8f32.v16f16(i1,<16 x half>,i1,<16 x half>,i16,<8 x float>,i1,i1); bf16 uses v16bf16 (<16 x bfloat>) directly, not the gfx11 <16 x i16> bitcast. Datalayout/triple byte-identical to gfx950/gfx1201. This catalog ships validated fp16/bf16 and fp8/bf8 WMMA atoms plus homogeneous FP8xFP8 and BF8xBF8 SCALE/SCALE16 with E8M0 scales; FP4/FP6/BF6 block-scaled forms remain intentionally omitted. Lane maps are a HYPOTHESIS verified by examples/gfx1250/wmma_probe.py. MEMORY: gfx1250 HAS async global->LDS DMA (its own GFX12 family: global_load_async_to_lds_b128 + s_wait_asynccnt -- NOT the gfx9 buffer_load_lds/global_load_lds, which do not select), so has_async_global_lds=true (GPU-validated via the DTLA path). It also HAS a hardware transpose-LDS read, ds_load_tr16_b128 (GPU-validated via compute_pv_dstr). The has_async_lds and has_ds_read_tr flags stay false on purpose: in this catalog they denote the gfx9/gfx950 ds_read_*_tr_* / async-LDS *ABI* used to admit the gfx950 MFMA tiled builder (attention_arch.py), which gfx1250 does NOT have -- its async/transpose use distinct GFX12 opcodes gated directly in the ISA backend, not via these MFMA-pipeline bits. has_tdm=false: TDM not verified on this part.", "family": "cdna", "target_family": "gfx12_cdna", "matrix_path": "wmma", @@ -177,8 +177,10 @@ {"family": "wmma", "a": "fp8", "b": "bf8", "c": "fp32", "m": 16, "n": 16, "k": 64, "op_id": "wmma_gfx1250_f32_16x16x64_fp8_bf8"}, {"family": "wmma", "a": "bf8", "b": "fp8", "c": "fp32", "m": 16, "n": 16, "k": 64, "op_id": "wmma_gfx1250_f32_16x16x64_bf8_fp8"}, {"family": "wmma", "a": "bf8", "b": "bf8", "c": "fp32", "m": 16, "n": 16, "k": 64, "op_id": "wmma_gfx1250_f32_16x16x64_bf8_bf8"}, - {"family": "wmma_scale", "a": "fp8", "b": "fp8", "c": "fp32", "m": 16, "n": 16, "k": 128, "op_id": "wmma_scale_f32_16x16x128_fp8_fp8"}, - {"family": "wmma_scale16", "a": "fp8", "b": "fp8", "c": "fp32", "m": 16, "n": 16, "k": 128, "op_id": "wmma_scale16_f32_16x16x128_fp8_fp8"} + {"family": "wmma_scaled", "a": "fp8", "b": "fp8", "c": "fp32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32, "m": 16, "n": 16, "k": 128, "op_id": "wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k32"}, + {"family": "wmma_scaled", "a": "bf8", "b": "bf8", "c": "fp32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 32, "m": 16, "n": 16, "k": 128, "op_id": "wmma_gfx1250_f32_16x16x128_bf8_bf8_scale_e8m0_e8m0_k32"}, + {"family": "wmma_scaled", "a": "fp8", "b": "fp8", "c": "fp32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 16, "m": 16, "n": 16, "k": 128, "op_id": "wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k16"}, + {"family": "wmma_scaled", "a": "bf8", "b": "bf8", "c": "fp32", "a_scale_dtype": "e8m0", "b_scale_dtype": "e8m0", "scale_block_k": 16, "m": 16, "n": 16, "k": 128, "op_id": "wmma_gfx1250_f32_16x16x128_bf8_bf8_scale_e8m0_e8m0_k16"} ] }, "gfx11-generic": { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/target.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/target.py index 1c95ac55af9a..fe94ac88f845 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 @@ -20,46 +20,15 @@ import json import re from dataclasses import dataclass, field +from enum import Enum, IntEnum from functools import lru_cache from pathlib import Path from typing import Any, Dict, List, Optional, Tuple -_DATA_FILE = Path(__file__).parent / "data" / "arch_specs.json" - -# Canonical dtype spellings used as catalog keys. Instance/spec dtype strings -# are normalised through this map so "f16"/"half" and "fp16" all resolve. -_DTYPE_ALIASES = { - "f16": "fp16", - "half": "fp16", - "fp16": "fp16", - "bf16": "bf16", - "bfloat16": "bf16", - "f32": "fp32", - "float": "fp32", - "fp32": "fp32", - "fp8": "fp8e4m3", - "fp8e4m3": "fp8e4m3", - "bf8": "bf8e5m2", - "bf8e5m2": "bf8e5m2", - # Integer WMMA: "iu8"/"iu4" are the RDNA WMMA integer operand families - # (signedness is an instruction operand, not the dtype); "i32" is the - # integer accumulator. Scalar int spellings pass through for completeness. - "iu8": "iu8", - "iu4": "iu4", - "i8": "i8", - "int8": "i8", - "i4": "i4", - "int4": "i4", - "i32": "i32", - "int32": "i32", -} - - -def normalize_dtype(name: str) -> str: - """Map a dtype spelling to its canonical catalog key.""" - key = name.strip().lower() - return _DTYPE_ALIASES.get(key, key) +# Re-export for callers using the original architecture entry point. +from ..dtypes import normalize_dtype +_DATA_FILE = Path(__file__).parent / "data" / "arch_specs.json" # A callable that, given an :class:`~rocke.core.ir.IRBuilder`, a runtime lane # ``Value`` (0..wave_size-1) and a compile-time fragment slot index, emits the @@ -94,7 +63,9 @@ class LayoutMap: role ``"acc"`` (accumulator C/D, coords are ``(row, col)``), ``"a"`` (A operand, coords are ``(row, k)``) or ``"b"`` (B operand, coords are - ``(k, col)``). + ``(k, col)``), ``"a_scale"`` (``(row, K-group)``), or + ``"b_scale"`` (``(K-group, col)``). Scale slots count logical elements, + independently of the backend register carrier. frag_len Number of fragment slots per lane for this role (the per-lane vector length: e.g. 4 for an MFMA 16x16x16 accumulator, 8 for the WMMA @@ -130,13 +101,56 @@ def coord(self, builder: Any, lane: Any, slot: int) -> Tuple[Any, Any]: return self.fn(builder, lane, slot) +class MmaScaleDType(str, Enum): + """Scale value formats, independent of matrix dtypes and target support. + + E5M3 is an unsigned scale format, distinct from the signed E5M2 matrix + format. Only E4M3 shares the accepted matrix spelling ``fp8e4m3``. + """ + + E8M0 = "e8m0" + E4M3 = "e4m3" + E5M3 = "e5m3" + + def __str__(self) -> str: + return self.value + + @classmethod + def _missing_(cls, value: object) -> MmaScaleDType | None: + if value == "fp8e4m3": + return cls.E4M3 + return None + + +class MmaScaleBlockK(IntEnum): + """Number of K elements sharing one scale, common to both matrix inputs.""" + + K16 = 16 + K32 = 32 + + +def _normalize_mma_scales( + a_dtype: str | None, b_dtype: str | None, block_k: int | None +) -> tuple[MmaScaleDType | None, MmaScaleDType | None, MmaScaleBlockK | None]: + """Validate the complete scale contract, independently of backend support.""" + if a_dtype is None and b_dtype is None and block_k is None: + return None, None, None + try: + a, b = MmaScaleDType(a_dtype), MmaScaleDType(b_dtype) + except ValueError: + raise ValueError("MMA scale dtype must be e8m0, e4m3, or e5m3") from None + if type(block_k) not in (int, MmaScaleBlockK) or block_k not in (16, 32): + raise ValueError("MMA scale_block_k must be an integer equal to 16 or 32") + return a, b, MmaScaleBlockK(block_k) + + @dataclass(frozen=True) class MmaOp: """A single supported matrix-multiply-accumulate atom on a target. - ``op_id`` is the opaque handle the backend consumes; today it matches the - ISA-named ``IRBuilder`` method (e.g. ``mfma_f32_16x16x16_f16``). It is **not** - LLVM intrinsic text — the backend maps ``op_id`` to an intrinsic. + ``op_id`` identifies a concrete operand contract. Scaled WMMA IDs extend + the target/accumulator/shape/input convention with A/B scale types and a + shared K-group size. The backend reads metadata, not the ID spelling. The ``*_frag_len`` fields and the layout-map accessors describe the *physical* register fragmentation of the atom: how many values of each @@ -148,7 +162,7 @@ class MmaOp: the accessor until a map is added. """ - family: str # "mma" | "wmma" + family: str # "mma" | "wmma" | "wmma_scaled" a_dtype: str b_dtype: str c_dtype: str @@ -167,6 +181,41 @@ class MmaOp: _b_layout: Optional[LayoutMap] = field(default=None, repr=False, compare=False) _c_layout: Optional[LayoutMap] = field(default=None, repr=False, compare=False) + # All absent means unscaled. The K-group size applies to both inputs; + # scale types describe values, not the backend's packed register carrier. + a_scale_dtype: MmaScaleDType | str | None = None + b_scale_dtype: MmaScaleDType | str | None = None + scale_block_k: MmaScaleBlockK | None = None + + a_scale_frag_len: int = 0 + b_scale_frag_len: int = 0 + _a_scale_layout: Optional[LayoutMap] = field( + default=None, repr=False, compare=False + ) + _b_scale_layout: Optional[LayoutMap] = field( + default=None, repr=False, compare=False + ) + + def __post_init__(self) -> None: + a, b, block_k = _normalize_mma_scales( + self.a_scale_dtype, self.b_scale_dtype, self.scale_block_k + ) + object.__setattr__(self, "a_scale_dtype", a) + object.__setattr__(self, "b_scale_dtype", b) + object.__setattr__(self, "scale_block_k", block_k) + for role in ("a_scale", "b_scale"): + count = getattr(self, f"{role}_frag_len") + layout = getattr(self, f"_{role}_layout") + if type(count) is not int or count < 0 or (block_k is None and count): + raise ValueError(f"invalid {role} fragment length: {count!r}") + if layout is not None and ( + not count + or layout.role != role + or layout.frag_len != count + or layout.wave_size != self.wave_size + ): + raise ValueError(f"{role} layout does not match its fragment metadata") + @property def shape(self) -> Tuple[int, int, int]: return (self.m, self.n, self.k) @@ -184,6 +233,14 @@ def c_layout(self) -> LayoutMap: """The accumulator (C/D) ``(row, col)`` lane/slot -> coordinate map.""" return self._require_layout(self._c_layout, "c") + def a_scale_layout(self) -> LayoutMap: + """A scale ``(row, K-group)`` coordinates for each logical scale slot.""" + return self._require_layout(self._a_scale_layout, "a_scale") + + def b_scale_layout(self) -> LayoutMap: + """B scale ``(K-group, col)`` coordinates for each logical scale slot.""" + return self._require_layout(self._b_scale_layout, "b_scale") + # Convenience aliases for the accumulator (the most-used map). def acc_layout(self) -> LayoutMap: return self.c_layout() @@ -569,10 +626,20 @@ def _wmma_gfx1250_b_16x16x32(builder, lane, slot): return k, col +def _wmma_gfx1250_a_scale(builder, lane, slot): + row = builder.mod(lane, builder.const_i32(16)) + return row, builder.const_i32(slot) + + +def _wmma_gfx1250_b_scale(builder, lane, slot): + col = builder.mod(lane, builder.const_i32(16)) + return builder.const_i32(slot), col + + @dataclass(frozen=True) class _FragInfo: """Per-op_id fragment metadata: per-lane vector lengths, wave size, and the - (optional) lane/slot coordinate functions for A / B / accumulator.""" + (optional) lane/slot coordinate functions for matrices and scales.""" a_frag_len: int b_frag_len: int @@ -581,6 +648,10 @@ class _FragInfo: a_fn: _LaneCoordFn = None b_fn: _LaneCoordFn = None c_fn: _LaneCoordFn = None + a_scale_frag_len: int = 0 + b_scale_frag_len: int = 0 + a_scale_fn: _LaneCoordFn = None + b_scale_fn: _LaneCoordFn = None # op_id -> physical fragment metadata. Frag lengths are populated for every @@ -750,10 +821,10 @@ class _FragInfo: None, _wmma_gfx12_acc_16x16, ), - # Native gfx1250 MX FP8 WMMA. A/B each carry 64 bytes as <16 x i32>; + # Native gfx1250 scaled WMMA. FP8/BF8 use 64 bytes per lane as <16 x i32>. # SCALE packs four K=32 E8M0 bytes in i32 and SCALE16 packs eight K=16 # bytes in i64. Both share the gfx12 column-distributed accumulator. - "wmma_scale_f32_16x16x128_fp8_fp8": _FragInfo( + "wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k32": _FragInfo( 16, 16, 8, @@ -761,8 +832,12 @@ class _FragInfo: None, None, _wmma_gfx12_acc_16x16, + a_scale_frag_len=4, + b_scale_frag_len=4, + a_scale_fn=_wmma_gfx1250_a_scale, + b_scale_fn=_wmma_gfx1250_b_scale, ), - "wmma_scale16_f32_16x16x128_fp8_fp8": _FragInfo( + "wmma_gfx1250_f32_16x16x128_bf8_bf8_scale_e8m0_e8m0_k32": _FragInfo( 16, 16, 8, @@ -770,6 +845,36 @@ class _FragInfo: None, None, _wmma_gfx12_acc_16x16, + a_scale_frag_len=4, + b_scale_frag_len=4, + a_scale_fn=_wmma_gfx1250_a_scale, + b_scale_fn=_wmma_gfx1250_b_scale, + ), + "wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k16": _FragInfo( + 16, + 16, + 8, + 32, + None, + None, + _wmma_gfx12_acc_16x16, + a_scale_frag_len=8, + b_scale_frag_len=8, + a_scale_fn=_wmma_gfx1250_a_scale, + b_scale_fn=_wmma_gfx1250_b_scale, + ), + "wmma_gfx1250_f32_16x16x128_bf8_bf8_scale_e8m0_e8m0_k16": _FragInfo( + 16, + 16, + 8, + 32, + None, + None, + _wmma_gfx12_acc_16x16, + a_scale_frag_len=8, + b_scale_frag_len=8, + a_scale_fn=_wmma_gfx1250_a_scale, + b_scale_fn=_wmma_gfx1250_b_scale, ), "wmma_gfx1250_f32_16x16x32_bf16": _FragInfo( 16, @@ -810,7 +915,13 @@ class ResourceLimits: class MmaCatalog: - """The arch-selected set of MMA atoms, with enumeration + best-K selection.""" + """The arch-selected MMA atoms, with optional exact A/B scale filtering. + + ``scales=(a_dtype, b_dtype, block_k)`` selects the complete scale contract. + Scale dtypes accept ``MmaScaleDType`` members or their string spellings. + ``None`` leaves scales unconstrained; ``(None, None, None)`` selects unscaled + atoms. Exact selection and largest-K ties must be unambiguous. + """ def __init__(self, ops: List[MmaOp]) -> None: self._ops = tuple(ops) @@ -826,9 +937,14 @@ def enumerate( a_dtype: str, b_dtype: str, c_dtype: str, + scales: tuple[str | None, str | None, int | None] | None = None, m: Optional[int] = None, n: Optional[int] = None, ) -> List[MmaOp]: + if scales is not None: + if len(scales) != 3: + raise ValueError("scales must contain exactly 3 entries") + scales = _normalize_mma_scales(*scales) a, b, c = ( normalize_dtype(a_dtype), normalize_dtype(b_dtype), @@ -840,6 +956,11 @@ def enumerate( continue if (op.a_dtype, op.b_dtype, op.c_dtype) != (a, b, c): continue + if ( + scales is not None + and (op.a_scale_dtype, op.b_scale_dtype, op.scale_block_k) != scales + ): + continue if m is not None and op.m != m: continue if n is not None and op.n != n: @@ -854,6 +975,7 @@ def has_shape( a_dtype: str, b_dtype: str, c_dtype: str, + scales: tuple[str | None, str | None, int | None] | None = None, m: int, n: int, k: int, @@ -865,6 +987,7 @@ def has_shape( a_dtype=a_dtype, b_dtype=b_dtype, c_dtype=c_dtype, + scales=scales, m=m, n=n, ) @@ -877,6 +1000,7 @@ def select_largest_k( a_dtype: str, b_dtype: str, c_dtype: str, + scales: tuple[str | None, str | None, int | None] | None = None, m: int, n: int, k_max: Optional[int] = None, @@ -888,6 +1012,7 @@ def select_largest_k( a_dtype=a_dtype, b_dtype=b_dtype, c_dtype=c_dtype, + scales=scales, m=m, n=n, ) @@ -895,7 +1020,14 @@ def select_largest_k( ] if not cands: return None - return max(cands, key=lambda op: op.k) + largest_k = max(op.k for op in cands) + return self._unique([op for op in cands if op.k == largest_k]) + + @staticmethod + def _unique(ops: list[MmaOp]) -> MmaOp | None: + if len(ops) > 1: + raise ValueError("ambiguous MMA query; specify the full operand contract") + return ops[0] if ops else None def by_op_id(self, op_id: str) -> Optional[MmaOp]: """Look up an atom by its ``op_id`` handle (the backend's MMA key).""" @@ -911,21 +1043,21 @@ def op_for_shape( a_dtype: str, b_dtype: str, c_dtype: str, + scales: tuple[str | None, str | None, int | None] | None = None, m: int, n: int, k: int, ) -> Optional[MmaOp]: - for op in self.enumerate( + candidates = self.enumerate( family=family, a_dtype=a_dtype, b_dtype=b_dtype, c_dtype=c_dtype, + scales=scales, m=m, n=n, - ): - if op.k == k: - return op - return None + ) + return self._unique([op for op in candidates if op.k == k]) @dataclass(frozen=True) @@ -1033,6 +1165,22 @@ def _op_id_c_dtype() -> Dict[str, str]: return out +@lru_cache(maxsize=1) +def _op_id_family() -> Dict[str, str]: + """Resolve operation families without selecting a target or parsing IDs.""" + out: Dict[str, str] = {} + for row in _load_specs().values(): + for op in row["mma"]: + op_id, family = op["op_id"], op["family"] + previous = out.setdefault(op_id, family) + if previous != family: + raise ValueError( + f"arch SSOT drift: op_id {op_id!r} has inconsistent family " + f"across arches ({previous!r} vs {family!r})" + ) + return out + + def _build_mma_op(o: dict) -> MmaOp: """Construct an :class:`MmaOp` from one catalog JSON row, attaching the physical fragment lengths and layout maps registered for its op_id.""" @@ -1060,6 +1208,13 @@ def _mk(role: str, frag_len: int, fn: _LaneCoordFn) -> Optional[LayoutMap]: _a_layout=_mk("a", info.a_frag_len, info.a_fn), _b_layout=_mk("b", info.b_frag_len, info.b_fn), _c_layout=_mk("c", info.c_frag_len, info.c_fn), + a_scale_dtype=o.get("a_scale_dtype"), + b_scale_dtype=o.get("b_scale_dtype"), + scale_block_k=o.get("scale_block_k"), + a_scale_frag_len=info.a_scale_frag_len, + b_scale_frag_len=info.b_scale_frag_len, + _a_scale_layout=_mk("a_scale", info.a_scale_frag_len, info.a_scale_fn), + _b_scale_layout=_mk("b_scale", info.b_scale_frag_len, info.b_scale_fn), ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/wmma_scale.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/wmma_scale.py new file mode 100644 index 000000000000..2febb7514b05 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/wmma_scale.py @@ -0,0 +1,103 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""gfx1250 scaled-WMMA operand contracts shared by packing and lowering.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from .target import ArchTarget, MmaOp, MmaScaleDType + + +@dataclass(frozen=True) +class ScalePacking: + """Pack consecutive K groups into a word, first group in its low byte. + + Each encoded scale occupies one byte. A and B use this ordering independently; + the integer word is only a carrier for the encoded floating-point scales. + """ + + count: int + block_k: int + + @property + def element_bits(self) -> int: + return 8 + + @property + def word_bits(self) -> int: + return self.count * self.element_bits + + @property + def llvm_type(self) -> str: + return f"i{self.word_bits}" + + +@dataclass(frozen=True) +class ScaledWmmaOp: + """Backend packing derived from a supported catalog operand contract.""" + + atom: MmaOp + matrix_formats: tuple[int, int] + scale_formats: tuple[int, int] + scales: ScalePacking + + @property + def op_id(self) -> str: + return self.atom.op_id + + @property + def scale16(self) -> bool: + return self.scales.block_k == 16 + + @property + def matrix_llvm_types(self) -> tuple[str, str]: + return tuple( + f"<{n} x i32>" for n in (self.atom.a_frag_len, self.atom.b_frag_len) + ) + + @property + def intrinsic_suffix(self) -> str: + atom = self.atom + return ( + f"f32.{atom.m}x{atom.n}x{atom.k}.f8f6f4.v{atom.c_frag_len}f32." + f"v{atom.a_frag_len}i32.v{atom.b_frag_len}i32" + ) + + @property + def intrinsic(self) -> str: + mode = "scale16" if self.scale16 else "scale" + return f"llvm.amdgcn.wmma.{mode}.{self.intrinsic_suffix}" + + @property + def declaration_key(self) -> str: + return f"wmma.scale.block{self.scales.block_k}.gfx1250.{self.intrinsic_suffix}" + + +def gfx1250_scaled_wmma(op_id: str) -> ScaledWmmaOp | None: + """Resolve a catalog contract; LLVM selectors are backend details.""" + atom = ArchTarget.from_gfx("gfx1250").mma.by_op_id(op_id.removeprefix("tile.")) + if atom is None or atom.family != "wmma_scaled": + return None + formats = {"fp8e4m3": 0, "bf8e5m2": 1} + # The current backend supports E8M0 for both inputs and a shared K-group size. + # Keep these restrictions here, independently of the catalog query model. + if ( + atom.a_dtype not in formats + or atom.b_dtype not in formats + or atom.a_scale_dtype != MmaScaleDType.E8M0 + or atom.b_scale_dtype != MmaScaleDType.E8M0 + or atom.scale_block_k not in (16, 32) + or atom.c_dtype != "fp32" + or atom.shape != (16, 16, 128) + ): + raise ValueError(f"unsupported scaled WMMA backend contract: {atom.op_id}") + counts = (atom.a_scale_frag_len, atom.b_scale_frag_len) + if any(count != atom.k // atom.scale_block_k for count in counts): + raise ValueError(f"unsupported scaled WMMA scale fragment lengths: {counts}") + return ScaledWmmaOp( + atom=atom, + matrix_formats=(formats[atom.a_dtype], formats[atom.b_dtype]), + scale_formats=(0, 0), # E8M0 for each source. + scales=ScalePacking(count=atom.a_scale_frag_len, block_k=atom.scale_block_k), + ) 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 new file mode 100644 index 000000000000..583b591b1390 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/dtypes.py @@ -0,0 +1,50 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Target-independent dtype names shared by specs and architecture catalogs. + +These names describe logical formats. Recognition does not imply an IR scalar +representation, conversion support, or a matrix instruction on a given target. +Target-specific support belongs to the architecture catalog. +""" + +from __future__ import annotations + +# Canonical dtype spellings used as catalog keys. Instance/spec dtype strings +# are normalised through this map so "f16"/"half" and "fp16" all resolve. +_DTYPE_ALIASES = { + "f16": "fp16", + "half": "fp16", + "fp16": "fp16", + "bf16": "bf16", + "bfloat16": "bf16", + "f32": "fp32", + "float": "fp32", + "fp32": "fp32", + "fp8": "fp8e4m3", + "fp8e4m3": "fp8e4m3", + "bf8": "bf8e5m2", + "bf8e5m2": "bf8e5m2", + "fp6": "fp6e2m3", + "fp6e2m3": "fp6e2m3", + "bf6": "fp6e3m2", + "fp6e3m2": "fp6e3m2", + "fp4": "fp4e2m1", + "fp4e2m1": "fp4e2m1", + # Integer WMMA: "iu8"/"iu4" are the RDNA WMMA integer operand families + # (signedness is an instruction operand, not the dtype); "i32" is the + # integer accumulator. Scalar int spellings pass through for completeness. + "iu8": "iu8", + "iu4": "iu4", + "i8": "i8", + "int8": "i8", + "i4": "i4", + "int4": "i4", + "i32": "i32", + "int32": "i32", +} + + +def normalize_dtype(name: str) -> str: + """Map a dtype spelling to its canonical catalog key.""" + key = name.strip().lower() + return _DTYPE_ALIASES.get(key, key) 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 8960d3196946..cab47560124e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py @@ -29,6 +29,8 @@ from dataclasses import dataclass, field from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple +from .arch import target as _arch + # ----------------------------- Types -------------------------------------- @@ -71,9 +73,8 @@ def __repr__(self) -> str: # Both are read from the arch SSOT (``core/arch/target``): fragment lengths from # ``_MMA_FRAGMENT_INFO`` and the accumulator dtype from the JSON catalog. ir.py # keeps *no* private copy of that data — an ``MmaOp`` object supplies both fields -# directly, and a bare op_id string is resolved through the lazy helpers below. -# The arch package is imported lazily (inside the helpers) so ir.py stays -# importable without eagerly loading the arch tree. +# directly, and a bare op_id string is resolved through the helpers below. +# The arch module loads its JSON catalog lazily when a lookup needs it. # op_id -> the ``result_name_hint`` the legacy ISA-named method used. This is # purely ir-side SSA naming (not arch data), kept here so the emitted value @@ -85,8 +86,6 @@ def __repr__(self) -> str: "mfma_f32_16x16x96_fp6": "acc6", "mfma_f32_16x16x128_fp8": "acc128", "mfma_scale_f32_16x16x128_f8f6f4": "mxacc", - "wmma_scale_f32_16x16x128_fp8_fp8": "mxacc", - "wmma_scale16_f32_16x16x128_fp8_fp8": "mxacc", } @@ -135,12 +134,10 @@ def _check_cachepolicy(op: str, value: int) -> int: def _mma_c_frag_len(op_id: str) -> int: """Accumulator fragment length for ``op_id`` from the arch SSOT. - Resolved through ``core/arch/target._MMA_FRAGMENT_INFO`` (imported lazily); + Resolved through ``core/arch/target._MMA_FRAGMENT_INFO``; ir.py holds no private copy. Unknown op_ids (frag length 0) raise, matching the strictness callers relied on. """ - from rocke.core.arch import target as _arch - frag_len = _arch._frag_info(op_id).c_frag_len if frag_len <= 0: raise ValueError( @@ -154,11 +151,9 @@ def _mma_c_is_int(op_id: str) -> bool: """True when ``op_id`` accumulates in i32 (integer WMMA). Sourced from the arch catalog's accumulator dtype - (``core/arch/data/arch_specs.json`` via ``target._op_id_c_dtype``), imported - lazily. Op_ids absent from the catalog default to the f32 accumulator. + (``core/arch/data/arch_specs.json`` via ``target._op_id_c_dtype``). + Op_ids absent from the catalog default to the f32 accumulator. """ - from rocke.core.arch import target as _arch - return _arch._op_id_c_dtype().get(op_id) == "i32" @@ -1859,7 +1854,11 @@ def mma( c_dtype = getattr(op, "c_dtype", None) is_int_acc = c_dtype == "i32" if c_dtype is not None else _mma_c_is_int(op_id) c_elem = I32 if is_int_acc else F32 - hint = _MMA_RESULT_HINT.get(op_id, "acc") + hint = ( + "mxacc" + if _arch._op_id_family().get(op_id) == "wmma_scaled" + else _MMA_RESULT_HINT.get(op_id, "acc") + ) return self._op( "tile.mma", [a, b, c, *extra], @@ -1963,38 +1962,6 @@ def wmma_gfx1250_f32_16x16x64_bf8_bf8(self, a: Value, b: Value, c: Value) -> Val """gfx1250 (gfx1250) BF8 K=64 WMMA. Thin wrapper over :meth:`mma`.""" return self.mma("wmma_gfx1250_f32_16x16x64_bf8_bf8", a, b, c) - def wmma_scale_f32_16x16x128_fp8_fp8( - self, - a: Value, - b: Value, - c: Value, - a_scale: Value, - b_scale: Value, - ) -> Value: - """gfx1250 native SCALE FP8 WMMA with packed E8M0 scale operands. - - A and B are ``<16 x i32>`` fragments (64 FP8 bytes per lane), C is - ``<8 x f32>``, and each scale operand is one i32 packing four E8M0 - bytes for the instruction's four K=32 scale blocks. - """ - return self.mma("wmma_scale_f32_16x16x128_fp8_fp8", a, b, c, a_scale, b_scale) - - def wmma_scale16_f32_16x16x128_fp8_fp8( - self, - a: Value, - b: Value, - c: Value, - a_scale: Value, - b_scale: Value, - ) -> Value: - """gfx1250 native SCALE16 FP8 WMMA with eight packed E8M0 scales. - - The matrix and accumulator fragments match - :meth:`wmma_scale_f32_16x16x128_fp8_fp8`; each scale operand is i64 - because SCALE16 carries eight K=16 E8M0 scale bytes. - """ - return self.mma("wmma_scale16_f32_16x16x128_fp8_fp8", a, b, c, a_scale, b_scale) - def mfma_f32_16x16x16_f16(self, a: Value, b: Value, c: Value) -> Value: return self.mma("mfma_f32_16x16x16_f16", a, b, c) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/isa/backend.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/isa/backend.py index 8970d949570e..e4a657583bc0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/isa/backend.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/isa/backend.py @@ -31,6 +31,7 @@ from typing import Callable, Dict, Tuple, Union from ..arch import ArchTarget +from ..arch.wmma_scale import gfx1250_scaled_wmma class ISABackend: @@ -332,23 +333,6 @@ class Gfx9MfmaBackend(ISABackend): } -# gfx1250 native MX FP8 WMMA, available with the LLVM 23 toolchain used by -# ROCm 7.13+. Both operations consume <16 x i32> matrix fragments; SCALE packs -# four E8M0 scale bytes in i32 while SCALE16 packs eight in i64. -_GFX1250_WMMA_FP8_SCALE = { - "tile.wmma_scale_f32_16x16x128_fp8_fp8": ( - "wmma.scale.gfx1250.f32.16x16x128.fp8.fp8", - "llvm.amdgcn.wmma.scale.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32", - "i32", - ), - "tile.wmma_scale16_f32_16x16x128_fp8_fp8": ( - "wmma.scale16.gfx1250.f32.16x16x128.fp8.fp8", - "llvm.amdgcn.wmma.scale16.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32", - "i64", - ), -} - - class Gfx11RdnaBackend(ISABackend): """RDNA3 / RDNA3.5 (gfx11, e.g. gfx1151 Strix Halo). **wave32**, **WMMA** (no MFMA), and a distinct ``s_waitcnt`` layout from gfx9/10. Datalayout + @@ -575,9 +559,9 @@ def emit_lds_barrier_drain(self, lowerer, *, drain_vmem: bool) -> None: lowerer._current().emit(" call void @llvm.amdgcn.s.wait.dscnt(i16 0)") def emit_wmma(self, lowerer, op) -> None: - scale_spec = _GFX1250_WMMA_FP8_SCALE.get(op.name) + scale_spec = gfx1250_scaled_wmma(op.name) if scale_spec is not None: - self._emit_wmma_scale(lowerer, op, scale_spec) + self._emit_wmma_scale(lowerer, op) return fp8_spec = _GFX1250_WMMA_FP8.get(op.name) if fp8_spec is not None: @@ -585,9 +569,14 @@ def emit_wmma(self, lowerer, op) -> None: return spec = _GFX1250_WMMA.get(op.name) if spec is None: + scaled_ops = [ + f"tile.{atom.op_id}" + for atom in self.arch.mma.ops + if gfx1250_scaled_wmma(atom.op_id) is not None + ] raise NotImplementedError( f"WMMA op {op.name!r} not yet wired for {self.arch.gfx}; " - f"known: {sorted(_GFX1250_WMMA) + sorted(_GFX1250_WMMA_FP8) + sorted(_GFX1250_WMMA_FP8_SCALE)}" + f"known: {sorted(_GFX1250_WMMA) + sorted(_GFX1250_WMMA_FP8) + sorted(scaled_ops)}" ) decl_key, intrinsic, elt = spec a, b, c = op.operands @@ -622,15 +611,22 @@ def _emit_wmma_fp8(self, lowerer, op, spec) -> None: f"i1 false, i1 false)" ) - def _emit_wmma_scale(self, lowerer, op, spec) -> None: - """Emit the ROCm 7.13+ gfx1250 SCALE/SCALE16 FP8 call.""" + def _emit_wmma_scale(self, lowerer, op) -> None: + """Emit the gfx1250 SCALE/SCALE16 intrinsic using the resolved operand contract.""" if lowerer._flavor != "llvm23": raise NotImplementedError( f"{op.name} requires llvm23 (ROCm 7.13+), got {lowerer._flavor}" ) if len(op.operands) != 5: raise ValueError(f"{op.name} expects 5 operands, got {len(op.operands)}") - decl_key, intrinsic, scale_ty = spec + spec = gfx1250_scaled_wmma(op.name) + if spec is None: + raise NotImplementedError(f"unsupported scaled WMMA op {op.name!r}") + # Declarations describe physical signatures, independently of matrix encodings. + decl_key = spec.declaration_key + intrinsic = spec.intrinsic + scale_ty = spec.scales.llvm_type + fmt0, fmt1 = spec.matrix_formats a, b, c, a_scale, b_scale = op.operands if a_scale.type.name != scale_ty or b_scale.type.name != scale_ty: raise ValueError( @@ -640,11 +636,11 @@ def _emit_wmma_scale(self, lowerer, op, spec) -> None: lowerer._need(decl_key) lowerer._current().emit( f" {op.result.name} = call <8 x float> @{intrinsic}(" - f"i32 0, <16 x i32> {lowerer._operand(a)}, " - f"i32 0, <16 x i32> {lowerer._operand(b)}, " + f"i32 {fmt0}, {spec.matrix_llvm_types[0]} {lowerer._operand(a)}, " + f"i32 {fmt1}, {spec.matrix_llvm_types[1]} {lowerer._operand(b)}, " f"i16 0, <8 x float> {lowerer._operand(c)}, " - f"i32 0, i32 0, {scale_ty} {lowerer._operand(a_scale)}, " - f"i32 0, i32 0, {scale_ty} {lowerer._operand(b_scale)}, " + f"i32 0, i32 {spec.scale_formats[0]}, {scale_ty} {lowerer._operand(a_scale)}, " + f"i32 0, i32 {spec.scale_formats[1]}, {scale_ty} {lowerer._operand(b_scale)}, " f"i1 false, i1 false)" ) 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 50fa21575fcc..f0628a9fdee8 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 @@ -17,6 +17,7 @@ from typing import List, Optional +from .arch.wmma_scale import gfx1250_scaled_wmma from .ir import ( KernelDef, Op, @@ -309,6 +310,9 @@ def _pop_indent(self) -> None: self._indent -= 1 def lower_op(self, op: Op) -> None: + if op.name.startswith("tile.") and gfx1250_scaled_wmma(op.name) is not None: + self._emit_wmma_gfx1250_scaled(op) + return method = getattr(self, f"_op_{op.name.replace('.', '_')}", None) if method is None: raise NotImplementedError(f"no HIP lowering for op {op.name!r}") @@ -526,7 +530,7 @@ def _op_tile_mma(self, op: Op) -> None: The HIP source path has no ISA backend (it emits ``__builtin_amdgcn_*`` directly), so the ``op_id`` is mapped back to the concrete - ``tile.`` op and dispatched through the existing per-op handler. + ``tile.`` op and passed through the shared operation dispatcher. This keeps the HIP emission identical to the legacy ISA-named path while the IRBuilder helpers route through :meth:`IRBuilder.mma`. """ @@ -652,12 +656,6 @@ def _op_tile_wmma_gfx1250_f32_16x16x64_bf8_fp8(self, op: Op) -> None: def _op_tile_wmma_gfx1250_f32_16x16x64_bf8_bf8(self, op: Op) -> None: self._emit_wmma_gfx1250_fp8(op, "bf8_bf8") - def _op_tile_wmma_scale_f32_16x16x128_fp8_fp8(self, op: Op) -> None: - self._emit_wmma_gfx1250_scaled(op, scale16=False) - - def _op_tile_wmma_scale16_f32_16x16x128_fp8_fp8(self, op: Op) -> None: - self._emit_wmma_gfx1250_scaled(op, scale16=True) - def _emit_wmma_gfx1250_fp8(self, op: Op, ab: str) -> None: # gfx1250 K=64 FP8/BF8 builtin: A/B are <8 x i32> (32 low-bit # bytes per lane), 6-operand form (A, B, fmt, C, reuseA, reuseB). @@ -669,23 +667,24 @@ def _emit_wmma_gfx1250_fp8(self, op: Op, ab: str) -> None: f"{_name(a)}, {_name(b)}, (int16_t)0, {_name(c)}, false, false);" ) - def _emit_wmma_gfx1250_scaled(self, op: Op, *, scale16: bool) -> None: - op_id = ( - "wmma_scale16_f32_16x16x128_fp8_fp8" - if scale16 - else "wmma_scale_f32_16x16x128_fp8_fp8" - ) + def _emit_wmma_gfx1250_scaled(self, op: Op) -> None: + spec = gfx1250_scaled_wmma(op.name) + if spec is None: + raise NotImplementedError(f"unsupported scaled WMMA op {op.name!r}") + op_id = spec.op_id + fmt0, fmt1 = spec.matrix_formats self._require_wmma_arch(op_id) a, b, c, a_scale, b_scale = op.operands builtin = ( "__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4" - if scale16 + if spec.scale16 else "__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4" ) self._emit( f"f32x8 {_name(op.result)} = {builtin}(" - f"0, {_name(a)}, 0, {_name(b)}, (int16_t)0, {_name(c)}, " - f"0, 0, {_name(a_scale)}, 0, 0, {_name(b_scale)}, false, false);" + f"{fmt0}, {_name(a)}, {fmt1}, {_name(b)}, (int16_t)0, {_name(c)}, " + f"0, {spec.scale_formats[0]}, {_name(a_scale)}, " + f"0, {spec.scale_formats[1]}, {_name(b_scale)}, false, false);" ) def _op_tile_wmma_gfx1250_f32_16x16x32_bf16(self, op: Op) -> None: 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 88e82f17efbd..4743a1335ea8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py @@ -608,13 +608,13 @@ def _resolve_llvm_flavor() -> str: "declare <8 x float> @llvm.amdgcn.wmma.f32.16x16x64.bf8.bf8.v8f32.v8i32(" "<8 x i32>, <8 x i32>, i16 immarg, <8 x float>, i1 immarg, i1 immarg)" ), - "wmma.scale.gfx1250.f32.16x16x128.fp8.fp8": ( + "wmma.scale.block32.gfx1250.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32": ( "declare <8 x float> @llvm.amdgcn.wmma.scale.f32.16x16x128.f8f6f4." "v8f32.v16i32.v16i32(i32 immarg, <16 x i32>, i32 immarg, " "<16 x i32>, i16 immarg, <8 x float>, i32 immarg, i32 immarg, i32, " "i32 immarg, i32 immarg, i32, i1 immarg, i1 immarg)" ), - "wmma.scale16.gfx1250.f32.16x16x128.fp8.fp8": ( + "wmma.scale.block16.gfx1250.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32": ( "declare <8 x float> @llvm.amdgcn.wmma.scale16.f32.16x16x128.f8f6f4." "v8f32.v16i32.v16i32(i32 immarg, <16 x i32>, i32 immarg, " "<16 x i32>, i16 immarg, <8 x float>, i32 immarg, i32 immarg, i64, " diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx1250/gemm/README.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx1250/gemm/README.md new file mode 100644 index 000000000000..1be444bf4917 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx1250/gemm/README.md @@ -0,0 +1,42 @@ +# gfx1250 block-scaled GEMM examples + +`mxfp8_gemm` runs homogeneous FP8 E4M3 and BF8 E5M2 separately by default. +Use `--dtype fp8` or `--dtype bf8` to run one encoding; `--dtype both` is +the explicit default. Canonical names are `fp8e4m3` and `bf8e5m2`. + +```bash +ROCKE_LLVM_FLAVOR=llvm23 python -m rocke.examples.gfx1250.gemm.mxfp8_gemm +ROCKE_LLVM_FLAVOR=llvm23 python -m rocke.examples.gfx1250.gemm.mxfp8_gemm --dtype bf8 --matrix-path wmma_scale16 +``` + +SCALE uses E8M0 blocks of 32 K elements; SCALE16 uses blocks of 16. +Use `--compile-route hip` for HIP compilation and `--case all` for neutral, +one-sided, combined, and isolated scale-group fixtures. The `mixed` fixture +means scale variation, not different A/B matrix formats. + +The examples share argument parsing, the spec-driven builder, compiler, +launcher, and numerical verifier. Numerical tests use opt-in fixtures in +`tests/instances/conftest.py` and are independently runnable by family. + +The target-independent aliases in `core/dtypes.py` normalize logical format +names. Recognition does not imply that a target supports an atom; use its +catalog to query supported operand formats and shapes. Each example defines +its accepted matrix and scale formats. + +`e8m0` identifies a scale format carried in bytes and packed integer operands. +It is not a general scalar IR type or conversion API. Native scaled WMMA uses +one `wmma_scaled` catalog family. Its operation IDs encode the matrix source +and accumulator dtype, each source's scale format, and a shared K-group size; the backend +selects the LLVM intrinsic and packed operand types from that contract. + +The native loader uses the atom's `a_scale_layout()` and `b_scale_layout()` for +lane ownership. Global scale tensors have shapes `[M, K/block_k]` and +`[K/block_k, N]`, respectively. K32 packs four E8M0 bytes per lane into i32; +K16 packs eight into i64, first K group in the low byte. Both half-waves carry +the same scales. Matrix A/B lane maps for these scaled atoms are not yet exposed. + +CPU tests independently check every scale coordinate, multi-tile addresses and +packed byte order. The opt-in numerical suite exercises both encodings and +scale group sizes, one-sided/group-isolated inputs, multiple tiles and K steps, +and HIP/COMGR compilation. These cases bound the validation; dtype recognition +alone does not establish support for other operand contracts. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx1250/gemm/_scaled_gemm_example.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx1250/gemm/_scaled_gemm_example.py new file mode 100644 index 000000000000..c19179a4b992 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx1250/gemm/_scaled_gemm_example.py @@ -0,0 +1,46 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Shared arguments and execution for focused scaled GEMM examples.""" + +from __future__ import annotations + +import argparse + +from .block_scaled_gemm_verify import run_cases + + +def argument_parser(description: str) -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=description) + parser.add_argument("--m", type=int, default=32) + parser.add_argument("--n", type=int, default=48) + parser.add_argument("--k", type=int, choices=(128, 256), default=256) + parser.add_argument( + "--matrix-path", + choices=("wmma_scale", "wmma_scale16"), + default="wmma_scale", + help="SCALE uses K=32 scale blocks; SCALE16 uses K=16 blocks", + ) + parser.add_argument("--compile-route", choices=("comgr", "hip"), default="comgr") + parser.add_argument( + "--case", + default="mixed", + help="neutral, a-only, b-only, mixed, group-N, or all", + ) + return parser + + +def verify(spec, args: argparse.Namespace) -> int: + """Prepare bounded inputs, compile, launch, and compare with the reference. + + The shared verifier owns packing, independent decoding, HIP allocations, + and result checks. Family examples select the public input contract. + """ + cases = ( + ("neutral", "a-only", "b-only", "mixed") + + tuple(f"group-{g}" for g in range(spec.K // spec.block_k)) + if args.case == "all" + else (args.case,) + ) + count = run_cases(spec, cases, compile_route=args.compile_route) + print(f"PASS: verified {count} cases", flush=True) + return 0 diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx1250/gemm/block_scaled_gemm_verify.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx1250/gemm/block_scaled_gemm_verify.py index f703106d894e..86252f0c3c9f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx1250/gemm/block_scaled_gemm_verify.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx1250/gemm/block_scaled_gemm_verify.py @@ -3,7 +3,7 @@ """Launch gfx1250 block-scaled GEMM and compare with an independent reference. The default invocation keeps the K=64 FP8/BF8 WMMA + FP32-scale verifier. -Native ``--matrix-path wmma_scale`` / ``wmma_scale16`` use FP8 and E8M0 +Native ``--matrix-path wmma_scale`` / ``wmma_scale16`` use FP8/BF8 and E8M0 scales with K=32 / K=16 groups. Native fixtures cover K=128 or 256 and use bounded dyadic values, permitting exact comparison after BF16 rounding. @@ -29,7 +29,9 @@ from ....helpers import compile_kernel from ....helpers.compile import compile_kernel_via_hipcc +from ....core.dtypes import normalize_dtype from ....instances.gfx1250.block_scaled_gemm import ( + _LOWBIT_DTYPES, BlockScaledGemmSpec, block_scaled_gemm_grid, build_block_scaled_gemm, @@ -211,7 +213,7 @@ def run_cases( inputs = make_case_inputs(spec, case) expected = reference_result(*inputs, spec.block_k, native=native) label = ( - f"{spec.resolved_matrix_path()}/{compile_route}/{case} " + f"{spec.resolved_matrix_path()}/{spec.dtype_a}/{compile_route}/{case} " f"{spec.M}x{spec.N}x{spec.K} bk{spec.block_k}" ) got = _launch(rt, fn, spec, inputs) @@ -233,7 +235,12 @@ def main(argv: list[str] | None = None) -> int: p.add_argument("--n", type=int, default=16) p.add_argument("--k", type=int, default=128) p.add_argument("--block-k", type=int, default=None) - p.add_argument("--dtype", default="fp8e4m3", choices=("fp8e4m3", "bf8e5m2")) + p.add_argument( + "--dtype", + default="fp8e4m3", + type=normalize_dtype, + choices=sorted(_LOWBIT_DTYPES), + ) p.add_argument("--tol", type=float, default=2e-2, help="legacy WMMA tolerance only") p.add_argument( "--matrix-path", default="wmma", choices=("wmma", "wmma_scale", "wmma_scale16") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx1250/gemm/mxfp8_gemm.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx1250/gemm/mxfp8_gemm.py new file mode 100644 index 000000000000..b48ccf7e43c9 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx1250/gemm/mxfp8_gemm.py @@ -0,0 +1,44 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Homogeneous FP8 E4M3 or BF8 E5M2 GEMM with E8M0 block scales on gfx1250.""" + +from __future__ import annotations + +from ....core.dtypes import normalize_dtype +from ....instances.gfx1250.block_scaled_gemm import BlockScaledGemmSpec +from ._scaled_gemm_example import argument_parser, verify + + +def make_spec(args) -> BlockScaledGemmSpec: + return BlockScaledGemmSpec( + name="mxfp8_gemm", + M=args.m, + N=args.n, + K=args.k, + dtype_a=args.dtype, + dtype_b=args.dtype, + dtype_c="bf16", + scale_dtype="e8m0", + matrix_path=args.matrix_path, + block_k=16 if args.matrix_path == "wmma_scale16" else 32, + ) + + +def main(argv: list[str] | None = None) -> int: + parser = argument_parser(__doc__) + parser.add_argument( + "--dtype", + type=normalize_dtype, + choices=("fp8e4m3", "bf8e5m2", "both"), + default="both", + ) + args = parser.parse_args(argv) + dtypes = ("fp8e4m3", "bf8e5m2") if args.dtype == "both" else (args.dtype,) + for dtype in dtypes: + args.dtype = dtype + verify(make_spec(args), args) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/gfx1250/block_scaled_gemm.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/gfx1250/block_scaled_gemm.py index 4a8b128d4169..9ee88cbdeb4a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/gfx1250/block_scaled_gemm.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/gfx1250/block_scaled_gemm.py @@ -13,6 +13,9 @@ from dataclasses import dataclass from typing import List, Tuple +from ...core.dtypes import normalize_dtype +from ...core.arch import ArchTarget +from ...core.arch.wmma_scale import gfx1250_scaled_wmma from ...core.ir import ( BF16, F16, @@ -29,7 +32,8 @@ from ...helpers.quant import quant_ir_type from ...helpers.spec import SignatureBuilder, ceil_div_grid, kernel_name_join -_LOWBIT_DTYPES = {"fp8", "fp8e4m3", "bf8", "bf8e5m2"} +_LOWBIT_FORMATS = {"fp8e4m3": "fp8", "bf8e5m2": "bf8"} +_LOWBIT_DTYPES = frozenset(_LOWBIT_FORMATS) _OUTPUT_DTYPES = {"fp16", "f16", "bf16"} _SCALE_DTYPES = {"fp16", "f16", "fp32", "f32"} _SUPPORTED_MATRIX_PATHS = { @@ -56,11 +60,11 @@ def _wmma_op_id(dtype_a: str, dtype_b: str) -> str: def _canon_lowbit(dtype: str) -> str: - if dtype in ("fp8", "fp8e4m3"): - return "fp8" - if dtype in ("bf8", "bf8e5m2"): - return "bf8" - raise ValueError(f"expected fp8/bf8 low-bit dtype, got {dtype!r}") + """Map a normalized catalog dtype to its stable instruction-format token.""" + try: + return _LOWBIT_FORMATS[normalize_dtype(dtype)] + except KeyError: + raise ValueError(f"unsupported low-bit matrix dtype: {dtype!r}") from None def _wire_scale_dtype(dtype: str) -> str: @@ -79,8 +83,8 @@ class BlockScaledGemmSpec: M: int N: int K: int - dtype_a: str = "fp8" - dtype_b: str = "fp8" + dtype_a: str = "fp8e4m3" + dtype_b: str = "fp8e4m3" dtype_c: str = "bf16" dtype_acc: str = "fp32" scale_dtype: str = "fp32" @@ -91,6 +95,10 @@ class BlockScaledGemmSpec: tile_n: int = 16 tile_k: int = 128 + def __post_init__(self) -> None: + object.__setattr__(self, "dtype_a", normalize_dtype(self.dtype_a)) + object.__setattr__(self, "dtype_b", normalize_dtype(self.dtype_b)) + @property def block_size(self) -> int: return _WAVE @@ -114,6 +122,20 @@ def resolved_matrix_path(self) -> str: return self.matrix_path +def _native_scaled_atom(spec: BlockScaledGemmSpec, target: ArchTarget): + scale_dtype = "e8m0" if spec.scale_dtype == "i8" else spec.scale_dtype + return target.mma.op_for_shape( + family="wmma_scaled", + a_dtype=spec.dtype_a, + b_dtype=spec.dtype_b, + c_dtype=spec.dtype_acc, + scales=(scale_dtype, scale_dtype, spec.block_k), + m=_BLOCK_M, + n=_BLOCK_N, + k=_WMMA_SCALE_K, + ) + + def is_valid_spec(spec: BlockScaledGemmSpec, arch: str = "gfx1250") -> Tuple[bool, str]: """Return ``(ok, reason)`` for the gfx1250 block-scaled GEMM contract.""" from ...core.arch import ArchTarget @@ -142,7 +164,10 @@ def is_valid_spec(spec: BlockScaledGemmSpec, arch: str = "gfx1250") -> Tuple[boo if spec.M <= 0 or spec.N <= 0 or spec.K <= 0: return False, f"M/N/K must be positive (got M={spec.M}, N={spec.N}, K={spec.K})" - if spec.dtype_a not in _LOWBIT_DTYPES or spec.dtype_b not in _LOWBIT_DTYPES: + if ( + normalize_dtype(spec.dtype_a) not in _LOWBIT_DTYPES + or normalize_dtype(spec.dtype_b) not in _LOWBIT_DTYPES + ): return False, ( f"A/B must be fp8 or bf8 (got A={spec.dtype_a!r}, B={spec.dtype_b!r})" ) @@ -150,11 +175,7 @@ def is_valid_spec(spec: BlockScaledGemmSpec, arch: str = "gfx1250") -> Tuple[boo native_scale = matrix_path in ("wmma_scale", "wmma_scale16") family = matrix_path if native_scale else "wmma" atom_k = _WMMA_SCALE_K if native_scale else _WMMA_K - if native_scale and ( - _canon_lowbit(spec.dtype_a) != "fp8" or _canon_lowbit(spec.dtype_b) != "fp8" - ): - return False, "native gfx1250 SCALE/SCALE16 currently supports fp8 x fp8 only" - if not target.mma.has_shape( + if not native_scale and not target.mma.has_shape( family=family, a_dtype=_canon_lowbit(spec.dtype_a), b_dtype=_canon_lowbit(spec.dtype_b), @@ -174,17 +195,22 @@ def is_valid_spec(spec: BlockScaledGemmSpec, arch: str = "gfx1250") -> Tuple[boo if spec.layout != "RCR": return False, f"block_scaled_gemm supports RCR only (got {spec.layout!r})" if native_scale: - if spec.scale_dtype not in ("e8m0", "i8"): - return False, ( - "native gfx1250 SCALE/SCALE16 requires packed E8M0 scale bytes " - f"(got {spec.scale_dtype!r})" - ) required_block_k = 16 if matrix_path == "wmma_scale16" else 32 if spec.block_k != required_block_k: return False, ( f"{matrix_path} requires block_k={required_block_k} E8M0 groups " f"(got {spec.block_k})" ) + try: + atom = _native_scaled_atom(spec, target) + except ValueError as exc: + return False, str(exc) + if atom is None: + return ( + False, + "no gfx1250 scaled WMMA atom for the requested operand and scale contract", + ) + else: try: _wire_scale_dtype(spec.scale_dtype) @@ -202,7 +228,10 @@ def is_valid_spec(spec: BlockScaledGemmSpec, arch: str = "gfx1250") -> Tuple[boo return False, "M and N must be multiples of 16" if native_scale: - return True, f"ok: gfx1250 K=128 native {matrix_path} FP8 GEMM" + return ( + True, + f"ok: gfx1250 K=128 native {matrix_path} {_canon_lowbit(spec.dtype_a)} GEMM", + ) return True, "ok: gfx1250 K=64 FP8/BF8 WMMA block-scaled GEMM" @@ -232,6 +261,7 @@ def block_scaled_gemm_grid(spec: BlockScaledGemmSpec) -> Tuple[int, int, int]: def _storage_type(dtype: str) -> Type: + dtype = normalize_dtype(dtype) if dtype in ("fp16", "f16"): return F16 if dtype == "bf16": @@ -250,20 +280,21 @@ def _as_f32(b: IRBuilder, v): def build_block_scaled_gemm( spec: BlockScaledGemmSpec, arch: str = "gfx1250" ) -> KernelDef: - """Build a gfx1250 FP8/BF8 block-scaled GEMM (RCR, ``C = A @ B^T``). + """Build a gfx1250 block-scaled GEMM (RCR, ``C = A @ B^T``). One wave (32 lanes) computes one 16x16 output tile without LDS. The legacy ``wmma`` path uses K=64 FP8/BF8 atoms, accumulates each ``block_k`` group, and applies FP16/FP32 A/B scales in software. The native ``wmma_scale`` and - ``wmma_scale16`` paths use K=128 FP8 atoms and pass packed E8M0 scale bytes + ``wmma_scale16`` paths use K=128 scaled WMMA atoms and pass packed E8M0 scale bytes directly to the instruction, with K=32 and K=16 scale groups respectively. Lane ``l`` owns output column ``l % 16`` and rows ``(l // 16) * 8 : (l // 16 + 1) * 8``. Legacy matrix fragments carry 32 - low-bit bytes per lane as ``<8 x i32>``. Native fragments carry 64 bytes as + low-bit bytes per lane as ``<8 x i32>``. For FP8/BF8, native fragments carry 64 bytes as ``<16 x i32>`` as four 16-byte K chunks, alternating chunks between lane halves. Both paths use the gfx12 column-distributed ``<8 x f32>`` accumulator layout. + Scale arrays remain A_scale[M, K/block_k] and B_scale[K/block_k, N]. """ ok, reason = is_valid_spec(spec, arch=arch) if not ok: @@ -275,13 +306,13 @@ def build_block_scaled_gemm( matrix_path = spec.resolved_matrix_path() native_scale = matrix_path in ("wmma_scale", "wmma_scale16") scale_ty = I8 if native_scale else _scale_type(spec.scale_dtype) - op_id = ( - f"{matrix_path}_f32_16x16x128_fp8_fp8" - if native_scale - else _wmma_op_id(spec.dtype_a, spec.dtype_b) + atom = ( + _native_scaled_atom(spec, ArchTarget.from_gfx(arch)) if native_scale else None ) + op_id = atom.op_id if atom is not None else _wmma_op_id(spec.dtype_a, spec.dtype_b) + scale_op = gfx1250_scaled_wmma(op_id) atom_k = _WMMA_SCALE_K if native_scale else _WMMA_K - frag_words = 16 if native_scale else _ACC + frag_words = atom.a_frag_len if atom is not None else _ACC a_frag_ty = VectorType(I32, frag_words) groups = spec.K // spec.block_k @@ -353,23 +384,29 @@ def _load_frag(ptr, base, storage_ty, k0): packed = ir.vec_concat(packed, chunk) return ir.bitcast(packed, a_frag_ty) - def _pack_strided_scales(ptr, row_or_col, call_idx, *, for_b): - count = 8 if matrix_path == "wmma_scale16" else 4 - word_ty = I64 if count == 8 else I32 - packed = ir.const_i64(0) if count == 8 else ir.const_i32(0) + def _pack_strided_scales(ptr, call_idx, *, for_b): + assert scale_op is not None and atom is not None + layout = atom.b_scale_layout() if for_b else atom.a_scale_layout() + packing = scale_op.scales + count = packing.count + word_ty = I64 if packing.word_bits == 64 else I32 + word_const = ir.const_i64 if packing.word_bits == 64 else ir.const_i32 + packed = word_const(0) scale_groups = spec.K // spec.block_k for j in range(count): - scale_group = call_idx * count + j + coord0, coord1 = layout.coord(ir, lane, j) + group_offset = ir.const_i32(call_idx * (atom.k // spec.block_k)) if for_b: - idx = ir.add(ir.mul(ir.const_i32(scale_group), cN), row_or_col) + group = ir.add(group_offset, coord0) + col = ir.add(n0, coord1) + idx = ir.add(ir.mul(group, cN), col) else: - idx = ir.add( - ir.mul(row_or_col, ir.const_i32(scale_groups)), - ir.const_i32(scale_group), - ) + row = ir.add(m0, coord0) + group = ir.add(group_offset, coord1) + idx = ir.add(ir.mul(row, ir.const_i32(scale_groups)), group) byte = ir.global_load(ptr, idx, I8, align=1) widened = ir.zext(byte, word_ty) - shift = ir.const_i64(j * 8) if count == 8 else ir.const_i32(j * 8) + shift = word_const(j * packing.element_bits) shifted = ir.shl(widened, shift) packed = ir.lor(packed, shifted) return packed @@ -380,8 +417,8 @@ def _pack_strided_scales(ptr, row_or_col, call_idx, *, for_b): k0 = step * _WMMA_SCALE_K a_frag = _load_frag(A, a_base, a_ty, k0) b_frag = _load_frag(B, b_base, b_ty, k0) - a_scale = _pack_strided_scales(AScale, a_row, step, for_b=False) - b_scale = _pack_strided_scales(BScale, b_row, step, for_b=True) + a_scale = _pack_strided_scales(AScale, step, for_b=False) + b_scale = _pack_strided_scales(BScale, step, for_b=True) acc = ir.mma(op_id, a_frag, b_frag, acc, a_scale, b_scale) out_col = ir.add(n0, frag) 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 b53d8f87b7df..8a27e62f01ef 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,33 @@ if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/core/smoke.cpp") endif() if(BUILD_TESTING) + add_executable(rocke_dtypes core/test_dtypes.cpp) + target_link_libraries(rocke_dtypes PRIVATE rocke_core) + add_test(NAME rocke_dtypes COMMAND rocke_dtypes) + set_tests_properties(rocke_dtypes PROPERTIES LABELS "ckc;unit;host") + rocke_cxx_test_installed_name(_rocke_installed_name rocke_dtypes) + set_target_properties(rocke_dtypes + PROPERTIES OUTPUT_NAME "${_rocke_installed_name}") + install(TARGETS rocke_dtypes + RUNTIME DESTINATION ${ROCKE_CXX_TEST_INSTALL_DIR}) + set_property(GLOBAL APPEND PROPERTY ROCKE_INSTALL_TEST_TARGETS rocke_dtypes) + + add_executable(rocke_mma_catalog_lookup core/test_mma_catalog_lookup.cpp) + find_package(Threads REQUIRED) + target_link_libraries(rocke_mma_catalog_lookup PRIVATE rocke_core Threads::Threads) + target_include_directories(rocke_mma_catalog_lookup PRIVATE + "${CMAKE_CURRENT_SOURCE_DIR}/../cpp/core/arch") + add_test(NAME rocke_mma_catalog_lookup COMMAND rocke_mma_catalog_lookup) + set_tests_properties(rocke_mma_catalog_lookup PROPERTIES LABELS + "ckc;unit;host") + rocke_cxx_test_installed_name(_rocke_installed_name rocke_mma_catalog_lookup) + set_target_properties(rocke_mma_catalog_lookup + PROPERTIES OUTPUT_NAME "${_rocke_installed_name}") + install(TARGETS rocke_mma_catalog_lookup + RUNTIME DESTINATION ${ROCKE_CXX_TEST_INSTALL_DIR}) + set_property(GLOBAL APPEND PROPERTY ROCKE_INSTALL_TEST_TARGETS + rocke_mma_catalog_lookup) + add_executable(rocke_arch_target_identity core/test_arch_target_identity.cpp) target_link_libraries(rocke_arch_target_identity PRIVATE rocke_core) add_test(NAME rocke_arch_target_identity COMMAND rocke_arch_target_identity) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp index af39a12bcf2d..4ef24a30488c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp @@ -19,6 +19,7 @@ #include #include +#include "rocke/arch_target.h" #include "rocke/ir.h" #include "rocke/lower_hip.h" #include "rocke/lower_llvm.h" @@ -793,10 +794,13 @@ void case_gfx1250_scaled_wmma() rocke_value_t* fragment = rocke_b_vec_concat(b, lo, hi); rocke_value_t* c = rocke_b_global_load_vN(b, accum, lane, rocke_f32(), 8, 0); rocke_value_t* scale = rocke_b_global_load(b, scales, lane, v.scale_type(), 1); - rocke_value_t* d = v.scale16 ? rocke_b_wmma_scale16_f32_16x16x128_fp8_fp8( - b, fragment, fragment, c, scale, scale) - : rocke_b_wmma_scale_f32_16x16x128_fp8_fp8( - b, fragment, fragment, c, scale, scale); + const auto block = v.scale16 ? ROCKE_MMA_SCALE_K16 : ROCKE_MMA_SCALE_K32; + const rocke_mma_scale_filter_t scales_query = {"e8m0", "e8m0", block}; + const rocke_arch_target_t* target = rocke_arch_target_from_gfx("gfx1250"); + const rocke_mma_op_t* atom = rocke_mma_catalog_op_for_shape( + &target->mma, "wmma_scaled", "fp8", "fp8", "fp32", 16, 16, 128, &scales_query); + rocke_value_t* extra[2] = {scale, scale}; + rocke_value_t* d = rocke_b_mma(b, atom->op_id, fragment, fragment, c, extra, 2); rocke_b_global_store(b, accum, lane, rocke_b_vec_extract(b, d, 0), 1); }; const char* name = v.scale16 ? "wmma_scale16" : "wmma_scale"; @@ -884,6 +888,59 @@ void case_opcode_names_are_aligned() } } +/* Malformed result lists must reach the public lowering error boundary. */ +void case_wmma_result_count() +{ + for(bool integer : {false, true}) + { + for(int count : {0, 2, 1}) + { + rocke_ir_builder_t b; + if(rocke_ir_builder_init(&b, "wmma_result_count") != ROCKE_OK) + { + fail("rocke_ir_builder_init", __LINE__); + return; + } + auto* a = rocke_b_zero_vec(&b, integer ? rocke_i32() : rocke_f16(), integer ? 4 : 16); + auto* c = rocke_b_zero_vec(&b, integer ? rocke_i32() : rocke_f32(), 8); + auto* d = rocke_b_mma(&b, + integer ? "wmma_i32_16x16x16_iu8" : "wmma_f32_16x16x16_f16", + a, + a, + c, + nullptr, + 0); + rocke_value_t* results[] = {d, d}; + d->op->num_results = count; + d->op->results = count ? results : nullptr; + rocke_b_ret(&b); + char* out = nullptr; + char err[ROCKE_ERR_MSG_CAP] = {}; + const auto status = rocke_lower_kernel_to_llvm_ex( + b.kernel, ROCKE_LLVM_FLAVOR_LLVM23, "gfx1151", &out, err, sizeof(err)); + if(count == 1) + { + if(status != ROCKE_OK || !out) + fail("valid WMMA failed", __LINE__); + else + EXPECT_IR(std::string(out), integer ? "call <8 x i32>" : "call <8 x float>"); + } + else + { + char expected[80]; + snprintf(expected, + sizeof(expected), + "tile.mma: expected exactly one result, got %d", + count); + if(status != ROCKE_ERR_VALUE || out || strcmp(err, expected) != 0) + fail("malformed WMMA did not return its result-count error", __LINE__); + } + std::free(out); + rocke_ir_builder_free(&b); + } + } +} + struct TestCase { const char* name; @@ -891,6 +948,7 @@ struct TestCase }; const TestCase k_cases[] = { + {"wmma_result_count", case_wmma_result_count}, {"ds_swizzle_raw_offset", case_ds_swizzle_raw_offset}, {"quad_perm_hip", case_quad_perm_hip}, {"quad_perm_rejects_invalid_input", case_quad_perm_rejects_invalid_input}, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/mma_frag_ssot.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/mma_frag_ssot.cpp index ceb797c3d5cc..9cad6977a3f5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/mma_frag_ssot.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/mma_frag_ssot.cpp @@ -99,8 +99,10 @@ int main(void) CHECK(gfx1250 != NULL, "gfx1250 target"); if(gfx1250) { - check_catalog_fragments(&gfx1250->mma, "wmma_scale_f32_16x16x128_fp8_fp8"); - check_catalog_fragments(&gfx1250->mma, "wmma_scale16_f32_16x16x128_fp8_fp8"); + check_catalog_fragments(&gfx1250->mma, + "wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k32"); + check_catalog_fragments(&gfx1250->mma, + "wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k16"); } /* Unknown op_id must be rejected. The engine's error path either returns diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_arch_mma_ssot.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_arch_mma_ssot.py index fe55778e8f03..34b7a8576513 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_arch_mma_ssot.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_arch_mma_ssot.py @@ -11,12 +11,20 @@ from __future__ import annotations +import json import unittest +from dataclasses import asdict, replace from unittest import mock +import pytest + +from rocke.core.arch import ArchTarget, MmaCatalog, MmaOp, MmaScaleBlockK, MmaScaleDType +from rocke.core.arch.wmma_scale import gfx1250_scaled_wmma + from rocke.core.arch.target import ( _load_specs, _op_id_c_dtype, + _op_id_family, normalize_dtype, ) @@ -60,5 +68,328 @@ def test_raises_on_cross_arch_disagreement(self): _op_id_c_dtype.cache_clear() +def _contracts(): + base = MmaOp( + family="wmma_scaled", + a_dtype="fp8e4m3", + b_dtype="bf8e5m2", + c_dtype="fp32", + a_scale_dtype="e8m0", + b_scale_dtype="e8m0", + scale_block_k=32, + m=16, + n=16, + k=128, + op_id="fixture", + ) + rows = [base] + for field in ("a_scale_dtype", "b_scale_dtype"): + for dtype in ("e4m3", "e5m3"): + rows.append(replace(base, **{field: dtype})) + rows.append(replace(base, scale_block_k=16)) + rows.append( + replace(base, a_scale_dtype=None, b_scale_dtype=None, scale_block_k=None) + ) + return [replace(row, op_id=f"fixture_{i}") for i, row in enumerate(rows)] + + +@pytest.mark.parametrize("use_strings", [False, True]) +def test_full_contract_queries_distinguish_each_scale_type_and_shared_block( + use_strings, +): + rows = _contracts() + assert len({row.op_id for row in rows}) == len(rows) + catalog = MmaCatalog(rows) + for row in rows: + query = dict( + family=row.family, + a_dtype=row.a_dtype, + b_dtype=row.b_dtype, + c_dtype=row.c_dtype, + scales=( + ( + str(row.a_scale_dtype) + if use_strings and row.a_scale_dtype + else row.a_scale_dtype + ), + ( + str(row.b_scale_dtype) + if use_strings and row.b_scale_dtype + else row.b_scale_dtype + ), + row.scale_block_k, + ), + m=row.m, + n=row.n, + ) + assert catalog.enumerate(**query) == [row] + assert catalog.has_shape(**query, k=row.k) + assert catalog.op_for_shape(**query, k=row.k) is row + assert catalog.select_largest_k(**query) is row + assert catalog.select_largest_k(**query, k_max=64) is None + + +def test_partial_query_rejects_ambiguity_but_enumeration_and_existence_are_valid(): + catalog = MmaCatalog(_contracts()) + query = dict( + family="wmma_scaled", + a_dtype="fp8", + b_dtype="bf8", + c_dtype="fp32", + m=16, + n=16, + ) + assert len(catalog.enumerate(**query)) == 7 + assert catalog.has_shape(**query, k=128) + with pytest.raises(ValueError, match="ambiguous MMA query"): + catalog.op_for_shape(**query, k=128) + with pytest.raises(ValueError, match="ambiguous MMA query"): + catalog.select_largest_k(**query) + assert len(catalog.enumerate(**query, scales=(None, None, None))) == 1 + assert catalog.op_for_shape(**query, scales=("e5m3", "e5m3", 32), k=128) is None + with pytest.raises(ValueError, match="exactly 3"): + catalog.enumerate(**query, scales=(None,)) + + +def test_scale_alias_selects_the_same_contract(): + catalog = MmaCatalog(_contracts()) + query = dict( + family="wmma_scaled", + a_dtype="fp8", + b_dtype="bf8", + c_dtype="f32", + m=16, + n=16, + k=128, + ) + for scales, field in ( + (("fp8e4m3", "e8m0", 32), "a_scale_dtype"), + (("e8m0", "fp8e4m3", 32), "b_scale_dtype"), + ): + row = catalog.op_for_shape(**query, scales=scales) + assert row is not None and getattr(row, field) is MmaScaleDType.E4M3 + assert getattr(replace(row, **{field: "fp8e4m3"}), field) is MmaScaleDType.E4M3 + + +@pytest.mark.parametrize("dtype", list(MmaScaleDType)) +def test_scale_dtype_preserves_string_and_json_behavior(dtype): + assert isinstance(dtype, str) + assert dtype == dtype.value + assert str(dtype) == f"{dtype}" == dtype.value + assert {dtype: 1}[dtype.value] == 1 + assert json.dumps(dtype) == json.dumps(dtype.value) + assert MmaScaleDType(json.loads(json.dumps(dtype))) is dtype + row = replace(_contracts()[0], a_scale_dtype=dtype.value, b_scale_dtype=dtype) + assert row.a_scale_dtype is row.b_scale_dtype is dtype + decoded = json.loads(json.dumps(asdict(row))) + assert decoded["a_scale_dtype"] == decoded["b_scale_dtype"] == dtype.value + restored = MmaOp(**decoded) + assert restored == row + assert restored.a_scale_dtype is restored.b_scale_dtype is dtype + + +def test_scale_dtype_members_and_alias(): + assert [dtype.value for dtype in MmaScaleDType] == ["e8m0", "e4m3", "e5m3"] + assert MmaScaleDType("fp8e4m3") is MmaScaleDType.E4M3 + + +@pytest.mark.parametrize("dtype", ["e5m2", "fp8e5m2", "bf8e5m2", "bf8"]) +@pytest.mark.parametrize("field", ["a_scale_dtype", "b_scale_dtype"]) +def test_e5m2_is_not_a_scale_dtype_alias(dtype, field): + with pytest.raises(ValueError): + MmaScaleDType(dtype) + scales = ["e8m0", "e8m0", 32] + scales[0 if field == "a_scale_dtype" else 1] = dtype + catalog = MmaCatalog(_contracts()) + query = dict( + family="wmma_scaled", + a_dtype="fp8", + b_dtype="bf8", + c_dtype="fp32", + scales=tuple(scales), + m=16, + n=16, + ) + for method, kwargs in ( + (catalog.enumerate, query), + (catalog.has_shape, {**query, "k": 128}), + (catalog.op_for_shape, {**query, "k": 128}), + (catalog.select_largest_k, query), + ): + with pytest.raises(ValueError, match="MMA scale dtype"): + method(**kwargs) + + +def test_scaled_catalog_identity_and_backend_contract(): + catalog = ArchTarget.from_gfx("gfx1250").mma + rows = [row for row in catalog.ops if row.family == "wmma_scaled"] + assert len(rows) == 4 + assert len({row.op_id for row in rows}) == 4 + for row in rows: + dtype = {"fp8e4m3": "fp8", "bf8e5m2": "bf8"}[row.a_dtype] + assert row.op_id == ( + f"wmma_gfx1250_f32_16x16x128_{dtype}_{dtype}" + f"_scale_e8m0_e8m0_k{row.scale_block_k}" + ) + assert row.a_scale_dtype == row.b_scale_dtype == "e8m0" + assert row.a_scale_dtype is row.b_scale_dtype is MmaScaleDType.E8M0 + assert isinstance(row.scale_block_k, MmaScaleBlockK) + packing = gfx1250_scaled_wmma(row.op_id) + assert packing.atom is row + assert packing.matrix_formats == ( + (0, 0) if row.a_dtype == "fp8e4m3" else (1, 1) + ) + assert packing.scales.count * packing.scales.block_k == row.k + assert packing.matrix_llvm_types == ("<16 x i32>", "<16 x i32>") + for family in ("wmma_scale", "wmma_scale16"): + old_id = f"{family}_f32_16x16x128_fp8_fp8" + assert catalog.by_op_id(old_id) is None + assert gfx1250_scaled_wmma(old_id) is None + + +def test_scale_block_k_is_a_two_value_enum(): + assert list(MmaScaleBlockK) == [MmaScaleBlockK.K16, MmaScaleBlockK.K32] + for block in MmaScaleBlockK: + assert replace(_contracts()[0], scale_block_k=block).scale_block_k is block + assert replace(_contracts()[0], scale_block_k=int(block)).scale_block_k is block + with pytest.raises(ValueError): + MmaScaleBlockK(64) + + +@pytest.mark.parametrize("field", ["a_scale_dtype", "b_scale_dtype"]) +@pytest.mark.parametrize( + "dtype", ["i32", "fp4", "e5m2", "fp8e5m2", "bf8e5m2", "bf8", "", None] +) +def test_invalid_scale_format(field, dtype): + with pytest.raises(ValueError, match="e8m0, e4m3, or e5m3"): + replace(_contracts()[0], **{field: dtype}) + + +@pytest.mark.parametrize("block", [0, 8, 64, 16.0, True, "32", None]) +def test_invalid_scale_block_size(block): + with pytest.raises(ValueError, match="integer equal to 16 or 32"): + replace(_contracts()[0], scale_block_k=block) + + +@pytest.mark.parametrize( + "scales", + [ + ("e8m0", None, 32), + (None, "e8m0", 32), + (None, None, 32), + ("e8m0", "e8m0", None), + ("i32", "e8m0", 32), + ("e8m0", "e8m0", 0), + ("e8m0", "e8m0", 16.0), + ("e8m0", "e8m0", True), + ], +) +def test_invalid_scale_query_even_for_empty_catalog(scales): + with pytest.raises(ValueError, match="MMA scale"): + MmaCatalog([]).enumerate( + family="wmma_scaled", + a_dtype="fp8", + b_dtype="fp8", + c_dtype="fp32", + scales=scales, + ) + + +def test_unscaled_defaults_have_no_scale_metadata(): + row = MmaOp("wmma", "fp8e4m3", "fp8e4m3", "fp32", 16, 16, 64, "fixture") + assert (row.a_scale_dtype, row.b_scale_dtype, row.scale_block_k) == ( + None, + None, + None, + ) + assert ( + MmaCatalog([row]).op_for_shape( + family="wmma", + a_dtype="fp8", + b_dtype="fp8", + c_dtype="fp32", + m=16, + n=16, + k=64, + scales=(None, None, None), + ) + is row + ) + + +def test_largest_k_only_rejects_ties_at_the_selected_k(): + rows = _contracts() + unique = replace(rows[0], k=256) + catalog = MmaCatalog([*rows, unique]) + query = dict( + family="wmma_scaled", a_dtype="fp8", b_dtype="bf8", c_dtype="fp32", m=16, n=16 + ) + assert catalog.select_largest_k(**query) is unique + with pytest.raises(ValueError, match="ambiguous MMA query"): + catalog.select_largest_k(**query, k_max=128) + + if __name__ == "__main__": # pragma: no cover unittest.main() + + +def test_op_id_family_matches_all_catalogs_and_rejects_conflicts(): + specs = _load_specs() + for row in specs.values(): + for op in row["mma"]: + assert _op_id_family()[op["op_id"]] == op["family"] + assert _op_id_family().get("unknown") is None + sample = next(op for row in specs.values() for op in row["mma"]) + drifted = {**specs, "synthetic": {"mma": [{**sample, "family": "conflicting"}]}} + _op_id_family.cache_clear() + try: + with mock.patch("rocke.core.arch.target._load_specs", return_value=drifted): + with pytest.raises(ValueError, match="inconsistent family"): + _op_id_family() + finally: + _op_id_family.cache_clear() + + +def test_mma_naming_without_target_lookup(): + from rocke.core.ir import F16, F32, I32, IRBuilder + + atoms = [ + ( + ArchTarget.from_gfx("gfx950").mma.by_op_id("mfma_f32_16x16x16_f16"), + F16, + "acc", + ), + ( + ArchTarget.from_gfx("gfx1151").mma.by_op_id("wmma_i32_16x16x16_iu8"), + I32, + "acc", + ), + ( + next( + op + for op in ArchTarget.from_gfx("gfx1250").mma.ops + if op.family == "wmma_scaled" + ), + I32, + "mxacc", + ), + ] + _op_id_family.cache_clear() + with mock.patch.object( + ArchTarget, + "from_gfx", + side_effect=AssertionError("target lookup in generic MMA"), + ): + for atom, elem, hint in atoms: + assert atom is not None + for arg in (atom, atom.op_id): + b = IRBuilder("neutral_mma") + a = b.zero_vec(elem, atom.a_frag_len) + c = b.zero_vec(I32 if atom.c_dtype == "i32" else F32, atom.c_frag_len) + extra = ( + (b.const_i32(0), b.const_i32(0)) + if atom.family == "wmma_scaled" + else () + ) + result = b.mma(arg, a, a, c, *extra) + assert result.name.lstrip("%").startswith(hint) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_dtypes.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_dtypes.cpp new file mode 100644 index 000000000000..84e94d2d63e1 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_dtypes.cpp @@ -0,0 +1,39 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* The public dtype API needs no architecture or IR-builder header. + * Mirrors tests/core/test_dtypes.py, including unknown-name pass-through. + */ +#include "rocke/dtypes.h" + +#include +#include + +int main() +{ + const char* cases[][2] = {{" HALF ", "fp16"}, + {"bfloat16", "bf16"}, + {" Float\t", "fp32"}, + {"FP8", "fp8e4m3"}, + {"BF8", "bf8e5m2"}, + {"FP6", "fp6e2m3"}, + {"BF6", "fp6e3m2"}, + {"FP4", "fp4e2m1"}, + {"int32", ROCKE_DTYPE_I32}, + {" Custom_Format ", "custom_format"}, + {"", ""}}; + for(const auto& entry : cases) + { + char scratch[64]; + if(strcmp(rocke_normalize_dtype(entry[0], scratch, sizeof(scratch)), entry[1]) != 0 + || strcmp(rocke_normalize_dtype(entry[1], scratch, sizeof(scratch)), entry[1]) != 0) + { + fprintf(stderr, "dtype normalization failed: %s\n", entry[0]); + return 1; + } + } + char scratch[64]; + if(strcmp(rocke_normalize_dtype(NULL, scratch, sizeof(scratch)), "") != 0 + || strcmp(rocke_normalize_dtype(" BF6 ", NULL, 0), "fp6e3m2") != 0) + return 1; + return 0; +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_dtypes.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_dtypes.py new file mode 100644 index 000000000000..47014e02e8a2 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_dtypes.py @@ -0,0 +1,37 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Core dtype names are independent of target support and retain public aliases.""" + +import pytest + +from rocke.core.dtypes import normalize_dtype + + +@pytest.mark.parametrize( + "spelling,canonical", + [ + (" HALF ", "fp16"), + ("bfloat16", "bf16"), + (" Float\t", "fp32"), + ("FP8", "fp8e4m3"), + ("BF8", "bf8e5m2"), + ("FP6", "fp6e2m3"), + ("BF6", "fp6e3m2"), + ("FP4", "fp4e2m1"), + ("int32", "i32"), + (" Custom_Format ", "custom_format"), + ("", ""), + ], +) +def test_normalize_dtype(spelling, canonical): + assert normalize_dtype(spelling) == canonical + assert normalize_dtype(canonical) == canonical + + +def test_architecture_entry_points_reexport_core_normalization(): + from rocke.core import normalize_dtype as core_normalize + from rocke.core.arch import normalize_dtype as arch_normalize + from rocke.core.arch.target import normalize_dtype as target_normalize + + assert normalize_dtype.__module__ == "rocke.core.dtypes" + assert core_normalize is arch_normalize is target_normalize is normalize_dtype diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_gfx1250_mxfp8_example.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_gfx1250_mxfp8_example.py new file mode 100644 index 000000000000..763e2b1020ed --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_gfx1250_mxfp8_example.py @@ -0,0 +1,78 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Both homogeneous eight-bit formats are usable from the focused example.""" + +import argparse + +import numpy as np +import pytest + +from rocke.core.lower_hip import lower_kernel_to_hip +from rocke.core.dtypes import normalize_dtype +from rocke.examples.gfx1250.gemm import _scaled_gemm_example, mxfp8_gemm +from rocke.examples.gfx1250.gemm.block_scaled_gemm_verify import ( + make_case_inputs, + reference_result, +) +from rocke.instances.gfx1250.block_scaled_gemm import build_block_scaled_gemm + + +@pytest.mark.parametrize( + "dtype,selector,ml_name", [("fp8", 0, "float8_e4m3fn"), ("bf8", 1, "float8_e5m2")] +) +@pytest.mark.parametrize("path", ["wmma_scale", "wmma_scale16"]) +def test_mxfp8_formats_and_hip_selectors(dtype, selector, ml_name, path): + ml = pytest.importorskip("ml_dtypes") + spec = mxfp8_gemm.make_spec( + argparse.Namespace(m=32, n=48, k=256, dtype=dtype, matrix_path=path) + ) + assert spec.dtype_a == spec.dtype_b == normalize_dtype(dtype) + kernel = build_block_scaled_gemm(spec) + hip = lower_kernel_to_hip(kernel, arch="gfx1250") + assert f"__builtin_amdgcn_{path}_f32_16x16x128_f8f6f4({selector}," in hip + inputs = make_case_inputs(spec, "mixed") + assert inputs[0].dtype == inputs[1].dtype == np.dtype(getattr(ml, ml_name)) + ref = reference_result(*inputs, spec.block_k, native=True) + assert ref.shape == (32, 48) and np.isfinite(ref).all() + + +@pytest.mark.parametrize( + "argv,dtypes", + [ + (["--dtype", "fp8e4m3"], ("fp8e4m3",)), + (["--dtype", " FP8 "], ("fp8e4m3",)), + (["--dtype", "bf8e5m2"], ("bf8e5m2",)), + (["--dtype", " BF8 "], ("bf8e5m2",)), + ([], ("fp8e4m3", "bf8e5m2")), + (["--dtype", "both"], ("fp8e4m3", "bf8e5m2")), + (["--dtype", "fp8"], ("fp8e4m3",)), + (["--dtype", "bf8"], ("bf8e5m2",)), + ], +) +def test_mxfp8_cli_selects_formats(monkeypatch, argv, dtypes): + calls = [] + + def run(spec, cases, **kwargs): + calls.append((spec, cases, kwargs)) + return len(cases) + + monkeypatch.setattr(_scaled_gemm_example, "run_cases", run) + assert ( + mxfp8_gemm.main( + argv + + [ + "--matrix-path", + "wmma_scale16", + "--compile-route", + "hip", + "--case", + "all", + ] + ) + == 0 + ) + assert tuple(spec.dtype_a for spec, _, _ in calls) == dtypes + for spec, cases, kwargs in calls: + assert spec.dtype_a == spec.dtype_b + assert spec.block_k == 16 and len(cases) == 20 + assert kwargs == {"compile_route": "hip"} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_gfx1250_scaled_wmma.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_gfx1250_scaled_wmma.py index 71028eed346a..5e66d0728d75 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_gfx1250_scaled_wmma.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_gfx1250_scaled_wmma.py @@ -5,9 +5,13 @@ from __future__ import annotations import unittest +from dataclasses import replace +from unittest import mock from rocke.core.arch import ArchTarget +from rocke.core.arch.wmma_scale import gfx1250_scaled_wmma from rocke.core.ir import F32, I32, I64, IRBuilder, PtrType +from rocke.core.ir_serialize import parse, serialize from rocke.core.lower_hip import lower_kernel_to_hip from rocke.core.lower_llvm import lower_kernel_to_llvm from rocke.instances.gfx1250.block_scaled_gemm import ( @@ -18,6 +22,22 @@ ) +def _scaled_atom(dtype="fp8", scale16=False): + block_k = 16 if scale16 else 32 + atom = ArchTarget.from_gfx("gfx1250").mma.op_for_shape( + family="wmma_scaled", + a_dtype=dtype, + b_dtype=dtype, + c_dtype="fp32", + scales=("e8m0", "e8m0", block_k), + m=16, + n=16, + k=128, + ) + assert atom is not None + return atom + + def _build_scaled_atom(*, scale16: bool): b = IRBuilder("gfx1250_scaled_wmma") matrix = b.param("matrix", PtrType(I32, "global"), readonly=True) @@ -30,20 +50,139 @@ def _build_scaled_atom(*, scale16: bool): fragment = b.vec_concat(lo, hi) c = b.global_load_vN(accum, lane, F32, 8) scale = b.global_load(scales, lane, scale_ty) - if scale16: - d = b.wmma_scale16_f32_16x16x128_fp8_fp8(fragment, fragment, c, scale, scale) - else: - d = b.wmma_scale_f32_16x16x128_fp8_fp8(fragment, fragment, c, scale, scale) + d = b.mma(_scaled_atom(scale16=scale16), fragment, fragment, c, scale, scale) b.global_store(accum, lane, b.vec_extract(d, 0)) return b.kernel class TestGfx1250ScaledWmma(unittest.TestCase): + def test_backend_rejects_unsupported_scale_type_on_either_input(self): + catalog = ArchTarget.from_gfx("gfx1250").mma + for field in ("a_scale_dtype", "b_scale_dtype"): + atom = replace(_scaled_atom(), **{field: "e4m3"}) + with ( + self.subTest(field=field), + mock.patch.object(catalog, "by_op_id", return_value=atom), + self.assertRaisesRegex(ValueError, "unsupported scaled WMMA"), + ): + gfx1250_scaled_wmma(atom.op_id) + + def test_scale_fragment_counts_are_checked_for_each_source(self): + catalog = ArchTarget.from_gfx("gfx1250").mma + for role in ("a_scale", "b_scale"): + for count in (0, 8): + atom = replace( + _scaled_atom(), + **{ + f"{role}_frag_len": count, + f"_{role}_layout": None, + }, + ) + with ( + self.subTest(role=role, count=count), + mock.patch.object(catalog, "by_op_id", return_value=atom), + self.assertRaisesRegex(ValueError, "scale fragment lengths"), + ): + gfx1250_scaled_wmma(atom.op_id) + + def test_loader_requires_both_scale_maps(self): + for role in ("a_scale", "b_scale"): + atom = replace(_scaled_atom(), **{f"_{role}_layout": None}) + spec = BlockScaledGemmSpec( + "missing_scale_map", + M=16, + N=16, + K=128, + block_k=32, + scale_dtype="e8m0", + matrix_path="wmma_scale", + ) + with ( + mock.patch( + "rocke.instances.gfx1250.block_scaled_gemm._native_scaled_atom", + return_value=atom, + ), + self.assertRaisesRegex(NotImplementedError, role), + ): + build_block_scaled_gemm(spec, arch="gfx1250") + # Lowering prepacked operands needs the carrier contract, not a loader map. + with mock.patch.object( + ArchTarget.from_gfx("gfx1250").mma, "by_op_id", return_value=atom + ): + self.assertEqual(gfx1250_scaled_wmma(atom.op_id).scales.count, 4) + + def test_backend_derives_each_matrix_carrier_and_selector(self): + # Synthetic metadata exercises the lowering contract without adding atoms. + atom = replace(_scaled_atom(), b_dtype="bf8e5m2", b_frag_len=8) + catalog = ArchTarget.from_gfx("gfx1250").mma + with mock.patch.object(catalog, "by_op_id", return_value=atom): + spec = gfx1250_scaled_wmma(atom.op_id) + self.assertEqual(spec.matrix_formats, (0, 1)) + self.assertEqual(spec.matrix_llvm_types, ("<16 x i32>", "<8 x i32>")) + self.assertTrue(spec.intrinsic.endswith("v16i32.v8i32")) + self.assertIn("wmma.scale.block32", spec.declaration_key) + + def test_matrix_formats_share_intrinsic_declarations(self): + b = IRBuilder("shared_scaled_wmma_declarations") + matrix = b.param("matrix", PtrType(I32, "global"), readonly=True) + accum = b.param("accum", PtrType(F32, "global")) + lane = b.thread_id_x() + lo = b.global_load_vN(matrix, lane, I32, 8) + hi = b.global_load_vN(matrix, b.add(lane, b.const_i32(8)), I32, 8) + fragment = b.vec_concat(lo, hi) + c = b.global_load_vN(accum, lane, F32, 8) + for mode, scale_ty in (("scale", I32), ("scale16", I64)): + scale_ptr = b.param( + f"{mode}_ptr", PtrType(scale_ty, "global"), readonly=True + ) + scale = b.global_load(scale_ptr, lane, scale_ty) + for dtype in ("fp8", "bf8"): + c = b.mma( + _scaled_atom(dtype, mode == "scale16"), + fragment, + fragment, + c, + scale, + scale, + ) + b.global_store(accum, lane, b.vec_extract(c, 0)) + llvm = lower_kernel_to_llvm(b.kernel, arch="gfx1250", llvm_flavor="llvm23") + for mode, scale_ty in (("scale", "i32"), ("scale16", "i64")): + intrinsic = ( + f"llvm.amdgcn.wmma.{mode}.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32" + ) + with self.subTest(mode=mode): + self.assertEqual(llvm.count(f"declare <8 x float> @{intrinsic}("), 1) + calls = [ + line + for line in llvm.splitlines() + if f"call <8 x float> @{intrinsic}(" in line + ] + self.assertEqual(len(calls), 2) + for selector, call in enumerate(calls): + self.assertEqual(call.count(f"i32 {selector}, <16 x i32>"), 2) + self.assertEqual(call.count(f", {scale_ty} %"), 2) + + def test_catalog_ids_survive_serialization(self): + for scale16 in (False, True): + kernel = _build_scaled_atom(scale16=scale16) + text = serialize(kernel) + self.assertIn(_scaled_atom(scale16=scale16).op_id, text) + restored = parse(text) + self.assertEqual( + lower_kernel_to_llvm(kernel, arch="gfx1250", llvm_flavor="llvm23"), + lower_kernel_to_llvm(restored, arch="gfx1250", llvm_flavor="llvm23"), + ) + self.assertEqual( + lower_kernel_to_hip(kernel, arch="gfx1250"), + lower_kernel_to_hip(restored, arch="gfx1250"), + ) + def test_catalog_fragment_lengths(self): target = ArchTarget.from_gfx("gfx1250") for op_id in ( - "wmma_scale_f32_16x16x128_fp8_fp8", - "wmma_scale16_f32_16x16x128_fp8_fp8", + "wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k32", + "wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k16", ): with self.subTest(op_id=op_id): op = target.mma.by_op_id(op_id) @@ -88,8 +227,9 @@ def test_scale_and_scale16_exact_llvm23_abi(self): def test_scaled_wmma_rejects_pre_llvm23_flavors(self): for flavor in ("llvm20", "llvm22"): - with self.subTest(flavor=flavor), self.assertRaisesRegex( - NotImplementedError, "requires llvm23" + with ( + self.subTest(flavor=flavor), + self.assertRaisesRegex(NotImplementedError, "requires llvm23"), ): lower_kernel_to_llvm( _build_scaled_atom(scale16=False), @@ -108,6 +248,27 @@ def test_scale_and_scale16_hip_builtin_lowering(self): ) self.assertIn(builtin, hip) + def test_scaled_hip_dispatch_preserves_neutral_and_concrete_ops(self): + for scale16 in (False, True): + mode = "wmma_scale16" if scale16 else "wmma_scale" + for dtype, selector in (("fp8", 0), ("bf8", 1)): + with self.subTest(mode=mode, dtype=dtype): + kernel = _build_scaled_atom(scale16=scale16) + call = next(op for op in kernel.body.ops if op.name == "tile.mma") + op_id = _scaled_atom(dtype, scale16).op_id + call.attrs["op_id"] = op_id + neutral = lower_kernel_to_hip(kernel, arch="gfx1250") + self.assertIn( + f"__builtin_amdgcn_{mode}_f32_16x16x128_f8f6f4({selector},", + neutral, + ) + call.name = f"tile.{call.attrs.pop('op_id')}" + self.assertEqual( + neutral, lower_kernel_to_hip(kernel, arch="gfx1250") + ) + with self.assertRaisesRegex(NotImplementedError, "not available"): + lower_kernel_to_hip(kernel, arch="gfx942") + def test_native_block_scaled_gemm_uses_packed_e8m0_in_instruction(self): for matrix_path, block_k, scale_ty, fragment_load, load_count in ( ("wmma_scale", 32, "i32", "load <16 x i8>", 8), @@ -181,7 +342,7 @@ def test_native_slice_rejects_unsupported_contracts(self): ) ok, why = is_valid_spec(bad_dtype) self.assertFalse(ok) - self.assertIn("fp8 x fp8 only", why) + self.assertIn("requested operand and scale contract", why) bad_block = BlockScaledGemmSpec( name="bad_block", @@ -196,6 +357,21 @@ def test_native_slice_rejects_unsupported_contracts(self): self.assertFalse(ok) self.assertIn("requires block_k=16", why) + def test_unknown_scaled_atom_reports_supported_operations(self): + for scale16 in (False, True): + with self.subTest(scale16=scale16): + kernel = _build_scaled_atom(scale16=scale16) + call = next(op for op in kernel.body.ops if op.name == "tile.mma") + mode = "wmma_scale16" if scale16 else "wmma_scale" + call.attrs["op_id"] = f"{mode}_f32_16x16x128_fp16_fp16" + with self.assertRaisesRegex( + NotImplementedError, "not yet wired for gfx1250" + ) as error: + lower_kernel_to_llvm(kernel, arch="gfx1250", llvm_flavor="llvm23") + self.assertIn(_scaled_atom("bf8", scale16).op_id, str(error.exception)) + with self.assertRaisesRegex(NotImplementedError, "no HIP lowering"): + lower_kernel_to_hip(kernel, arch="gfx1250") + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_lowbit_dtype_aliases.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_lowbit_dtype_aliases.py new file mode 100644 index 000000000000..fe590a754013 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_lowbit_dtype_aliases.py @@ -0,0 +1,98 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Public dtype aliases resolve to explicit catalog keys without changing ops.""" + +import pytest +from rocke.core.arch import ArchTarget +from rocke.core.dtypes import normalize_dtype + + +@pytest.mark.parametrize( + "alias,canonical", + [ + ("fp8", "fp8e4m3"), + ("bf8", "bf8e5m2"), + ("fp6", "fp6e2m3"), + ("bf6", "fp6e3m2"), + ("fp4", "fp4e2m1"), + ], +) +def test_alias_catalog_lookup(alias, canonical): + assert normalize_dtype(alias) == canonical + assert normalize_dtype(" " + alias.upper() + " ") == canonical + assert normalize_dtype(canonical) == canonical + found = False + for gfx in ("gfx950", "gfx1250"): + catalog = ArchTarget.from_gfx(gfx).mma + for op in catalog.ops: + assert op.a_dtype not in ("fp8", "bf8", "fp6", "bf6", "fp4") + assert op.b_dtype not in ("fp8", "bf8", "fp6", "bf6", "fp4") + if op.a_dtype != canonical: + continue + found = True + assert catalog.has_shape( + family=op.family, + scales=(op.a_scale_dtype, op.b_scale_dtype, op.scale_block_k), + a_dtype=alias, + b_dtype=op.b_dtype, + c_dtype=op.c_dtype, + m=op.m, + n=op.n, + k=op.k, + ) + # BF6 name recognition does not require a corresponding catalog atom. + if alias != "bf6": + assert found + + +@pytest.mark.parametrize("dtype", ["fp4", "fp4e2m1", "fp6", "fp6e2m3", "fp16"]) +def test_fp8_example_rejects_other_families(dtype): + from rocke.examples.gfx1250.gemm import mxfp8_gemm + + with pytest.raises(SystemExit) as exc: + mxfp8_gemm.main(["--dtype", dtype]) + assert exc.value.code == 2 + + +def test_supported_aliases_preserve_lowered_kernel(): + from dataclasses import replace + from rocke.core.lower_llvm import lower_kernel_to_llvm + from rocke.core.lower_hip import lower_kernel_to_hip + from rocke.instances.gfx1250.block_scaled_gemm import ( + _LOWBIT_DTYPES, + BlockScaledGemmSpec, + build_block_scaled_gemm, + block_scaled_gemm_signature, + ) + + for short in ("fp8", "bf8", "fp6", "bf6", "fp4"): + canonical = normalize_dtype(short) + if canonical not in _LOWBIT_DTYPES: + continue + for mode, block_k in (("wmma_scale", 32), ("wmma_scale16", 16)): + spec = BlockScaledGemmSpec( + name="alias_identity", + M=32, + N=48, + K=256, + dtype_a=short, + dtype_b=short, + scale_dtype="e8m0", + matrix_path=mode, + block_k=block_k, + ) + explicit = replace(spec, dtype_a=canonical, dtype_b=canonical) + mixed_case = replace(spec, dtype_a=" " + short.upper() + " ", dtype_b=short) + assert spec == explicit == mixed_case + assert spec.dtype_a == spec.dtype_b == canonical + assert hash(spec) == hash(explicit) == hash(mixed_case) + a, b = build_block_scaled_gemm(spec), build_block_scaled_gemm(explicit) + assert block_scaled_gemm_signature(spec) == block_scaled_gemm_signature( + explicit + ) + assert lower_kernel_to_llvm( + a, arch="gfx1250", llvm_flavor="llvm23" + ) == lower_kernel_to_llvm(b, arch="gfx1250", llvm_flavor="llvm23") + assert lower_kernel_to_hip(a, arch="gfx1250") == lower_kernel_to_hip( + b, arch="gfx1250" + ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_catalog_lookup.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_catalog_lookup.cpp new file mode 100644 index 000000000000..0ca6d1c65b9e --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_catalog_lookup.cpp @@ -0,0 +1,460 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* Exercise public catalog queries against stored keys, without normalizing + * away malformed table entries. Alias queries must find the same atom. */ +#include "mma_family_index.h" +#include "rocke/arch_target.h" +#include "rocke/error.hpp" +#include "rocke/helper_rocke.core.arch.h" +#include "rocke/wmma_scale_internal.h" + +#include +#include +#include +#include +#include + +static const char* short_dtype(const char* dtype) +{ + const char* aliases[][2] = {{"fp8e4m3", "fp8"}, + {"bf8e5m2", "bf8"}, + {"fp6e2m3", "fp6"}, + {"fp6e3m2", "bf6"}, + {"fp4e2m1", "fp4"}}; + for(const auto& pair : aliases) + { + if(strcmp(dtype, pair[0]) == 0) + return pair[1]; + } + return dtype; +} + +#define CHECK(condition) \ + do \ + { \ + if(!(condition)) \ + { \ + fprintf(stderr, "line %d: %s\\n", __LINE__, #condition); \ + return 1; \ + } \ + } while(0) + +template +static bool rejects_query(F query) +{ + try + { + query(); + } + catch(const ckc::Error& e) + { + return e.code() == ROCKE_ERR_VALUE; + } + return false; +} + +static int test_family_index_first_use() +{ + std::atomic ready{0}; + std::atomic start{false}; + std::thread workers[8]; + bool passed[8] = {}; + for(int i = 0; i < 8; ++i) + workers[i] = std::thread([&, i] { + ++ready; + while(!start.load()) + std::this_thread::yield(); + int count = 0; + const auto* arches = rocke_known_arches(&count); + bool ok = true; + for(int a = 0; a < count; ++a) + { + const auto* arch = rocke_arch_target_from_gfx(arches[(a + i) % count]); + for(int j = 0; j < arch->mma.num_ops; ++j) + { + const auto& op = arch->mma.ops[j]; + char copied_id[256]; + snprintf(copied_id, sizeof(copied_id), "%s", op.op_id); + const char* family = rocke_arch_mma_op_id_family(copied_id); + ok &= family && strcmp(family, op.family) == 0; + ok &= family == rocke_arch_mma_op_id_family(op.op_id); + } + } + passed[i] = ok && !rocke_arch_mma_op_id_family(NULL) && !rocke_arch_mma_op_id_family("") + && !rocke_arch_mma_op_id_family("unknown"); + }); + while(ready.load() != 8) + std::this_thread::yield(); + start = true; + for(auto& worker : workers) + worker.join(); + for(bool ok : passed) + CHECK(ok); + return 0; +} + +static int test_family_index_duplicates() +{ + rocke_mma_op_t ops[6] = {}; + const char* ids[] = {"same", "conflict", "z", "same", "conflict", "a"}; + const char* families[] = {"wmma", "mma", "wmma_scaled", "wmma", "wmma", "mma"}; + for(int i = 0; i < 6; ++i) + { + ops[i].op_id = ids[i]; + ops[i].family = families[i]; + } + rocke_arch_target_t targets[2] = {}; + targets[0].mma = {ops, 3}; + targets[1].mma = {ops + 3, 3}; + const rocke_ati_arch_row_t registry[] + = {{"first", &targets[0]}, {"missing", NULL}, {"second", &targets[1]}}; + const rocke_ati_mma_family_index index(registry, 3); + CHECK(strcmp(index.lookup("same"), "wmma") == 0); + CHECK(strcmp(index.lookup("a"), "mma") == 0); + CHECK(strcmp(index.lookup("z"), "wmma_scaled") == 0); + CHECK(!index.lookup(NULL) && !index.lookup("") && !index.lookup("between")); + bool conflict = false; + try + { + index.lookup("conflict"); + } + catch(const ckc::Error& e) + { + conflict + = e.code() == ROCKE_ERR_VALUE + && strcmp(e.what(), "arch SSOT drift: op_id has inconsistent family across arches") + == 0; + } + CHECK(conflict); + CHECK(strcmp(index.lookup("a"), "mma") == 0); + const rocke_ati_mma_family_index empty(NULL, 0); + CHECK(!empty.lookup("same")); + return 0; +} + +static int test_mma_result_names() +{ + const char* ids[] = {"mfma_f32_16x16x16_f16", + "wmma_f32_16x16x16_f16", + "wmma_gfx1250_f32_16x16x64_fp8_fp8", + "wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k32", + "wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k16", + "wmma_gfx1250_f32_16x16x128_bf8_bf8_scale_e8m0_e8m0_k32", + "wmma_gfx1250_f32_16x16x128_bf8_bf8_scale_e8m0_e8m0_k16"}; + for(int i = 0; i < 7; ++i) + { + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "mma_names") == ROCKE_OK); + auto* value = rocke_b_const_i32(&b, 0); + auto* result = rocke_b_mma(&b, ids[i], value, value, value, NULL, 0); + CHECK(result && result->name); + CHECK(strncmp(result->name, i < 3 ? "%acc" : "%mxacc", i < 3 ? 4 : 6) == 0); + rocke_ir_builder_free(&b); + } + return 0; +} + +static int test_scale_contracts() +{ + const auto* arch = rocke_arch_target_from_gfx("gfx1250"); + const rocke_mma_scale_filter_t e8 = {"e8m0", "e8m0", ROCKE_MMA_SCALE_K32}; + const auto* base = rocke_mma_catalog_op_for_shape( + &arch->mma, "wmma_scaled", "fp8", "fp8", "fp32", 16, 16, 128, &e8); + CHECK(base); + rocke_mma_op_t rows[8]; + rows[0] = *base; + rows[0].b_dtype = "bf8e5m2"; + auto packing_atom = rows[0]; + packing_atom.b_frag_len = 8; + const auto packing = rocke_scaled_wmma_contract(&packing_atom); + CHECK(packing.matrix_formats[0] == 0 && packing.matrix_formats[1] == 1); + CHECK(packing.matrix_words[0] == 16 && packing.matrix_words[1] == 8); + CHECK(strstr(packing.intrinsic, "v16i32.v8i32")); + for(int input = 0; input < 2; ++input) + { + auto unsupported = *base; + (input == 0 ? unsupported.a_scale_dtype : unsupported.b_scale_dtype) = "e4m3"; + CHECK(rejects_query([&] { rocke_scaled_wmma_contract(&unsupported); })); + } + int count = 1; + for(int input = 0; input < 2; ++input) + { + for(const char* dtype : {"e4m3", "e5m3"}) + { + rows[count] = rows[0]; + (input == 0 ? rows[count].a_scale_dtype : rows[count].b_scale_dtype) = dtype; + ++count; + } + } + rows[count] = rows[0]; + rows[count++].scale_block_k = ROCKE_MMA_SCALE_K16; + rows[count] = rows[0]; + rows[count].a_scale_dtype = NULL; + rows[count].b_scale_dtype = NULL; + rows[count++].scale_block_k = ROCKE_MMA_SCALE_NONE; + rocke_mma_catalog_t cat = {rows, count}; + for(int i = 0; i < count; ++i) + { + const auto& op = rows[i]; + const rocke_mma_scale_filter_t scales + = {op.a_scale_dtype, op.b_scale_dtype, op.scale_block_k}; + CHECK(rocke_mma_catalog_enumerate( + &cat, "wmma_scaled", "fp8", "bf8", "f32", 16, 16, NULL, 0, &scales) + == 1); + CHECK(rocke_mma_catalog_has_shape( + &cat, "wmma_scaled", "fp8", "bf8", "f32", 16, 16, 128, &scales)); + CHECK(rocke_mma_catalog_op_for_shape( + &cat, "wmma_scaled", "fp8", "bf8", "f32", 16, 16, 128, &scales) + == &op); + CHECK(rocke_mma_catalog_select_largest_k( + &cat, "wmma_scaled", "fp8", "bf8", "f32", 16, 16, -1, &scales) + == &op); + CHECK(!rocke_mma_catalog_select_largest_k( + &cat, "wmma_scaled", "fp8", "bf8", "f32", 16, 16, 64, &scales)); + } + CHECK(rocke_mma_catalog_enumerate(&cat, "wmma_scaled", "fp8", "bf8", "fp32", 16, 16, NULL, 0) + == count); + CHECK(rocke_mma_catalog_has_shape(&cat, "wmma_scaled", "fp8", "bf8", "fp32", 16, 16, 128)); + CHECK(rejects_query([&] { + rocke_mma_catalog_op_for_shape(&cat, "wmma_scaled", "fp8", "bf8", "fp32", 16, 16, 128); + })); + CHECK(rejects_query([&] { + rocke_mma_catalog_select_largest_k(&cat, "wmma_scaled", "fp8", "bf8", "fp32", 16, 16, -1); + })); + const rocke_mma_scale_filter_t unscaled = {NULL, NULL, ROCKE_MMA_SCALE_NONE}; + CHECK(rocke_mma_catalog_op_for_shape( + &cat, "wmma_scaled", "fp8", "bf8", "fp32", 16, 16, 128, &unscaled) + == &rows[6]); + const rocke_mma_scale_filter_t alias = {"fp8e4m3", "e8m0", ROCKE_MMA_SCALE_K32}; + CHECK(rocke_mma_catalog_op_for_shape( + &cat, "wmma_scaled", "fp8", "bf8", "fp32", 16, 16, 128, &alias) + == &rows[1]); + const rocke_mma_scale_filter_t alias_b = {"e8m0", "fp8e4m3", ROCKE_MMA_SCALE_K32}; + CHECK(rocke_mma_catalog_op_for_shape( + &cat, "wmma_scaled", "fp8", "bf8", "fp32", 16, 16, 128, &alias_b) + == &rows[3]); + rows[7] = rows[0]; + rows[7].k = 256; + cat.num_ops = 8; + CHECK(rocke_mma_catalog_select_largest_k(&cat, "wmma_scaled", "fp8", "bf8", "fp32", 16, 16, -1) + == &rows[7]); + CHECK(rejects_query([&] { + rocke_mma_catalog_select_largest_k(&cat, "wmma_scaled", "fp8", "bf8", "fp32", 16, 16, 128); + })); + for(const auto scales : + {rocke_mma_scale_filter_t{"i32", "e8m0", ROCKE_MMA_SCALE_K32}, + rocke_mma_scale_filter_t{"e8m0", "i32", ROCKE_MMA_SCALE_K32}, + rocke_mma_scale_filter_t{"e8m0", "e8m0", static_cast(8)}, + rocke_mma_scale_filter_t{NULL, "e8m0", ROCKE_MMA_SCALE_K32}, + rocke_mma_scale_filter_t{"e8m0", NULL, ROCKE_MMA_SCALE_K32}, + rocke_mma_scale_filter_t{NULL, NULL, ROCKE_MMA_SCALE_K32}, + rocke_mma_scale_filter_t{"e8m0", "e8m0", ROCKE_MMA_SCALE_NONE}}) + { + CHECK(rejects_query([&] { + rocke_mma_catalog_enumerate( + &cat, "wmma_scaled", "fp8", "bf8", "fp32", 16, 16, NULL, 0, &scales); + })); + } + for(const char* dtype : {"e5m2", "fp8e5m2", "bf8e5m2", "bf8"}) + { + for(int input = 0; input < 2; ++input) + { + rocke_mma_scale_filter_t scales = {"e8m0", "e8m0", ROCKE_MMA_SCALE_K32}; + (input == 0 ? scales.a_scale_dtype : scales.b_scale_dtype) = dtype; + CHECK(rejects_query([&] { + rocke_mma_catalog_enumerate( + &cat, "wmma_scaled", "fp8", "bf8", "fp32", 16, 16, NULL, 0, &scales); + })); + CHECK(rejects_query([&] { + rocke_mma_catalog_has_shape( + &cat, "wmma_scaled", "fp8", "bf8", "fp32", 16, 16, 128, &scales); + })); + CHECK(rejects_query([&] { + rocke_mma_catalog_op_for_shape( + &cat, "wmma_scaled", "fp8", "bf8", "fp32", 16, 16, 128, &scales); + })); + CHECK(rejects_query([&] { + rocke_mma_catalog_select_largest_k( + &cat, "wmma_scaled", "fp8", "bf8", "fp32", 16, 16, -1, &scales); + })); + } + } + const rocke_mma_scale_filter_t valid_missing = {"e5m3", "e5m3", ROCKE_MMA_SCALE_K32}; + CHECK(!rocke_mma_catalog_op_for_shape( + &cat, "wmma_scaled", "fp8", "bf8", "fp32", 16, 16, 128, &valid_missing)); + char id[256]; + int packed_rows = 0, scaled_rows = 0; + for(int i = 0; i < arch->mma.num_ops; ++i) + { + const auto& op = arch->mma.ops[i]; + if(strncmp(op.op_id, "wmma_gfx1250_f32_16x16x64_", 25) == 0) + { + CHECK(op.a_frag_len == 8 && op.b_frag_len == 8 && op.c_frag_len == 8); + ++packed_rows; + } + if(strcmp(op.family, "wmma_scaled") == 0) + { + snprintf(id, + sizeof(id), + "wmma_gfx1250_f32_16x16x128_%s_%s_scale_%s_%s_k%d", + short_dtype(op.a_dtype), + short_dtype(op.b_dtype), + op.a_scale_dtype, + op.b_scale_dtype, + op.scale_block_k); + CHECK(strcmp(id, op.op_id) == 0); + CHECK(rocke_mma_catalog_by_op_id(&arch->mma, id) == &op); + ++scaled_rows; + } + } + CHECK(packed_rows == 4 && scaled_rows == 4); + return 0; +} + +static int test_scale_layouts_and_families() +{ + CHECK(!rocke_arch_mma_op_id_family(NULL)); + CHECK(!rocke_arch_mma_op_id_family("unknown")); + int num_arches = 0; + const char* const* arches = rocke_known_arches(&num_arches); + int scaled = 0; + for(int i = 0; i < num_arches; ++i) + { + const auto* arch = rocke_arch_target_from_gfx(arches[i]); + for(int j = 0; j < arch->mma.num_ops; ++j) + { + const auto* op = &arch->mma.ops[j]; + CHECK(strcmp(rocke_arch_mma_op_id_family(op->op_id), op->family) == 0); + if(!op->a_scale_dtype) + { + CHECK(op->a_scale_frag_len == 0 && op->b_scale_frag_len == 0); + CHECK(!op->a_scale_layout && !op->b_scale_layout); + for(auto getter : {rocke_mma_op_a_scale_layout, rocke_mma_op_b_scale_layout}) + { + bool rejected = false; + try + { + getter(op, NULL); + } + catch(const ckc::Error& e) + { + rejected = e.code() == ROCKE_ERR_NOTIMPL; + } + CHECK(rejected); + } + continue; + } + ++scaled; + CHECK(strcmp(arches[i], "gfx1250") == 0); + const int count = op->scale_block_k == ROCKE_MMA_SCALE_K16 ? 8 : 4; + CHECK(op->a_scale_frag_len == count && op->b_scale_frag_len == count); + for(int source = 0; source < 2; ++source) + { + const auto* map = source ? rocke_mma_op_b_scale_layout(op, NULL) + : rocke_mma_op_a_scale_layout(op, NULL); + CHECK(map && map->frag_len == count && map->wave_size == 32); + CHECK(map->role == (source ? ROCKE_MMA_ROLE_B_SCALE : ROCKE_MMA_ROLE_A_SCALE)); + for(int slot = -1; slot <= count; ++slot) + { + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "scale_layout") == ROCKE_OK); + auto* lane = rocke_b_thread_id_x(&b); + rocke_value_t *x = NULL, *y = NULL; + const bool ok = rocke_layout_map_coord(map, &b, lane, slot, &x, &y); + if(slot < 0 || slot == count) + { + CHECK(!ok && !x && !y && b.status == ROCKE_ERR_VALUE); + CHECK(strstr(b.err, source ? "'b_scale'" : "'a_scale'")); + } + else + CHECK(ok && x && y && b.status == ROCKE_OK); + rocke_ir_builder_free(&b); + } + auto invalid = *op; + (source ? invalid.b_scale_frag_len : invalid.a_scale_frag_len) = count / 2; + CHECK(rejects_query([&] { rocke_scaled_wmma_contract(&invalid); })); + } + } + } + CHECK(scaled == 4); + return 0; +} + +int main() +{ + // Keep this first: no previous query may prime the shared family index. + if(test_family_index_first_use() || test_family_index_duplicates() || test_mma_result_names() + || test_scale_contracts() || test_scale_layouts_and_families()) + return 1; + int checked = 0; + for(const char* gfx : {"gfx950", "gfx1250"}) + { + const rocke_arch_target_t* arch = rocke_arch_target_from_gfx(gfx); + if(!arch) + return 1; + for(int i = 0; i < arch->mma.num_ops; ++i) + { + const rocke_mma_op_t* op = &arch->mma.ops[i]; + for(const char* dtype : {op->a_dtype, op->b_dtype, op->c_dtype}) + { + char scratch[64]; + if(strcmp(dtype, rocke_normalize_dtype(dtype, scratch, sizeof(scratch))) != 0) + { + fprintf(stderr, "%s: noncanonical stored dtype %s\n", op->op_id, dtype); + return 1; + } + } + const rocke_mma_scale_filter_t scales + = {op->a_scale_dtype, op->b_scale_dtype, op->scale_block_k}; + for(const char* a : {op->a_dtype, short_dtype(op->a_dtype)}) + { + for(const char* b : {op->b_dtype, short_dtype(op->b_dtype)}) + { + if(!rocke_mma_catalog_has_shape( + &arch->mma, op->family, a, b, op->c_dtype, op->m, op->n, op->k, &scales) + || rocke_mma_catalog_op_for_shape(&arch->mma, + op->family, + a, + b, + op->c_dtype, + op->m, + op->n, + op->k, + &scales) + != op + || rocke_archtarget_op_for_shape( + arch, op->family, a, b, op->c_dtype, op->m, op->n, op->k, &scales) + != op + || rocke_mma_catalog_select_largest_k(&arch->mma, + op->family, + a, + b, + op->c_dtype, + op->m, + op->n, + op->k, + &scales) + != op + || rocke_mma_catalog_enumerate(&arch->mma, + op->family, + a, + b, + op->c_dtype, + op->m, + op->n, + NULL, + 0, + &scales) + < 1) + { + fprintf(stderr, "%s: catalog query missed %s/%s\n", op->op_id, a, b); + return 1; + } + } + } + ++checked; + } + } + printf("catalog lookup: %d rows passed canonical and alias queries\n", checked); + return 0; +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_scale_layouts.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_scale_layouts.py new file mode 100644 index 000000000000..42b2dcbbe176 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_scale_layouts.py @@ -0,0 +1,174 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Independent lane-coordinate, tensor-address and packed-scale oracles.""" + +from collections import Counter +from dataclasses import replace +import operator + +import pytest + +from rocke.core.arch import ArchTarget, known_arches +from rocke.core.ir import IRBuilder +from rocke.instances.gfx1250.block_scaled_gemm import ( + BlockScaledGemmSpec, + build_block_scaled_gemm, +) + + +def _atom(dtype, block_k): + return ArchTarget.from_gfx("gfx1250").mma.op_for_shape( + family="wmma_scaled", + a_dtype=dtype, + b_dtype=dtype, + c_dtype="fp32", + m=16, + n=16, + k=128, + scales=("e8m0", "e8m0", block_k), + ) + + +def _evaluate(value, *, lane=0, tile=(0, 0), loads=None): + """Evaluate only the integer SSA dependency graph, never the MMA itself.""" + op = value.op + if op.name == "arith.constant": + return op.attrs["value"] + if op.name == "gpu.thread_id": + return lane + if op.name == "gpu.block_id": + return tile[0 if op.attrs["axis"] == "x" else 1] + + def evaluate(v): + return _evaluate(v, lane=lane, tile=tile, loads=loads) + + if op.name == "memref.global_load_typed": + ptr, index = op.operands + offset = evaluate(index) + role = ptr.name.lstrip("%") + assert role in ("A_scale", "B_scale") + loads.append((role, offset)) + # Distinct address-sensitive bytes expose stride, group and order bugs. + return (offset * 37 + (11 if role == "A_scale" else 73)) % 256 + if op.name == "arith.zext": + return evaluate(op.operands[0]) & 255 + operations = { + "arith.add": operator.add, + "arith.mul": operator.mul, + "arith.div": operator.floordiv, + "arith.mod": operator.mod, + "arith.shl": operator.lshift, + "arith.or": operator.or_, + } + assert op.name in operations, op.name + return operations[op.name](*(evaluate(v) for v in op.operands)) + + +@pytest.mark.parametrize("dtype", ["fp8", "bf8"]) +@pytest.mark.parametrize("block_k", [16, 32]) +@pytest.mark.parametrize("role", ["a_scale", "b_scale"]) +def test_scale_coordinates(dtype, block_k, role): + atom = _atom(dtype, block_k) + layout = getattr(atom, f"{role}_layout")() + count = 128 // block_k + assert getattr(atom, f"{role}_frag_len") == count + assert (layout.role, layout.wave_size, layout.frag_len) == (role, 32, count) + b = IRBuilder("scale_coordinates") + lane = b.thread_id_x() + coords = [layout.coord(b, lane, slot) for slot in range(count)] + seen = Counter() + for l in range(32): + for j, pair in enumerate(coords): + actual = tuple(_evaluate(v, lane=l) for v in pair) + expected = (l % 16, j) if role == "a_scale" else (j, l % 16) + assert actual == expected + seen[actual] += 1 + assert len(seen) == 16 * count + assert set(seen.values()) == {2} + for slot in (-1, count): + with pytest.raises(ValueError, match="fragment slot"): + layout.coord(b, lane, slot) + + +def test_scale_defaults_and_unavailable_maps(): + for gfx in known_arches(): + for atom in ArchTarget.from_gfx(gfx).mma.ops: + if atom.a_scale_dtype is None: + assert atom.a_scale_frag_len == atom.b_scale_frag_len == 0 + for accessor in (atom.a_scale_layout, atom.b_scale_layout): + with pytest.raises(NotImplementedError, match="no verified"): + accessor() + atom = _atom("fp8", 32) + unknown = replace(atom, _a_scale_layout=None, _b_scale_layout=None) + with pytest.raises(NotImplementedError, match="a_scale"): + unknown.a_scale_layout() + with pytest.raises(NotImplementedError, match="b_scale"): + unknown.b_scale_layout() + # Adding scale maps does not invent unverified matrix operand maps. + for accessor in (atom.a_layout, atom.b_layout): + with pytest.raises(NotImplementedError): + accessor() + assert ( + replace( + atom, _a_scale_layout=replace(atom.a_scale_layout(), fn=lambda *_: None) + ) + == atom + ) + + +@pytest.mark.parametrize("role", ["a_scale", "b_scale"]) +def test_scale_metadata_rejects_inconsistent_maps(role): + atom = _atom("fp8", 32) + layout = getattr(atom, f"{role}_layout")() + for change in ({"role": "a"}, {"frag_len": 8}, {"wave_size": 64}): + with pytest.raises(ValueError, match="layout does not match"): + replace(atom, **{f"_{role}_layout": replace(layout, **change)}) + for count in (-1, True, 1.5): + with pytest.raises(ValueError, match="fragment length"): + replace(atom, **{f"{role}_frag_len": count}) + + +@pytest.mark.parametrize("dtype", ["fp8", "bf8"]) +@pytest.mark.parametrize("block_k", [16, 32]) +def test_instance_scale_addresses_and_packed_words(dtype, block_k): + spec = BlockScaledGemmSpec( + "scale_addresses", + M=32, + N=48, + K=256, + dtype_a=dtype, + dtype_b=dtype, + scale_dtype="e8m0", + block_k=block_k, + matrix_path="wmma_scale16" if block_k == 16 else "wmma_scale", + ) + kernel = build_block_scaled_gemm(spec, arch="gfx1250") + mmas = [op for op in kernel.body.ops if op.name == "tile.mma"] + assert len(mmas) == 2 + count = 128 // block_k + for x, y in ((0, 0), (2, 1)): + for lane in range(32): + for step, mma in enumerate(mmas): + for source, value in enumerate(mma.operands[3:]): + offsets = [ + ( + (16 * y + lane % 16) * (256 // block_k) + step * count + j + if source == 0 + else (step * count + j) * 48 + 16 * x + lane % 16 + ) + for j in range(count) + ] + role = "A_scale" if source == 0 else "B_scale" + expected = int.from_bytes( + bytes( + (offset * 37 + (11 if source == 0 else 73)) % 256 + for offset in offsets + ), + "little", + ) + loads = [] + assert ( + _evaluate(value, lane=lane, tile=(x, y), loads=loads) + == expected + ) + assert loads == [(role, offset) for offset in offsets] diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json index c27564225f0f..83f2ff5ccb49 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json @@ -1155,7 +1155,24 @@ "status": "ok" } }, - "expected_failures": {}, + "expected_failures": { + "gemm/gfx1250/mxfp8/wmma_scale/bf8e5m2": { + "message": "tile.wmma_gfx1250_f32_16x16x128_bf8_bf8_scale_e8m0_e8m0_k32 requires llvm23 (ROCm 7.13+), got llvm20", + "type": "NotImplementedError" + }, + "gemm/gfx1250/mxfp8/wmma_scale/fp8e4m3": { + "message": "tile.wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k32 requires llvm23 (ROCm 7.13+), got llvm20", + "type": "NotImplementedError" + }, + "gemm/gfx1250/mxfp8/wmma_scale16/bf8e5m2": { + "message": "tile.wmma_gfx1250_f32_16x16x128_bf8_bf8_scale_e8m0_e8m0_k16 requires llvm23 (ROCm 7.13+), got llvm20", + "type": "NotImplementedError" + }, + "gemm/gfx1250/mxfp8/wmma_scale16/fp8e4m3": { + "message": "tile.wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k16 requires llvm23 (ROCm 7.13+), got llvm20", + "type": "NotImplementedError" + } + }, "summary": { "arches": { "gfx11-generic": 1, @@ -1166,7 +1183,7 @@ "gfx942": 32, "gfx950": 85 }, - "expected_failures": 0, + "expected_failures": 4, "families": { "attention": 14, "attention_d128_swa": 1, @@ -2343,7 +2360,24 @@ "status": "ok" } }, - "expected_failures": {}, + "expected_failures": { + "gemm/gfx1250/mxfp8/wmma_scale/bf8e5m2": { + "message": "tile.wmma_gfx1250_f32_16x16x128_bf8_bf8_scale_e8m0_e8m0_k32 requires llvm23 (ROCm 7.13+), got llvm22", + "type": "NotImplementedError" + }, + "gemm/gfx1250/mxfp8/wmma_scale/fp8e4m3": { + "message": "tile.wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k32 requires llvm23 (ROCm 7.13+), got llvm22", + "type": "NotImplementedError" + }, + "gemm/gfx1250/mxfp8/wmma_scale16/bf8e5m2": { + "message": "tile.wmma_gfx1250_f32_16x16x128_bf8_bf8_scale_e8m0_e8m0_k16 requires llvm23 (ROCm 7.13+), got llvm22", + "type": "NotImplementedError" + }, + "gemm/gfx1250/mxfp8/wmma_scale16/fp8e4m3": { + "message": "tile.wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k16 requires llvm23 (ROCm 7.13+), got llvm22", + "type": "NotImplementedError" + } + }, "summary": { "arches": { "gfx11-generic": 1, @@ -2354,7 +2388,7 @@ "gfx942": 32, "gfx950": 85 }, - "expected_failures": 0, + "expected_failures": 4, "families": { "attention": 14, "attention_d128_swa": 1, @@ -3298,6 +3332,38 @@ "sha256": "4602cb57a8b57916659d91de9dd76756616541649600d7c8596d4073c39f48c4", "status": "ok" }, + "gemm/gfx1250/mxfp8/wmma_scale/bf8e5m2": { + "arch": "gfx1250", + "bytes": 89055, + "family": "gemm", + "kernel_name": "irhash_mxfp8_bf8e5m2_wmma_scale_block_scaled_bf8_bf8_M32N48K256_bk32_t16x16x128_wmma_scale", + "sha256": "f9e29be4f431fcbebcdecf60db75024cbd55874a05c3b6d706a998c942335682", + "status": "ok" + }, + "gemm/gfx1250/mxfp8/wmma_scale/fp8e4m3": { + "arch": "gfx1250", + "bytes": 89055, + "family": "gemm", + "kernel_name": "irhash_mxfp8_fp8e4m3_wmma_scale_block_scaled_fp8_fp8_M32N48K256_bk32_t16x16x128_wmma_scale", + "sha256": "e1fcadf213c82d7941bcbd37561d3f8106841517a018d04239593678782fe284", + "status": "ok" + }, + "gemm/gfx1250/mxfp8/wmma_scale16/bf8e5m2": { + "arch": "gfx1250", + "bytes": 95649, + "family": "gemm", + "kernel_name": "irhash_mxfp8_bf8e5m2_wmma_scale16_block_scaled_bf8_bf8_M32N48K256_bk16_t16x16x128_wmma_scale16", + "sha256": "8ca807c7e2b4bc8b9de57841d00b7175a28de970197b10aefadb75612885b74c", + "status": "ok" + }, + "gemm/gfx1250/mxfp8/wmma_scale16/fp8e4m3": { + "arch": "gfx1250", + "bytes": 95649, + "family": "gemm", + "kernel_name": "irhash_mxfp8_fp8e4m3_wmma_scale16_block_scaled_fp8_fp8_M32N48K256_bk16_t16x16x128_wmma_scale16", + "sha256": "43d60d203f40c6d191e8cbf67792cb2fdfe264eec0cbd891d5e19868709bab6c", + "status": "ok" + }, "gemm/gfx1250/t32x32x32": { "arch": "gfx1250", "bytes": 26610, @@ -3537,7 +3603,7 @@ "gfx11-generic": 1, "gfx1151": 10, "gfx1201": 7, - "gfx1250": 4, + "gfx1250": 8, "gfx90a": 5, "gfx942": 32, "gfx950": 85 @@ -3556,12 +3622,12 @@ "conv_wgrad_reduce": 3, "conv_wgrad_two_stage": 2, "deep_fused_conv": 9, - "gemm": 13, + "gemm": 17, "grouped_gemm": 1, "kda_chunkwise": 13, "moe": 9 }, - "ok_cases": 144 + "ok_cases": 148 } } }, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx1250_scaled_wmma_numeric.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/conftest.py similarity index 80% rename from dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx1250_scaled_wmma_numeric.py rename to dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/conftest.py index 9ceb92484625..c6cd4d321b8f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx1250_scaled_wmma_numeric.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/conftest.py @@ -53,24 +53,14 @@ def gpu_env(): return env -@pytest.mark.parametrize( - "matrix_path,route,m,n,k,case,count", - [ - ("wmma_scale", "comgr", 16, 16, 128, "all", 8), - ("wmma_scale16", "comgr", 16, 16, 128, "all", 12), - ("wmma_scale", "comgr", 32, 48, 256, "mixed", 1), - ("wmma_scale16", "comgr", 32, 48, 256, "mixed", 1), - ("wmma_scale", "hip", 32, 48, 256, "mixed", 1), - ("wmma_scale16", "hip", 32, 48, 256, "mixed", 1), - ("wmma", "comgr", 16, 16, 128, "mixed", 1), - ], -) -def test_scaled_wmma_numeric(gpu_env, matrix_path, route, m, n, k, case, count): +def _run_numeric(gpu_env, dtype, matrix_path, route, m, n, k, case, count): result = subprocess.run( [ sys.executable, "-m", _MODULE, + "--dtype", + dtype, "--matrix-path", matrix_path, "--compile-route", @@ -94,3 +84,11 @@ def test_scaled_wmma_numeric(gpu_env, matrix_path, route, m, n, k, case, count): assert f"PASS: verified {count} cases" in output, output assert output.count("bad=0") == count, output print(output, end="") + + +@pytest.fixture +def numeric_case(gpu_env): + """Run one bounded scaled-GEMM case with the validated GPU environment.""" + from functools import partial + + return partial(_run_numeric, gpu_env) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx1250_lowering_emit.c b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx1250_lowering_emit.c index bc5f0b71ea66..77993e4bd0dc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx1250_lowering_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx1250_lowering_emit.c @@ -12,6 +12,7 @@ #include #include +#include "rocke/arch_target.h" #include "rocke/ir.h" #include "rocke/ir_serialize.h" #include "rocke/lower_llvm.h" @@ -119,8 +120,14 @@ static void build_wmma_k64_bf8_bf8(rocke_ir_builder_t* b) /* K=128 FP8 SCALE/SCALE16 WMMA. Matrix fragments are <16 x i32>; packed * E8M0 scale operands are i32 for SCALE and i64 for SCALE16. */ -static void wmma_scaled(rocke_ir_builder_t* b, bool scale16) +static void wmma_scaled(rocke_ir_builder_t* b, bool scale16, const char* dtype) { + const rocke_mma_scale_block_k_t block = scale16 ? ROCKE_MMA_SCALE_K16 : ROCKE_MMA_SCALE_K32; + const rocke_mma_scale_filter_t query = {"e8m0", "e8m0", block}; + const rocke_arch_target_t* target = rocke_arch_target_from_gfx("gfx1250"); + const rocke_mma_op_t* atom = rocke_mma_catalog_op_for_shape( + &target->mma, "wmma_scaled", dtype, dtype, "fp32", 16, 16, 128, &query); + const char* op_id = atom->op_id; const rocke_type_t* scale_ty = scale16 ? rocke_i64() : rocke_i32(); rocke_value_t* a_ptr = frag_param(b, "A", rocke_i32(), true); rocke_value_t* b_ptr = frag_param(b, "B", rocke_i32(), true); @@ -137,21 +144,30 @@ static void wmma_scaled(rocke_ir_builder_t* b, bool scale16) rocke_value_t* bb = rocke_b_vec_concat(b, b_lo, b_hi); rocke_value_t* c = rocke_b_global_load_vN(b, c_ptr, tid, rocke_f32(), 8, /*align=*/0); rocke_value_t* scale = rocke_b_global_load(b, scale_ptr, tid, scale_ty, /*align=*/1); - rocke_value_t* d = scale16 - ? rocke_b_wmma_scale16_f32_16x16x128_fp8_fp8(b, a, bb, c, scale, scale) - : rocke_b_wmma_scale_f32_16x16x128_fp8_fp8(b, a, bb, c, scale, scale); + rocke_value_t* scales[] = {scale, scale}; + rocke_value_t* d = rocke_b_mma(b, op_id, a, bb, c, scales, 2); rocke_b_global_store(b, c_ptr, tid, d, /*align=*/1); rocke_b_ret(b); } static void build_wmma_scale(rocke_ir_builder_t* b) { - wmma_scaled(b, false); + wmma_scaled(b, false, "fp8e4m3"); +} + +static void build_wmma_scale_bf8(rocke_ir_builder_t* b) +{ + wmma_scaled(b, false, "bf8e5m2"); } static void build_wmma_scale16(rocke_ir_builder_t* b) { - wmma_scaled(b, true); + wmma_scaled(b, true, "fp8e4m3"); +} + +static void build_wmma_scale16_bf8(rocke_ir_builder_t* b) +{ + wmma_scaled(b, true, "bf8e5m2"); } /* ds_read_b128_tr_b16. gfx950 has one type-agnostic opcode returning @@ -328,6 +344,39 @@ static void build_tensor_transfers(rocke_ir_builder_t* b) rocke_b_ret(b); } +static void build_scale_coordinates(rocke_ir_builder_t* b, rocke_mma_scale_block_k_t block) +{ + const auto* arch = rocke_arch_target_from_gfx("gfx1250"); + const rocke_mma_scale_filter_t scales = {"e8m0", "e8m0", block}; + const auto* atom = rocke_mma_catalog_op_for_shape( + &arch->mma, "wmma_scaled", "fp8", "fp8", "fp32", 16, 16, 128, &scales); + rocke_value_t* out = frag_param(b, "coords", rocke_i32(), false); + rocke_value_t* lane = rocke_b_thread_id_x(b); + const rocke_layout_map_t* maps[] + = {rocke_mma_op_a_scale_layout(atom, b), rocke_mma_op_b_scale_layout(atom, b)}; + for(int source = 0; source < 2; ++source) + { + for(int slot = 0; slot < maps[source]->frag_len; ++slot) + { + rocke_value_t *x = NULL, *y = NULL; + rocke_layout_map_coord(maps[source], b, lane, slot, &x, &y); + rocke_b_global_store(b, out, lane, x, 1); + rocke_b_global_store(b, out, lane, y, 1); + } + } + rocke_b_ret(b); +} + +static void build_scale_coordinates_k32(rocke_ir_builder_t* b) +{ + build_scale_coordinates(b, ROCKE_MMA_SCALE_K32); +} + +static void build_scale_coordinates_k16(rocke_ir_builder_t* b) +{ + build_scale_coordinates(b, ROCKE_MMA_SCALE_K16); +} + typedef void (*build_fn_t)(rocke_ir_builder_t*); typedef struct config @@ -362,6 +411,10 @@ static const config_t CONFIGS[] = { {build_global_tr16_bf16, "gfx1250"}, {build_global_tr16_i16, "gfx1250"}, {build_tensor_transfers, "gfx1250"}, + {build_wmma_scale_bf8, "gfx1250"}, + {build_wmma_scale16_bf8, "gfx1250"}, + {build_scale_coordinates_k32, "gfx1250"}, + {build_scale_coordinates_k16, "gfx1250"}, }; static const int NUM_CONFIGS = (int)(sizeof(CONFIGS) / sizeof(CONFIGS[0])); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx1250_lowering_emit.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx1250_lowering_emit.py index 80ca5f0f3075..7dd837947469 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx1250_lowering_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/gfx1250_lowering_emit.py @@ -19,6 +19,7 @@ # each with its gfx950 counterpart pins both sides of the choice. # # arch is per-config (see _spec), llvm_flavor = AUTO, matching the C side. +from rocke.core.arch import ArchTarget from rocke.core.ir import BF16, F16, F32, I16, I32, I64, IRBuilder, KernelDef, PtrType from _emit_common import run_emit @@ -77,7 +78,19 @@ def build(b: IRBuilder) -> None: def _wmma_scaled(a_kind, b_kind, scale_mode): """K=128 scaled WMMA, parameterized by operand dtypes and scale mode.""" scale_ty = {"scale": I32, "scale16": I64}[scale_mode] - op_id = f"wmma_{scale_mode}_f32_16x16x128_{a_kind}_{b_kind}" + block_k = 16 if scale_mode == "scale16" else 32 + atom = ArchTarget.from_gfx("gfx1250").mma.op_for_shape( + family="wmma_scaled", + a_dtype=a_kind, + b_dtype=b_kind, + c_dtype="fp32", + scales=("e8m0", "e8m0", block_k), + m=16, + n=16, + k=128, + ) + assert atom is not None + op_id = atom.op_id def build(b: IRBuilder) -> None: a_ptr = b.param( @@ -254,6 +267,40 @@ def build_tensor_transfers(b: IRBuilder) -> None: ] +# Homogeneous BF8 belongs to the eight-bit example contract. +CONFIGS.extend( + (_wmma_scaled("bf8", "bf8", mode), "gfx1250") for mode in ("scale", "scale16") +) + + +def _scale_coordinates(block_k): + atom = ArchTarget.from_gfx("gfx1250").mma.op_for_shape( + family="wmma_scaled", + a_dtype="fp8", + b_dtype="fp8", + c_dtype="fp32", + m=16, + n=16, + k=128, + scales=("e8m0", "e8m0", block_k), + ) + + def build(b): + out = b.param("coords", PtrType(I32, "global"), noalias=True, align=16) + lane = b.thread_id_x() + for layout in (atom.a_scale_layout(), atom.b_scale_layout()): + for slot in range(layout.frag_len): + x, y = layout.coord(b, lane, slot) + b.global_store(out, lane, x) + b.global_store(out, lane, y) + b.ret() + + return build + + +CONFIGS.extend((_scale_coordinates(block), "gfx1250") for block in (32, 16)) + + def _spec(idx: int): """Config selector: the (builder, arch) pair the shared driver expects.""" if not 0 <= idx < len(CONFIGS): diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py index cb2e0c24051f..6eb13309c1ea 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py @@ -1091,6 +1091,31 @@ def _build(): return _build +def build_mxfp8_gemm_case(dtype, matrix_path): + def _build(): + from rocke.instances.gfx1250.block_scaled_gemm import ( + BlockScaledGemmSpec, + build_block_scaled_gemm, + ) + + return build_block_scaled_gemm( + BlockScaledGemmSpec( + name=f"irhash_mxfp8_{dtype}_{matrix_path}", + M=32, + N=48, + K=256, + dtype_a=dtype, + dtype_b=dtype, + dtype_c="bf16", + scale_dtype="e8m0", + matrix_path=matrix_path, + block_k=16 if matrix_path == "wmma_scale16" else 32, + ) + ) + + return _build + + def cases(): out = [] @@ -1366,6 +1391,16 @@ def add(family, case_id, arch, build): ), ) + # Homogeneous FP8/BF8 with E8M0 scales, shared by source and installed gates. + for dtype in ("fp8e4m3", "bf8e5m2"): + for matrix_path in ("wmma_scale", "wmma_scale16"): + add( + "gemm", + f"gemm/gfx1250/mxfp8/{matrix_path}/{dtype}", + "gfx1250", + build_mxfp8_gemm_case(dtype, matrix_path), + ) + # Conv: problem-shape and arch variants. conv1 = (1, 8, 8, 16, 32, 3, 3, 1, 1, 1, 1, 1, 1) conv2 = (2, 16, 16, 32, 32, 1, 1, 1, 1, 0, 0, 1, 1) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx1250_mxfp8_numeric.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx1250_mxfp8_numeric.py new file mode 100644 index 000000000000..12e4b5ebced0 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx1250_mxfp8_numeric.py @@ -0,0 +1,26 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""gfx1250 MXFP8 numerical cases.""" + +import pytest + + +@pytest.mark.parametrize( + "dtype,matrix_path,route,m,n,k,case,count", + [ + (dtype, *case) + for dtype in ("fp8", "bf8") + for case in [ + ("wmma_scale", "comgr", 16, 16, 128, "all", 8), + ("wmma_scale16", "comgr", 16, 16, 128, "all", 12), + ("wmma_scale", "comgr", 32, 48, 256, "mixed", 1), + ("wmma_scale16", "comgr", 32, 48, 256, "mixed", 1), + ("wmma_scale", "hip", 32, 48, 256, "mixed", 1), + ("wmma_scale16", "hip", 32, 48, 256, "mixed", 1), + ("wmma", "comgr", 16, 16, 128, "mixed", 1), + ] + if dtype == "fp8" or case[0] != "wmma" + ], +) +def test_numeric(numeric_case, dtype, matrix_path, route, m, n, k, case, count): + numeric_case(dtype, matrix_path, route, m, n, k, case, count) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx1250_mxfp_examples.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx1250_mxfp_examples.py new file mode 100644 index 000000000000..91e6513daf8c --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/test_gfx1250_mxfp_examples.py @@ -0,0 +1,30 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Focused examples build real kernels and constrain their input contracts.""" + +import argparse +import importlib + +import pytest + +from rocke.core.arch.target import normalize_dtype +from rocke.core.lower_llvm import lower_kernel_to_llvm +from rocke.instances.gfx1250.block_scaled_gemm import build_block_scaled_gemm + + +@pytest.mark.parametrize("family,dtype", [("mxfp8", "fp8")]) +@pytest.mark.parametrize("path,block_k", [("wmma_scale", 32), ("wmma_scale16", 16)]) +def test_example_contract_and_lowering(family, dtype, path, block_k): + example = importlib.import_module(f"rocke.examples.gfx1250.gemm.{family}_gemm") + args = argparse.Namespace(m=32, n=48, k=256, matrix_path=path, dtype=dtype) + spec = example.make_spec(args) + assert (spec.dtype_a, spec.dtype_b, spec.scale_dtype) == ( + normalize_dtype(dtype), + normalize_dtype(dtype), + "e8m0", + ) + assert spec.block_k == block_k + llvm = lower_kernel_to_llvm( + build_block_scaled_gemm(spec), arch="gfx1250", llvm_flavor="llvm23" + ) + assert "@llvm.amdgcn.wmma.scale" in llvm