From d7fb2080895b43829d3880d86978722a8cbbffce Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Tue, 22 Sep 2026 20:34:55 -0400 Subject: [PATCH 01/24] feat(rocke): share logical dtype and packed operand contracts --- .../rocke/platform/cpp/core/dtypes.cpp | 25 ++ .../rocke/platform/cpp/core/ir/core_types.cpp | 32 +++ .../platform/cpp/core/lower_hip/core.cpp | 57 ++++ .../rocke/platform/cpp/core/storage.cpp | 265 ++++++++++++++++++ .../rocke/platform/cpp/helpers/mma_io.cpp | 206 ++++++++++++++ .../rocke/platform/cpp/helpers/quant.cpp | 45 +-- .../rocke/platform/cpp/include/rocke/dtypes.h | 18 ++ .../rocke/helper_rocke.helpers.mma_io.h | 41 +++ .../rocke/helper_rocke.helpers.quant.h | 44 +-- .../rocke/platform/cpp/include/rocke/ir.h | 15 + .../platform/cpp/include/rocke/storage.h | 114 ++++++++ .../cpp/include/rocke/wmma_scale_internal.h | 56 +++- .../architecture/dtype_storage_packing.md | 124 ++++++++ .../platform/python/rocke/core/__init__.py | 19 +- .../python/rocke/core/arch/wmma_scale.py | 68 ++++- .../platform/python/rocke/core/dtypes.py | 67 +++++ .../rocke/platform/python/rocke/core/ir.py | 36 +++ .../platform/python/rocke/core/lower_hip.py | 37 +++ .../platform/python/rocke/core/storage.py | 229 +++++++++++++++ .../platform/python/rocke/helpers/mma_io.py | 162 +++++++++++ .../platform/python/rocke/helpers/quant.py | 37 ++- .../rocke/platform/tests/CMakeLists.txt | 9 + .../rocke/platform/tests/core/test_dtypes.py | 49 ++++ .../rocke/platform/tests/core/test_mma_io.py | 147 ++++++++++ .../platform/tests/core/test_storage.cpp | 186 ++++++++++++ .../rocke/platform/tests/core/test_storage.py | 94 +++++++ 26 files changed, 2093 insertions(+), 89 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/storage.cpp create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.mma_io.h create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/storage.h create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/storage.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/dtypes.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/dtypes.cpp index 0f717a76e35f..b5ab7d5507ae 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/dtypes.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/dtypes.cpp @@ -137,3 +137,28 @@ const char* rocke_normalize_dtype(const char* name, char* lowered, size_t cap) } return lowered; } + +static const rocke_dtype_info_t k_dtype_info[] = { + {"fp32", ROCKE_DTYPE_FLOAT, 32}, {"fp16", ROCKE_DTYPE_FLOAT, 16}, + {"bf16", ROCKE_DTYPE_FLOAT, 16}, {"fp8e4m3", ROCKE_DTYPE_FLOAT, 8}, + {"bf8e5m2", ROCKE_DTYPE_FLOAT, 8}, {"fp6e2m3", ROCKE_DTYPE_FLOAT, 6}, + {"fp6e3m2", ROCKE_DTYPE_FLOAT, 6}, {"fp4e2m1", ROCKE_DTYPE_FLOAT, 4}, + {"e8m0", ROCKE_DTYPE_FLOAT, 8}, {"e4m3", ROCKE_DTYPE_FLOAT, 8}, + {"e5m3", ROCKE_DTYPE_FLOAT, 8}, {"i4", ROCKE_DTYPE_INTEGER, 4}, + {"iu4", ROCKE_DTYPE_INTEGER, 4}, {"i8", ROCKE_DTYPE_INTEGER, 8}, + {"iu8", ROCKE_DTYPE_INTEGER, 8}, {"i16", ROCKE_DTYPE_INTEGER, 16}, + {"i32", ROCKE_DTYPE_INTEGER, 32}, {"i64", ROCKE_DTYPE_INTEGER, 64}, + {"i1", ROCKE_DTYPE_PREDICATE, 1}, +}; + +const rocke_dtype_info_t* rocke_dtype_info(const char* name) +{ + if(!name) + return NULL; + char scratch[64]; + const char* key = rocke_normalize_dtype(name, scratch, sizeof(scratch)); + for(const auto& info : k_dtype_info) + if(strcmp(key, info.name) == 0) + return &info; + return NULL; +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp index bafe79e1f905..d5e5bb48ee58 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp @@ -21,6 +21,7 @@ #include #include "rocke/arena.h" +#include "rocke/dtypes.h" #include "rocke/ir.h" #include "rocke/ir_internal.h" @@ -55,6 +56,13 @@ ROCKE_SCALAR_SINGLETON(rocke_f32, ROCKE_SCALAR_F32, "f32") ROCKE_SCALAR_SINGLETON(rocke_fp8e4m3, ROCKE_SCALAR_FP8E4M3, "fp8e4m3") ROCKE_SCALAR_SINGLETON(rocke_bf8e5m2, ROCKE_SCALAR_BF8E5M2, "bf8e5m2") +ROCKE_SCALAR_SINGLETON(rocke_fp4e2m1, ROCKE_SCALAR_FP4E2M1, "fp4e2m1") +ROCKE_SCALAR_SINGLETON(rocke_fp6e2m3, ROCKE_SCALAR_FP6E2M3, "fp6e2m3") +ROCKE_SCALAR_SINGLETON(rocke_fp6e3m2, ROCKE_SCALAR_FP6E3M2, "fp6e3m2") +ROCKE_SCALAR_SINGLETON(rocke_e8m0, ROCKE_SCALAR_E8M0, "e8m0") +ROCKE_SCALAR_SINGLETON(rocke_e4m3, ROCKE_SCALAR_E4M3, "e4m3") +ROCKE_SCALAR_SINGLETON(rocke_e5m3, ROCKE_SCALAR_E5M3, "e5m3") + #undef ROCKE_SCALAR_SINGLETON const rocke_type_t* rocke_scalar_by_name(const char* name) @@ -81,9 +89,33 @@ const rocke_type_t* rocke_scalar_by_name(const char* name) return rocke_fp8e4m3(); if(strcmp(name, "bf8e5m2") == 0) return rocke_bf8e5m2(); + if(strcmp(name, "fp4e2m1") == 0) + return rocke_fp4e2m1(); + if(strcmp(name, "fp6e2m3") == 0) + return rocke_fp6e2m3(); + if(strcmp(name, "fp6e3m2") == 0) + return rocke_fp6e3m2(); + if(strcmp(name, "e8m0") == 0) + return rocke_e8m0(); + if(strcmp(name, "e4m3") == 0) + return rocke_e4m3(); + if(strcmp(name, "e5m3") == 0) + return rocke_e5m3(); return NULL; } +const rocke_type_t* rocke_dtype_to_ir_type(const char* dtype) +{ + const rocke_dtype_info_t* info = rocke_dtype_info(dtype); + if(!info) + return NULL; + if(strcmp(info->name, "fp16") == 0) + return rocke_f16(); + if(strcmp(info->name, "fp32") == 0) + return rocke_f32(); + return rocke_scalar_by_name(info->name); +} + /* VectorType(elem, count) -> "vec<{elem}x{count}>" */ const rocke_type_t* rocke_vector_type(rocke_ir_builder_t* b, const rocke_type_t* elem, int count) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp index 02cb5a9176cd..039b225b394e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp @@ -895,6 +895,60 @@ rocke_hip_arch_t rocke_hip_arch_from_gfx(const char* gfx) /* ============================== public entry ======================== */ +/* Mirrors _extra_vector_declarations: preserve the static prologue and add + * only widths actually encountered, in parameter/operand/result walk order. */ +static void + h_extra_vector_type(ckc::rocke_h_lowerer_t* lw, const rocke_type_t* t, rocke_strbuf_t* out) +{ + if(!t) + return; + if(t->kind == ROCKE_TYPE_PTR) + h_extra_vector_type(lw, t->pointee, out); + else if(t->kind == ROCKE_TYPE_SMEM) + h_extra_vector_type(lw, t->elem, out); + else if(t->kind == ROCKE_TYPE_VECTOR) + { + const char* name = ckc::rocke_h_type_to_hip(lw, t); + const char* prefix + = strcmp(t->elem->name, "i1") == 0 + ? "boolx" + : ckc::rocke_h_vec_prefix_checked(lw, t->elem->name, true, "vector type"); + const char* scalar + = strcmp(prefix, "i8x") == 0 ? "int8_t" : ckc::rocke_h_hip_scalar(t->elem->name); + const char* macro + = rocke_arena_printf(&lw->b->arena, "_ROCKE_VEC(%s, %s, %d)", scalar, prefix, t->count); + const char* declaration + = rocke_arena_printf(&lw->b->arena, + "using %s = %s __attribute__((ext_vector_type(%d)));", + name, + scalar, + t->count); + if(!macro || !declaration) + ckc::rocke_h_fail(lw, ROCKE_ERR_OOM, "vector declaration allocation failed"); + if(!strstr(ROCKE_HIP_PROLOGUE, macro) && !strstr(rocke_strbuf_cstr(out), declaration)) + { + rocke_strbuf_append(out, declaration); + rocke_strbuf_append_char(out, '\n'); + } + } +} + +static void h_extra_vector_region(ckc::rocke_h_lowerer_t* lw, + const rocke_region_t* region, + rocke_strbuf_t* out) +{ + for(int i = 0; i < region->num_ops; ++i) + { + const auto* op = region->ops[i]; + for(int j = 0; j < op->num_operands; ++j) + h_extra_vector_type(lw, op->operands[j]->type, out); + for(int j = 0; j < op->num_results; ++j) + h_extra_vector_type(lw, op->results[j]->type, out); + for(int j = 0; j < op->num_regions; ++j) + h_extra_vector_region(lw, op->regions[j], out); + } +} + rocke_status_t rocke_lower_kernel_to_hip(rocke_ir_builder_t* b, const rocke_kernel_def_t* kernel, const rocke_lower_hip_opts_t* opts, @@ -993,6 +1047,9 @@ rocke_status_t rocke_lower_kernel_to_hip(rocke_ir_builder_t* b, { rocke_strbuf_append(out, ROCKE_HIP_PROLOGUE); rocke_strbuf_append_char(out, '\n'); + for(int j = 0; j < kernel->num_params; ++j) + h_extra_vector_type(&lw, kernel->params[j]->type, out); + h_extra_vector_region(&lw, kernel->body, out); } /* head */ rocke_strbuf_appendf(out, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/storage.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/storage.cpp new file mode 100644 index 000000000000..b1a2dda9cb6c --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/storage.cpp @@ -0,0 +1,265 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* Mirrors core/storage.py: little-endian patterns, checked sizes and coordinates. + * Matrix and scale operands share these mechanics without sharing semantics. + */ +#include "rocke/storage.h" + +#include + +static bool valid(const rocke_bit_packing_t* p) +{ + return p && p->element_bits >= 1 && p->element_bits <= p->slot_bits && p->slot_bits <= 64; +} + +static bool add(uint64_t a, uint64_t b, uint64_t* out) +{ + if(!out || a > UINT64_MAX - b) + return false; + *out = a + b; + return true; +} + +static bool mul(uint64_t a, uint64_t b, uint64_t* out) +{ + if(!out || (b && a > UINT64_MAX / b)) + return false; + *out = a * b; + return true; +} + +bool rocke_bit_packing_init(rocke_bit_packing_t* out, int element_bits, int slot_bits) +{ + rocke_bit_packing_t p = {element_bits, slot_bits ? slot_bits : element_bits}; + if(!out || !valid(&p)) + return false; + *out = p; + return true; +} + +bool rocke_bit_packing_group(const rocke_bit_packing_t* p, + int carrier_bits, + uint64_t* values, + uint64_t* carriers) +{ + if(!valid(p) || !values || !carriers + || (carrier_bits != 8 && carrier_bits != 16 && carrier_bits != 32 && carrier_bits != 64)) + return false; + int a = p->slot_bits, b = carrier_bits; + while(b) + { + int next = a % b; + a = b; + b = next; + } + *values = carrier_bits / a; + *carriers = p->slot_bits / a; + return true; +} + +bool rocke_bit_packing_offset(const rocke_bit_packing_t* p, uint64_t index, uint64_t* bits) +{ + return valid(p) && mul(index, p->slot_bits, bits); +} + +bool rocke_bit_packing_bytes(const rocke_bit_packing_t* p, + uint64_t count, + uint64_t bit_offset, + uint64_t* bytes) +{ + uint64_t bits; + if(!bytes || !rocke_bit_packing_offset(p, count, &bits)) + return false; + if(!count) + { + *bytes = 0; + return true; + } + if(!add(bits, bit_offset, &bits)) + return false; + *bytes = bits / 8 + (bits % 8 != 0); + return true; +} + +bool rocke_bit_pack(const rocke_bit_packing_t* p, + const uint64_t* patterns, + size_t count, + uint64_t bit_offset, + uint8_t* dst, + size_t dst_size) +{ + uint64_t bytes; + if(!rocke_bit_packing_bytes(p, count, bit_offset, &bytes) || bytes > dst_size + || (count && (!patterns || !dst))) + return false; + for(size_t i = 0; i < count; ++i) + if(p->element_bits < 64 && patterns[i] >= (uint64_t(1) << p->element_bits)) + return false; + if(bytes) + memset(dst, 0, (size_t)bytes); + for(size_t i = 0; i < count; ++i) + { + const uint64_t pos = bit_offset + i * uint64_t(p->slot_bits); + for(int bit = 0; bit < p->element_bits; ++bit) + { + const uint64_t at = pos + bit; + dst[at / 8] |= uint8_t(((patterns[i] >> bit) & 1) << (at % 8)); + } + } + return true; +} + +bool rocke_bit_unpack(const rocke_bit_packing_t* p, + const uint8_t* src, + size_t src_size, + size_t count, + uint64_t bit_offset, + uint64_t* patterns) +{ + uint64_t bytes; + if(!rocke_bit_packing_bytes(p, count, bit_offset, &bytes) || bytes > src_size + || (count && (!patterns || !src))) + return false; + for(size_t i = 0; i < count; ++i) + { + patterns[i] = 0; + const uint64_t pos = bit_offset + i * uint64_t(p->slot_bits); + for(int bit = 0; bit < p->element_bits; ++bit) + { + const uint64_t at = pos + bit; + patterns[i] |= uint64_t((src[at / 8] >> (at % 8)) & 1) << bit; + } + } + return true; +} + +bool rocke_fragment_packing_init(rocke_fragment_packing_t* out, + const rocke_bit_packing_t* p, + uint64_t count, + int carrier_bits, + uint64_t carrier_count) +{ + uint64_t values, carriers, bits, capacity; + if(!out || !rocke_bit_packing_group(p, carrier_bits, &values, &carriers) + || !rocke_bit_packing_offset(p, count, &bits) || !mul(carrier_count, carrier_bits, &capacity) + || bits > capacity) + return false; + *out = {*p, count, carrier_bits, carrier_count}; + return true; +} + +bool rocke_fragment_pack(const rocke_fragment_packing_t* p, + const uint64_t* patterns, + size_t count, + uint64_t* carriers, + size_t carrier_count) +{ + rocke_fragment_packing_t checked; + if(!p + || !rocke_fragment_packing_init( + &checked, &p->packing, p->count, p->carrier_bits, p->carrier_count) + || count != p->count || carrier_count != p->carrier_count || (count && !patterns) + || (carrier_count && !carriers)) + return false; + for(size_t i = 0; i < count; ++i) + if(p->packing.element_bits < 64 && patterns[i] >= (uint64_t(1) << p->packing.element_bits)) + return false; + for(size_t i = 0; i < carrier_count; ++i) + carriers[i] = 0; + for(size_t i = 0; i < count; ++i) + for(int bit = 0; bit < p->packing.element_bits; ++bit) + { + uint64_t at = i * uint64_t(p->packing.slot_bits) + bit; + carriers[at / p->carrier_bits] |= ((patterns[i] >> bit) & 1) << (at % p->carrier_bits); + } + return true; +} + +bool rocke_tensor_storage_init(rocke_tensor_storage_t* out, + const char* dtype, + uint64_t rows, + uint64_t cols, + uint64_t row_stride_bytes, + int slot_bits, + uint64_t base_bit_offset, + uint64_t alignment_bytes) +{ + rocke_tensor_storage_t p = {}; + p.dtype = rocke_dtype_info(dtype); + uint64_t row_bytes, bytes; + if(!out || !p.dtype || !alignment_bytes || (alignment_bytes & (alignment_bytes - 1)) + || !rocke_bit_packing_init(&p.packing, p.dtype->encoded_bits, slot_bits) + || !rocke_bit_packing_bytes(&p.packing, cols, base_bit_offset, &row_bytes)) + return false; + p.rows = rows; + p.cols = cols; + p.row_stride_bytes = row_stride_bytes == UINT64_MAX ? row_bytes : row_stride_bytes; + p.base_bit_offset = base_bit_offset; + p.alignment_bytes = alignment_bytes; + if(p.row_stride_bytes < row_bytes || !rocke_tensor_storage_bytes(&p, &bytes)) + return false; + *out = p; + return true; +} + +bool rocke_tensor_storage_bytes(const rocke_tensor_storage_t* p, uint64_t* bytes) +{ + if(!p || !bytes || !valid(&p->packing)) + return false; + if(!p->rows || !p->cols) + { + *bytes = 0; + return true; + } + uint64_t row_bytes, prefix; + return rocke_bit_packing_bytes(&p->packing, p->cols, p->base_bit_offset, &row_bytes) + && mul(p->rows - 1, p->row_stride_bytes, &prefix) && add(prefix, row_bytes, bytes); +} + +bool rocke_tensor_storage_address( + const rocke_tensor_storage_t* p, uint64_t row, uint64_t col, uint64_t* byte, int* shift) +{ + if(!p || !byte || !shift || row >= p->rows || col >= p->cols) + return false; + uint64_t bits, prefix; + if(!rocke_bit_packing_offset(&p->packing, col, &bits) || !add(bits, p->base_bit_offset, &bits) + || !mul(row, p->row_stride_bytes, &prefix) || !add(prefix, bits / 8, byte)) + return false; + *shift = bits % 8; + return true; +} + +bool rocke_matrix_fragment_layout_init(rocke_matrix_fragment_layout_t* out, + const rocke_fragment_packing_t* fragment, + int chunk_elements, + int lane_groups, + int lanes_per_group) +{ + rocke_fragment_packing_t checked; + uint64_t bits; + if(!out || !fragment || chunk_elements <= 0 || lane_groups <= 0 || lanes_per_group <= 0 + || fragment->count % chunk_elements + || !rocke_fragment_packing_init(&checked, + &fragment->packing, + fragment->count, + fragment->carrier_bits, + fragment->carrier_count) + || !rocke_bit_packing_offset(&fragment->packing, chunk_elements, &bits) || bits % 8) + return false; + *out = {*fragment, chunk_elements, lane_groups, lanes_per_group}; + return true; +} + +bool rocke_matrix_fragment_coord( + const rocke_matrix_fragment_layout_t* p, int lane, uint64_t slot, uint64_t* row, uint64_t* k) +{ + if(!p || !row || !k || lane < 0 || p->lanes_per_group <= 0 || p->lane_groups <= 0 + || p->chunk_elements <= 0 || uint64_t(lane) >= uint64_t(p->lane_groups) * p->lanes_per_group + || slot >= p->fragment.count) + return false; + uint64_t chunk; + *row = lane % p->lanes_per_group; + return mul(slot / p->chunk_elements, p->lane_groups, &chunk) + && add(chunk, lane / p->lanes_per_group, &chunk) && mul(chunk, p->chunk_elements, &chunk) + && add(chunk, slot % p->chunk_elements, k); +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp new file mode 100644 index 000000000000..ecb70efaefb2 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp @@ -0,0 +1,206 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* Descriptor-driven transport. Mirrors helpers/mma_io.py in emission order. */ +#include "rocke/helper_rocke.helpers.mma_io.h" + +#include +#include + +#include "rocke/error_boundary.hpp" + +const rocke_type_t* rocke_storage_ir_type(const char* dtype) +{ + const rocke_dtype_info_t* info = rocke_dtype_info(dtype); + if(!info) + return NULL; + if(info->encoded_bits % 8 || strcmp(info->name, "e8m0") == 0 || strcmp(info->name, "e4m3") == 0 + || strcmp(info->name, "e5m3") == 0) + return rocke_i8(); + return rocke_dtype_to_ir_type(info->name); +} + +static uint64_t gcd(uint64_t a, uint64_t b) +{ + while(b) + { + uint64_t next = a % b; + a = b; + b = next; + } + return a; +} + +rocke_value_t* rocke_h_load_matrix_fragment(rocke_ir_builder_t* b, + rocke_value_t* ptr, + rocke_value_t* row_base, + rocke_value_t* lane_group, + uint64_t k0, + const rocke_tensor_storage_t* storage, + const rocke_matrix_fragment_layout_t* layout, + const rocke_type_t* carrier_type) +{ + return ckc::guard_builder(b, [&]() -> rocke_value_t* { + if(!rocke_i_live(b)) + return NULL; + if(!storage || !storage->dtype || !layout || !carrier_type || !ptr || !row_base + || !lane_group) + ckc::raise_status(ROCKE_ERR_VALUE, "null matrix fragment argument"); + const auto& packing = layout->fragment; + rocke_matrix_fragment_layout_t checked; + if(!rocke_matrix_fragment_layout_init(&checked, + &packing, + layout->chunk_elements, + layout->lane_groups, + layout->lanes_per_group) + || !packing.count || packing.count > INT_MAX || packing.carrier_count > INT_MAX) + ckc::raise_status(ROCKE_ERR_VALUE, "invalid matrix fragment chunk layout"); + const rocke_type_t* unit_type = rocke_storage_ir_type(storage->dtype->name); + if(!unit_type || ptr->type->kind != ROCKE_TYPE_PTR + || !rocke_type_eq(ptr->type->pointee, unit_type)) + ckc::raise_status(ROCKE_ERR_VALUE, "matrix pointer storage type mismatch"); + const uint64_t unit_bytes = rocke_dtype_info(unit_type->name)->encoded_bits / 8; + if(storage->packing.element_bits != packing.packing.element_bits + || storage->packing.slot_bits != packing.packing.slot_bits || storage->base_bit_offset) + ckc::raise_status(ROCKE_ERR_VALUE, + "matrix fragment requires matching packing and a zero bit origin"); + const uint64_t span = packing.count * layout->lane_groups; + if(k0 > storage->cols || span > storage->cols - k0) + ckc::raise_status(ROCKE_ERR_VALUE, "matrix fragment exceeds the packed row"); + uint64_t origin_bits, chunk_bytes; + if(!rocke_bit_packing_offset(&storage->packing, k0, &origin_bits) + || !rocke_bit_packing_bytes(&packing.packing, layout->chunk_elements, 0, &chunk_bytes) + || origin_bits % (8 * unit_bytes) || chunk_bytes % unit_bytes) + ckc::raise_status(ROCKE_ERR_VALUE, + "matrix fragment is not aligned to pointer storage units"); + const rocke_dtype_info_t* carrier_info = rocke_dtype_info(carrier_type->name); + if(!carrier_info || carrier_info->encoded_bits != packing.carrier_bits) + ckc::raise_status(ROCKE_ERR_VALUE, "matrix carrier type width mismatch"); + const uint64_t payload_bits = packing.count * packing.packing.slot_bits; + if(payload_bits % packing.carrier_bits) + ckc::raise_status(ROCKE_ERR_VALUE, + "matrix fragment payload must occupy whole carriers"); + const int live_carriers = payload_bits / packing.carrier_bits; + const int padding = packing.carrier_count - live_carriers; + if(padding && !rocke_type_eq(carrier_type, rocke_i32())) + ckc::raise_status(ROCKE_ERR_VALUE, + "padded matrix fragments currently require i32 carriers"); + const uint64_t chunk_units = chunk_bytes / unit_bytes; + const uint64_t origin_bytes = origin_bits / 8; + if(origin_bytes / unit_bytes > INT_MAX || chunk_units * layout->lane_groups > INT_MAX) + ckc::raise_status(ROCKE_ERR_VALUE, "matrix fragment offset exceeds i32 range"); + int alignment = gcd(gcd(storage->alignment_bytes, storage->row_stride_bytes), + gcd(chunk_bytes, origin_bytes)); + auto* lane_chunk = rocke_b_mul(b, lane_group, rocke_b_const_i32(b, chunk_units)); + auto* step_base = rocke_b_add(b, row_base, rocke_b_const_i32(b, origin_bytes / unit_bytes)); + /* Collect loads before concatenation to preserve Python SSA numbering. */ + rocke_value_t** chunks = (rocke_value_t**)rocke_arena_alloc( + &b->arena, size_t(packing.count) * sizeof(rocke_value_t*)); + if(!chunks) + ckc::raise_status(ROCKE_ERR_OOM, "matrix fragment allocation failed"); + int num_chunks = 0; + for(uint64_t j = 0; j < packing.count / layout->chunk_elements; ++j) + { + auto* offset = rocke_b_add( + b, + rocke_b_add( + b, step_base, rocke_b_const_i32(b, j * layout->lane_groups * chunk_units)), + lane_chunk); + uint64_t remaining = chunk_units, consumed = 0; + const int max_width = unit_bytes == 4 ? 8 : 16; + while(remaining) + { + int width = 1; + while(width < max_width && uint64_t(width * 2) <= remaining) + width *= 2; + auto* at + = consumed ? rocke_b_add(b, offset, rocke_b_const_i32(b, consumed)) : offset; + int load_align = gcd(alignment, consumed * unit_bytes); + auto* value + = width == 1 ? rocke_b_vector_splat( + b, rocke_b_global_load(b, ptr, at, unit_type, load_align), 1) + : rocke_b_global_load_vN(b, ptr, at, unit_type, width, load_align); + chunks[num_chunks++] = value; + consumed += width; + remaining -= width; + } + } + auto* payload = chunks[0]; + for(int i = 1; i < num_chunks; ++i) + payload = rocke_b_vec_concat(b, payload, chunks[i]); + const auto* payload_type = rocke_vector_type(b, carrier_type, live_carriers); + if(!rocke_type_eq(payload->type, payload_type)) + payload = rocke_b_bitcast(b, payload, payload_type); + if(padding) + payload = rocke_b_vec_concat( + b, payload, rocke_b_vector_splat(b, rocke_b_const_i32(b, 0), padding)); + return payload; + }); +} + +rocke_status_t rocke_h_pack_fragment_bits(rocke_ir_builder_t* b, + rocke_load_bits_fn load_bits, + void* ctx, + const rocke_fragment_packing_t* fragment, + rocke_value_t** words, + size_t word_count) +{ + return ckc::guard_status(b, [&]() -> rocke_status_t { + if(!rocke_i_live(b)) + return ROCKE_ERR_VALUE; + rocke_fragment_packing_t checked; + if(!fragment || !load_bits || !words + || !rocke_fragment_packing_init(&checked, + &fragment->packing, + fragment->count, + fragment->carrier_bits, + fragment->carrier_count) + || word_count != fragment->carrier_count || fragment->count > INT_MAX) + ckc::raise_status(ROCKE_ERR_VALUE, "invalid pattern packing arguments"); + const auto& f = *fragment; + if(f.carrier_bits != 32 && f.carrier_bits != 64) + ckc::raise_status(ROCKE_ERR_VALUE, + "IR pattern packing currently requires i32 or i64 carriers"); + const auto* word_type = f.carrier_bits == 64 ? rocke_i64() : rocke_i32(); + auto constant = [&](uint64_t value) { + return f.carrier_bits == 64 ? rocke_b_const_i64(b, int64_t(value)) + : rocke_b_const_i32(b, int32_t(value)); + }; + for(size_t j = 0; j < word_count; ++j) + words[j] = constant(0); + for(uint64_t j = 0; j < f.count; ++j) + { + auto* pattern = load_bits(b, (int)j, ctx); + if(!pattern + || (!rocke_type_eq(pattern->type, rocke_i8()) + && !rocke_type_eq(pattern->type, rocke_i16()) + && !rocke_type_eq(pattern->type, rocke_i32()) + && !rocke_type_eq(pattern->type, rocke_i64()))) + ckc::raise_status(ROCKE_ERR_VALUE, + "pattern packing requires unsigned patterns in integer carriers"); + const int bits = rocke_dtype_info(pattern->type->name)->encoded_bits; + if(bits < f.packing.element_bits) + ckc::raise_status(ROCKE_ERR_VALUE, "pattern carrier is smaller than encoded width"); + if(bits < f.carrier_bits) + pattern = rocke_b_zext(b, pattern, word_type); + else if(bits > f.carrier_bits) + ckc::raise_status(ROCKE_ERR_VALUE, "pattern carrier is wider than output carrier"); + uint64_t start = j * f.packing.slot_bits; + int remaining = f.packing.element_bits, consumed = 0; + while(remaining) + { + const uint64_t word = start / f.carrier_bits; + const int shift = start % f.carrier_bits; + const int take + = remaining < f.carrier_bits - shift ? remaining : f.carrier_bits - shift; + auto* part = consumed ? rocke_b_lshr(b, pattern, constant(consumed)) : pattern; + if(take < remaining) + part = rocke_b_land(b, part, constant((uint64_t(1) << take) - 1)); + words[word] = rocke_b_lor(b, words[word], rocke_b_shl(b, part, constant(shift))); + start += take; + consumed += take; + remaining -= take; + } + } + return ROCKE_OK; + }); +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/quant.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/quant.cpp index e6e92d1206e2..fcb13a085464 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/quant.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/quant.cpp @@ -1,37 +1,8 @@ // Copyright (c) Advanced Micro Devices, Inc., or its affiliates. // SPDX-License-Identifier: MIT -/* - * helper_rocke.helpers.quant.c -- C99 port of - * rocke.helpers.quant.quant_ir_type. - * - * Faithful translation of: - * - * _QDTYPE_ALIAS = { - * "i8": "i8", "int8": "i8", - * "fp8e4m3": "fp8e4m3", "fp8": "fp8e4m3", "fp8_e4m3": "fp8e4m3", - * "bf8e5m2": "bf8e5m2", "bf8": "bf8e5m2", "fp8_e5m2": "bf8e5m2", - * } - * - * def _canon(qdtype: str) -> QDType: - * if qdtype not in _QDTYPE_ALIAS: - * raise ValueError( - * f"unsupported quant dtype {qdtype!r}; expected one of " - * f"{sorted(_QDTYPE_ALIAS)}") - * return _QDTYPE_ALIAS[qdtype] - * - * def quant_ir_type(qdtype: str) -> Type: - * canon = _canon(qdtype) - * if canon == "i8": return I8 - * if canon == "fp8e4m3": return FP8E4M3 - * if canon == "bf8e5m2": return BF8E5M2 - * raise ValueError(f"unreachable: canon={canon!r}") - * - * Mapping invariants (must stay byte-identical to the Python so downstream IR is - * identical): - * "i8" / "int8" -> rocke_i8() (Python I8) - * "fp8e4m3" / "fp8" / "fp8_e4m3" -> rocke_fp8e4m3() (Python FP8E4M3) - * "bf8e5m2" / "bf8" / "fp8_e5m2" -> rocke_bf8e5m2() (Python BF8E5M2) - * else -> NULL (Python ValueError) +/* Mirrors helpers/quant.py::quant_ir_type. Preserve the existing quant aliases + * and delegate logical type resolution to rocke_dtype_to_ir_type. FP4/FP6 + * recognition is independent of scalar quantization/conversion support. */ #include "rocke/helper_rocke.helpers.quant.h" @@ -49,20 +20,24 @@ const rocke_type_t* rocke_quant_ir_type(const char* qdtype) /* `_canon` -> "i8" : keys "i8", "int8". */ if(strcmp(qdtype, "i8") == 0 || strcmp(qdtype, "int8") == 0) { - return rocke_i8(); + return rocke_dtype_to_ir_type("i8"); } /* `_canon` -> "fp8e4m3" : keys "fp8e4m3", "fp8", "fp8_e4m3". */ if(strcmp(qdtype, "fp8e4m3") == 0 || strcmp(qdtype, "fp8") == 0 || strcmp(qdtype, "fp8_e4m3") == 0) { - return rocke_fp8e4m3(); + return rocke_dtype_to_ir_type("fp8e4m3"); } /* `_canon` -> "bf8e5m2" : keys "bf8e5m2", "bf8", "fp8_e5m2". */ if(strcmp(qdtype, "bf8e5m2") == 0 || strcmp(qdtype, "bf8") == 0 || strcmp(qdtype, "fp8_e5m2") == 0) { - return rocke_bf8e5m2(); + return rocke_dtype_to_ir_type("bf8e5m2"); } + if(strcmp(qdtype, "fp4") == 0 || strcmp(qdtype, "fp4e2m1") == 0 || strcmp(qdtype, "fp6") == 0 + || strcmp(qdtype, "fp6e2m3") == 0 || strcmp(qdtype, "bf6") == 0 + || strcmp(qdtype, "fp6e3m2") == 0) + return rocke_dtype_to_ir_type(qdtype); /* Python: _canon raises ValueError. No builder here, so signal via NULL. */ return NULL; } 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 index f7da21efad31..6d814029b952 100644 --- 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 @@ -28,6 +28,24 @@ extern "C" { * 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); +/* Borrowed immutable encoding descriptor, valid for the process lifetime. */ +typedef enum rocke_dtype_category +{ + ROCKE_DTYPE_FLOAT, + ROCKE_DTYPE_INTEGER, + ROCKE_DTYPE_PREDICATE +} rocke_dtype_category_t; + +typedef struct rocke_dtype_info +{ + const char* name; + rocke_dtype_category_t category; + int encoded_bits; +} rocke_dtype_info_t; + +/* NULL for unknown encodings. No architecture or operation support implied. */ +const rocke_dtype_info_t* rocke_dtype_info(const char* name); + #ifdef __cplusplus } /* extern "C" */ #endif diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.mma_io.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.mma_io.h new file mode 100644 index 000000000000..8f5db63aa960 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.mma_io.h @@ -0,0 +1,41 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* Mirrors helpers/mma_io.py; descriptors expand to ordinary IR before serialization. + * storage_ir_type -> rocke_storage_ir_type + * load_matrix_fragment -> rocke_h_load_matrix_fragment + * pack_fragment_bits -> rocke_h_pack_fragment_bits + */ +#ifndef ROCKE_HELPER_MMA_IO_H +#define ROCKE_HELPER_MMA_IO_H + +#include "rocke/ir.h" +#include "rocke/storage.h" + +#ifdef __cplusplus +extern "C" { +#endif + +const rocke_type_t* rocke_storage_ir_type(const char* dtype); +rocke_value_t* rocke_h_load_matrix_fragment(rocke_ir_builder_t* b, + rocke_value_t* ptr, + rocke_value_t* row_base, + rocke_value_t* lane_group, + uint64_t k0, + const rocke_tensor_storage_t* storage, + const rocke_matrix_fragment_layout_t* layout, + const rocke_type_t* carrier_type); + +/* Callback returns a canonical unsigned pattern in an integer carrier. + * ctx and output pointer array are borrowed for the duration of this call. */ +typedef rocke_value_t* (*rocke_load_bits_fn)(rocke_ir_builder_t* b, int index, void* ctx); +rocke_status_t rocke_h_pack_fragment_bits(rocke_ir_builder_t* b, + rocke_load_bits_fn load_bits, + void* ctx, + const rocke_fragment_packing_t* fragment, + rocke_value_t** words, + size_t word_count); + +#ifdef __cplusplus +} +#endif +#endif /* ROCKE_HELPER_MMA_IO_H */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.quant.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.quant.h index 540e797219ea..2b24e6d2e172 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.quant.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.quant.h @@ -3,42 +3,11 @@ * * rocke/helper_rocke.helpers.quant.h -- C99 port of rocke.helpers.quant. * - * Scope of THIS file: only `quant_ir_type` is ported here (the other quant - * helpers in helpers/quant.py -- quantize_scalar_f32, dequantize_scalar_to_f32, - * pack_quant_chunk_local_f32, store_packed_chunk_local, quant_max_abs, - * ir_to_qdtype -- are out of scope for this phase). The function maps a - * quant-dtype alias spelling string to the canonical IR scalar type, - * byte-identically to the Python: - * - * def quant_ir_type(qdtype: str) -> Type: - * canon = _canon(qdtype) - * if canon == "i8": return I8 - * if canon == "fp8e4m3": return FP8E4M3 - * if canon == "bf8e5m2": return BF8E5M2 - * raise ValueError(...) - * - * with `_canon` normalising the accepted aliases: - * - * "i8" / "int8" -> "i8" - * "fp8e4m3" / "fp8" / "fp8_e4m3" -> "fp8e4m3" - * "bf8e5m2" / "bf8" / "fp8_e5m2" -> "bf8e5m2" - * - * and raising ValueError for any other spelling. - * - * The Python helper takes no IRBuilder and raises ValueError on an unsupported - * dtype. C99 has no exceptions, so we expose two faithful spellings, mirroring - * the sibling helper_rocke.helpers.io port: - * - * rocke_quant_ir_type(qdtype) - * Pure map. Returns the interned scalar singleton - * (rocke_i8()/rocke_fp8e4m3()/rocke_bf8e5m2()) for any accepted alias; returns - * NULL for anything else (the analog of "raise ValueError"). No error - * state is set because there is no builder. - * - * rocke_b_quant_ir_type(b, qdtype) - * Builder-aware spelling matching the rest of the C port's sticky-error - * model: on an unsupported dtype it records ROCKE_ERR_VALUE + a message on - * the builder (mirroring the Python ValueError text) and returns NULL. + * Type resolution only; scalar quantization and conversion remain separate. + * rocke_quant_ir_type delegates to rocke_dtype_to_ir_type for nominal types, + * preserving the existing quant aliases and adding FP4/FP6 recognition. + * Returns a borrowed static singleton or NULL for an unsupported name. + * The builder-aware variant also records the corresponding error. */ #ifndef ROCKE_HELPER_ROCKE_HELPERS_QUANT_H #define ROCKE_HELPER_ROCKE_HELPERS_QUANT_H @@ -55,6 +24,9 @@ extern "C" { * "i8", "int8" -> rocke_i8() * "fp8e4m3", "fp8", "fp8_e4m3" -> rocke_fp8e4m3() * "bf8e5m2", "bf8", "fp8_e5m2" -> rocke_bf8e5m2() + * "fp4", "fp4e2m1" -> rocke_fp4e2m1() + * "fp6", "fp6e2m3" -> rocke_fp6e2m3() + * "bf6", "fp6e3m2" -> rocke_fp6e3m2() * Returns NULL for any other value (the Python ValueError path). */ const rocke_type_t* rocke_quant_ir_type(const char* qdtype); 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 59ae31feced3..131ea637e341 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 @@ -105,6 +105,12 @@ typedef enum rocke_scalar_kind ROCKE_SCALAR_F32, ROCKE_SCALAR_FP8E4M3, ROCKE_SCALAR_BF8E5M2, + ROCKE_SCALAR_FP4E2M1, + ROCKE_SCALAR_FP6E2M3, + ROCKE_SCALAR_FP6E3M2, + ROCKE_SCALAR_E8M0, + ROCKE_SCALAR_E4M3, + ROCKE_SCALAR_E5M3, ROCKE_SCALAR__COUNT } rocke_scalar_kind_t; @@ -574,6 +580,15 @@ const rocke_type_t* rocke_f16(void); const rocke_type_t* rocke_f32(void); const rocke_type_t* rocke_fp8e4m3(void); const rocke_type_t* rocke_bf8e5m2(void); +const rocke_type_t* rocke_fp4e2m1(void); +const rocke_type_t* rocke_fp6e2m3(void); +const rocke_type_t* rocke_fp6e3m2(void); +const rocke_type_t* rocke_e8m0(void); +const rocke_type_t* rocke_e4m3(void); +const rocke_type_t* rocke_e5m3(void); + +/* Logical dtype resolver; NULL for unknown or unrepresented encodings. */ +const rocke_type_t* rocke_dtype_to_ir_type(const char* dtype); /* Look up a scalar singleton by canonical name ("i32",...); NULL if unknown. */ const rocke_type_t* rocke_scalar_by_name(const char* name); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/storage.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/storage.h new file mode 100644 index 000000000000..3ea1794d6369 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/storage.h @@ -0,0 +1,114 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* Shared bit storage; mirrors core/storage.py. + * BitPacking -> rocke_bit_packing_t + * FragmentPacking -> rocke_fragment_packing_t + * TensorStorage -> rocke_tensor_storage_t + * All descriptors are caller-owned. Functions return false for invalid inputs, + * overflow or insufficient buffers; no exceptions cross this C API. + */ +#ifndef ROCKE_STORAGE_H +#define ROCKE_STORAGE_H + +#include +#include +#include + +#include "rocke/dtypes.h" + +#ifdef __cplusplus +extern "C" { +#endif + +typedef struct rocke_bit_packing +{ + int element_bits; + int slot_bits; /* 0 selects dense packing at construction. */ +} rocke_bit_packing_t; + +typedef struct rocke_fragment_packing +{ + rocke_bit_packing_t packing; + uint64_t count; + int carrier_bits; + uint64_t carrier_count; +} rocke_fragment_packing_t; + +typedef struct rocke_matrix_fragment_layout +{ + rocke_fragment_packing_t fragment; + int chunk_elements; + int lane_groups; + int lanes_per_group; +} rocke_matrix_fragment_layout_t; + +bool rocke_matrix_fragment_layout_init(rocke_matrix_fragment_layout_t* out, + const rocke_fragment_packing_t* fragment, + int chunk_elements, + int lane_groups, + int lanes_per_group); +bool rocke_matrix_fragment_coord( + const rocke_matrix_fragment_layout_t* p, int lane, uint64_t slot, uint64_t* row, uint64_t* k); + +typedef struct rocke_tensor_storage +{ + const rocke_dtype_info_t* dtype; /* Borrowed static descriptor. */ + uint64_t rows; + uint64_t cols; + uint64_t row_stride_bytes; + uint64_t base_bit_offset; + uint64_t alignment_bytes; + rocke_bit_packing_t packing; +} rocke_tensor_storage_t; + +bool rocke_bit_packing_init(rocke_bit_packing_t* out, int element_bits, int slot_bits); +bool rocke_bit_packing_group(const rocke_bit_packing_t* p, + int carrier_bits, + uint64_t* values, + uint64_t* carriers); +bool rocke_bit_packing_offset(const rocke_bit_packing_t* p, uint64_t index, uint64_t* bits); +bool rocke_bit_packing_bytes(const rocke_bit_packing_t* p, + uint64_t count, + uint64_t bit_offset, + uint64_t* bytes); +/* Pack writes a fresh zero-padded stream; caller must exclusively own dst. + * Unpack reads only the allocated bytes, including a partial final byte. */ +bool rocke_bit_pack(const rocke_bit_packing_t* p, + const uint64_t* patterns, + size_t count, + uint64_t bit_offset, + uint8_t* dst, + size_t dst_size); +bool rocke_bit_unpack(const rocke_bit_packing_t* p, + const uint8_t* src, + size_t src_size, + size_t count, + uint64_t bit_offset, + uint64_t* patterns); +bool rocke_fragment_packing_init(rocke_fragment_packing_t* out, + const rocke_bit_packing_t* p, + uint64_t count, + int carrier_bits, + uint64_t carrier_count); +bool rocke_fragment_pack(const rocke_fragment_packing_t* p, + const uint64_t* patterns, + size_t count, + uint64_t* carriers, + size_t carrier_count); +/* row_stride_bytes=UINT64_MAX selects a minimal packed row stride. */ +bool rocke_tensor_storage_init(rocke_tensor_storage_t* out, + const char* dtype, + uint64_t rows, + uint64_t cols, + uint64_t row_stride_bytes, + int slot_bits, + uint64_t base_bit_offset, + uint64_t alignment_bytes); +bool rocke_tensor_storage_bytes(const rocke_tensor_storage_t* p, uint64_t* bytes); +bool rocke_tensor_storage_address( + const rocke_tensor_storage_t* p, uint64_t row, uint64_t col, uint64_t* byte, int* shift); + +#ifdef __cplusplus +} +#endif +#endif /* ROCKE_STORAGE_H */ 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 index c85df088e922..08585dd6b734 100644 --- 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 @@ -14,6 +14,7 @@ #include "rocke/arch_target.h" #include "rocke/error.hpp" #include "rocke/ir.h" +#include "rocke/storage.h" typedef struct rocke_scale_packing { @@ -23,6 +24,52 @@ typedef struct rocke_scale_packing int block_k; } rocke_scale_packing_t; +typedef struct rocke_scale_association +{ + int block_k; +} rocke_scale_association_t; + +static inline rocke_scale_association_t rocke_scale_association(const rocke_scale_packing_t* p) +{ + if(p->block_k <= 0) + ckc::raise_status(ROCKE_ERR_VALUE, "scale block_k must be positive"); + return {p->block_k}; +} + +static inline rocke_fragment_packing_t rocke_scale_fragment(const rocke_scale_packing_t* p) +{ + rocke_bit_packing_t bits; + rocke_fragment_packing_t result; + rocke_scale_association(p); + if(p->count != 1 && p->count != 2 && p->count != 4 && p->count != 8) + ckc::raise_status(ROCKE_ERR_VALUE, "carrier_bits must be 8, 16, 32, or 64"); + if(!rocke_bit_packing_init(&bits, rocke_dtype_info("e8m0")->encoded_bits, 0) + || !rocke_fragment_packing_init(&result, &bits, p->count, p->count * 8, 1)) + ckc::raise_status(ROCKE_ERR_VALUE, "invalid scale packing"); + return result; +} + +static inline rocke_matrix_fragment_layout_t rocke_scaled_matrix_layout(const char* dtype, + int abi_words) +{ + const rocke_dtype_info_t* info = rocke_dtype_info(dtype); + int chunk = 0; + if(info && (strcmp(info->name, "fp8e4m3") == 0 || strcmp(info->name, "bf8e5m2") == 0)) + chunk = 16; + else if(info + && (strcmp(info->name, "fp4e2m1") == 0 || strcmp(info->name, "fp6e2m3") == 0 + || strcmp(info->name, "fp6e3m2") == 0)) + chunk = 32; + rocke_bit_packing_t bits; + rocke_fragment_packing_t fragment; + rocke_matrix_fragment_layout_t result; + if(!chunk || abi_words < 0 || !rocke_bit_packing_init(&bits, info->encoded_bits, 0) + || !rocke_fragment_packing_init(&fragment, &bits, 64, 32, abi_words) + || !rocke_matrix_fragment_layout_init(&result, &fragment, chunk, 2, 16)) + ckc::raise_status(ROCKE_ERR_VALUE, "unsupported scaled matrix layout"); + return result; +} + typedef struct rocke_scaled_wmma_op { const char* op_id; @@ -45,6 +92,13 @@ static inline const rocke_mma_op_t* rocke_gfx1250_scaled_wmma(const char* op_id) return atom && strcmp(atom->family, "wmma_scaled") == 0 ? atom : NULL; } +static inline rocke_matrix_fragment_layout_t + rocke_scaled_wmma_matrix_layout(const rocke_mma_op_t* atom, bool for_b) +{ + return rocke_scaled_matrix_layout(for_b ? atom->b_dtype : atom->a_dtype, + for_b ? atom->b_frag_len : atom->a_frag_len); +} + static inline rocke_scaled_wmma_op_t rocke_scaled_wmma_contract(const rocke_mma_op_t* atom) { rocke_scaled_wmma_op_t spec = {}; @@ -104,7 +158,7 @@ static inline const rocke_mma_op_t* rocke_gfx1250_scaled_wmma_from_op(const rock static inline int rocke_scale_word_bits(const rocke_scale_packing_t* packing) { - return packing->count * 8; + return rocke_scale_fragment(packing).carrier_bits; } #endif /* ROCKE_WMMA_SCALE_INTERNAL_H */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md new file mode 100644 index 000000000000..37d8d51b91f0 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md @@ -0,0 +1,124 @@ +# Logical dtypes, storage, and operand packing + +Matrix inputs and scales use three separate contracts: + +1. A logical dtype identifies the numerical encoding and its bit width. +2. Tensor storage describes how elements occupy an addressable buffer. +3. An instruction fragment describes lane mapping, carrier capacity, and padding. + +The selected atom determines supported operand combinations and instruction +selectors. An integer carrier does not identify the numerical format. + +## Contents + +- [Logical dtypes](#logical-dtypes) +- [Common packing](#common-packing) +- [Tensor storage](#tensor-storage) +- [Matrix and scale layouts](#matrix-and-scale-layouts) +- [Mirroring and follow-up work](#mirroring-and-follow-up-work) + +## Logical dtypes + +[`dtype_info`](../../python/rocke/core/dtypes.py) and +[`rocke_dtype_info`](../../cpp/include/rocke/dtypes.h) return target-independent +encoding descriptors. `dtype_to_ir_type` and `rocke_dtype_to_ir_type` resolve +logical types, including distinct FP4 E2M1, FP6 E2M3, FP6 E3M2, and scale types. +`quant_ir_type` delegates type resolution to this common resolver. + +Recognition and serialization do not enable scalar conversion, arithmetic, or +an instruction on a target. In this implementation, FP4/FP6 and the nominal scale +types are used in authoring metadata; storage helpers expand their transport into +existing integer operations before serialization. Raw scalar lowering of these +types remains unsupported. Existing FP16/BF16 and FP8/BF8 lowering is unchanged. + +Scale-format validation remains separate from matrix alias normalization: +E5M3 is not BF8 E5M2. Registering a scale dtype does not enable a backend selector. + +## Common packing + +`BitPacking` / `rocke_bit_packing_t` describes unsigned bit patterns with an +encoded width and slot width. Dense packing uses equal widths. A wider slot has +unused high bits, which packers zero. The current encoding is least-significant +bit first in a little-endian byte stream. + +| Format | Dense group occupying whole 32-bit words | +|---|---| +| FP16/BF16 | 2 elements in 1 word | +| FP8/BF8 | 4 elements in 1 word | +| FP6/BF6 | 16 elements across 3 words | +| FP4 | 8 elements in 1 word | + +FP6 elements can cross byte and word boundaries. Packing cannot be expressed +generally as `32 // element_bits` elements per word. + +`FragmentPacking` / `rocke_fragment_packing_t` combines bit packing with a logical +count and a carrier width/count. Capacity is checked independently of payload +size; unused carriers are zero. Carriers may hold integer bits or typed values, +depending on the instruction ABI. + +Host pack/unpack routines process bit patterns only. They do not round, clamp, +decode floats, or apply scales. Their tests use literal expected bits and +independent integer arithmetic; numerical tests remain separate. + +## Tensor storage + +`TensorStorage` / `rocke_tensor_storage_t` initially supports a two-dimensional +logical shape with a contiguous packed inner axis, an explicit byte row stride, +bit offset, and base alignment. Nonnegative sizes and offsets are checked against +the native uint64 range. A/B descriptors are independent. + +For example, dense FP6 uses four elements in three bytes; `slot_bits=8` describes +one six-bit pattern in each byte instead. Both have the same logical dtype. +Address calculation includes the bit offset, and allocation bounds include the +partial tail byte. Empty views require no accesses. + +[`storage_ir_type`](../../python/rocke/helpers/mma_io.py) chooses an addressable unit. FP4/FP6 use I8 +units while retaining their logical identities in the storage descriptor. +This function is deliberately separate from `dtype_to_ir_type`. + +The fragment loader accepts complete, byte-aligned chunks. It derives load +alignment from base alignment, row stride, chunk spacing, and origin. A 24-byte +FP6 chunk becomes 16-byte and 8-byte loads without reading past the chunk. +Partial fragments and nonzero bit origins are rejected by this initial loader. +The host bit packer writes a fresh, exclusively owned output buffer; concurrent +packed stores are not provided. + +## Matrix and scale layouts + +`MatrixFragmentLayout` describes contiguous K chunks interleaved between lane +groups. Coordinate mapping belongs to the atom-specific layout, not to the +generic bit packer. `ScaledWmmaOp.matrix_layout(operand)` selects independent +A/B layouts. Existing `a_frag_len`/`b_frag_len` retain their ABI-vector meaning. + +The current gfx1250 scaled matrix layouts have 64 elements per lane and sixteen +i32 carriers. FP8 occupies all sixteen words. FP4 occupies eight and pads eight. +The FP6 transport descriptor occupies twelve and pads four, without enabling +FP6 catalog entries or numerical conversions. + +`ScaleAssociation` records `block_k`: how many source K elements share one +scale. This is separate from the bit layout. Existing `ScalePacking(count, +block_k)` remains a compatibility adapter, exposing both the association and a +shared `FragmentPacking`. + +Current scaled-WMMA scale fragments hold four eight-bit patterns in i32 for K32, +or eight in i64 for K16. The first K group occupies the low byte. The A/B scale +coordinate maps determine which scales each lane loads. The same +`pack_fragment_bits` helper can pack six-bit fields that cross word boundaries. + +## Mirroring and follow-up work + +The Python and C++ helpers expand descriptors before serialization. Their emitted +IR and HIP source are tested byte for byte, including FP16/BF16 typed carriers, +FP4/FP6/FP8 matrix payloads, and scale word packing. Existing GEMM signatures +remain compatible. HIP emits declarations for any encountered vector widths +absent from its fixed compatibility prologue. + +The shared descriptor/helper changes can be consumed by both FP4 and FP6 +branches. FP4/FP8 builder migration can use them directly. Future FP6 numerical +integration should consume the six-bit storage and fragment descriptors and +retain separate target and conversion validation. + +First-class tensor-view/fragment IR nodes, arbitrary packed axes, masked partial +tiles, concurrent packed stores, and scalar low-bit conversions are separate +extensions. Any descriptor retained in future IR nodes must serialize, validate, +participate in cache keys, and have a native implementation. 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 bd0f96fc8493..57596e6ba255 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 @@ -45,12 +45,19 @@ apply_codegen_policy, codegen_policy_for_kernel, ) -from .dtypes import normalize_dtype +from .dtypes import DTypeCategory, DTypeInfo, dtype_info, normalize_dtype from .ir import ( BF16, F16, F32, FP8E4M3, + FP4E2M1, + FP6E2M3, + FP6E3M2, + E8M0, + E4M3, + E5M3, + dtype_to_ir_type, I1, I8, I32, @@ -96,6 +103,16 @@ __all__ = [ "normalize_dtype", + "DTypeCategory", + "DTypeInfo", + "dtype_info", + "dtype_to_ir_type", + "FP4E2M1", + "FP6E2M3", + "FP6E3M2", + "E8M0", + "E4M3", + "E5M3", "CodegenPolicy", "SchedulerStrategy", "apply_codegen_policy", 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 index 2febb7514b05..7f0284ed983c 100644 --- 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 @@ -7,6 +7,36 @@ from dataclasses import dataclass from .target import ArchTarget, MmaOp, MmaScaleDType +from ..dtypes import dtype_info +from ..storage import BitPacking, FragmentPacking, MatrixFragmentLayout + + +@dataclass(frozen=True) +class ScaleAssociation: + """Number of source K elements sharing one scale.""" + + block_k: int + + def __post_init__(self) -> None: + if self.block_k <= 0: + raise ValueError("scale block_k must be positive") + + +def scaled_matrix_layout(dtype: str, abi_words: int) -> MatrixFragmentLayout: + """gfx1250 scaled operand layout, independent of atom availability. + + FP6 describes the transport contract for future catalog integration. + """ + info = dtype_info(dtype) + chunks = {"fp8e4m3": 16, "bf8e5m2": 16, "fp6e2m3": 32, "fp6e3m2": 32, "fp4e2m1": 32} + if info.name not in chunks: + raise ValueError(f"unsupported scaled matrix layout {dtype!r}") + return MatrixFragmentLayout( + FragmentPacking(BitPacking(info.encoded_bits), 64, 32, abi_words), + chunks[info.name], + 2, + 16, + ) @dataclass(frozen=True) @@ -20,13 +50,31 @@ class ScalePacking: count: int block_k: int + def __post_init__(self) -> None: + self.association + self.fragment + + @property + def association(self) -> ScaleAssociation: + return ScaleAssociation(self.block_k) + + @property + def packing(self) -> BitPacking: + return BitPacking(dtype_info("e8m0").encoded_bits) + + @property + def fragment(self) -> FragmentPacking: + return FragmentPacking( + self.packing, self.count, self.count * self.packing.element_bits, 1 + ) + @property def element_bits(self) -> int: - return 8 + return self.packing.element_bits @property def word_bits(self) -> int: - return self.count * self.element_bits + return self.fragment.carrier_bits @property def llvm_type(self) -> str: @@ -42,6 +90,22 @@ class ScaledWmmaOp: scale_formats: tuple[int, int] scales: ScalePacking + def matrix_layout(self, operand: str) -> MatrixFragmentLayout: + if operand == "a": + return scaled_matrix_layout(self.atom.a_dtype, self.atom.a_frag_len) + if operand == "b": + return scaled_matrix_layout(self.atom.b_dtype, self.atom.b_frag_len) + raise ValueError("matrix operand must be 'a' or 'b'") + + def scale_packing(self, operand: str) -> ScalePacking: + if operand == "a": + count = self.atom.a_scale_frag_len + elif operand == "b": + count = self.atom.b_scale_frag_len + else: + raise ValueError("scale operand must be 'a' or 'b'") + return ScalePacking(count, self.atom.scale_block_k) + @property def op_id(self) -> str: return self.atom.op_id diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/dtypes.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/dtypes.py index 583b591b1390..3808f87f3ad0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/dtypes.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/dtypes.py @@ -9,6 +9,9 @@ from __future__ import annotations +from dataclasses import dataclass +from enum import Enum + # 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 = { @@ -48,3 +51,67 @@ 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) + + +class DTypeCategory(str, Enum): + FLOAT = "float" + INTEGER = "integer" + PREDICATE = "predicate" + + +@dataclass(frozen=True) +class DTypeInfo: + """Logical encoding, independent of storage slots and instruction support.""" + + name: str + category: DTypeCategory + encoded_bits: int + + +_DTYPE_INFO = { + name: DTypeInfo(name, category, bits) + for category, entries in ( + ( + DTypeCategory.FLOAT, + ( + ("fp32", 32), + ("fp16", 16), + ("bf16", 16), + ("fp8e4m3", 8), + ("bf8e5m2", 8), + ("fp6e2m3", 6), + ("fp6e3m2", 6), + ("fp4e2m1", 4), + ("e8m0", 8), + ("e4m3", 8), + ("e5m3", 8), + ), + ), + ( + DTypeCategory.INTEGER, + ( + ("i4", 4), + ("iu4", 4), + ("i8", 8), + ("iu8", 8), + ("i16", 16), + ("i32", 32), + ("i64", 64), + ), + ), + (DTypeCategory.PREDICATE, (("i1", 1),)), + ) + for name, bits in entries +} + + +def dtype_info(name: str) -> DTypeInfo: + """Resolve an encoding; recognition alone does not enable an operation. + + Scale names remain distinct keys. In particular e5m3 is not bf8e5m2. + Operand-role validation belongs to the selected atom. + """ + key = normalize_dtype(name) + if key not in _DTYPE_INFO: + raise ValueError(f"unknown dtype {name!r}") + return _DTYPE_INFO[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 cab47560124e..6dc0175c467d 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 @@ -30,6 +30,7 @@ from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple from .arch import target as _arch +from .dtypes import dtype_info # ----------------------------- Types -------------------------------------- @@ -52,6 +53,41 @@ def __repr__(self) -> str: F32 = Type("f32") FP8E4M3 = Type("fp8e4m3") BF8E5M2 = Type("bf8e5m2") +FP4E2M1 = Type("fp4e2m1") +FP6E2M3 = Type("fp6e2m3") +FP6E3M2 = Type("fp6e3m2") +E8M0 = Type("e8m0") +E4M3 = Type("e4m3") +E5M3 = Type("e5m3") + + +def dtype_to_ir_type(dtype: str) -> Type: + """Resolve a logical type without claiming scalar operation support. + + Packed memory is described separately; a low-bit type is never an I8 alias. + """ + info = dtype_info(dtype) + types = { + "i1": I1, + "i8": I8, + "i16": I16, + "i32": I32, + "i64": I64, + "fp16": F16, + "bf16": BF16, + "fp32": F32, + "fp8e4m3": FP8E4M3, + "bf8e5m2": BF8E5M2, + "fp4e2m1": FP4E2M1, + "fp6e2m3": FP6E2M3, + "fp6e3m2": FP6E3M2, + "e8m0": E8M0, + "e4m3": E4M3, + "e5m3": E5M3, + } + if info.name not in types: + raise ValueError(f"no logical IR type for dtype {info.name!r}") + return types[info.name] # AMDGPU buffer-load AUX-byte cache-coherency hints. The AUX field of 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 f0628a9fdee8..22e38bca4581 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 @@ -24,6 +24,7 @@ PtrType, Region, SmemType, + Type, Value, VectorType, ) @@ -2440,6 +2441,41 @@ def _find_enclosing_for(region: Region, target: Op) -> Optional[Op]: return None +def _extra_vector_declarations(kernel: KernelDef) -> list[str]: + """Declare encountered widths absent from the fixed compatibility prologue.""" + declarations: dict[str, str] = {} + + def visit_type(t: Type) -> None: + if isinstance(t, PtrType): + visit_type(t.pointee) + elif isinstance(t, SmemType): + visit_type(t.elem) + elif isinstance(t, VectorType): + name = _type_to_hip(t) + prefix = ( + "boolx" + if t.elem.name == "i1" + else _vec_prefix(t.elem.name, "vector type") + ) + scalar = "int8_t" if prefix == "i8x" else _HIP_TYPE[t.elem.name] + if f"_ROCKE_VEC({scalar}, {prefix}, {t.count})" not in HIP_PROLOGUE: + declarations[name] = ( + f"using {name} = {scalar} __attribute__((ext_vector_type({t.count})));" + ) + + def visit_region(region: Region) -> None: + for op in region.ops: + for value in (*op.operands, *op.results): + visit_type(value.type) + for child in op.regions: + visit_region(child) + + for param in kernel.params: + visit_type(param.type) + visit_region(kernel.body) + return list(declarations.values()) + + def lower_kernel_to_hip( kernel: KernelDef, *, @@ -2507,6 +2543,7 @@ def lower_kernel_to_hip( parts: List[str] = [] if include_prologue: parts.append(HIP_PROLOGUE) + parts.extend(_extra_vector_declarations(kernel)) parts.append(head) if smem_block: parts.append(smem_block) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/storage.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/storage.py new file mode 100644 index 000000000000..7c59be34fc0f --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/storage.py @@ -0,0 +1,229 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Bit storage shared by matrix elements and scales, independent of an ISA.""" + +from __future__ import annotations + +from dataclasses import dataclass +from math import gcd +from collections.abc import Sequence + +from .dtypes import DTypeInfo, dtype_info + +_MAX_BITS = (1 << 64) - 1 + + +def _checked(value: int) -> int: + if not isinstance(value, int) or not 0 <= value <= _MAX_BITS: + raise ValueError("storage size or offset exceeds uint64 range") + return value + + +@dataclass(frozen=True) +class BitPacking: + """Low-bit-first patterns in a little-endian bit stream. + + slot_bits=None means dense packing. Wider slots have zero unused high bits. + No operand role, K grouping, or numeric conversion is implied. + """ + + element_bits: int + slot_bits: int | None = None + + def __post_init__(self) -> None: + slot = self.element_bits if self.slot_bits is None else self.slot_bits + if not (1 <= self.element_bits <= slot <= 64): + raise ValueError("packing requires 1 <= element_bits <= slot_bits <= 64") + object.__setattr__(self, "slot_bits", slot) + + def group(self, carrier_bits: int) -> tuple[int, int]: + """Smallest whole-carrier group: (logical slots, carriers).""" + if carrier_bits not in (8, 16, 32, 64): + raise ValueError("carrier_bits must be 8, 16, 32, or 64") + common = gcd(self.slot_bits, carrier_bits) + return carrier_bits // common, self.slot_bits // common + + def bit_offset(self, index: int) -> int: + return _checked(_checked(index) * self.slot_bits) + + def byte_size(self, count: int, bit_offset: int = 0) -> int: + _checked(bit_offset) + bits = self.bit_offset(count) + if count == 0: + return 0 + return (_checked(bits + bit_offset) + 7) // 8 + + def pack(self, patterns: Sequence[int], *, bit_offset: int = 0) -> bytes: + """Pack encoded unsigned patterns, without interpreting their values.""" + out = bytearray(self.byte_size(len(patterns), bit_offset)) + for i, pattern in enumerate(patterns): + if not 0 <= pattern < (1 << self.element_bits): + raise ValueError("pattern does not fit element_bits") + pos = bit_offset + self.bit_offset(i) + # Byte-sized pieces avoid an implicit assumption about word boundaries. + remaining = self.element_bits + while remaining: + byte, shift = divmod(pos, 8) + take = min(remaining, 8 - shift) + out[byte] |= (pattern & ((1 << take) - 1)) << shift + pattern >>= take + remaining -= take + pos += take + return bytes(out) + + def unpack(self, data: bytes, count: int, *, bit_offset: int = 0) -> list[int]: + if len(data) < self.byte_size(count, bit_offset): + raise ValueError("packed buffer is too small") + result = [] + for i in range(count): + pos = bit_offset + self.bit_offset(i) + value = 0 + for bit in range(self.element_bits): + at = pos + bit + value |= ((data[at // 8] >> (at % 8)) & 1) << bit + result.append(value) + return result + + +@dataclass(frozen=True) +class FragmentPacking: + """A payload and its carrier capacity; all unused bits are zero.""" + + packing: BitPacking + count: int + carrier_bits: int + carrier_count: int + + def __post_init__(self) -> None: + self.packing.group(self.carrier_bits) + capacity = _checked(_checked(self.carrier_count) * self.carrier_bits) + if self.packing.bit_offset(self.count) > capacity: + raise ValueError("fragment payload exceeds carrier capacity") + + @property + def payload_bits(self) -> int: + return self.packing.bit_offset(self.count) + + @property + def live_carriers(self) -> int: + return (self.payload_bits + self.carrier_bits - 1) // self.carrier_bits + + @property + def padding_bits(self) -> int: + return self.carrier_count * self.carrier_bits - self.payload_bits + + def pack(self, patterns: Sequence[int]) -> tuple[int, ...]: + if len(patterns) != self.count: + raise ValueError("fragment logical count mismatch") + width = self.carrier_bits // 8 + data = self.packing.pack(patterns) + data += bytes(self.carrier_count * width - len(data)) + return tuple( + int.from_bytes(data[i : i + width], "little") + for i in range(0, len(data), width) + ) + + +@dataclass(frozen=True) +class MatrixFragmentLayout: + """Contiguous K chunks interleaved between lane groups. + + The atom selects this mapping. Storage and carrier widths do not determine it. + """ + + fragment: FragmentPacking + chunk_elements: int + lane_groups: int + lanes_per_group: int + + def __post_init__(self) -> None: + if ( + self.chunk_elements <= 0 + or self.lane_groups <= 0 + or self.lanes_per_group <= 0 + or self.fragment.count % self.chunk_elements + ): + raise ValueError("invalid matrix fragment chunk layout") + if self.fragment.packing.bit_offset(self.chunk_elements) % 8: + raise ValueError("matrix fragment chunks must occupy whole bytes") + + @property + def chunks_per_lane(self) -> int: + return self.fragment.count // self.chunk_elements + + @property + def chunk_bytes(self) -> int: + return self.fragment.packing.byte_size(self.chunk_elements) + + def coord(self, lane: int, slot: int) -> tuple[int, int]: + if not ( + 0 <= lane < self.lane_groups * self.lanes_per_group + and 0 <= slot < self.fragment.count + ): + raise ValueError("matrix fragment coordinate out of bounds") + chunk, element = divmod(slot, self.chunk_elements) + return ( + lane % self.lanes_per_group, + (chunk * self.lane_groups + lane // self.lanes_per_group) + * self.chunk_elements + + element, + ) + + +@dataclass(frozen=True) +class TensorStorage: + """Two-dimensional storage with a contiguous packed inner axis. + + Shape is logical. Row strides are bytes. Generic axis permutations and + negative strides are intentionally rejected by this initial interface. + """ + + dtype: str + shape: tuple[int, int] + row_stride_bytes: int | None = None + slot_bits: int | None = None + base_bit_offset: int = 0 + alignment_bytes: int = 1 + + def __post_init__(self) -> None: + info = dtype_info(self.dtype) + object.__setattr__(self, "dtype", info.name) + if len(self.shape) != 2: + raise ValueError("storage requires a two-dimensional logical shape") + for dim in self.shape: + _checked(dim) + _checked(self.base_bit_offset) + if self.alignment_bytes <= 0 or self.alignment_bytes & ( + self.alignment_bytes - 1 + ): + raise ValueError("alignment_bytes must be a positive power of two") + row_bytes = self.packing.byte_size(self.shape[1], self.base_bit_offset) + stride = row_bytes if self.row_stride_bytes is None else self.row_stride_bytes + if _checked(stride) < row_bytes: + raise ValueError("row stride is smaller than the packed row") + object.__setattr__(self, "row_stride_bytes", stride) + _checked(self.byte_size) + + @property + def info(self) -> DTypeInfo: + return dtype_info(self.dtype) + + @property + def packing(self) -> BitPacking: + return BitPacking(self.info.encoded_bits, self.slot_bits) + + @property + def byte_size(self) -> int: + rows, cols = self.shape + if rows == 0 or cols == 0: + return 0 + return _checked( + (rows - 1) * self.row_stride_bytes + + self.packing.byte_size(cols, self.base_bit_offset) + ) + + def address(self, row: int, col: int) -> tuple[int, int]: + if not (0 <= row < self.shape[0] and 0 <= col < self.shape[1]): + raise ValueError("storage coordinate out of bounds") + bits = _checked(self.base_bit_offset + self.packing.bit_offset(col)) + return _checked(row * self.row_stride_bytes + bits // 8), bits % 8 diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py new file mode 100644 index 000000000000..8d0bae29b5a2 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py @@ -0,0 +1,162 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Descriptor-driven packed transport, expanded before IR serialization.""" + +from __future__ import annotations + +from collections.abc import Callable +from math import gcd + +from ..core.dtypes import dtype_info +from ..core.ir import ( + I8, + I16, + I32, + I64, + IRBuilder, + Type, + Value, + VectorType, + dtype_to_ir_type, +) +from ..core.storage import FragmentPacking, MatrixFragmentLayout, TensorStorage + + +def storage_ir_type(dtype: str) -> Type: + """Addressable storage unit, distinct from the logical dtype resolver.""" + info = dtype_info(dtype) + if info.encoded_bits % 8 or info.name in ("e8m0", "e4m3", "e5m3"): + return I8 + return dtype_to_ir_type(info.name) + + +def load_matrix_fragment( + b: IRBuilder, + ptr: Value, + row_base: Value, + lane_group: Value, + k0: int, + *, + storage: TensorStorage, + layout: MatrixFragmentLayout, + carrier_type: Type = I32, +) -> Value: + """Load whole-byte chunks directly into the atom's carrier vector. + + row_base counts pointer storage units. lane_group is derived from the atom's + lane mapping by the caller. Partial chunks and nonzero bit origins reject. + """ + packing = layout.fragment + if not (0 < packing.count <= 0x7FFFFFFF and packing.carrier_count <= 0x7FFFFFFF): + raise ValueError("invalid matrix fragment chunk layout") + unit_type = storage_ir_type(storage.dtype) + unit_bytes = dtype_info(unit_type.name).encoded_bits // 8 + if ptr.type.pointee != unit_type: + raise ValueError("matrix pointer storage type mismatch") + if storage.packing != packing.packing or storage.base_bit_offset: + raise ValueError( + "matrix fragment requires matching packing and a zero bit origin" + ) + span = layout.chunk_elements * layout.lane_groups * layout.chunks_per_lane + if k0 < 0 or k0 + span > storage.shape[1]: + raise ValueError("matrix fragment exceeds the packed row") + origin_bits = storage.packing.bit_offset(k0) + if origin_bits % (8 * unit_bytes) or layout.chunk_bytes % unit_bytes: + raise ValueError("matrix fragment is not aligned to pointer storage units") + if dtype_info(carrier_type.name).encoded_bits != packing.carrier_bits: + raise ValueError("matrix carrier type width mismatch") + if packing.payload_bits % packing.carrier_bits: + raise ValueError("matrix fragment payload must occupy whole carriers") + chunk_units = layout.chunk_bytes // unit_bytes + origin_bytes = origin_bits // 8 + if ( + origin_bytes // unit_bytes > 0x7FFFFFFF + or chunk_units * layout.lane_groups > 0x7FFFFFFF + ): + raise ValueError("matrix fragment offset exceeds i32 range") + alignment = gcd( + storage.alignment_bytes, + storage.row_stride_bytes, + layout.chunk_bytes, + origin_bytes, + ) + lane_chunk = b.mul(lane_group, b.const_i32(chunk_units)) + step_base = b.add(row_base, b.const_i32(origin_bytes // unit_bytes)) + chunks = [] + for j in range(layout.chunks_per_lane): + offset = b.add( + b.add(step_base, b.const_i32(j * layout.lane_groups * chunk_units)), + lane_chunk, + ) + remaining = chunk_units + consumed = 0 + max_width = 8 if unit_bytes == 4 else 16 + while remaining: + width = min(max_width, 1 << (remaining.bit_length() - 1)) + at = b.add(offset, b.const_i32(consumed)) if consumed else offset + load_align = gcd(alignment, consumed * unit_bytes) + if width == 1: + value = b.vector_splat( + b.global_load(ptr, at, unit_type, align=load_align), 1 + ) + else: + value = b.global_load_vN(ptr, at, unit_type, width, align=load_align) + chunks.append(value) + consumed += width + remaining -= width + payload = chunks[0] + for chunk in chunks[1:]: + payload = b.vec_concat(payload, chunk) + payload_type = VectorType(carrier_type, packing.live_carriers) + if payload.type != payload_type: + payload = b.bitcast(payload, payload_type) + padding = packing.carrier_count - packing.live_carriers + if padding: + if carrier_type != I32: + raise ValueError("padded matrix fragments currently require i32 carriers") + payload = b.vec_concat(payload, b.vector_splat(b.const_i32(0), padding)) + return payload + + +def pack_fragment_bits( + b: IRBuilder, + load_bits: Callable[[int], Value], + fragment: FragmentPacking, +) -> list[Value]: + """Pack unsigned encoded patterns into integer carriers, including split fields. + + The loader must return canonical patterns with zero high bits. Tensor decoding + and numeric quantization are separate operations. + """ + if fragment.carrier_bits not in (32, 64): + raise ValueError("IR pattern packing currently requires i32 or i64 carriers") + word_type = I64 if fragment.carrier_bits == 64 else I32 + constant = b.const_i64 if fragment.carrier_bits == 64 else b.const_i32 + words = [constant(0) for _ in range(fragment.carrier_count)] + for j in range(fragment.count): + pattern = load_bits(j) + if pattern.type not in (I8, I16, I32, I64): + raise ValueError( + "pattern packing requires unsigned patterns in integer carriers" + ) + pattern_bits = dtype_info(pattern.type.name).encoded_bits + if pattern_bits < fragment.packing.element_bits: + raise ValueError("pattern carrier is smaller than encoded width") + if pattern_bits < fragment.carrier_bits: + pattern = b.zext(pattern, word_type) + elif pattern_bits > fragment.carrier_bits: + raise ValueError("pattern carrier is wider than output carrier") + start = fragment.packing.bit_offset(j) + remaining = fragment.packing.element_bits + consumed = 0 + while remaining: + word, shift = divmod(start, fragment.carrier_bits) + take = min(remaining, fragment.carrier_bits - shift) + part = b.lshr(pattern, constant(consumed)) if consumed else pattern + if take < remaining: + part = b.land(part, constant((1 << take) - 1)) + words[word] = b.lor(words[word], b.shl(part, constant(shift))) + start += take + consumed += take + remaining -= take + return words diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/quant.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/quant.py index f8ead583eb19..1ae76ffb0d33 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/quant.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/quant.py @@ -39,7 +39,18 @@ from typing import List, Literal -from ..core.ir import BF8E5M2, F32, FP8E4M3, I8, I32, I64, IRBuilder, Type, Value +from ..core.ir import ( + BF8E5M2, + F32, + FP8E4M3, + I8, + I32, + I64, + IRBuilder, + Type, + Value, + dtype_to_ir_type, +) __all__ = [ @@ -55,7 +66,7 @@ ] -QDType = Literal["i8", "fp8e4m3", "bf8e5m2"] +QDType = Literal["i8", "fp8e4m3", "bf8e5m2", "fp4e2m1", "fp6e2m3", "fp6e3m2"] # Per-dtype clamp magnitude (the largest representable absolute value). @@ -89,6 +100,9 @@ "i8": "i8", "fp8e4m3": "fp8e4m3", "bf8e5m2": "bf8e5m2", + "fp4e2m1": "fp4e2m1", + "fp6e2m3": "fp6e2m3", + "fp6e3m2": "fp6e3m2", } @@ -104,19 +118,14 @@ def _canon(qdtype: str) -> QDType: def quant_ir_type(qdtype: str) -> Type: """Map a quant-dtype alias string to the canonical IR :class:`Type`. - Accepts ``"i8"`` / ``"int8"`` / ``"fp8e4m3"`` / ``"fp8"`` / - ``"bf8e5m2"`` / ``"bf8"``. Raises :class:`ValueError` for anything - else; this is the single point of truth for the alias map so - everywhere in the codebase agrees on the canonical names. + Accepts the existing i8/FP8/BF8 aliases and the FP4/FP6 aliases from + the common dtype resolver. Type recognition does not enable scalar + quantization or dequantization for the low-bit formats. """ - canon = _canon(qdtype) - if canon == "i8": - return I8 - if canon == "fp8e4m3": - return FP8E4M3 - if canon == "bf8e5m2": - return BF8E5M2 - raise ValueError(f"unreachable: canon={canon!r}") + # Type recognition includes low-bit formats before scalar conversions exist. + if qdtype in ("fp4", "fp4e2m1", "fp6", "fp6e2m3", "bf6", "fp6e3m2"): + return dtype_to_ir_type(qdtype) + return dtype_to_ir_type(_canon(qdtype)) def quant_max_abs(qdtype: str) -> float: 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 8a27e62f01ef..4704a66de77a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt @@ -14,6 +14,15 @@ if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/core/smoke.cpp") endif() if(BUILD_TESTING) + add_executable(rocke_storage core/test_storage.cpp) + target_link_libraries(rocke_storage PRIVATE rocke_core) + add_test(NAME rocke_storage COMMAND rocke_storage) + set_tests_properties(rocke_storage PROPERTIES LABELS "ckc;unit;host") + rocke_cxx_test_installed_name(_rocke_installed_name rocke_storage) + set_target_properties(rocke_storage PROPERTIES OUTPUT_NAME "${_rocke_installed_name}") + install(TARGETS rocke_storage RUNTIME DESTINATION ${ROCKE_CXX_TEST_INSTALL_DIR}) + set_property(GLOBAL APPEND PROPERTY ROCKE_INSTALL_TEST_TARGETS rocke_storage) + add_executable(rocke_dtypes core/test_dtypes.cpp) target_link_libraries(rocke_dtypes PRIVATE rocke_core) add_test(NAME rocke_dtypes COMMAND rocke_dtypes) 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 index 47014e02e8a2..cee48c279c6e 100644 --- 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 @@ -5,6 +5,55 @@ import pytest from rocke.core.dtypes import normalize_dtype +from rocke.core.dtypes import dtype_info +from rocke.core.ir import I8, IRBuilder, dtype_to_ir_type +from rocke.core.ir_serialize import parse, serialize +from rocke.helpers.quant import quant_ir_type, ir_to_qdtype, dequantize_scalar_to_f32 + + +@pytest.mark.parametrize( + "dtype,bits", + [ + ("fp4", 4), + ("fp6", 6), + ("bf6", 6), + ("fp8", 8), + ("bf8", 8), + ("f16", 16), + ("bf16", 16), + ("f32", 32), + ("e8m0", 8), + ("e4m3", 8), + ("e5m3", 8), + ], +) +def test_logical_dtype_and_serialization(dtype, bits): + info = dtype_info(dtype) + logical = dtype_to_ir_type(dtype) + assert info.encoded_bits == bits + assert logical != I8 + b = IRBuilder("dtype_transport") + b.param("code", logical) + assert serialize(parse(serialize(b.kernel))) == serialize(b.kernel) + + +@pytest.mark.parametrize("dtype", ["fp4", "fp6", "bf6"]) +def test_quant_type_recognition_does_not_enable_integer_conversion(dtype): + logical = quant_ir_type(dtype) + assert logical == dtype_to_ir_type(dtype) + assert ir_to_qdtype(logical) == dtype_info(dtype).name + b = IRBuilder("no_lowbit_conversion") + value = b.param("value", logical) + scale = b.const_f32(1) + with pytest.raises(ValueError, match="unsupported input type"): + dequantize_scalar_to_f32(b, value, scale=scale) + + +def test_unknown_dtype_and_unrepresented_integer_family(): + with pytest.raises(ValueError, match="unknown dtype"): + dtype_info("custom_format") + with pytest.raises(ValueError, match="no logical IR type"): + dtype_to_ir_type("iu4") @pytest.mark.parametrize( diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py new file mode 100644 index 000000000000..79af622a235c --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py @@ -0,0 +1,147 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Mirrored authoring and bit-pattern placement, independent of MMA support.""" + +import os +import subprocess +from pathlib import Path + +import pytest + +from rocke.core.arch.wmma_scale import scaled_matrix_layout +from rocke.core.ir import I8, I32, I64, IRBuilder, PtrType +from rocke.core.ir_serialize import serialize +from rocke.core.lower_llvm import lower_kernel_to_llvm +from rocke.core.lower_hip import lower_kernel_to_hip +from rocke.core.storage import ( + BitPacking, + FragmentPacking, + MatrixFragmentLayout, + TensorStorage, +) +from rocke.helpers.mma_io import ( + load_matrix_fragment, + pack_fragment_bits, + storage_ir_type, +) + + +CASES = [ + "fp4", + "fp6", + "bf6", + "fp8", + "bf8", + "f16", + "bf16", + "pack_fp6_cross_word", + "pack_scale_bytes_i32", + "pack_scale_bytes_i64", +] + + +def build_transport(dtype): + b = IRBuilder("transport") + patterns = dtype.startswith("pack_") + unit = I8 if patterns else storage_ir_type(dtype) + typed = dtype in ("f16", "bf16") + carrier = unit if typed else I64 if dtype == "pack_scale_bytes_i64" else I32 + a = b.param("A", PtrType(unit, "global")) + o = b.param("O", PtrType(carrier, "global")) + if patterns: + bits = 6 if dtype == "pack_fp6_cross_word" else 8 + count = 16 if bits == 6 else 8 if dtype == "pack_scale_bytes_i64" else 4 + words = 3 if bits == 6 else 1 + fragment = FragmentPacking( + BitPacking(bits), count, 64 if carrier == I64 else 32, words + ) + values = pack_fragment_bits( + b, lambda j: b.global_load(a, b.const_i32(j), I8, align=1), fragment + ) + for j, value in enumerate(values): + b.global_store(o, b.const_i32(j), value, align=4) + else: + storage = TensorStorage(dtype, (16, 128), alignment_bytes=16) + layout = ( + MatrixFragmentLayout(FragmentPacking(BitPacking(16), 32, 16, 32), 16, 2, 16) + if typed + else scaled_matrix_layout(dtype, 16) + ) + base = b.const_i32(0) + lane = b.mod(b.thread_id_x(), b.const_i32(32)) + group = b.div(lane, b.const_i32(16)) + value = load_matrix_fragment( + b, a, base, group, 0, storage=storage, layout=layout, carrier_type=carrier + ) + for j in range(value.type.count): + index = b.const_i32(j) + element = b.vec_extract(value, j) + b.global_store(o, index, element, align=2 if typed else 4) + return b.kernel + + +@pytest.mark.parametrize("dtype", CASES) +def test_descriptor_authoring_lowers_in_both_backends(dtype): + kernel = build_transport(dtype) + assert "define" in lower_kernel_to_llvm( + kernel, arch="gfx1250", llvm_flavor="llvm23" + ) + assert "transport" in lower_kernel_to_hip(kernel, arch="gfx1250") + + +@pytest.mark.parametrize("dtype", CASES) +@pytest.mark.parametrize("route", ["ir", "hip"]) +def test_native_helper_emits_identical_ir(dtype, route): + executable = os.environ.get("ROCKE_STORAGE_TEST") + if not executable: + pytest.skip("set ROCKE_STORAGE_TEST to the built native storage test") + assert Path(executable).is_file() + actual = subprocess.run( + [executable, "--hip" if route == "hip" else "--emit", dtype], + check=True, + capture_output=True, + text=True, + ) + kernel = build_transport(dtype) + expected = ( + serialize(kernel) + if route == "ir" + else lower_kernel_to_hip(kernel, arch="gfx1250") + ) + assert actual.stdout == expected + + +def test_reject_partial_fragment_and_wrong_pointer(): + b = IRBuilder("invalid_transport") + ptr = b.param("A", PtrType(I8, "global")) + zero = b.const_i32(0) + with pytest.raises(ValueError, match="exceeds"): + load_matrix_fragment( + b, + ptr, + zero, + zero, + 1, + storage=TensorStorage("fp4", (16, 128)), + layout=scaled_matrix_layout("fp4", 16), + ) + with pytest.raises(ValueError, match="pointer storage type"): + load_matrix_fragment( + b, + ptr, + zero, + zero, + 0, + storage=TensorStorage("fp8", (16, 128)), + layout=scaled_matrix_layout("fp8", 16), + ) + + +def test_hip_declares_only_encountered_missing_vector_widths(): + source = lower_kernel_to_hip(build_transport("fp6"), arch="gfx1250") + for name in ("i8x24", "i8x40", "i32x12"): + assert source.count(f"using {name} =") == 1 + assert "using f16x32 =" not in source + source = lower_kernel_to_hip(build_transport("f16"), arch="gfx1250") + assert source.count("using f16x32 =") == 1 + assert "using i8x24 =" not in source diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp new file mode 100644 index 000000000000..ef2b578de597 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp @@ -0,0 +1,186 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* Host packing fixtures and native authoring emission for test_mma_io.py. */ +#include "rocke/storage.h" + +#include +#include +#include +#include + +#include "rocke/helper_rocke.helpers.mma_io.h" +#include "rocke/helper_rocke.helpers.quant.h" +#include "rocke/ir_serialize.h" +#include "rocke/lower_hip.h" +#include "rocke/wmma_scale_internal.h" + +#define CHECK(x) \ + do \ + { \ + if(!(x)) \ + { \ + fprintf(stderr, "check failed at %d: %s\n", __LINE__, #x); \ + return 1; \ + } \ + } while(0) + +static rocke_value_t* load_bits(rocke_ir_builder_t* b, int j, void* ctx) +{ + return rocke_b_global_load(b, (rocke_value_t*)ctx, rocke_b_const_i32(b, j), rocke_i8(), 1); +} + +static int emit(const char* dtype, bool hip) +{ + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "transport") == ROCKE_OK); + bool patterns = strncmp(dtype, "pack_", 5) == 0; + const auto* unit = patterns ? rocke_i8() : rocke_storage_ir_type(dtype); + CHECK(unit); + const bool typed = strcmp(dtype, "f16") == 0 || strcmp(dtype, "bf16") == 0; + const auto* carrier = typed ? unit : rocke_i32(); + if(strcmp(dtype, "pack_scale_bytes_i64") == 0) + carrier = rocke_i64(); + auto* a = rocke_b_param(&b, "A", rocke_ptr_type(&b, unit, "global"), NULL); + auto* o = rocke_b_param(&b, "O", rocke_ptr_type(&b, carrier, "global"), NULL); + if(patterns) + { + int bits = strcmp(dtype, "pack_fp6_cross_word") == 0 ? 6 : 8; + int count = bits == 6 ? 16 : (strcmp(dtype, "pack_scale_bytes_i64") == 0 ? 8 : 4); + rocke_bit_packing_t packing; + rocke_fragment_packing_t fragment; + CHECK(rocke_bit_packing_init(&packing, bits, 0)); + int words = bits == 6 ? 3 : 1; + CHECK(rocke_fragment_packing_init( + &fragment, &packing, count, carrier == rocke_i64() ? 64 : 32, words)); + rocke_value_t* values[3]; + CHECK(rocke_h_pack_fragment_bits(&b, load_bits, a, &fragment, values, words) == ROCKE_OK); + for(int j = 0; j < words; ++j) + rocke_b_global_store(&b, o, rocke_b_const_i32(&b, j), values[j], 4); + } + else + { + rocke_tensor_storage_t storage; + CHECK(rocke_tensor_storage_init(&storage, dtype, 16, 128, UINT64_MAX, 0, 0, 16)); + rocke_matrix_fragment_layout_t layout; + if(typed) + { + rocke_bit_packing_t packing; + rocke_fragment_packing_t fragment; + CHECK(rocke_bit_packing_init(&packing, 16, 0)); + CHECK(rocke_fragment_packing_init(&fragment, &packing, 32, 16, 32)); + CHECK(rocke_matrix_fragment_layout_init(&layout, &fragment, 16, 2, 16)); + } + else + layout = rocke_scaled_matrix_layout(dtype, 16); + auto* base = rocke_b_const_i32(&b, 0); + auto* thread = rocke_b_thread_id_x(&b); + auto* lane = rocke_b_mod(&b, thread, rocke_b_const_i32(&b, 32)); + auto* group = rocke_b_div(&b, lane, rocke_b_const_i32(&b, 16)); + auto* value + = rocke_h_load_matrix_fragment(&b, a, base, group, 0, &storage, &layout, carrier); + CHECK(value && rocke_ir_builder_ok(&b)); + for(int j = 0; j < value->type->count; ++j) + { + auto* index = rocke_b_const_i32(&b, j); + auto* element = rocke_b_vec_extract(&b, value, j); + rocke_b_global_store(&b, o, index, element, typed ? 2 : 4); + } + } + CHECK(rocke_ir_builder_ok(&b)); + if(hip) + { + rocke_strbuf_t text; + CHECK(rocke_strbuf_init(&text, 256) == 0); + rocke_lower_hip_opts_t opts = {}; + opts.arch = "gfx1250"; + CHECK(rocke_lower_kernel_to_hip(&b, b.kernel, &opts, &text) == ROCKE_OK); + fputs(rocke_strbuf_cstr(&text), stdout); + rocke_strbuf_free(&text); + } + else + { + char* text = NULL; + CHECK(rocke_ir_serialize(b.kernel, &text) == ROCKE_OK); + fputs(text, stdout); + free(text); + } + rocke_ir_builder_free(&b); + return 0; +} + +int main(int argc, char** argv) +{ + if(argc == 3 && strcmp(argv[1], "--emit") == 0) + return emit(argv[2], false); + if(argc == 3 && strcmp(argv[1], "--hip") == 0) + return emit(argv[2], true); + for(int bits : {4, 6, 8, 16, 32, 64}) + { + rocke_bit_packing_t packing; + CHECK(rocke_bit_packing_init(&packing, bits, 0)); + uint64_t patterns[256], decoded[256]; + size_t count = bits <= 8 ? size_t(1) << bits : 4; + for(size_t i = 0; i < count; ++i) + patterns[i] = bits <= 8 ? i : (i == 3 ? UINT64_MAX >> (64 - bits) : i); + for(uint64_t offset : {uint64_t(0), uint64_t(1), uint64_t(7), uint64_t(9)}) + { + uint64_t size; + CHECK(rocke_bit_packing_bytes(&packing, count, offset, &size)); + uint8_t data[2049]; + memset(data, 0xA5, sizeof(data)); + CHECK(rocke_bit_pack(&packing, patterns, count, offset, data, size)); + CHECK(data[size] == 0xA5); // No tail overrun. + CHECK(!rocke_bit_unpack(&packing, data, size - 1, count, offset, decoded)); + CHECK(rocke_bit_unpack(&packing, data, size, count, offset, decoded)); + CHECK(memcmp(patterns, decoded, count * sizeof(uint64_t)) == 0); + } + } + rocke_bit_packing_t six; + CHECK(rocke_bit_packing_init(&six, 6, 0)); + uint64_t values, carriers; + CHECK(rocke_bit_packing_group(&six, 32, &values, &carriers)); + CHECK(values == 16 && carriers == 3); + uint64_t patterns[16] = {}, words[3]; + patterns[5] = patterns[10] = 63; + rocke_fragment_packing_t fragment; + CHECK(rocke_fragment_packing_init(&fragment, &six, 16, 32, 3)); + CHECK(rocke_fragment_pack(&fragment, patterns, 16, words, 3)); + CHECK(words[0] == 0xC0000000 && words[1] == 0xF000000F && words[2] == 3); + CHECK(!rocke_bit_packing_bytes(&six, UINT64_MAX, 0, &values)); + CHECK(!rocke_fragment_packing_init(&fragment, &six, 16, 32, 2)); + rocke_tensor_storage_t storage; + CHECK(rocke_tensor_storage_init(&storage, "bf6", 3, 5, 8, 0, 3, 1)); + CHECK(rocke_tensor_storage_bytes(&storage, &values) && values == 21); + int shift; + CHECK(rocke_tensor_storage_address(&storage, 2, 4, &values, &shift)); + CHECK(values == 19 && shift == 3); + CHECK(!rocke_tensor_storage_address(&storage, 3, 0, &values, &shift)); + CHECK(!rocke_tensor_storage_init(&storage, "fp6", 3, 5, 3, 0, 0, 1)); + CHECK(!rocke_tensor_storage_init(&storage, "fp8", UINT64_MAX, 3, 3, 0, 0, 1)); + CHECK(rocke_tensor_storage_init(&storage, "fp4", 3, 0, 0, 0, 7, 1)); + CHECK(rocke_tensor_storage_bytes(&storage, &values) && values == 0); + for(const char* dtype : {"fp4", "fp6", "bf6", "e8m0", "e4m3", "e5m3"}) + { + const auto* logical = rocke_dtype_to_ir_type(dtype); + CHECK(logical && logical != rocke_i8()); + CHECK(rocke_scalar_by_name(logical->name) == logical); + if(dtype[0] != 'e') + CHECK(rocke_quant_ir_type(dtype) == logical); + rocke_ir_builder_t b, parsed; + CHECK(rocke_ir_builder_init(&b, "types") == ROCKE_OK); + CHECK(rocke_ir_builder_init(&parsed, "parsed") == ROCKE_OK); + rocke_b_param(&b, "pattern", logical, NULL); + char* text = NULL; + CHECK(rocke_ir_serialize(b.kernel, &text) == ROCKE_OK); + rocke_kernel_def_t* kernel = NULL; + CHECK(rocke_ir_parse(text, &parsed, &kernel) == ROCKE_OK); + char* again = NULL; + CHECK(rocke_ir_serialize(kernel, &again) == ROCKE_OK); + CHECK(strcmp(text, again) == 0); + free(text); + free(again); + rocke_ir_builder_free(&b); + rocke_ir_builder_free(&parsed); + } + return 0; +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py new file mode 100644 index 000000000000..515430599388 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py @@ -0,0 +1,94 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Independent bit-stream fixtures and storage bounds for matrices and scales.""" + +import pytest + +from rocke.core.arch.wmma_scale import ScalePacking, scaled_matrix_layout +from rocke.core.storage import BitPacking, FragmentPacking, TensorStorage + + +@pytest.mark.parametrize("bits", [4, 6, 8, 16, 32, 64]) +@pytest.mark.parametrize("offset", [0, 1, 7, 9]) +def test_dense_patterns_against_integer_oracle(bits, offset): + patterns = list(range(1 << bits)) if bits <= 8 else [0, 1, (1 << bits) - 1, 0x1234] + packing = BitPacking(bits) + encoded = packing.pack(patterns, bit_offset=offset) + expected = sum(v << (offset + i * bits) for i, v in enumerate(patterns)) + assert int.from_bytes(encoded, "little") == expected + assert packing.unpack(encoded, len(patterns), bit_offset=offset) == patterns + with pytest.raises(ValueError, match="too small"): + packing.unpack(encoded[:-1], len(patterns), bit_offset=offset) + + +def test_fp6_crosses_words_without_per_word_padding(): + packing = BitPacking(6) + assert packing.group(8) == (4, 3) + assert packing.group(32) == (16, 3) + patterns = [0] * 16 + patterns[5] = patterns[10] = 63 + # slot 5 occupies bits 30..35; slot 10 occupies bits 60..65. + expected = (0xC0000000, 0xF000000F, 0x00000003) + assert FragmentPacking(packing, 16, 32, 3).pack(patterns) == expected + assert BitPacking(6, 8).pack(patterns) == bytes(patterns) + assert packing.pack([1, 2, 3, 4]) == bytes([0x81, 0x30, 0x10]) + + +@pytest.mark.parametrize( + "dtype,live", [("fp8", 16), ("fp6", 12), ("bf6", 12), ("fp4", 8)] +) +def test_matrix_payload_and_padding(dtype, live): + layout = scaled_matrix_layout(dtype, 16) + assert layout.fragment.count == 64 + assert layout.fragment.live_carriers == live + assert layout.fragment.padding_bits == (16 - live) * 32 + patterns = [(1 << layout.fragment.packing.element_bits) - 1] * 64 + assert layout.fragment.pack(patterns) == (0xFFFFFFFF,) * live + (0,) * (16 - live) + for lane in range(32): + coords = [layout.coord(lane, slot) for slot in range(64)] + assert all(row == lane % 16 and 0 <= k < 128 for row, k in coords) + other = {layout.coord(lane ^ 16, slot)[1] for slot in range(64)} + assert {k for _, k in coords}.isdisjoint(other) + assert {k for _, k in coords} | other == set(range(128)) + + +@pytest.mark.parametrize( + "count,block_k,word", [(4, 32, 0x04030201), (8, 16, 0x0807060504030201)] +) +def test_scales_use_common_bit_packing(count, block_k, word): + scales = ScalePacking(count, block_k) + assert scales.association.block_k == block_k + assert scales.fragment.pack(list(range(1, count + 1))) == (word,) + assert scales.packing == BitPacking(8) + assert scales.word_bits == count * 8 + + +def test_storage_rows_offsets_and_empty_views(): + storage = TensorStorage("bf6", (3, 5), row_stride_bytes=8, base_bit_offset=3) + assert storage.dtype == "fp6e3m2" + assert storage.byte_size == 21 + assert storage.address(2, 4) == (19, 3) + assert TensorStorage("fp4", (0, 5)).byte_size == 0 + assert TensorStorage("fp4", (3, 0), base_bit_offset=7).byte_size == 0 + assert BitPacking(6).pack([], bit_offset=7) == b"" + assert TensorStorage("e8m0", (2, 8)).byte_size == 16 + with pytest.raises(ValueError, match="out of bounds"): + storage.address(3, 0) + with pytest.raises(ValueError, match="row stride"): + TensorStorage("fp6", (3, 5), row_stride_bytes=3) + with pytest.raises(ValueError, match="power of two"): + TensorStorage("fp4", (1, 4), alignment_bytes=3) + + +def test_invalid_packing_and_overflow(): + for bits, slots in [(0, None), (65, None), (6, 4)]: + with pytest.raises(ValueError): + BitPacking(bits, slots) + with pytest.raises(ValueError, match="capacity"): + FragmentPacking(BitPacking(6), 16, 32, 2) + with pytest.raises(ValueError, match="uint64"): + BitPacking(6).byte_size(1 << 63) + with pytest.raises(ValueError, match="uint64"): + TensorStorage("fp8", (1 << 63, 3)) + with pytest.raises(ValueError, match="fit"): + BitPacking(6).pack([64]) From 239e0c66f45d33a313cefe9f97ef638932590148 Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Tue, 22 Sep 2026 21:22:35 -0400 Subject: [PATCH 02/24] fix(rocke): preserve HIP vector load alignment --- .../cpp/core/lower_hip/lower_hip_mem.cpp | 19 ++++++++++ .../platform/python/rocke/core/lower_hip.py | 13 +++++++ .../rocke/platform/tests/core/test_mma_io.py | 37 ++++++++++++++++++- .../platform/tests/core/test_storage.cpp | 6 ++- 4 files changed, 73 insertions(+), 2 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp index 0b0584e37904..198da0fc0c2d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp @@ -15,6 +15,7 @@ * rocke_h_name, rocke_h_type_to_hip, rocke_h_hip_scalar, rocke_h_vec_prefix, * rocke_h_smem_set_storage/_storage, rocke_h_fail, rocke_h_live) are NOT defined here. */ +#include "rocke/dtypes.h" #include "rocke/lower_hip_internal.h" #include @@ -530,6 +531,24 @@ static rocke_status_t _op_memref_global_load_vN(rocke_h_lowerer_t* lw, const roc elem_name = mem_attr_str(op, "elem_type", "f16"); prefix = rocke_h_vec_prefix_checked(lw, elem_name, /*full_map=*/true, "global_load_vN"); res = rocke_h_name(lw, op->results[0]); + const int64_t byte_count = vec * (rocke_dtype_info(elem_name)->encoded_bits / 8); + const int64_t align = mem_attr_int(op, "align", vec * 2); + if(align < byte_count) + { + /* Do not strengthen the IR's alignment through a vector-pointer cast. */ + rocke_h_emitf( + lw, + "%s%lld %s; __builtin_memcpy(&%s, __builtin_assume_aligned(%s + %s, %lld), %lld);", + prefix, + (long long)vec, + res, + res, + rocke_h_name(lw, ptr), + rocke_h_name(lw, idx), + (long long)align, + (long long)byte_count); + return lw->status; + } rocke_h_emitf(lw, "%s%lld %s = *reinterpret_cast(%s + %s);", prefix, 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 22e38bca4581..cf601237d07f 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 @@ -18,6 +18,7 @@ from typing import List, Optional from .arch.wmma_scale import gfx1250_scaled_wmma +from .dtypes import dtype_info from .ir import ( KernelDef, Op, @@ -486,6 +487,18 @@ def _op_memref_global_load_vN(self, op: Op) -> None: vec = int(op.attrs["vec"]) elem_name = op.attrs.get("elem_type", "f16") prefix = _vec_prefix(elem_name, "global_load_vN") + byte_count = vec * (dtype_info(elem_name).encoded_bits // 8) + align = int(op.attrs.get("align", vec * 2)) + if align < byte_count: + # A vector-pointer dereference would promise natural alignment. + # Copy the payload using only the alignment guaranteed by the IR. + self._emit( + f"{prefix}{vec} {_name(op.result)}; " + f"__builtin_memcpy(&{_name(op.result)}, " + f"__builtin_assume_aligned({_name(ptr)} + {_name(idx)}, {align}), " + f"{byte_count});" + ) + return self._emit( f"{prefix}{vec} {_name(op.result)} = " f"*reinterpret_cast({_name(ptr)} + {_name(idx)});" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py index 79af622a235c..b711ecf45d26 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py @@ -3,6 +3,8 @@ """Mirrored authoring and bit-pattern placement, independent of MMA support.""" import os +import re +import shutil import subprocess from pathlib import Path @@ -29,6 +31,7 @@ CASES = [ "fp4", "fp6", + "fp6_padded", "bf6", "fp8", "bf8", @@ -42,6 +45,8 @@ def build_transport(dtype): b = IRBuilder("transport") + padded = dtype == "fp6_padded" + dtype = "fp6" if padded else dtype patterns = dtype.startswith("pack_") unit = I8 if patterns else storage_ir_type(dtype) typed = dtype in ("f16", "bf16") @@ -61,7 +66,12 @@ def build_transport(dtype): for j, value in enumerate(values): b.global_store(o, b.const_i32(j), value, align=4) else: - storage = TensorStorage(dtype, (16, 128), alignment_bytes=16) + storage = TensorStorage( + dtype, + (16, 128), + row_stride_bytes=97 if padded else None, + alignment_bytes=16, + ) layout = ( MatrixFragmentLayout(FragmentPacking(BitPacking(16), 32, 16, 32), 16, 2, 16) if typed @@ -145,3 +155,28 @@ def test_hip_declares_only_encountered_missing_vector_widths(): source = lower_kernel_to_hip(build_transport("f16"), arch="gfx1250") assert source.count("using f16x32 =") == 1 assert "using i8x24 =" not in source + + +@pytest.mark.parametrize( + "dtype,copy_bytes,alignment", + [("fp6", 16, 8), ("fp6_padded", 16, 1), ("f16", 32, 16)], +) +def test_hip_compiler_preserves_fragment_load_alignment(dtype, copy_bytes, alignment): + """Inspect compiler IR: source parity cannot detect a shared alignment bug.""" + if not shutil.which("hipcc"): + pytest.skip("hipcc not in PATH") + from rocke.helpers.compile import emit_device_llvm_ir_via_hipcc + + llvm = emit_device_llvm_ir_via_hipcc( + build_transport(dtype), arch="gfx950", extra_flags=["-O0"] + ) + source_alignments = [ + int(align) + for align, size in re.findall( + r"@llvm\.memcpy[^\n]*?\([^,]+,\s*ptr[^,]*\balign (\d+) [^,]+,\s*i64 (\d+)", + llvm, + ) + if int(size) == copy_bytes + ] + assert source_alignments, "expected unaligned-safe fragment payload copies" + assert all(value <= alignment for value in source_alignments), source_alignments diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp index ef2b578de597..47d4749920fb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp @@ -33,6 +33,9 @@ static int emit(const char* dtype, bool hip) { rocke_ir_builder_t b; CHECK(rocke_ir_builder_init(&b, "transport") == ROCKE_OK); + const bool padded = strcmp(dtype, "fp6_padded") == 0; + if(padded) + dtype = "fp6"; bool patterns = strncmp(dtype, "pack_", 5) == 0; const auto* unit = patterns ? rocke_i8() : rocke_storage_ir_type(dtype); CHECK(unit); @@ -60,7 +63,8 @@ static int emit(const char* dtype, bool hip) else { rocke_tensor_storage_t storage; - CHECK(rocke_tensor_storage_init(&storage, dtype, 16, 128, UINT64_MAX, 0, 0, 16)); + CHECK(rocke_tensor_storage_init( + &storage, dtype, 16, 128, padded ? 97 : UINT64_MAX, 0, 0, 16)); rocke_matrix_fragment_layout_t layout; if(typed) { From 70ec174ef8559f90fca7092416a10071fc708810 Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Tue, 22 Sep 2026 22:37:50 -0400 Subject: [PATCH 03/24] fix(rocke): validate load alignment and fragment coordinates --- .../cpp/core/lower_hip/lower_hip_mem.cpp | 3 ++ .../rocke/platform/cpp/core/storage.cpp | 8 +++- .../platform/python/rocke/core/lower_hip.py | 4 ++ .../rocke/platform/tests/core/test_mma_io.py | 11 +++++ .../platform/tests/core/test_storage.cpp | 48 +++++++++++++++++++ 5 files changed, 72 insertions(+), 2 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp index 198da0fc0c2d..cdfee832aad9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp @@ -533,6 +533,9 @@ static rocke_status_t _op_memref_global_load_vN(rocke_h_lowerer_t* lw, const roc res = rocke_h_name(lw, op->results[0]); const int64_t byte_count = vec * (rocke_dtype_info(elem_name)->encoded_bits / 8); const int64_t align = mem_attr_int(op, "align", vec * 2); + if(align <= 0 || (align & (align - 1))) + return rocke_h_fail( + lw, ROCKE_ERR_VALUE, "global_load_vN: alignment must be a positive power of two"); if(align < byte_count) { /* Do not strengthen the IR's alignment through a vector-pointer cast. */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/storage.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/storage.cpp index b1a2dda9cb6c..5ad91cf3287a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/storage.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/storage.cpp @@ -253,8 +253,12 @@ bool rocke_matrix_fragment_layout_init(rocke_matrix_fragment_layout_t* out, bool rocke_matrix_fragment_coord( const rocke_matrix_fragment_layout_t* p, int lane, uint64_t slot, uint64_t* row, uint64_t* k) { - if(!p || !row || !k || lane < 0 || p->lanes_per_group <= 0 || p->lane_groups <= 0 - || p->chunk_elements <= 0 || uint64_t(lane) >= uint64_t(p->lane_groups) * p->lanes_per_group + // Public C descriptors can be constructed without going through init(). + rocke_matrix_fragment_layout_t checked; + if(!p || !row || !k || lane < 0 + || !rocke_matrix_fragment_layout_init( + &checked, &p->fragment, p->chunk_elements, p->lane_groups, p->lanes_per_group) + || uint64_t(lane) >= uint64_t(p->lane_groups) * p->lanes_per_group || slot >= p->fragment.count) return false; uint64_t chunk; 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 cf601237d07f..1352f07217ff 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 @@ -489,6 +489,10 @@ def _op_memref_global_load_vN(self, op: Op) -> None: prefix = _vec_prefix(elem_name, "global_load_vN") byte_count = vec * (dtype_info(elem_name).encoded_bits // 8) align = int(op.attrs.get("align", vec * 2)) + if align <= 0 or align & (align - 1): + raise ValueError( + "global_load_vN: alignment must be a positive power of two" + ) if align < byte_count: # A vector-pointer dereference would promise natural alignment. # Copy the payload using only the alignment guaranteed by the IR. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py index b711ecf45d26..5ac95d34e755 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py @@ -180,3 +180,14 @@ def test_hip_compiler_preserves_fragment_load_alignment(dtype, copy_bytes, align ] assert source_alignments, "expected unaligned-safe fragment payload copies" assert all(value <= alignment for value in source_alignments), source_alignments + + +@pytest.mark.parametrize("alignment", [0, -1, -16, 3, 24]) +def test_hip_rejects_invalid_vector_load_alignment(alignment): + b = IRBuilder("invalid_alignment") + ptr = b.param("A", PtrType(I8, "global")) + value = b.global_load_vN(ptr, b.const_i32(0), I8, 16) + # Serialized/custom IR can bypass the builder's default alignment handling. + value.op.attrs["align"] = alignment + with pytest.raises(ValueError, match="alignment must be a positive power of two"): + lower_kernel_to_hip(b.kernel, arch="gfx1250") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp index 47d4749920fb..fd9b9ffbd7a3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp @@ -118,6 +118,54 @@ int main(int argc, char** argv) return emit(argv[2], false); if(argc == 3 && strcmp(argv[1], "--hip") == 0) return emit(argv[2], true); + for(int alignment : {0, -1, -16, 3, 24}) + { + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "invalid_alignment") == ROCKE_OK); + auto* ptr = rocke_b_param(&b, "A", rocke_ptr_type(&b, rocke_i8(), "global"), NULL); + auto* value = rocke_b_global_load_vN(&b, ptr, rocke_b_const_i32(&b, 0), rocke_i8(), 16, 16); + CHECK(value); + // Exercise raw IR, including alignments normalized by the builder. + rocke_attr_set_int(&b, &value->op->attrs, "align", alignment); + rocke_strbuf_t text; + CHECK(rocke_strbuf_init(&text, 256) == 0); + rocke_lower_hip_opts_t opts = {}; + opts.arch = "gfx1250"; + CHECK(rocke_lower_kernel_to_hip(&b, b.kernel, &opts, &text) == ROCKE_ERR_VALUE); + rocke_strbuf_free(&text); + rocke_ir_builder_free(&b); + } + const auto layout = rocke_scaled_matrix_layout("fp6", 16); + uint64_t row, k; + CHECK(rocke_matrix_fragment_coord(&layout, 31, 63, &row, &k)); + CHECK(row == 15 && k == 127); + // C callers can construct or mutate public descriptors without init(). + for(int defect = 0; defect < 6; ++defect) + { + auto invalid = layout; + switch(defect) + { + case 0: + invalid.fragment.packing.element_bits = 0; + break; + case 1: + invalid.fragment.packing.slot_bits = 4; + break; + case 2: + invalid.fragment.carrier_bits = 7; + break; + case 3: + invalid.fragment.carrier_count = 1; + break; + case 4: + invalid.chunk_elements = 3; + break; + case 5: + invalid.chunk_elements = 2; + break; // 12 bits: not byte aligned. + } + CHECK(!rocke_matrix_fragment_coord(&invalid, 0, 0, &row, &k)); + } for(int bits : {4, 6, 8, 16, 32, 64}) { rocke_bit_packing_t packing; From 29b5454ee2a7434ec8236d1cc7f2bf8b91ed1ae0 Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Thu, 24 Sep 2026 14:08:31 -0400 Subject: [PATCH 04/24] fix(rocke): validate storage descriptors and typed row strides --- .../rocke/platform/cpp/core/storage.cpp | 18 ++- .../rocke/platform/cpp/helpers/mma_io.cpp | 6 +- .../platform/cpp/include/rocke/storage.h | 1 + .../architecture/dtype_storage_packing.md | 7 + .../platform/python/rocke/helpers/mma_io.py | 9 +- .../rocke/platform/tests/core/test_mma_io.py | 61 ++++++++- .../platform/tests/core/test_storage.cpp | 125 +++++++++++++++++- .../rocke/platform/tests/core/test_storage.py | 7 + 8 files changed, 217 insertions(+), 17 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/storage.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/storage.cpp index 5ad91cf3287a..131c00e3989d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/storage.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/storage.cpp @@ -196,7 +196,7 @@ bool rocke_tensor_storage_init(rocke_tensor_storage_t* out, p.row_stride_bytes = row_stride_bytes == UINT64_MAX ? row_bytes : row_stride_bytes; p.base_bit_offset = base_bit_offset; p.alignment_bytes = alignment_bytes; - if(p.row_stride_bytes < row_bytes || !rocke_tensor_storage_bytes(&p, &bytes)) + if(!rocke_tensor_storage_bytes(&p, &bytes)) return false; *out = p; return true; @@ -204,22 +204,28 @@ bool rocke_tensor_storage_init(rocke_tensor_storage_t* out, bool rocke_tensor_storage_bytes(const rocke_tensor_storage_t* p, uint64_t* bytes) { - if(!p || !bytes || !valid(&p->packing)) + uint64_t row_bytes, prefix; + if(!p || !bytes || !p->dtype || !valid(&p->packing) + || rocke_dtype_info(p->dtype->name) != p->dtype + || p->packing.element_bits != p->dtype->encoded_bits || !p->alignment_bytes + || (p->alignment_bytes & (p->alignment_bytes - 1)) + || !rocke_bit_packing_bytes(&p->packing, p->cols, p->base_bit_offset, &row_bytes) + || p->row_stride_bytes < row_bytes) return false; if(!p->rows || !p->cols) { *bytes = 0; return true; } - uint64_t row_bytes, prefix; - return rocke_bit_packing_bytes(&p->packing, p->cols, p->base_bit_offset, &row_bytes) - && mul(p->rows - 1, p->row_stride_bytes, &prefix) && add(prefix, row_bytes, bytes); + return mul(p->rows - 1, p->row_stride_bytes, &prefix) && add(prefix, row_bytes, bytes); } bool rocke_tensor_storage_address( const rocke_tensor_storage_t* p, uint64_t row, uint64_t col, uint64_t* byte, int* shift) { - if(!p || !byte || !shift || row >= p->rows || col >= p->cols) + uint64_t bytes; + if(!byte || !shift || !rocke_tensor_storage_bytes(p, &bytes) || row >= p->rows + || col >= p->cols) return false; uint64_t bits, prefix; if(!rocke_bit_packing_offset(&p->packing, col, &bits) || !add(bits, p->base_bit_offset, &bits) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp index ecb70efaefb2..c7d01678a04e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp @@ -45,6 +45,9 @@ rocke_value_t* rocke_h_load_matrix_fragment(rocke_ir_builder_t* b, if(!storage || !storage->dtype || !layout || !carrier_type || !ptr || !row_base || !lane_group) ckc::raise_status(ROCKE_ERR_VALUE, "null matrix fragment argument"); + uint64_t storage_bytes; + if(!rocke_tensor_storage_bytes(storage, &storage_bytes)) + ckc::raise_status(ROCKE_ERR_VALUE, "invalid tensor storage descriptor"); const auto& packing = layout->fragment; rocke_matrix_fragment_layout_t checked; if(!rocke_matrix_fragment_layout_init(&checked, @@ -69,7 +72,8 @@ rocke_value_t* rocke_h_load_matrix_fragment(rocke_ir_builder_t* b, uint64_t origin_bits, chunk_bytes; if(!rocke_bit_packing_offset(&storage->packing, k0, &origin_bits) || !rocke_bit_packing_bytes(&packing.packing, layout->chunk_elements, 0, &chunk_bytes) - || origin_bits % (8 * unit_bytes) || chunk_bytes % unit_bytes) + || origin_bits % (8 * unit_bytes) || chunk_bytes % unit_bytes + || storage->row_stride_bytes % unit_bytes) ckc::raise_status(ROCKE_ERR_VALUE, "matrix fragment is not aligned to pointer storage units"); const rocke_dtype_info_t* carrier_info = rocke_dtype_info(carrier_type->name); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/storage.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/storage.h index 3ea1794d6369..133c1310c9ac 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/storage.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/storage.h @@ -104,6 +104,7 @@ bool rocke_tensor_storage_init(rocke_tensor_storage_t* out, int slot_bits, uint64_t base_bit_offset, uint64_t alignment_bytes); +/* Size and address queries validate public descriptors; invalid input returns false. */ bool rocke_tensor_storage_bytes(const rocke_tensor_storage_t* p, uint64_t* bytes); bool rocke_tensor_storage_address( const rocke_tensor_storage_t* p, uint64_t row, uint64_t col, uint64_t* byte, int* shift); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md index 37d8d51b91f0..ecab80f8458e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md @@ -71,15 +71,22 @@ For example, dense FP6 uses four elements in three bytes; `slot_bits=8` describe one six-bit pattern in each byte instead. Both have the same logical dtype. Address calculation includes the bit offset, and allocation bounds include the partial tail byte. Empty views require no accesses. +Native size queries, address queries, and fragment loads revalidate the public +storage descriptor, including dtype/packing consistency, stride, alignment, and +allocation bounds, because C callers can construct or mutate these structs. [`storage_ir_type`](../../python/rocke/helpers/mma_io.py) chooses an addressable unit. FP4/FP6 use I8 units while retaining their logical identities in the storage descriptor. This function is deliberately separate from `dtype_to_ir_type`. +FP8/BF8 retain nominal FP8/BF8 pointer types despite occupying one byte per value. The fragment loader accepts complete, byte-aligned chunks. It derives load alignment from base alignment, row stride, chunk spacing, and origin. A 24-byte FP6 chunk becomes 16-byte and 8-byte loads without reading past the chunk. Partial fragments and nonzero bit origins are rejected by this initial loader. +Row strides must also be divisible by the pointer storage unit: an odd-byte +FP16/BF16 stride is valid generic storage but cannot be loaded through a typed +two-byte pointer. `row_base` is an element index in those storage units. The host bit packer writes a fresh, exclusively owned output buffer; concurrent packed stores are not provided. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py index 8d0bae29b5a2..38bcbe29e7e1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py @@ -43,7 +43,8 @@ def load_matrix_fragment( ) -> Value: """Load whole-byte chunks directly into the atom's carrier vector. - row_base counts pointer storage units. lane_group is derived from the atom's + row_base counts pointer storage units, which must divide the row stride. + lane_group is derived from the atom's lane mapping by the caller. Partial chunks and nonzero bit origins reject. """ packing = layout.fragment @@ -61,7 +62,11 @@ def load_matrix_fragment( if k0 < 0 or k0 + span > storage.shape[1]: raise ValueError("matrix fragment exceeds the packed row") origin_bits = storage.packing.bit_offset(k0) - if origin_bits % (8 * unit_bytes) or layout.chunk_bytes % unit_bytes: + if ( + origin_bits % (8 * unit_bytes) + or layout.chunk_bytes % unit_bytes + or storage.row_stride_bytes % unit_bytes + ): raise ValueError("matrix fragment is not aligned to pointer storage units") if dtype_info(carrier_type.name).encoded_bits != packing.carrier_bits: raise ValueError("matrix carrier type width mismatch") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py index 5ac95d34e755..190beabfe0e9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py @@ -37,6 +37,8 @@ "bf8", "f16", "bf16", + "f16_padded", + "bf16_padded", "pack_fp6_cross_word", "pack_scale_bytes_i32", "pack_scale_bytes_i64", @@ -45,8 +47,8 @@ def build_transport(dtype): b = IRBuilder("transport") - padded = dtype == "fp6_padded" - dtype = "fp6" if padded else dtype + padded = dtype.endswith("_padded") + dtype = dtype.removesuffix("_padded") patterns = dtype.startswith("pack_") unit = I8 if patterns else storage_ir_type(dtype) typed = dtype in ("f16", "bf16") @@ -69,7 +71,7 @@ def build_transport(dtype): storage = TensorStorage( dtype, (16, 128), - row_stride_bytes=97 if padded else None, + row_stride_bytes=(258 if typed else 97) if padded else None, alignment_bytes=16, ) layout = ( @@ -77,7 +79,7 @@ def build_transport(dtype): if typed else scaled_matrix_layout(dtype, 16) ) - base = b.const_i32(0) + base = b.const_i32((129 if typed else 97) if padded else 0) lane = b.mod(b.thread_id_x(), b.const_i32(32)) group = b.div(lane, b.const_i32(16)) value = load_matrix_fragment( @@ -157,6 +159,57 @@ def test_hip_declares_only_encountered_missing_vector_widths(): assert "using i8x24 =" not in source +@pytest.mark.parametrize("dtype", ["f16", "bf16"]) +def test_reject_row_stride_not_in_pointer_units(dtype): + b = IRBuilder("odd_stride") + unit = storage_ir_type(dtype) + ptr = b.param("A", PtrType(unit, "global")) + zero = b.const_i32(0) + before = serialize(b.kernel) + with pytest.raises(ValueError, match="aligned to pointer storage units"): + load_matrix_fragment( + b, + ptr, + zero, + zero, + 0, + storage=TensorStorage(dtype, (2, 64), row_stride_bytes=129), + layout=MatrixFragmentLayout( + FragmentPacking(BitPacking(16), 32, 16, 32), 16, 2, 16 + ), + carrier_type=unit, + ) + assert serialize(b.kernel) == before + + +@pytest.mark.parametrize("dtype", ["fp8", "bf8"]) +def test_byte_width_does_not_erase_pointer_identity(dtype): + assert storage_ir_type(dtype) != I8 + b = IRBuilder("nominal_pointer") + ptr = b.param("A", PtrType(I8, "global")) + zero = b.const_i32(0) + with pytest.raises(ValueError, match="pointer storage type mismatch"): + load_matrix_fragment( + b, + ptr, + zero, + zero, + 0, + storage=TensorStorage(dtype, (16, 128)), + layout=scaled_matrix_layout(dtype, 16), + ) + + +@pytest.mark.parametrize("elem_type", ["unknown", "fp4e2m1"]) +def test_hip_rejects_unsupported_vector_element(elem_type): + b = IRBuilder("invalid_element") + ptr = b.param("A", PtrType(I8, "global")) + value = b.global_load_vN(ptr, b.const_i32(0), I8, 16) + value.op.attrs["elem_type"] = elem_type + with pytest.raises(NotImplementedError, match="unsupported element type"): + lower_kernel_to_hip(b.kernel, arch="gfx1250") + + @pytest.mark.parametrize( "dtype,copy_bytes,alignment", [("fp6", 16, 8), ("fp6_padded", 16, 1), ("f16", 32, 16)], diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp index fd9b9ffbd7a3..cd210412cf9b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp @@ -33,9 +33,11 @@ static int emit(const char* dtype, bool hip) { rocke_ir_builder_t b; CHECK(rocke_ir_builder_init(&b, "transport") == ROCKE_OK); - const bool padded = strcmp(dtype, "fp6_padded") == 0; + const bool padded = strstr(dtype, "_padded") != NULL; if(padded) - dtype = "fp6"; + dtype = strcmp(dtype, "f16_padded") == 0 ? "f16" + : strcmp(dtype, "bf16_padded") == 0 ? "bf16" + : "fp6"; bool patterns = strncmp(dtype, "pack_", 5) == 0; const auto* unit = patterns ? rocke_i8() : rocke_storage_ir_type(dtype); CHECK(unit); @@ -64,7 +66,7 @@ static int emit(const char* dtype, bool hip) { rocke_tensor_storage_t storage; CHECK(rocke_tensor_storage_init( - &storage, dtype, 16, 128, padded ? 97 : UINT64_MAX, 0, 0, 16)); + &storage, dtype, 16, 128, padded ? (typed ? 258 : 97) : UINT64_MAX, 0, 0, 16)); rocke_matrix_fragment_layout_t layout; if(typed) { @@ -76,7 +78,7 @@ static int emit(const char* dtype, bool hip) } else layout = rocke_scaled_matrix_layout(dtype, 16); - auto* base = rocke_b_const_i32(&b, 0); + auto* base = rocke_b_const_i32(&b, padded ? (typed ? 129 : 97) : 0); auto* thread = rocke_b_thread_id_x(&b); auto* lane = rocke_b_mod(&b, thread, rocke_b_const_i32(&b, 32)); auto* group = rocke_b_div(&b, lane, rocke_b_const_i32(&b, 16)); @@ -112,12 +114,127 @@ static int emit(const char* dtype, bool hip) return 0; } +static int test_storage_validation() +{ + rocke_tensor_storage_t storage; + CHECK(rocke_tensor_storage_init(&storage, "fp6", 16, 128, 97, 0, 0, 16)); + const auto layout = rocke_scaled_matrix_layout("fp6", 16); + for(int defect = 0; defect < 10; ++defect) + { + auto invalid = storage; + switch(defect) + { + case 0: + invalid.dtype = NULL; + break; + case 1: + invalid.dtype = rocke_dtype_info("fp8"); + break; + case 2: + invalid.packing.element_bits = 0; + break; + case 3: + invalid.packing.slot_bits = 0; + break; + case 4: + invalid.packing.slot_bits = 4; + break; + case 5: + invalid.row_stride_bytes = 95; + break; + case 6: + invalid.alignment_bytes = 0; + break; + case 7: + invalid.alignment_bytes = 3; + break; + case 8: + invalid.rows = UINT64_MAX; + break; + case 9: + invalid.base_bit_offset = UINT64_MAX; + break; + } + uint64_t bytes; + int shift; + CHECK(!rocke_tensor_storage_bytes(&invalid, &bytes)); + CHECK(!rocke_tensor_storage_address(&invalid, 0, 0, &bytes, &shift)); + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "invalid_storage") == ROCKE_OK); + auto* ptr = rocke_b_param(&b, "A", rocke_ptr_type(&b, rocke_i8(), "global"), NULL); + auto* zero = rocke_b_const_i32(&b, 0); + CHECK( + !rocke_h_load_matrix_fragment(&b, ptr, zero, zero, 0, &invalid, &layout, rocke_i32())); + CHECK(!rocke_ir_builder_ok(&b)); + rocke_ir_builder_free(&b); + if(defect < 8) + { + invalid.rows = 0; + CHECK(!rocke_tensor_storage_bytes(&invalid, &bytes)); + } + } + for(const char* dtype : {"f16", "bf16", "fp8", "bf8"}) + { + const auto* unit = rocke_storage_ir_type(dtype); + CHECK(unit && !rocke_type_eq(unit, rocke_i8())); + const bool typed = strcmp(dtype, "f16") == 0 || strcmp(dtype, "bf16") == 0; + rocke_tensor_storage_t direct = {rocke_dtype_info(dtype), + 2, + 128, + typed ? 257u : 129u, + 0, + 1, + {typed ? 16 : 8, typed ? 16 : 8}}; + uint64_t bytes; + int shift; + CHECK(rocke_tensor_storage_bytes(&direct, &bytes)); + CHECK(rocke_tensor_storage_address(&direct, 1, 0, &bytes, &shift)); + CHECK(bytes == direct.row_stride_bytes && shift == 0); + rocke_matrix_fragment_layout_t fragment; + if(typed) + { + rocke_fragment_packing_t packing; + CHECK(rocke_fragment_packing_init(&packing, &direct.packing, 32, 16, 32)); + CHECK(rocke_matrix_fragment_layout_init(&fragment, &packing, 16, 2, 16)); + } + else + fragment = rocke_scaled_matrix_layout(dtype, 16); + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "invalid_pointer_units") == ROCKE_OK); + auto* ptr + = rocke_b_param(&b, "A", rocke_ptr_type(&b, typed ? unit : rocke_i8(), "global"), NULL); + auto* zero = rocke_b_const_i32(&b, 0); + CHECK(!rocke_h_load_matrix_fragment( + &b, ptr, zero, zero, 0, &direct, &fragment, typed ? unit : rocke_i32())); + CHECK(!rocke_ir_builder_ok(&b)); + rocke_ir_builder_free(&b); + } + return 0; +} + int main(int argc, char** argv) { if(argc == 3 && strcmp(argv[1], "--emit") == 0) return emit(argv[2], false); if(argc == 3 && strcmp(argv[1], "--hip") == 0) return emit(argv[2], true); + CHECK(test_storage_validation() == 0); + for(const char* elem_type : {"unknown", "fp4e2m1"}) + { + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "invalid_element") == ROCKE_OK); + auto* ptr = rocke_b_param(&b, "A", rocke_ptr_type(&b, rocke_i8(), "global"), NULL); + auto* value = rocke_b_global_load_vN(&b, ptr, rocke_b_const_i32(&b, 0), rocke_i8(), 16, 16); + CHECK(value); + rocke_attr_set_str(&b, &value->op->attrs, "elem_type", elem_type); + rocke_strbuf_t text; + CHECK(rocke_strbuf_init(&text, 256) == 0); + rocke_lower_hip_opts_t opts = {}; + opts.arch = "gfx1250"; + CHECK(rocke_lower_kernel_to_hip(&b, b.kernel, &opts, &text) == ROCKE_ERR_KEY); + rocke_strbuf_free(&text); + rocke_ir_builder_free(&b); + } for(int alignment : {0, -1, -16, 3, 24}) { rocke_ir_builder_t b; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py index 515430599388..6824621d8f9c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py @@ -80,6 +80,13 @@ def test_storage_rows_offsets_and_empty_views(): TensorStorage("fp4", (1, 4), alignment_bytes=3) +@pytest.mark.parametrize("dtype", ["f16", "bf16"]) +def test_storage_allows_byte_strides(dtype): + storage = TensorStorage(dtype, (2, 64), row_stride_bytes=129) + assert storage.address(1, 0) == (129, 0) + assert storage.byte_size == 257 + + def test_invalid_packing_and_overflow(): for bits, slots in [(0, None), (65, None), (6, 4)]: with pytest.raises(ValueError): From 40b9b825ab57695a79808a538eb68e7ec54aca8c Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:01:58 -0400 Subject: [PATCH 05/24] fix(rocke): unify E4M3 dtype and storage identity --- .../rocke/platform/cpp/core/dtypes.cpp | 27 ++++++++++++------ .../rocke/platform/cpp/core/ir/core_types.cpp | 5 +--- .../rocke/platform/cpp/helpers/mma_io.cpp | 3 +- .../rocke/platform/cpp/helpers/quant.cpp | 6 ++-- .../rocke/helper_rocke.helpers.quant.h | 2 +- .../rocke/platform/cpp/include/rocke/ir.h | 2 -- .../architecture/dtype_storage_packing.md | 12 ++++++-- .../platform/python/rocke/core/__init__.py | 2 -- .../platform/python/rocke/core/dtypes.py | 4 +-- .../rocke/platform/python/rocke/core/ir.py | 2 -- .../python/rocke/core/ir_serialize.py | 2 +- .../platform/python/rocke/helpers/mma_io.py | 2 +- .../platform/python/rocke/helpers/quant.py | 1 + .../rocke/platform/tests/core/test_dtypes.py | 28 ++++++++++++++++++- .../rocke/platform/tests/core/test_mma_io.py | 22 +++++++++++++-- .../platform/tests/core/test_storage.cpp | 21 +++++++++++++- 16 files changed, 105 insertions(+), 36 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/dtypes.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/dtypes.cpp index b5ab7d5507ae..dcc32275b741 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/dtypes.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/dtypes.cpp @@ -44,6 +44,7 @@ static const rocke_dtype_alias_t k_dtype_aliases[] = { {"int32", ROCKE_DTYPE_I32}, {"fp8", "fp8e4m3"}, {"fp8e4m3", "fp8e4m3"}, + {"e4m3", "fp8e4m3"}, {"bf8", "bf8e5m2"}, {"bf8e5m2", "bf8e5m2"}, {"fp6", "fp6e2m3"}, @@ -139,15 +140,23 @@ const char* rocke_normalize_dtype(const char* name, char* lowered, size_t cap) } static const rocke_dtype_info_t k_dtype_info[] = { - {"fp32", ROCKE_DTYPE_FLOAT, 32}, {"fp16", ROCKE_DTYPE_FLOAT, 16}, - {"bf16", ROCKE_DTYPE_FLOAT, 16}, {"fp8e4m3", ROCKE_DTYPE_FLOAT, 8}, - {"bf8e5m2", ROCKE_DTYPE_FLOAT, 8}, {"fp6e2m3", ROCKE_DTYPE_FLOAT, 6}, - {"fp6e3m2", ROCKE_DTYPE_FLOAT, 6}, {"fp4e2m1", ROCKE_DTYPE_FLOAT, 4}, - {"e8m0", ROCKE_DTYPE_FLOAT, 8}, {"e4m3", ROCKE_DTYPE_FLOAT, 8}, - {"e5m3", ROCKE_DTYPE_FLOAT, 8}, {"i4", ROCKE_DTYPE_INTEGER, 4}, - {"iu4", ROCKE_DTYPE_INTEGER, 4}, {"i8", ROCKE_DTYPE_INTEGER, 8}, - {"iu8", ROCKE_DTYPE_INTEGER, 8}, {"i16", ROCKE_DTYPE_INTEGER, 16}, - {"i32", ROCKE_DTYPE_INTEGER, 32}, {"i64", ROCKE_DTYPE_INTEGER, 64}, + {"fp32", ROCKE_DTYPE_FLOAT, 32}, + {"fp16", ROCKE_DTYPE_FLOAT, 16}, + {"bf16", ROCKE_DTYPE_FLOAT, 16}, + {"fp8e4m3", ROCKE_DTYPE_FLOAT, 8}, + {"bf8e5m2", ROCKE_DTYPE_FLOAT, 8}, + {"fp6e2m3", ROCKE_DTYPE_FLOAT, 6}, + {"fp6e3m2", ROCKE_DTYPE_FLOAT, 6}, + {"fp4e2m1", ROCKE_DTYPE_FLOAT, 4}, + {"e8m0", ROCKE_DTYPE_FLOAT, 8}, + {"e5m3", ROCKE_DTYPE_FLOAT, 8}, + {"i4", ROCKE_DTYPE_INTEGER, 4}, + {"iu4", ROCKE_DTYPE_INTEGER, 4}, + {"i8", ROCKE_DTYPE_INTEGER, 8}, + {"iu8", ROCKE_DTYPE_INTEGER, 8}, + {"i16", ROCKE_DTYPE_INTEGER, 16}, + {"i32", ROCKE_DTYPE_INTEGER, 32}, + {"i64", ROCKE_DTYPE_INTEGER, 64}, {"i1", ROCKE_DTYPE_PREDICATE, 1}, }; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp index d5e5bb48ee58..6062cadb30c0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp @@ -60,7 +60,6 @@ ROCKE_SCALAR_SINGLETON(rocke_fp4e2m1, ROCKE_SCALAR_FP4E2M1, "fp4e2m1") ROCKE_SCALAR_SINGLETON(rocke_fp6e2m3, ROCKE_SCALAR_FP6E2M3, "fp6e2m3") ROCKE_SCALAR_SINGLETON(rocke_fp6e3m2, ROCKE_SCALAR_FP6E3M2, "fp6e3m2") ROCKE_SCALAR_SINGLETON(rocke_e8m0, ROCKE_SCALAR_E8M0, "e8m0") -ROCKE_SCALAR_SINGLETON(rocke_e4m3, ROCKE_SCALAR_E4M3, "e4m3") ROCKE_SCALAR_SINGLETON(rocke_e5m3, ROCKE_SCALAR_E5M3, "e5m3") #undef ROCKE_SCALAR_SINGLETON @@ -85,7 +84,7 @@ const rocke_type_t* rocke_scalar_by_name(const char* name) return rocke_f16(); if(strcmp(name, "f32") == 0) return rocke_f32(); - if(strcmp(name, "fp8e4m3") == 0) + if(strcmp(name, "fp8e4m3") == 0 || strcmp(name, "e4m3") == 0) return rocke_fp8e4m3(); if(strcmp(name, "bf8e5m2") == 0) return rocke_bf8e5m2(); @@ -97,8 +96,6 @@ const rocke_type_t* rocke_scalar_by_name(const char* name) return rocke_fp6e3m2(); if(strcmp(name, "e8m0") == 0) return rocke_e8m0(); - if(strcmp(name, "e4m3") == 0) - return rocke_e4m3(); if(strcmp(name, "e5m3") == 0) return rocke_e5m3(); return NULL; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp index c7d01678a04e..521ccfa0c2fc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp @@ -13,8 +13,7 @@ const rocke_type_t* rocke_storage_ir_type(const char* dtype) const rocke_dtype_info_t* info = rocke_dtype_info(dtype); if(!info) return NULL; - if(info->encoded_bits % 8 || strcmp(info->name, "e8m0") == 0 || strcmp(info->name, "e4m3") == 0 - || strcmp(info->name, "e5m3") == 0) + if(info->encoded_bits % 8 || strcmp(info->name, "e8m0") == 0 || strcmp(info->name, "e5m3") == 0) return rocke_i8(); return rocke_dtype_to_ir_type(info->name); } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/quant.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/quant.cpp index fcb13a085464..89a7c90df38c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/quant.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/quant.cpp @@ -22,9 +22,9 @@ const rocke_type_t* rocke_quant_ir_type(const char* qdtype) { return rocke_dtype_to_ir_type("i8"); } - /* `_canon` -> "fp8e4m3" : keys "fp8e4m3", "fp8", "fp8_e4m3". */ + /* `_canon` -> "fp8e4m3" : keys "fp8e4m3", "fp8", "fp8_e4m3", "e4m3". */ if(strcmp(qdtype, "fp8e4m3") == 0 || strcmp(qdtype, "fp8") == 0 - || strcmp(qdtype, "fp8_e4m3") == 0) + || strcmp(qdtype, "fp8_e4m3") == 0 || strcmp(qdtype, "e4m3") == 0) { return rocke_dtype_to_ir_type("fp8e4m3"); } @@ -63,7 +63,7 @@ const rocke_type_t* rocke_b_quant_ir_type(rocke_ir_builder_t* b, const char* qdt b, ROCKE_ERR_VALUE, "unsupported quant dtype %s%s%s; expected one of " - "['bf8', 'bf8e5m2', 'fp8', 'fp8_e4m3', 'fp8_e5m2', 'fp8e4m3', " + "['bf8', 'bf8e5m2', 'e4m3', 'fp8', 'fp8_e4m3', 'fp8_e5m2', 'fp8e4m3', " "'i8', 'int8']", qdtype ? "'" : "", qdtype ? qdtype : "None", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.quant.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.quant.h index 2b24e6d2e172..12b21e1a2ee0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.quant.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.quant.h @@ -22,7 +22,7 @@ extern "C" { * * Accepts (matching Python `_QDTYPE_ALIAS`): * "i8", "int8" -> rocke_i8() - * "fp8e4m3", "fp8", "fp8_e4m3" -> rocke_fp8e4m3() + * "fp8e4m3", "fp8", "fp8_e4m3", "e4m3" -> rocke_fp8e4m3() * "bf8e5m2", "bf8", "fp8_e5m2" -> rocke_bf8e5m2() * "fp4", "fp4e2m1" -> rocke_fp4e2m1() * "fp6", "fp6e2m3" -> rocke_fp6e2m3() 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 131ea637e341..a6fcf74b8ab5 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 @@ -109,7 +109,6 @@ typedef enum rocke_scalar_kind ROCKE_SCALAR_FP6E2M3, ROCKE_SCALAR_FP6E3M2, ROCKE_SCALAR_E8M0, - ROCKE_SCALAR_E4M3, ROCKE_SCALAR_E5M3, ROCKE_SCALAR__COUNT } rocke_scalar_kind_t; @@ -584,7 +583,6 @@ const rocke_type_t* rocke_fp4e2m1(void); const rocke_type_t* rocke_fp6e2m3(void); const rocke_type_t* rocke_fp6e3m2(void); const rocke_type_t* rocke_e8m0(void); -const rocke_type_t* rocke_e4m3(void); const rocke_type_t* rocke_e5m3(void); /* Logical dtype resolver; NULL for unknown or unrepresented encodings. */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md index ecab80f8458e..0d68276de961 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md @@ -22,15 +22,21 @@ selectors. An integer carrier does not identify the numerical format. [`dtype_info`](../../python/rocke/core/dtypes.py) and [`rocke_dtype_info`](../../cpp/include/rocke/dtypes.h) return target-independent encoding descriptors. `dtype_to_ir_type` and `rocke_dtype_to_ir_type` resolve -logical types, including distinct FP4 E2M1, FP6 E2M3, FP6 E3M2, and scale types. +logical types, including distinct FP4 E2M1, FP6 E2M3, FP6 E3M2, E8M0, and E5M3. `quant_ir_type` delegates type resolution to this common resolver. Recognition and serialization do not enable scalar conversion, arithmetic, or -an instruction on a target. In this implementation, FP4/FP6 and the nominal scale -types are used in authoring metadata; storage helpers expand their transport into +an instruction on a target. In this implementation, FP4/FP6, E8M0, and E5M3 +are used in authoring metadata; storage helpers expand their transport into existing integer operations before serialization. Raw scalar lowering of these types remains unsupported. Existing FP16/BF16 and FP8/BF8 lowering is unchanged. +`e4m3` is an alias for the canonical `fp8e4m3` encoding, logical IR type, and +nominal FP8 storage pointer. Both spellings serialize as `fp8e4m3`. The +architecture scale selector `MmaScaleDType.E4M3` describes its operand role; it +does not introduce a second scalar type. Scale words can still be assembled +from raw byte patterns with the common bit packer. + Scale-format validation remains separate from matrix alias normalization: E5M3 is not BF8 E5M2. Registering a scale dtype does not enable a backend selector. 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 57596e6ba255..c6b12bbdb49e 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 @@ -55,7 +55,6 @@ FP6E2M3, FP6E3M2, E8M0, - E4M3, E5M3, dtype_to_ir_type, I1, @@ -111,7 +110,6 @@ "FP6E2M3", "FP6E3M2", "E8M0", - "E4M3", "E5M3", "CodegenPolicy", "SchedulerStrategy", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/dtypes.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/dtypes.py index 3808f87f3ad0..b0745c226610 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/dtypes.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/dtypes.py @@ -25,6 +25,7 @@ "fp32": "fp32", "fp8": "fp8e4m3", "fp8e4m3": "fp8e4m3", + "e4m3": "fp8e4m3", "bf8": "bf8e5m2", "bf8e5m2": "bf8e5m2", "fp6": "fp6e2m3", @@ -83,7 +84,6 @@ class DTypeInfo: ("fp6e3m2", 6), ("fp4e2m1", 4), ("e8m0", 8), - ("e4m3", 8), ("e5m3", 8), ), ), @@ -108,7 +108,7 @@ class DTypeInfo: def dtype_info(name: str) -> DTypeInfo: """Resolve an encoding; recognition alone does not enable an operation. - Scale names remain distinct keys. In particular e5m3 is not bf8e5m2. + E4M3 shares the fp8e4m3 encoding; e5m3 remains distinct from bf8e5m2. Operand-role validation belongs to the selected atom. """ key = normalize_dtype(name) 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 6dc0175c467d..73f01b156cfc 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 @@ -57,7 +57,6 @@ def __repr__(self) -> str: FP6E2M3 = Type("fp6e2m3") FP6E3M2 = Type("fp6e3m2") E8M0 = Type("e8m0") -E4M3 = Type("e4m3") E5M3 = Type("e5m3") @@ -82,7 +81,6 @@ def dtype_to_ir_type(dtype: str) -> Type: "fp6e2m3": FP6E2M3, "fp6e3m2": FP6E3M2, "e8m0": E8M0, - "e4m3": E4M3, "e5m3": E5M3, } if info.name not in types: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir_serialize.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir_serialize.py index 18c9c5cdbc5a..542a14d5c580 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir_serialize.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir_serialize.py @@ -136,7 +136,7 @@ def _parse_type(s: str) -> Type: shape_str = inner[lb + 1 : rb].strip() shape = tuple(int(x) for x in shape_str.split("x")) if shape_str else () return SmemType(elem, shape) - return Type(s) + return Type("fp8e4m3" if s == "e4m3" else s) # -------------------------------------------------------------------------- diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py index 38bcbe29e7e1..1e0745da422a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py @@ -25,7 +25,7 @@ def storage_ir_type(dtype: str) -> Type: """Addressable storage unit, distinct from the logical dtype resolver.""" info = dtype_info(dtype) - if info.encoded_bits % 8 or info.name in ("e8m0", "e4m3", "e5m3"): + if info.encoded_bits % 8 or info.name in ("e8m0", "e5m3"): return I8 return dtype_to_ir_type(info.name) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/quant.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/quant.py index 1ae76ffb0d33..07e7cc1158a9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/quant.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/quant.py @@ -90,6 +90,7 @@ "fp8e4m3": "fp8e4m3", "fp8": "fp8e4m3", "fp8_e4m3": "fp8e4m3", + "e4m3": "fp8e4m3", "bf8e5m2": "bf8e5m2", "bf8": "bf8e5m2", "fp8_e5m2": "bf8e5m2", 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 index cee48c279c6e..ca6de16b1fe8 100644 --- 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 @@ -6,7 +6,10 @@ from rocke.core.dtypes import normalize_dtype from rocke.core.dtypes import dtype_info -from rocke.core.ir import I8, IRBuilder, dtype_to_ir_type +from rocke.core.arch.target import MmaScaleDType +from rocke.core.ir import FP8E4M3, I8, IRBuilder, dtype_to_ir_type +from rocke.core.storage import TensorStorage +from rocke.helpers.mma_io import storage_ir_type from rocke.core.ir_serialize import parse, serialize from rocke.helpers.quant import quant_ir_type, ir_to_qdtype, dequantize_scalar_to_f32 @@ -63,6 +66,7 @@ def test_unknown_dtype_and_unrepresented_integer_family(): ("bfloat16", "bf16"), (" Float\t", "fp32"), ("FP8", "fp8e4m3"), + (" E4M3 ", "fp8e4m3"), ("BF8", "bf8e5m2"), ("FP6", "fp6e2m3"), ("BF6", "fp6e3m2"), @@ -84,3 +88,25 @@ def test_architecture_entry_points_reexport_core_normalization(): assert normalize_dtype.__module__ == "rocke.core.dtypes" assert core_normalize is arch_normalize is target_normalize is normalize_dtype + + +@pytest.mark.parametrize("spelling", ["e4m3", "fp8e4m3", "fp8", MmaScaleDType.E4M3]) +def test_e4m3_shares_encoding_type_and_storage(spelling): + assert dtype_info(spelling) is dtype_info("fp8e4m3") + assert dtype_to_ir_type(spelling) is FP8E4M3 + assert storage_ir_type(spelling) is FP8E4M3 + assert TensorStorage(spelling, (2, 128)) == TensorStorage("fp8e4m3", (2, 128)) + assert quant_ir_type(spelling) is FP8E4M3 + b = IRBuilder("e4m3_conversion") + value = b.param("value", dtype_to_ir_type(spelling)) + b.cvt_fp8_to_f32(value) + assert serialize(b.kernel) == serialize(parse(serialize(b.kernel))) + + +def test_scale_role_remains_independent_of_encoding_aliases(): + assert MmaScaleDType("fp8e4m3") is MmaScaleDType.E4M3 + assert str(MmaScaleDType.E4M3) == "e4m3" + assert dtype_info("e5m3") is not dtype_info("bf8e5m2") + assert dtype_to_ir_type("e5m3") != dtype_to_ir_type("bf8e5m2") + with pytest.raises(ValueError): + MmaScaleDType("bf8e5m2") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py index 190beabfe0e9..e1e1e3d74ece 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py @@ -11,8 +11,8 @@ import pytest from rocke.core.arch.wmma_scale import scaled_matrix_layout -from rocke.core.ir import I8, I32, I64, IRBuilder, PtrType -from rocke.core.ir_serialize import serialize +from rocke.core.ir import FP8E4M3, I8, I32, I64, IRBuilder, PtrType, VectorType +from rocke.core.ir_serialize import parse, serialize from rocke.core.lower_llvm import lower_kernel_to_llvm from rocke.core.lower_hip import lower_kernel_to_hip from rocke.core.storage import ( @@ -34,6 +34,7 @@ "fp6_padded", "bf6", "fp8", + "e4m3", "bf8", "f16", "bf16", @@ -244,3 +245,20 @@ def test_hip_rejects_invalid_vector_load_alignment(alignment): value.op.attrs["align"] = alignment with pytest.raises(ValueError, match="alignment must be a positive power of two"): lower_kernel_to_hip(b.kernel, arch="gfx1250") + + +def test_e4m3_alias_serialization_matches_native(): + executable = os.environ.get("ROCKE_STORAGE_TEST") + if not executable: + pytest.skip("set ROCKE_STORAGE_TEST to the built native storage test") + b = IRBuilder("canonical_e4m3") + b.param("scalar", FP8E4M3) + b.param("pointer", PtrType(FP8E4M3, "global")) + b.param("vector", VectorType(FP8E4M3, 4)) + canonical = serialize(b.kernel) + alias = canonical.replace("fp8e4m3", "e4m3") + assert serialize(parse(alias)) == canonical + result = subprocess.run( + [executable, "--parse", alias], check=True, capture_output=True, text=True + ) + assert result.stdout == canonical diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp index cd210412cf9b..0d4ce1fcf0f4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp @@ -173,7 +173,7 @@ static int test_storage_validation() CHECK(!rocke_tensor_storage_bytes(&invalid, &bytes)); } } - for(const char* dtype : {"f16", "bf16", "fp8", "bf8"}) + for(const char* dtype : {"f16", "bf16", "fp8", "e4m3", "bf8"}) { const auto* unit = rocke_storage_ir_type(dtype); CHECK(unit && !rocke_type_eq(unit, rocke_i8())); @@ -218,7 +218,26 @@ int main(int argc, char** argv) return emit(argv[2], false); if(argc == 3 && strcmp(argv[1], "--hip") == 0) return emit(argv[2], true); + if(argc == 3 && strcmp(argv[1], "--parse") == 0) + { + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "parse") == ROCKE_OK); + rocke_kernel_def_t* kernel = NULL; + CHECK(rocke_ir_parse(argv[2], &b, &kernel) == ROCKE_OK); + char* text = NULL; + CHECK(rocke_ir_serialize(kernel, &text) == ROCKE_OK); + fputs(text, stdout); + free(text); + rocke_ir_builder_free(&b); + return 0; + } CHECK(test_storage_validation() == 0); + CHECK(rocke_dtype_info("e4m3") == rocke_dtype_info("fp8e4m3")); + CHECK(rocke_dtype_to_ir_type("e4m3") == rocke_fp8e4m3()); + CHECK(rocke_storage_ir_type("e4m3") == rocke_fp8e4m3()); + CHECK(rocke_quant_ir_type("e4m3") == rocke_fp8e4m3()); + CHECK(rocke_scalar_by_name("e4m3") == rocke_fp8e4m3()); + CHECK(rocke_dtype_info("e5m3") != rocke_dtype_info("bf8e5m2")); for(const char* elem_type : {"unknown", "fp4e2m1"}) { rocke_ir_builder_t b; From d21f2833df3b17e72fc77ce8c12399b178f4ab8d Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Thu, 24 Sep 2026 16:03:26 -0400 Subject: [PATCH 06/24] fix(rocke): wire native parity into the standard runner Report complete quant type aliases while preserving scalar quantization restrictions. --- .../rocke/platform/cpp/helpers/quant.cpp | 8 +- .../platform/python/rocke/helpers/quant.py | 21 +++- .../rocke/platform/tests/README.md | 13 ++ .../rocke/platform/tests/core/test_dtypes.py | 43 ++++++- .../rocke/platform/tests/core/test_mma_io.py | 17 +++ .../platform/tests/core/test_storage.cpp | 18 +++ .../rocke/platform/tests/run_all.py | 75 ++++++++++- .../platform/tests/tools/test_run_all.py | 118 ++++++++++++++++++ 8 files changed, 298 insertions(+), 15 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/tools/test_run_all.py diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/quant.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/quant.cpp index 89a7c90df38c..e7ff6241752c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/quant.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/quant.cpp @@ -55,16 +55,16 @@ const rocke_type_t* rocke_b_quant_ir_type(rocke_ir_builder_t* b, const char* qdt ty = rocke_quant_ir_type(qdtype); if(ty == NULL) { - /* Mirror the Python `_canon` ValueError, including the {qdtype!r} + /* Mirror the Python `quant_ir_type` ValueError, including the {qdtype!r} * single-quote repr for the (non-NULL) string case and the sorted - * key list `sorted(_QDTYPE_ALIAS)`. NULL is reported as "None" to + * key list `sorted(_QUANT_TYPE_ALIAS)`. NULL is reported as "None" to * match Python's repr(None). */ return (const rocke_type_t*)rocke_i_set_err( b, ROCKE_ERR_VALUE, "unsupported quant dtype %s%s%s; expected one of " - "['bf8', 'bf8e5m2', 'e4m3', 'fp8', 'fp8_e4m3', 'fp8_e5m2', 'fp8e4m3', " - "'i8', 'int8']", + "['bf6', 'bf8', 'bf8e5m2', 'e4m3', 'fp4', 'fp4e2m1', 'fp6', 'fp6e2m3', " + "'fp6e3m2', 'fp8', 'fp8_e4m3', 'fp8_e5m2', 'fp8e4m3', 'i8', 'int8']", qdtype ? "'" : "", qdtype ? qdtype : "None", qdtype ? "'" : ""); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/quant.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/quant.py index 07e7cc1158a9..eb4604c48809 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/quant.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/quant.py @@ -97,6 +97,17 @@ } +# Type recognition includes formats without scalar quantization support. +_QUANT_TYPE_ALIAS = _QDTYPE_ALIAS | { + "fp4": "fp4e2m1", + "fp4e2m1": "fp4e2m1", + "fp6": "fp6e2m3", + "fp6e2m3": "fp6e2m3", + "bf6": "fp6e3m2", + "fp6e3m2": "fp6e3m2", +} + + _IR_TO_QDTYPE = { "i8": "i8", "fp8e4m3": "fp8e4m3", @@ -123,10 +134,12 @@ def quant_ir_type(qdtype: str) -> Type: the common dtype resolver. Type recognition does not enable scalar quantization or dequantization for the low-bit formats. """ - # Type recognition includes low-bit formats before scalar conversions exist. - if qdtype in ("fp4", "fp4e2m1", "fp6", "fp6e2m3", "bf6", "fp6e3m2"): - return dtype_to_ir_type(qdtype) - return dtype_to_ir_type(_canon(qdtype)) + if qdtype not in _QUANT_TYPE_ALIAS: + raise ValueError( + f"unsupported quant dtype {qdtype!r}; expected one of " + f"{sorted(_QUANT_TYPE_ALIAS)}" + ) + return dtype_to_ir_type(_QUANT_TYPE_ALIAS[qdtype]) def quant_max_abs(qdtype: str) -> float: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/README.md b/dnn-providers/hip-kernel-provider/rocke/platform/tests/README.md index 28c98405ffe1..38da77e7b588 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/README.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/README.md @@ -104,3 +104,16 @@ divergence - the Python builder correctly rejects wave32 WMMA on gfx942.) - EXCLUDED from rocKE: `test_gen_instances.py` (imports `ck4inductor`, a separate package) and `test_rocke_examples.py` (drives the external `example/ck_tile/dsl` tree, not part of rocKE) stay in `composablekernel/python/test`. + +### Native storage parity in the standard runner + +`run_all.py --build-root ` builds the `rocke_storage` target in a configured +build and obtains its executable path from CTest. Both pytest passes receive that +path, so storage IR/HIP parity and serialization tests run automatically. +`--config` selects the native test configuration (default `Release`). A build or +fixture-discovery failure stops the runner instead of silently skipping coverage. + +An explicit `ROCKE_STORAGE_TEST` overrides discovery and must name an existing +executable. With no configured build or explicit override, the runner reports +native storage parity as skipped; direct pytest invocations can use the same +override. 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 index ca6de16b1fe8..2222d6ed4727 100644 --- 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 @@ -2,6 +2,8 @@ # SPDX-License-Identifier: MIT """Core dtype names are independent of target support and retain public aliases.""" +import ast + import pytest from rocke.core.dtypes import normalize_dtype @@ -11,7 +13,12 @@ from rocke.core.storage import TensorStorage from rocke.helpers.mma_io import storage_ir_type from rocke.core.ir_serialize import parse, serialize -from rocke.helpers.quant import quant_ir_type, ir_to_qdtype, dequantize_scalar_to_f32 +from rocke.helpers.quant import ( + quant_ir_type, + ir_to_qdtype, + dequantize_scalar_to_f32, + quant_max_abs, +) @pytest.mark.parametrize( @@ -110,3 +117,37 @@ def test_scale_role_remains_independent_of_encoding_aliases(): assert dtype_to_ir_type("e5m3") != dtype_to_ir_type("bf8e5m2") with pytest.raises(ValueError): MmaScaleDType("bf8e5m2") + + +def test_quant_type_diagnostic_lists_every_accepted_spelling(): + expected = { + "i8", + "int8", + "fp8", + "fp8e4m3", + "fp8_e4m3", + "e4m3", + "bf8", + "bf8e5m2", + "fp8_e5m2", + "fp4", + "fp4e2m1", + "fp6", + "fp6e2m3", + "bf6", + "fp6e3m2", + } + with pytest.raises(ValueError) as error: + quant_ir_type("invalid") + accepted = ast.literal_eval(str(error.value).split("expected one of ")[1]) + assert accepted == sorted(expected) + assert all(quant_ir_type(name) is not None for name in accepted) + + +@pytest.mark.parametrize( + "dtype", ["fp4", "fp4e2m1", "fp6", "fp6e2m3", "bf6", "fp6e3m2"] +) +def test_type_alias_does_not_enable_scalar_quantization(dtype): + assert quant_ir_type(dtype) is not None + with pytest.raises(ValueError, match="unsupported quant dtype"): + quant_max_abs(dtype) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py index e1e1e3d74ece..558b3ff0ddae 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py @@ -262,3 +262,20 @@ def test_e4m3_alias_serialization_matches_native(): [executable, "--parse", alias], check=True, capture_output=True, text=True ) assert result.stdout == canonical + + +def test_quant_type_diagnostic_matches_native(): + from rocke.helpers.quant import quant_ir_type + + executable = os.environ.get("ROCKE_STORAGE_TEST") + if not executable: + pytest.skip("set ROCKE_STORAGE_TEST to the built native storage test") + with pytest.raises(ValueError) as error: + quant_ir_type("invalid") + result = subprocess.run( + [executable, "--quant-error", "invalid"], + check=True, + capture_output=True, + text=True, + ) + assert result.stdout == str(error.value) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp index 0d4ce1fcf0f4..edc54e0e51d7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp @@ -8,6 +8,7 @@ #include #include +#include "rocke/error.hpp" #include "rocke/helper_rocke.helpers.mma_io.h" #include "rocke/helper_rocke.helpers.quant.h" #include "rocke/ir_serialize.h" @@ -231,6 +232,23 @@ int main(int argc, char** argv) rocke_ir_builder_free(&b); return 0; } + if(argc == 3 && strcmp(argv[1], "--quant-error") == 0) + { + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "quant_error") == ROCKE_OK); + try + { + rocke_b_quant_ir_type(&b, argv[2]); + CHECK(false); + } + catch(const ckc::Error& error) + { + CHECK(error.code() == ROCKE_ERR_VALUE); + fputs(error.what(), stdout); + } + rocke_ir_builder_free(&b); + return 0; + } CHECK(test_storage_validation() == 0); CHECK(rocke_dtype_info("e4m3") == rocke_dtype_info("fp8e4m3")); CHECK(rocke_dtype_to_ir_type("e4m3") == rocke_fp8e4m3()); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/run_all.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/run_all.py index a888b1b7ffb8..7f60753b8175 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/run_all.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/run_all.py @@ -11,11 +11,12 @@ # # Usage: # python rocke/platform/tests/run_all.py [--no-guard] [--no-gate] [--no-pytest] -# [--no-both] [--only SUBSTR] [--build-root DIR] +# [--no-both] [--only SUBSTR] [--build-root DIR] [--config CONFIG] from __future__ import annotations import argparse +import json import os import re import subprocess @@ -86,7 +87,59 @@ def relative_path_guard() -> int: return 0 -def differential_pytest_pass() -> int: +def native_pytest_env(build_root: Path, config: str) -> dict[str, str]: + """Make the native storage comparisons available to both pytest passes.""" + env = dict(os.environ) + executable = env.get("ROCKE_STORAGE_TEST") + if not executable and (build_root / "CMakeCache.txt").is_file(): + subprocess.run( + [ + "cmake", + "--build", + str(build_root), + "--config", + config, + "--target", + "rocke_storage", + ], + check=True, + ) + listing = subprocess.run( + [ + "ctest", + "--test-dir", + str(build_root), + "-C", + config, + "--show-only=json-v1", + "-R", + "^rocke_storage$", + ], + check=True, + capture_output=True, + text=True, + ) + tests = json.loads(listing.stdout)["tests"] + commands = [ + test.get("command", []) for test in tests if test["name"] == "rocke_storage" + ] + if len(commands) != 1 or not commands[0]: + raise ValueError("CTest did not resolve the built rocke_storage executable") + executable = commands[0][0] + if executable: + path = Path(executable).resolve() + if not path.is_file(): + raise ValueError(f"native storage test executable does not exist: {path}") + env["ROCKE_STORAGE_TEST"] = str(path) + print(f"\n== native storage parity: {path} ==") + else: + print( + "\n== native storage parity: SKIPPED (no configured build or ROCKE_STORAGE_TEST) ==" + ) + return env + + +def differential_pytest_pass(env: dict[str, str]) -> int: """Re-run pytest with ``ROCKE_BACKEND=both`` (the cross-engine gate). The default pass exercises one engine per assertion, so two engines that @@ -110,6 +163,7 @@ def differential_pytest_pass() -> int: [sys.executable, "-c", "import rocke_engine"], capture_output=True, cwd=str(TESTS), + env=env, ) if probe.returncode != 0: print( @@ -133,7 +187,7 @@ def differential_pytest_pass() -> int: "either compared byte-for-byte or fails; remaining skips are " "environmental (torch / GPU)" ) - env = dict(os.environ, ROCKE_BACKEND="both") + env = dict(env, ROCKE_BACKEND="both") return subprocess.run( [sys.executable, "-m", "pytest", str(TESTS), "-rs"], cwd=str(TESTS), env=env ).returncode @@ -157,6 +211,9 @@ def main() -> int: ap.add_argument( "--build-root", default=str(Path(tempfile.gettempdir()) / "rocke_verify") ) + ap.add_argument( + "--config", default="Release", help="native test build/CTest configuration" + ) args = ap.parse_args() status = 0 @@ -177,13 +234,18 @@ def main() -> int: status |= subprocess.run(gate).returncode if not args.no_pytest: + try: + pytest_env = native_pytest_env(Path(args.build_root).resolve(), args.config) + except (OSError, ValueError, subprocess.CalledProcessError) as exc: + print(f"native storage parity setup failed: {exc}", file=sys.stderr) + return 1 print("\n== pytest ==") status |= subprocess.run( - [sys.executable, "-m", "pytest", str(TESTS)], cwd=str(TESTS) + [sys.executable, "-m", "pytest", str(TESTS)], cwd=str(TESTS), env=pytest_env ).returncode if not args.no_pytest and not args.no_both: - status |= differential_pytest_pass() + status |= differential_pytest_pass(pytest_env) build_root = Path(args.build_root) # Only ctest when the CTest-registered binaries were actually built (the @@ -200,7 +262,8 @@ def main() -> int: ): print("\n== ctest ==") status |= subprocess.run( - ["ctest", "--output-on-failure", "--no-tests=ignore"], cwd=str(build_root) + ["ctest", "-C", args.config, "--output-on-failure", "--no-tests=ignore"], + cwd=str(build_root), ).returncode print("\nRESULT:", "GREEN" if status == 0 else "RED") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/tools/test_run_all.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/tools/test_run_all.py new file mode 100644 index 000000000000..28e06b7b5c46 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/tools/test_run_all.py @@ -0,0 +1,118 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""The standard runner must supply native fixture coverage to both pytest passes.""" + +import importlib.util +import json +from pathlib import Path +import subprocess + +import pytest + + +@pytest.fixture +def runner(monkeypatch): + path = Path(__file__).resolve().parents[1] / "run_all.py" + spec = importlib.util.spec_from_file_location("rocke_run_all", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + monkeypatch.delenv("ROCKE_STORAGE_TEST", raising=False) + monkeypatch.delenv("ROCKE_BACKEND", raising=False) + return module + + +@pytest.mark.parametrize( + "filename", ["rocke_storage", "provider_rocke_storage_test.exe"] +) +def test_runner_passes_registered_fixture_to_both_backends( + runner, monkeypatch, tmp_path, filename +): + (tmp_path / "CMakeCache.txt").touch() + executable = tmp_path / "Debug" / filename + executable.parent.mkdir() + executable.touch() + calls = [] + + def run(command, **kwargs): + calls.append((command, kwargs)) + listing = {"tests": [{"name": "rocke_storage", "command": [str(executable)]}]} + return subprocess.CompletedProcess(command, 0, stdout=json.dumps(listing)) + + monkeypatch.setattr(runner.subprocess, "run", run) + monkeypatch.setattr( + runner.sys, + "argv", + [ + "run_all.py", + "--no-guard", + "--no-gate", + "--build-root", + str(tmp_path), + "--config", + "Debug", + ], + ) + assert runner.main() == 0 + children = [ + (cmd, kw["env"]) for cmd, kw in calls if cmd[0] == runner.sys.executable + ] + assert len(children) == 3 # default pytest, extension import probe, both pytest + assert all(env["ROCKE_STORAGE_TEST"] == str(executable) for _, env in children) + assert "ROCKE_BACKEND" not in children[0][1] + assert children[-1][1]["ROCKE_BACKEND"] == "both" + assert calls[0][0] == [ + "cmake", + "--build", + str(tmp_path), + "--config", + "Debug", + "--target", + "rocke_storage", + ] + + +def test_explicit_fixture_override_is_resolved_before_pytest_changes_directory( + runner, monkeypatch, tmp_path +): + executable = tmp_path / "fixture" + executable.touch() + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("ROCKE_STORAGE_TEST", "fixture") + env = runner.native_pytest_env(tmp_path, "Release") + assert env["ROCKE_STORAGE_TEST"] == str(executable) + + +def test_missing_build_reports_native_coverage_skip(runner, tmp_path, capsys): + assert "ROCKE_STORAGE_TEST" not in runner.native_pytest_env(tmp_path, "Release") + assert "native storage parity: SKIPPED" in capsys.readouterr().out + + +def test_invalid_override_is_an_error(runner, monkeypatch, tmp_path): + monkeypatch.setenv("ROCKE_STORAGE_TEST", str(tmp_path / "missing")) + with pytest.raises(ValueError, match="does not exist"): + runner.native_pytest_env(tmp_path, "Release") + + +@pytest.mark.parametrize("failure", ["build", "registration", "executable"]) +def test_native_setup_failure_prevents_silently_skipped_pytest( + runner, monkeypatch, tmp_path, failure +): + (tmp_path / "CMakeCache.txt").touch() + calls = [] + + def run(command, **kwargs): + calls.append(command) + if command[0] == "cmake" and failure == "build": + raise subprocess.CalledProcessError(1, command) + command_path = [str(tmp_path / "missing")] if failure == "executable" else [] + listing = {"tests": [{"name": "rocke_storage", "command": command_path}]} + return subprocess.CompletedProcess(command, 0, stdout=json.dumps(listing)) + + monkeypatch.setattr(runner.subprocess, "run", run) + monkeypatch.setattr( + runner.sys, + "argv", + ["run_all.py", "--no-guard", "--no-gate", "--build-root", str(tmp_path)], + ) + assert runner.main() == 1 + assert all(command[0] != runner.sys.executable for command in calls) From 1956a3402e58f2aebb2ee06ffd4d35ec7dc5bc1d Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:54:50 -0400 Subject: [PATCH 07/24] fix(rocke): reject empty storage in fragment loaders Keep empty descriptors valid while rejecting fragment loads before IR emission in both engines. Cover empty shapes and nonempty controls in Python and native tests. --- .../rocke/platform/cpp/helpers/mma_io.cpp | 2 + .../platform/python/rocke/helpers/mma_io.py | 2 + .../rocke/platform/tests/core/test_mma_io.py | 39 +++++++++++++++ .../platform/tests/core/test_storage.cpp | 49 +++++++++++++++++++ 4 files changed, 92 insertions(+) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp index 521ccfa0c2fc..d34f32a589b9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp @@ -47,6 +47,8 @@ rocke_value_t* rocke_h_load_matrix_fragment(rocke_ir_builder_t* b, uint64_t storage_bytes; if(!rocke_tensor_storage_bytes(storage, &storage_bytes)) ckc::raise_status(ROCKE_ERR_VALUE, "invalid tensor storage descriptor"); + if(storage_bytes == 0) + ckc::raise_status(ROCKE_ERR_VALUE, "matrix fragment requires nonempty tensor storage"); const auto& packing = layout->fragment; rocke_matrix_fragment_layout_t checked; if(!rocke_matrix_fragment_layout_init(&checked, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py index 1e0745da422a..15e7893550cc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py @@ -47,6 +47,8 @@ def load_matrix_fragment( lane_group is derived from the atom's lane mapping by the caller. Partial chunks and nonzero bit origins reject. """ + if storage.byte_size == 0: + raise ValueError("matrix fragment requires nonempty tensor storage") packing = layout.fragment if not (0 < packing.count <= 0x7FFFFFFF and packing.carrier_count <= 0x7FFFFFFF): raise ValueError("invalid matrix fragment chunk layout") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py index 558b3ff0ddae..caa74ce0cc67 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py @@ -150,6 +150,45 @@ def test_reject_partial_fragment_and_wrong_pointer(): ) +@pytest.mark.parametrize("dtype", ["fp4", "fp6", "fp8", "f16", "bf16"]) +@pytest.mark.parametrize("shape", [(0, 128), (16, 0), (0, 0), (16, 128)]) +def test_fragment_load_requires_nonempty_storage(dtype, shape): + storage = TensorStorage(dtype, shape) + typed = dtype in ("f16", "bf16") + unit = storage_ir_type(dtype) + layout = ( + MatrixFragmentLayout(FragmentPacking(BitPacking(16), 32, 16, 32), 16, 2, 16) + if typed + else scaled_matrix_layout(dtype, 16) + ) + b = IRBuilder("empty_storage") + ptr = b.param("A", PtrType(unit, "global")) + zero = b.const_i32(0) + before = serialize(b.kernel) + + def load(): + return load_matrix_fragment( + b, + ptr, + zero, + zero, + 0, + storage=storage, + layout=layout, + carrier_type=unit if typed else I32, + ) + + if 0 in shape: + assert storage.byte_size == 0 + with pytest.raises(ValueError, match="nonempty tensor storage"): + load() + assert serialize(b.kernel) == before + else: + assert storage.byte_size > 0 + assert load() is not None + assert "memref.global_load" in serialize(b.kernel) + + def test_hip_declares_only_encountered_missing_vector_widths(): source = lower_kernel_to_hip(build_transport("fp6"), arch="gfx1250") for name in ("i8x24", "i8x40", "i32x12"): diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp index edc54e0e51d7..4114583b8320 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp @@ -213,6 +213,54 @@ static int test_storage_validation() return 0; } +static int test_empty_storage_load() +{ + for(const char* dtype : {"fp4", "fp6", "fp8", "f16", "bf16"}) + { + const auto* unit = rocke_storage_ir_type(dtype); + const bool typed = strcmp(dtype, "f16") == 0 || strcmp(dtype, "bf16") == 0; + for(int shape = 0; shape < 4; ++shape) + { + const uint64_t rows = shape == 0 || shape == 2 ? 0 : 16; + const uint64_t cols = shape == 1 || shape == 2 ? 0 : 128; + rocke_tensor_storage_t storage; + CHECK(rocke_tensor_storage_init(&storage, dtype, rows, cols, UINT64_MAX, 0, 0, 1)); + uint64_t bytes; + CHECK(rocke_tensor_storage_bytes(&storage, &bytes)); + CHECK((bytes == 0) == (shape != 3)); + rocke_matrix_fragment_layout_t layout; + if(typed) + { + rocke_fragment_packing_t fragment; + CHECK(rocke_fragment_packing_init(&fragment, &storage.packing, 32, 16, 32)); + CHECK(rocke_matrix_fragment_layout_init(&layout, &fragment, 16, 2, 16)); + } + else + layout = rocke_scaled_matrix_layout(dtype, 16); + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "empty_storage") == ROCKE_OK); + auto* ptr = rocke_b_param(&b, "A", rocke_ptr_type(&b, unit, "global"), NULL); + auto* zero = rocke_b_const_i32(&b, 0); + const int before = b.kernel->body->num_ops; + auto* value = rocke_h_load_matrix_fragment( + &b, ptr, zero, zero, 0, &storage, &layout, typed ? unit : rocke_i32()); + if(bytes == 0) + { + CHECK(!value && rocke_ir_builder_status(&b) == ROCKE_ERR_VALUE); + CHECK(strstr(rocke_ir_builder_error(&b), "nonempty tensor storage")); + CHECK(b.kernel->body->num_ops == before); + } + else + { + CHECK(value && rocke_ir_builder_ok(&b)); + CHECK(b.kernel->body->num_ops > before); + } + rocke_ir_builder_free(&b); + } + } + return 0; +} + int main(int argc, char** argv) { if(argc == 3 && strcmp(argv[1], "--emit") == 0) @@ -250,6 +298,7 @@ int main(int argc, char** argv) return 0; } CHECK(test_storage_validation() == 0); + CHECK(test_empty_storage_load() == 0); CHECK(rocke_dtype_info("e4m3") == rocke_dtype_info("fp8e4m3")); CHECK(rocke_dtype_to_ir_type("e4m3") == rocke_fp8e4m3()); CHECK(rocke_storage_ir_type("e4m3") == rocke_fp8e4m3()); From 699e6b8b34b7e09971f4c28d37970c2d302c7c27 Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:47:22 -0400 Subject: [PATCH 08/24] fix(rocke): bound fragment offsets and discover configured CTests Validate complete static fragment displacements before emitting i32 offsets in both engines. Resolve CTest readiness from registered commands and cover multi-config execution and failure propagation. --- .../rocke/platform/cpp/helpers/mma_io.cpp | 7 +- .../platform/python/rocke/helpers/mma_io.py | 3 + .../rocke/platform/tests/README.md | 8 ++- .../rocke/platform/tests/core/test_mma_io.py | 41 +++++++++++ .../platform/tests/core/test_storage.cpp | 56 +++++++++++++++ .../rocke/platform/tests/run_all.py | 65 +++++++++-------- .../platform/tests/tools/test_run_all.py | 72 +++++++++++++++++++ 7 files changed, 220 insertions(+), 32 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp index d34f32a589b9..c4baddbb2557 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp @@ -91,7 +91,12 @@ rocke_value_t* rocke_h_load_matrix_fragment(rocke_ir_builder_t* b, "padded matrix fragments currently require i32 carriers"); const uint64_t chunk_units = chunk_bytes / unit_bytes; const uint64_t origin_bytes = origin_bits / 8; - if(origin_bytes / unit_bytes > INT_MAX || chunk_units * layout->lane_groups > INT_MAX) + // Bound the last loaded unit without overflowing intermediate products. + if(origin_bytes / unit_bytes > INT_MAX + || chunk_units > uint64_t(INT_MAX) / layout->lane_groups + || packing.count / layout->chunk_elements + > (uint64_t(INT_MAX) + 1 - origin_bytes / unit_bytes) / chunk_units + / layout->lane_groups) ckc::raise_status(ROCKE_ERR_VALUE, "matrix fragment offset exceeds i32 range"); int alignment = gcd(gcd(storage->alignment_bytes, storage->row_stride_bytes), gcd(chunk_bytes, origin_bytes)); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py index 15e7893550cc..b47f120a0beb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py @@ -76,9 +76,12 @@ def load_matrix_fragment( raise ValueError("matrix fragment payload must occupy whole carriers") chunk_units = layout.chunk_bytes // unit_bytes origin_bytes = origin_bits // 8 + # Bound the last loaded unit across all chunks and lane groups, including k0. if ( origin_bytes // unit_bytes > 0x7FFFFFFF or chunk_units * layout.lane_groups > 0x7FFFFFFF + or layout.chunks_per_lane + > (0x80000000 - origin_bytes // unit_bytes) // chunk_units // layout.lane_groups ): raise ValueError("matrix fragment offset exceeds i32 range") alignment = gcd( diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/README.md b/dnn-providers/hip-kernel-provider/rocke/platform/tests/README.md index 38da77e7b588..870e008c00f7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/README.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/README.md @@ -45,9 +45,11 @@ Four distinct things run here; don't conflate them: 1. **Default runner** (`python tests/run_all.py`): relative-path guard -> byte-identity gate (`tools/check_byte_identity.py`) -> `pytest` (the `test_*.py` modules - above) -> `ctest` **iff** the registered binaries (`rocke_ir_serialize_roundtrip`, - `rocke_tiled_attention_2d_reentrancy`) are present in `--build-root`. This is the - only set that is gated. + above) -> `ctest` when any registered test executable is built in `--build-root` + for the selected `--config`. CTest supplies the executable paths, including + configuration directories and platform suffixes. The entire registered suite + runs, so a partial build exposes missing tests as failures. No configured test + build or no built registered tests is reported explicitly as a skipped stage. 2. **Diagnostics** (opt-in, not in the gate): `run_diff.py --ir` (IR-canonical diff), `fuzz_diff.py`, `ir_artifact_diff.py`. 3. **GPU / manual numeric lanes** (need a HIP device; skipped/not-collected diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py index caa74ce0cc67..23515da8a3cd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py @@ -189,6 +189,47 @@ def load(): assert "memref.global_load" in serialize(b.kernel) +@pytest.mark.parametrize("dtype", ["e8m0", "fp6", "f16"]) +@pytest.mark.parametrize( + "count,groups,k0,valid", + [ + (64, 100_000_000, 0, False), + (64, 1 << 25, 0, True), + (64, 1 << 25, 1, False), + (64, 2, (1 << 31) - 128, True), + (64, 2, (1 << 31) - 127, False), + (16, 100_000_000, 1_500_000_000, False), + (64, 0x7FFFFFFF, 0, False), + ], +) +def test_fragment_offset_bounds(dtype, count, groups, k0, valid): + # Padded FP6 and E8M0 use byte units; FP16 uses two-byte pointer units. + bits = 16 if dtype == "f16" else 6 if dtype == "fp6" else 8 + packing = BitPacking(bits, 16 if dtype == "f16" else 8) + storage = TensorStorage( + dtype, (1, k0 + count * groups), slot_bits=packing.slot_bits + ) + layout = MatrixFragmentLayout( + FragmentPacking(packing, count, 32, count * packing.slot_bits // 32), + 16, + groups, + 1, + ) + b = IRBuilder("offset_bounds") + ptr = b.param("A", PtrType(storage_ir_type(dtype), "global")) + zero = b.const_i32(0) + before = serialize(b.kernel) + if valid: + assert ( + load_matrix_fragment(b, ptr, zero, zero, k0, storage=storage, layout=layout) + is not None + ) + else: + with pytest.raises(ValueError, match="offset exceeds i32 range"): + load_matrix_fragment(b, ptr, zero, zero, k0, storage=storage, layout=layout) + assert serialize(b.kernel) == before + + def test_hip_declares_only_encountered_missing_vector_widths(): source = lower_kernel_to_hip(build_transport("fp6"), arch="gfx1250") for name in ("i8x24", "i8x40", "i32x12"): diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp index 4114583b8320..981597c52560 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp @@ -261,6 +261,61 @@ static int test_empty_storage_load() return 0; } +static int test_fragment_offset_bounds() +{ + const struct + { + int count, groups; + uint64_t k0; + bool valid; + } cases[] = {{64, 100000000, 0, false}, + {64, 1 << 25, 0, true}, + {64, 1 << 25, 1, false}, + {64, 2, (uint64_t(1) << 31) - 128, true}, + {64, 2, (uint64_t(1) << 31) - 127, false}, + {16, 100000000, 1500000000, false}, + {64, INT32_MAX, 0, false}}; + for(const char* dtype : {"e8m0", "fp6", "f16"}) + { + const int slot_bits = strcmp(dtype, "f16") == 0 ? 16 : 8; + for(const auto& c : cases) + { + rocke_tensor_storage_t storage; + CHECK(rocke_tensor_storage_init(&storage, + dtype, + 1, + c.k0 + uint64_t(c.count) * c.groups, + UINT64_MAX, + slot_bits, + 0, + 1)); + rocke_fragment_packing_t fragment; + CHECK(rocke_fragment_packing_init( + &fragment, &storage.packing, c.count, 32, c.count * slot_bits / 32)); + rocke_matrix_fragment_layout_t layout; + CHECK(rocke_matrix_fragment_layout_init(&layout, &fragment, 16, c.groups, 1)); + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "offset_bounds") == ROCKE_OK); + auto* ptr = rocke_b_param( + &b, "A", rocke_ptr_type(&b, rocke_storage_ir_type(dtype), "global"), NULL); + auto* zero = rocke_b_const_i32(&b, 0); + const int before = b.kernel->body->num_ops; + auto* value = rocke_h_load_matrix_fragment( + &b, ptr, zero, zero, c.k0, &storage, &layout, rocke_i32()); + if(c.valid) + CHECK(value && rocke_ir_builder_ok(&b)); + else + { + CHECK(!value && rocke_ir_builder_status(&b) == ROCKE_ERR_VALUE); + CHECK(strstr(rocke_ir_builder_error(&b), "offset exceeds i32 range")); + CHECK(b.kernel->body->num_ops == before); + } + rocke_ir_builder_free(&b); + } + } + return 0; +} + int main(int argc, char** argv) { if(argc == 3 && strcmp(argv[1], "--emit") == 0) @@ -299,6 +354,7 @@ int main(int argc, char** argv) } CHECK(test_storage_validation() == 0); CHECK(test_empty_storage_load() == 0); + CHECK(test_fragment_offset_bounds() == 0); CHECK(rocke_dtype_info("e4m3") == rocke_dtype_info("fp8e4m3")); CHECK(rocke_dtype_to_ir_type("e4m3") == rocke_fp8e4m3()); CHECK(rocke_storage_ir_type("e4m3") == rocke_fp8e4m3()); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/run_all.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/run_all.py index 7f60753b8175..7f4b90c6124e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/run_all.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/run_all.py @@ -87,6 +87,34 @@ def relative_path_guard() -> int: return 0 +def ctest_tests( + build_root: Path, config: str, pattern: str | None = None +) -> list[dict]: + command = [ + "ctest", + "--test-dir", + str(build_root), + "-C", + config, + "--show-only=json-v1", + ] + if pattern: + command += ["-R", pattern] + listing = subprocess.run(command, check=True, capture_output=True, text=True) + return json.loads(listing.stdout)["tests"] + + +def ctest_ready(build_root: Path, config: str) -> bool: + if not (build_root / "CTestTestfile.cmake").is_file(): + print("\n== ctest: SKIPPED (no configured test build) ==") + return False + commands = [test.get("command", []) for test in ctest_tests(build_root, config)] + if any(command and Path(command[0]).is_file() for command in commands): + return True + print(f"\n== ctest: SKIPPED (no built registered tests for {config}) ==") + return False + + def native_pytest_env(build_root: Path, config: str) -> dict[str, str]: """Make the native storage comparisons available to both pytest passes.""" env = dict(os.environ) @@ -104,22 +132,7 @@ def native_pytest_env(build_root: Path, config: str) -> dict[str, str]: ], check=True, ) - listing = subprocess.run( - [ - "ctest", - "--test-dir", - str(build_root), - "-C", - config, - "--show-only=json-v1", - "-R", - "^rocke_storage$", - ], - check=True, - capture_output=True, - text=True, - ) - tests = json.loads(listing.stdout)["tests"] + tests = ctest_tests(build_root, config, "^rocke_storage$") commands = [ test.get("command", []) for test in tests if test["name"] == "rocke_storage" ] @@ -248,18 +261,14 @@ def main() -> int: status |= differential_pytest_pass(pytest_env) build_root = Path(args.build_root) - # Only ctest when the CTest-registered binaries were actually built (the - # byte-identity gate builds just `rocke_core`, so a gate-only build dir has the - # registration file but no test executables -> running ctest there would - # spuriously fail). Gate on the registered tests only; `rocke_smoke` is an - # optional build-only target (not an add_test target) so it is not a signal. - test_bins = [ - build_root / "tests" / b - for b in ("rocke_ir_serialize_roundtrip", "rocke_tiled_attention_2d_reentrancy") - ] - if (build_root / "CTestTestfile.cmake").exists() and any( - b.exists() for b in test_bins - ): + # Run the entire registered suite once any executable is built; partial + # builds must expose their missing tests rather than silently lose coverage. + try: + ready = ctest_ready(build_root, args.config) + except (OSError, ValueError, subprocess.CalledProcessError) as exc: + print(f"CTest discovery failed: {exc}", file=sys.stderr) + return 1 + if ready: print("\n== ctest ==") status |= subprocess.run( ["ctest", "-C", args.config, "--output-on-failure", "--no-tests=ignore"], diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/tools/test_run_all.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/tools/test_run_all.py index 28e06b7b5c46..d9dbd4f08bdd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/tools/test_run_all.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/tools/test_run_all.py @@ -28,6 +28,7 @@ def test_runner_passes_registered_fixture_to_both_backends( runner, monkeypatch, tmp_path, filename ): (tmp_path / "CMakeCache.txt").touch() + (tmp_path / "CTestTestfile.cmake").touch() executable = tmp_path / "Debug" / filename executable.parent.mkdir() executable.touch() @@ -69,6 +70,77 @@ def run(command, **kwargs): "--target", "rocke_storage", ] + assert any(cmd[:3] == ["ctest", "-C", "Debug"] for cmd, _ in calls) + + +@pytest.mark.parametrize("built", ["none", "partial", "all"]) +@pytest.mark.parametrize("ctest_rc", [0, 8]) +def test_ctest_readiness_and_failure_propagation( + runner, monkeypatch, tmp_path, capsys, built, ctest_rc +): + (tmp_path / "CTestTestfile.cmake").touch() + executable = tmp_path / "Debug" / "renamed.exe" + executable.parent.mkdir() + if built != "none": + executable.touch() + commands = [[str(executable)], [] if built != "all" else [str(executable)]] + calls = [] + + def run(command, **kwargs): + calls.append(command) + listing = { + "tests": [ + {"name": str(i), "command": cmd} for i, cmd in enumerate(commands) + ] + } + return subprocess.CompletedProcess( + command, + 0 if "--show-only=json-v1" in command else ctest_rc, + stdout=json.dumps(listing), + ) + + monkeypatch.setattr(runner.subprocess, "run", run) + monkeypatch.setattr( + runner.sys, + "argv", + [ + "run_all.py", + "--no-guard", + "--no-gate", + "--no-pytest", + "--build-root", + str(tmp_path), + "--config", + "Debug", + ], + ) + assert runner.main() == (0 if built == "none" else ctest_rc) + executions = [cmd for cmd in calls if cmd[:3] == ["ctest", "-C", "Debug"]] + assert len(executions) == (0 if built == "none" else 1) + if built == "none": + assert "ctest: SKIPPED" in capsys.readouterr().out + + +def test_ctest_discovery_failure_is_an_error(runner, monkeypatch, tmp_path): + (tmp_path / "CTestTestfile.cmake").touch() + + def run(command, **kwargs): + raise subprocess.CalledProcessError(1, command) + + monkeypatch.setattr(runner.subprocess, "run", run) + monkeypatch.setattr( + runner.sys, + "argv", + [ + "run_all.py", + "--no-guard", + "--no-gate", + "--no-pytest", + "--build-root", + str(tmp_path), + ], + ) + assert runner.main() == 1 def test_explicit_fixture_override_is_resolved_before_pytest_changes_directory( From 4b57104778b7b76ee4d8061fa43bfe240c90954e Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:35:42 -0400 Subject: [PATCH 09/24] refactor(rocke): keep bit packing without a duplicate tensor descriptor --- .../rocke/platform/cpp/core/storage.cpp | 60 ------ .../rocke/platform/cpp/helpers/mma_io.cpp | 34 ++- .../rocke/helper_rocke.helpers.mma_io.h | 9 +- .../platform/cpp/include/rocke/storage.h | 28 --- .../architecture/dtype_storage_packing.md | 65 +++--- .../platform/python/rocke/core/storage.py | 61 ------ .../platform/python/rocke/helpers/mma_io.py | 45 ++-- .../rocke/platform/tests/core/test_dtypes.py | 2 - .../rocke/platform/tests/core/test_mma_io.py | 107 +++------- .../platform/tests/core/test_storage.cpp | 193 +++--------------- .../rocke/platform/tests/core/test_storage.py | 31 +-- 11 files changed, 129 insertions(+), 506 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/storage.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/storage.cpp index 131c00e3989d..c951bf74096f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/storage.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/storage.cpp @@ -175,66 +175,6 @@ bool rocke_fragment_pack(const rocke_fragment_packing_t* p, return true; } -bool rocke_tensor_storage_init(rocke_tensor_storage_t* out, - const char* dtype, - uint64_t rows, - uint64_t cols, - uint64_t row_stride_bytes, - int slot_bits, - uint64_t base_bit_offset, - uint64_t alignment_bytes) -{ - rocke_tensor_storage_t p = {}; - p.dtype = rocke_dtype_info(dtype); - uint64_t row_bytes, bytes; - if(!out || !p.dtype || !alignment_bytes || (alignment_bytes & (alignment_bytes - 1)) - || !rocke_bit_packing_init(&p.packing, p.dtype->encoded_bits, slot_bits) - || !rocke_bit_packing_bytes(&p.packing, cols, base_bit_offset, &row_bytes)) - return false; - p.rows = rows; - p.cols = cols; - p.row_stride_bytes = row_stride_bytes == UINT64_MAX ? row_bytes : row_stride_bytes; - p.base_bit_offset = base_bit_offset; - p.alignment_bytes = alignment_bytes; - if(!rocke_tensor_storage_bytes(&p, &bytes)) - return false; - *out = p; - return true; -} - -bool rocke_tensor_storage_bytes(const rocke_tensor_storage_t* p, uint64_t* bytes) -{ - uint64_t row_bytes, prefix; - if(!p || !bytes || !p->dtype || !valid(&p->packing) - || rocke_dtype_info(p->dtype->name) != p->dtype - || p->packing.element_bits != p->dtype->encoded_bits || !p->alignment_bytes - || (p->alignment_bytes & (p->alignment_bytes - 1)) - || !rocke_bit_packing_bytes(&p->packing, p->cols, p->base_bit_offset, &row_bytes) - || p->row_stride_bytes < row_bytes) - return false; - if(!p->rows || !p->cols) - { - *bytes = 0; - return true; - } - return mul(p->rows - 1, p->row_stride_bytes, &prefix) && add(prefix, row_bytes, bytes); -} - -bool rocke_tensor_storage_address( - const rocke_tensor_storage_t* p, uint64_t row, uint64_t col, uint64_t* byte, int* shift) -{ - uint64_t bytes; - if(!byte || !shift || !rocke_tensor_storage_bytes(p, &bytes) || row >= p->rows - || col >= p->cols) - return false; - uint64_t bits, prefix; - if(!rocke_bit_packing_offset(&p->packing, col, &bits) || !add(bits, p->base_bit_offset, &bits) - || !mul(row, p->row_stride_bytes, &prefix) || !add(prefix, bits / 8, byte)) - return false; - *shift = bits % 8; - return true; -} - bool rocke_matrix_fragment_layout_init(rocke_matrix_fragment_layout_t* out, const rocke_fragment_packing_t* fragment, int chunk_elements, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp index c4baddbb2557..92a4e6a93396 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp @@ -34,21 +34,18 @@ rocke_value_t* rocke_h_load_matrix_fragment(rocke_ir_builder_t* b, rocke_value_t* row_base, rocke_value_t* lane_group, uint64_t k0, - const rocke_tensor_storage_t* storage, + const char* dtype, const rocke_matrix_fragment_layout_t* layout, - const rocke_type_t* carrier_type) + const rocke_type_t* carrier_type, + int alignment_bytes) { return ckc::guard_builder(b, [&]() -> rocke_value_t* { if(!rocke_i_live(b)) return NULL; - if(!storage || !storage->dtype || !layout || !carrier_type || !ptr || !row_base - || !lane_group) + if(!dtype || !layout || !carrier_type || !ptr || !row_base || !lane_group) ckc::raise_status(ROCKE_ERR_VALUE, "null matrix fragment argument"); - uint64_t storage_bytes; - if(!rocke_tensor_storage_bytes(storage, &storage_bytes)) - ckc::raise_status(ROCKE_ERR_VALUE, "invalid tensor storage descriptor"); - if(storage_bytes == 0) - ckc::raise_status(ROCKE_ERR_VALUE, "matrix fragment requires nonempty tensor storage"); + if(alignment_bytes <= 0 || (alignment_bytes & (alignment_bytes - 1))) + ckc::raise_status(ROCKE_ERR_VALUE, "alignment_bytes must be a positive power of two"); const auto& packing = layout->fragment; rocke_matrix_fragment_layout_t checked; if(!rocke_matrix_fragment_layout_init(&checked, @@ -58,23 +55,17 @@ rocke_value_t* rocke_h_load_matrix_fragment(rocke_ir_builder_t* b, layout->lanes_per_group) || !packing.count || packing.count > INT_MAX || packing.carrier_count > INT_MAX) ckc::raise_status(ROCKE_ERR_VALUE, "invalid matrix fragment chunk layout"); - const rocke_type_t* unit_type = rocke_storage_ir_type(storage->dtype->name); + const rocke_type_t* unit_type = rocke_storage_ir_type(dtype); if(!unit_type || ptr->type->kind != ROCKE_TYPE_PTR || !rocke_type_eq(ptr->type->pointee, unit_type)) ckc::raise_status(ROCKE_ERR_VALUE, "matrix pointer storage type mismatch"); const uint64_t unit_bytes = rocke_dtype_info(unit_type->name)->encoded_bits / 8; - if(storage->packing.element_bits != packing.packing.element_bits - || storage->packing.slot_bits != packing.packing.slot_bits || storage->base_bit_offset) - ckc::raise_status(ROCKE_ERR_VALUE, - "matrix fragment requires matching packing and a zero bit origin"); - const uint64_t span = packing.count * layout->lane_groups; - if(k0 > storage->cols || span > storage->cols - k0) - ckc::raise_status(ROCKE_ERR_VALUE, "matrix fragment exceeds the packed row"); + if(rocke_dtype_info(dtype)->encoded_bits != packing.packing.element_bits) + ckc::raise_status(ROCKE_ERR_VALUE, "matrix dtype and packing width mismatch"); uint64_t origin_bits, chunk_bytes; - if(!rocke_bit_packing_offset(&storage->packing, k0, &origin_bits) + if(!rocke_bit_packing_offset(&packing.packing, k0, &origin_bits) || !rocke_bit_packing_bytes(&packing.packing, layout->chunk_elements, 0, &chunk_bytes) - || origin_bits % (8 * unit_bytes) || chunk_bytes % unit_bytes - || storage->row_stride_bytes % unit_bytes) + || origin_bits % (8 * unit_bytes) || chunk_bytes % unit_bytes) ckc::raise_status(ROCKE_ERR_VALUE, "matrix fragment is not aligned to pointer storage units"); const rocke_dtype_info_t* carrier_info = rocke_dtype_info(carrier_type->name); @@ -98,8 +89,7 @@ rocke_value_t* rocke_h_load_matrix_fragment(rocke_ir_builder_t* b, > (uint64_t(INT_MAX) + 1 - origin_bytes / unit_bytes) / chunk_units / layout->lane_groups) ckc::raise_status(ROCKE_ERR_VALUE, "matrix fragment offset exceeds i32 range"); - int alignment = gcd(gcd(storage->alignment_bytes, storage->row_stride_bytes), - gcd(chunk_bytes, origin_bytes)); + int alignment = gcd(alignment_bytes, gcd(chunk_bytes, origin_bytes)); auto* lane_chunk = rocke_b_mul(b, lane_group, rocke_b_const_i32(b, chunk_units)); auto* step_base = rocke_b_add(b, row_base, rocke_b_const_i32(b, origin_bytes / unit_bytes)); /* Collect loads before concatenation to preserve Python SSA numbering. */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.mma_io.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.mma_io.h index 8f5db63aa960..7dea7828c6f6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.mma_io.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.mma_io.h @@ -8,6 +8,7 @@ #ifndef ROCKE_HELPER_MMA_IO_H #define ROCKE_HELPER_MMA_IO_H +#include "rocke/dtypes.h" #include "rocke/ir.h" #include "rocke/storage.h" @@ -16,14 +17,18 @@ extern "C" { #endif const rocke_type_t* rocke_storage_ir_type(const char* dtype); +/* row_base is in pointer storage units. alignment_bytes is guaranteed at that + * row address, before k0 and lane/chunk offsets. Caller owns tensor indexing, + * allocation bounds (including empty/partial rows), and lane_group validity. */ rocke_value_t* rocke_h_load_matrix_fragment(rocke_ir_builder_t* b, rocke_value_t* ptr, rocke_value_t* row_base, rocke_value_t* lane_group, uint64_t k0, - const rocke_tensor_storage_t* storage, + const char* dtype, const rocke_matrix_fragment_layout_t* layout, - const rocke_type_t* carrier_type); + const rocke_type_t* carrier_type, + int alignment_bytes); /* Callback returns a canonical unsigned pattern in an integer carrier. * ctx and output pointer array are borrowed for the duration of this call. */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/storage.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/storage.h index 133c1310c9ac..9634dbcc2e1f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/storage.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/storage.h @@ -3,7 +3,6 @@ /* Shared bit storage; mirrors core/storage.py. * BitPacking -> rocke_bit_packing_t * FragmentPacking -> rocke_fragment_packing_t - * TensorStorage -> rocke_tensor_storage_t * All descriptors are caller-owned. Functions return false for invalid inputs, * overflow or insufficient buffers; no exceptions cross this C API. */ @@ -14,8 +13,6 @@ #include #include -#include "rocke/dtypes.h" - #ifdef __cplusplus extern "C" { #endif @@ -50,17 +47,6 @@ bool rocke_matrix_fragment_layout_init(rocke_matrix_fragment_layout_t* out, bool rocke_matrix_fragment_coord( const rocke_matrix_fragment_layout_t* p, int lane, uint64_t slot, uint64_t* row, uint64_t* k); -typedef struct rocke_tensor_storage -{ - const rocke_dtype_info_t* dtype; /* Borrowed static descriptor. */ - uint64_t rows; - uint64_t cols; - uint64_t row_stride_bytes; - uint64_t base_bit_offset; - uint64_t alignment_bytes; - rocke_bit_packing_t packing; -} rocke_tensor_storage_t; - bool rocke_bit_packing_init(rocke_bit_packing_t* out, int element_bits, int slot_bits); bool rocke_bit_packing_group(const rocke_bit_packing_t* p, int carrier_bits, @@ -95,20 +81,6 @@ bool rocke_fragment_pack(const rocke_fragment_packing_t* p, size_t count, uint64_t* carriers, size_t carrier_count); -/* row_stride_bytes=UINT64_MAX selects a minimal packed row stride. */ -bool rocke_tensor_storage_init(rocke_tensor_storage_t* out, - const char* dtype, - uint64_t rows, - uint64_t cols, - uint64_t row_stride_bytes, - int slot_bits, - uint64_t base_bit_offset, - uint64_t alignment_bytes); -/* Size and address queries validate public descriptors; invalid input returns false. */ -bool rocke_tensor_storage_bytes(const rocke_tensor_storage_t* p, uint64_t* bytes); -bool rocke_tensor_storage_address( - const rocke_tensor_storage_t* p, uint64_t row, uint64_t col, uint64_t* byte, int* shift); - #ifdef __cplusplus } #endif diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md index 0d68276de961..aeb0a863af35 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md @@ -3,7 +3,7 @@ Matrix inputs and scales use three separate contracts: 1. A logical dtype identifies the numerical encoding and its bit width. -2. Tensor storage describes how elements occupy an addressable buffer. +2. Bit packing describes how encoded patterns occupy bytes and carrier words. 3. An instruction fragment describes lane mapping, carrier capacity, and padding. The selected atom determines supported operand combinations and instruction @@ -13,7 +13,7 @@ selectors. An integer carrier does not identify the numerical format. - [Logical dtypes](#logical-dtypes) - [Common packing](#common-packing) -- [Tensor storage](#tensor-storage) +- [Fragment loads](#fragment-loads) - [Matrix and scale layouts](#matrix-and-scale-layouts) - [Mirroring and follow-up work](#mirroring-and-follow-up-work) @@ -66,35 +66,36 @@ Host pack/unpack routines process bit patterns only. They do not round, clamp, decode floats, or apply scales. Their tests use literal expected bits and independent integer arithmetic; numerical tests remain separate. -## Tensor storage - -`TensorStorage` / `rocke_tensor_storage_t` initially supports a two-dimensional -logical shape with a contiguous packed inner axis, an explicit byte row stride, -bit offset, and base alignment. Nonnegative sizes and offsets are checked against -the native uint64 range. A/B descriptors are independent. - -For example, dense FP6 uses four elements in three bytes; `slot_bits=8` describes -one six-bit pattern in each byte instead. Both have the same logical dtype. -Address calculation includes the bit offset, and allocation bounds include the -partial tail byte. Empty views require no accesses. -Native size queries, address queries, and fragment loads revalidate the public -storage descriptor, including dtype/packing consistency, stride, alignment, and -allocation bounds, because C callers can construct or mutate these structs. - -[`storage_ir_type`](../../python/rocke/helpers/mma_io.py) chooses an addressable unit. FP4/FP6 use I8 -units while retaining their logical identities in the storage descriptor. -This function is deliberately separate from `dtype_to_ir_type`. -FP8/BF8 retain nominal FP8/BF8 pointer types despite occupying one byte per value. - -The fragment loader accepts complete, byte-aligned chunks. It derives load -alignment from base alignment, row stride, chunk spacing, and origin. A 24-byte -FP6 chunk becomes 16-byte and 8-byte loads without reading past the chunk. -Partial fragments and nonzero bit origins are rejected by this initial loader. -Row strides must also be divisible by the pointer storage unit: an odd-byte -FP16/BF16 stride is valid generic storage but cannot be loaded through a typed -two-byte pointer. `row_base` is an element index in those storage units. -The host bit packer writes a fresh, exclusively owned output buffer; concurrent -packed stores are not provided. +## Fragment loads + +[`storage_ir_type`](../../python/rocke/helpers/mma_io.py) chooses an addressable +unit from the logical dtype. FP4/FP6 use I8; FP8/BF8 retain their nominal pointer +types despite occupying one byte per value. This function is separate from +`dtype_to_ir_type`. + +`load_matrix_fragment` takes a dtype, fragment layout, caller-selected row base, +lane group, and K origin. `row_base` counts pointer storage units. The optional +`alignment_bytes` (default 1) guarantees alignment at that row address, before +the loader adds K and lane/chunk offsets. For a buffer aligned to 16 bytes with +97-byte row spacing, arbitrary rows only guarantee alignment 1; a typed FP16 +buffer with 258-byte row spacing guarantees alignment 2 and uses row offsets +in units of two bytes. + +The loader checks dtype/packing width, pointer identity, whole-unit chunk and +origin alignment, carrier capacity, and static i32 displacement limits. It +reduces the supplied alignment for chunk spacing and K origin. A 24-byte FP6 +chunk becomes 16-byte and 8-byte loads without reading past the chunk. + +Tensor shape, strides, row selection, and allocation bounds belong to the +caller. This helper always loads a complete fragment: the caller must guard +empty or partial rows and provide a valid lane group. It is not a bounds-checked +tensor view. Existing `TensorDescriptor`/`TensorView` integration is deferred; +no separate tensor descriptor is introduced here. + +`BitPacking(6)` describes dense FP6; `BitPacking(6, 8)` describes one six-bit +pattern per byte. Bit offsets and stream sizes are checked against uint64. +Host packers support partial tail bytes and write fresh, exclusively owned +buffers; concurrent packed stores are not provided. ## Matrix and scale layouts @@ -128,7 +129,7 @@ absent from its fixed compatibility prologue. The shared descriptor/helper changes can be consumed by both FP4 and FP6 branches. FP4/FP8 builder migration can use them directly. Future FP6 numerical -integration should consume the six-bit storage and fragment descriptors and +integration should consume the six-bit packing and fragment descriptors and retain separate target and conversion validation. First-class tensor-view/fragment IR nodes, arbitrary packed axes, masked partial diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/storage.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/storage.py index 7c59be34fc0f..c9c1ce3282e9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/storage.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/storage.py @@ -8,8 +8,6 @@ from math import gcd from collections.abc import Sequence -from .dtypes import DTypeInfo, dtype_info - _MAX_BITS = (1 << 64) - 1 @@ -168,62 +166,3 @@ def coord(self, lane: int, slot: int) -> tuple[int, int]: * self.chunk_elements + element, ) - - -@dataclass(frozen=True) -class TensorStorage: - """Two-dimensional storage with a contiguous packed inner axis. - - Shape is logical. Row strides are bytes. Generic axis permutations and - negative strides are intentionally rejected by this initial interface. - """ - - dtype: str - shape: tuple[int, int] - row_stride_bytes: int | None = None - slot_bits: int | None = None - base_bit_offset: int = 0 - alignment_bytes: int = 1 - - def __post_init__(self) -> None: - info = dtype_info(self.dtype) - object.__setattr__(self, "dtype", info.name) - if len(self.shape) != 2: - raise ValueError("storage requires a two-dimensional logical shape") - for dim in self.shape: - _checked(dim) - _checked(self.base_bit_offset) - if self.alignment_bytes <= 0 or self.alignment_bytes & ( - self.alignment_bytes - 1 - ): - raise ValueError("alignment_bytes must be a positive power of two") - row_bytes = self.packing.byte_size(self.shape[1], self.base_bit_offset) - stride = row_bytes if self.row_stride_bytes is None else self.row_stride_bytes - if _checked(stride) < row_bytes: - raise ValueError("row stride is smaller than the packed row") - object.__setattr__(self, "row_stride_bytes", stride) - _checked(self.byte_size) - - @property - def info(self) -> DTypeInfo: - return dtype_info(self.dtype) - - @property - def packing(self) -> BitPacking: - return BitPacking(self.info.encoded_bits, self.slot_bits) - - @property - def byte_size(self) -> int: - rows, cols = self.shape - if rows == 0 or cols == 0: - return 0 - return _checked( - (rows - 1) * self.row_stride_bytes - + self.packing.byte_size(cols, self.base_bit_offset) - ) - - def address(self, row: int, col: int) -> tuple[int, int]: - if not (0 <= row < self.shape[0] and 0 <= col < self.shape[1]): - raise ValueError("storage coordinate out of bounds") - bits = _checked(self.base_bit_offset + self.packing.bit_offset(col)) - return _checked(row * self.row_stride_bytes + bits // 8), bits % 8 diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py index b47f120a0beb..28bfdf60c316 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py @@ -19,7 +19,7 @@ VectorType, dtype_to_ir_type, ) -from ..core.storage import FragmentPacking, MatrixFragmentLayout, TensorStorage +from ..core.storage import FragmentPacking, MatrixFragmentLayout def storage_ir_type(dtype: str) -> Type: @@ -37,38 +37,32 @@ def load_matrix_fragment( lane_group: Value, k0: int, *, - storage: TensorStorage, + dtype: str, layout: MatrixFragmentLayout, carrier_type: Type = I32, + alignment_bytes: int = 1, ) -> Value: - """Load whole-byte chunks directly into the atom's carrier vector. + """Load complete chunks from a caller-selected row into carrier registers. - row_base counts pointer storage units, which must divide the row stride. - lane_group is derived from the atom's - lane mapping by the caller. Partial chunks and nonzero bit origins reject. + row_base counts pointer storage units; alignment_bytes is the guaranteed + alignment at that row address, before k0 and lane/chunk offsets. The caller + owns tensor indexing and allocation bounds, including empty/partial rows. + lane_group comes from the atom's lane mapping. k0 must address a whole + pointer unit. TensorDescriptor/TensorView integration is deferred. """ - if storage.byte_size == 0: - raise ValueError("matrix fragment requires nonempty tensor storage") + if alignment_bytes <= 0 or alignment_bytes & (alignment_bytes - 1): + raise ValueError("alignment_bytes must be a positive power of two") packing = layout.fragment if not (0 < packing.count <= 0x7FFFFFFF and packing.carrier_count <= 0x7FFFFFFF): raise ValueError("invalid matrix fragment chunk layout") - unit_type = storage_ir_type(storage.dtype) + unit_type = storage_ir_type(dtype) unit_bytes = dtype_info(unit_type.name).encoded_bits // 8 if ptr.type.pointee != unit_type: raise ValueError("matrix pointer storage type mismatch") - if storage.packing != packing.packing or storage.base_bit_offset: - raise ValueError( - "matrix fragment requires matching packing and a zero bit origin" - ) - span = layout.chunk_elements * layout.lane_groups * layout.chunks_per_lane - if k0 < 0 or k0 + span > storage.shape[1]: - raise ValueError("matrix fragment exceeds the packed row") - origin_bits = storage.packing.bit_offset(k0) - if ( - origin_bits % (8 * unit_bytes) - or layout.chunk_bytes % unit_bytes - or storage.row_stride_bytes % unit_bytes - ): + if dtype_info(dtype).encoded_bits != packing.packing.element_bits: + raise ValueError("matrix dtype and packing width mismatch") + origin_bits = packing.packing.bit_offset(k0) + if origin_bits % (8 * unit_bytes) or layout.chunk_bytes % unit_bytes: raise ValueError("matrix fragment is not aligned to pointer storage units") if dtype_info(carrier_type.name).encoded_bits != packing.carrier_bits: raise ValueError("matrix carrier type width mismatch") @@ -84,12 +78,7 @@ def load_matrix_fragment( > (0x80000000 - origin_bytes // unit_bytes) // chunk_units // layout.lane_groups ): raise ValueError("matrix fragment offset exceeds i32 range") - alignment = gcd( - storage.alignment_bytes, - storage.row_stride_bytes, - layout.chunk_bytes, - origin_bytes, - ) + alignment = gcd(alignment_bytes, layout.chunk_bytes, origin_bytes) lane_chunk = b.mul(lane_group, b.const_i32(chunk_units)) step_base = b.add(row_base, b.const_i32(origin_bytes // unit_bytes)) chunks = [] 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 index 2222d6ed4727..2fd2eb0ac7e9 100644 --- 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 @@ -10,7 +10,6 @@ from rocke.core.dtypes import dtype_info from rocke.core.arch.target import MmaScaleDType from rocke.core.ir import FP8E4M3, I8, IRBuilder, dtype_to_ir_type -from rocke.core.storage import TensorStorage from rocke.helpers.mma_io import storage_ir_type from rocke.core.ir_serialize import parse, serialize from rocke.helpers.quant import ( @@ -102,7 +101,6 @@ def test_e4m3_shares_encoding_type_and_storage(spelling): assert dtype_info(spelling) is dtype_info("fp8e4m3") assert dtype_to_ir_type(spelling) is FP8E4M3 assert storage_ir_type(spelling) is FP8E4M3 - assert TensorStorage(spelling, (2, 128)) == TensorStorage("fp8e4m3", (2, 128)) assert quant_ir_type(spelling) is FP8E4M3 b = IRBuilder("e4m3_conversion") value = b.param("value", dtype_to_ir_type(spelling)) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py index 23515da8a3cd..67caa9541dcf 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py @@ -19,7 +19,6 @@ BitPacking, FragmentPacking, MatrixFragmentLayout, - TensorStorage, ) from rocke.helpers.mma_io import ( load_matrix_fragment, @@ -69,12 +68,6 @@ def build_transport(dtype): for j, value in enumerate(values): b.global_store(o, b.const_i32(j), value, align=4) else: - storage = TensorStorage( - dtype, - (16, 128), - row_stride_bytes=(258 if typed else 97) if padded else None, - alignment_bytes=16, - ) layout = ( MatrixFragmentLayout(FragmentPacking(BitPacking(16), 32, 16, 32), 16, 2, 16) if typed @@ -84,7 +77,15 @@ def build_transport(dtype): lane = b.mod(b.thread_id_x(), b.const_i32(32)) group = b.div(lane, b.const_i32(16)) value = load_matrix_fragment( - b, a, base, group, 0, storage=storage, layout=layout, carrier_type=carrier + b, + a, + base, + group, + 0, + dtype=dtype, + layout=layout, + carrier_type=carrier, + alignment_bytes=(2 if typed else 1) if padded else 16, ) for j in range(value.type.count): index = b.const_i32(j) @@ -124,69 +125,43 @@ def test_native_helper_emits_identical_ir(dtype, route): assert actual.stdout == expected -def test_reject_partial_fragment_and_wrong_pointer(): +def test_reject_bit_origin(): b = IRBuilder("invalid_transport") ptr = b.param("A", PtrType(I8, "global")) zero = b.const_i32(0) - with pytest.raises(ValueError, match="exceeds"): + with pytest.raises(ValueError, match="aligned to pointer storage units"): load_matrix_fragment( b, ptr, zero, zero, 1, - storage=TensorStorage("fp4", (16, 128)), + dtype="fp4", layout=scaled_matrix_layout("fp4", 16), ) - with pytest.raises(ValueError, match="pointer storage type"): - load_matrix_fragment( - b, - ptr, - zero, - zero, - 0, - storage=TensorStorage("fp8", (16, 128)), - layout=scaled_matrix_layout("fp8", 16), - ) -@pytest.mark.parametrize("dtype", ["fp4", "fp6", "fp8", "f16", "bf16"]) -@pytest.mark.parametrize("shape", [(0, 128), (16, 0), (0, 0), (16, 128)]) -def test_fragment_load_requires_nonempty_storage(dtype, shape): - storage = TensorStorage(dtype, shape) - typed = dtype in ("f16", "bf16") - unit = storage_ir_type(dtype) - layout = ( - MatrixFragmentLayout(FragmentPacking(BitPacking(16), 32, 16, 32), 16, 2, 16) - if typed - else scaled_matrix_layout(dtype, 16) - ) - b = IRBuilder("empty_storage") - ptr = b.param("A", PtrType(unit, "global")) +@pytest.mark.parametrize( + "dtype,alignment,error", + [("fp6", 3, "positive power of two"), ("fp4", 1, "packing width mismatch")], +) +def test_fragment_input_contract(dtype, alignment, error): + b = IRBuilder("invalid_fragment") + ptr = b.param("A", PtrType(I8, "global")) zero = b.const_i32(0) before = serialize(b.kernel) - - def load(): - return load_matrix_fragment( + with pytest.raises(ValueError, match=error): + load_matrix_fragment( b, ptr, zero, zero, 0, - storage=storage, - layout=layout, - carrier_type=unit if typed else I32, + dtype=dtype, + layout=scaled_matrix_layout("fp6", 16), + alignment_bytes=alignment, ) - - if 0 in shape: - assert storage.byte_size == 0 - with pytest.raises(ValueError, match="nonempty tensor storage"): - load() - assert serialize(b.kernel) == before - else: - assert storage.byte_size > 0 - assert load() is not None - assert "memref.global_load" in serialize(b.kernel) + assert serialize(b.kernel) == before @pytest.mark.parametrize("dtype", ["e8m0", "fp6", "f16"]) @@ -206,9 +181,6 @@ def test_fragment_offset_bounds(dtype, count, groups, k0, valid): # Padded FP6 and E8M0 use byte units; FP16 uses two-byte pointer units. bits = 16 if dtype == "f16" else 6 if dtype == "fp6" else 8 packing = BitPacking(bits, 16 if dtype == "f16" else 8) - storage = TensorStorage( - dtype, (1, k0 + count * groups), slot_bits=packing.slot_bits - ) layout = MatrixFragmentLayout( FragmentPacking(packing, count, 32, count * packing.slot_bits // 32), 16, @@ -221,12 +193,12 @@ def test_fragment_offset_bounds(dtype, count, groups, k0, valid): before = serialize(b.kernel) if valid: assert ( - load_matrix_fragment(b, ptr, zero, zero, k0, storage=storage, layout=layout) + load_matrix_fragment(b, ptr, zero, zero, k0, dtype=dtype, layout=layout) is not None ) else: with pytest.raises(ValueError, match="offset exceeds i32 range"): - load_matrix_fragment(b, ptr, zero, zero, k0, storage=storage, layout=layout) + load_matrix_fragment(b, ptr, zero, zero, k0, dtype=dtype, layout=layout) assert serialize(b.kernel) == before @@ -240,29 +212,6 @@ def test_hip_declares_only_encountered_missing_vector_widths(): assert "using i8x24 =" not in source -@pytest.mark.parametrize("dtype", ["f16", "bf16"]) -def test_reject_row_stride_not_in_pointer_units(dtype): - b = IRBuilder("odd_stride") - unit = storage_ir_type(dtype) - ptr = b.param("A", PtrType(unit, "global")) - zero = b.const_i32(0) - before = serialize(b.kernel) - with pytest.raises(ValueError, match="aligned to pointer storage units"): - load_matrix_fragment( - b, - ptr, - zero, - zero, - 0, - storage=TensorStorage(dtype, (2, 64), row_stride_bytes=129), - layout=MatrixFragmentLayout( - FragmentPacking(BitPacking(16), 32, 16, 32), 16, 2, 16 - ), - carrier_type=unit, - ) - assert serialize(b.kernel) == before - - @pytest.mark.parametrize("dtype", ["fp8", "bf8"]) def test_byte_width_does_not_erase_pointer_identity(dtype): assert storage_ir_type(dtype) != I8 @@ -276,7 +225,7 @@ def test_byte_width_does_not_erase_pointer_identity(dtype): zero, zero, 0, - storage=TensorStorage(dtype, (16, 128)), + dtype=dtype, layout=scaled_matrix_layout(dtype, 16), ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp index 981597c52560..de507eba1c34 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp @@ -65,9 +65,6 @@ static int emit(const char* dtype, bool hip) } else { - rocke_tensor_storage_t storage; - CHECK(rocke_tensor_storage_init( - &storage, dtype, 16, 128, padded ? (typed ? 258 : 97) : UINT64_MAX, 0, 0, 16)); rocke_matrix_fragment_layout_t layout; if(typed) { @@ -83,8 +80,8 @@ static int emit(const char* dtype, bool hip) auto* thread = rocke_b_thread_id_x(&b); auto* lane = rocke_b_mod(&b, thread, rocke_b_const_i32(&b, 32)); auto* group = rocke_b_div(&b, lane, rocke_b_const_i32(&b, 16)); - auto* value - = rocke_h_load_matrix_fragment(&b, a, base, group, 0, &storage, &layout, carrier); + auto* value = rocke_h_load_matrix_fragment( + &b, a, base, group, 0, dtype, &layout, carrier, padded ? (typed ? 2 : 1) : 16); CHECK(value && rocke_ir_builder_ok(&b)); for(int j = 0; j < value->type->count; ++j) { @@ -115,152 +112,38 @@ static int emit(const char* dtype, bool hip) return 0; } -static int test_storage_validation() +static int test_fragment_inputs() { - rocke_tensor_storage_t storage; - CHECK(rocke_tensor_storage_init(&storage, "fp6", 16, 128, 97, 0, 0, 16)); - const auto layout = rocke_scaled_matrix_layout("fp6", 16); - for(int defect = 0; defect < 10; ++defect) + const struct { - auto invalid = storage; - switch(defect) - { - case 0: - invalid.dtype = NULL; - break; - case 1: - invalid.dtype = rocke_dtype_info("fp8"); - break; - case 2: - invalid.packing.element_bits = 0; - break; - case 3: - invalid.packing.slot_bits = 0; - break; - case 4: - invalid.packing.slot_bits = 4; - break; - case 5: - invalid.row_stride_bytes = 95; - break; - case 6: - invalid.alignment_bytes = 0; - break; - case 7: - invalid.alignment_bytes = 3; - break; - case 8: - invalid.rows = UINT64_MAX; - break; - case 9: - invalid.base_bit_offset = UINT64_MAX; - break; - } - uint64_t bytes; - int shift; - CHECK(!rocke_tensor_storage_bytes(&invalid, &bytes)); - CHECK(!rocke_tensor_storage_address(&invalid, 0, 0, &bytes, &shift)); - rocke_ir_builder_t b; - CHECK(rocke_ir_builder_init(&b, "invalid_storage") == ROCKE_OK); - auto* ptr = rocke_b_param(&b, "A", rocke_ptr_type(&b, rocke_i8(), "global"), NULL); - auto* zero = rocke_b_const_i32(&b, 0); - CHECK( - !rocke_h_load_matrix_fragment(&b, ptr, zero, zero, 0, &invalid, &layout, rocke_i32())); - CHECK(!rocke_ir_builder_ok(&b)); - rocke_ir_builder_free(&b); - if(defect < 8) - { - invalid.rows = 0; - CHECK(!rocke_tensor_storage_bytes(&invalid, &bytes)); - } - } - for(const char* dtype : {"f16", "bf16", "fp8", "e4m3", "bf8"}) + const char* dtype; + const char* layout_dtype; + uint64_t k0; + int alignment; + const char* error; + } cases[] = {{"fp6", "fp6", 0, 3, "positive power of two"}, + {"fp4", "fp6", 0, 1, "packing width mismatch"}, + {"fp4", "fp4", 1, 1, "aligned to pointer storage units"}, + {"fp8", "fp8", 0, 1, "pointer storage type mismatch"}, + {"bf8", "bf8", 0, 1, "pointer storage type mismatch"}}; + for(const auto& c : cases) { - const auto* unit = rocke_storage_ir_type(dtype); - CHECK(unit && !rocke_type_eq(unit, rocke_i8())); - const bool typed = strcmp(dtype, "f16") == 0 || strcmp(dtype, "bf16") == 0; - rocke_tensor_storage_t direct = {rocke_dtype_info(dtype), - 2, - 128, - typed ? 257u : 129u, - 0, - 1, - {typed ? 16 : 8, typed ? 16 : 8}}; - uint64_t bytes; - int shift; - CHECK(rocke_tensor_storage_bytes(&direct, &bytes)); - CHECK(rocke_tensor_storage_address(&direct, 1, 0, &bytes, &shift)); - CHECK(bytes == direct.row_stride_bytes && shift == 0); - rocke_matrix_fragment_layout_t fragment; - if(typed) - { - rocke_fragment_packing_t packing; - CHECK(rocke_fragment_packing_init(&packing, &direct.packing, 32, 16, 32)); - CHECK(rocke_matrix_fragment_layout_init(&fragment, &packing, 16, 2, 16)); - } - else - fragment = rocke_scaled_matrix_layout(dtype, 16); rocke_ir_builder_t b; - CHECK(rocke_ir_builder_init(&b, "invalid_pointer_units") == ROCKE_OK); - auto* ptr - = rocke_b_param(&b, "A", rocke_ptr_type(&b, typed ? unit : rocke_i8(), "global"), NULL); + CHECK(rocke_ir_builder_init(&b, "invalid_fragment") == ROCKE_OK); + auto* ptr = rocke_b_param(&b, "A", rocke_ptr_type(&b, rocke_i8(), "global"), NULL); auto* zero = rocke_b_const_i32(&b, 0); + const auto layout = rocke_scaled_matrix_layout(c.layout_dtype, 16); + const int before = b.kernel->body->num_ops; CHECK(!rocke_h_load_matrix_fragment( - &b, ptr, zero, zero, 0, &direct, &fragment, typed ? unit : rocke_i32())); - CHECK(!rocke_ir_builder_ok(&b)); + &b, ptr, zero, zero, c.k0, c.dtype, &layout, rocke_i32(), c.alignment)); + CHECK(rocke_ir_builder_status(&b) == ROCKE_ERR_VALUE); + CHECK(strstr(rocke_ir_builder_error(&b), c.error)); + CHECK(b.kernel->body->num_ops == before); rocke_ir_builder_free(&b); } return 0; } -static int test_empty_storage_load() -{ - for(const char* dtype : {"fp4", "fp6", "fp8", "f16", "bf16"}) - { - const auto* unit = rocke_storage_ir_type(dtype); - const bool typed = strcmp(dtype, "f16") == 0 || strcmp(dtype, "bf16") == 0; - for(int shape = 0; shape < 4; ++shape) - { - const uint64_t rows = shape == 0 || shape == 2 ? 0 : 16; - const uint64_t cols = shape == 1 || shape == 2 ? 0 : 128; - rocke_tensor_storage_t storage; - CHECK(rocke_tensor_storage_init(&storage, dtype, rows, cols, UINT64_MAX, 0, 0, 1)); - uint64_t bytes; - CHECK(rocke_tensor_storage_bytes(&storage, &bytes)); - CHECK((bytes == 0) == (shape != 3)); - rocke_matrix_fragment_layout_t layout; - if(typed) - { - rocke_fragment_packing_t fragment; - CHECK(rocke_fragment_packing_init(&fragment, &storage.packing, 32, 16, 32)); - CHECK(rocke_matrix_fragment_layout_init(&layout, &fragment, 16, 2, 16)); - } - else - layout = rocke_scaled_matrix_layout(dtype, 16); - rocke_ir_builder_t b; - CHECK(rocke_ir_builder_init(&b, "empty_storage") == ROCKE_OK); - auto* ptr = rocke_b_param(&b, "A", rocke_ptr_type(&b, unit, "global"), NULL); - auto* zero = rocke_b_const_i32(&b, 0); - const int before = b.kernel->body->num_ops; - auto* value = rocke_h_load_matrix_fragment( - &b, ptr, zero, zero, 0, &storage, &layout, typed ? unit : rocke_i32()); - if(bytes == 0) - { - CHECK(!value && rocke_ir_builder_status(&b) == ROCKE_ERR_VALUE); - CHECK(strstr(rocke_ir_builder_error(&b), "nonempty tensor storage")); - CHECK(b.kernel->body->num_ops == before); - } - else - { - CHECK(value && rocke_ir_builder_ok(&b)); - CHECK(b.kernel->body->num_ops > before); - } - rocke_ir_builder_free(&b); - } - } - return 0; -} - static int test_fragment_offset_bounds() { const struct @@ -280,18 +163,12 @@ static int test_fragment_offset_bounds() const int slot_bits = strcmp(dtype, "f16") == 0 ? 16 : 8; for(const auto& c : cases) { - rocke_tensor_storage_t storage; - CHECK(rocke_tensor_storage_init(&storage, - dtype, - 1, - c.k0 + uint64_t(c.count) * c.groups, - UINT64_MAX, - slot_bits, - 0, - 1)); + rocke_bit_packing_t packing; + CHECK(rocke_bit_packing_init( + &packing, strcmp(dtype, "fp6") == 0 ? 6 : slot_bits, slot_bits)); rocke_fragment_packing_t fragment; CHECK(rocke_fragment_packing_init( - &fragment, &storage.packing, c.count, 32, c.count * slot_bits / 32)); + &fragment, &packing, c.count, 32, c.count * slot_bits / 32)); rocke_matrix_fragment_layout_t layout; CHECK(rocke_matrix_fragment_layout_init(&layout, &fragment, 16, c.groups, 1)); rocke_ir_builder_t b; @@ -301,7 +178,7 @@ static int test_fragment_offset_bounds() auto* zero = rocke_b_const_i32(&b, 0); const int before = b.kernel->body->num_ops; auto* value = rocke_h_load_matrix_fragment( - &b, ptr, zero, zero, c.k0, &storage, &layout, rocke_i32()); + &b, ptr, zero, zero, c.k0, dtype, &layout, rocke_i32(), 1); if(c.valid) CHECK(value && rocke_ir_builder_ok(&b)); else @@ -352,8 +229,7 @@ int main(int argc, char** argv) rocke_ir_builder_free(&b); return 0; } - CHECK(test_storage_validation() == 0); - CHECK(test_empty_storage_load() == 0); + CHECK(test_fragment_inputs() == 0); CHECK(test_fragment_offset_bounds() == 0); CHECK(rocke_dtype_info("e4m3") == rocke_dtype_info("fp8e4m3")); CHECK(rocke_dtype_to_ir_type("e4m3") == rocke_fp8e4m3()); @@ -459,17 +335,6 @@ int main(int argc, char** argv) CHECK(words[0] == 0xC0000000 && words[1] == 0xF000000F && words[2] == 3); CHECK(!rocke_bit_packing_bytes(&six, UINT64_MAX, 0, &values)); CHECK(!rocke_fragment_packing_init(&fragment, &six, 16, 32, 2)); - rocke_tensor_storage_t storage; - CHECK(rocke_tensor_storage_init(&storage, "bf6", 3, 5, 8, 0, 3, 1)); - CHECK(rocke_tensor_storage_bytes(&storage, &values) && values == 21); - int shift; - CHECK(rocke_tensor_storage_address(&storage, 2, 4, &values, &shift)); - CHECK(values == 19 && shift == 3); - CHECK(!rocke_tensor_storage_address(&storage, 3, 0, &values, &shift)); - CHECK(!rocke_tensor_storage_init(&storage, "fp6", 3, 5, 3, 0, 0, 1)); - CHECK(!rocke_tensor_storage_init(&storage, "fp8", UINT64_MAX, 3, 3, 0, 0, 1)); - CHECK(rocke_tensor_storage_init(&storage, "fp4", 3, 0, 0, 0, 7, 1)); - CHECK(rocke_tensor_storage_bytes(&storage, &values) && values == 0); for(const char* dtype : {"fp4", "fp6", "bf6", "e8m0", "e4m3", "e5m3"}) { const auto* logical = rocke_dtype_to_ir_type(dtype); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py index 6824621d8f9c..5284914cc389 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py @@ -1,11 +1,11 @@ # Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -"""Independent bit-stream fixtures and storage bounds for matrices and scales.""" +"""Independent bit-stream fixtures for matrices and scales.""" import pytest from rocke.core.arch.wmma_scale import ScalePacking, scaled_matrix_layout -from rocke.core.storage import BitPacking, FragmentPacking, TensorStorage +from rocke.core.storage import BitPacking, FragmentPacking @pytest.mark.parametrize("bits", [4, 6, 8, 16, 32, 64]) @@ -23,6 +23,7 @@ def test_dense_patterns_against_integer_oracle(bits, offset): def test_fp6_crosses_words_without_per_word_padding(): packing = BitPacking(6) + assert packing.pack([], bit_offset=7) == b"" assert packing.group(8) == (4, 3) assert packing.group(32) == (16, 3) patterns = [0] * 16 @@ -63,30 +64,6 @@ def test_scales_use_common_bit_packing(count, block_k, word): assert scales.word_bits == count * 8 -def test_storage_rows_offsets_and_empty_views(): - storage = TensorStorage("bf6", (3, 5), row_stride_bytes=8, base_bit_offset=3) - assert storage.dtype == "fp6e3m2" - assert storage.byte_size == 21 - assert storage.address(2, 4) == (19, 3) - assert TensorStorage("fp4", (0, 5)).byte_size == 0 - assert TensorStorage("fp4", (3, 0), base_bit_offset=7).byte_size == 0 - assert BitPacking(6).pack([], bit_offset=7) == b"" - assert TensorStorage("e8m0", (2, 8)).byte_size == 16 - with pytest.raises(ValueError, match="out of bounds"): - storage.address(3, 0) - with pytest.raises(ValueError, match="row stride"): - TensorStorage("fp6", (3, 5), row_stride_bytes=3) - with pytest.raises(ValueError, match="power of two"): - TensorStorage("fp4", (1, 4), alignment_bytes=3) - - -@pytest.mark.parametrize("dtype", ["f16", "bf16"]) -def test_storage_allows_byte_strides(dtype): - storage = TensorStorage(dtype, (2, 64), row_stride_bytes=129) - assert storage.address(1, 0) == (129, 0) - assert storage.byte_size == 257 - - def test_invalid_packing_and_overflow(): for bits, slots in [(0, None), (65, None), (6, 4)]: with pytest.raises(ValueError): @@ -95,7 +72,5 @@ def test_invalid_packing_and_overflow(): FragmentPacking(BitPacking(6), 16, 32, 2) with pytest.raises(ValueError, match="uint64"): BitPacking(6).byte_size(1 << 63) - with pytest.raises(ValueError, match="uint64"): - TensorStorage("fp8", (1 << 63, 3)) with pytest.raises(ValueError, match="fit"): BitPacking(6).pack([64]) From 2bc6e0ff5d2a4cc12331f139ac0366e052def007 Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:11:52 -0400 Subject: [PATCH 10/24] fix(rocke): tighten packing scope and prepare complete native tests --- .../rocke/platform/cpp/helpers/mma_io.cpp | 3 ++ .../rocke/helper_rocke.helpers.mma_io.h | 1 + .../cpp/include/rocke/wmma_scale_internal.h | 15 ++------ .../architecture/dtype_storage_packing.md | 9 ++--- .../python/rocke/core/arch/wmma_scale.py | 18 ++-------- .../platform/python/rocke/helpers/mma_io.py | 9 +++-- .../rocke/platform/tests/core/test_mma_io.py | 9 +++++ .../platform/tests/core/test_storage.cpp | 14 ++++++++ .../rocke/platform/tests/core/test_storage.py | 2 +- .../rocke/platform/tests/run_all.py | 4 +-- .../platform/tests/tools/test_run_all.py | 34 +++++++++++++++++-- 11 files changed, 77 insertions(+), 41 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp index 92a4e6a93396..b0a1c5e7ba94 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp @@ -160,6 +160,9 @@ rocke_status_t rocke_h_pack_fragment_bits(rocke_ir_builder_t* b, if(f.carrier_bits != 32 && f.carrier_bits != 64) ckc::raise_status(ROCKE_ERR_VALUE, "IR pattern packing currently requires i32 or i64 carriers"); + if(f.packing.element_bits > 32) + ckc::raise_status(ROCKE_ERR_VALUE, + "IR pattern packing supports encoded fields of at most 32 bits"); const auto* word_type = f.carrier_bits == 64 ? rocke_i64() : rocke_i32(); auto constant = [&](uint64_t value) { return f.carrier_bits == 64 ? rocke_b_const_i64(b, int64_t(value)) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.mma_io.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.mma_io.h index 7dea7828c6f6..7e70c670cf08 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.mma_io.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.mma_io.h @@ -33,6 +33,7 @@ rocke_value_t* rocke_h_load_matrix_fragment(rocke_ir_builder_t* b, /* Callback returns a canonical unsigned pattern in an integer carrier. * ctx and output pointer array are borrowed for the duration of this call. */ typedef rocke_value_t* (*rocke_load_bits_fn)(rocke_ir_builder_t* b, int index, void* ctx); +/* Encoded fields are at most 32 bits; i64 carriers support eight-byte scales. */ rocke_status_t rocke_h_pack_fragment_bits(rocke_ir_builder_t* b, rocke_load_bits_fn load_bits, void* ctx, 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 index 08585dd6b734..dcd0e389b07a 100644 --- 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 @@ -24,23 +24,12 @@ typedef struct rocke_scale_packing int block_k; } rocke_scale_packing_t; -typedef struct rocke_scale_association -{ - int block_k; -} rocke_scale_association_t; - -static inline rocke_scale_association_t rocke_scale_association(const rocke_scale_packing_t* p) -{ - if(p->block_k <= 0) - ckc::raise_status(ROCKE_ERR_VALUE, "scale block_k must be positive"); - return {p->block_k}; -} - static inline rocke_fragment_packing_t rocke_scale_fragment(const rocke_scale_packing_t* p) { rocke_bit_packing_t bits; rocke_fragment_packing_t result; - rocke_scale_association(p); + if(p->block_k <= 0) + ckc::raise_status(ROCKE_ERR_VALUE, "scale block_k must be positive"); if(p->count != 1 && p->count != 2 && p->count != 4 && p->count != 8) ckc::raise_status(ROCKE_ERR_VALUE, "carrier_bits must be 8, 16, 32, or 64"); if(!rocke_bit_packing_init(&bits, rocke_dtype_info("e8m0")->encoded_bits, 0) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md index aeb0a863af35..0508701c82fe 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md @@ -109,15 +109,16 @@ i32 carriers. FP8 occupies all sixteen words. FP4 occupies eight and pads eight. The FP6 transport descriptor occupies twelve and pads four, without enabling FP6 catalog entries or numerical conversions. -`ScaleAssociation` records `block_k`: how many source K elements share one -scale. This is separate from the bit layout. Existing `ScalePacking(count, -block_k)` remains a compatibility adapter, exposing both the association and a -shared `FragmentPacking`. +`ScalePacking(count, block_k)` records how many source K elements share one +scale in `block_k` and exposes the byte layout through a shared `FragmentPacking`. +Scale association and bit layout remain separate properties of this contract. Current scaled-WMMA scale fragments hold four eight-bit patterns in i32 for K32, or eight in i64 for K16. The first K group occupies the low byte. The A/B scale coordinate maps determine which scales each lane loads. The same `pack_fragment_bits` helper can pack six-bit fields that cross word boundaries. +Its IR contract accepts encoded fields of at most 32 bits in i32 or i64 carriers; +wide encoded fields are outside the current operand requirements. ## Mirroring and follow-up work 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 index 7f0284ed983c..a506ad183e71 100644 --- 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 @@ -11,17 +11,6 @@ from ..storage import BitPacking, FragmentPacking, MatrixFragmentLayout -@dataclass(frozen=True) -class ScaleAssociation: - """Number of source K elements sharing one scale.""" - - block_k: int - - def __post_init__(self) -> None: - if self.block_k <= 0: - raise ValueError("scale block_k must be positive") - - def scaled_matrix_layout(dtype: str, abi_words: int) -> MatrixFragmentLayout: """gfx1250 scaled operand layout, independent of atom availability. @@ -51,13 +40,10 @@ class ScalePacking: block_k: int def __post_init__(self) -> None: - self.association + if self.block_k <= 0: + raise ValueError("scale block_k must be positive") self.fragment - @property - def association(self) -> ScaleAssociation: - return ScaleAssociation(self.block_k) - @property def packing(self) -> BitPacking: return BitPacking(dtype_info("e8m0").encoded_bits) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py index 28bfdf60c316..6972121da6a3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py @@ -124,11 +124,16 @@ def pack_fragment_bits( ) -> list[Value]: """Pack unsigned encoded patterns into integer carriers, including split fields. - The loader must return canonical patterns with zero high bits. Tensor decoding - and numeric quantization are separate operations. + Encoded fields are at most 32 bits; i64 carriers support eight-byte scale + words. The loader must return canonical patterns with zero high bits. Tensor + decoding and numeric quantization are separate operations. """ if fragment.carrier_bits not in (32, 64): raise ValueError("IR pattern packing currently requires i32 or i64 carriers") + if fragment.packing.element_bits > 32: + raise ValueError( + "IR pattern packing supports encoded fields of at most 32 bits" + ) word_type = I64 if fragment.carrier_bits == 64 else I32 constant = b.const_i64 if fragment.carrier_bits == 64 else b.const_i32 words = [constant(0) for _ in range(fragment.carrier_count)] diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py index 67caa9541dcf..c084230d9773 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py @@ -308,3 +308,12 @@ def test_quant_type_diagnostic_matches_native(): text=True, ) assert result.stdout == str(error.value) + + +def test_reject_wide_encoded_fields_before_emission(): + b = IRBuilder("unsupported_packing") + fragment = FragmentPacking(BitPacking(40), 2, 64, 2) + before = len(b.kernel.body.ops) + with pytest.raises(ValueError, match="encoded fields of at most 32 bits"): + pack_fragment_bits(b, lambda j: b.const_i64(1 << 35), fragment) + assert len(b.kernel.body.ops) == before diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp index de507eba1c34..c224e5d9a24a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp @@ -229,6 +229,20 @@ int main(int argc, char** argv) rocke_ir_builder_free(&b); return 0; } + { + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "unsupported_packing") == ROCKE_OK); + rocke_bit_packing_t packing; + rocke_fragment_packing_t fragment; + rocke_value_t* words[2]; + CHECK(rocke_bit_packing_init(&packing, 40, 0)); + CHECK(rocke_fragment_packing_init(&fragment, &packing, 2, 64, 2)); + CHECK(rocke_h_pack_fragment_bits(&b, load_bits, NULL, &fragment, words, 2) + == ROCKE_ERR_VALUE); + CHECK(strstr(rocke_ir_builder_error(&b), "encoded fields of at most 32 bits")); + CHECK(b.kernel->body->num_ops == 0); + rocke_ir_builder_free(&b); + } CHECK(test_fragment_inputs() == 0); CHECK(test_fragment_offset_bounds() == 0); CHECK(rocke_dtype_info("e4m3") == rocke_dtype_info("fp8e4m3")); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py index 5284914cc389..cb572e5af990 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py @@ -58,7 +58,7 @@ def test_matrix_payload_and_padding(dtype, live): ) def test_scales_use_common_bit_packing(count, block_k, word): scales = ScalePacking(count, block_k) - assert scales.association.block_k == block_k + assert scales.block_k == block_k assert scales.fragment.pack(list(range(1, count + 1))) == (word,) assert scales.packing == BitPacking(8) assert scales.word_bits == count * 8 diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/run_all.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/run_all.py index 7f4b90c6124e..51e12f36ea9a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/run_all.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/run_all.py @@ -116,7 +116,7 @@ def ctest_ready(build_root: Path, config: str) -> bool: def native_pytest_env(build_root: Path, config: str) -> dict[str, str]: - """Make the native storage comparisons available to both pytest passes.""" + """Build the configured suite and supply native parity to both pytest passes.""" env = dict(os.environ) executable = env.get("ROCKE_STORAGE_TEST") if not executable and (build_root / "CMakeCache.txt").is_file(): @@ -127,8 +127,6 @@ def native_pytest_env(build_root: Path, config: str) -> dict[str, str]: str(build_root), "--config", config, - "--target", - "rocke_storage", ], check=True, ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/tools/test_run_all.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/tools/test_run_all.py index d9dbd4f08bdd..484503960776 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/tools/test_run_all.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/tools/test_run_all.py @@ -6,6 +6,7 @@ import json from pathlib import Path import subprocess +import shutil import pytest @@ -67,8 +68,6 @@ def run(command, **kwargs): str(tmp_path), "--config", "Debug", - "--target", - "rocke_storage", ] assert any(cmd[:3] == ["ctest", "-C", "Debug"] for cmd, _ in calls) @@ -188,3 +187,34 @@ def run(command, **kwargs): ) assert runner.main() == 1 assert all(command[0] != runner.sys.executable for command in calls) + + +def test_fresh_build_prepares_entire_ctest_suite(runner, tmp_path): + if not shutil.which("cmake") or not shutil.which("ctest"): + pytest.skip("CMake and CTest required for fresh-build regression") + source = tmp_path / "source" + source.mkdir() + (source / "main.c").write_text("int main(void) { return 0; }\n") + (source / "CMakeLists.txt").write_text( + "cmake_minimum_required(VERSION 3.20)\n" + "project(runner_fixture C)\n" + "enable_testing()\n" + "add_library(rocke_core STATIC main.c)\n" + "add_executable(rocke_storage main.c)\n" + "add_executable(rocke_dtypes main.c)\n" + "add_test(NAME rocke_storage COMMAND rocke_storage)\n" + "add_test(NAME rocke_dtypes COMMAND rocke_dtypes)\n" + ) + build = tmp_path / "build" + subprocess.run(["cmake", "-S", str(source), "-B", str(build)], check=True) + # Reproduce the byte-identity gate's partial build before pytest setup. + subprocess.run( + ["cmake", "--build", str(build), "--target", "rocke_core"], check=True + ) + env = runner.native_pytest_env(build, "Release") + assert Path(env["ROCKE_STORAGE_TEST"]).is_file() + assert runner.ctest_ready(build, "Release") + subprocess.run( + ["ctest", "--test-dir", str(build), "-C", "Release", "--output-on-failure"], + check=True, + ) From 2732998629265d810793a677f4974733f7c00fea Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:33:39 -0400 Subject: [PATCH 11/24] docs(rocke): clarify complete native suite setup in test runner --- .../hip-kernel-provider/rocke/platform/tests/README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/README.md b/dnn-providers/hip-kernel-provider/rocke/platform/tests/README.md index 870e008c00f7..82ffdd97c593 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/README.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/README.md @@ -109,9 +109,9 @@ divergence - the Python builder correctly rejects wave32 WMMA on gfx942.) ### Native storage parity in the standard runner -`run_all.py --build-root ` builds the `rocke_storage` target in a configured -build and obtains its executable path from CTest. Both pytest passes receive that -path, so storage IR/HIP parity and serialization tests run automatically. +`run_all.py --build-root ` builds all configured targets before pytest, +then obtains the `rocke_storage` executable path from CTest. Both pytest passes +receive that path, so storage IR/HIP parity and serialization tests run automatically. `--config` selects the native test configuration (default `Release`). A build or fixture-discovery failure stops the runner instead of silently skipping coverage. From af73df5ea70b7eaff40602a476ceeec755540ae6 Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Thu, 24 Sep 2026 21:43:09 -0400 Subject: [PATCH 12/24] feat(rocke): support contiguous 96-bit global and LDS loads --- .../rocke/platform/cpp/core/ir/ir_mem.cpp | 12 ++-- .../rocke/platform/cpp/core/ir/ir_tile.cpp | 10 +-- .../cpp/core/lower_hip/lower_hip_mem.cpp | 19 +++++- .../platform/cpp/core/lower_llvm/mem.cpp | 3 +- .../architecture/dtype_storage_packing.md | 8 +++ .../rocke/platform/python/rocke/core/ir.py | 41 ++++++------ .../platform/python/rocke/core/lower_hip.py | 14 +++- .../platform/python/rocke/core/lower_llvm.py | 3 +- .../rocke/platform/tests/core/test_mma_io.py | 64 ++++++++++++++++++- .../platform/tests/core/test_storage.cpp | 36 ++++++++++- 10 files changed, 169 insertions(+), 41 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp index 2d2f89d72065..b07e81cb6b16 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp @@ -227,14 +227,14 @@ rocke_value_t* rocke_b_global_load_vN(rocke_ir_builder_t* b, || rocke_i_type_is(dtype, "i16")) { elem_bytes = 2; - if(n != 2 && n != 4 && n != 8 && n != 16) + if(n != 2 && n != 4 && n != 6 && n != 8 && n != 16) return (rocke_value_t*)rocke_i_set_err( b, ROCKE_ERR_VALUE, "unsupported vector width for global_load_vN: %d", n); } else if(rocke_i_type_is(dtype, "f32") || rocke_i_type_is(dtype, "i32")) { elem_bytes = 4; - if(n != 2 && n != 4 && n != 8) + if(n != 2 && n != 3 && n != 4 && n != 8) return (rocke_value_t*)rocke_i_set_err( b, ROCKE_ERR_VALUE, "unsupported vector width for %s global_load_vN: %d", en, n); } @@ -242,7 +242,7 @@ rocke_value_t* rocke_b_global_load_vN(rocke_ir_builder_t* b, || rocke_i_type_is(dtype, "i8")) { elem_bytes = 1; - if(n != 2 && n != 4 && n != 8 && n != 16) + if(n != 2 && n != 4 && n != 8 && n != 12 && n != 16) return (rocke_value_t*)rocke_i_set_err( b, ROCKE_ERR_VALUE, "unsupported vector width for %s global_load_vN: %d", en, n); } @@ -262,7 +262,11 @@ rocke_value_t* rocke_b_global_load_vN(rocke_ir_builder_t* b, a = rocke_i_attrs(b); rocke_attr_set_str(b, &a, "elem_type", en); rocke_attr_set_int(b, &a, "vec", (int64_t)n); - rocke_attr_set_int(b, &a, "align", (int64_t)(align > 0 ? align : n * elem_bytes)); + rocke_attr_set_int( + b, + &a, + "align", + (int64_t)(align > 0 ? align : (n * elem_bytes == 12 ? elem_bytes : n * elem_bytes))); { char hint[16]; /* result_name_hint = "gv{n}" */ 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 8ebb8bb6191e..0585df2a5aab 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 @@ -349,8 +349,9 @@ rocke_value_t* rocke_b_smem_load_vN(rocke_ir_builder_t* b, const rocke_type_t* vt; rocke_attr_map_t attrs; const char* dn; - static const int allowed_8bit[] = {1, 2, 4, 8, 16}; - static const int allowed_other[] = {1, 2, 4, 8}; + static const int allowed_8bit[] = {1, 2, 4, 8, 12, 16}; + static const int allowed_16bit[] = {1, 2, 4, 6, 8}; + static const int allowed_32bit[] = {1, 2, 3, 4, 8}; char hint[16]; if(!rocke_i_live(b)) { @@ -375,8 +376,9 @@ rocke_value_t* rocke_b_smem_load_vN(rocke_ir_builder_t* b, { bool eight = (strcmp(dn, "fp8e4m3") == 0 || strcmp(dn, "bf8e5m2") == 0 || strcmp(dn, "i8") == 0); - const int* allowed = eight ? allowed_8bit : allowed_other; - int acount = eight ? 5 : 4; + bool half = strcmp(dn, "f16") == 0 || strcmp(dn, "bf16") == 0; + const int* allowed = eight ? allowed_8bit : half ? allowed_16bit : allowed_32bit; + int acount = eight ? 6 : 5; if(!rocke_n_in(n, allowed, acount)) { return (rocke_value_t*)rocke_i_set_err( diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp index cdfee832aad9..60129168ffcd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp @@ -335,6 +335,21 @@ static rocke_status_t _op_tile_smem_load_vN(rocke_h_lowerer_t* lw, const rocke_o } idx_str = mem_idx_join(lw, &op->operands[1], op->num_operands - 1); res = rocke_h_name(lw, op->results[0]); + const int64_t byte_count = n * (rocke_dtype_info(elem_name)->encoded_bits / 8); + if(byte_count & (byte_count - 1)) + { + /* Clang pads vector objects; copy only the actual LDS payload. */ + rocke_h_emitf(lw, + "%s%lld %s; __builtin_memcpy(&%s, &%s[%s], %lld);", + prefix, + (long long)n, + res, + res, + storage, + idx_str, + (long long)byte_count); + return lw->status; + } rocke_h_emitf(lw, "%s%lld %s = *reinterpret_cast(&%s[%s]);", prefix, @@ -536,9 +551,9 @@ static rocke_status_t _op_memref_global_load_vN(rocke_h_lowerer_t* lw, const roc if(align <= 0 || (align & (align - 1))) return rocke_h_fail( lw, ROCKE_ERR_VALUE, "global_load_vN: alignment must be a positive power of two"); - if(align < byte_count) + if(align < byte_count || (byte_count & (byte_count - 1))) { - /* Do not strengthen the IR's alignment through a vector-pointer cast. */ + /* Copy only the payload, not vector padding, with the IR's alignment. */ rocke_h_emitf( lw, "%s%lld %s; __builtin_memcpy(&%s, __builtin_assume_aligned(%s + %s, %lld), %lld);", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp index 1bb057802dab..5cf265459213 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp @@ -570,7 +570,8 @@ static void op_tile_smem_load_vN(rocke_lower_t* L, const rocke_op_t* op) elem_bytes = 8; } } - int64_t align = vec * elem_bytes; + /* New 96-bit widths guarantee only element alignment, including FP8. */ + int64_t align = (vec == 3 || vec == 6 || vec == 12) ? 12 / vec : vec * elem_bytes; /* gfx1250: vec==8 loads are marked volatile to block the WMMA-aware backend * pass from substituting ds_load_tr16_b128 (transposed) for the plain * sequential ds_read_b128. Mirrors Python _op_tile_smem_load_vN lines diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md index 0508701c82fe..5d15503e83d1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md @@ -86,6 +86,14 @@ origin alignment, carrier capacity, and static i32 displacement limits. It reduces the supplied alignment for chunk spacing and K origin. A 24-byte FP6 chunk becomes 16-byte and 8-byte loads without reading past the chunk. +The existing `global_load_vN` and `smem_load_vN` builders also accept contiguous +96-bit payloads: twelve byte elements, six 16-bit elements, or three 32-bit +elements of their supported types. These loads default to element alignment; +HIP copies exactly twelve bytes, excluding vector-object padding. A single +96-bit machine instruction depends on target and alignment. This does not +change the fragment loader's chunk selection. FP6 transpose loads are deferred +until tensor-descriptor integration defines their lane/layout contract. + Tensor shape, strides, row selection, and allocation bounds belong to the caller. This helper always loads a complete fragment: the caller must guard empty or partial rows and provide a valid lane group. It is not a bounds-checked 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 73f01b156cfc..0ebe58891330 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 @@ -1564,34 +1564,30 @@ def global_load_vN( ) -> Value: """Vectorised global load of N consecutive values. - Supports the full element-type catalog the LLVM lowering already - accepts: ``f16`` / ``bf16`` (N in {2, 4, 8}), ``f32`` / ``i32`` - (N in {2, 4, 8}), ``i16`` (N in {2, 4, 8}), ``fp8e4m3`` / - ``bf8e5m2`` / ``i8`` (N in {2, 4, 8, 16}). - - Lowers to a single ``load `` from ``addrspace(1)``; - AMDGPU's backend coalesces these into a single VMEM transaction - (``global_load_dwordxN``) when the address is naturally aligned. - - The per-element size is folded into the default alignment so the - common case (8 fp8 → 8-byte load, 4 f32 → 16-byte load) does - not need an explicit ``align=`` kwarg. + Supports f16/bf16/i16 (N in {2, 4, 6, 8, 16}), f32/i32 + (N in {2, 3, 4, 8}), and fp8e4m3/bf8e5m2/i8 (N in {2, 4, 8, 12, 16}). + Loads exactly N elements. Instruction selection depends on target and + alignment; 96-bit payloads do not require a 96-bit scalar type. + + Default alignment is the payload size for power-of-two loads, and + element alignment for 12-byte loads. An explicit alignment is a caller + guarantee about the address after adding idx. """ if dtype.name in ("f16", "bf16", "i16"): elem_bytes = 2 # n=16 (32-byte `global_load_dwordx8`) is needed for the RDNA WMMA # <16 x half> operand fragment; AMDGPU coalesces it when aligned. - if n not in (2, 4, 8, 16): + if n not in (2, 4, 6, 8, 16): raise ValueError(f"unsupported vector width for global_load_vN: {n}") elif dtype.name in ("f32", "i32"): elem_bytes = 4 - if n not in (2, 4, 8): + if n not in (2, 3, 4, 8): raise ValueError( f"unsupported vector width for {dtype.name} global_load_vN: {n}" ) elif dtype.name in ("fp8e4m3", "bf8e5m2", "i8"): elem_bytes = 1 - if n not in (2, 4, 8, 16): + if n not in (2, 4, 8, 12, 16): raise ValueError( f"unsupported vector width for {dtype.name} global_load_vN: {n}" ) @@ -1607,7 +1603,9 @@ def global_load_vN( attrs={ "elem_type": dtype.name, "vec": n, - "align": int(align or (n * elem_bytes)), + "align": int( + align or (elem_bytes if n * elem_bytes == 12 else n * elem_bytes) + ), }, result_name_hint=f"gv{n}", ).result @@ -1814,10 +1812,9 @@ def smem_load_vN_f16(self, smem: Value, *indices, n: int = 0) -> Value: def smem_load_vN(self, smem: Value, *indices, dtype: Type, n: int = 0) -> Value: """LDS load of ````. Supports 8-bit (fp8e4m3 / bf8e5m2 / i8), 16-bit (f16 / bf16) and 32-bit (f32 / i32) element types; - AMDGPU lowers vector LDS loads to ``ds_read_b{8, 16, 32, 64, 128}`` - based on total payload size. The 8-bit variants must use ``n in {1, - 2, 4, 8, 16}`` so the resulting payload still maps to a single - ``ds_read_b*`` instruction (n=16 → ds_read_b128). + loads exactly N elements. In addition to power-of-two widths, accepts + 96-bit payloads (12 bytes, six halfwords, or three words), using element + alignment. The target and alignment determine instruction selection. """ if dtype.name not in ("f16", "bf16", "f32", "i32", "fp8e4m3", "bf8e5m2", "i8"): raise ValueError( @@ -1825,9 +1822,9 @@ def smem_load_vN(self, smem: Value, *indices, dtype: Type, n: int = 0) -> Value: f"bf8e5m2 / i8, got {dtype.name}" ) allowed_n = ( - (1, 2, 4, 8, 16) + (1, 2, 4, 8, 12, 16) if dtype.name in ("fp8e4m3", "bf8e5m2", "i8") - else (1, 2, 4, 8) + else (1, 2, 4, 6, 8) if dtype.name in ("f16", "bf16") else (1, 2, 3, 4, 8) ) if n not in allowed_n: raise ValueError( 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 1352f07217ff..18ecd1a9b322 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 @@ -493,9 +493,9 @@ def _op_memref_global_load_vN(self, op: Op) -> None: raise ValueError( "global_load_vN: alignment must be a positive power of two" ) - if align < byte_count: - # A vector-pointer dereference would promise natural alignment. - # Copy the payload using only the alignment guaranteed by the IR. + if align < byte_count or byte_count & (byte_count - 1): + # Non-power-of-two vector objects include padding. Copy only the + # payload, using only the alignment guaranteed by the IR. self._emit( f"{prefix}{vec} {_name(op.result)}; " f"__builtin_memcpy(&{_name(op.result)}, " @@ -1839,6 +1839,14 @@ def _op_tile_smem_load_vN(self, op: Op) -> None: if storage is None: raise RuntimeError("smem load_vN before smem_alloc was lowered") idx_str = "][".join(_name(i) for i in indices) + byte_count = n * (dtype_info(elem_name).encoded_bits // 8) + if byte_count & (byte_count - 1): + # Clang rounds vector object sizes up; LDS payloads have no padding. + self._emit( + f"{prefix}{n} {_name(op.result)}; " + f"__builtin_memcpy(&{_name(op.result)}, &{storage}[{idx_str}], {byte_count});" + ) + return self._emit( f"{prefix}{n} {_name(op.result)} = " f"*reinterpret_cast(&{storage}[{idx_str}]);" 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 4743a1335ea8..ebf67fd5af12 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 @@ -3101,7 +3101,8 @@ def _op_tile_smem_load_vN(self, op: Op) -> None: op.result.type.elem.name, 2, # type: ignore[attr-defined] ) - align = vec * elem_bytes + # New 96-bit widths guarantee only element alignment, including FP8. + align = 12 // vec if vec in (3, 6, 12) else vec * elem_bytes # gfx1250: mark 8-wide (128-bit) LDS loads volatile to block the WMMA-aware # pass from substituting ds_load_tr16_b128 (transposed) in place of the plain # sequential ds_read_b128. Only 8-wide loads feed the 16x16x32 WMMA fragment diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py index c084230d9773..8b2897f0edef 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py @@ -11,7 +11,7 @@ import pytest from rocke.core.arch.wmma_scale import scaled_matrix_layout -from rocke.core.ir import FP8E4M3, I8, I32, I64, IRBuilder, PtrType, VectorType +from rocke.core.ir import F16, FP8E4M3, I8, I32, I64, IRBuilder, PtrType, VectorType from rocke.core.ir_serialize import parse, serialize from rocke.core.lower_llvm import lower_kernel_to_llvm from rocke.core.lower_hip import lower_kernel_to_hip @@ -28,6 +28,9 @@ CASES = [ + "load96_i8", + "load96_f16", + "load96_i32", "fp4", "fp6", "fp6_padded", @@ -47,6 +50,28 @@ def build_transport(dtype): b = IRBuilder("transport") + if dtype.startswith("load96_"): + unit, n = { + "load96_i8": (I8, 12), + "load96_f16": (F16, 6), + "load96_i32": (I32, 3), + }[dtype] + a = b.param("A", PtrType(unit, "global")) + o = b.param("O", PtrType(unit, "global")) + one = b.const_i32(1) + value = b.global_load_vN(a, one, unit, n) + smem = b.smem_alloc(unit, (n + 1,), name_hint="payload") + for j in range(n): + b.smem_store_vN(smem, [b.const_i32(j + 1)], b.vec_extract(value, j), 1) + b.s_barrier_bare() + value = b.smem_load_vN(smem, one, dtype=unit, n=n) + aligned = b.global_load_vN(a, b.const_i32(0), unit, n, align=16) + for k, vector in enumerate((value, aligned)): + for j in range(n): + b.global_store( + o, b.const_i32(k * n + j), b.vec_extract(vector, j), align=12 // n + ) + return b.kernel padded = dtype.endswith("_padded") dtype = dtype.removesuffix("_padded") patterns = dtype.startswith("pack_") @@ -317,3 +342,40 @@ def test_reject_wide_encoded_fields_before_emission(): with pytest.raises(ValueError, match="encoded fields of at most 32 bits"): pack_fragment_bits(b, lambda j: b.const_i64(1 << 35), fragment) assert len(b.kernel.body.ops) == before + + +@pytest.mark.parametrize("dtype", ["load96_i8", "load96_f16", "load96_i32"]) +def test_96bit_loads_copy_exact_payload(dtype): + """Offset and over-aligned global loads plus an LDS tail must not read padding.""" + kernel = build_transport(dtype) + llvm = lower_kernel_to_llvm(kernel, arch="gfx1250", llvm_flavor="llvm23") + n, elem, alignment = { + "load96_i8": (12, "i8", 1), + "load96_f16": (6, "half", 2), + "load96_i32": (3, "i32", 4), + }[dtype] + for space in (1, 3): + assert re.search( + rf"load <{n} x {elem}>, ptr addrspace\({space}\).*align {alignment}\b", llvm + ) + hip = lower_kernel_to_hip(kernel, arch="gfx1250") + assert len(re.findall(r"__builtin_memcpy\([^\n]*, 12\);", hip)) == 3 + if not shutil.which("hipcc"): + pytest.skip("hipcc not in PATH") + from rocke.helpers.compile import emit_device_llvm_ir_via_hipcc + + compiled = emit_device_llvm_ir_via_hipcc(kernel, arch="gfx950", extra_flags=["-O0"]) + # Clang's vector object is 16 bytes; the source copy must remain 12 bytes. + compiled = compiled.split("define protected amdgpu_kernel void @transport(", 1)[ + 1 + ].split("\n}", 1)[0] + source_copies = re.findall( + r"@llvm\.memcpy[^\n]*?\([^,]+,\s*ptr[^,]*\balign (\d+) [^,]+,\s*i64 (\d+)", + compiled, + ) + assert len(source_copies) == 3 + for (actual_align, size), max_align in zip( + source_copies, (alignment, alignment, 16) + ): + assert int(size) == 12 + assert int(actual_align) <= max_align diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp index c224e5d9a24a..716855565fb7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp @@ -40,15 +40,45 @@ static int emit(const char* dtype, bool hip) : strcmp(dtype, "bf16_padded") == 0 ? "bf16" : "fp6"; bool patterns = strncmp(dtype, "pack_", 5) == 0; - const auto* unit = patterns ? rocke_i8() : rocke_storage_ir_type(dtype); + const bool load96 = strncmp(dtype, "load96_", 7) == 0; + const auto* unit = load96 ? (strcmp(dtype, "load96_i8") == 0 ? rocke_i8() + : strcmp(dtype, "load96_f16") == 0 ? rocke_f16() + : rocke_i32()) + : patterns ? rocke_i8() + : rocke_storage_ir_type(dtype); CHECK(unit); const bool typed = strcmp(dtype, "f16") == 0 || strcmp(dtype, "bf16") == 0; - const auto* carrier = typed ? unit : rocke_i32(); + const auto* carrier = (typed || load96) ? unit : rocke_i32(); if(strcmp(dtype, "pack_scale_bytes_i64") == 0) carrier = rocke_i64(); auto* a = rocke_b_param(&b, "A", rocke_ptr_type(&b, unit, "global"), NULL); auto* o = rocke_b_param(&b, "O", rocke_ptr_type(&b, carrier, "global"), NULL); - if(patterns) + if(load96) + { + const int n = unit == rocke_i8() ? 12 : unit == rocke_f16() ? 6 : 3; + auto* one = rocke_b_const_i32(&b, 1); + auto* value = rocke_b_global_load_vN(&b, a, one, unit, n, 0); + int shape[] = {n + 1}; + auto* smem = rocke_b_smem_alloc(&b, unit, shape, 1, "payload"); + for(int j = 0; j < n; ++j) + { + rocke_value_t* indices[] = {rocke_b_const_i32(&b, j + 1)}; + auto* element = rocke_b_vec_extract(&b, value, j); + rocke_b_smem_store_vN(&b, smem, indices, 1, element, 1); + } + rocke_b_s_barrier_bare(&b); + value = rocke_b_smem_load_vN(&b, smem, &one, 1, unit, n); + auto* aligned = rocke_b_global_load_vN(&b, a, rocke_b_const_i32(&b, 0), unit, n, 16); + rocke_value_t* vectors[] = {value, aligned}; + for(int k = 0; k < 2; ++k) + for(int j = 0; j < n; ++j) + { + auto* index = rocke_b_const_i32(&b, k * n + j); + auto* element = rocke_b_vec_extract(&b, vectors[k], j); + rocke_b_global_store(&b, o, index, element, 12 / n); + } + } + else if(patterns) { int bits = strcmp(dtype, "pack_fp6_cross_word") == 0 ? 6 : 8; int count = bits == 6 ? 16 : (strcmp(dtype, "pack_scale_bytes_i64") == 0 ? 8 : 4); From 33f4d963c4e2db2bb4268456e400004526d47362 Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Thu, 24 Sep 2026 22:25:40 -0400 Subject: [PATCH 13/24] fix(rocke): reject unsupported fragment padding before emission --- .../platform/python/rocke/helpers/mma_io.py | 6 +++--- .../rocke/platform/tests/core/test_mma_io.py | 17 ++++++++++++++--- 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py index 6972121da6a3..aa0a498b3c38 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py @@ -68,6 +68,9 @@ def load_matrix_fragment( raise ValueError("matrix carrier type width mismatch") if packing.payload_bits % packing.carrier_bits: raise ValueError("matrix fragment payload must occupy whole carriers") + padding = packing.carrier_count - packing.live_carriers + if padding and carrier_type != I32: + raise ValueError("padded matrix fragments currently require i32 carriers") chunk_units = layout.chunk_bytes // unit_bytes origin_bytes = origin_bits // 8 # Bound the last loaded unit across all chunks and lane groups, including k0. @@ -109,10 +112,7 @@ def load_matrix_fragment( payload_type = VectorType(carrier_type, packing.live_carriers) if payload.type != payload_type: payload = b.bitcast(payload, payload_type) - padding = packing.carrier_count - packing.live_carriers if padding: - if carrier_type != I32: - raise ValueError("padded matrix fragments currently require i32 carriers") payload = b.vec_concat(payload, b.vector_splat(b.const_i32(0), padding)) return payload diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py index 8b2897f0edef..b5099a2085a1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py @@ -168,11 +168,15 @@ def test_reject_bit_origin(): @pytest.mark.parametrize( "dtype,alignment,error", - [("fp6", 3, "positive power of two"), ("fp4", 1, "packing width mismatch")], + [ + ("fp6", 3, "positive power of two"), + ("fp4", 1, "packing width mismatch"), + ("f16", 2, "padded matrix fragments currently require i32 carriers"), + ], ) def test_fragment_input_contract(dtype, alignment, error): b = IRBuilder("invalid_fragment") - ptr = b.param("A", PtrType(I8, "global")) + ptr = b.param("A", PtrType(storage_ir_type(dtype), "global")) zero = b.const_i32(0) before = serialize(b.kernel) with pytest.raises(ValueError, match=error): @@ -183,7 +187,14 @@ def test_fragment_input_contract(dtype, alignment, error): zero, 0, dtype=dtype, - layout=scaled_matrix_layout("fp6", 16), + layout=( + MatrixFragmentLayout( + FragmentPacking(BitPacking(16), 8, 16, 16), 8, 1, 1 + ) + if dtype == "f16" + else scaled_matrix_layout("fp6", 16) + ), + carrier_type=F16 if dtype == "f16" else I32, alignment_bytes=alignment, ) assert serialize(b.kernel) == before From b03f86c95cac8122b875896cbd61dd479a29076c Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Thu, 24 Sep 2026 23:18:18 -0400 Subject: [PATCH 14/24] fix(rocke): preserve vector aliases in bare HIP output --- .../platform/cpp/core/lower_hip/core.cpp | 6 +-- .../platform/python/rocke/core/lower_hip.py | 8 ++-- .../rocke/platform/tests/core/test_mma_io.py | 43 ++++++++++++++++++- .../platform/tests/core/test_storage.cpp | 6 ++- 4 files changed, 54 insertions(+), 9 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp index 039b225b394e..de9b2d5cb110 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp @@ -1047,10 +1047,10 @@ rocke_status_t rocke_lower_kernel_to_hip(rocke_ir_builder_t* b, { rocke_strbuf_append(out, ROCKE_HIP_PROLOGUE); rocke_strbuf_append_char(out, '\n'); - for(int j = 0; j < kernel->num_params; ++j) - h_extra_vector_type(&lw, kernel->params[j]->type, out); - h_extra_vector_region(&lw, kernel->body, out); } + for(int j = 0; j < kernel->num_params; ++j) + h_extra_vector_type(&lw, kernel->params[j]->type, out); + h_extra_vector_region(&lw, kernel->body, out); /* head */ rocke_strbuf_appendf(out, "extern \"C\" __global__ __launch_bounds__(%d)\n" 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 18ecd1a9b322..ab29bc41586b 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 @@ -2513,9 +2513,9 @@ def lower_kernel_to_hip( The output is: 1. The :data:`HIP_PROLOGUE` (typedefs + ```` include + AMDGPU vector typedefs). Disable with - ``include_prologue=False`` when you want only the body text - (e.g. for embedding into a larger TU that already has these - typedefs). + ``include_prologue=False`` for embedding into a larger TU that already + has the shared prologue. Required per-kernel vector typedefs are still + emitted. 2. The kernel's ``__global__`` signature, derived from :attr:`KernelDef.params`. Pointer params get ``__restrict__``; ``__launch_bounds__`` is taken from @@ -2568,7 +2568,7 @@ def lower_kernel_to_hip( parts: List[str] = [] if include_prologue: parts.append(HIP_PROLOGUE) - parts.extend(_extra_vector_declarations(kernel)) + parts.extend(_extra_vector_declarations(kernel)) parts.append(head) if smem_block: parts.append(smem_block) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py index b5099a2085a1..9821489b72c5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py @@ -14,7 +14,7 @@ from rocke.core.ir import F16, FP8E4M3, I8, I32, I64, IRBuilder, PtrType, VectorType from rocke.core.ir_serialize import parse, serialize from rocke.core.lower_llvm import lower_kernel_to_llvm -from rocke.core.lower_hip import lower_kernel_to_hip +from rocke.core.lower_hip import HIP_PROLOGUE, lower_kernel_to_hip from rocke.core.storage import ( BitPacking, FragmentPacking, @@ -248,6 +248,47 @@ def test_hip_declares_only_encountered_missing_vector_widths(): assert "using i8x24 =" not in source +@pytest.mark.parametrize("route", ["python", "native"]) +def test_bare_hip_keeps_required_vector_declarations(route, tmp_path): + source = lower_kernel_to_hip( + build_transport("load96_i8"), arch="gfx1250", include_prologue=False + ) + if route == "native": + executable = os.environ.get("ROCKE_STORAGE_TEST") + if not executable: + pytest.skip("set ROCKE_STORAGE_TEST to the built native storage test") + actual = subprocess.run( + [executable, "--hip-bare", "load96_i8"], + check=True, + capture_output=True, + text=True, + ).stdout + assert actual == source + source = actual + assert "#include" not in source + assert source.count("using i8x12 =") == 1 + if not shutil.which("hipcc"): + pytest.skip("hipcc not in PATH") + path = tmp_path / "composed.hip" + path.write_text( + HIP_PROLOGUE + + source + + source.replace("void transport(", "void transport_second(") + ) + subprocess.run( + [ + "hipcc", + "--offload-arch=gfx950", + "--cuda-device-only", + "-fsyntax-only", + str(path), + ], + check=True, + capture_output=True, + text=True, + ) + + @pytest.mark.parametrize("dtype", ["fp8", "bf8"]) def test_byte_width_does_not_erase_pointer_identity(dtype): assert storage_ir_type(dtype) != I8 diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp index 716855565fb7..53c530399a0d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp @@ -30,7 +30,7 @@ static rocke_value_t* load_bits(rocke_ir_builder_t* b, int j, void* ctx) return rocke_b_global_load(b, (rocke_value_t*)ctx, rocke_b_const_i32(b, j), rocke_i8(), 1); } -static int emit(const char* dtype, bool hip) +static int emit(const char* dtype, bool hip, bool include_prologue = true) { rocke_ir_builder_t b; CHECK(rocke_ir_builder_init(&b, "transport") == ROCKE_OK); @@ -127,6 +127,8 @@ static int emit(const char* dtype, bool hip) CHECK(rocke_strbuf_init(&text, 256) == 0); rocke_lower_hip_opts_t opts = {}; opts.arch = "gfx1250"; + opts.include_prologue = include_prologue; + opts.include_prologue_set = true; CHECK(rocke_lower_kernel_to_hip(&b, b.kernel, &opts, &text) == ROCKE_OK); fputs(rocke_strbuf_cstr(&text), stdout); rocke_strbuf_free(&text); @@ -229,6 +231,8 @@ int main(int argc, char** argv) return emit(argv[2], false); if(argc == 3 && strcmp(argv[1], "--hip") == 0) return emit(argv[2], true); + if(argc == 3 && strcmp(argv[1], "--hip-bare") == 0) + return emit(argv[2], true, false); if(argc == 3 && strcmp(argv[1], "--parse") == 0) { rocke_ir_builder_t b; From e47c76de700871a801e12a04b282832bd9dbd268 Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Fri, 25 Sep 2026 00:19:47 -0400 Subject: [PATCH 15/24] fix(rocke): load packed matrix fragments into word carriers --- .../rocke/platform/cpp/helpers/mma_io.cpp | 26 +++++++++++++---- .../architecture/dtype_storage_packing.md | 7 +++++ .../platform/python/rocke/helpers/mma_io.py | 28 +++++++++++++++---- .../rocke/platform/tests/core/test_mma_io.py | 23 +++++++++++++-- .../platform/tests/core/test_storage.cpp | 2 ++ 5 files changed, 74 insertions(+), 12 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp index b0a1c5e7ba94..b19383b4668f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp @@ -90,6 +90,11 @@ rocke_value_t* rocke_h_load_matrix_fragment(rocke_ir_builder_t* b, / layout->lane_groups) ckc::raise_status(ROCKE_ERR_VALUE, "matrix fragment offset exceeds i32 range"); int alignment = gcd(alignment_bytes, gcd(chunk_bytes, origin_bytes)); + // Use 2--4-word vectors; retain storage-element loads for a one-word tail. + const bool word_loads = unit_bytes == 1 && rocke_type_eq(carrier_type, rocke_i32()) + && chunk_bytes % 4 == 0 && chunk_bytes % 16 != 4; + const auto* load_type = word_loads ? rocke_i32() : unit_type; + const int load_step = word_loads ? 4 : 1; auto* lane_chunk = rocke_b_mul(b, lane_group, rocke_b_const_i32(b, chunk_units)); auto* step_base = rocke_b_add(b, row_base, rocke_b_const_i32(b, origin_bytes / unit_bytes)); /* Collect loads before concatenation to preserve Python SSA numbering. */ @@ -112,16 +117,27 @@ rocke_value_t* rocke_h_load_matrix_fragment(rocke_ir_builder_t* b, int width = 1; while(width < max_width && uint64_t(width * 2) <= remaining) width *= 2; + if(word_loads) + width = remaining / 4 < 4 ? int(remaining / 4) : 4; auto* at = consumed ? rocke_b_add(b, offset, rocke_b_const_i32(b, consumed)) : offset; int load_align = gcd(alignment, consumed * unit_bytes); + // Form the byte address before loading words; preserve signed i32 offsets. + auto* load_ptr + = word_loads ? rocke_b_global_ptr_add(b, ptr, rocke_b_sext(b, at, rocke_i64())) + : ptr; + auto* load_at = word_loads ? rocke_b_const_i32(b, 0) : at; auto* value - = width == 1 ? rocke_b_vector_splat( - b, rocke_b_global_load(b, ptr, at, unit_type, load_align), 1) - : rocke_b_global_load_vN(b, ptr, at, unit_type, width, load_align); + = width == 1 + ? rocke_b_vector_splat( + b, + rocke_b_global_load(b, load_ptr, load_at, load_type, load_align), + 1) + : rocke_b_global_load_vN( + b, load_ptr, load_at, load_type, width, load_align); chunks[num_chunks++] = value; - consumed += width; - remaining -= width; + consumed += width * load_step; + remaining -= width * load_step; } } auto* payload = chunks[0]; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md index 5d15503e83d1..0180817073e3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md @@ -85,6 +85,13 @@ The loader checks dtype/packing width, pointer identity, whole-unit chunk and origin alignment, carrier capacity, and static i32 displacement limits. It reduces the supplied alignment for chunk spacing and K origin. A 24-byte FP6 chunk becomes 16-byte and 8-byte loads without reading past the chunk. +For byte-stored operands with i32 carriers, the loader selects up to four words +per load. It produces i32 vectors directly when every load has at least two +words. A 12-byte chunk produces three i32 lanes (`i32x3` in HIP), with the +original byte-address alignment. Address calculation remains in bytes, and the +loaded words concatenate into the operand fragment without intermediate byte +vectors. Typed chunks, partial words, and one-word tails retain their storage +element loads and final carrier bitcast when needed. The existing `global_load_vN` and `smem_load_vN` builders also accept contiguous 96-bit payloads: twelve byte elements, six 16-bit elements, or three 32-bit diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py index aa0a498b3c38..ebcfc0e21a42 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py @@ -82,6 +82,15 @@ def load_matrix_fragment( ): raise ValueError("matrix fragment offset exceeds i32 range") alignment = gcd(alignment_bytes, layout.chunk_bytes, origin_bytes) + # Use 2--4-word vectors; retain storage-element loads for a one-word tail. + word_loads = ( + unit_bytes == 1 + and carrier_type == I32 + and layout.chunk_bytes % 4 == 0 + and layout.chunk_bytes % 16 != 4 + ) + load_type = I32 if word_loads else unit_type + load_step = 4 if word_loads else 1 lane_chunk = b.mul(lane_group, b.const_i32(chunk_units)) step_base = b.add(row_base, b.const_i32(origin_bytes // unit_bytes)) chunks = [] @@ -94,18 +103,27 @@ def load_matrix_fragment( consumed = 0 max_width = 8 if unit_bytes == 4 else 16 while remaining: - width = min(max_width, 1 << (remaining.bit_length() - 1)) + width = ( + min(4, remaining // 4) + if word_loads + else min(max_width, 1 << (remaining.bit_length() - 1)) + ) at = b.add(offset, b.const_i32(consumed)) if consumed else offset load_align = gcd(alignment, consumed * unit_bytes) + # Form the byte address before loading words; preserve signed i32 offsets. + load_ptr = b.global_ptr_add(ptr, b.sext(at, I64)) if word_loads else ptr + load_at = b.const_i32(0) if word_loads else at if width == 1: value = b.vector_splat( - b.global_load(ptr, at, unit_type, align=load_align), 1 + b.global_load(load_ptr, load_at, load_type, align=load_align), 1 ) else: - value = b.global_load_vN(ptr, at, unit_type, width, align=load_align) + value = b.global_load_vN( + load_ptr, load_at, load_type, width, align=load_align + ) chunks.append(value) - consumed += width - remaining -= width + consumed += width * load_step + remaining -= width * load_step payload = chunks[0] for chunk in chunks[1:]: payload = b.vec_concat(payload, chunk) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py index 9821489b72c5..59e1bbf1b433 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py @@ -98,6 +98,8 @@ def build_transport(dtype): if typed else scaled_matrix_layout(dtype, 16) ) + if dtype == "fp6" and padded: + layout = MatrixFragmentLayout(layout.fragment, 16, 2, 16) base = b.const_i32((129 if typed else 97) if padded else 0) lane = b.mod(b.thread_id_x(), b.const_i32(32)) group = b.div(lane, b.const_i32(16)) @@ -239,9 +241,24 @@ def test_fragment_offset_bounds(dtype, count, groups, k0, valid): def test_hip_declares_only_encountered_missing_vector_widths(): + for case, widths in (("fp6", [4, 2, 4, 2]), ("fp6_padded", [3, 3, 3, 3])): + kernel = build_transport(case) + loads = [op for op in kernel.body.ops if op.name == "memref.global_load_vN"] + assert [op.result.type for op in loads] == [VectorType(I32, n) for n in widths] + assert not any(op.name == "arith.bitcast" for op in kernel.body.ops) + for count in (4, 20): + b = IRBuilder("scalar_tail") + ptr = b.param("A", PtrType(I8, "global")) + zero = b.const_i32(0) + layout = MatrixFragmentLayout( + FragmentPacking(BitPacking(8), count, 32, count // 4), count, 1, 1 + ) + load_matrix_fragment(b, ptr, zero, zero, 0, dtype="e8m0", layout=layout) + assert not any(op.name == "tile.global_ptr_add" for op in b.kernel.body.ops) source = lower_kernel_to_hip(build_transport("fp6"), arch="gfx1250") - for name in ("i8x24", "i8x40", "i32x12"): + for name in ("i32x6", "i32x10", "i32x12"): assert source.count(f"using {name} =") == 1 + assert "using i8x" not in source assert "using f16x32 =" not in source source = lower_kernel_to_hip(build_transport("f16"), arch="gfx1250") assert source.count("using f16x32 =") == 1 @@ -319,7 +336,7 @@ def test_hip_rejects_unsupported_vector_element(elem_type): @pytest.mark.parametrize( "dtype,copy_bytes,alignment", - [("fp6", 16, 8), ("fp6_padded", 16, 1), ("f16", 32, 16)], + [("fp6", 16, 8), ("fp6_padded", 12, 1), ("f16", 32, 16)], ) def test_hip_compiler_preserves_fragment_load_alignment(dtype, copy_bytes, alignment): """Inspect compiler IR: source parity cannot detect a shared alignment bug.""" @@ -339,6 +356,8 @@ def test_hip_compiler_preserves_fragment_load_alignment(dtype, copy_bytes, align if int(size) == copy_bytes ] assert source_alignments, "expected unaligned-safe fragment payload copies" + if dtype == "fp6_padded": + assert len(source_alignments) == 4 assert all(value <= alignment for value in source_alignments), source_alignments diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp index 53c530399a0d..d170d81d4dd2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp @@ -106,6 +106,8 @@ static int emit(const char* dtype, bool hip, bool include_prologue = true) } else layout = rocke_scaled_matrix_layout(dtype, 16); + if(strcmp(dtype, "fp6") == 0 && padded) + CHECK(rocke_matrix_fragment_layout_init(&layout, &layout.fragment, 16, 2, 16)); auto* base = rocke_b_const_i32(&b, padded ? (typed ? 129 : 97) : 0); auto* thread = rocke_b_thread_id_x(&b); auto* lane = rocke_b_mod(&b, thread, rocke_b_const_i32(&b, 32)); From 4c7816df47ed5ed6bc60b4e98fc5ae51dc2d27d6 Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Fri, 25 Sep 2026 02:26:05 -0400 Subject: [PATCH 16/24] fix(rocke): preserve override builds and scalar quant dtype contract --- .../platform/python/rocke/helpers/__init__.py | 2 ++ .../platform/python/rocke/helpers/quant.py | 11 ++++--- .../rocke/platform/tests/README.md | 6 ++-- .../rocke/platform/tests/core/test_dtypes.py | 16 ++++++++++ .../rocke/platform/tests/run_all.py | 21 +++++++----- .../platform/tests/tools/test_run_all.py | 32 ++++++++++++++++--- 6 files changed, 68 insertions(+), 20 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py index 0398a767bd09..4e015d6d7411 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py @@ -318,6 +318,7 @@ host_preshuffle_layout, ) from .quant import ( + LogicalQDType, QDType, QUANT_MAX_ABS, dequantize_scalar_to_f32, @@ -622,6 +623,7 @@ "mfma_attention_bwd_dq_dk_dv_inner_body", "streamk_num_macro_tiles", # Quantisation (f32 <-> {i8, fp8e4m3, bf8e5m2}) + "LogicalQDType", "QDType", "QUANT_MAX_ABS", "dequantize_scalar_to_f32", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/quant.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/quant.py index eb4604c48809..4b59538287c2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/quant.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/quant.py @@ -21,9 +21,10 @@ This module exposes the small set of helpers every quantised op needs: * :class:`QDType` — the output dtype literal alias used in spec dataclasses. +* :class:`LogicalQDType` — recognized types, including FP4/FP6 without scalar conversion. * :data:`QUANT_MAX_ABS` — clamp bound per output type (matches CK Tile's ``ComputeDataType{127.0f}`` / ``448.0f`` / ``57344.0f`` literals). -* :func:`quant_ir_type` — map a ``QDType`` to the IR :class:`Type`. +* :func:`quant_ir_type` — map a recognized dtype string to the IR :class:`Type`. * :func:`quantize_scalar_f32` — the one-element quant fast path (``cvt_f32_to_(clamp(x * inv_scale, -max, max))``). * :func:`dequantize_scalar_to_f32` — the dual; auto-dispatches on the @@ -37,7 +38,7 @@ from __future__ import annotations -from typing import List, Literal +from typing import List, Literal, TypeAlias from ..core.ir import ( BF8E5M2, @@ -54,6 +55,7 @@ __all__ = [ + "LogicalQDType", "QDType", "QUANT_MAX_ABS", "dequantize_scalar_to_f32", @@ -66,7 +68,8 @@ ] -QDType = Literal["i8", "fp8e4m3", "bf8e5m2", "fp4e2m1", "fp6e2m3", "fp6e3m2"] +QDType: TypeAlias = Literal["i8", "fp8e4m3", "bf8e5m2"] +LogicalQDType: TypeAlias = Literal[QDType, "fp4e2m1", "fp6e2m3", "fp6e3m2"] # Per-dtype clamp magnitude (the largest representable absolute value). @@ -154,7 +157,7 @@ def quant_max_abs(qdtype: str) -> float: return QUANT_MAX_ABS[canon] -def ir_to_qdtype(t: Type) -> QDType: +def ir_to_qdtype(t: Type) -> LogicalQDType: """Inverse of :func:`quant_ir_type`. Rejects non-quant types.""" if t.name in _IR_TO_QDTYPE: return _IR_TO_QDTYPE[t.name] # type: ignore[return-value] diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/README.md b/dnn-providers/hip-kernel-provider/rocke/platform/tests/README.md index 82ffdd97c593..e4e5da4de8bd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/README.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/README.md @@ -116,6 +116,6 @@ receive that path, so storage IR/HIP parity and serialization tests run automati fixture-discovery failure stops the runner instead of silently skipping coverage. An explicit `ROCKE_STORAGE_TEST` overrides discovery and must name an existing -executable. With no configured build or explicit override, the runner reports -native storage parity as skipped; direct pytest invocations can use the same -override. +executable; it does not skip the configured build. With no configured build or +explicit override, the runner reports native storage parity as skipped; direct +pytest invocations can use the same override. 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 index 2fd2eb0ac7e9..f5018fbeef28 100644 --- 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 @@ -3,6 +3,7 @@ """Core dtype names are independent of target support and retain public aliases.""" import ast +from typing import get_args, get_type_hints import pytest @@ -149,3 +150,18 @@ def test_type_alias_does_not_enable_scalar_quantization(dtype): assert quant_ir_type(dtype) is not None with pytest.raises(ValueError, match="unsupported quant dtype"): quant_max_abs(dtype) + + +def test_quant_dtype_annotations_match_scalar_and_logical_contracts(): + from rocke import helpers + from rocke.helpers import quant + + scalar = {"i8", "fp8e4m3", "bf8e5m2"} + assert set(get_args(quant.QDType)) == scalar + assert all(quant_max_abs(dtype) > 0 for dtype in get_args(quant.QDType)) + logical = scalar | {"fp4e2m1", "fp6e2m3", "fp6e3m2"} + assert set(get_args(quant.LogicalQDType)) == logical + assert get_type_hints(ir_to_qdtype)["return"] == quant.LogicalQDType + assert all(ir_to_qdtype(quant_ir_type(dtype)) == dtype for dtype in logical) + assert helpers.QDType is quant.QDType + assert helpers.LogicalQDType is quant.LogicalQDType diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/run_all.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/run_all.py index 51e12f36ea9a..e39834b58a34 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/run_all.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/run_all.py @@ -119,7 +119,7 @@ def native_pytest_env(build_root: Path, config: str) -> dict[str, str]: """Build the configured suite and supply native parity to both pytest passes.""" env = dict(os.environ) executable = env.get("ROCKE_STORAGE_TEST") - if not executable and (build_root / "CMakeCache.txt").is_file(): + if (build_root / "CMakeCache.txt").is_file(): subprocess.run( [ "cmake", @@ -130,13 +130,18 @@ def native_pytest_env(build_root: Path, config: str) -> dict[str, str]: ], check=True, ) - tests = ctest_tests(build_root, config, "^rocke_storage$") - commands = [ - test.get("command", []) for test in tests if test["name"] == "rocke_storage" - ] - if len(commands) != 1 or not commands[0]: - raise ValueError("CTest did not resolve the built rocke_storage executable") - executable = commands[0][0] + if not executable: + tests = ctest_tests(build_root, config, "^rocke_storage$") + commands = [ + test.get("command", []) + for test in tests + if test["name"] == "rocke_storage" + ] + if len(commands) != 1 or not commands[0]: + raise ValueError( + "CTest did not resolve the built rocke_storage executable" + ) + executable = commands[0][0] if executable: path = Path(executable).resolve() if not path.is_file(): diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/tools/test_run_all.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/tools/test_run_all.py index 484503960776..eed98eb130ac 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/tools/test_run_all.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/tools/test_run_all.py @@ -25,14 +25,19 @@ def runner(monkeypatch): @pytest.mark.parametrize( "filename", ["rocke_storage", "provider_rocke_storage_test.exe"] ) +@pytest.mark.parametrize("override", [False, True]) def test_runner_passes_registered_fixture_to_both_backends( - runner, monkeypatch, tmp_path, filename + runner, monkeypatch, tmp_path, filename, override ): (tmp_path / "CMakeCache.txt").touch() (tmp_path / "CTestTestfile.cmake").touch() executable = tmp_path / "Debug" / filename executable.parent.mkdir() executable.touch() + selected = tmp_path / "override" if override else executable + if override: + selected.touch() + monkeypatch.setenv("ROCKE_STORAGE_TEST", str(selected)) calls = [] def run(command, **kwargs): @@ -59,7 +64,7 @@ def run(command, **kwargs): (cmd, kw["env"]) for cmd, kw in calls if cmd[0] == runner.sys.executable ] assert len(children) == 3 # default pytest, extension import probe, both pytest - assert all(env["ROCKE_STORAGE_TEST"] == str(executable) for _, env in children) + assert all(env["ROCKE_STORAGE_TEST"] == str(selected) for _, env in children) assert "ROCKE_BACKEND" not in children[0][1] assert children[-1][1]["ROCKE_BACKEND"] == "both" assert calls[0][0] == [ @@ -70,6 +75,7 @@ def run(command, **kwargs): "Debug", ] assert any(cmd[:3] == ["ctest", "-C", "Debug"] for cmd, _ in calls) + assert any("-R" in cmd for cmd, _ in calls) is not override @pytest.mark.parametrize("built", ["none", "partial", "all"]) @@ -164,16 +170,22 @@ def test_invalid_override_is_an_error(runner, monkeypatch, tmp_path): runner.native_pytest_env(tmp_path, "Release") -@pytest.mark.parametrize("failure", ["build", "registration", "executable"]) +@pytest.mark.parametrize( + "failure", ["build", "build_override", "registration", "executable"] +) def test_native_setup_failure_prevents_silently_skipped_pytest( runner, monkeypatch, tmp_path, failure ): (tmp_path / "CMakeCache.txt").touch() + if failure == "build_override": + executable = tmp_path / "override" + executable.touch() + monkeypatch.setenv("ROCKE_STORAGE_TEST", str(executable)) calls = [] def run(command, **kwargs): calls.append(command) - if command[0] == "cmake" and failure == "build": + if command[0] == "cmake" and failure.startswith("build"): raise subprocess.CalledProcessError(1, command) command_path = [str(tmp_path / "missing")] if failure == "executable" else [] listing = {"tests": [{"name": "rocke_storage", "command": command_path}]} @@ -189,7 +201,10 @@ def run(command, **kwargs): assert all(command[0] != runner.sys.executable for command in calls) -def test_fresh_build_prepares_entire_ctest_suite(runner, tmp_path): +@pytest.mark.parametrize("override", [False, True]) +def test_fresh_build_prepares_entire_ctest_suite( + runner, monkeypatch, tmp_path, override +): if not shutil.which("cmake") or not shutil.which("ctest"): pytest.skip("CMake and CTest required for fresh-build regression") source = tmp_path / "source" @@ -211,8 +226,15 @@ def test_fresh_build_prepares_entire_ctest_suite(runner, tmp_path): subprocess.run( ["cmake", "--build", str(build), "--target", "rocke_core"], check=True ) + if override: + # An external fixture must not suppress the configured suite's build. + executable = tmp_path / "override" + executable.touch() + monkeypatch.setenv("ROCKE_STORAGE_TEST", str(executable)) env = runner.native_pytest_env(build, "Release") assert Path(env["ROCKE_STORAGE_TEST"]).is_file() + if override: + assert env["ROCKE_STORAGE_TEST"] == str(executable) assert runner.ctest_ready(build, "Release") subprocess.run( ["ctest", "--test-dir", str(build), "-C", "Release", "--output-on-failure"], From 6e536d15c1a3d27abb10f0313063a41922bc43a9 Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Fri, 25 Sep 2026 12:35:21 -0400 Subject: [PATCH 17/24] fix(rocke): bound fragment load scratch capacity Size temporary load-pointer storage by addressable units so padded slots cannot overrun the array. Preserve IR order and add mirrored transport coverage for a layout requiring more loads than logical elements. Simplify scale-width queries while retaining the shared packing adapter, and clarify the contiguous-K matrix layout contract. --- .../rocke/platform/cpp/helpers/mma_io.cpp | 5 ++-- .../platform/cpp/include/rocke/storage.h | 2 ++ .../cpp/include/rocke/wmma_scale_internal.h | 21 ++++++------- .../architecture/dtype_storage_packing.md | 4 ++- .../python/rocke/core/arch/wmma_scale.py | 13 ++++---- .../platform/python/rocke/core/storage.py | 2 ++ .../rocke/platform/tests/core/test_mma_io.py | 9 ++++++ .../platform/tests/core/test_storage.cpp | 30 +++++++++++++++++++ .../rocke/platform/tests/core/test_storage.py | 12 +++++++- 9 files changed, 77 insertions(+), 21 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp index b19383b4668f..9d8230329102 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp @@ -97,9 +97,10 @@ rocke_value_t* rocke_h_load_matrix_fragment(rocke_ir_builder_t* b, const int load_step = word_loads ? 4 : 1; auto* lane_chunk = rocke_b_mul(b, lane_group, rocke_b_const_i32(b, chunk_units)); auto* step_base = rocke_b_add(b, row_base, rocke_b_const_i32(b, origin_bytes / unit_bytes)); - /* Collect loads before concatenation to preserve Python SSA numbering. */ + /* Collect loads before concatenation to preserve Python SSA numbering. + * Each load consumes at least one storage unit, including padded slots. */ rocke_value_t** chunks = (rocke_value_t**)rocke_arena_alloc( - &b->arena, size_t(packing.count) * sizeof(rocke_value_t*)); + &b->arena, size_t(payload_bits / (8 * unit_bytes)) * sizeof(rocke_value_t*)); if(!chunks) ckc::raise_status(ROCKE_ERR_OOM, "matrix fragment allocation failed"); int num_chunks = 0; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/storage.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/storage.h index 9634dbcc2e1f..37a3ee9be0c9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/storage.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/storage.h @@ -31,6 +31,8 @@ typedef struct rocke_fragment_packing uint64_t carrier_count; } rocke_fragment_packing_t; +/* Interleaved contiguous K chunks, not arbitrary matrix distributions, + * transposed axes, or swizzled tensor addresses. */ typedef struct rocke_matrix_fragment_layout { rocke_fragment_packing_t fragment; 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 index dcd0e389b07a..34ae13074d7b 100644 --- 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 @@ -24,16 +24,22 @@ typedef struct rocke_scale_packing int block_k; } rocke_scale_packing_t; -static inline rocke_fragment_packing_t rocke_scale_fragment(const rocke_scale_packing_t* p) +static inline int rocke_scale_word_bits(const rocke_scale_packing_t* p) { - rocke_bit_packing_t bits; - rocke_fragment_packing_t result; if(p->block_k <= 0) ckc::raise_status(ROCKE_ERR_VALUE, "scale block_k must be positive"); if(p->count != 1 && p->count != 2 && p->count != 4 && p->count != 8) ckc::raise_status(ROCKE_ERR_VALUE, "carrier_bits must be 8, 16, 32, or 64"); - if(!rocke_bit_packing_init(&bits, rocke_dtype_info("e8m0")->encoded_bits, 0) - || !rocke_fragment_packing_init(&result, &bits, p->count, p->count * 8, 1)) + return p->count * 8; +} + +static inline rocke_fragment_packing_t rocke_scale_fragment(const rocke_scale_packing_t* p) +{ + const int word_bits = rocke_scale_word_bits(p); + rocke_bit_packing_t bits; + rocke_fragment_packing_t result; + if(!rocke_bit_packing_init(&bits, 8, 0) + || !rocke_fragment_packing_init(&result, &bits, p->count, word_bits, 1)) ckc::raise_status(ROCKE_ERR_VALUE, "invalid scale packing"); return result; } @@ -145,9 +151,4 @@ static inline const rocke_mma_op_t* rocke_gfx1250_scaled_wmma_from_op(const rock 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 rocke_scale_fragment(packing).carrier_bits; -} - #endif /* ROCKE_WMMA_SCALE_INTERNAL_H */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md index 0180817073e3..ab5a8349a6f6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md @@ -116,7 +116,9 @@ buffers; concurrent packed stores are not provided. `MatrixFragmentLayout` describes contiguous K chunks interleaved between lane groups. Coordinate mapping belongs to the atom-specific layout, not to the -generic bit packer. `ScaledWmmaOp.matrix_layout(operand)` selects independent +generic bit packer. This is a contiguous-K mapping, not a general matrix +distribution or a description of transposed axes and swizzled tensor addresses. +`ScaledWmmaOp.matrix_layout(operand)` selects independent A/B layouts. Existing `a_frag_len`/`b_frag_len` retain their ABI-vector meaning. The current gfx1250 scaled matrix layouts have 64 elements per lane and sixteen 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 index a506ad183e71..3b327b1b4133 100644 --- 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 @@ -42,25 +42,24 @@ class ScalePacking: def __post_init__(self) -> None: if self.block_k <= 0: raise ValueError("scale block_k must be positive") - self.fragment + if not isinstance(self.count, int) or self.count not in (1, 2, 4, 8): + raise ValueError("carrier_bits must be 8, 16, 32, or 64") @property def packing(self) -> BitPacking: - return BitPacking(dtype_info("e8m0").encoded_bits) + return BitPacking(self.element_bits) @property def fragment(self) -> FragmentPacking: - return FragmentPacking( - self.packing, self.count, self.count * self.packing.element_bits, 1 - ) + return FragmentPacking(self.packing, self.count, self.word_bits, 1) @property def element_bits(self) -> int: - return self.packing.element_bits + return 8 @property def word_bits(self) -> int: - return self.fragment.carrier_bits + return self.count * self.element_bits @property def llvm_type(self) -> str: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/storage.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/storage.py index c9c1ce3282e9..e0770a939c1c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/storage.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/storage.py @@ -127,6 +127,8 @@ class MatrixFragmentLayout: """Contiguous K chunks interleaved between lane groups. The atom selects this mapping. Storage and carrier widths do not determine it. + This describes interleaved contiguous K chunks, not arbitrary matrix + distributions, transposed axes, or swizzled tensor addresses. """ fragment: FragmentPacking diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py index 59e1bbf1b433..96c3e7e930fe 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py @@ -32,6 +32,7 @@ "load96_f16", "load96_i32", "fp4", + "fp4_slot40", "fp6", "fp6_padded", "bf6", @@ -72,6 +73,9 @@ def build_transport(dtype): o, b.const_i32(k * n + j), b.vec_extract(vector, j), align=12 // n ) return b.kernel + padded_slots = dtype == "fp4_slot40" + if padded_slots: + dtype = "fp4" padded = dtype.endswith("_padded") dtype = dtype.removesuffix("_padded") patterns = dtype.startswith("pack_") @@ -100,6 +104,11 @@ def build_transport(dtype): ) if dtype == "fp6" and padded: layout = MatrixFragmentLayout(layout.fragment, 16, 2, 16) + if padded_slots: + # Four logical elements need eight loads (4 + 1 bytes per slot). + layout = MatrixFragmentLayout( + FragmentPacking(BitPacking(4, 40), 4, 32, 5), 1, 2, 16 + ) base = b.const_i32((129 if typed else 97) if padded else 0) lane = b.mod(b.thread_id_x(), b.const_i32(32)) group = b.div(lane, b.const_i32(16)) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp index d170d81d4dd2..185c6f12d3db 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.cpp @@ -34,6 +34,9 @@ static int emit(const char* dtype, bool hip, bool include_prologue = true) { rocke_ir_builder_t b; CHECK(rocke_ir_builder_init(&b, "transport") == ROCKE_OK); + const bool padded_slots = strcmp(dtype, "fp4_slot40") == 0; + if(padded_slots) + dtype = "fp4"; const bool padded = strstr(dtype, "_padded") != NULL; if(padded) dtype = strcmp(dtype, "f16_padded") == 0 ? "f16" @@ -108,6 +111,15 @@ static int emit(const char* dtype, bool hip, bool include_prologue = true) layout = rocke_scaled_matrix_layout(dtype, 16); if(strcmp(dtype, "fp6") == 0 && padded) CHECK(rocke_matrix_fragment_layout_init(&layout, &layout.fragment, 16, 2, 16)); + if(padded_slots) + { + // Four logical elements need eight loads (4 + 1 bytes per slot). + rocke_bit_packing_t packing; + rocke_fragment_packing_t fragment; + CHECK(rocke_bit_packing_init(&packing, 4, 40)); + CHECK(rocke_fragment_packing_init(&fragment, &packing, 4, 32, 5)); + CHECK(rocke_matrix_fragment_layout_init(&layout, &fragment, 1, 2, 16)); + } auto* base = rocke_b_const_i32(&b, padded ? (typed ? 129 : 97) : 0); auto* thread = rocke_b_thread_id_x(&b); auto* lane = rocke_b_mod(&b, thread, rocke_b_const_i32(&b, 32)); @@ -281,6 +293,24 @@ int main(int argc, char** argv) } CHECK(test_fragment_inputs() == 0); CHECK(test_fragment_offset_bounds() == 0); + for(int count : {1, 2, 4, 8}) + { + const rocke_scale_packing_t scales = {count, 32}; + CHECK(rocke_scale_word_bits(&scales) == count * 8); + CHECK(rocke_scale_fragment(&scales).carrier_bits == count * 8); + } + for(const auto scales : {rocke_scale_packing_t{0, 32}, {-1, 32}, {3, 32}, {16, 32}, {4, 0}}) + { + try + { + rocke_scale_word_bits(&scales); + CHECK(false); + } + catch(const ckc::Error& error) + { + CHECK(error.code() == ROCKE_ERR_VALUE); + } + } CHECK(rocke_dtype_info("e4m3") == rocke_dtype_info("fp8e4m3")); CHECK(rocke_dtype_to_ir_type("e4m3") == rocke_fp8e4m3()); CHECK(rocke_storage_ir_type("e4m3") == rocke_fp8e4m3()); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py index cb572e5af990..476276e2fc1e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py @@ -54,7 +54,8 @@ def test_matrix_payload_and_padding(dtype, live): @pytest.mark.parametrize( - "count,block_k,word", [(4, 32, 0x04030201), (8, 16, 0x0807060504030201)] + "count,block_k,word", + [(1, 32, 0x01), (2, 32, 0x0201), (4, 32, 0x04030201), (8, 16, 0x0807060504030201)], ) def test_scales_use_common_bit_packing(count, block_k, word): scales = ScalePacking(count, block_k) @@ -62,6 +63,15 @@ def test_scales_use_common_bit_packing(count, block_k, word): assert scales.fragment.pack(list(range(1, count + 1))) == (word,) assert scales.packing == BitPacking(8) assert scales.word_bits == count * 8 + assert scales.llvm_type == f"i{count * 8}" + + +@pytest.mark.parametrize( + "count,block_k", [(0, 32), (-1, 32), (3, 32), (16, 32), (4.0, 32), (4, 0)] +) +def test_invalid_scale_packing(count, block_k): + with pytest.raises(ValueError): + ScalePacking(count, block_k) def test_invalid_packing_and_overflow(): From a8241a9e41e29aaaa6272071dda30e28acd34cce Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Fri, 25 Sep 2026 15:11:19 -0400 Subject: [PATCH 18/24] feat(rocke): add logical TF32 and native gfx942 atoms --- .../platform/cpp/bindings/rocke_engine.cpp | 20 + .../rocke/platform/cpp/core/arch/data.cpp | 94 +++++ .../rocke/platform/cpp/core/dtypes.cpp | 63 +--- .../rocke/platform/cpp/core/ir/core_types.cpp | 3 + .../rocke/platform/cpp/core/ir/ir_arith.cpp | 42 +++ .../rocke/platform/cpp/core/ir/ir_tile.cpp | 13 +- .../platform/cpp/core/lower_hip/core.cpp | 17 +- .../cpp/core/lower_hip/lower_hip_mma.cpp | 31 ++ .../platform/cpp/core/lower_llvm/core.cpp | 12 +- .../platform/cpp/core/lower_llvm/data.cpp | 6 + .../platform/cpp/core/lower_llvm/mem.cpp | 4 +- .../platform/cpp/core/lower_llvm/mma.cpp | 16 + .../rocke/platform/cpp/core/verify.cpp | 4 + .../rocke/platform/cpp/helpers/mma_io.cpp | 2 + .../include/rocke/instance_tf32_mma_probe.h | 15 + .../rocke/platform/cpp/include/rocke/ir.h | 4 + .../cpp/include/rocke/tf32_internal.h | 65 ++++ .../cpp/instances/gfx942/tf32_mma_probe.cpp | 138 +++++++ .../architecture/dtype_storage_packing.md | 7 + .../architecture/ir_serialization_format.md | 6 + .../dsl_docs/reference/mfma_atom_catalog.md | 7 + .../dsl_docs/reference/mma_metadata.md | 12 + .../dsl_docs/reference/op_vocabulary.md | 7 + .../rocke/core/arch/data/arch_specs.json | 2 + .../platform/python/rocke/core/arch/target.py | 38 ++ .../platform/python/rocke/core/dtypes.py | 3 + .../rocke/platform/python/rocke/core/ir.py | 67 +++- .../platform/python/rocke/core/lower_hip.py | 22 +- .../platform/python/rocke/core/lower_llvm.py | 62 ++- .../rocke/platform/python/rocke/core/tf32.py | 43 +++ .../platform/python/rocke/core/verify.py | 5 + .../rocke/examples/gfx942/tf32_numerics.py | 352 ++++++++++++++++++ .../examples/gfx942/tf32_numerics/README.md | 72 ++++ .../platform/python/rocke/helpers/mma_io.py | 2 + .../rocke/instances/gfx942/tf32_mma_probe.py | 120 ++++++ .../rocke/platform/tests/CMakeLists.txt | 9 + .../rocke/platform/tests/core/test_tf32.cpp | 73 ++++ .../rocke/platform/tests/core/test_tf32.py | 247 ++++++++++++ .../differential/golden/llvm_gfx_all.json | 12 + .../tests/instances/parity/tf32_mma_emit.c | 74 ++++ .../tests/instances/parity/tf32_mma_emit.py | 20 + .../tests/numeric/test_gfx942_tf32_numeric.py | 17 + 42 files changed, 1759 insertions(+), 69 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_tf32_mma_probe.h create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/tf32_internal.h create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/tf32_mma_probe.cpp create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/tf32.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics/README.md create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/gfx942/tf32_mma_probe.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.cpp create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/tf32_mma_emit.c create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/tf32_mma_emit.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/numeric/test_gfx942_tf32_numeric.py diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp index 542f7d3698bd..e462a030a029 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine.cpp @@ -52,6 +52,7 @@ extern "C" { #include "rocke/instance_mfma_gemm.h" #include "rocke/instance_mx_gemm.h" #include "rocke/instance_streamk_gemm.h" +#include "rocke/instance_tf32_mma_probe.h" #include "rocke/ir.h" #include "rocke/ir_serialize.h" #include "rocke/lower_llvm.h" @@ -3512,6 +3513,25 @@ PYBIND11_MODULE(rocke_engine, m) []() { return std::string(rocke_engine_version()); }, "Human-readable engine version of this module."); + m.def( + "tf32_mma_probe_serialize_ir", + [](int shape, const std::string& preparation) { + rocke_ir_builder_t b; + rocke_kernel_def_t* kernel = rocke_build_tf32_mma_probe(&b, shape, preparation.c_str()); + char* text = nullptr; + rocke_status_t status + = kernel ? rocke_ir_serialize(kernel, &text) : rocke_ir_builder_status(&b); + std::string error = rocke_ir_builder_error(&b); + std::string result = text ? text : ""; + free(text); + rocke_ir_builder_free(&b); + if(status != ROCKE_OK) + throw std::runtime_error(error); + return result; + }, + py::arg("m") = 16, + py::arg("preparation") = "raw"); + /* ---- family-agnostic lower-from-serialized-IR ---- * The keystone of the ROCKE_BACKEND=cpp default for Python-authored * kernels: a Python front end serializes any KernelDef to ck.dsl.ir/v1 diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/data.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/data.cpp index 3774425da6ac..5faf47273f37 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/data.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/arch/data.cpp @@ -708,6 +708,67 @@ static const rocke_layout_map_t lm_mfma_32x32x16_a = {ROCKE_MMA_ROLE_A, 8, 64, _ static const rocke_layout_map_t lm_mfma_32x32x16_b = {ROCKE_MMA_ROLE_B, 8, 64, _mfma_b_32x32x16}; static const rocke_layout_map_t lm_mfma_32x32x16_c = {ROCKE_MMA_ROLE_ACC, 16, 64, _mfma_acc_32x32}; +static void _mfma_a_16x16x8_xf32( + rocke_ir_builder_t* b, rocke_value_t* lane, int slot, rocke_value_t** c0, rocke_value_t** c1) +{ + rocke_value_t* c = rocke_b_const_i32(b, 16); + rocke_value_t* axis = rocke_b_mod(b, lane, c); + rocke_value_t* group = rocke_b_div(b, lane, c); + rocke_value_t* two = rocke_b_const_i32(b, 2); + rocke_value_t* base = rocke_b_mul(b, group, two); + rocke_value_t* index = rocke_b_const_i32(b, slot); + rocke_value_t* k = rocke_b_add(b, base, index); + *c0 = axis; + *c1 = k; +} + +static void _mfma_b_16x16x8_xf32( + rocke_ir_builder_t* b, rocke_value_t* lane, int slot, rocke_value_t** c0, rocke_value_t** c1) +{ + rocke_value_t* c = rocke_b_const_i32(b, 16); + rocke_value_t* axis = rocke_b_mod(b, lane, c); + rocke_value_t* group = rocke_b_div(b, lane, c); + rocke_value_t* two = rocke_b_const_i32(b, 2); + rocke_value_t* base = rocke_b_mul(b, group, two); + rocke_value_t* index = rocke_b_const_i32(b, slot); + rocke_value_t* k = rocke_b_add(b, base, index); + *c0 = k; + *c1 = axis; +} + +static void _mfma_a_32x32x4_xf32( + rocke_ir_builder_t* b, rocke_value_t* lane, int slot, rocke_value_t** c0, rocke_value_t** c1) +{ + rocke_value_t* c = rocke_b_const_i32(b, 32); + rocke_value_t* axis = rocke_b_mod(b, lane, c); + rocke_value_t* group = rocke_b_div(b, lane, c); + rocke_value_t* two = rocke_b_const_i32(b, 2); + rocke_value_t* base = rocke_b_mul(b, group, two); + rocke_value_t* index = rocke_b_const_i32(b, slot); + rocke_value_t* k = rocke_b_add(b, base, index); + *c0 = axis; + *c1 = k; +} + +static void _mfma_b_32x32x4_xf32( + rocke_ir_builder_t* b, rocke_value_t* lane, int slot, rocke_value_t** c0, rocke_value_t** c1) +{ + rocke_value_t* c = rocke_b_const_i32(b, 32); + rocke_value_t* axis = rocke_b_mod(b, lane, c); + rocke_value_t* group = rocke_b_div(b, lane, c); + rocke_value_t* two = rocke_b_const_i32(b, 2); + rocke_value_t* base = rocke_b_mul(b, group, two); + rocke_value_t* index = rocke_b_const_i32(b, slot); + rocke_value_t* k = rocke_b_add(b, base, index); + *c0 = k; + *c1 = axis; +} + +static const rocke_layout_map_t lm_xf32_16_a = {ROCKE_MMA_ROLE_A, 2, 64, _mfma_a_16x16x8_xf32}; +static const rocke_layout_map_t lm_xf32_16_b = {ROCKE_MMA_ROLE_B, 2, 64, _mfma_b_16x16x8_xf32}; +static const rocke_layout_map_t lm_xf32_32_a = {ROCKE_MMA_ROLE_A, 2, 64, _mfma_a_32x32x4_xf32}; +static const rocke_layout_map_t lm_xf32_32_b = {ROCKE_MMA_ROLE_B, 2, 64, _mfma_b_32x32x4_xf32}; + /* --- mfma_f32_16x16x4_f32: a/b/c present (frag 1/1/4, wave64) (#8348) --- */ static const rocke_layout_map_t lm_mfma_16x16x4_f32_a = {ROCKE_MMA_ROLE_A, 1, 64, _mfma_a_16x16x4_f32}; @@ -890,6 +951,8 @@ static_assert(offsetof(rocke_ati_mma_frag_row_t, c_frag_len) == sizeof(const cha "c_frag_len must immediately follow op_id (no leading padding)"); static const rocke_ati_mma_frag_row_t rocke_ati_mma_frag[] = { + {"mfma_f32_32x32x4_xf32", 16}, + {"mfma_f32_16x16x8_xf32", 4}, /* --- MFMA fp32 (wave64) --- */ {"mfma_f32_16x16x4_f32", 4}, {"mfma_f32_32x32x2_f32", 16}, @@ -1040,6 +1103,37 @@ static const rocke_mma_op_t k_mma_gfx90a[] = { /* ----------------------------- gfx942 (CDNA) ----------------------------- */ static const rocke_mma_op_t k_mma_gfx942[] = { + {"mma", + "tf32", + "tf32", + "fp32", + 16, + 16, + 8, + "mfma_f32_16x16x8_xf32", + 2, + 2, + 4, + 64, + &lm_xf32_16_a, + &lm_xf32_16_b, + &lm_mfma_16x16x4_f32_c}, + {"mma", + "tf32", + "tf32", + "fp32", + 32, + 32, + 4, + "mfma_f32_32x32x4_xf32", + 2, + 2, + 16, + 64, + &lm_xf32_32_a, + &lm_xf32_32_b, + &lm_mfma_32x32x2_f32_c}, + {"mma", "fp32", "fp32", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/dtypes.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/dtypes.cpp index dcc32275b741..362d722b9e26 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/dtypes.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/dtypes.cpp @@ -26,33 +26,16 @@ typedef struct rocke_dtype_alias /* Byte-for-byte the _DTYPE_ALIASES map (insertion order is irrelevant: lookup is * by exact lowercased key). */ static const rocke_dtype_alias_t k_dtype_aliases[] = { - {"f16", "fp16"}, - {"half", "fp16"}, - {"fp16", "fp16"}, - {"bf16", "bf16"}, - {"bfloat16", "bf16"}, - {"f32", "fp32"}, - {"float", "fp32"}, - {"fp32", "fp32"}, - {"iu8", "iu8"}, - {"iu4", "iu4"}, - {"i8", "i8"}, - {"int8", "i8"}, - {"i4", "i4"}, - {"int4", "i4"}, - {"i32", ROCKE_DTYPE_I32}, - {"int32", ROCKE_DTYPE_I32}, - {"fp8", "fp8e4m3"}, - {"fp8e4m3", "fp8e4m3"}, - {"e4m3", "fp8e4m3"}, - {"bf8", "bf8e5m2"}, - {"bf8e5m2", "bf8e5m2"}, - {"fp6", "fp6e2m3"}, - {"fp6e2m3", "fp6e2m3"}, - {"bf6", "fp6e3m2"}, - {"fp6e3m2", "fp6e3m2"}, - {"fp4", "fp4e2m1"}, - {"fp4e2m1", "fp4e2m1"}, + {"f16", "fp16"}, {"half", "fp16"}, {"fp16", "fp16"}, + {"bf16", "bf16"}, {"bfloat16", "bf16"}, {"f32", "fp32"}, + {"float", "fp32"}, {"fp32", "fp32"}, {"tf32", "tf32"}, + {"xf32", "tf32"}, {"iu8", "iu8"}, {"iu4", "iu4"}, + {"i8", "i8"}, {"int8", "i8"}, {"i4", "i4"}, + {"int4", "i4"}, {"i32", ROCKE_DTYPE_I32}, {"int32", ROCKE_DTYPE_I32}, + {"fp8", "fp8e4m3"}, {"fp8e4m3", "fp8e4m3"}, {"e4m3", "fp8e4m3"}, + {"bf8", "bf8e5m2"}, {"bf8e5m2", "bf8e5m2"}, {"fp6", "fp6e2m3"}, + {"fp6e2m3", "fp6e2m3"}, {"bf6", "fp6e3m2"}, {"fp6e3m2", "fp6e3m2"}, + {"fp4", "fp4e2m1"}, {"fp4e2m1", "fp4e2m1"}, }; #define K_NUM_DTYPE_ALIASES ((int)(sizeof(k_dtype_aliases) / sizeof(k_dtype_aliases[0]))) @@ -140,23 +123,15 @@ const char* rocke_normalize_dtype(const char* name, char* lowered, size_t cap) } static const rocke_dtype_info_t k_dtype_info[] = { - {"fp32", ROCKE_DTYPE_FLOAT, 32}, - {"fp16", ROCKE_DTYPE_FLOAT, 16}, - {"bf16", ROCKE_DTYPE_FLOAT, 16}, - {"fp8e4m3", ROCKE_DTYPE_FLOAT, 8}, - {"bf8e5m2", ROCKE_DTYPE_FLOAT, 8}, - {"fp6e2m3", ROCKE_DTYPE_FLOAT, 6}, - {"fp6e3m2", ROCKE_DTYPE_FLOAT, 6}, - {"fp4e2m1", ROCKE_DTYPE_FLOAT, 4}, - {"e8m0", ROCKE_DTYPE_FLOAT, 8}, - {"e5m3", ROCKE_DTYPE_FLOAT, 8}, - {"i4", ROCKE_DTYPE_INTEGER, 4}, - {"iu4", ROCKE_DTYPE_INTEGER, 4}, - {"i8", ROCKE_DTYPE_INTEGER, 8}, - {"iu8", ROCKE_DTYPE_INTEGER, 8}, - {"i16", ROCKE_DTYPE_INTEGER, 16}, - {"i32", ROCKE_DTYPE_INTEGER, 32}, - {"i64", ROCKE_DTYPE_INTEGER, 64}, + {"fp32", ROCKE_DTYPE_FLOAT, 32}, {"tf32", ROCKE_DTYPE_FLOAT, 32}, + {"fp16", ROCKE_DTYPE_FLOAT, 16}, {"bf16", ROCKE_DTYPE_FLOAT, 16}, + {"fp8e4m3", ROCKE_DTYPE_FLOAT, 8}, {"bf8e5m2", ROCKE_DTYPE_FLOAT, 8}, + {"fp6e2m3", ROCKE_DTYPE_FLOAT, 6}, {"fp6e3m2", ROCKE_DTYPE_FLOAT, 6}, + {"fp4e2m1", ROCKE_DTYPE_FLOAT, 4}, {"e8m0", ROCKE_DTYPE_FLOAT, 8}, + {"e5m3", ROCKE_DTYPE_FLOAT, 8}, {"i4", ROCKE_DTYPE_INTEGER, 4}, + {"iu4", ROCKE_DTYPE_INTEGER, 4}, {"i8", ROCKE_DTYPE_INTEGER, 8}, + {"iu8", ROCKE_DTYPE_INTEGER, 8}, {"i16", ROCKE_DTYPE_INTEGER, 16}, + {"i32", ROCKE_DTYPE_INTEGER, 32}, {"i64", ROCKE_DTYPE_INTEGER, 64}, {"i1", ROCKE_DTYPE_PREDICATE, 1}, }; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp index 6062cadb30c0..cd34e3b984c0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp @@ -53,6 +53,7 @@ ROCKE_SCALAR_SINGLETON(rocke_i64, ROCKE_SCALAR_I64, "i64") ROCKE_SCALAR_SINGLETON(rocke_bf16, ROCKE_SCALAR_BF16, "bf16") ROCKE_SCALAR_SINGLETON(rocke_f16, ROCKE_SCALAR_F16, "f16") ROCKE_SCALAR_SINGLETON(rocke_f32, ROCKE_SCALAR_F32, "f32") +ROCKE_SCALAR_SINGLETON(rocke_tf32, ROCKE_SCALAR_TF32, "tf32") ROCKE_SCALAR_SINGLETON(rocke_fp8e4m3, ROCKE_SCALAR_FP8E4M3, "fp8e4m3") ROCKE_SCALAR_SINGLETON(rocke_bf8e5m2, ROCKE_SCALAR_BF8E5M2, "bf8e5m2") @@ -82,6 +83,8 @@ const rocke_type_t* rocke_scalar_by_name(const char* name) return rocke_bf16(); if(strcmp(name, "f16") == 0) return rocke_f16(); + if(strcmp(name, "tf32") == 0) + return rocke_tf32(); if(strcmp(name, "f32") == 0) return rocke_f32(); if(strcmp(name, "fp8e4m3") == 0 || strcmp(name, "e4m3") == 0) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_arith.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_arith.cpp index 910fe60eb2fc..5338a4d9b659 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_arith.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_arith.cpp @@ -591,6 +591,48 @@ rocke_value_t* rocke_b_cast_to_f32(rocke_ir_builder_t* b, rocke_value_t* v) return rocke_i_op1(b, ROCKE_OP_ARITH_CAST_TO_F32, operands, 1, rocke_f32(), NULL, "f32"); } +rocke_value_t* rocke_b_cvt_f32_to_tf32(rocke_ir_builder_t* b, rocke_value_t* v) +{ + if(!rocke_i_live(b)) + return NULL; + if(!v || !rocke_type_eq(v->type, rocke_f32())) + return (rocke_value_t*)rocke_i_set_err( + b, ROCKE_ERR_VALUE, "cvt_f32_to_tf32 expects f32 input"); + rocke_value_t* bits = rocke_b_bitcast(b, v, rocke_i32()); + rocke_value_t* magnitude_mask = rocke_b_const_i32(b, 2147483647); + rocke_value_t* mag = rocke_b_land(b, bits, magnitude_mask); + rocke_value_t* sign_mask = rocke_b_const_i32(b, -2147483648); + rocke_value_t* sign = rocke_b_land(b, bits, sign_mask); + rocke_value_t* exponent_mask = rocke_b_const_i32(b, 2139095040); + rocke_value_t* exp = rocke_b_land(b, mag, exponent_mask); + rocke_value_t* exponent_all_ones = rocke_b_const_i32(b, 2139095040); + rocke_value_t* special = rocke_b_cmp_eq(b, exp, exponent_all_ones); + rocke_value_t* zero = rocke_b_const_i32(b, 0); + rocke_value_t* safe = rocke_b_select(b, special, zero, mag); + rocke_value_t* discarded_bits = rocke_b_const_i32(b, 13); + rocke_value_t* shift = rocke_b_lshr(b, safe, discarded_bits); + rocke_value_t* low_bit = rocke_b_const_i32(b, 1); + rocke_value_t* odd = rocke_b_land(b, shift, low_bit); + rocke_value_t* rounding_bias = rocke_b_const_i32(b, 4095); + rocke_value_t* bias = rocke_b_add(b, odd, rounding_bias); + rocke_value_t* rounded = rocke_b_add(b, safe, bias); + rocke_value_t* precision_mask = rocke_b_const_i32(b, -8192); + rounded = rocke_b_land(b, rounded, precision_mask); + rounded = rocke_b_lor(b, rounded, sign); + rocke_value_t* fraction_mask = rocke_b_const_i32(b, 8388607); + rocke_value_t* frac = rocke_b_land(b, mag, fraction_mask); + rocke_value_t* zero_fraction = rocke_b_const_i32(b, 0); + rocke_value_t* is_nan = rocke_b_cmp_ne(b, frac, zero_fraction); + rocke_value_t* quiet_bit = rocke_b_const_i32(b, 4194304); + rocke_value_t* quiet = rocke_b_lor(b, bits, quiet_bit); + rocke_value_t* nan_precision_mask = rocke_b_const_i32(b, -8192); + quiet = rocke_b_land(b, quiet, nan_precision_mask); + rocke_value_t* nonfinite = rocke_b_select(b, is_nan, quiet, bits); + rocke_value_t* result = rocke_b_select(b, special, nonfinite, rounded); + result = rocke_b_bitcast(b, result, rocke_tf32()); + return result; +} + rocke_value_t* rocke_b_cast_f32_to(rocke_ir_builder_t* b, rocke_value_t* v, const rocke_type_t* target) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_tile.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_tile.cpp index 0585df2a5aab..3311d8e603ba 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_tile.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_tile.cpp @@ -13,6 +13,7 @@ * rocke_i_attrs / type helpers) lives in bucket 0 (ir_core.c) and is declared * in rocke/ir_internal.h. */ +#include "rocke/tf32_internal.h" #include #include @@ -98,7 +99,8 @@ static int rocke_elem_bytes_name(const char* elem_name) { return 1; } - if(strcmp(elem_name, "f32") == 0 || strcmp(elem_name, "i32") == 0) + if(strcmp(elem_name, "f32") == 0 || strcmp(elem_name, "i32") == 0 + || strcmp(elem_name, "tf32") == 0) { return 4; } @@ -363,13 +365,13 @@ rocke_value_t* rocke_b_smem_load_vN(rocke_ir_builder_t* b, } dn = dtype->name; if(!(strcmp(dn, "f16") == 0 || strcmp(dn, "bf16") == 0 || strcmp(dn, "f32") == 0 - || strcmp(dn, "i32") == 0 || strcmp(dn, "fp8e4m3") == 0 || strcmp(dn, "bf8e5m2") == 0 - || strcmp(dn, "i8") == 0)) + || strcmp(dn, "tf32") == 0 || strcmp(dn, "i32") == 0 || strcmp(dn, "fp8e4m3") == 0 + || strcmp(dn, "bf8e5m2") == 0 || strcmp(dn, "i8") == 0)) { return (rocke_value_t*)rocke_i_set_err( b, ROCKE_ERR_VALUE, - "smem_load_vN supports f16 / bf16 / f32 / i32 / fp8e4m3 / " + "smem_load_vN supports f16 / bf16 / f32 / i32 / tf32 / fp8e4m3 / " "bf8e5m2 / i8, got %s", dn); } @@ -607,6 +609,9 @@ rocke_value_t* rocke_b_mma(rocke_ir_builder_t* b, { ops[3 + i] = extra[i]; } + const char* tf32_error = rocke_tf32_mma_error(op_id, ops, nops); + if(tf32_error) + return (rocke_value_t*)rocke_i_set_err(b, ROCKE_ERR_VALUE, "%s", tf32_error); attrs = rocke_i_attrs(b); rocke_attr_set_str(b, &attrs, "op_id", op_id); return rocke_i_op1(b, ROCKE_OP_TILE_MMA, ops, nops, vt, &attrs, hint); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp index de9b2d5cb110..47c5931704ec 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/core.cpp @@ -28,6 +28,7 @@ * arena-backed table keyed by the producing Value pointer, since the frozen IR * attrs must not be mutated. */ +#include "rocke/tf32_internal.h" #include #include #include @@ -295,7 +296,7 @@ const char* rocke_h_hip_scalar(const char* ir_scalar_name) { return "int16_t"; } - if(strcmp(ir_scalar_name, "i32") == 0) + if(strcmp(ir_scalar_name, "i32") == 0 || strcmp(ir_scalar_name, "tf32") == 0) { return "int"; } @@ -345,7 +346,7 @@ const char* rocke_h_vec_prefix(const char* ir_scalar_name, bool full_map) { return "f32x"; } - if(strcmp(ir_scalar_name, "i32") == 0) + if(strcmp(ir_scalar_name, "i32") == 0 || strcmp(ir_scalar_name, "tf32") == 0) { return "i32x"; } @@ -383,8 +384,8 @@ static bool rocke_h_scalar_in_vec_map(const char* name, bool full_map) { return false; } - return strcmp(name, "f32") == 0 || strcmp(name, "i32") == 0 || strcmp(name, "i16") == 0 - || strcmp(name, "i8") == 0 || strcmp(name, "fp8e4m3") == 0 + return strcmp(name, "tf32") == 0 || strcmp(name, "f32") == 0 || strcmp(name, "i32") == 0 + || strcmp(name, "i16") == 0 || strcmp(name, "i8") == 0 || strcmp(name, "fp8e4m3") == 0 || strcmp(name, "bf8e5m2") == 0; } @@ -450,8 +451,9 @@ const char* rocke_h_type_to_hip(rocke_h_lowerer_t* lw, const rocke_type_t* t) * through to the KeyError. Detect the listed set explicitly so an * unknown vector elem is an error rather than silently "f16x". */ if(strcmp(elem, "f16") != 0 && strcmp(elem, "bf16") != 0 && strcmp(elem, "f32") != 0 - && strcmp(elem, "i32") != 0 && strcmp(elem, "i16") != 0 && strcmp(elem, "i8") != 0 - && strcmp(elem, "fp8e4m3") != 0 && strcmp(elem, "bf8e5m2") != 0) + && strcmp(elem, "tf32") != 0 && strcmp(elem, "i32") != 0 && strcmp(elem, "i16") != 0 + && strcmp(elem, "i8") != 0 && strcmp(elem, "fp8e4m3") != 0 + && strcmp(elem, "bf8e5m2") != 0) { rocke_h_fail(lw, ROCKE_ERR_KEY, "type_to_hip: unmappable vector elem '%s'", elem); return ""; @@ -796,6 +798,9 @@ rocke_status_t rocke_h_lower_op(rocke_h_lowerer_t* lw, const rocke_op_t* op) { return rocke_h_fail(lw, ROCKE_ERR_VALUE, "lower_op: NULL op"); } + const char* tf32_error = rocke_tf32_op_error(op); + if(tf32_error) + return rocke_h_fail(lw, ROCKE_ERR_VALUE, "%s", tf32_error); fn = rocke_h_dispatch(op->opcode); if(!fn) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mma.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mma.cpp index 3abfd360a3c5..6abf236d5547 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mma.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mma.cpp @@ -30,6 +30,7 @@ #include "rocke/ir.h" #include "rocke/lower_hip.h" #include "rocke/lower_hip_internal.h" +#include "rocke/tf32_internal.h" #include "rocke/wmma_scale_internal.h" #include /* snprintf */ @@ -134,6 +135,36 @@ static rocke_status_t rocke_h_op_tile_mma(rocke_h_lowerer_t* lw, const rocke_op_ { return rocke_h_fail(lw, ROCKE_ERR_KEY, "tile.mma: missing 'op_id' attr"); } + int tf32_count = rocke_tf32_mma_count(op_id); + if(tf32_count) + { + if(strcmp(lw->arch.gfx, "gfx942") != 0) + return rocke_h_fail(lw, ROCKE_ERR_VALUE, "XF32 MMA requires gfx942"); + rocke_h_emitf(lw, + "f32x%d %s = __builtin_amdgcn_%s(__builtin_bit_cast(f32x2, %s), " + "__builtin_bit_cast(f32x2, %s), %s, 0, 0, 0);", + tf32_count, + rocke_h_name(lw, op->results[0]), + op_id, + rocke_h_name(lw, op->operands[0]), + rocke_h_name(lw, op->operands[1]), + rocke_h_name(lw, op->operands[2])); + return ROCKE_OK; + } + /* Ordinary FP32 controls use catalog IDs that have no legacy opcode. */ + if(strcmp(op_id, "mfma_f32_16x16x4_f32") == 0 || strcmp(op_id, "mfma_f32_32x32x2_f32") == 0) + { + bool small = strcmp(op_id, "mfma_f32_16x16x4_f32") == 0; + rocke_h_emitf(lw, + "f32x%d %s = __builtin_amdgcn_%s(%s, %s, %s, 0, 0, 0);", + small ? 4 : 16, + rocke_h_name(lw, op->results[0]), + small ? "mfma_f32_16x16x4f32" : "mfma_f32_32x32x2f32", + rocke_h_name(lw, op->operands[0]), + rocke_h_name(lw, op->operands[1]), + rocke_h_name(lw, op->operands[2])); + return ROCKE_OK; + } if(rocke_gfx1250_scaled_wmma(op_id)) { return h_emit_gfx1250_scaled_wmma(lw, op); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/core.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/core.cpp index 0fec791f2b90..979b5d492cfa 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/core.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/core.cpp @@ -20,6 +20,7 @@ * stub bodies remain in this file. */ #include "rocke/lower_llvm_internal.h" +#include "rocke/tf32_internal.h" #include #include @@ -978,7 +979,7 @@ const char* rocke_ll_llvm_type(rocke_lower_t* L, const rocke_type_t* t) return "i8"; if(strcmp(n, "i16") == 0) return "i16"; - if(strcmp(n, "i32") == 0) + if(strcmp(n, "i32") == 0 || strcmp(n, "tf32") == 0) return "i32"; if(strcmp(n, "i64") == 0) return "i64"; @@ -1076,7 +1077,7 @@ const char* rocke_ll_llvm_type_from_name(rocke_lower_t* L, const char* name) { rocke_ll_fail(L, ROCKE_ERR_NOTIMPL, "no LLVM type for (null)"); } - if(strcmp(name, "i32") == 0) + if(strcmp(name, "i32") == 0 || strcmp(name, "tf32") == 0) return "i32"; if(strcmp(name, "i64") == 0) return "i64"; @@ -1114,7 +1115,7 @@ const char* rocke_ll_llvm_type_from_name(rocke_lower_t* L, const char* name) lelem = "half"; else if(strcmp(elem, "bf16") == 0) lelem = "bfloat"; - else if(strcmp(elem, "i32") == 0) + else if(strcmp(elem, "i32") == 0 || strcmp(elem, "tf32") == 0) lelem = "i32"; else { @@ -1629,7 +1630,7 @@ static int ll_smem_seg_size(const rocke_type_t* stype) eb = 1; else if(strcmp(n, "f16") == 0 || strcmp(n, "bf16") == 0) eb = 2; - else if(strcmp(n, "i32") == 0 || strcmp(n, "f32") == 0) + else if(strcmp(n, "i32") == 0 || strcmp(n, "tf32") == 0 || strcmp(n, "f32") == 0) eb = 4; else if(strcmp(n, "i64") == 0) eb = 8; @@ -1932,6 +1933,9 @@ void rocke_ll_lower_op(rocke_lower_t* L, const rocke_op_t* op) { return; } + const char* tf32_error = rocke_tf32_op_error(op); + if(tf32_error) + rocke_ll_fail(L, ROCKE_ERR_VALUE, "%s", tf32_error); rocke_opcode_t oc = op->opcode; rocke_ll_op_fn fn = NULL; if(oc > ROCKE_OP_INVALID && oc < ROCKE_OP__COUNT) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/data.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/data.cpp index e6fd436551cb..f6e3e0150288 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 @@ -242,6 +242,12 @@ const rocke_ll_decl_t ROCKE_LL_INTRINSIC_DECLS[] = { {"mfma.f32.32x32x16.f16", "declare <16 x float> @llvm.amdgcn.mfma.f32.32x32x16.f16(<8 x half>, <8 x half>, <16 x " "float>, i32 immarg, i32 immarg, i32 immarg)"}, + {"mfma.f32.16x16x8.xf32", + "declare <4 x float> @llvm.amdgcn.mfma.f32.16x16x8.xf32(<2 x float>, <2 x float>, <4 x " + "float>, i32 immarg, i32 immarg, i32 immarg)"}, + {"mfma.f32.32x32x4.xf32", + "declare <16 x float> @llvm.amdgcn.mfma.f32.32x32x4.xf32(<2 x float>, <2 x float>, <16 x " + "float>, i32 immarg, i32 immarg, i32 immarg)"}, {"mfma.f32.16x16x4f32", "declare <4 x float> @llvm.amdgcn.mfma.f32.16x16x4f32(float, float, <4 x float>, i32 immarg, " "i32 immarg, i32 immarg)"}, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp index 5cf265459213..d3d53a814422 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp @@ -45,7 +45,7 @@ static int ll_elem_bytes(const char* name) { return 2; } - if(strcmp(name, "i32") == 0 || strcmp(name, "f32") == 0) + if(strcmp(name, "i32") == 0 || strcmp(name, "tf32") == 0 || strcmp(name, "f32") == 0) { return 4; } @@ -561,7 +561,7 @@ static void op_tile_smem_load_vN(rocke_lower_t* L, const rocke_op_t* op) { elem_bytes = 2; } - else if(strcmp(en, "i32") == 0 || strcmp(en, "f32") == 0) + else if(strcmp(en, "i32") == 0 || strcmp(en, "tf32") == 0 || strcmp(en, "f32") == 0) { elem_bytes = 4; } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mma.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mma.cpp index b998e9263138..94589e3766c9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mma.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mma.cpp @@ -19,6 +19,7 @@ * reached from _op_tile_mma's op_id routing, which mirrors the Python CDNA * ISABackend.emit_mma rebuilding ``tile.`` and re-dispatching. */ +#include "rocke/tf32_internal.h" #include #include @@ -90,6 +91,9 @@ static void _op_tile_mma(rocke_lower_t* L, const rocke_op_t* op) rocke_ll_fail(L, ROCKE_ERR_KEY, "tile.mma: missing op_id attribute"); } + if(rocke_tf32_mma_count(op_id) && strcmp(L->backend->gfx, "gfx942") != 0) + rocke_ll_fail(L, ROCKE_ERR_VALUE, "XF32 MMA requires gfx942"); + /* f16 / bf16 / f32 dense atoms resolve from the table; the scaled / fp4 / * fp6 / fp8-bf8 / hero atoms keep their dedicated bodies below. */ if(_try_emit_mfma_table(L, op, op_id)) @@ -623,6 +627,18 @@ typedef struct _mfma_spec } _mfma_spec_t; static const _mfma_spec_t MFMA_SPECS[] = { + {"mfma_f32_32x32x4_xf32", + "mfma.f32.32x32x4.xf32", + "llvm.amdgcn.mfma.f32.32x32x4.xf32", + "<2 x i32>", + "<16 x float>", + "<2 x float>"}, + {"mfma_f32_16x16x8_xf32", + "mfma.f32.16x16x8.xf32", + "llvm.amdgcn.mfma.f32.16x16x8.xf32", + "<2 x i32>", + "<4 x float>", + "<2 x float>"}, {"mfma_f32_16x16x16_f16", "mfma.f32.16x16x16f16", "llvm.amdgcn.mfma.f32.16x16x16f16", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp index d6520a8b7a2a..7ffca80cec6b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/verify.cpp @@ -23,6 +23,7 @@ */ #include "rocke/verify.h" +#include "rocke/tf32_internal.h" #include #include @@ -321,6 +322,9 @@ static void check_op(verifier_t* v, const rocke_op_t* op); static void check_contract(verifier_t* v, const rocke_op_t* op) { + const char* tf32_error = rocke_tf32_op_error(op); + if(tf32_error) + v_errf(v, op, "%s", tf32_error); rocke_opcode_t o = op->opcode; const char* name = op->name; if(is_binary_same_type(o)) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp index 9d8230329102..fa0e634d0aab 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/mma_io.cpp @@ -13,6 +13,8 @@ const rocke_type_t* rocke_storage_ir_type(const char* dtype) const rocke_dtype_info_t* info = rocke_dtype_info(dtype); if(!info) return NULL; + if(strcmp(info->name, "tf32") == 0) + return rocke_i32(); if(info->encoded_bits % 8 || strcmp(info->name, "e8m0") == 0 || strcmp(info->name, "e5m3") == 0) return rocke_i8(); return rocke_dtype_to_ir_type(info->name); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_tf32_mma_probe.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_tf32_mma_probe.h new file mode 100644 index 000000000000..9431f548329d --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_tf32_mma_probe.h @@ -0,0 +1,15 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +#ifndef ROCKE_INSTANCE_TF32_MMA_PROBE_H +#define ROCKE_INSTANCE_TF32_MMA_PROBE_H +#include "rocke/ir.h" +#ifdef __cplusplus +extern "C" { +#endif +/* Initializes b; returned kernel is owned by b. Mirrors Tf32MmaProbeSpec. */ +rocke_kernel_def_t* + rocke_build_tf32_mma_probe(rocke_ir_builder_t* b, int m, const char* preparation); +#ifdef __cplusplus +} +#endif +#endif diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h index a6fcf74b8ab5..8c6b09ed045a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir.h @@ -110,6 +110,7 @@ typedef enum rocke_scalar_kind ROCKE_SCALAR_FP6E3M2, ROCKE_SCALAR_E8M0, ROCKE_SCALAR_E5M3, + ROCKE_SCALAR_TF32, ROCKE_SCALAR__COUNT } rocke_scalar_kind_t; @@ -577,6 +578,7 @@ const rocke_type_t* rocke_i64(void); const rocke_type_t* rocke_bf16(void); const rocke_type_t* rocke_f16(void); const rocke_type_t* rocke_f32(void); +const rocke_type_t* rocke_tf32(void); const rocke_type_t* rocke_fp8e4m3(void); const rocke_type_t* rocke_bf8e5m2(void); const rocke_type_t* rocke_fp4e2m1(void); @@ -785,6 +787,8 @@ rocke_value_t* rocke_b_cvt_scalef32_pk_f32_fp8x4(rocke_ir_builder_t* b, rocke_value_t* rocke_b_cvt_scalef32_pk_f32_bf8x4(rocke_ir_builder_t* b, rocke_value_t* v, rocke_value_t* scale); +/* Explicit RNE conversion; bitcast preserves raw FP32 payloads instead. */ +rocke_value_t* rocke_b_cvt_f32_to_tf32(rocke_ir_builder_t* b, rocke_value_t* v); rocke_value_t* rocke_b_cvt_f32_to_fp8(rocke_ir_builder_t* b, rocke_value_t* v); rocke_value_t* rocke_b_cvt_f32_to_bf8(rocke_ir_builder_t* b, rocke_value_t* v); rocke_value_t* rocke_b_cvt_f32_to_i8_sat(rocke_ir_builder_t* b, rocke_value_t* v); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/tf32_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/tf32_internal.h new file mode 100644 index 000000000000..158f66df4e07 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/tf32_internal.h @@ -0,0 +1,65 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* Mirrors core/tf32.py. Validation sees logical types, never carrier equality. */ +#ifndef ROCKE_TF32_INTERNAL_H +#define ROCKE_TF32_INTERNAL_H +#include "rocke/ir.h" +#include +static inline int rocke_tf32_mma_count(const char* id) +{ + if(!id) + return 0; + if(strcmp(id, "mfma_f32_16x16x8_xf32") == 0) + return 4; + if(strcmp(id, "mfma_f32_32x32x4_xf32") == 0) + return 16; + return 0; +} +static inline const char* rocke_tf32_mma_error(const char* id, rocke_value_t* const* args, int n) +{ + int count = rocke_tf32_mma_count(id); + if(!count) + { + for(int i = 0; i < n; ++i) + if(args[i] + && (strcmp(args[i]->type->name, "tf32") == 0 + || strncmp(args[i]->type->name, "vec" : "vec"; + if(n != 3 || !args[0] || !args[1] || !args[2] || strcmp(args[0]->type->name, "vec") != 0 + || strcmp(args[1]->type->name, "vec") != 0 || strcmp(args[2]->type->name, acc) != 0) + return "XF32 MMA requires two vec operands and its FP32 accumulator"; + return NULL; +} +static inline const char* rocke_tf32_op_error(const rocke_op_t* op) +{ + const char* id = rocke_attr_get_str(&op->attrs, "op_id"); + if(!id && strncmp(op->name, "tile.", 5) == 0) + id = op->name + 5; + const char* error = (strcmp(op->name, "tile.mma") == 0 || strncmp(op->name, "tile.mfma", 9) == 0 + || strncmp(op->name, "tile.wmma", 9) == 0) + ? rocke_tf32_mma_error(id, op->operands, op->num_operands) + : NULL; + if(error) + return error; + int count = rocke_tf32_mma_count(id); + if(count + && (op->num_results != 1 + || strcmp(op->results[0]->type->name, count == 4 ? "vec" : "vec") != 0)) + return "XF32 MMA result must match its FP32 accumulator"; + if((strncmp(op->name, "arith.", 6) == 0 || strncmp(op->name, "math.", 5) == 0) + && strcmp(op->name, "arith.bitcast") != 0 && strcmp(op->name, "arith.select") != 0) + { + for(int i = 0; i < op->num_operands + op->num_results; ++i) + { + const rocke_value_t* v + = i < op->num_operands ? op->operands[i] : op->results[i - op->num_operands]; + if(strcmp(v->type->name, "tf32") == 0 || strncmp(v->type->name, "vec +#include + +rocke_kernel_def_t* rocke_build_tf32_mma_probe(rocke_ir_builder_t* b, int m, const char* mode) +{ + char name[96]; + snprintf(name, sizeof(name), "tf32_probe_%d_%s", m, mode ? mode : "invalid"); + if(rocke_ir_builder_init(b, name) != ROCKE_OK) + return NULL; + return ckc::guard_builder(b, [&]() -> rocke_kernel_def_t* { + if((m != 16 && m != 32) || !mode + || (strcmp(mode, "raw") && strcmp(mode, "carrier") && strcmp(mode, "rne") + && strcmp(mode, "prepacked") && strcmp(mode, "fp32"))) + return (rocke_kernel_def_t*)rocke_i_set_err( + b, ROCKE_ERR_VALUE, "TF32 probe requires m=16/32 and a known preparation"); + int k = 128 / m; + const rocke_arch_target_t* target = rocke_arch_target_from_gfx("gfx942"); + const rocke_mma_op_t* atom = rocke_mma_catalog_op_for_shape( + &target->mma, "mma", "tf32", "tf32", "fp32", m, m, k, NULL); + const rocke_type_t* ty = strcmp(mode, "prepacked") == 0 ? rocke_i32() : rocke_f32(); + rocke_param_opts_t opts = {}; + opts.align = 4; + opts.align_set = true; + const rocke_type_t* input_ptr = rocke_ptr_type(b, ty, "global"); + rocke_value_t* a = rocke_b_param(b, "A", input_ptr, &opts); + rocke_value_t* bb = rocke_b_param(b, "B", input_ptr, &opts); + const rocke_type_t* float_ptr = rocke_ptr_type(b, rocke_f32(), "global"); + rocke_value_t* cc = rocke_b_param(b, "C", float_ptr, &opts); + rocke_value_t* d = rocke_b_param(b, "D", float_ptr, &opts); + const rocke_type_t* bits_ptr = rocke_ptr_type(b, rocke_i32(), "global"); + rocke_value_t* pa = rocke_b_param(b, "PA", bits_ptr, &opts); + rocke_value_t* pb = rocke_b_param(b, "PB", bits_ptr, &opts); + rocke_value_t* lane = rocke_b_thread_id_x(b); + rocke_value_t* batch = rocke_b_block_id_x(b); + rocke_value_t* cm = rocke_b_const_i32(b, m); + rocke_value_t* axis = rocke_b_mod(b, lane, cm); + rocke_value_t* group = rocke_b_div(b, lane, cm); + rocke_value_t* ck = rocke_b_const_i32(b, k); + rocke_value_t* row = rocke_b_mul(b, axis, ck); + rocke_value_t* size = rocke_b_const_i32(b, m * k); + rocke_value_t* batch_input = rocke_b_mul(b, batch, size); + rocke_value_t* row_base = rocke_b_add(b, batch_input, row); + size = rocke_b_const_i32(b, m * m); + rocke_value_t* batch_output = rocke_b_mul(b, batch, size); + rocke_value_t* indices[16]; + rocke_value_t* c_values[16]; + for(int slot = 0; slot < atom->c_frag_len; ++slot) + { + rocke_value_t *r, *col; + rocke_layout_map_coord(atom->c_layout, b, lane, slot, &r, &col); + rocke_value_t* offset = rocke_b_mul(b, r, cm); + offset = rocke_b_add(b, offset, col); + indices[slot] = rocke_b_add(b, batch_output, offset); + c_values[slot] = rocke_b_global_load(b, cc, indices[slot], rocke_f32(), 4); + } + rocke_value_t* acc = rocke_b_vec_pack(b, c_values, atom->c_frag_len, rocke_f32()); + if(strcmp(mode, "fp32") == 0) + { + const rocke_mma_op_t* full = rocke_mma_catalog_op_for_shape( + &target->mma, "mma", "fp32", "fp32", "fp32", m, m, k / 2, NULL); + for(int step = 0; step < 2; ++step) + { + rocke_value_t* delta = rocke_b_const_i32(b, step * (64 / m)); + delta = rocke_b_add(b, group, delta); + rocke_value_t* index = rocke_b_add(b, row_base, delta); + rocke_value_t* av = rocke_b_global_load(b, a, index, rocke_f32(), 4); + rocke_value_t* bv = rocke_b_global_load(b, bb, index, rocke_f32(), 4); + rocke_value_t* bits = rocke_b_bitcast(b, av, rocke_i32()); + rocke_b_global_store(b, pa, index, bits, 4); + bits = rocke_b_bitcast(b, bv, rocke_i32()); + rocke_b_global_store(b, pb, index, bits, 4); + acc = rocke_b_mma(b, full->op_id, av, bv, acc, NULL, 0); + } + } + else + { + rocke_matrix_fragment_layout_t layout = {{{32, 32}, 2, 32, 2}, 2, 64 / m, m}; + rocke_value_t* inputs[] = {a, bb}; + rocke_value_t* prepared[] = {pa, pb}; + rocke_value_t* fragments[2]; + for(int role = 0; role < 2; ++role) + { + bool packed = strcmp(mode, "prepacked") == 0; + const rocke_type_t* carrier + = packed || strcmp(mode, "carrier") == 0 ? rocke_i32() : rocke_f32(); + rocke_value_t* values = rocke_h_load_matrix_fragment(b, + inputs[role], + row_base, + group, + 0, + packed ? "tf32" : "fp32", + &layout, + carrier, + 4); + if(strcmp(mode, "rne") == 0) + { + rocke_value_t* elems[2]; + for(int i = 0; i < 2; ++i) + { + rocke_value_t* e = rocke_b_vec_extract(b, values, i); + elems[i] = rocke_b_cvt_f32_to_tf32(b, e); + } + values = rocke_b_vec_pack(b, elems, 2, rocke_tf32()); + } + else + values = rocke_b_bitcast(b, values, rocke_vector_type(b, rocke_tf32(), 2)); + for(int slot = 0; slot < 2; ++slot) + { + rocke_value_t* two = rocke_b_const_i32(b, 2); + rocke_value_t* offset = rocke_b_mul(b, group, two); + rocke_value_t* index = rocke_b_const_i32(b, slot); + offset = rocke_b_add(b, offset, index); + index = rocke_b_add(b, row_base, offset); + rocke_value_t* e = rocke_b_vec_extract(b, values, slot); + rocke_value_t* bits = rocke_b_bitcast(b, e, rocke_i32()); + rocke_b_global_store(b, prepared[role], index, bits, 4); + } + fragments[role] = values; + } + acc = rocke_b_mma(b, atom->op_id, fragments[0], fragments[1], acc, NULL, 0); + } + for(int slot = 0; slot < atom->c_frag_len; ++slot) + { + rocke_value_t* value = rocke_b_vec_extract(b, acc, slot); + rocke_b_global_store(b, d, indices[slot], value, 4); + } + rocke_b_ret(b); + return rocke_ir_builder_kernel(b); + }); +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md index ab5a8349a6f6..e88fbe70f20b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/dtype_storage_packing.md @@ -40,6 +40,12 @@ from raw byte patterns with the common bit packer. Scale-format validation remains separate from matrix alias normalization: E5M3 is not BF8 E5M2. Registering a scale dtype does not enable a backend selector. +Logical `tf32` (`xf32` alias) is a distinct 32-bit type carried as I32. +Fragment loading can reinterpret FP32 storage into I32 without conversion; +`bitcast` then wraps those bits as logical TF32. Explicit `cvt_f32_to_tf32` +performs RNE preparation. Packing itself never rounds. See the +[TF32 numerical example](../../python/rocke/examples/gfx942/tf32_numerics/README.md). + ## Common packing `BitPacking` / `rocke_bit_packing_t` describes unsigned bit patterns with an @@ -49,6 +55,7 @@ bit first in a little-endian byte stream. | Format | Dense group occupying whole 32-bit words | |---|---| +| TF32 | 1 element in 1 word | | FP16/BF16 | 2 elements in 1 word | | FP8/BF8 | 4 elements in 1 word | | FP6/BF6 | 16 elements across 3 words | diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/ir_serialization_format.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/ir_serialization_format.md index 6e4cc3199de5..a96210026896 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/ir_serialization_format.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/ir_serialization_format.md @@ -386,3 +386,9 @@ kernel @add_one { (`parse` of that text rebuilds a `KernelDef` whose `lower_kernel_to_llvm` output is identical to the original's.) + +The scalar name `tf32` denotes logical TF32 with a 32-bit I32 carrier. It may +appear in vectors, pointers, shared-memory types and loop carries, using the +existing type grammar (for example, `vec`). Serialization preserves its +logical identity independently of `f32` and `i32`; no format-version change is +needed. The public dtype alias `xf32` normalizes to `tf32` before IR construction. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/mfma_atom_catalog.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/mfma_atom_catalog.md index 24ccdf131042..c2d53b3c61db 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/mfma_atom_catalog.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/mfma_atom_catalog.md @@ -137,3 +137,10 @@ A `tile_m=128, tile_n=128, warp_m=2, warp_n=2, warp_tile_m=32, warp_tile_n=32` G - Using `f16_16x16x32` or `f16_32x32x16` on gfx940/gfx942: link-time failure (intrinsic absent). - Using bf16 atoms without going through `mfma_16x16x16_for_dtype`: declaration mismatch (plain `bf16.16x16x16` doesn't exist; must use the `_1k` variant). - `4x4x4` epilogue stores not compositing `batch_idx` into the per-group offset: 16 batches collide on the same outputs. + +Logical `tf32` on gfx942 has two native atoms: `mfma_f32_16x16x8_xf32` and +`mfma_f32_32x32x4_xf32`. Both use two TF32 operands per lane for each input, +stored as two I32 words; accumulator sizes are four and sixteen FP32 values. +The [numerical example](../../python/rocke/examples/gfx942/tf32_numerics/README.md) +demonstrates raw input transport and explicit RNE preparation. These atoms do +not alter ordinary FP32 selection or enable TF32 production GEMM dispatch. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/mma_metadata.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/mma_metadata.md index 41642b1170ce..3f9597229488 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/mma_metadata.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/mma_metadata.md @@ -128,3 +128,15 @@ Existing `wmma_scale*_f32_*` and dotted `wmma.scaled.*` IDs are retired, as are the dedicated scaled builder wrappers; serialized IR using those IDs must be regenerated. Use `tile.mma` with a resolved catalog atom. Other MMA operation IDs retain their existing spelling. + +## Logical TF32 operands + +On gfx942, `tf32` (`xf32` alias) selects `mfma_f32_16x16x8_xf32` or +`mfma_f32_32x32x4_xf32`. Both take two logical TF32 A/B elements per lane; +C/D are FP32 with four or sixteen elements. FP32 queries retain the ordinary +FP32 atoms. The logical type uses one I32 storage word per element and is +bitcast to the intrinsic's float-vector ABI during lowering. + +See the [TF32 numerical example](../../python/rocke/examples/gfx942/tf32_numerics/README.md) +for the distinction between raw FP32 payloads and explicit RNE preparation, +and for a reproducible comparison through both engines. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/op_vocabulary.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/op_vocabulary.md index 54305867692f..ee9a663cea9a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/op_vocabulary.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/op_vocabulary.md @@ -179,3 +179,10 @@ static_if(host_bool, body) # raises TypeError on SSA Value - `ir_lowering/ir_model.md` for the IRBuilder API by category. - `ir_lowering/lowering_pipeline.md` for the op-to-LLVM mapping. - `ir_lowering/backend_details.md` for purity classification and waitcnt encoding. + +`IRBuilder.cvt_f32_to_tf32` / `rocke_b_cvt_f32_to_tf32` is a scalar F32-to-TF32 +RNE recipe built from existing integer arithmetic, select, and bitcast ops; it +introduces no serialized opcode. It quiets NaNs and retains their upper payload +bits. `bitcast` to logical TF32 only reinterprets the payload. Logical TF32 uses +I32 storage; general arithmetic requires explicitly reinterpreting it as F32. +See [TF32 numerics](../../python/rocke/examples/gfx942/tf32_numerics/README.md). diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/arch_specs.json b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/arch_specs.json index 3d12a94fae34..76a3d5a62031 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/arch_specs.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/arch_specs.json @@ -45,6 +45,8 @@ "sgprs": 102 }, "mma": [ + {"family": "mma", "a": "tf32", "b": "tf32", "c": "fp32", "m": 16, "n": 16, "k": 8, "op_id": "mfma_f32_16x16x8_xf32"}, + {"family": "mma", "a": "tf32", "b": "tf32", "c": "fp32", "m": 32, "n": 32, "k": 4, "op_id": "mfma_f32_32x32x4_xf32"}, {"family": "mma", "a": "fp32", "b": "fp32", "c": "fp32", "m": 16, "n": 16, "k": 4, "op_id": "mfma_f32_16x16x4_f32"}, {"family": "mma", "a": "fp32", "b": "fp32", "c": "fp32", "m": 32, "n": 32, "k": 2, "op_id": "mfma_f32_32x32x2_f32"}, {"family": "mma", "a": "fp16", "b": "fp16", "c": "fp32", "m": 16, "n": 16, "k": 16, "op_id": "mfma_f32_16x16x16_f16"}, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/target.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/target.py index fe94ac88f845..0c98b45c999d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/target.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/target.py @@ -330,6 +330,38 @@ def _mfma_b_16x16(builder, lane, slot): return k, n_in_atom +def _mfma_a_16x16x8_xf32(builder, lane, slot): + c = builder.const_i32(16) + axis = builder.mod(lane, c) + group = builder.div(lane, c) + k = builder.add(builder.mul(group, builder.const_i32(2)), builder.const_i32(slot)) + return axis, k + + +def _mfma_b_16x16x8_xf32(builder, lane, slot): + c = builder.const_i32(16) + axis = builder.mod(lane, c) + group = builder.div(lane, c) + k = builder.add(builder.mul(group, builder.const_i32(2)), builder.const_i32(slot)) + return k, axis + + +def _mfma_a_32x32x4_xf32(builder, lane, slot): + c = builder.const_i32(32) + axis = builder.mod(lane, c) + group = builder.div(lane, c) + k = builder.add(builder.mul(group, builder.const_i32(2)), builder.const_i32(slot)) + return axis, k + + +def _mfma_b_32x32x4_xf32(builder, lane, slot): + c = builder.const_i32(32) + axis = builder.mod(lane, c) + group = builder.div(lane, c) + k = builder.add(builder.mul(group, builder.const_i32(2)), builder.const_i32(slot)) + return k, axis + + def _mfma_a_16x16x4_f32(builder, lane, slot): """MFMA 16x16x4 fp32 A operand: each lane holds one fp32 scalar. @@ -659,6 +691,12 @@ class _FragInfo: # layout-map functions are populated for the atoms whose lane math is verified. # Adding a new atom is one row here. _MMA_FRAGMENT_INFO: Dict[str, _FragInfo] = { + "mfma_f32_32x32x4_xf32": _FragInfo( + 2, 2, 16, 64, _mfma_a_32x32x4_xf32, _mfma_b_32x32x4_xf32, _mfma_acc_32x32 + ), + "mfma_f32_16x16x8_xf32": _FragInfo( + 2, 2, 4, 64, _mfma_a_16x16x8_xf32, _mfma_b_16x16x8_xf32, _mfma_acc_16x16 + ), # --- MFMA fp32 (wave64) ----------------------------------------------- # A/B are scalar float per lane (a_frag_len=b_frag_len=1); accumulator # shares the standard 16x16 / 32x32 layout (c_frag_len=4 / 16). diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/dtypes.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/dtypes.py index b0745c226610..731eafe5efa6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/dtypes.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/dtypes.py @@ -23,6 +23,8 @@ "f32": "fp32", "float": "fp32", "fp32": "fp32", + "tf32": "tf32", + "xf32": "tf32", "fp8": "fp8e4m3", "fp8e4m3": "fp8e4m3", "e4m3": "fp8e4m3", @@ -76,6 +78,7 @@ class DTypeInfo: DTypeCategory.FLOAT, ( ("fp32", 32), + ("tf32", 32), ("fp16", 16), ("bf16", 16), ("fp8e4m3", 8), diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py index 0ebe58891330..b5a436cff555 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py @@ -51,6 +51,7 @@ def __repr__(self) -> str: BF16 = Type("bf16") F16 = Type("f16") F32 = Type("f32") +TF32 = Type("tf32") FP8E4M3 = Type("fp8e4m3") BF8E5M2 = Type("bf8e5m2") FP4E2M1 = Type("fp4e2m1") @@ -75,6 +76,7 @@ def dtype_to_ir_type(dtype: str) -> Type: "fp16": F16, "bf16": BF16, "fp32": F32, + "tf32": TF32, "fp8e4m3": FP8E4M3, "bf8e5m2": BF8E5M2, "fp4e2m1": FP4E2M1, @@ -911,6 +913,49 @@ def cast_to_f32(self, v: Value) -> Value: raise ValueError(f"cast_to_f32 unsupported from {v.type.name}") return self._op("arith.cast_to_f32", [v], [F32], result_name_hint="f32").result + def cvt_f32_to_tf32(self, v: Value) -> Value: + """Round f32 to TF32 precision (RNE), carried as i32 bits. + + Finite encodings have 13 zero low bits. Infinities and signed zero are + preserved; NaNs are quieted and retain their upper payload bits. + Use bitcast(v, TF32) to request native XF32 truncation without rounding. + """ + if v.type != F32: + raise ValueError("cvt_f32_to_tf32 expects f32 input") + bits = self.bitcast(v, I32) + magnitude_mask = self.const_i32(2147483647) + mag = self.land(bits, magnitude_mask) + sign_mask = self.const_i32(-2147483648) + sign = self.land(bits, sign_mask) + exponent_mask = self.const_i32(2139095040) + exp = self.land(mag, exponent_mask) + exponent_all_ones = self.const_i32(2139095040) + special = self.cmp_eq(exp, exponent_all_ones) + zero = self.const_i32(0) + safe = self.select(special, zero, mag) + discarded_bits = self.const_i32(13) + shift = self.lshr(safe, discarded_bits) + low_bit = self.const_i32(1) + odd = self.land(shift, low_bit) + rounding_bias = self.const_i32(4095) + bias = self.add(odd, rounding_bias) + rounded = self.add(safe, bias) + precision_mask = self.const_i32(-8192) + rounded = self.land(rounded, precision_mask) + rounded = self.lor(rounded, sign) + fraction_mask = self.const_i32(8388607) + frac = self.land(mag, fraction_mask) + zero_fraction = self.const_i32(0) + is_nan = self.cmp_ne(frac, zero_fraction) + quiet_bit = self.const_i32(4194304) + quiet = self.lor(bits, quiet_bit) + nan_precision_mask = self.const_i32(-8192) + quiet = self.land(quiet, nan_precision_mask) + nonfinite = self.select(is_nan, quiet, bits) + result = self.select(special, nonfinite, rounded) + result = self.bitcast(result, TF32) + return result + def cast_f32_to(self, v: Value, target: Type) -> Value: if v.type.name != "f32": raise ValueError("cast_f32_to expects f32 input") @@ -1778,7 +1823,7 @@ def smem_store_vN( elem_bytes = ( 1 if elem_name in ("i8", "fp8e4m3", "bf8e5m2") - else 4 if elem_name in ("f32", "i32") else 2 + else 4 if elem_name in ("f32", "i32", "tf32") else 2 ) self._op( "tile.smem_store_vN", @@ -1816,9 +1861,18 @@ def smem_load_vN(self, smem: Value, *indices, dtype: Type, n: int = 0) -> Value: 96-bit payloads (12 bytes, six halfwords, or three words), using element alignment. The target and alignment determine instruction selection. """ - if dtype.name not in ("f16", "bf16", "f32", "i32", "fp8e4m3", "bf8e5m2", "i8"): + if dtype.name not in ( + "f16", + "bf16", + "f32", + "i32", + "tf32", + "fp8e4m3", + "bf8e5m2", + "i8", + ): raise ValueError( - "smem_load_vN supports f16 / bf16 / f32 / i32 / fp8e4m3 / " + "smem_load_vN supports f16 / bf16 / f32 / i32 / tf32 / fp8e4m3 / " f"bf8e5m2 / i8, got {dtype.name}" ) allowed_n = ( @@ -1874,6 +1928,11 @@ def mma( (``a_scale``, ``b_scale``); ordinary atoms take exactly ``a, b, c``. """ op_id = op.op_id if hasattr(op, "op_id") else str(op) + from .tf32 import tf32_mma_error + + error = tf32_mma_error(op_id, [a, b, c, *extra]) + if error: + raise ValueError(error) c_frag_len = ( op.c_frag_len if hasattr(op, "c_frag_len") and op.c_frag_len @@ -4187,7 +4246,7 @@ def global_store_vN( elem_bytes = 2 if n == 16: raise ValueError(f"global_store_vN n=16 not supported for {elem_name}") - elif elem_name in ("f32", "i32"): + elif elem_name in ("f32", "i32", "tf32"): elem_bytes = 4 if n == 16: raise ValueError(f"global_store_vN n=16 not supported for {elem_name}") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py index ab29bc41586b..c572abc0c198 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py @@ -36,6 +36,7 @@ "i8": "int8_t", "i16": "int16_t", "i32": "int", + "tf32": "int", "i64": "int64_t", "f16": "fp16", "bf16": "bf16", @@ -67,7 +68,7 @@ def _type_to_hip(t) -> str: return f"bf16x{t.count}" if elem == "f32": return f"f32x{t.count}" - if elem == "i32": + if elem in ("i32", "tf32"): return f"i32x{t.count}" if elem == "i16": return f"i16x{t.count}" @@ -182,6 +183,7 @@ def _name(v: Value) -> str: "bf16": "bf16x", "f32": "f32x", "i32": "i32x", + "tf32": "i32x", "i16": "i16x", "i8": "i8x", "fp8e4m3": "i8x", @@ -312,6 +314,11 @@ def _pop_indent(self) -> None: self._indent -= 1 def lower_op(self, op: Op) -> None: + from .tf32 import tf32_op_error + + error = tf32_op_error(op) + if error: + raise ValueError(error) if op.name.startswith("tile.") and gfx1250_scaled_wmma(op.name) is not None: self._emit_wmma_gfx1250_scaled(op) return @@ -552,6 +559,19 @@ def _op_tile_mma(self, op: Op) -> None: This keeps the HIP emission identical to the legacy ISA-named path while the IRBuilder helpers route through :meth:`IRBuilder.mma`. """ + from .tf32 import TF32_MMA + + op_id = op.attrs["op_id"] + if op_id in TF32_MMA: + if self.arch.gfx != "gfx942": + raise ValueError("XF32 MMA requires gfx942") + a, b, c = op.operands + self._emit( + f"f32x{TF32_MMA[op_id]} {_name(op.result)} = __builtin_amdgcn_{op_id}(" + f"__builtin_bit_cast(f32x2, {_name(a)}), " + f"__builtin_bit_cast(f32x2, {_name(b)}), {_name(c)}, 0, 0, 0);" + ) + return op_id = op.attrs["op_id"] legacy = Op( name=f"tile.{op_id}", 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 ebf67fd5af12..8d0b5ae00516 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 @@ -661,6 +661,8 @@ def _resolve_llvm_flavor() -> str: "<8 x half>, <8 x half>, <16 x float>, " "i32 immarg, i32 immarg, i32 immarg)" ), + "mfma.f32.16x16x8.xf32": "declare <4 x float> @llvm.amdgcn.mfma.f32.16x16x8.xf32(<2 x float>, <2 x float>, <4 x float>, i32 immarg, i32 immarg, i32 immarg)", + "mfma.f32.32x32x4.xf32": "declare <16 x float> @llvm.amdgcn.mfma.f32.32x32x4.xf32(<2 x float>, <2 x float>, <16 x float>, i32 immarg, i32 immarg, i32 immarg)", "mfma.f32.16x16x4f32": ( "declare <4 x float> @llvm.amdgcn.mfma.f32.16x16x4f32(" "float, float, <4 x float>, " @@ -1112,7 +1114,7 @@ def _llvm_type(t: Type) -> str: return "i8" if t.name == "i16": return "i16" - if t.name == "i32": + if t.name in ("i32", "tf32"): return "i32" if t.name == "i64": return "i64" @@ -1801,6 +1803,7 @@ def _compute_smem_layout(self) -> None: "f16": 2, "bf16": 2, "i32": 4, + "tf32": 4, "f32": 4, "i64": 8, } @@ -1982,6 +1985,11 @@ def _emit_smem_base_ptr(self, gname: str, stype: SmemType) -> str: # ----- per-op lowerings ----- def lower_op(self, op: Op) -> None: + from .tf32 import tf32_op_error + + error = tf32_op_error(op) + if error: + raise ValueError(error) method = getattr(self, f"_op_{op.name.replace('.', '_')}", None) if method is None: raise NotImplementedError(f"no LLVM lowering for op {op.name!r}") @@ -2956,6 +2964,7 @@ def _op_tile_smem_store(self, op: Op) -> None: "f16": 2, "bf16": 2, "i32": 4, + "tf32": 4, "f32": 4, "i64": 8, }.get(value.type.name, 2) @@ -3020,6 +3029,7 @@ def _op_tile_smem_store_vN(self, op: Op) -> None: "f16": 2, "bf16": 2, "i32": 4, + "tf32": 4, "f32": 4, "i64": 8, }.get( @@ -3097,7 +3107,15 @@ def _op_tile_smem_load_vN(self, op: Op) -> None: elem_ty = _llvm_type(op.result.type.elem) # type: ignore[attr-defined] # Element byte size drives the vector alignment. 16-bit # (f16 / bf16): 2 bytes; 32-bit (f32 / i32): 4 bytes. - elem_bytes = {"i8": 1, "f16": 2, "bf16": 2, "i32": 4, "f32": 4, "i64": 8}.get( + elem_bytes = { + "i8": 1, + "f16": 2, + "bf16": 2, + "i32": 4, + "tf32": 4, + "f32": 4, + "i64": 8, + }.get( op.result.type.elem.name, 2, # type: ignore[attr-defined] ) @@ -3140,6 +3158,10 @@ def _op_tile_wmma_f32_16x16x16_bf16(self, op: Op) -> None: self._backend.emit_wmma(self, op) def _op_tile_mma(self, op: Op) -> None: + from .tf32 import TF32_MMA + + if op.attrs.get("op_id") in TF32_MMA and self._backend.arch.gfx != "gfx942": + raise ValueError("XF32 MMA requires gfx942") # Target-neutral MMA: the ISA backend maps ``op.attrs["op_id"]`` to the # matching MFMA (CDNA) or WMMA (RDNA) emission. CDNA backends reuse the # existing ``_op_tile_`` handler verbatim, so the output is @@ -3199,6 +3221,31 @@ def _op_tile_mfma_f32_16x16x32_bf16(self, op: Op) -> None: f"i32 0, i32 0, i32 0)" ) + def _op_tile_mfma_f32_16x16x8_xf32(self, op: Op) -> None: + self._emit_xf32(op, "mfma.f32.16x16x8.xf32", 4) + + def _op_tile_mfma_f32_32x32x4_xf32(self, op: Op) -> None: + self._emit_xf32(op, "mfma.f32.32x32x4.xf32", 16) + + def _emit_xf32(self, op: Op, intrinsic: str, count: int) -> None: + if self._backend.arch.gfx != "gfx942": + raise ValueError("XF32 MMA requires gfx942") + a, b, c = op.operands + self._need(intrinsic) + a_cast = self._fresh("mfma_a_i16") + b_cast = self._fresh("mfma_b_i16") + self._current().emit( + f" {a_cast} = bitcast <2 x i32> {self._operand(a)} to <2 x float>" + ) + self._current().emit( + f" {b_cast} = bitcast <2 x i32> {self._operand(b)} to <2 x float>" + ) + self._current().emit( + f" {op.result.name} = call <{count} x float> @llvm.amdgcn.{intrinsic}(" + f"<2 x float> {a_cast}, <2 x float> {b_cast}, " + f"<{count} x float> {self._operand(c)}, i32 0, i32 0, i32 0)" + ) + def _op_tile_mfma_f32_16x16x4_f32(self, op: Op) -> None: a, b, c = op.operands self._need("mfma.f32.16x16x4f32") @@ -5292,6 +5339,7 @@ def _op_memref_global_store_vN(self, op: Op) -> None: "f16": 2, "bf16": 2, "i32": 4, + "tf32": 4, "f32": 4, "i64": 8, }.get(elem_name, 2) @@ -6081,7 +6129,7 @@ def _format_agpr_alloc(value: object) -> str: def _llvm_type_from_name(name: str) -> str: """Map our IR type-name string (from op.attrs) back to LLVM IR text.""" - if name == "i32": + if name in ("i32", "tf32"): return "i32" if name == "i64": return "i64" @@ -6100,7 +6148,13 @@ def _llvm_type_from_name(name: str) -> str: inner = name[4:-1] elem, _, count = inner.partition("x") count = int(count) - elem_map = {"f32": "float", "f16": "half", "bf16": "bfloat", "i32": "i32"} + elem_map = { + "tf32": "i32", + "f32": "float", + "f16": "half", + "bf16": "bfloat", + "i32": "i32", + } return f"<{count} x {elem_map[elem]}>" raise NotImplementedError(f"no LLVM type for {name!r}") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/tf32.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/tf32.py new file mode 100644 index 000000000000..b836d53b0b57 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/tf32.py @@ -0,0 +1,43 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Logical TF32 operand contracts, independent of the physical i32 carrier.""" + +TF32_MMA = {"mfma_f32_16x16x8_xf32": 4, "mfma_f32_32x32x4_xf32": 16} + + +def tf32_mma_error(op_id, operands, results=None): + count = TF32_MMA.get(op_id) + if count is None: + if any( + v.type.name == "tf32" or v.type.name.startswith("vec", "vec", f"vec"] + if [v.type.name for v in operands] != expected: + return "XF32 MMA requires two vec operands and its FP32 accumulator" + if results is not None and [v.type.name for v in results] != [expected[2]]: + return "XF32 MMA result must match its FP32 accumulator" + return None + + +def tf32_op_error(op): + op_id = op.attrs.get("op_id", op.name.removeprefix("tile.")) + error = ( + tf32_mma_error(op_id, op.operands, op.results) + if op.name == "tile.mma" or op.name.startswith(("tile.mfma", "tile.wmma")) + else None + ) + if error: + return error + if op.name.startswith(("arith.", "math.")) and op.name not in ( + "arith.bitcast", + "arith.select", + ): + if any( + v.type.name == "tf32" or v.type.name.startswith("vec None: # ---- contracts ---- def _check_contract(self, op: Op, scope: Dict[str, Value]) -> None: + from .tf32 import tf32_op_error + + error = tf32_op_error(op) + if error: + self.err(error, op) name = op.name if name in _BINARY_SAME_TYPE: if len(op.operands) != 2 or len(op.results) != 1: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics.py new file mode 100644 index 000000000000..83bb36dd3e49 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics.py @@ -0,0 +1,352 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Compare native XF32 input preparation with independent NumPy references.""" + +import argparse +import ctypes +import json +import os +from pathlib import Path +import struct + +import numpy as np + +from rocke.core.ir_serialize import serialize +from rocke.core.lower_hip import lower_kernel_to_hip +from rocke.core.lower_llvm import _lower_kernel_to_llvm_python +from rocke.instances.gfx942.tf32_mma_probe import ( + PREPARATIONS, + Tf32MmaProbeSpec, + build_tf32_mma_probe, +) + + +def rne_bits(bits): + """Independent integer quotient/remainder reference; quiet NaNs.""" + u = np.asarray(bits, dtype=np.uint32) + sign = u & np.uint32(0x80000000) + magnitude = u & np.uint32(0x7FFFFFFF) + quotient = magnitude.astype(np.uint64) // 8192 + remainder = magnitude % 8192 + rounded = ( + quotient + ((remainder > 4096) | ((remainder == 4096) & (quotient % 2 == 1))) + ) * 8192 + finite = rounded.astype(np.uint32) | sign + special = u & np.uint32(0xFFFFE000) + special = np.where(magnitude > 0x7F800000, special | np.uint32(0x400000), special) + return np.where(magnitude < 0x7F800000, finite, special).astype(np.uint32) + + +def corpus(m): + k = 128 // m + # Exact values, ties (even/odd), either side, carry, normals, subnormals, + # infinities, signed zeros, signaling and quiet NaNs, both signs. + bits = np.array( + [ + 0, + 1, + 0x1000, + 0x7FFFFF, + 0x800000, + 0x801001, + 0x3F800000, + 0x3F800001, + 0x3F800FFF, + 0x3F801000, + 0x3F801001, + 0x3F802FFF, + 0x3F803000, + 0x3F803001, + 0x3FFFFFFF, + 0x7F7FFFFF, + 0x7F800000, + 0x7F800001, + 0x7FC00001, + ], + dtype=np.uint32, + ) + bits = np.concatenate((bits, bits | np.uint32(0x80000000))) + rng = np.random.default_rng(1714) + finite = rng.uniform(-8, 8, 1024).astype(np.float32).view(np.uint32) + bits = np.concatenate((bits, finite)) + single = len(bits) * 2 + a = np.zeros((single + 32, m, k), dtype=np.float32) + b = np.zeros_like(a) + c = np.zeros((len(a), m, m), dtype=np.float32) + a[: len(bits), :, 0] = bits.view(np.float32)[:, None] + b[: len(bits), :, 0] = 1 + a[len(bits) : single, :, 0] = 1 + b[len(bits) : single, :, 0] = bits.view(np.float32)[:, None] + a[single:] = rng.uniform(-2, 2, a[single:].shape) + b[single:] = rng.uniform(-2, 2, b[single:].shape) + c[single:] = rng.uniform(-1, 1, c[single:].shape) + # Asymmetric dense inputs and cancellation with a nonzero accumulator. + a[-1, :, 1::2] = -a[-1, :, ::2] + b[-1, :, 1::2] = b[-1, :, ::2] + return a, b, c, bits, single + + +def same_values(a, b): + # Compare every finite bit (including zero sign); NaN payloads are not an + # arithmetic-output contract. Prepared input payloads are checked separately. + return (a.view(np.uint32) == b.view(np.uint32)) | (np.isnan(a) & np.isnan(b)) + + +def launch(blob, name, a, b, c): + from rocke.runtime.hip_module import Runtime + + runtime = Runtime() + module = runtime.load_module(blob) + arrays = [ + a.copy(), + b.copy(), + c.copy(), + np.full_like(c, np.nan), + np.full(a.shape, 0xDEADBEEF, np.uint32), + np.full(b.shape, 0xDEADBEEF, np.uint32), + ] + ptrs = [] + try: + for array in arrays: + ptr = runtime.alloc(array.nbytes) + ptrs.append(ptr) + runtime.memcpy_h2d( + ptr, (ctypes.c_ubyte * array.nbytes).from_buffer(array), array.nbytes + ) + runtime.launch_blocking( + module.get_function(name), + (len(a), 1, 1), + (64, 1, 1), + struct.pack("6Q", *ptrs), + ) + for array, ptr in zip(arrays[3:], ptrs[3:]): + runtime.memcpy_d2h( + (ctypes.c_ubyte * array.nbytes).from_buffer(array), ptr, array.nbytes + ) + return arrays[3:] + finally: + for ptr in ptrs: + runtime.free(ptr) + module.unload() + + +def run(output_dir, backend="both", shapes=(16, 32)): + from rocke.runtime.comgr import build_hsaco_from_llvm_ir, resolved_lib_path + from rocke.runtime.hip_module import get_device_arch + + arch = get_device_arch() + if not arch or arch.split(":")[0] != "gfx942": + raise RuntimeError(f"TF32 numerical validation requires gfx942, got {arch}") + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + engines = ("python", "cpp") if backend == "both" else (backend,) + native = None + if "cpp" in engines: + import rocke_engine as native # Required: no fallback. + summary = { + "arch": arch, + "numpy": np.__version__, + "flavor": os.getenv("ROCKE_LLVM_FLAVOR"), + "results": [], + } + for m in shapes: + a, b, c, bits, single = corpus(m) + np.savez(output_dir / f"inputs_{m}.npz", a=a, b=b, c=c) + all_results = {} + for engine in engines: + results = {} + for variant in (*PREPARATIONS, "masked"): + mode = "prepacked" if variant == "masked" else variant + spec = Tf32MmaProbeSpec(m, mode) + kernel = build_tf32_mma_probe(spec) + ir = serialize(kernel) + if engine == "cpp": + ir_native = native.tf32_mma_probe_serialize_ir(m, mode) + assert ir == ir_native, (m, mode, "builder parity") + ll = native.lower_serialized_ir( + ir_native, + arch="gfx942", + flavor=os.getenv("ROCKE_LLVM_FLAVOR", "llvm23"), + ) + assert ll == _lower_kernel_to_llvm_python(kernel, arch="gfx942"), ( + m, + mode, + "lowerer parity", + ) + else: + ll = _lower_kernel_to_llvm_python(kernel, arch="gfx942") + stem = output_dir / f"{engine}_{m}_{variant}" + stem.with_suffix(".ll").write_text(ll) + stem.with_suffix(".hip").write_text( + lower_kernel_to_hip(kernel, arch="gfx942") + ) + blob, _ = build_hsaco_from_llvm_ir(ll, isa="amdgcn-amd-amdhsa--gfx942") + stem.with_suffix(".hsaco").write_bytes(blob) + aa, bb = a, b + if variant == "prepacked": + aa, bb = rne_bits(a.view(np.uint32)), rne_bits(b.view(np.uint32)) + elif variant == "masked": + aa, bb = a.view(np.uint32) & np.uint32(0xFFFFE000), b.view( + np.uint32 + ) & np.uint32(0xFFFFE000) + d, pa, pb = launch(blob, spec.name, aa, bb, c) + expected_a = ( + rne_bits(a.view(np.uint32)) + if variant in ("rne", "prepacked") + else aa.view(np.uint32) + ) + expected_b = ( + rne_bits(b.view(np.uint32)) + if variant in ("rne", "prepacked") + else bb.view(np.uint32) + ) + assert np.array_equal(pa, expected_a), ( + engine, + m, + variant, + "A preparation", + ) + assert np.array_equal(pb, expected_b), ( + engine, + m, + variant, + "B preparation", + ) + np.savez(stem.with_suffix(".npz"), d=d, pa=pa, pb=pb) + results[variant] = d + # Dense reference uses float64 scalar products and a conservative + # FP32 accumulation bound, proportional to sum(abs(products))+abs(C). + ra, rb = expected_a, expected_b + if variant in ("raw", "carrier", "masked"): + ra, rb = ra & np.uint32(0xFFFFE000), rb & np.uint32(0xFFFFE000) + da, db = ra.view(np.float32)[single:].astype(np.float64), rb.view( + np.float32 + )[single:].astype(np.float64) + ref = np.einsum("bik,bjk->bij", da, db) + c[single:].astype(np.float64) + scale = np.einsum("bik,bjk->bij", abs(da), abs(db)) + abs(c[single:]) + error = abs(d[single:].astype(np.float64) - ref) + assert np.all( + error <= 4 * (128 // m) * np.finfo(np.float32).eps * scale + 1e-7 + ), (engine, m, variant, "dense reference", float(error.max())) + # Exact single-product normal results avoid reduction-order noise. + input_bits = np.concatenate((bits, bits)) + expected = ( + input_bits + if variant == "fp32" + else ( + rne_bits(input_bits) + if variant in ("rne", "prepacked") + else input_bits & np.uint32(0xFFFFE000) + ) + ) + # +0 accumulation turns a negative-zero product into +0. + expected = np.where( + (expected & 0x7FFFFFFF) == 0, np.uint32(0), expected + ).astype(np.uint32) + assert np.all( + same_values( + d[:single], + np.broadcast_to( + expected.view(np.float32)[:, None, None], d[:single].shape + ), + ) + ), (engine, m, variant, "single product") + summary["results"].append( + { + "engine": engine, + "m": m, + "variant": variant, + "batches": len(a), + "prepared_words": int(pa.size + pb.size), + "dense_max_abs_error": float(error.max()), + "status": "pass", + } + ) + for left, right in ( + ("raw", "carrier"), + ("raw", "masked"), + ("rne", "prepacked"), + ): + assert np.all(same_values(results[left], results[right])), ( + engine, + m, + left, + right, + ) + # Preserve all boundary observations, including denormal and NaN behavior. + observations = [ + { + "input_bits": f"{int(u):08x}", + "rne_bits": f"{int(rne_bits(u)):08x}", + **{ + v: f"{int(d[i,0,0].view(np.uint32)):08x}" + for v, d in results.items() + }, + } + for i, u in enumerate(bits) + ] + for row in observations: + original = np.array(int(row["input_bits"], 16), dtype=np.uint32).view( + np.float32 + ) + row["input_value"] = str(original) + for variant in ("raw", "rne"): + actual_bits = int(row[variant], 16) + actual = np.array(actual_bits, dtype=np.uint32).view(np.float32) + if np.isfinite(original) and np.isfinite(actual): + row[f"{variant}_abs_error_vs_fp32"] = abs( + float(actual) - float(original) + ) + # Within one sign, IEEE encoding order gives ULP distance. + if (actual_bits ^ int(row["input_bits"], 16)) < 0x80000000: + row[f"{variant}_ulp_error_vs_fp32"] = abs( + actual_bits - int(row["input_bits"], 16) + ) + (output_dir / f"{engine}_{m}_boundaries.json").write_text( + json.dumps(observations, indent=2) + "\n" + ) + counterexample = next( + row + for row in observations + if row["raw"] != row["rne"] and row["input_bits"] == "3f801001" + ) + print(engine, m, json.dumps(counterexample), flush=True) + all_results[engine] = results + if backend == "both": + for variant in all_results["python"]: + assert np.all( + same_values( + all_results["python"][variant], all_results["cpp"][variant] + ) + ) + summary["comgr"] = resolved_lib_path() + summary["status"] = "pass" + (output_dir / "results.json").write_text(json.dumps(summary, indent=2) + "\n") + return summary + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--backend", choices=("python", "cpp", "both"), default="both") + parser.add_argument("--shape", choices=("16", "32", "all"), default="all") + args = parser.parse_args() + args.output_dir.mkdir(parents=True, exist_ok=True) + result_path = args.output_dir / "results.json" + result_path.write_text(json.dumps({"status": "running"}) + "\n") + try: + result = run( + args.output_dir, + args.backend, + (16, 32) if args.shape == "all" else (int(args.shape),), + ) + except Exception as error: + result_path.write_text( + json.dumps({"status": "fail", "error": str(error)}) + "\n" + ) + raise + print(json.dumps({"status": result["status"], "variants": len(result["results"])})) + + +if __name__ == "__main__": + main() diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics/README.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics/README.md new file mode 100644 index 000000000000..8650aeec453c --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics/README.md @@ -0,0 +1,72 @@ +# Native TF32 numerics on gfx942 + +Run the [numerical example](../tf32_numerics.py) with NumPy and a built +`rocke_engine` on `PYTHONPATH`, using a ROCm runtime/compiler pair that supports +gfx942: + +```console +python -m rocke.examples.gfx942.tf32_numerics --backend both --shape all --output-dir results +``` + +Set `ROCKE_LLVM_FLAVOR` to the installed COMGR release (`llvm20` for ROCm <7.2, +`llvm22` for 7.2–7.12, `llvm23` for 7.13+). The default runs both engines and both +native shapes. Missing gfx942 hardware or the native extension is an error. +The [spec-driven probe](../../../instances/gfx942/tf32_mma_probe.py) is also +available through the native builder and the paired test emitters. + +The example loads dense A and transposed B tiles, accumulates into FP32 C, and +saves D plus the actual prepared A/B bits. Each shape runs raw FP32, the same bits +loaded through I32, explicit GPU RNE preparation, independent host RNE preparation, +host masking, and ordinary FP32 MFMA over the same total K. The C++ lane builds +its own IR and lowers it natively; it also checks exact Python/C++ byte identity. + +The corpus includes both signs, even/odd ties and their neighbors, exponent +carry, subnormals, signed zero, infinities, signaling/quiet NaNs, deterministic +dense data, cancellation, asymmetric matrices, and nonzero C. Single products +are checked exactly (NaNs by classification, zero after positive-zero +accumulation). Dense results use float64 products with an explicit FP32 +accumulation bound based on the sum of absolute products. No BLAS TF32 policy +can influence the reference. + +The output directory contains inputs, prepared words and outputs (`.npz`), LLVM, +HIP source, HSACO, full boundary observations, and `results.json`. The program +exits nonzero on compilation, launch, preparation, parity, or numerical failure. +Boundary JSON preserves hexadecimal results for every variant; the console +prints a representative RNE counterexample. This is a correctness example, not +a performance benchmark. + +## Input preparation and logical type + +`TF32` is a distinct logical IR type with an **I32 carrier**, one 32-bit word per +element. The dtype names `tf32` and `xf32` select the same catalog format. +`bitcast(fp32_value, TF32)` preserves the raw payload; it does not round. +`cvt_f32_to_tf32(fp32_value)` explicitly rounds to nearest, ties to even, setting +the low 13 fraction bits to zero. It preserves infinities and signed zero and +quiets NaNs while retaining the upper payload bits. A bitcast back to F32 +reinterprets the stored bits exactly; it does not truncate a raw payload. + +On gfx942 with ROCm 7.14, both native XF32 shapes accepted raw FP32 payloads. +For the tested corpus their outputs matched masking the low 13 bits; I32 +transport was bit-preserving. RNE preparation changed rounding-boundary results. +For a single product with B=1 and C=0: + +| Path | Input bits | Output bits | +|---|---|---| +| Ordinary FP32 MFMA | `3f801001` | `3f801001` | +| Raw XF32 / I32 carrier | `3f801001` | `3f800000` | +| GPU / host RNE → XF32 | `3f802000` | `3f802000` | + +Thus preparation is unnecessary for **instruction acceptance**, but explicit +preparation is necessary when the caller requires **RNE inputs**. Raw payloads +retain native XF32 behavior, including special encodings: in the tested corpus +a signaling NaN whose payload occupied only discarded bits became infinity. +The RNE conversion prevents that by quieting NaNs before clearing low bits. +These observations describe this corpus and toolchain; they are not a guarantee +that raw payloads have general FP32 arithmetic semantics. + +The catalog exposes only gfx942 `16x16x8` and `32x32x4` TF32 atoms, with two +logical TF32 elements per lane in A/B and FP32 accumulation. The LLVM intrinsic +ABI takes two floats, so lowering bitcasts the I32 carrier at the intrinsic +boundary. Ordinary FP32 catalog selection is unchanged. Generic TF32 arithmetic +is rejected; bitcasts, vector transport, selection, and the native atoms are +supported. Production GEMM dispatch and cross-target emulation are separate work. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py index ebcfc0e21a42..dcb5898824d8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py @@ -25,6 +25,8 @@ def storage_ir_type(dtype: str) -> Type: """Addressable storage unit, distinct from the logical dtype resolver.""" info = dtype_info(dtype) + if info.name == "tf32": + return I32 if info.encoded_bits % 8 or info.name in ("e8m0", "e5m3"): return I8 return dtype_to_ir_type(info.name) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/gfx942/tf32_mma_probe.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/gfx942/tf32_mma_probe.py new file mode 100644 index 000000000000..3bce6406cc0c --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/instances/gfx942/tf32_mma_probe.py @@ -0,0 +1,120 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""One wave per dense atom, with explicit FP32 input-preparation variants.""" + +from dataclasses import dataclass + +from rocke.core.arch import ArchTarget +from rocke.core.ir import F32, I32, TF32, IRBuilder, PtrType, VectorType +from rocke.core.storage import BitPacking, FragmentPacking, MatrixFragmentLayout +from rocke.helpers.mma_io import load_matrix_fragment + +PREPARATIONS = ("raw", "carrier", "rne", "prepacked", "fp32") + + +@dataclass(frozen=True) +class Tf32MmaProbeSpec: + m: int = 16 + preparation: str = "raw" + + def __post_init__(self): + if self.m not in (16, 32) or self.preparation not in PREPARATIONS: + raise ValueError("TF32 probe requires m=16/32 and a known preparation") + + @property + def name(self): + return f"tf32_probe_{self.m}_{self.preparation}" + + +def build_tf32_mma_probe(spec: Tf32MmaProbeSpec): + m, mode = spec.m, spec.preparation + k = 128 // m + target = ArchTarget.from_gfx("gfx942") + atom = target.mma.op_for_shape( + family="mma", + a_dtype="tf32", + b_dtype="tf32", + c_dtype="fp32", + m=m, + n=m, + k=k, + ) + b = IRBuilder(spec.name) + ty = I32 if mode == "prepacked" else F32 + a = b.param("A", PtrType(ty, "global"), align=4) + bb = b.param("B", PtrType(ty, "global"), align=4) + cc = b.param("C", PtrType(F32, "global"), align=4) + d = b.param("D", PtrType(F32, "global"), align=4) + pa = b.param("PA", PtrType(I32, "global"), align=4) + pb = b.param("PB", PtrType(I32, "global"), align=4) + lane = b.thread_id_x() + batch = b.block_id_x() + cm = b.const_i32(m) + axis = b.mod(lane, cm) + group = b.div(lane, cm) + ck = b.const_i32(k) + row = b.mul(axis, ck) + batch_input = b.mul(batch, b.const_i32(m * k)) + row_base = b.add(batch_input, row) + batch_output = b.mul(batch, b.const_i32(m * m)) + c_values, out_indices = [], [] + for slot in range(atom.c_frag_len): + r, col = atom.c_layout().coord(b, lane, slot) + index = b.add(batch_output, b.add(b.mul(r, cm), col)) + out_indices.append(index) + c_values.append(b.global_load(cc, index, F32, align=4)) + acc = b.vec_pack(c_values, F32) + if mode == "fp32": + full = target.mma.op_for_shape( + family="mma", + a_dtype="fp32", + b_dtype="fp32", + c_dtype="fp32", + m=m, + n=m, + k=k // 2, + ) + for step in range(2): + index = b.add(row_base, b.add(group, b.const_i32(step * (64 // m)))) + av = b.global_load(a, index, F32, align=4) + bv = b.global_load(bb, index, F32, align=4) + b.global_store(pa, index, b.bitcast(av, I32), align=4) + b.global_store(pb, index, b.bitcast(bv, I32), align=4) + acc = b.mma(full, av, bv, acc) + else: + layout = MatrixFragmentLayout( + FragmentPacking(BitPacking(32), 2, 32, 2), 2, 64 // m, m + ) + fragments = [] + for ptr, prepared in ((a, pa), (bb, pb)): + carrier = I32 if mode in ("carrier", "prepacked") else F32 + values = load_matrix_fragment( + b, + ptr, + row_base, + group, + 0, + dtype="tf32" if mode == "prepacked" else "fp32", + layout=layout, + carrier_type=carrier, + alignment_bytes=4, + ) + if mode == "rne": + values = b.vec_pack( + [b.cvt_f32_to_tf32(b.vec_extract(values, i)) for i in range(2)], + TF32, + ) + else: + values = b.bitcast(values, VectorType(TF32, 2)) + for slot in range(2): + index = b.add( + row_base, b.add(b.mul(group, b.const_i32(2)), b.const_i32(slot)) + ) + bits = b.bitcast(b.vec_extract(values, slot), I32) + b.global_store(prepared, index, bits, align=4) + fragments.append(values) + acc = b.mma(atom, *fragments, acc) + for slot, index in enumerate(out_indices): + b.global_store(d, index, b.vec_extract(acc, slot), align=4) + b.ret() + return b.kernel diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt b/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt index 4704a66de77a..a4da38d3557b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt @@ -14,6 +14,15 @@ if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/core/smoke.cpp") endif() if(BUILD_TESTING) + add_executable(rocke_tf32 core/test_tf32.cpp) + target_link_libraries(rocke_tf32 PRIVATE rocke_core) + add_test(NAME rocke_tf32 COMMAND rocke_tf32) + set_tests_properties(rocke_tf32 PROPERTIES LABELS "ckc;unit;host") + rocke_cxx_test_installed_name(_rocke_installed_name rocke_tf32) + set_target_properties(rocke_tf32 PROPERTIES OUTPUT_NAME "${_rocke_installed_name}") + install(TARGETS rocke_tf32 RUNTIME DESTINATION ${ROCKE_CXX_TEST_INSTALL_DIR}) + set_property(GLOBAL APPEND PROPERTY ROCKE_INSTALL_TEST_TARGETS rocke_tf32) + add_executable(rocke_storage core/test_storage.cpp) target_link_libraries(rocke_storage PRIVATE rocke_core) add_test(NAME rocke_storage COMMAND rocke_storage) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.cpp new file mode 100644 index 000000000000..744ed05993aa --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.cpp @@ -0,0 +1,73 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +#include "rocke/error.hpp" +#include "rocke/instance_tf32_mma_probe.h" +#include "rocke/ir_serialize.h" +#include "rocke/lower_hip.h" +#include "rocke/lower_llvm.h" +#include "rocke/verify.h" +#include +#include +#include +#include + +#define CHECK(expr) \ + do \ + { \ + if(!(expr)) \ + { \ + std::fprintf(stderr, "line %d: %s\n", __LINE__, #expr); \ + return 1; \ + } \ + } while(0) + +int main() +{ + CHECK(!rocke_type_eq(rocke_tf32(), rocke_i32())); + CHECK(!rocke_type_eq(rocke_tf32(), rocke_f32())); + CHECK(rocke_scalar_by_name("tf32") == rocke_tf32()); + const char* modes[] = {"raw", "carrier", "rne", "prepacked", "fp32"}; + for(int m : {16, 32}) + { + for(const char* mode : modes) + { + rocke_ir_builder_t b; + auto* kernel = rocke_build_tf32_mma_probe(&b, m, mode); + CHECK(kernel); + rocke_diag_t* diagnostics = nullptr; + size_t count = 0; + rocke_verify(kernel, &diagnostics, &count); + CHECK(count == 0); + rocke_diags_free(diagnostics, count); + rocke_lower_hip_opts_t opts = {}; + opts.arch = "gfx942"; + rocke_strbuf_t hip; + CHECK(rocke_strbuf_init(&hip, 0) == 0); + CHECK(rocke_lower_kernel_to_hip(&b, kernel, &opts, &hip) == ROCKE_OK); + CHECK(std::strstr(rocke_strbuf_cstr(&hip), "__builtin_amdgcn_mfma_f32_")); + rocke_strbuf_free(&hip); + if(std::strcmp(mode, "fp32") != 0) + { + char* ll = nullptr; + CHECK(rocke_lower_kernel_to_llvm(kernel, ROCKE_LLVM_FLAVOR_AUTO, "gfx950", &ll) + != ROCKE_OK); + std::free(ll); + } + rocke_ir_builder_free(&b); + } + } + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "invalid") == ROCKE_OK); + auto* integer = rocke_b_const_i32(&b, 1); + try + { + rocke_b_cvt_f32_to_tf32(&b, integer); + CHECK(false); + } + catch(const ckc::Error& error) + { + CHECK(error.code() == ROCKE_ERR_VALUE); + } + rocke_ir_builder_free(&b); + return 0; +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py new file mode 100644 index 000000000000..753f65d546e0 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py @@ -0,0 +1,247 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Logical TF32 contracts and strict native parity, independent of GPU access.""" +import numpy as np +import pytest + +from rocke.core.arch import ArchTarget +from rocke.core.dtypes import dtype_info, normalize_dtype +from rocke.core.ir import ( + F32, + I32, + TF32, + IRBuilder, + PtrType, + VectorType, + dtype_to_ir_type, +) +from rocke.core.ir_serialize import serialize, parse +from rocke.core.lower_hip import lower_kernel_to_hip +from rocke.core.lower_llvm import _lower_kernel_to_llvm_python +from rocke.examples.gfx942.tf32_numerics import rne_bits +from rocke.instances.gfx942.tf32_mma_probe import ( + PREPARATIONS, + Tf32MmaProbeSpec, + build_tf32_mma_probe, +) + + +def test_dtype_identity(): + assert TF32 != F32 and TF32 != I32 + for alias in ("tf32", "xf32", " XF32 "): + assert normalize_dtype(alias) == "tf32" + assert dtype_to_ir_type(alias) == TF32 + assert dtype_info(alias).encoded_bits == 32 + + +@pytest.mark.parametrize("m,k,acc", [(16, 8, 4), (32, 4, 16)]) +def test_catalog_isolation(m, k, acc): + catalog = ArchTarget.from_gfx("gfx942").mma + atom = catalog.op_for_shape( + family="mma", a_dtype="xf32", b_dtype="tf32", c_dtype="fp32", m=m, n=m, k=k + ) + assert (atom.a_frag_len, atom.b_frag_len, atom.c_frag_len) == (2, 2, acc) + assert atom.op_id == f"mfma_f32_{m}x{m}x{k}_xf32" + full = catalog.op_for_shape( + family="mma", a_dtype="fp32", b_dtype="fp32", c_dtype="fp32", m=m, n=m, k=k // 2 + ) + assert full.op_id == f"mfma_f32_{m}x{m}x{k//2}_f32" + assert not catalog.has_shape( + family="mma", a_dtype="fp32", b_dtype="fp32", c_dtype="fp32", m=m, n=m, k=k + ) + for gfx in ("gfx950", "gfx1151", "gfx1250"): + assert not any(op.a_dtype == "tf32" for op in ArchTarget.from_gfx(gfx).mma.ops) + + class Integers: + const_i32 = staticmethod(int) + mod = staticmethod(lambda x, y: x % y) + div = staticmethod(lambda x, y: x // y) + mul = staticmethod(lambda x, y: x * y) + add = staticmethod(lambda x, y: x + y) + + calc = Integers() + for layout, shape in ( + (atom.a_layout(), (m, k)), + (atom.b_layout(), (k, m)), + (atom.c_layout(), (m, m)), + ): + count = acc if shape == (m, m) else 2 + coords = [ + layout.coord(calc, lane, slot) + for lane in range(64) + for slot in range(count) + ] + assert len(set(coords)) == shape[0] * shape[1] == len(coords) + assert set(coords) == {(i, j) for i in range(shape[0]) for j in range(shape[1])} + + +@pytest.mark.parametrize("m", [16, 32]) +@pytest.mark.parametrize("mode", PREPARATIONS) +def test_both_builders_and_lowerers(m, mode): + native = pytest.importorskip("rocke_engine") + kernel = build_tf32_mma_probe(Tf32MmaProbeSpec(m, mode)) + ir = serialize(kernel) + assert serialize(parse(ir)) == ir + assert native.tf32_mma_probe_serialize_ir(m, mode) == ir + for flavor in ("llvm20", "llvm22", "llvm23"): + ll = _lower_kernel_to_llvm_python(kernel, arch="gfx942", llvm_flavor=flavor) + assert native.lower_serialized_ir(ir, arch="gfx942", flavor=flavor) == ll + if mode != "fp32": + assert f"llvm.amdgcn.mfma.f32.{m}x{m}x{128//m}.xf32" in ll + assert "sitofp" not in ll and "uitofp" not in ll + if mode != "fp32": + assert "__builtin_bit_cast(f32x2," in lower_kernel_to_hip(kernel, arch="gfx942") + + +@pytest.mark.parametrize("dtype,n", [(F32, 2), (I32, 2), (TF32, 1), (TF32, 4)]) +def test_mma_requires_logical_tf32(dtype, n): + b = IRBuilder("bad") + a = b.param("a", VectorType(dtype, n)) + c = b.param("c", VectorType(F32, 4)) + with pytest.raises(ValueError, match="XF32 MMA requires"): + b.mma("mfma_f32_16x16x8_xf32", a, a, c) + + +@pytest.mark.parametrize("gfx", ["gfx950", "gfx1151", "gfx1250"]) +def test_rejects_unsupported_target(gfx): + kernel = build_tf32_mma_probe(Tf32MmaProbeSpec()) + with pytest.raises(ValueError, match="gfx942"): + _lower_kernel_to_llvm_python(kernel, arch=gfx) + with pytest.raises(ValueError, match="gfx942"): + lower_kernel_to_hip(kernel, arch=gfx) + native = pytest.importorskip("rocke_engine") + with pytest.raises(Exception): + native.lower_serialized_ir(serialize(kernel), arch=gfx) + + +def test_arithmetic_rejected_after_serialization(): + b = IRBuilder("bad_arithmetic") + v = b.param("x", TF32) + b.add(v, v) + b.ret() + with pytest.raises(ValueError, match="TF32 arithmetic"): + _lower_kernel_to_llvm_python(b.kernel) + native = pytest.importorskip("rocke_engine") + with pytest.raises(Exception): + native.lower_serialized_ir(serialize(b.kernel)) + + +def test_transport(): + b = IRBuilder("transport") + p = b.param("p", PtrType(TF32, "global")) + i = b.const_i32(0) + v = b.global_load(p, i, TF32, align=4) + pair = b.vec_pack([v, v], TF32) + v = b.vec_extract(pair, 1) + b.global_store(p, i, v, align=4) + b.ret() + ir = serialize(b.kernel) + assert serialize(parse(ir)) == ir + ll = _lower_kernel_to_llvm_python(b.kernel) + assert "load i32" in ll and "store i32" in ll + native = pytest.importorskip("rocke_engine") + assert native.lower_serialized_ir(ir) == ll + + +def test_rne_reference_edges(): + inputs = np.array( + [ + 0, + 0x80000000, + 0x3F801000, + 0x3F803000, + 0x3F801001, + 0xBF801001, + 0x7F7FFFFF, + 0xFF7FFFFF, + 0x7F800000, + 0xFF800000, + 0x7F800001, + 0xFF800001, + 0x1000, + 0x1001, + 0x7FFFFF, + ], + np.uint32, + ) + expected = np.array( + [ + 0, + 0x80000000, + 0x3F800000, + 0x3F804000, + 0x3F802000, + 0xBF802000, + 0x7F800000, + 0xFF800000, + 0x7F800000, + 0xFF800000, + 0x7FC00000, + 0xFFC00000, + 0, + 0x2000, + 0x800000, + ], + np.uint32, + ) + assert np.array_equal(rne_bits(inputs), expected) + assert np.array_equal(rne_bits(expected), expected) + + +@pytest.mark.parametrize("n", [1, 2, 4]) +def test_lds_and_loop_transport(n): + b = IRBuilder("tf32_lds_loop") + p = b.param("p", PtrType(TF32, "global")) + zero = b.const_i32(0) + one = b.const_i32(1) + v = b.global_load(p, zero, TF32, align=4) + pair = b.vec_pack([v] * n, TF32) + shared = b.smem_alloc(TF32, (64,)) + b.smem_store_vN(shared, [zero], v if n == 1 else pair, n) + loaded = b.smem_load_vN(shared, zero, dtype=TF32, n=n) + loop = b.scf_for_iter(zero, one, one, [("payload", loaded)]) + with loop as (_, [payload]): + b.scf_yield(payload) + b.global_store(p, zero, b.vec_extract(loop.results[0], 0), align=4) + b.ret() + ir = serialize(b.kernel) + assert serialize(parse(ir)) == ir + ll = _lower_kernel_to_llvm_python(b.kernel) + assert "[256 x i8]" in ll + assert f"phi <{n} x i32>" in ll + assert "int smem" in lower_kernel_to_hip( + b.kernel + ) or "int shared" in lower_kernel_to_hip(b.kernel) + native = pytest.importorskip("rocke_engine") + assert native.lower_serialized_ir(ir) == ll + + +def test_fp32_atom_rejects_tf32_carrier(): + b = IRBuilder("wrong_atom") + a = b.param("a", TF32) + c = b.param("c", VectorType(F32, 4)) + with pytest.raises(ValueError, match="TF32 operands require"): + b.mma("mfma_f32_16x16x4_f32", a, a, c) + + +def test_parsed_mma_rejects_wrong_result_and_operands(): + native = pytest.importorskip("rocke_engine") + for bad in ("vec", "vec", "vec"): + ir = serialize(build_tf32_mma_probe(Tf32MmaProbeSpec())).replace( + "vec", bad + ) + with pytest.raises(ValueError, match="XF32 MMA requires"): + _lower_kernel_to_llvm_python(parse(ir), arch="gfx942") + with pytest.raises(Exception): + native.lower_serialized_ir(ir, arch="gfx942") + + +@pytest.mark.parametrize("m", [16, 32]) +@pytest.mark.parametrize("mode", PREPARATIONS) +def test_recipe_recording_preserves_ir(m, mode): + from rocke.portable_ir.src.recording_builder import record_kernel + + spec = Tf32MmaProbeSpec(m, mode) + recorded, recipe = record_kernel(lambda: build_tf32_mma_probe(spec)) + assert serialize(recorded) == serialize(build_tf32_mma_probe(spec)) + assert recipe diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json index 251b898c7701..33bc058de419 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json @@ -786,6 +786,18 @@ "14": "9997bc5e87821bf254534c0b7e2c55b928c0d2420828e52ca09b46a149c0dc93", "15": "17efbd8d97270032cb45822ab296ee7b546323e0158f32bde502c9592b2a7259", "16": "eed1a46189fb8296576cd4e1f840c8f8c9c137f4aadf2193331bc74dd92d2d51" + }, + "tf32_mma": { + "0": "5c55f4ced229ea0f5e87fa2fd95d79ca163cab8e74a0e561e416f067ec7b5c12", + "1": "a04ceff3d0aa833d139f5552b6dcd3348e8ba57f81b1d6903bb19f4388024822", + "2": "27f2533b40d0a64fc63ecf2338c9c9a1d7aa8952386558986f88db12ce92c573", + "3": "a56f37a803b2506e7223c2f41cf10bad35cc67dd2f5b2f534f2bc7d74b9a7404", + "4": "fac881d8ca2ca0cac250de093201c31d271976a10a958aa66e6e4698f45ecd2b", + "5": "71ef4ae6c1e94e84d20801fbbcf4519c0a5578cacb7223d8230dab7fc349c976", + "6": "6eb563ec23871f49e4858d7bf7c1a97218093c70d2a758551754934fbb02a887", + "7": "aeb835b4a246d15f8e55c705a5f109d4422eb27bb4fd4320e6efecd7dbdc8526", + "8": "18d9204ef8738449e215fe36bdb58e68c055a9b947611f02dba05204b9aa2821", + "9": "913879a7f860a5c9c0eb61a967a16efc0732e957d674833b8b532457173ec834" } } } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/tf32_mma_emit.c b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/tf32_mma_emit.c new file mode 100644 index 000000000000..73308b4d8b43 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/tf32_mma_emit.c @@ -0,0 +1,74 @@ +/* Copyright (c) Advanced Micro Devices, Inc., or its affiliates. + * SPDX-License-Identifier: MIT */ +#include "rocke/instance_tf32_mma_probe.h" +#include "rocke/ir_serialize.h" +#include "rocke/lower_hip.h" +#include "rocke/lower_llvm.h" +#include "rocke/verify.h" +#include +#include +#include + +int main(int argc, char** argv) +{ + int idx = argc > 1 ? atoi(argv[1]) : -1; + if(idx < 0 || idx >= 10) + { + fprintf(stderr, "unknown config index %d\n", idx); + return 2; + } + const char* modes[] = {"raw", "carrier", "rne", "prepacked", "fp32"}; + const char* mode = argc > 2 ? argv[2] : "ll"; + rocke_ir_builder_t b; + rocke_kernel_def_t* kernel = rocke_build_tf32_mma_probe(&b, idx < 5 ? 16 : 32, modes[idx % 5]); + if(!kernel) + { + fprintf(stderr, "%s\n", b.err); + rocke_ir_builder_free(&b); + return 1; + } + char* text = NULL; + rocke_status_t status = ROCKE_OK; + if(strcmp(mode, "ll") == 0) + status = rocke_lower_kernel_to_llvm(kernel, ROCKE_LLVM_FLAVOR_AUTO, "gfx942", &text); + else if(strcmp(mode, "ir") == 0) + status = rocke_ir_serialize(kernel, &text); + else if(strcmp(mode, "hip") == 0) + { + rocke_strbuf_t buf; + rocke_strbuf_init(&buf, 0); + rocke_lower_hip_opts_t opts = {0}; + opts.arch = "gfx942"; + status = rocke_lower_kernel_to_hip(&b, kernel, &opts, &buf); + if(status == ROCKE_OK) + fputs(buf.data, stdout); + rocke_strbuf_free(&buf); + } + else if(strcmp(mode, "verify") == 0) + { + rocke_diag_t* diagnostics = NULL; + size_t count = 0; + rocke_verify(kernel, &diagnostics, &count); + for(size_t i = 0; i < count; ++i) + { + char* line = rocke_diag_to_string(&diagnostics[i]); + if(line) + { + puts(line); + free(line); + } + } + rocke_diags_free(diagnostics, count); + } + else + status = ROCKE_ERR_VALUE; + if(text) + { + fputs(text, stdout); + free(text); + } + if(status != ROCKE_OK) + fprintf(stderr, "emit failed: %d %s\n", status, b.err); + rocke_ir_builder_free(&b); + return status != ROCKE_OK; +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/tf32_mma_emit.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/tf32_mma_emit.py new file mode 100644 index 000000000000..f58bcda3e7a8 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/tf32_mma_emit.py @@ -0,0 +1,20 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +from rocke.instances.gfx942.tf32_mma_probe import ( + PREPARATIONS, + Tf32MmaProbeSpec, + build_tf32_mma_probe, +) +from _emit_common import run_emit + + +def _spec(idx): + if not 0 <= idx < 10: + raise SystemExit(f"unknown config index {idx}") + return Tf32MmaProbeSpec(16 if idx < 5 else 32, PREPARATIONS[idx % 5]) + + +if __name__ == "__main__": + raise SystemExit( + run_emit(_spec, lambda spec, arch: build_tf32_mma_probe(spec), arch="gfx942") + ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/numeric/test_gfx942_tf32_numeric.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/numeric/test_gfx942_tf32_numeric.py new file mode 100644 index 000000000000..ae7d2589485a --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/numeric/test_gfx942_tf32_numeric.py @@ -0,0 +1,17 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +import os +import pytest +from rocke.runtime.hip_module import get_device_arch +from rocke.examples.gfx942.tf32_numerics import run + + +def test_gfx942_tf32_numeric(tmp_path): + arch = get_device_arch() + if not arch or arch.split(":")[0] != "gfx942": + if os.getenv("ROCKE_REQUIRE_GFX942") == "1": + pytest.fail(f"Required gfx942 device absent: {arch}") + pytest.skip("requires gfx942") + result = run(tmp_path, backend="both") + assert result["status"] == "pass" + assert len(result["results"]) == 24 From 98bfae88e119b5e207ce3df1246c5a7e275dc76c Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Fri, 25 Sep 2026 16:23:54 -0400 Subject: [PATCH 19/24] fix(hip-kernel-provider): tighten TF32 API parity and validation --- .../rocke/platform/cpp/core/ir/ir_mem.cpp | 5 +- .../rocke/platform/python/rocke/core/ir.py | 4 +- .../rocke/examples/gfx942/tf32_numerics.py | 16 ++-- .../examples/gfx942/tf32_numerics/README.md | 8 +- .../rocke/platform/tests/core/test_tf32.cpp | 45 ++++++++++ .../rocke/platform/tests/core/test_tf32.py | 86 ++++++++++++++++++- 6 files changed, 148 insertions(+), 16 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp index b07e81cb6b16..f1ba0d6334e8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp @@ -339,7 +339,8 @@ void rocke_b_global_store_vN(rocke_ir_builder_t* b, return; } } - else if(rocke_i_type_is(et, "f32") || rocke_i_type_is(et, "i32")) + else if(rocke_i_type_is(et, "f32") || rocke_i_type_is(et, "i32") + || rocke_i_type_is(et, "tf32")) { elem_bytes = 4; if(n == 16) @@ -359,7 +360,7 @@ void rocke_b_global_store_vN(rocke_ir_builder_t* b, (void)rocke_i_set_err( b, ROCKE_ERR_VALUE, - "global_store_vN supports f16/bf16/i16/f32/i32/i8/fp8e4m3/bf8e5m2, got %s", + "global_store_vN supports f16/bf16/i16/f32/i32/tf32/i8/fp8e4m3/bf8e5m2, got %s", en); return; } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py index b5a436cff555..03976023959c 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 @@ -4230,7 +4230,7 @@ def global_store_vN( """Vectorised global store of N consecutive elements. Supports the full element-type catalog the LLVM lowering already - emits: ``f16`` / ``bf16`` / ``i16`` (2-byte), ``f32`` / ``i32`` + emits: ``f16`` / ``bf16`` / ``i16`` (2-byte), ``f32`` / ``i32`` / ``tf32`` (4-byte), ``i8`` / ``fp8e4m3`` / ``bf8e5m2`` (1-byte). Lowers to a single ``store `` and AMDGPU coalesces into one ``global_store_dwordxN`` transaction. @@ -4254,7 +4254,7 @@ def global_store_vN( elem_bytes = 1 else: raise ValueError( - "global_store_vN supports f16/bf16/i16/f32/i32/i8/fp8e4m3/bf8e5m2, " + "global_store_vN supports f16/bf16/i16/f32/i32/tf32/i8/fp8e4m3/bf8e5m2, " f"got {elem_name}" ) self._op( diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics.py index 83bb36dd3e49..977e71e265a8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics.py @@ -5,7 +5,6 @@ import argparse import ctypes import json -import os from pathlib import Path import struct @@ -13,7 +12,7 @@ from rocke.core.ir_serialize import serialize from rocke.core.lower_hip import lower_kernel_to_hip -from rocke.core.lower_llvm import _lower_kernel_to_llvm_python +from rocke.core.lower_llvm import _lower_kernel_to_llvm_python, _resolve_llvm_flavor from rocke.instances.gfx942.tf32_mma_probe import ( PREPARATIONS, Tf32MmaProbeSpec, @@ -143,10 +142,11 @@ def run(output_dir, backend="both", shapes=(16, 32)): native = None if "cpp" in engines: import rocke_engine as native # Required: no fallback. + flavor = _resolve_llvm_flavor() summary = { "arch": arch, "numpy": np.__version__, - "flavor": os.getenv("ROCKE_LLVM_FLAVOR"), + "flavor": flavor, "results": [], } for m in shapes: @@ -166,15 +166,19 @@ def run(output_dir, backend="both", shapes=(16, 32)): ll = native.lower_serialized_ir( ir_native, arch="gfx942", - flavor=os.getenv("ROCKE_LLVM_FLAVOR", "llvm23"), + flavor=flavor, ) - assert ll == _lower_kernel_to_llvm_python(kernel, arch="gfx942"), ( + assert ll == _lower_kernel_to_llvm_python( + kernel, arch="gfx942", llvm_flavor=flavor + ), ( m, mode, "lowerer parity", ) else: - ll = _lower_kernel_to_llvm_python(kernel, arch="gfx942") + ll = _lower_kernel_to_llvm_python( + kernel, arch="gfx942", llvm_flavor=flavor + ) stem = output_dir / f"{engine}_{m}_{variant}" stem.with_suffix(".ll").write_text(ll) stem.with_suffix(".hip").write_text( diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics/README.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics/README.md index 8650aeec453c..389971de7609 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics/README.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics/README.md @@ -8,9 +8,11 @@ gfx942: python -m rocke.examples.gfx942.tf32_numerics --backend both --shape all --output-dir results ``` -Set `ROCKE_LLVM_FLAVOR` to the installed COMGR release (`llvm20` for ROCm <7.2, -`llvm22` for 7.2–7.12, `llvm23` for 7.13+). The default runs both engines and both -native shapes. Missing gfx942 hardware or the native extension is an error. +The example resolves one LLVM flavor from the installed COMGR release for both +engines and records it in `results.json`. Override with `ROCKE_LLVM_FLAVOR` when +needed (`llvm20` for ROCm <7.2, `llvm22` for 7.2–7.12, `llvm23` for 7.13+). +The default runs both engines and both native shapes. Missing gfx942 hardware +or the native extension is an error. The [spec-driven probe](../../../instances/gfx942/tf32_mma_probe.py) is also available through the native builder and the paired test emitters. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.cpp index 744ed05993aa..83812784aff9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.cpp @@ -56,6 +56,51 @@ int main() rocke_ir_builder_free(&b); } } + // Mirror test_global_vector_store: native authoring must accept the same widths. + for(int n : {1, 2, 4, 8, 16}) + { + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "tf32_vector_store") == ROCKE_OK); + auto* p = rocke_b_param(&b, "p", rocke_ptr_type(&b, rocke_tf32(), "global"), nullptr); + auto* index = rocke_b_const_i32(&b, 0); + auto* value = rocke_b_bitcast(&b, index, rocke_tf32()); + rocke_value_t* components[16]; + for(int i = 0; i < n; ++i) + components[i] = value; + auto* values = rocke_b_vec_pack(&b, components, n, rocke_tf32()); + if(n == 16) + { + try + { + rocke_b_global_store_vN(&b, p, index, values, n, 0); + CHECK(false); + } + catch(const ckc::Error& error) + { + CHECK(error.code() == ROCKE_ERR_VALUE); + CHECK(std::strstr(error.what(), "n=16 not supported for tf32")); + } + } + else + { + rocke_b_global_store_vN(&b, p, index, values, n, 0); + rocke_b_ret(&b); + for(int f = 0; f < rocke_llvm_flavor_count(); ++f) + { + char* ll = nullptr; + auto flavor = rocke_llvm_flavor_from_name(rocke_llvm_flavor_at(f)); + CHECK(rocke_lower_kernel_to_llvm(rocke_ir_builder_kernel(&b), flavor, "gfx950", &ll) + == ROCKE_OK); + char store[64], align[32]; + std::snprintf(store, sizeof(store), "store <%d x i32>", n); + std::snprintf(align, sizeof(align), "align %d", n * 4); + CHECK(std::strstr(ll, store)); + CHECK(std::strstr(ll, align)); + std::free(ll); + } + } + rocke_ir_builder_free(&b); + } rocke_ir_builder_t b; CHECK(rocke_ir_builder_init(&b, "invalid") == ROCKE_OK); auto* integer = rocke_b_const_i32(&b, 1); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py index 753f65d546e0..dfea2d82832a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py @@ -224,7 +224,7 @@ def test_fp32_atom_rejects_tf32_carrier(): b.mma("mfma_f32_16x16x4_f32", a, a, c) -def test_parsed_mma_rejects_wrong_result_and_operands(): +def test_parsed_mma_rejects_wrong_operands(): native = pytest.importorskip("rocke_engine") for bad in ("vec", "vec", "vec"): ir = serialize(build_tf32_mma_probe(Tf32MmaProbeSpec())).replace( @@ -238,10 +238,90 @@ def test_parsed_mma_rejects_wrong_result_and_operands(): @pytest.mark.parametrize("m", [16, 32]) @pytest.mark.parametrize("mode", PREPARATIONS) -def test_recipe_recording_preserves_ir(m, mode): +def test_recipe_replay_preserves_ir(m, mode, monkeypatch): from rocke.portable_ir.src.recording_builder import record_kernel + from rocke.portable_ir.src import online, recipe_bundle spec = Tf32MmaProbeSpec(m, mode) recorded, recipe = record_kernel(lambda: build_tf32_mma_probe(spec)) assert serialize(recorded) == serialize(build_tf32_mma_probe(spec)) - assert recipe + for flavor in ("llvm20", "llvm22", "llvm23"): + monkeypatch.setenv("ROCKE_LLVM_FLAVOR", flavor) + replayed, _ = online.recipe_cbor_to_llvm( + recipe_bundle.cbor_encode(recipe), arch="gfx942" + ) + assert replayed == _lower_kernel_to_llvm_python( + recorded, arch="gfx942", llvm_flavor=flavor + ) + + +@pytest.mark.parametrize("n", [1, 2, 4, 8, 16]) +def test_global_vector_store(n): + b = IRBuilder("tf32_vector_store") + p = b.param("p", PtrType(TF32, "global")) + i = b.const_i32(0) + v = b.global_load(p, i, TF32, align=4) + values = b.vec_pack([v] * n, TF32) + if n == 16: + with pytest.raises(ValueError, match="n=16 not supported for tf32"): + b.global_store_vN(p, i, values, n) + return + b.global_store_vN(p, i, values, n) + b.ret() + ir = serialize(b.kernel) + assert serialize(parse(ir)) == ir + native = pytest.importorskip("rocke_engine") + for flavor in ("llvm20", "llvm22", "llvm23"): + ll = _lower_kernel_to_llvm_python(b.kernel, llvm_flavor=flavor) + assert f"store <{n} x i32>" in ll + assert f"align {n * 4}" in ll + assert native.lower_serialized_ir(ir, flavor=flavor) == ll + + +@pytest.mark.parametrize("m", [16, 32]) +@pytest.mark.parametrize("role", ["accumulator", "result"]) +@pytest.mark.parametrize("bad", ["integer", "tf32", "width"]) +def test_parsed_mma_rejects_wrong_accumulator_and_result(m, role, bad): + native = pytest.importorskip("rocke_engine") + kernel = build_tf32_mma_probe(Tf32MmaProbeSpec(m)) + mma = next(op for op in kernel.body.ops if op.name == "tile.mma") + value = mma.operands[2] if role == "accumulator" else mma.result + count = 4 if m == 16 else 16 + value.type = VectorType( + {"integer": I32, "tf32": TF32, "width": F32}[bad], + count + 1 if bad == "width" else count, + ) + ir = serialize(kernel) + message = "XF32 MMA requires" if role == "accumulator" else "XF32 MMA result" + with pytest.raises(ValueError, match=message): + _lower_kernel_to_llvm_python(parse(ir), arch="gfx942") + with pytest.raises(Exception, match=message): + native.lower_serialized_ir(ir, arch="gfx942") + + +@pytest.mark.parametrize("flavor", ["llvm20", "llvm22", "llvm23"]) +def test_numeric_harness_auto_flavor(flavor, monkeypatch, tmp_path): + pytest.importorskip("rocke_engine") + from rocke.examples.gfx942 import tf32_numerics + from rocke.runtime import comgr, hip_module + + monkeypatch.delenv("ROCKE_LLVM_FLAVOR", raising=False) + monkeypatch.setattr(tf32_numerics, "_resolve_llvm_flavor", lambda: flavor) + monkeypatch.setattr(hip_module, "get_device_arch", lambda: "gfx942") + expected = _lower_kernel_to_llvm_python( + build_tf32_mma_probe(Tf32MmaProbeSpec()), + arch="gfx942", + llvm_flavor=flavor, + ) + + class ReachedCompiler(Exception): + pass + + def check_compiler_input(ll, *, isa): + assert ll == expected + raise ReachedCompiler + + # Exercise the actual builders/lowerers, then stop before GPU compilation. + monkeypatch.setattr(comgr, "build_hsaco_from_llvm_ir", check_compiler_input) + with pytest.raises(ReachedCompiler): + tf32_numerics.run(tmp_path, backend="cpp", shapes=(16,)) From 25446dfc8e28a4fccc86ca97fbe24bdff05a90ff Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Tue, 29 Sep 2026 21:58:15 -0400 Subject: [PATCH 20/24] fix(hip-kernel-provider): tighten rocKE TF32 vector contracts --- .../rocke/platform/cpp/core/ir/ir_mem.cpp | 5 +- .../cpp/include/rocke/tf32_internal.h | 10 +- .../dsl_docs/reference/dtype_support.md | 78 +++++++---- .../rocke/platform/python/rocke/core/ir.py | 6 +- .../rocke/platform/python/rocke/core/tf32.py | 11 +- .../rocke/platform/tests/core/test_tf32.cpp | 124 ++++++++++++++++++ .../rocke/platform/tests/core/test_tf32.py | 111 ++++++++++++++++ 7 files changed, 310 insertions(+), 35 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp index 0b29c220109a..05608fbafc5d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_mem.cpp @@ -231,7 +231,8 @@ rocke_value_t* rocke_b_global_load_vN(rocke_ir_builder_t* b, return (rocke_value_t*)rocke_i_set_err( b, ROCKE_ERR_VALUE, "unsupported vector width for global_load_vN: %d", n); } - else if(rocke_i_type_is(dtype, "f32") || rocke_i_type_is(dtype, "i32")) + else if(rocke_i_type_is(dtype, "f32") || rocke_i_type_is(dtype, "i32") + || rocke_i_type_is(dtype, "tf32")) { elem_bytes = 4; if(n != 2 && n != 3 && n != 4 && n != 8) @@ -251,7 +252,7 @@ rocke_value_t* rocke_b_global_load_vN(rocke_ir_builder_t* b, return (rocke_value_t*)rocke_i_set_err( b, ROCKE_ERR_VALUE, - "global_load_vN supports f16/bf16/i16/f32/i32/fp8e4m3/bf8e5m2/i8, got %s", + "global_load_vN supports f16/bf16/i16/f32/i32/tf32/fp8e4m3/bf8e5m2/i8, got %s", en); } vt = rocke_vector_type(b, dtype, n); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/tf32_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/tf32_internal.h index 158f66df4e07..4d98d6b97f35 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/tf32_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/tf32_internal.h @@ -49,8 +49,14 @@ static inline const char* rocke_tf32_op_error(const rocke_op_t* op) && (op->num_results != 1 || strcmp(op->results[0]->type->name, count == 4 ? "vec" : "vec") != 0)) return "XF32 MMA result must match its FP32 accumulator"; - if((strncmp(op->name, "arith.", 6) == 0 || strncmp(op->name, "math.", 5) == 0) - && strcmp(op->name, "arith.bitcast") != 0 && strcmp(op->name, "arith.select") != 0) + /* Only payload transport, reinterpretation, and selection preserve TF32. */ + if((strncmp(op->name, "arith.", 6) == 0 || strncmp(op->name, "math.", 5) == 0 + || strncmp(op->name, "vector.", 7) == 0) + && strcmp(op->name, "arith.bitcast") != 0 && strcmp(op->name, "arith.select") != 0 + && strcmp(op->name, "vector.bitcast") != 0 && strcmp(op->name, "vector.extract") != 0 + && strcmp(op->name, "vector.pack") != 0 && strcmp(op->name, "vector.concat") != 0 + && strcmp(op->name, "vector.insert") != 0 && strcmp(op->name, "vector.splat") != 0 + && strcmp(op->name, "vector.select") != 0) { for(int i = 0; i < op->num_operands + op->num_results; ++i) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/dtype_support.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/dtype_support.md index 8e8b4921351a..6fbc855e458b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/dtype_support.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/dtype_support.md @@ -23,12 +23,18 @@ estimates, shipping order, and acceptance checklists are maintained separately. | **Out of scope** | Runtime/dispatch dtype validation, the provider's C-api surface, and any measurement. | | **Keeping it current** | When a dtype's support changes at any layer, update the affected row(s) and the §2.6 anchors in the same change. The criteria (§2.2) and rubric (§2.3) are the stable part; the scores are not. **Line anchors and the §4B atom tables are hand-maintained mirrors of the basis tree and go stale on rebase** — an arch entry gaining or losing atoms upstream silently invalidates §4B without touching this file. Re-check both against the new base before merging any rebase. | +**TF32 coverage update.** The TF32 entries in §1, §4B, and §4D reflect the +logical dtype and native gfx942 atoms added after the original basis snapshot. +Other scores and source line numbers retain their stated historical basis. + --- ## 1. The canonical dtype set -rocKE's compiler knows exactly **10 scalar IR types**, defined identically in -both engines: +The original basis snapshot listed ten scalar IR types. The table below also +includes the subsequently added logical TF32 type. For the newer logical +low-bit and scale types, see [dtype storage and packing](../architecture/dtype_storage_packing.md). +The definitions in both engines are: - Python: `platform/python/rocke/core/ir.py:43-52` - C++: `platform/cpp/include/rocke/ir.h:93-108`, @@ -44,6 +50,7 @@ both engines: | `f16` | `half` | 16 | float | | `bf16` | `bfloat` | 16 | float | | `f32` | `float` | 32 | float | +| `tf32` | `i32` (carrier) | 32 | logical float; native gfx942 XF32 input | | `fp8e4m3` | `i8` (storage) | 8 | float (OCP E4M3) | | `bf8e5m2` | `i8` (storage) | 8 | float (OCP E5M2) | @@ -52,9 +59,10 @@ LLVM mapping: `lower_llvm.py:1093-1129` (Python, `_llvm_type`) ↔ storage and materialized through `llvm.amdgcn.cvt.*` intrinsics — there is no native `<8 x fp8>` LLVM type in the emitter. -**Sub-byte / block formats are NOT scalar IR types.** `i4`, `fp4`, `fp6`, and -the `e8m0` MX scale exist only as *packed encodings* consumed by dedicated -helpers/atoms, never as first-class `Type` objects (see §4). +**Sub-byte / block formats in the original snapshot.** The basis tree treated +`i4`, `fp4`, `fp6`, and the `e8m0` MX scale as packed encodings consumed by +dedicated helpers/atoms. Those historical scores do not describe the newer +logical-type and storage APIs linked above. > CK-Tile gives these formats named logical types backed by a packed storage > representation: `pk_fp4_t = pk_float4_e2m1_t` @@ -67,15 +75,24 @@ helpers/atoms, never as first-class `Type` objects (see §4). > rocKE's dedicated packed-format paths can provide support without a first-class > IR `Type`; C5 separately measures whether kernel families accept the format. -**Reduced-precision compute (the "tf32 request") is a mode, not a storage type.** -It never appears as a tensor dtype: storage stays `f32` and the *precision* is -reduced before the MMA. What is missing is the **mode itself**: no spec knob selects -reduced-precision compute, so a caller's TF32 request cannot be expressed at all. - -Native XF32 is available on gfx942, a supported rocKE target, but absent on -gfx950 and gfx1250 (§4D). The catalog has no XF32 atom, including on gfx942. -Native instruction absence does not rule out software implementations of TF32 -semantics. +**TF32 is a logical dtype with a 32-bit I32 carrier.** `tf32` and its `xf32` +alias select two native gfx942 XF32 atoms independently of ordinary FP32 +selection. The logical type survives serialization; only A/B are bitcast to +the intrinsic's float-vector ABI, while C/D remain FP32. + +`bitcast(value, TF32)` wraps an FP32 or I32 payload without rounding. +`cvt_f32_to_tf32(value)` explicitly prepares FP32 inputs with round-to-nearest, +ties-to-even; bitcasting back to F32 reinterprets the prepared bits. Generic +TF32 scalar and vector arithmetic is rejected. Bitcasts, vector transport, +selection, and the native atoms are supported. Direct global vector loads +accept 2, 3, 4, or 8 TF32 elements, with the same alignment rules as I32. +The fragment loader uses `storage_ir_type("tf32") == I32` and returns carriers +that the caller wraps as logical TF32. + +Native XF32 is available on gfx942, but absent on gfx950 and gfx1250 (§4D). +Production GEMM/convolution selectors and cross-target TF32 emulation remain +outside this support. The [numerical probe](../../python/rocke/examples/gfx942/tf32_numerics/README.md) +covers raw payloads, explicit RNE preparation, and ordinary FP32 controls. A **BF16 compute mode over `f32` storage** could convert operands to `bf16` and use existing bf16 MFMA/WMMA atoms. Its precision differs from TF32: the reference @@ -85,14 +102,17 @@ of an `f32`, `GpuRefTypes.h:92-103`, gated by `USE_TF32`), while `bf16` keeps `1.0` in bf16. Widening validation tolerances does not make these contracts equivalent; `USE_TF32` output is not a bit-accuracy oracle for BF16 compute. -The matrix below therefore labels this possible mode `bf16-mode`. No API selector -or TF32-request mapping currently exists. +The matrix below therefore labels this possible mode `bf16-mode`. It is separate +from logical TF32; no production selector maps TF32 requests to BF16 compute. **Types that are entirely absent from the vocabulary:** `f64`, unsigned `u8/u16/u32/u64`, and a distinct `bool` (predicates reuse `i1`). No engine references them. -Both engines define the same set of **10 scalar IR types**. +TF32 source anchors: [Python type and builders](../../python/rocke/core/ir.py), +[native type declarations](../../cpp/include/rocke/ir.h), +[Python contracts](../../python/rocke/core/tf32.py), and +[native contracts](../../cpp/include/rocke/tf32_internal.h). --- @@ -162,6 +182,8 @@ planning questions and deliberately live outside this document (see Scores 0–3 per criterion; `—` = N/A (excluded from the denominator). Rows are grouped by type family, not ranked — see "What the scores are not" in §2.3. +TF32 was added after this scored snapshot; its current layer boundaries are +described in §1 and §4D rather than assigned a retrospective percentage here. Gap categories: **[C]** conversion gap · **[P]** partial/plumbing gap · **[M]** missing entirely. @@ -406,13 +428,13 @@ to gfx1251, so this is not a missing gfx1250 atom. `core/arch/data/arch_specs.json` is rocKE's architecture SSOT. An atom that is not declared there is unusable by any kernel regardless of what the silicon provides. -Verbatim at the basis commit: +The original basis snapshot, with the gfx942 XF32 addition noted below: | Arch | Declared matrix atoms | dtypes reachable | |---|---|---| | `gfx90a` | `mfma_f32_{16x16x16,32x32x8}_{f16,bf16}` | f16, bf16 | -| `gfx942` | the gfx90a set + `mfma_f32_{16x16x4,32x32x2}_f32` + `mfma_f32_{16x16x32,32x32x16}_{fp8,bf8}` | f32, f16, bf16, fp8, bf8 | -| `gfx950` | the gfx942 set + K-packed `{16x16x32,32x32x16}_{f16,bf16}` + `mfma_f32_16x16x128_fp4` + `mfma_f32_16x16x96_fp6` | f32, f16, bf16, fp8, bf8, **fp4, fp6** | +| `gfx942` | the gfx90a set + `mfma_f32_{16x16x4,32x32x2}_f32` + `mfma_f32_{16x16x32,32x32x16}_{fp8,bf8}` + `mfma_f32_{16x16x8,32x32x4}_xf32` | f32, f16, bf16, fp8, bf8, **tf32** | +| `gfx950` | the gfx942 set **excluding XF32** + K-packed `{16x16x32,32x32x16}_{f16,bf16}` + `mfma_f32_16x16x128_fp4` + `mfma_f32_16x16x96_fp6` | f32, f16, bf16, fp8, bf8, **fp4, fp6** | | `gfx1151` | `wmma_f32_16x16x16_{f16,bf16}`, `wmma_i32_16x16x16_{iu8,iu4}` | f16, bf16, **iu8, iu4** | | `gfx1201` | `wmma_gfx12_f32_16x16x16_{f16,bf16}` | f16, bf16 | | `gfx1250` | `wmma_gfx1250_f32_16x16x4_f32`, `..._16x16x32_{f16,bf16}`, `..._16x16x64_{fp8_fp8,fp8_bf8,bf8_fp8,bf8_bf8}`, `wmma_scale_f32_16x16x128_fp8_fp8`, `wmma_scale16_f32_16x16x128_fp8_fp8` | f32, f16, bf16, fp8, bf8 (+ fp8 block-scaled) | @@ -428,8 +450,8 @@ Verbatim at the basis commit: ### 4C. Deltas — silicon capability vs. rocKE's declared atoms The following gaps compare §4A with §4B. The first five concern atom declarations -only; f64 (4C.6) also lacks an IR type. The gfx942 XF32 declaration gap is covered -separately in §4D. +only; f64 (4C.6) also lacks an IR type. The gfx942 XF32 declaration gap is now +closed; §4D describes its logical dtype and remaining integration boundaries. #### 4C.1 gfx1250 — fp4/fp6/bf6 block-scaled atoms are omitted (deliberately) The part carries `V_WMMA_F32_16X16X128_F8F6F4`, a dedicated @@ -478,11 +500,13 @@ gfx1250 has no native F64 matrix instruction and is excluded from this gap. | gfx1151 / gfx1201 (RDNA3.5 / RDNA4) | ❌ never present | | gfx1250 (CDNA5) | ❌ absent from the public XML and LLVM target features | -**Consequence for rocKE.** gfx942 is a supported target: it has a catalog entry -and block-scale GEMM accepts it for FP8/BF8. It provides native XF32, but rocKE -does not declare XF32 atoms. gfx950/gfx1250 have no native XF32 instruction. -See [§1](#1-the-canonical-dtype-set) for the distinction between native XF32, -software TF32 implementations, and a BF16 compute mode. +**Consequence for rocKE.** The gfx942 catalog declares +`mfma_f32_16x16x8_xf32` and `mfma_f32_32x32x4_xf32`, with logical TF32 A/B and +FP32 C/D. Both Python and C++ lowerers support them and reject these atoms on +other targets. gfx950/gfx1250 have no native XF32 instruction. This atom support +does not add production GEMM/convolution dispatch or cross-target emulation. +See [§1](#1-the-canonical-dtype-set) for the distinction between logical TF32, +raw payload wrapping, explicit RNE preparation, and a BF16 compute mode. ### 4E. The FP8 encoding dialect (FNUZ vs OCP) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py index 8b1b2312b5ea..2802a0460be2 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 @@ -1617,7 +1617,7 @@ def global_load_vN( ) -> Value: """Vectorised global load of N consecutive values. - Supports f16/bf16/i16 (N in {2, 4, 6, 8, 16}), f32/i32 + Supports f16/bf16/i16 (N in {2, 4, 6, 8, 16}), f32/i32/tf32 (N in {2, 3, 4, 8}), and fp8e4m3/bf8e5m2/i8 (N in {2, 4, 8, 12, 16}). Loads exactly N elements. Instruction selection depends on target and alignment; 96-bit payloads do not require a 96-bit scalar type. @@ -1632,7 +1632,7 @@ def global_load_vN( # <16 x half> operand fragment; AMDGPU coalesces it when aligned. if n not in (2, 4, 6, 8, 16): raise ValueError(f"unsupported vector width for global_load_vN: {n}") - elif dtype.name in ("f32", "i32"): + elif dtype.name in ("f32", "i32", "tf32"): elem_bytes = 4 if n not in (2, 3, 4, 8): raise ValueError( @@ -1646,7 +1646,7 @@ def global_load_vN( ) else: raise ValueError( - "global_load_vN supports f16/bf16/i16/f32/i32/fp8e4m3/bf8e5m2/i8, " + "global_load_vN supports f16/bf16/i16/f32/i32/tf32/fp8e4m3/bf8e5m2/i8, " f"got {dtype.name}" ) return self._op( diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/tf32.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/tf32.py index b836d53b0b57..ffb375fd05c6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/tf32.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/tf32.py @@ -31,9 +31,18 @@ def tf32_op_error(op): ) if error: return error - if op.name.startswith(("arith.", "math.")) and op.name not in ( + # Vector arithmetic must not operate on the integer carrier bits either. + # Keep the operations that only move, reinterpret, or select payloads. + if op.name.startswith(("arith.", "math.", "vector.")) and op.name not in ( "arith.bitcast", "arith.select", + "vector.bitcast", + "vector.extract", + "vector.pack", + "vector.concat", + "vector.insert", + "vector.splat", + "vector.select", ): if any( v.type.name == "tf32" or v.type.name.startswith("vec", n); + std::snprintf(align, sizeof(align), "align %d", expected_align); + CHECK(std::strstr(ll, load)); + CHECK(std::strstr(ll, align)); + std::free(ll); + } + rocke_strbuf_t hip; + CHECK(rocke_strbuf_init(&hip, 0) == 0); + rocke_lower_hip_opts_t opts = {}; + opts.arch = "gfx942"; + CHECK(rocke_lower_kernel_to_hip(&b, b.kernel, &opts, &hip) == ROCKE_OK); + if(expected_align < n * 4 || n == 3) + { + char size[32]; + std::snprintf(size, sizeof(size), "%d);", n * 4); + CHECK(std::strstr(hip.data, "__builtin_memcpy")); + CHECK(std::strstr(hip.data, size)); + } + rocke_strbuf_free(&hip); + rocke_ir_builder_free(&b); + } + return 0; +} + int main() { + CHECK(test_vector_arithmetic_rejected() == 0); + CHECK(test_vector_load() == 0); CHECK(!rocke_type_eq(rocke_tf32(), rocke_i32())); CHECK(!rocke_type_eq(rocke_tf32(), rocke_f32())); CHECK(rocke_scalar_by_name("tf32") == rocke_tf32()); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py index dfea2d82832a..f4e5a8d26f0c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py @@ -8,6 +8,7 @@ from rocke.core.dtypes import dtype_info, normalize_dtype from rocke.core.ir import ( F32, + I1, I32, TF32, IRBuilder, @@ -18,6 +19,7 @@ from rocke.core.ir_serialize import serialize, parse from rocke.core.lower_hip import lower_kernel_to_hip from rocke.core.lower_llvm import _lower_kernel_to_llvm_python +from rocke.core.verify import verify from rocke.examples.gfx942.tf32_numerics import rne_bits from rocke.instances.gfx942.tf32_mma_probe import ( PREPARATIONS, @@ -126,6 +128,115 @@ def test_arithmetic_rejected_after_serialization(): native.lower_serialized_ir(serialize(b.kernel)) +@pytest.mark.parametrize( + "operation", + [ + "add", + "sub", + "mul", + "and", + "or", + "shl", + "lshr", + "smax", + "smin", + "max", + "fma", + "sum", + "reduce_max", + "cmp", + "trunc", + "sext", + "trunc_result", + ], +) +def test_vector_arithmetic_rejected_after_serialization(operation): + b = IRBuilder("invalid_tf32_vector_arithmetic") + elem = I32 if operation == "trunc_result" else TF32 + v = b.param("v", VectorType(elem, 2)) + if operation == "fma": + b.vector_fma(v, v, v) + elif operation in ("sum", "reduce_max"): + getattr(b, f"vector_{operation}")(v) + elif operation == "cmp": + b.vector_cmp("lt", v, v) + elif operation in ("trunc", "sext", "trunc_result"): + if operation == "trunc_result": + b.vector_trunc(v, TF32) + else: + getattr(b, f"vector_{operation}")(v, I32) + else: + getattr(b, f"vector_{operation}")(v, v) + b.ret() + ir = serialize(b.kernel) + kernel = parse(ir) + message = "TF32 arithmetic requires an explicit conversion to f32" + assert any(message in str(d) for d in verify(kernel)) + with pytest.raises(ValueError, match=message): + _lower_kernel_to_llvm_python(kernel) + with pytest.raises(ValueError, match=message): + lower_kernel_to_hip(kernel) + native = pytest.importorskip("rocke_engine") + with pytest.raises(Exception, match=message): + native.lower_serialized_ir(ir) + + +def test_vector_payload_operations_remain_supported(): + b = IRBuilder("tf32_vector_payload") + p = b.param("p", PtrType(TF32, "global")) + mask = b.param("mask", I1) + zero = b.const_i32(0) + v = b.global_load(p, zero, TF32, align=4) + pair = b.vec_pack([v, v], TF32) + pair = b.vec_insert(pair, b.vec_extract(pair, 0), 1) + pair = b.vec_bitcast(b.vec_bitcast(pair, VectorType(I32, 2)), VectorType(TF32, 2)) + pair = b.vector_select(mask, pair, b.vector_splat(v, 2)) + pair = b.select(mask, pair, pair) + b.global_store_vN(p, zero, b.vec_concat(pair, pair), 4) + b.ret() + ir = serialize(b.kernel) + kernel = parse(ir) + assert verify(kernel) == [] + assert "i32x4" in lower_kernel_to_hip(kernel) + native = pytest.importorskip("rocke_engine") + for flavor in ("llvm20", "llvm22", "llvm23"): + ll = _lower_kernel_to_llvm_python(kernel, llvm_flavor=flavor) + assert native.lower_serialized_ir(ir, flavor=flavor) == ll + + +@pytest.mark.parametrize("n", [2, 3, 4, 8]) +@pytest.mark.parametrize("align", [None, 4]) +def test_global_vector_load(n, align): + b = IRBuilder("tf32_vector_load") + p = b.param("p", PtrType(TF32, "global")) + zero = b.const_i32(0) + values = b.global_load_vN(p, zero, TF32, n, align=align) + b.global_store(p, zero, b.vec_extract(values, n - 1), align=4) + b.ret() + ir = serialize(b.kernel) + kernel = parse(ir) + assert verify(kernel) == [] + expected_align = align or (4 if n == 3 else n * 4) + hip = lower_kernel_to_hip(kernel) + if expected_align < n * 4 or n == 3: + assert f"{n * 4});" in hip + assert "__builtin_memcpy" in hip + native = pytest.importorskip("rocke_engine") + for flavor in ("llvm20", "llvm22", "llvm23"): + ll = _lower_kernel_to_llvm_python(kernel, llvm_flavor=flavor) + assert f"load <{n} x i32>" in ll + assert f"align {expected_align}" in ll + assert native.lower_serialized_ir(ir, flavor=flavor) == ll + + +@pytest.mark.parametrize("n", [0, 1, 6, 16]) +def test_global_vector_load_rejects_unsupported_width(n): + b = IRBuilder("invalid_tf32_vector_load") + p = b.param("p", PtrType(TF32, "global")) + with pytest.raises(ValueError, match="unsupported vector width for tf32"): + b.global_load_vN(p, b.const_i32(0), TF32, n) + + def test_transport(): b = IRBuilder("transport") p = b.param("p", PtrType(TF32, "global")) From 75508e81e332c8e0953f3a241b8058fc635e493d Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Tue, 29 Sep 2026 23:40:03 -0400 Subject: [PATCH 21/24] fix(hip-kernel-provider): preserve TF32 store alignment and MMA result contracts --- .../cpp/core/lower_hip/lower_hip_mem.cpp | 16 +++ .../platform/cpp/core/lower_llvm/mem.cpp | 2 +- .../cpp/include/rocke/tf32_internal.h | 12 +- .../dsl_docs/reference/dtype_support.md | 5 +- .../rocke/platform/python/rocke/core/ir.py | 5 +- .../platform/python/rocke/core/lower_hip.py | 13 ++ .../platform/python/rocke/core/lower_llvm.py | 2 +- .../rocke/platform/python/rocke/core/tf32.py | 5 + .../rocke/platform/tests/core/test_tf32.cpp | 117 +++++++++++------- .../rocke/platform/tests/core/test_tf32.py | 82 ++++++++++-- 10 files changed, 196 insertions(+), 63 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp index 60129168ffcd..ced6af4ad0a9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_mem.cpp @@ -641,6 +641,22 @@ static rocke_status_t _op_memref_global_store_vN(rocke_h_lowerer_t* lw, const ro n = mem_attr_int(op, "vec", 0); elem_name = mem_attr_str(op, "elem_type", "f16"); prefix = rocke_h_vec_prefix_checked(lw, elem_name, /*full_map=*/true, "global_store_vN"); + const int64_t byte_count = n * (rocke_dtype_info(elem_name)->encoded_bits / 8); + const int64_t align = mem_attr_int(op, "align", byte_count); + if(align <= 0 || (align & (align - 1))) + return rocke_h_fail( + lw, ROCKE_ERR_VALUE, "global_store_vN: alignment must be a positive power of two"); + if(align < byte_count || (byte_count & (byte_count - 1))) + { + rocke_h_emitf(lw, + "__builtin_memcpy(__builtin_assume_aligned(%s + %s, %lld), &%s, %lld);", + rocke_h_name(lw, ptr), + rocke_h_name(lw, idx), + (long long)align, + rocke_h_name(lw, val), + (long long)byte_count); + return lw->status; + } rocke_h_emitf(lw, "*reinterpret_cast<%s%lld*>(%s + %s) = %s;", prefix, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp index d3d53a814422..e4ad9ce9e627 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp @@ -338,7 +338,7 @@ static void op_memref_global_store_vN(rocke_lower_t* L, const rocke_op_t* op) : rocke_ll_llvm_type(L, val->type); const char* elem_name = ll_is_vec(val->type) ? val->type->elem->name : val->type->name; int elem_bytes = ll_elem_bytes(elem_name); - int64_t align = vec * elem_bytes; + int64_t align = ll_attr_int(op, "align", vec * elem_bytes); const char* ty = rocke_ll_llvm_type(L, val->type); rocke_ll_emitf(L, " %s = getelementptr inbounds %s, ptr addrspace(1) %s, i32 %s", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/tf32_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/tf32_internal.h index 4d98d6b97f35..bbbc8ca68721 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/tf32_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/tf32_internal.h @@ -38,13 +38,17 @@ static inline const char* rocke_tf32_op_error(const rocke_op_t* op) const char* id = rocke_attr_get_str(&op->attrs, "op_id"); if(!id && strncmp(op->name, "tile.", 5) == 0) id = op->name + 5; - const char* error = (strcmp(op->name, "tile.mma") == 0 || strncmp(op->name, "tile.mfma", 9) == 0 - || strncmp(op->name, "tile.wmma", 9) == 0) - ? rocke_tf32_mma_error(id, op->operands, op->num_operands) - : NULL; + const int is_mma = strcmp(op->name, "tile.mma") == 0 || strncmp(op->name, "tile.mfma", 9) == 0 + || strncmp(op->name, "tile.wmma", 9) == 0; + const char* error = is_mma ? rocke_tf32_mma_error(id, op->operands, op->num_operands) : NULL; if(error) return error; int count = rocke_tf32_mma_count(id); + if(is_mma && !count) + for(int i = 0; i < op->num_results; ++i) + if(strcmp(op->results[i]->type->name, "tf32") == 0 + || strncmp(op->results[i]->type->name, "vecnum_results != 1 || strcmp(op->results[0]->type->name, count == 4 ? "vec" : "vec") != 0)) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/dtype_support.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/dtype_support.md index 6fbc855e458b..079a77a13115 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/dtype_support.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/dtype_support.md @@ -85,7 +85,10 @@ the intrinsic's float-vector ABI, while C/D remain FP32. ties-to-even; bitcasting back to F32 reinterprets the prepared bits. Generic TF32 scalar and vector arithmetic is rejected. Bitcasts, vector transport, selection, and the native atoms are supported. Direct global vector loads -accept 2, 3, 4, or 8 TF32 elements, with the same alignment rules as I32. +accept 2, 3, 4, or 8 TF32 elements; global vector stores accept 1, 2, 4, or 8. +Both use the same alignment rules as I32. Payload width and address alignment +are independent: underaligned HIP stores copy exactly the payload bytes, and +target/alignment determine the machine transfer instructions. The fragment loader uses `storage_ir_type("tf32") == I32` and returns carriers that the caller wraps as logical TF32. 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 2802a0460be2..d08e3a6b54e8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py @@ -4240,8 +4240,9 @@ def global_store_vN( Supports the full element-type catalog the LLVM lowering already emits: ``f16`` / ``bf16`` / ``i16`` (2-byte), ``f32`` / ``i32`` / ``tf32`` (4-byte), ``i8`` / ``fp8e4m3`` / ``bf8e5m2`` (1-byte). Lowers to - a single ``store `` and AMDGPU coalesces into one - ``global_store_dwordxN`` transaction. + a single ``store `` with the supplied address alignment. + Payload width and address alignment are independent; target and + alignment determine whether the transfer uses one machine instruction. """ if n not in (1, 2, 4, 8, 16): raise ValueError(f"global_store_vN n must be 1, 2, 4, 8, or 16 (got {n})") 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 c572abc0c198..32249f619793 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py @@ -1837,6 +1837,19 @@ def _op_memref_global_store_vN(self, op: Op) -> None: n = int(op.attrs["vec"]) elem_name = op.attrs.get("elem_type", "f16") prefix = _vec_prefix(elem_name, "global_store_vN") + byte_count = n * (dtype_info(elem_name).encoded_bits // 8) + align = int(op.attrs.get("align", byte_count)) + if align <= 0 or align & (align - 1): + raise ValueError( + "global_store_vN: alignment must be a positive power of two" + ) + if align < byte_count or byte_count & (byte_count - 1): + self._emit( + f"__builtin_memcpy(" + f"__builtin_assume_aligned({_name(ptr)} + {_name(idx)}, {align}), " + f"&{_name(val)}, {byte_count});" + ) + return self._emit( f"*reinterpret_cast<{prefix}{n}*>({_name(ptr)} + {_name(idx)}) = " f"{_name(val)};" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py index 975a87e3b4d8..d625699dbaaf 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 @@ -5518,7 +5518,7 @@ def _op_memref_global_store_vN(self, op: Op) -> None: "f32": 4, "i64": 8, }.get(elem_name, 2) - align = vec * elem_bytes + align = int(op.attrs.get("align", vec * elem_bytes)) ty = _llvm_type(val.type) self._current().emit( f" store {ty} {self._operand(val)}, ptr addrspace(1) {gep}, align {align}" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/tf32.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/tf32.py index ffb375fd05c6..6aea0c65f4ee 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/tf32.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/tf32.py @@ -13,6 +13,11 @@ def tf32_mma_error(op_id, operands, results=None): for v in operands ): return "TF32 operands require an XF32 MMA" + if results is not None and any( + v.type.name == "tf32" or v.type.name.startswith("vec", "vec", f"vec"] if [v.type.name for v in operands] != expected: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.cpp index 26538d6c1f9e..b3c049a6ec6f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.cpp @@ -21,9 +21,9 @@ } \ } while(0) -static int test_vector_arithmetic_rejected() +static int test_invalid_tf32_ops_rejected() { - for(int variant = 0; variant < 6; ++variant) + for(int variant = 0; variant < 8; ++variant) { rocke_ir_builder_t b, parsed; CHECK(rocke_ir_builder_init(&b, "invalid_tf32_vector") == ROCKE_OK); @@ -51,7 +51,22 @@ static int test_vector_arithmetic_rejected() rocke_b_vector_trunc(&b, v, rocke_tf32()); break; } - rocke_b_ret(&b); + if(variant >= 6) + { + rocke_ir_builder_free(&b); + auto* probe = rocke_build_tf32_mma_probe(&b, variant == 6 ? 16 : 32, "fp32"); + CHECK(probe); + for(int i = 0; i < probe->body->num_ops; ++i) + { + auto* op = probe->body->ops[i]; + if(std::strcmp(op->name, "tile.mma") == 0) + op->results[0]->type + = rocke_vector_type(&b, rocke_tf32(), variant == 6 ? 4 : 16); + } + } + else + rocke_b_ret(&b); + const char* message = variant >= 6 ? "MMA results must not use TF32" : "TF32 arithmetic"; char* text = nullptr; CHECK(rocke_ir_serialize(b.kernel, &text) == ROCKE_OK); rocke_kernel_def_t* kernel = nullptr; @@ -62,7 +77,7 @@ static int test_vector_arithmetic_rejected() CHECK(rocke_verify(kernel, &diagnostics, &count) == ROCKE_OK); bool rejected = false; for(size_t i = 0; i < count; ++i) - rejected |= std::strstr(diagnostics[i].message, "TF32 arithmetic") != nullptr; + rejected |= std::strstr(diagnostics[i].message, message) != nullptr; CHECK(rejected); rocke_diags_free(diagnostics, count); for(int f = 0; f < rocke_llvm_flavor_count(); ++f) @@ -145,7 +160,7 @@ static int test_vector_load() int main() { - CHECK(test_vector_arithmetic_rejected() == 0); + CHECK(test_invalid_tf32_ops_rejected() == 0); CHECK(test_vector_load() == 0); CHECK(!rocke_type_eq(rocke_tf32(), rocke_i32())); CHECK(!rocke_type_eq(rocke_tf32(), rocke_f32())); @@ -182,49 +197,67 @@ int main() } // Mirror test_global_vector_store: native authoring must accept the same widths. for(int n : {1, 2, 4, 8, 16}) - { - rocke_ir_builder_t b; - CHECK(rocke_ir_builder_init(&b, "tf32_vector_store") == ROCKE_OK); - auto* p = rocke_b_param(&b, "p", rocke_ptr_type(&b, rocke_tf32(), "global"), nullptr); - auto* index = rocke_b_const_i32(&b, 0); - auto* value = rocke_b_bitcast(&b, index, rocke_tf32()); - rocke_value_t* components[16]; - for(int i = 0; i < n; ++i) - components[i] = value; - auto* values = rocke_b_vec_pack(&b, components, n, rocke_tf32()); - if(n == 16) + for(int alignment : {0, 4, 64}) { - try - { - rocke_b_global_store_vN(&b, p, index, values, n, 0); - CHECK(false); - } - catch(const ckc::Error& error) + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "tf32_vector_store") == ROCKE_OK); + auto* p = rocke_b_param(&b, "p", rocke_ptr_type(&b, rocke_tf32(), "global"), nullptr); + auto* index = rocke_b_const_i32(&b, 0); + auto* value = rocke_b_bitcast(&b, index, rocke_tf32()); + rocke_value_t* components[16]; + for(int i = 0; i < n; ++i) + components[i] = value; + auto* values = rocke_b_vec_pack(&b, components, n, rocke_tf32()); + if(n == 16) { - CHECK(error.code() == ROCKE_ERR_VALUE); - CHECK(std::strstr(error.what(), "n=16 not supported for tf32")); + try + { + rocke_b_global_store_vN(&b, p, index, values, n, alignment); + CHECK(false); + } + catch(const ckc::Error& error) + { + CHECK(error.code() == ROCKE_ERR_VALUE); + CHECK(std::strstr(error.what(), "n=16 not supported for tf32")); + } } - } - else - { - rocke_b_global_store_vN(&b, p, index, values, n, 0); - rocke_b_ret(&b); - for(int f = 0; f < rocke_llvm_flavor_count(); ++f) + else { - char* ll = nullptr; - auto flavor = rocke_llvm_flavor_from_name(rocke_llvm_flavor_at(f)); - CHECK(rocke_lower_kernel_to_llvm(rocke_ir_builder_kernel(&b), flavor, "gfx950", &ll) - == ROCKE_OK); - char store[64], align[32]; - std::snprintf(store, sizeof(store), "store <%d x i32>", n); - std::snprintf(align, sizeof(align), "align %d", n * 4); - CHECK(std::strstr(ll, store)); - CHECK(std::strstr(ll, align)); - std::free(ll); + rocke_b_global_store_vN(&b, p, index, values, n, alignment); + rocke_b_ret(&b); + for(int f = 0; f < rocke_llvm_flavor_count(); ++f) + { + char* ll = nullptr; + auto flavor = rocke_llvm_flavor_from_name(rocke_llvm_flavor_at(f)); + CHECK(rocke_lower_kernel_to_llvm( + rocke_ir_builder_kernel(&b), flavor, "gfx950", &ll) + == ROCKE_OK); + char store[64], align[32]; + std::snprintf(store, sizeof(store), "store <%d x i32>", n); + std::snprintf(align, sizeof(align), "align %d", alignment ? alignment : n * 4); + CHECK(std::strstr(ll, store)); + CHECK(std::strstr(ll, align)); + std::free(ll); + } + rocke_strbuf_t hip; + CHECK(rocke_strbuf_init(&hip, 0) == 0); + rocke_lower_hip_opts_t opts = {}; + opts.arch = "gfx942"; + CHECK(rocke_lower_kernel_to_hip(&b, b.kernel, &opts, &hip) == ROCKE_OK); + if(alignment && alignment < n * 4) + { + char size[32]; + std::snprintf(size, sizeof(size), ", %d);", n * 4); + CHECK(std::strstr(hip.data, "__builtin_memcpy(__builtin_assume_aligned(")); + CHECK(std::strstr(hip.data, size)); + CHECK(!std::strstr(hip.data, "*reinterpret_cast<")); + } + else + CHECK(std::strstr(hip.data, "*reinterpret_cast<")); + rocke_strbuf_free(&hip); } + rocke_ir_builder_free(&b); } - rocke_ir_builder_free(&b); - } rocke_ir_builder_t b; CHECK(rocke_ir_builder_init(&b, "invalid") == ROCKE_OK); auto* integer = rocke_b_const_i32(&b, 1); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py index f4e5a8d26f0c..5dc4d02ab94c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py @@ -367,28 +367,86 @@ def test_recipe_replay_preserves_ir(m, mode, monkeypatch): @pytest.mark.parametrize("n", [1, 2, 4, 8, 16]) -def test_global_vector_store(n): - b = IRBuilder("tf32_vector_store") - p = b.param("p", PtrType(TF32, "global")) +@pytest.mark.parametrize("align", [None, 4, 64, "legacy"]) +@pytest.mark.parametrize("dtype", [TF32, I32, F32]) +def test_global_vector_store(n, align, dtype): + b = IRBuilder("vector_store_alignment") + p = b.param("p", PtrType(dtype, "global")) i = b.const_i32(0) - v = b.global_load(p, i, TF32, align=4) - values = b.vec_pack([v] * n, TF32) + v = b.global_load(p, i, dtype, align=4) + values = b.vec_pack([v] * n, dtype) if n == 16: - with pytest.raises(ValueError, match="n=16 not supported for tf32"): - b.global_store_vN(p, i, values, n) + with pytest.raises(ValueError, match="n=16 not supported"): + b.global_store_vN(p, i, values, n, align=align) return - b.global_store_vN(p, i, values, n) + # Legacy serialized stores without an align attribute retain their default. + legacy = align == "legacy" + align = None if legacy else align + # The guarantee applies to the accessed address, including the index. + index = b.const_i32(1 if align == 4 else 0) + b.global_store_vN(p, index, values, n, align=align) + if legacy: + b.kernel.body.ops[-1].attrs.pop("align") b.ret() ir = serialize(b.kernel) - assert serialize(parse(ir)) == ir + kernel = parse(ir) + assert serialize(kernel) == ir + assert verify(kernel) == [] + expected_align = align or n * 4 + hip = lower_kernel_to_hip(kernel) + if expected_align < n * 4: + assert ( + f"__builtin_assume_aligned(p + {index.name.lstrip('%')}, {expected_align})" + in hip + ) + assert f"&{values.name.lstrip('%')}, {n * 4});" in hip + else: + assert "*reinterpret_cast<" in hip native = pytest.importorskip("rocke_engine") for flavor in ("llvm20", "llvm22", "llvm23"): - ll = _lower_kernel_to_llvm_python(b.kernel, llvm_flavor=flavor) - assert f"store <{n} x i32>" in ll - assert f"align {n * 4}" in ll + ll = _lower_kernel_to_llvm_python(kernel, llvm_flavor=flavor) + store = next(line for line in ll.splitlines() if "store <" in line) + assert f"store <{n} x {'float' if dtype == F32 else 'i32'}>" in store + assert store.endswith(f"align {expected_align}") assert native.lower_serialized_ir(ir, flavor=flavor) == ll +@pytest.mark.parametrize("m", [16, 32]) +@pytest.mark.parametrize("scalar_result", [False, True]) +def test_non_xf32_mma_rejects_tf32_result(m, scalar_result): + b = IRBuilder("non_xf32_tf32_result") + out = b.param("out", PtrType(TF32, "global")) + a = b.param("a", F32) + c = b.param("c", VectorType(F32, 4 if m == 16 else 16)) + atom = ArchTarget.from_gfx("gfx942").mma.op_for_shape( + family="mma", + a_dtype="fp32", + b_dtype="fp32", + c_dtype="fp32", + m=m, + n=m, + k=64 // m, + ) + result = b.mma(atom, a, a, c) + # Model malformed serialized IR while keeping all downstream types consistent. + result.type = TF32 if scalar_result else VectorType(TF32, 4 if m == 16 else 16) + value = result if scalar_result else b.vec_extract(result, 0) + b.global_store(out, b.const_i32(0), value, align=4) + b.ret() + ir = serialize(b.kernel) + kernel = parse(ir) + message = "MMA results must not use TF32" + assert any(message in str(d) for d in verify(kernel)) + with pytest.raises(ValueError, match=message): + lower_kernel_to_hip(kernel, arch="gfx942") + native = pytest.importorskip("rocke_engine") + for flavor in ("llvm20", "llvm22", "llvm23"): + with pytest.raises(ValueError, match=message): + _lower_kernel_to_llvm_python(kernel, arch="gfx942", llvm_flavor=flavor) + with pytest.raises(Exception, match=message): + native.lower_serialized_ir(ir, arch="gfx942", flavor=flavor) + + @pytest.mark.parametrize("m", [16, 32]) @pytest.mark.parametrize("role", ["accumulator", "result"]) @pytest.mark.parametrize("bad", ["integer", "tf32", "width"]) From a09dd3378f2b6b3fe72b71cea9ae9a86e2df9623 Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Wed, 30 Sep 2026 02:48:56 -0400 Subject: [PATCH 22/24] fix(hip-kernel-provider): reject invalid vector store alignment --- .../platform/cpp/core/lower_llvm/mem.cpp | 3 ++ .../platform/python/rocke/core/lower_llvm.py | 4 ++ .../rocke/platform/tests/core/test_tf32.cpp | 41 +++++++++++++++++++ .../rocke/platform/tests/core/test_tf32.py | 23 +++++++++++ 4 files changed, 71 insertions(+) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp index e4ad9ce9e627..ab59c3061643 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/mem.cpp @@ -339,6 +339,9 @@ static void op_memref_global_store_vN(rocke_lower_t* L, const rocke_op_t* op) const char* elem_name = ll_is_vec(val->type) ? val->type->elem->name : val->type->name; int elem_bytes = ll_elem_bytes(elem_name); int64_t align = ll_attr_int(op, "align", vec * elem_bytes); + if(align <= 0 || (align & (align - 1))) + rocke_ll_fail( + L, ROCKE_ERR_VALUE, "global_store_vN: alignment must be a positive power of two"); const char* ty = rocke_ll_llvm_type(L, val->type); rocke_ll_emitf(L, " %s = getelementptr inbounds %s, ptr addrspace(1) %s, i32 %s", 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 d625699dbaaf..ea3b0b47c9ee 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py @@ -5519,6 +5519,10 @@ def _op_memref_global_store_vN(self, op: Op) -> None: "i64": 8, }.get(elem_name, 2) align = int(op.attrs.get("align", vec * elem_bytes)) + if align <= 0 or align & (align - 1): + raise ValueError( + "global_store_vN: alignment must be a positive power of two" + ) ty = _llvm_type(val.type) self._current().emit( f" store {ty} {self._operand(val)}, ptr addrspace(1) {gep}, align {align}" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.cpp index b3c049a6ec6f..33b64dd7b735 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.cpp @@ -158,8 +158,49 @@ static int test_vector_load() return 0; } +static int test_invalid_store_alignment() +{ + for(int align : {0, -4, 3, 12}) + { + rocke_ir_builder_t b, parsed; + CHECK(rocke_ir_builder_init(&b, "invalid_store_alignment") == ROCKE_OK); + CHECK(rocke_ir_builder_init(&parsed, "parsed") == ROCKE_OK); + auto* p = rocke_b_param(&b, "p", rocke_ptr_type(&b, rocke_tf32(), "global"), nullptr); + auto* zero = rocke_b_const_i32(&b, 0); + auto* value = rocke_b_bitcast(&b, zero, rocke_tf32()); + rocke_value_t* lanes[] = {value, value, value, value}; + auto* values = rocke_b_vec_pack(&b, lanes, 4, rocke_tf32()); + rocke_b_global_store_vN(&b, p, zero, values, 4, align); + auto* body = b.kernel->body; + rocke_attr_set_int(&b, &body->ops[body->num_ops - 1]->attrs, "align", align); + rocke_b_ret(&b); + char* text = nullptr; + CHECK(rocke_ir_serialize(b.kernel, &text) == ROCKE_OK); + rocke_kernel_def_t* kernel = nullptr; + CHECK(rocke_ir_parse(text, &parsed, &kernel) == ROCKE_OK); + std::free(text); + for(int f = 0; f < rocke_llvm_flavor_count(); ++f) + { + char* ll = nullptr; + auto flavor = rocke_llvm_flavor_from_name(rocke_llvm_flavor_at(f)); + CHECK(rocke_lower_kernel_to_llvm(kernel, flavor, "gfx942", &ll) == ROCKE_ERR_VALUE); + std::free(ll); + } + rocke_strbuf_t hip; + CHECK(rocke_strbuf_init(&hip, 0) == 0); + rocke_lower_hip_opts_t opts = {}; + opts.arch = "gfx942"; + CHECK(rocke_lower_kernel_to_hip(&parsed, kernel, &opts, &hip) == ROCKE_ERR_VALUE); + rocke_strbuf_free(&hip); + rocke_ir_builder_free(&parsed); + rocke_ir_builder_free(&b); + } + return 0; +} + int main() { + CHECK(test_invalid_store_alignment() == 0); CHECK(test_invalid_tf32_ops_rejected() == 0); CHECK(test_vector_load() == 0); CHECK(!rocke_type_eq(rocke_tf32(), rocke_i32())); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py index 5dc4d02ab94c..bf583e6f9d2c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py @@ -411,6 +411,29 @@ def test_global_vector_store(n, align, dtype): assert native.lower_serialized_ir(ir, flavor=flavor) == ll +@pytest.mark.parametrize("align", [0, -4, 3, 12]) +def test_global_vector_store_rejects_invalid_alignment(align): + b = IRBuilder("invalid_store_alignment") + p = b.param("p", PtrType(TF32, "global")) + zero = b.const_i32(0) + value = b.bitcast(zero, TF32) + b.global_store_vN(p, zero, b.vec_pack([value] * 4, TF32), 4, align=align) + # The builder treats zero as its default; serialized IR can contain it literally. + b.kernel.body.ops[-1].attrs["align"] = align + b.ret() + ir = serialize(b.kernel) + kernel = parse(ir) + message = "global_store_vN: alignment must be a positive power of two" + with pytest.raises(ValueError, match=message): + lower_kernel_to_hip(kernel) + native = pytest.importorskip("rocke_engine") + for flavor in ("llvm20", "llvm22", "llvm23"): + with pytest.raises(ValueError, match=message): + _lower_kernel_to_llvm_python(kernel, llvm_flavor=flavor) + with pytest.raises(Exception, match=message): + native.lower_serialized_ir(ir, flavor=flavor) + + @pytest.mark.parametrize("m", [16, 32]) @pytest.mark.parametrize("scalar_result", [False, True]) def test_non_xf32_mma_rejects_tf32_result(m, scalar_result): From 1c75ffdf3c0172f94ed651aa15d0437ca717420d Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Wed, 30 Sep 2026 13:26:01 -0400 Subject: [PATCH 23/24] fix(rocke): keep TF32 numerical validation active under optimization --- .../rocke/examples/gfx942/tf32_numerics.py | 64 ++++----- .../core/test_tf32_numerics_validation.py | 132 ++++++++++++++++++ 2 files changed, 164 insertions(+), 32 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32_numerics_validation.py diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics.py index 977e71e265a8..ae21ad10248a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/examples/gfx942/tf32_numerics.py @@ -5,8 +5,8 @@ import argparse import ctypes import json -from pathlib import Path import struct +from pathlib import Path import numpy as np @@ -162,19 +162,21 @@ def run(output_dir, backend="both", shapes=(16, 32)): ir = serialize(kernel) if engine == "cpp": ir_native = native.tf32_mma_probe_serialize_ir(m, mode) - assert ir == ir_native, (m, mode, "builder parity") + if ir != ir_native: + raise RuntimeError( + f"m={m} mode={mode}: builder parity mismatch" + ) ll = native.lower_serialized_ir( ir_native, arch="gfx942", flavor=flavor, ) - assert ll == _lower_kernel_to_llvm_python( + if ll != _lower_kernel_to_llvm_python( kernel, arch="gfx942", llvm_flavor=flavor - ), ( - m, - mode, - "lowerer parity", - ) + ): + raise RuntimeError( + f"m={m} mode={mode}: lowerer parity mismatch" + ) else: ll = _lower_kernel_to_llvm_python( kernel, arch="gfx942", llvm_flavor=flavor @@ -204,18 +206,11 @@ def run(output_dir, backend="both", shapes=(16, 32)): if variant in ("rne", "prepacked") else bb.view(np.uint32) ) - assert np.array_equal(pa, expected_a), ( - engine, - m, - variant, - "A preparation", - ) - assert np.array_equal(pb, expected_b), ( - engine, - m, - variant, - "B preparation", - ) + context = f"{engine} m={m} variant={variant}" + if not np.array_equal(pa, expected_a): + raise RuntimeError(f"{context}: A preparation mismatch") + if not np.array_equal(pb, expected_b): + raise RuntimeError(f"{context}: B preparation mismatch") np.savez(stem.with_suffix(".npz"), d=d, pa=pa, pb=pb) results[variant] = d # Dense reference uses float64 scalar products and a conservative @@ -229,9 +224,12 @@ def run(output_dir, backend="both", shapes=(16, 32)): ref = np.einsum("bik,bjk->bij", da, db) + c[single:].astype(np.float64) scale = np.einsum("bik,bjk->bij", abs(da), abs(db)) + abs(c[single:]) error = abs(d[single:].astype(np.float64) - ref) - assert np.all( + if not np.all( error <= 4 * (128 // m) * np.finfo(np.float32).eps * scale + 1e-7 - ), (engine, m, variant, "dense reference", float(error.max())) + ): + raise RuntimeError( + f"{context}: dense reference mismatch (max error {float(error.max())})" + ) # Exact single-product normal results avoid reduction-order noise. input_bits = np.concatenate((bits, bits)) expected = ( @@ -247,14 +245,15 @@ def run(output_dir, backend="both", shapes=(16, 32)): expected = np.where( (expected & 0x7FFFFFFF) == 0, np.uint32(0), expected ).astype(np.uint32) - assert np.all( + if not np.all( same_values( d[:single], np.broadcast_to( expected.view(np.float32)[:, None, None], d[:single].shape ), ) - ), (engine, m, variant, "single product") + ): + raise RuntimeError(f"{context}: single product mismatch") summary["results"].append( { "engine": engine, @@ -271,12 +270,10 @@ def run(output_dir, backend="both", shapes=(16, 32)): ("raw", "masked"), ("rne", "prepacked"), ): - assert np.all(same_values(results[left], results[right])), ( - engine, - m, - left, - right, - ) + if not np.all(same_values(results[left], results[right])): + raise RuntimeError( + f"{engine} m={m}: variant parity mismatch ({left}, {right})" + ) # Preserve all boundary observations, including denormal and NaN behavior. observations = [ { @@ -318,11 +315,14 @@ def run(output_dir, backend="both", shapes=(16, 32)): all_results[engine] = results if backend == "both": for variant in all_results["python"]: - assert np.all( + if not np.all( same_values( all_results["python"][variant], all_results["cpp"][variant] ) - ) + ): + raise RuntimeError( + f"m={m} variant={variant}: backend parity mismatch (python, cpp)" + ) summary["comgr"] = resolved_lib_path() summary["status"] = "pass" (output_dir / "results.json").write_text(json.dumps(summary, indent=2) + "\n") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32_numerics_validation.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32_numerics_validation.py new file mode 100644 index 000000000000..aac9de10985b --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32_numerics_validation.py @@ -0,0 +1,132 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Host-only CLI checks with injected outputs; no GPU correctness evidence.""" + +import json +import os +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import numpy as np +import pytest +from rocke.examples.gfx942 import tf32_numerics as numerics +from rocke.runtime import comgr, hip_module + + +def _run_mocked_cli(fault: str, output_dir: str) -> None: + launches = 0 + + def corpus(m): + bits = np.array([0x3F801001], dtype=np.uint32) + a = np.zeros((3, m, 128 // m), dtype=np.float32) + b = np.zeros_like(a) + a[:, :, 0] = bits.view(np.float32)[0] + b[:, :, 0] = 1 + c = np.zeros((3, m, m), dtype=np.float32) + return a, b, c, bits, 2 + + def launch(blob, name, a, b, c): + nonlocal launches + launches += 1 + mode = name.rsplit("_", 1)[1] + pa, pb = a.view(np.uint32).copy(), b.view(np.uint32).copy() + if mode == "rne": + pa, pb = numerics.rne_bits(pa), numerics.rne_bits(pb) + ra, rb = pa, pb + if mode != "fp32": + ra, rb = pa & np.uint32(0xFFFFE000), pb & np.uint32(0xFFFFE000) + d = np.einsum("bik,bjk->bij", ra.view(np.float32), rb.view(np.float32)) + c + if fault == "A preparation": + pa.flat[0] ^= np.uint32(1) + elif fault == "B preparation": + pb.flat[0] ^= np.uint32(1) + elif fault == "dense reference": + d[-1, 0, 0] += 1 + elif fault == "single product": + d[0, 0, 0] += 1 + elif (fault == "variant parity" and mode == "carrier") or ( + fault == "backend parity" and launches > 6 + ): + # Stay within the dense-reference tolerance while breaking exact parity. + d[-1, 0, 0] = np.nextafter(d[-1, 0, 0], np.float32(np.inf)) + return d, pa, pb + + def native_ir(m, mode): + kernel = numerics.build_tf32_mma_probe(numerics.Tf32MmaProbeSpec(m, mode)) + return "wrong IR" if fault == "builder parity" else numerics.serialize(kernel) + + native = SimpleNamespace( + tf32_mma_probe_serialize_ir=native_ir, + lower_serialized_ir=lambda *a, **k: ( + "wrong LLVM" if fault == "lowerer parity" else "mock LLVM" + ), + ) + argv = ["tf32_numerics", "--output-dir", output_dir, "--shape", "16"] + with ( + patch.object(numerics, "corpus", corpus), + patch.object(numerics, "launch", launch), + patch.object(numerics, "_resolve_llvm_flavor", return_value="llvm23"), + patch.object( + numerics, "_lower_kernel_to_llvm_python", return_value="mock LLVM" + ), + patch.object(numerics, "lower_kernel_to_hip", return_value="mock HIP"), + patch.object(hip_module, "get_device_arch", return_value="gfx942"), + patch.object(comgr, "build_hsaco_from_llvm_ir", return_value=(b"mock", None)), + patch.object(comgr, "resolved_lib_path", return_value="mock COMGR"), + patch.dict(sys.modules, {"rocke_engine": native}), + patch.object(sys, "argv", argv), + ): + numerics.main() + + +@pytest.mark.parametrize("optimization", ["normal", "flag", "environment"]) +@pytest.mark.parametrize( + "fault", + [ + "none", + "A preparation", + "B preparation", + "dense reference", + "single product", + "builder parity", + "lowerer parity", + "variant parity", + "backend parity", + ], +) +def test_cli_validation_survives_optimization(tmp_path, optimization, fault): + env = os.environ.copy() + env.pop("PYTHONOPTIMIZE", None) + env["PYTHONPATH"] = os.pathsep.join(sys.path) + env["ROCKE_BACKEND"] = "python" + if optimization == "environment": + env["PYTHONOPTIMIZE"] = "1" + command = [sys.executable] + if optimization == "flag": + command.append("-O") + command += [ + "-c", + ( + "import runpy, sys; " + "runpy.run_path(sys.argv[1])['_run_mocked_cli'](sys.argv[2], sys.argv[3])" + ), + str(Path(__file__).resolve()), + fault, + str(tmp_path), + ] + result = subprocess.run( + command, env=env, capture_output=True, text=True, timeout=60, check=False + ) + report = json.loads((tmp_path / "results.json").read_text()) + if fault == "none": + assert result.returncode == 0, result.stderr + assert report["status"] == "pass" + assert len(report["results"]) == 12 + else: + assert result.returncode != 0, result.stdout + assert report["status"] == "fail" + assert fault in report["error"] + assert fault in result.stderr From 4f18490d8ce6314ef06014a5f951a38312caa411 Mon Sep 17 00:00:00 2001 From: Max Podkorytov <4273004+tenpercent@users.noreply.github.com> Date: Wed, 30 Sep 2026 13:57:31 -0400 Subject: [PATCH 24/24] fix(rocke): export TF32 through public Python APIs --- .../rocke/platform/python/rocke/__init__.py | 2 ++ .../rocke/platform/python/rocke/core/__init__.py | 2 ++ .../rocke/platform/tests/core/test_tf32.py | 11 +++++++++++ 3 files changed, 15 insertions(+) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/__init__.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/__init__.py index ffbea41535ac..cf129d2b4945 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/__init__.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/__init__.py @@ -68,6 +68,7 @@ I16, I32, I64, + TF32, IRBuilder, KernelDef, Op, @@ -174,6 +175,7 @@ "I16", "I32", "I64", + "TF32", "IRBuilder", "KernelDef", "Op", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/__init__.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/__init__.py index c6b12bbdb49e..7abcf1acbf25 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/__init__.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/__init__.py @@ -61,6 +61,7 @@ I8, I32, I64, + TF32, IRBuilder, KernelDef, Op, @@ -123,6 +124,7 @@ "I8", "I32", "I64", + "TF32", "IRBuilder", "KernelDef", "Op", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py index bf583e6f9d2c..07cdf2eb1bf4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_tf32.py @@ -28,6 +28,17 @@ ) +def test_public_tf32_exports(): + import rocke + import rocke.core + from rocke import TF32 as public_tf32 + from rocke.core import TF32 as core_tf32 + + assert public_tf32 is core_tf32 is TF32 + assert "TF32" in rocke.__all__ + assert "TF32" in rocke.core.__all__ + + def test_dtype_identity(): assert TF32 != F32 and TF32 != I32 for alias in ("tf32", "xf32", " XF32 "):