From dd891e882ce48fb68af633a286ea7eb1d3a0fc1e Mon Sep 17 00:00:00 2001 From: Chris Millette Date: Thu, 8 Oct 2026 19:37:06 -0400 Subject: [PATCH 01/11] fix(rocke): look in the installed ROCm before PATH for the IR-validity compiler check_ir_validity._llvm_tool looked for clang beside the comgr rocke had loaded, then in $ROCM_PATH/$ROCM_HOME, then on PATH. When the loaded comgr is one bundled with another package (the ROCm torch wheel ships its own), no LLVM sits beside it and the lookup fell through to PATH, typically a distro clang that cannot link AMDGPU code objects. The lookup now also searches the real ROCm installs rocke's runtime already discovers before falling back to PATH. Co-Authored-By: Claude --- .../tests/core/test_arch_domain_artifact.py | 44 +++++++++++++++++++ .../rocke/platform/tools/check_ir_validity.py | 15 +++++++ 2 files changed, 59 insertions(+) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_arch_domain_artifact.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_arch_domain_artifact.py index 64819d0d7165..ddc710d52026 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_arch_domain_artifact.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_arch_domain_artifact.py @@ -874,6 +874,50 @@ def test_both_gates_ask_the_question_at_O0(self): self.assertIn("-O0", argv) self.assertNotIn("-O3", argv) + def test_a_bundled_comgr_does_not_send_the_lookup_to_path(self): + """A comgr bundled with another package (the ROCm torch wheel) has no LLVM + beside it. The lookup must then find the installed ROCm toolchain, not the + first clang on PATH, which is typically a distro clang that cannot link + AMDGPU code objects.""" + from unittest import mock + + import check_ir_validity as V + from rocke.runtime.runtime_coexistence import _rocm_root_libdirs + + # Only installs that ship a compiler can answer: a runtime-only ROCm has + # no /llvm/bin/clang, and the lookup then rightly finds none. + installs = [ + d + for d in _rocm_root_libdirs() + if any( + (root / "llvm" / "bin" / "clang").is_file() + for root in list(Path(d).resolve().parents)[:3] + ) + ] + if not installs: + self.skipTest("no ROCm install with an LLVM toolchain on this host") + with tempfile.TemporaryDirectory() as bundled: + fake_comgr = Path(bundled) / "lib" / "libamd_comgr.so" + fake_comgr.parent.mkdir() + fake_comgr.touch() + env = {k: v for k, v in os.environ.items() if k != "ROCKE_LLVM_BIN"} + with ( + mock.patch.dict(os.environ, env, clear=True), + mock.patch( + "rocke.runtime.comgr.resolved_lib_path", + return_value=str(fake_comgr), + ), + ): + found = V._llvm_tool("clang") + self.assertIsNotNone(found) + self.assertTrue( + any( + Path(found).resolve().is_relative_to(Path(d).resolve().parents[1]) + for d in installs + ), + f"{found} is not inside a ROCm install {installs}", + ) + def test_a_target_this_clang_cannot_build_is_not_an_emission_failure(self): """rocke wires arches ahead of the toolchains that can build for them. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tools/check_ir_validity.py b/dnn-providers/hip-kernel-provider/rocke/platform/tools/check_ir_validity.py index 55d587e4bcd0..fc8d088da98e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tools/check_ir_validity.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tools/check_ir_validity.py @@ -150,6 +150,21 @@ def _is_file(path: Path) -> bool: if raw: roots.append(Path(raw)) + # Real ROCm installs, found without going through a loaded comgr. When the + # comgr that loaded is one bundled with another package (the ROCm torch + # wheel ships its own), no LLVM sits beside it, and stopping here would fall + # through to PATH -- typically a distro clang with no AMDGPU linker, the + # wrong-vintage toolchain this lookup exists to avoid. + try: + from rocke.runtime.runtime_coexistence import _rocm_root_libdirs + + install_libdirs = _rocm_root_libdirs() + except (ImportError, OSError): # rocke not importable here, or an unreadable /opt + install_libdirs = [] + for libdir in install_libdirs: + libp = Path(libdir).resolve() + roots.extend(libp.parents[i] for i in range(min(3, len(libp.parents)))) + for root in roots: cand = root / "llvm" / "bin" / exe if _is_file(cand): From 4580b24b4c43b4d9793fcb7884bd4c4f26739749 Mon Sep 17 00:00:00 2001 From: Chris Millette Date: Thu, 8 Oct 2026 17:41:19 -0400 Subject: [PATCH 02/11] test(rocke): add a compile-and-link check for lowered IR rocke.tests.ir_validity lets a test ask whether lowered LLVM IR compiles and links for a target, using the same compiler lookup and compile command as platform/tools/check_ir_validity.py. It reports accepted, rejected (with clang's first error line) or cannot_judge (no tools, no clang, a clang that does not know the target, or a clang that cannot run its linker); assert_target_accepts turns cannot_judge into a named skip. Co-Authored-By: Claude --- .../python/rocke/tests/ir_validity.py | 95 +++++++++++++++++ .../python/rocke/tests/test_ir_validity.py | 100 ++++++++++++++++++ 2 files changed, 195 insertions(+) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/ir_validity.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/test_ir_validity.py diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/ir_validity.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/ir_validity.py new file mode 100644 index 000000000000..3bfa3a91535c --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/ir_validity.py @@ -0,0 +1,95 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Test helper: judge whether lowered LLVM IR compiles and links for one target. + +It reuses rocke's IR-validity tool (``platform/tools/check_ir_validity.py``) so a +test and the corpus gate use the same compiler, the same flags and the same +definition of "valid": the module must compile *and link* to a code object. The +link is what catches a call to an intrinsic that does not exist; see that tool's +module docstring. + +``arbiter_verdict`` reports one of three verdicts: + +- ``"accepted"``: clang compiled and linked the module for the target. +- ``"rejected"``: clang read the module and refused it; the text is clang's first + error line. +- ``"cannot_judge"``: no verdict is possible here (no tools, no clang, a clang that + does not know the target, or a clang that cannot run its linker); the text says + which. + +``assert_target_accepts`` turns ``"cannot_judge"`` into a named pytest skip, so a +host whose toolchain predates a target reports the gap instead of passing. +""" + +from __future__ import annotations + +import sys +import tempfile +from pathlib import Path +from types import ModuleType +from typing import Literal + +import pytest + +from rocke.assets import platform_root + +Verdict = Literal["accepted", "rejected", "cannot_judge"] + +__all__ = ["Verdict", "arbiter_verdict", "assert_target_accepts"] + + +def _validity_tool() -> ModuleType | None: + """The IR-validity tool module, or None when the platform tools are not + shipped alongside this package (e.g. an installed wheel).""" + tools = platform_root() / "tools" + if not (tools / "check_ir_validity.py").is_file(): + return None + if str(tools) not in sys.path: + sys.path.insert(0, str(tools)) + import check_ir_validity + + return check_ir_validity + + +def arbiter_verdict(ir_text: str, target: str) -> tuple[Verdict, str]: + """Compile and link ``ir_text`` for ``target`` with the host clang. + + Returns ``(verdict, text)``; ``text`` is empty when the module is accepted. + """ + tool = _validity_tool() + if tool is None: + return "cannot_judge", "rocke platform tools not found" + clang = tool._llvm_tool("clang") + if clang is None: + return "cannot_judge", "no host clang to compile lowered IR" + with tempfile.TemporaryDirectory() as work: + source = Path(work) / "module.ll" + source.write_text(ir_text) + accepted, diagnostic = tool._compile( + clang, source, target, Path(work) / "module.hsaco" + ) + if accepted: + return "accepted", "" + if tool._target_unsupported(diagnostic): + return "cannot_judge", f"host clang cannot target {target}" + if _toolchain_incomplete(diagnostic): + return ( + "cannot_judge", + f"host clang {clang} cannot link: {diagnostic.splitlines()[0]}", + ) + return "rejected", tool._first_diagnostic(diagnostic) + + +def _toolchain_incomplete(diagnostic: str) -> bool: + """Whether clang failed because a tool it drives is missing, before judging the module.""" + return "unable to execute command" in diagnostic + + +def assert_target_accepts(ir_text: str, target: str) -> None: + """Fail unless ``ir_text`` compiles and links for ``target``; skip, naming + the reason, when no verdict is possible on this host.""" + verdict, text = arbiter_verdict(ir_text, target) + if verdict == "cannot_judge": + pytest.skip(text) + assert verdict == "accepted", f"LLVM rejected the lowered IR for {target}: {text}" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/test_ir_validity.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/test_ir_validity.py new file mode 100644 index 000000000000..120d76572607 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/test_ir_validity.py @@ -0,0 +1,100 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""The compile-and-link test helper gives the three verdicts it promises. + +Each case is a small hand-written module, so the expected verdict is fixed by the +module itself and not by the current state of rocke's lowering. A host without a +usable clang skips, which is the helper's own contract. +""" + +from __future__ import annotations + +import pytest + +from rocke.tests.ir_validity import arbiter_verdict, assert_target_accepts + +_VALID = """\ +target triple = "amdgcn-amd-amdhsa" +define amdgpu_kernel void @k(ptr addrspace(1) %p) { + store float 1.0, ptr addrspace(1) %p + ret void +} +""" + +# An MFMA operand printed as `float` but defined as `<1 x float>`: the parser +# rejects it before any code generation. +_WRONG_OPERAND_TYPE = """\ +target triple = "amdgcn-amd-amdhsa" +declare <4 x float> @llvm.amdgcn.mfma.f32.16x16x4f32(float, float, <4 x float>, i32, i32, i32) +define amdgpu_kernel void @k(ptr addrspace(1) %p, <1 x float> %a) { + %r = call <4 x float> @llvm.amdgcn.mfma.f32.16x16x4f32(float %a, float 1.0, <4 x float> zeroinitializer, i32 0, i32 0, i32 0) + store <4 x float> %r, ptr addrspace(1) %p + ret void +} +""" + +# A call to an intrinsic that does not exist compiles to an object; only the +# link step reports it. +_UNDEFINED_INTRINSIC = """\ +target triple = "amdgcn-amd-amdhsa" +declare float @llvm.amdgcn.not.an.intrinsic(float) +define amdgpu_kernel void @k(ptr addrspace(1) %p) { + %r = call float @llvm.amdgcn.not.an.intrinsic(float 1.0) + store float %r, ptr addrspace(1) %p + ret void +} +""" + + +def _verdict_or_skip(ir_text: str, target: str) -> tuple[str, str]: + verdict, text = arbiter_verdict(ir_text, target) + if verdict == "cannot_judge": + pytest.skip(text) + return verdict, text + + +def test_a_valid_module_is_accepted(): + assert _verdict_or_skip(_VALID, "gfx90a") == ("accepted", "") + assert_target_accepts(_VALID, "gfx90a") + + +def test_a_mistyped_operand_is_rejected_with_the_parser_error(): + verdict, text = _verdict_or_skip(_WRONG_OPERAND_TYPE, "gfx90a") + assert verdict == "rejected" + assert "defined with type '<1 x float>' but expected 'float'" in text + + +def test_a_call_to_a_nonexistent_intrinsic_is_rejected_at_link(): + verdict, text = _verdict_or_skip(_UNDEFINED_INTRINSIC, "gfx90a") + assert verdict == "rejected" + assert "undefined symbol: llvm.amdgcn.not.an.intrinsic" in text + + +def test_a_target_the_host_clang_does_not_know_cannot_be_judged(): + verdict, text = arbiter_verdict(_VALID, "gfx9999") + if verdict == "cannot_judge" and "no host clang" in text: + pytest.skip(text) + assert (verdict, text) == ("cannot_judge", "host clang cannot target gfx9999") + + +def test_assert_target_accepts_fails_on_a_rejected_module(): + _verdict_or_skip(_WRONG_OPERAND_TYPE, "gfx90a") + with pytest.raises(AssertionError, match="LLVM rejected the lowered IR for gfx90a"): + assert_target_accepts(_WRONG_OPERAND_TYPE, "gfx90a") + + +def test_a_clang_that_cannot_run_its_linker_is_not_a_verdict_on_the_module(): + # The message a distro clang without an AMDGPU linker prints. It says nothing about the + # module, so it must not be reported as a rejection. + from rocke.tests.ir_validity import _toolchain_incomplete + + assert _toolchain_incomplete( + 'clang: error: unable to execute command: Executable "ld.lld" doesn\'t exist!' + ) + assert not _toolchain_incomplete( + "ld.lld: error: undefined symbol: llvm.amdgcn.not.an.intrinsic" + ) + assert not _toolchain_incomplete( + "module.ll:4:64: error: '%a' defined with type '<1 x float>' but expected 'float'" + ) From cf64c627eefab56cc10ad76f927cc81117809cf4 Mon Sep 17 00:00:00 2001 From: Chris Millette Date: Wed, 14 Oct 2026 11:22:01 -0400 Subject: [PATCH 03/11] feat(rocke): add explicit f32 and f64 conversion ops cvt_f32_to_f64 (exact, fpext) and cvt_f64_to_f32 (round to nearest even, fptrunc) are builder methods on their own opcodes, arith.cvt_f32_to_f64 and arith.cvt_f64_to_f32, in both engines: the Python and C++ builders, the LLVM lowerings and the HIP lowerings. Each lowers to one instruction or one C cast. cast_to_f32 and cast_f32_to remain the f16/bf16 <-> f32 hub and keep rejecting f64, so no code built on them converts a double by accident; a conversion given the wrong source type fails with the same message in both builders and emits no op. zero_vec accepts F64 and I16 in both builders, so fill_fragment can zero an f64 fragment, and smem_load_vN accepts i16 in both builders, as smem_store_vN already did. rocke.tests.gpu runs a numeric test on the current GPU, the one rocke's runtime launches on, which hip_module now reports through a public current_device(), only when its arch is one the test names as supporting its data types, lowering and compiling with the active engine; on such an arch a lowering error or compiler rejection fails the test. On the GPU, copy, neg, abs, add, sub and mul of f64 and both conversions are checked bit-exact over ordinary values and the IEEE edge values (out-of-f32-range, subnormal, signed zeros, infinities, NaN), except where IEEE 754 leaves the bits open. The C++ builder side is covered by the rocke_f64_lowering ctest. The two opcodes are appended to rocke_opcode_t, so no existing opcode value changes, and the rocke_opcode_values ctest pins the value, name and purity of every opcode that existed before them. Co-Authored-By: Claude --- ...ROCKE_ENGINE_test_categories_external.yaml | 2 + .../rocke/platform/cpp/core/ir/core_types.cpp | 12 +- .../rocke/platform/cpp/core/ir/ir_arith.cpp | 36 +- .../rocke/platform/cpp/core/ir/ir_tile.cpp | 11 +- .../cpp/core/lower_hip/lower_hip_cast.cpp | 24 ++ .../platform/cpp/core/lower_llvm/convert.cpp | 17 + .../rocke/platform/cpp/include/rocke/ir.h | 14 +- .../platform/dsl_docs/ir_lowering/ir_model.md | 6 +- .../dsl_docs/reference/dtype_support.md | 6 +- .../dsl_docs/reference/op_vocabulary.md | 2 + .../rocke/platform/python/rocke/core/ir.py | 35 +- .../platform/python/rocke/core/lower_hip.py | 10 + .../platform/python/rocke/core/lower_llvm.py | 13 + .../python/rocke/runtime/hip_module.py | 7 +- .../rocke/platform/python/rocke/tests/gpu.py | 106 ++++++ .../platform/python/rocke/tests/test_gpu.py | 47 +++ .../rocke/platform/tests/CMakeLists.txt | 16 + .../tests/core/test_f64_conversions.py | 204 +++++++++++ .../platform/tests/core/test_f64_lowering.cpp | 95 ++++- .../tests/core/test_opcode_values.cpp | 334 ++++++++++++++++++ .../test_zero_vec_and_lds_element_types.py | 92 +++++ .../tests/runtime/test_cluster_launch.py | 4 +- 22 files changed, 1068 insertions(+), 25 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/gpu.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/test_gpu.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_f64_conversions.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_opcode_values.cpp create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_zero_vec_and_lds_element_types.py diff --git a/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml b/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml index 6e36eb870342..9a7d95abbfc9 100644 --- a/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml +++ b/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml @@ -33,6 +33,7 @@ test_categories: - "rocke_wgrad_streamk_host" - "rocke_wmma_fmha_decode_host" - "rocke_future_intrinsic_lowering" + - "rocke_ir_types" - "rocke_tiled_attention_2d_reentrancy" - "rocke_unified_attention_scalar_lower_to_llvm" - "rocke_python_import_smoke" @@ -64,6 +65,7 @@ test_categories: - "rocke_wgrad_streamk_host" - "rocke_wmma_fmha_decode_host" - "rocke_future_intrinsic_lowering" + - "rocke_ir_types" - "rocke_tiled_attention_2d_reentrancy" - "rocke_unified_attention_scalar_lower_to_llvm" - "rocke_python_import_smoke" 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 a9bcf7cea7de..cf0557cff26a 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 @@ -704,7 +704,11 @@ static const char* const rocke_opcode_names[ROCKE_OP__COUNT] = { "memref.global_flag_wait_eq", "vector.undef", - "tile.buffer_load_f16_d16"}; + "tile.buffer_load_f16_d16", + + /* arith.* f32 <-> f64 conversions */ + "arith.cvt_f32_to_f64", + "arith.cvt_f64_to_f32"}; const char* rocke_opcode_name(rocke_opcode_t op) { @@ -973,7 +977,11 @@ static const bool rocke_opcode_pure[ROCKE_OP__COUNT] = { /* memref.global_flag_wait_eq */ false, /* vector.undef */ false, - /* tile.buffer_load_f16_d16 */ false}; + /* tile.buffer_load_f16_d16 */ false, + + /* arith.* f32 <-> f64 conversions */ + /* arith.cvt_f32_to_f64 */ true, + /* arith.cvt_f64_to_f32 */ true}; bool rocke_opcode_is_pure(rocke_opcode_t op) { 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 1e2f7c367ef8..1244fed1bda2 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 @@ -124,8 +124,9 @@ rocke_value_t* rocke_b_zero_vec(rocke_ir_builder_t* b, const rocke_type_t* elem, if(rocke_type_eq(elem, rocke_f32())) return rocke_b_zero_vec_f32(b, n); if(rocke_type_eq(elem, rocke_f16()) || rocke_type_eq(elem, rocke_bf16()) - || rocke_type_eq(elem, rocke_fp8e4m3()) || rocke_type_eq(elem, rocke_bf8e5m2()) - || rocke_type_eq(elem, rocke_i8()) || rocke_type_eq(elem, rocke_i32())) + || rocke_type_eq(elem, rocke_f64()) || rocke_type_eq(elem, rocke_fp8e4m3()) + || rocke_type_eq(elem, rocke_bf8e5m2()) || rocke_type_eq(elem, rocke_i8()) + || rocke_type_eq(elem, rocke_i16()) || rocke_type_eq(elem, rocke_i32())) { hint = cz_hint(b, n); if(!hint) @@ -700,6 +701,37 @@ rocke_value_t* rocke_b_sitofp_f32(rocke_ir_builder_t* b, rocke_value_t* v) return rocke_i_op1(b, ROCKE_OP_ARITH_SITOFP_F32, operands, 1, rocke_f32(), NULL, "sitof"); } +/* f32 <-> f64 has its own opcode per direction rather than riding on + * cast_to_f32 / cast_f32_to: those two are the f16/bf16 <-> f32 hub and keep + * rejecting f64, so no caller widens or narrows a double by accident. */ +rocke_value_t* rocke_b_cvt_f32_to_f64(rocke_ir_builder_t* b, rocke_value_t* v) +{ + rocke_value_t* operands[1]; + if(!rocke_i_live(b)) + return NULL; + if(!v) + return (rocke_value_t*)rocke_i_set_err(b, ROCKE_ERR_VALUE, "cvt_f32_to_f64 NULL value"); + if(!rocke_i_type_is(v->type, "f32")) + return (rocke_value_t*)rocke_i_set_err( + b, ROCKE_ERR_VALUE, "cvt_f32_to_f64 expects f32 input, got %s", v->type->name); + operands[0] = v; + return rocke_i_op1(b, ROCKE_OP_ARITH_CVT_F32_TO_F64, operands, 1, rocke_f64(), NULL, "f64"); +} + +rocke_value_t* rocke_b_cvt_f64_to_f32(rocke_ir_builder_t* b, rocke_value_t* v) +{ + rocke_value_t* operands[1]; + if(!rocke_i_live(b)) + return NULL; + if(!v) + return (rocke_value_t*)rocke_i_set_err(b, ROCKE_ERR_VALUE, "cvt_f64_to_f32 NULL value"); + if(!rocke_i_type_is(v->type, "f64")) + return (rocke_value_t*)rocke_i_set_err( + b, ROCKE_ERR_VALUE, "cvt_f64_to_f32 expects f64 input, got %s", v->type->name); + operands[0] = v; + return rocke_i_op1(b, ROCKE_OP_ARITH_CVT_F64_TO_F32, operands, 1, rocke_f32(), NULL, "f32"); +} + rocke_value_t* rocke_b_cvt_fp8_to_f32(rocke_ir_builder_t* b, rocke_value_t* v) { rocke_value_t* operands[1]; 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 f18114e2c112..6787e4e033b3 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 @@ -380,14 +380,15 @@ rocke_value_t* rocke_b_smem_load_vN(rocke_ir_builder_t* b, return (rocke_value_t*)rocke_i_set_err(b, ROCKE_ERR_VALUE, "smem_load_vN dtype is NULL"); } dn = dtype->name; - if(!(strcmp(dn, "f16") == 0 || strcmp(dn, "bf16") == 0 || strcmp(dn, "f32") == 0 - || strcmp(dn, "tf32") == 0 || strcmp(dn, "i32") == 0 || strcmp(dn, "f64") == 0 - || strcmp(dn, "fp8e4m3") == 0 || strcmp(dn, "bf8e5m2") == 0 || strcmp(dn, "i8") == 0)) + if(!(strcmp(dn, "f16") == 0 || strcmp(dn, "bf16") == 0 || strcmp(dn, "i16") == 0 + || strcmp(dn, "f32") == 0 || strcmp(dn, "tf32") == 0 || strcmp(dn, "i32") == 0 + || strcmp(dn, "f64") == 0 || strcmp(dn, "fp8e4m3") == 0 || strcmp(dn, "bf8e5m2") == 0 + || strcmp(dn, "i8") == 0)) { return (rocke_value_t*)rocke_i_set_err( b, ROCKE_ERR_VALUE, - "smem_load_vN supports f16 / bf16 / f32 / i32 / tf32 / f64 / " + "smem_load_vN supports f16 / bf16 / i16 / f32 / i32 / tf32 / f64 / " "fp8e4m3 / bf8e5m2 / i8, got %s", dn); } @@ -399,7 +400,7 @@ rocke_value_t* rocke_b_smem_load_vN(rocke_ir_builder_t* b, allowed = allowed_8bit; acount = 6; } - else if(strcmp(dn, "f16") == 0 || strcmp(dn, "bf16") == 0) + else if(strcmp(dn, "f16") == 0 || strcmp(dn, "bf16") == 0 || strcmp(dn, "i16") == 0) { allowed = allowed_16bit; acount = 5; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_cast.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_cast.cpp index 67b6cf66bbf8..03150adab205 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_cast.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_cast.cpp @@ -171,6 +171,28 @@ static rocke_status_t rocke_h_op_arith_cast_f32_to(rocke_h_lowerer_t* lw, const return lw->status; } +/* def _op_arith_cvt_f32_to_f64(self, op): + * (v,) = op.operands + * self._emit(f"double {_name(op.result)} = (double){_name(v)};") */ +static rocke_status_t rocke_h_op_arith_cvt_f32_to_f64(rocke_h_lowerer_t* lw, const rocke_op_t* op) +{ + const rocke_value_t* v = op->operands[0]; + const rocke_value_t* r = h_res(op); + rocke_h_emitf(lw, "double %s = (double)%s;", rocke_h_name(lw, r), rocke_h_name(lw, v)); + return lw->status; +} + +/* def _op_arith_cvt_f64_to_f32(self, op): + * (v,) = op.operands + * self._emit(f"float {_name(op.result)} = (float){_name(v)};") */ +static rocke_status_t rocke_h_op_arith_cvt_f64_to_f32(rocke_h_lowerer_t* lw, const rocke_op_t* op) +{ + const rocke_value_t* v = op->operands[0]; + const rocke_value_t* r = h_res(op); + rocke_h_emitf(lw, "float %s = (float)%s;", rocke_h_name(lw, r), rocke_h_name(lw, v)); + return lw->status; +} + /* def _op_arith_sitofp_f32(self, op): * (v,) = op.operands * self._emit(f"float {_name(op.result)} = (float){_name(v)};") */ @@ -455,6 +477,8 @@ const rocke_h_handler_entry_t* rocke_h_handlers_cast(void) {ROCKE_OP_ARITH_BITCAST, rocke_h_op_arith_bitcast}, {ROCKE_OP_ARITH_CAST_TO_F32, rocke_h_op_arith_cast_to_f32}, {ROCKE_OP_ARITH_CAST_F32_TO, rocke_h_op_arith_cast_f32_to}, + {ROCKE_OP_ARITH_CVT_F32_TO_F64, rocke_h_op_arith_cvt_f32_to_f64}, + {ROCKE_OP_ARITH_CVT_F64_TO_F32, rocke_h_op_arith_cvt_f64_to_f32}, {ROCKE_OP_ARITH_SITOFP_F32, rocke_h_op_arith_sitofp_f32}, {ROCKE_OP_ARITH_CVT_FP8_TO_F32, rocke_h_op_arith_cvt_fp8_to_f32}, {ROCKE_OP_ARITH_CVT_BF8_TO_F32, rocke_h_op_arith_cvt_bf8_to_f32}, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/convert.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/convert.cpp index 0e1ac465b3b1..a85cd62fe14b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/convert.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/convert.cpp @@ -133,6 +133,21 @@ static void _op_arith_cast_f32_to(rocke_lower_t* L, const rocke_op_t* op) rocke_ll_llvm_type(L, op->results[0]->type)); } +static void _op_arith_cvt_f32_to_f64(rocke_lower_t* L, const rocke_op_t* op) +{ + const rocke_value_t* v = op->operands[0]; + rocke_ll_emitf( + L, " %s = fpext float %s to double", rocke_ll_result_name(op), rocke_ll_operand(L, v)); +} + +/* LLVM's fptrunc rounds to nearest even under the default FP environment. */ +static void _op_arith_cvt_f64_to_f32(rocke_lower_t* L, const rocke_op_t* op) +{ + const rocke_value_t* v = op->operands[0]; + rocke_ll_emitf( + L, " %s = fptrunc double %s to float", rocke_ll_result_name(op), rocke_ll_operand(L, v)); +} + static void _op_arith_sitofp_f32(rocke_lower_t* L, const rocke_op_t* op) { const rocke_value_t* v = op->operands[0]; @@ -506,6 +521,8 @@ void rocke_ll_register_convert(void) rocke_ll_set_handler(ROCKE_OP_ARITH_TRUNC_F32_TO_BF16, _op_arith_trunc_f32_to_bf16); rocke_ll_set_handler(ROCKE_OP_ARITH_CAST_TO_F32, _op_arith_cast_to_f32); rocke_ll_set_handler(ROCKE_OP_ARITH_CAST_F32_TO, _op_arith_cast_f32_to); + rocke_ll_set_handler(ROCKE_OP_ARITH_CVT_F32_TO_F64, _op_arith_cvt_f32_to_f64); + rocke_ll_set_handler(ROCKE_OP_ARITH_CVT_F64_TO_F32, _op_arith_cvt_f64_to_f32); rocke_ll_set_handler(ROCKE_OP_ARITH_SITOFP_F32, _op_arith_sitofp_f32); rocke_ll_set_handler(ROCKE_OP_ARITH_RINT_F32, _op_arith_rint_f32); 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 d642a45aa873..7e781d9f76e0 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 @@ -198,7 +198,12 @@ typedef struct rocke_attr_map /* One enumerator per distinct op name string in rocke.core.ir. Every lowerer * dispatches on this enum instead of Python getattr(self, "_op_"+name). The - * canonical dotted name string is recovered with rocke_opcode_name(). */ + * canonical dotted name string is recovered with rocke_opcode_name(). + * + * A new opcode goes at the end, just before ROCKE_OP__COUNT. Inserting one + * anywhere else renumbers every opcode after it, which breaks the C ABI (see + * "when to bump" in rocke/abi.h). tests/core/test_opcode_values.cpp pins the + * value of every existing opcode. */ typedef enum rocke_opcode { ROCKE_OP_INVALID = 0, @@ -447,6 +452,10 @@ typedef enum rocke_opcode ROCKE_OP_VECTOR_UNDEF, ROCKE_OP_TILE_BUFFER_LOAD_F16_D16, + /* arith.* f32 <-> f64 conversions */ + ROCKE_OP_ARITH_CVT_F32_TO_F64, + ROCKE_OP_ARITH_CVT_F64_TO_F32, + ROCKE_OP__COUNT } rocke_opcode_t; @@ -841,6 +850,9 @@ rocke_value_t* rocke_b_cast_to_f32(rocke_ir_builder_t* b, rocke_value_t* v); rocke_value_t* rocke_b_cast_f32_to(rocke_ir_builder_t* b, rocke_value_t* v, const rocke_type_t* target); rocke_value_t* rocke_b_sitofp_f32(rocke_ir_builder_t* b, rocke_value_t* v); +/* f32 <-> f64 have their own opcodes; cast_to_f32 / cast_f32_to reject f64. */ +rocke_value_t* rocke_b_cvt_f32_to_f64(rocke_ir_builder_t* b, rocke_value_t* v); +rocke_value_t* rocke_b_cvt_f64_to_f32(rocke_ir_builder_t* b, rocke_value_t* v); rocke_value_t* rocke_b_cvt_fp8_to_f32(rocke_ir_builder_t* b, rocke_value_t* v); rocke_value_t* rocke_b_cvt_bf8_to_f32(rocke_ir_builder_t* b, rocke_value_t* v); rocke_value_t* rocke_b_cvt_pk_f32_fp8x4(rocke_ir_builder_t* b, rocke_value_t* v); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/ir_lowering/ir_model.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/ir_lowering/ir_model.md index 570729eb16f3..9acf379d2c79 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/ir_lowering/ir_model.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/ir_lowering/ir_model.md @@ -92,7 +92,7 @@ const_f64(value) # printed as the exact IEEE-754 bit pattern fp16_zero() zero_vec_f32(n) # , all zeros zero_vec_f32_4() # <4 x f32>, all zeros -zero_vec(elem, n) # f32 / f16 / bf16 +zero_vec(elem, n) # f32 / f16 / bf16 / f64 / fp8e4m3 / bf8e5m2 / i8 / i16 / i32 ``` ### Arithmetic / compare / cast @@ -111,6 +111,8 @@ masked_select # alias for select sitofp_f32(v) # i32 -> f32 (LLVM sitofp) cast_to_f32(v) # f16/bf16 -> f32 cast_f32_to(v, target) # f32 -> f16/bf16 +cvt_f32_to_f64(v) # f32 -> f64 (fpext, exact) +cvt_f64_to_f32(v) # f64 -> f32 (fptrunc, round to nearest even) trunc_f32_to_f16(v) # f32 -> f16 (fptrunc with attrs) bitcast(v, target) # scalar bitcast of equal size vec_bitcast(v, target_vec_type) # vector bitcast (e.g. <2xi32> <-> <4xf16>) @@ -198,7 +200,7 @@ smem_ptr_add(lds_addr, byte_off) # i64 LDS address arithmetic smem_load_f16(smem, indices) smem_load_vN_f16(smem, *indices, n in {1,2,4,8}) smem_load_v4_f16(smem, row, col) -smem_load_vN(smem, *indices, dtype, n) # f16/bf16 +smem_load_vN(smem, *indices, dtype, n) # 8/16/32/64-bit dtypes; see the docstring smem_store_f16(smem, indices, value) smem_store_vN_f16(smem, indices, value, n) # ds_write_b{16,32,64,128} smem_store_vN(smem, indices, value, n) # n in {1,2,4,8} 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 109b94c46e96..9166a6a2f499 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 @@ -357,8 +357,10 @@ mode. The missing selector and precision contract are described in §2.5. (non-transcendental ops) and reduce2d accept f64 with `vec=2` ([SUPPORT_MATRIX](../../python/rocke/instances/SUPPORT_MATRIX.md)). `exp2`, `exp2_fast`, `log2`, `sqrt`, `rsqrt`, `rcp_fast`, `tanh`, `fabs`, - `fma`, `fmax3`, `fmin3`, and `vector_fma` reject f64 in both lowerings; there - are no conversions to or from f64, MMA atoms, or atomics. f64 is in the + `fma`, `fmax3`, `fmin3`, and `vector_fma` reject f64 in both lowerings. + `cvt_f32_to_f64` (exact) and `cvt_f64_to_f32` (round to nearest even) are the + only conversions; `cast_to_f32` / `cast_f32_to` refuse f64, and there are no + direct f16/bf16 conversions, MMA atoms, or atomics. f64 is in the hipDNN enum (`DOUBLE`) but rejected by the MIOpen provider. `V_MFMA_F64_*` is present on gfx90a/gfx942/gfx950 but undeclared (§4C.6). gfx1250 has no F64 matrix instruction: LLVM gates `V_WMMA_F64_16X16X4_F64` on 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 b90b994e3dff..a4980daf17c5 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 @@ -35,6 +35,8 @@ Complete reference of operations recognized by `core/ir.py` and lowered to AMDGP | `arith.cast_f32_to` | `cast_f32_to` | `fptrunc` to f16/bf16 | | `arith.trunc_f32_to_f16` | `trunc_f32_to_f16` | `fptrunc` (vec-packed in some emissions) | | `arith.sitofp_f32` | `sitofp_f32` | `sitofp i32 to float` | +| `arith.cvt_f32_to_f64` | `cvt_f32_to_f64` | `fpext float to double` (exact) | +| `arith.cvt_f64_to_f32` | `cvt_f64_to_f32` | `fptrunc double to float` (round to nearest even) | | `arith.cvt_fp8_to_f32` | `cvt_fp8_to_f32` | `llvm.amdgcn.cvt.f32.fp8` | | `arith.cvt_bf8_to_f32` | `cvt_bf8_to_f32` | `llvm.amdgcn.cvt.f32.bf8` | | `arith.cvt_f32_to_fp8` | `cvt_f32_to_fp8` | `llvm.amdgcn.cvt.pk.fp8.f32` (low byte) | 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 d9fdf3ba8103..da1e9d1eab52 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 @@ -1252,6 +1252,28 @@ def sitofp_f32(self, v: Value) -> Value: raise ValueError(f"sitofp_f32 expects i32 input, got {v.type.name}") return self._op("arith.sitofp_f32", [v], [F32], result_name_hint="sitof").result + # f32 <-> f64 has its own opcode per direction rather than riding on + # cast_to_f32 / cast_f32_to: those two are the f16/bf16 <-> f32 hub and + # keep rejecting f64, so no caller widens or narrows a double by accident. + def cvt_f32_to_f64(self, v: Value) -> Value: + """Convert an f32 to f64. Exact; lowers to one ``fpext float to double``.""" + if v.type != F32: + raise ValueError(f"cvt_f32_to_f64 expects f32 input, got {v.type.name}") + return self._op( + "arith.cvt_f32_to_f64", [v], [F64], result_name_hint="f64" + ).result + + def cvt_f64_to_f32(self, v: Value) -> Value: + """Convert an f64 to f32, rounding to nearest even. + + Lowers to one ``fptrunc double to float``. + """ + if v.type != F64: + raise ValueError(f"cvt_f64_to_f32 expects f64 input, got {v.type.name}") + return self._op( + "arith.cvt_f64_to_f32", [v], [F32], result_name_hint="f32" + ).result + def cvt_fp8_to_f32(self, v: Value) -> Value: """Convert one fp8e4m3 element to f32. @@ -1795,7 +1817,7 @@ def zero_vec_f32(self, n: int) -> Value: def zero_vec(self, elem: Type, n: int) -> Value: if elem == F32: return self.zero_vec_f32(n) - if elem in (F16, BF16, FP8E4M3, BF8E5M2, I8, I32): + if elem in (F16, BF16, F64, FP8E4M3, BF8E5M2, I8, I16, I32): return self._op( "arith.constant_vec", result_types=[VectorType(elem, n)], @@ -2311,8 +2333,8 @@ 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), 32-bit (f32 / i32) and 64-bit (f64, ``n in - {1, 2}``) element types; loads exactly N elements. In addition to + i8), 16-bit (f16 / bf16 / i16), 32-bit (f32 / i32) and 64-bit (f64, + ``n in {1, 2}``) element types; 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. @@ -2320,6 +2342,7 @@ def smem_load_vN(self, smem: Value, *indices, dtype: Type, n: int = 0) -> Value: if dtype.name not in ( "f16", "bf16", + "i16", "f32", "i32", "tf32", @@ -2329,12 +2352,12 @@ def smem_load_vN(self, smem: Value, *indices, dtype: Type, n: int = 0) -> Value: "i8", ): raise ValueError( - "smem_load_vN supports f16 / bf16 / f32 / i32 / tf32 / f64 / " + "smem_load_vN supports f16 / bf16 / i16 / f32 / i32 / tf32 / f64 / " f"fp8e4m3 / bf8e5m2 / i8, got {dtype.name}" ) if dtype.name in ("fp8e4m3", "bf8e5m2", "i8"): allowed_n = (1, 2, 4, 8, 12, 16) - elif dtype.name in ("f16", "bf16"): + elif dtype.name in ("f16", "bf16", "i16"): allowed_n = (1, 2, 4, 6, 8) elif dtype.name == "f64": allowed_n = (1, 2) @@ -5399,6 +5422,8 @@ def scf_if_else(self, cond: Value): "arith.cast_to_f32", "arith.cast_f32_to", "arith.sitofp_f32", + "arith.cvt_f32_to_f64", + "arith.cvt_f64_to_f32", "arith.cvt_fp8_to_f32", "math.exp2", "math.exp2_fast", 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 b2b08f068cd9..8c10f09cebca 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 @@ -1575,6 +1575,16 @@ def _op_arith_cast_f32_to(self, op: Op) -> None: cpp_t = _type_to_hip(op.result.type) self._emit(f"{cpp_t} {_name(op.result)} = ({cpp_t}){_name(v)};") + def _op_arith_cvt_f32_to_f64(self, op: Op) -> None: + # f32 -> f64 is exact; the C conversion is the same fpext. + (v,) = op.operands + self._emit(f"double {_name(op.result)} = (double){_name(v)};") + + def _op_arith_cvt_f64_to_f32(self, op: Op) -> None: + # f64 -> f32 rounds to nearest even, as the C conversion does. + (v,) = op.operands + self._emit(f"float {_name(op.result)} = (float){_name(v)};") + def _op_arith_sitofp_f32(self, op: Op) -> None: # i32 -> f32. C-style cast is sufficient. (v,) = op.operands 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 a0cc6db467ce..877b82649784 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 @@ -2625,6 +2625,19 @@ def _op_arith_cast_f32_to(self, op: Op) -> None: f" {op.result.name} = fptrunc float {self._operand(v)} to {_llvm_type(op.result.type)}" ) + def _op_arith_cvt_f32_to_f64(self, op: Op) -> None: + (v,) = op.operands + self._current().emit( + f" {op.result.name} = fpext float {self._operand(v)} to double" + ) + + def _op_arith_cvt_f64_to_f32(self, op: Op) -> None: + # LLVM's fptrunc rounds to nearest even under the default FP environment. + (v,) = op.operands + self._current().emit( + f" {op.result.name} = fptrunc double {self._operand(v)} to float" + ) + def _op_arith_sitofp_f32(self, op: Op) -> None: (v,) = op.operands if v.type.name != "i32": diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/hip_module.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/hip_module.py index ce97d3580943..5580c573e67d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/hip_module.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/hip_module.py @@ -339,7 +339,10 @@ def device_supports_cluster_launch(device: int = 0) -> Optional[bool]: return _cluster_launch_support[device] -def _current_device() -> int: +def current_device() -> int: + """HIP's current device for this thread: the one rocke's runtime launches + kernels on. Returns ``0``, HIP's default device, when the runtime cannot be + queried.""" cur = ctypes.c_int(0) try: rc = _hipGetDevice(ctypes.byref(cur)) @@ -685,7 +688,7 @@ def _enqueue( return () shape = check_launch_cluster(grid, cluster) - device = _current_device() + device = current_device() if device_supports_cluster_launch(device) is False: raise HipError( f"{where}: device {device} does not support cluster launch " diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/gpu.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/gpu.py new file mode 100644 index 000000000000..c6c221b50ffe --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/gpu.py @@ -0,0 +1,106 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Test helper: run a small kernel on the current GPU, only on arches that support it. + +Each numeric test names the arches whose hardware supports its data types. On any other +arch, or with no GPU, the test skips with the reason. On a supported arch every failure is +loud: a lowering error or a compiler rejection fails the test, it is never turned into a +skip. The one exception is a host toolchain that cannot judge the kernel (no clang, a clang +that does not know the arch, or one that cannot run its linker), which is a missing +toolchain, not a kernel defect. +""" + +from __future__ import annotations + +from collections.abc import Collection, Mapping + +import numpy as np +import pytest + +from rocke.core.ir import KernelDef +from rocke.tests.ir_validity import arbiter_verdict + +__all__ = ["current_device_arch", "run_on_device", "supported_device_arch"] + +# Host array dtype -> rocke signature dtype name. +_SIGNATURE_DTYPE = { + np.float64: "f64", + np.float32: "f32", + np.float16: "f16", + np.int32: "i32", + np.int16: "i16", + np.int8: "i8", +} + + +def current_device_arch() -> str: + """The arch of the GPU this process launches on: HIP's current device, the one + rocke's runtime keeps and launches kernels on. Skips when there is no GPU.""" + hip = pytest.importorskip("rocke.runtime.hip_module") + try: + arch = hip.get_device_arch(hip.current_device()) + except Exception as exc: # noqa: BLE001 - any HIP failure means "no GPU here" + pytest.skip(f"no GPU available: {exc}") + if not arch: + pytest.skip("no GPU available") + return arch + + +def supported_device_arch(kernel: KernelDef, arches: Collection[str]) -> str: + """The current GPU's arch when it is one of `arches` (the arches whose hardware + supports `kernel`'s data types) and the IR the active engine lowers `kernel` to + compiles and links for it; otherwise skip, naming why, or fail on a compiler + rejection.""" + arch = current_device_arch() + if arch not in arches: + pytest.skip( + f"{kernel.name} needs one of {sorted(arches)}; the current GPU is {arch}" + ) + from rocke.core.lower_llvm import lower_kernel_to_llvm + + text = lower_kernel_to_llvm(kernel, arch=arch) + verdict, reason = arbiter_verdict(text, arch) + if verdict == "cannot_judge": + pytest.skip(reason) + assert verdict == "accepted", f"LLVM rejected {kernel.name} for {arch}: {reason}" + return arch + + +def run_on_device( + kernel: KernelDef, + buffers: Mapping[str, np.ndarray], + threads: int, + arches: Collection[str], +) -> None: + """Launch `kernel` on the current GPU when it is one of `arches`, as one block of + `threads` lanes with one pointer parameter per entry of `buffers` (in order), then copy + every buffer back into its host array.""" + arch = supported_device_arch(kernel, arches) + from rocke.helpers.compile import compile_kernel + from rocke.helpers.spec import sig_param + from rocke.runtime.hip_module import Runtime + from rocke.runtime.host_buffers import as_u8_buffer + from rocke.runtime.launcher import ( + DeviceMem, + KernelLauncher, + LaunchConfig, + synchronize_and_release, + ) + + artifact = compile_kernel(kernel, arch=arch) + signature = [ + sig_param(name, _SIGNATURE_DTYPE[host.dtype.type]) + for name, host in buffers.items() + ] + launcher = KernelLauncher( + hsaco=artifact.hsaco, kernel_name=artifact.kernel_name, signature=signature + ) + rt = Runtime() + device = {name: DeviceMem(host.nbytes) for name, host in buffers.items()} + for name, host in buffers.items(): + rt.memcpy_h2d(device[name].ptr(), as_u8_buffer(host), host.nbytes) + launcher(device, config=LaunchConfig(grid=(1, 1, 1), block=(threads, 1, 1))) + synchronize_and_release() + for name, host in buffers.items(): + rt.memcpy_d2h(as_u8_buffer(host), device[name].ptr(), host.nbytes) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/test_gpu.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/test_gpu.py new file mode 100644 index 000000000000..25ca64608611 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/test_gpu.py @@ -0,0 +1,47 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""The GPU test helper checks the arch of the GPU it launches on. + +rocke's runtime launches on HIP's current device, so the arch check must read that +device and not a fixed device 0. With no GPU the helper skips, which is its own +contract. +""" + +from __future__ import annotations + +import pytest + +from rocke.tests.gpu import current_device_arch + +hip = pytest.importorskip("rocke.runtime.hip_module") + + +@pytest.fixture(autouse=True) +def _gpu_present() -> None: + """Skip without a HIP runtime or GPU: the HIP library loads on first use, so the + import above succeeds even where the device queries below would raise.""" + current_device_arch() + + +def _select_device(device: int) -> None: + hip._check(hip._hipSetDevice(device), "hipSetDevice") + + +def test_reports_the_only_devices_arch(): + if hip.get_device_count() != 1: + pytest.skip("needs exactly one GPU") + assert current_device_arch() == hip.get_device_arch(0) + + +def test_follows_a_change_of_current_device(): + if hip.get_device_count() < 2: + pytest.skip("needs two GPUs") + previous = hip.current_device() + try: + _select_device(1) + assert current_device_arch() == hip.get_device_arch(1) + _select_device(0) + assert current_device_arch() == hip.get_device_arch(0) + finally: + _select_device(previous) 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 356cb5b1de73..615def4e59f7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt @@ -188,6 +188,22 @@ if(BUILD_TESTING) rocke_f64_lowering) endif() + # Pins every opcode's value, name and purity: rocke_opcode_t is C ABI. + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/core/test_opcode_values.cpp") + add_executable(rocke_opcode_values core/test_opcode_values.cpp) + target_link_libraries(rocke_opcode_values PRIVATE rocke_core) + add_test(NAME rocke_opcode_values COMMAND rocke_opcode_values) + set_tests_properties(rocke_opcode_values + PROPERTIES LABELS "ckc;unit;host") + rocke_cxx_test_installed_name(_rocke_installed_name rocke_opcode_values) + set_target_properties(rocke_opcode_values + PROPERTIES OUTPUT_NAME "${_rocke_installed_name}") + install(TARGETS rocke_opcode_values + RUNTIME DESTINATION ${ROCKE_CXX_TEST_INSTALL_DIR}) + set_property(GLOBAL APPEND PROPERTY ROCKE_INSTALL_TEST_TARGETS + rocke_opcode_values) + endif() + # Holds the ISA-layer and lower_llvm-layer flavor->datalayout tables together. if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/core/flavor_datalayout_ssot.cpp") add_executable(rocke_flavor_datalayout_ssot core/flavor_datalayout_ssot.cpp) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_f64_conversions.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_f64_conversions.py new file mode 100644 index 000000000000..1f4d3d0feb1f --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_f64_conversions.py @@ -0,0 +1,204 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""The explicit f32 <-> f64 conversions, and f64 arithmetic on IEEE edge values. + +``cvt_f32_to_f64`` (exact) and ``cvt_f64_to_f32`` (round to nearest even) are the only +conversions between the two types: ``cast_to_f32`` / ``cast_f32_to`` stay the f16 / bf16 +hub and refuse f64. Lowering goes through the active engine, so a strict both-engines run +(``ROCKE_BACKEND=both ROCKE_CPP_STRICT=1``) also proves the two engines agree. The numeric +test runs only on a visible GPU whose arch it names, and skips otherwise. +""" + +import numpy as np +import pytest +from rocke.core.ir import F32, F64, IRBuilder, PtrType +from rocke.core.isa.backend import wired_arches +from rocke.core.lower_hip import lower_kernel_to_hip +from rocke.core.lower_llvm import lower_kernel_to_llvm +from rocke.tests.gpu import run_on_device +from rocke.tests.ir_validity import assert_target_accepts + + +def _round_trip_kernel(): + """NARROW[i] = f32(X[i]), WIDE[i] = f64(NARROW[i]).""" + b = IRBuilder("f64_round_trip") + x = b.param("X", PtrType(F64, "global")) + narrow = b.param("NARROW", PtrType(F32, "global")) + wide = b.param("WIDE", PtrType(F64, "global")) + tid = b.thread_id_x() + value = b.global_load(x, tid, F64, align=8) + narrowed = b.cvt_f64_to_f32(value) + widened = b.cvt_f32_to_f64(narrowed) + b.global_store(narrow, tid, narrowed, align=4) + b.global_store(wide, tid, widened, align=8) + return b, value, narrowed, widened + + +def test_each_conversion_lowers_to_one_llvm_instruction(): + b, value, narrowed, widened = _round_trip_kernel() + text = lower_kernel_to_llvm(b.kernel, arch="gfx90a") + lines = [line.strip() for line in text.splitlines()] + assert lines.count(f"{narrowed.name} = fptrunc double {value.name} to float") == 1 + assert lines.count(f"{widened.name} = fpext float {narrowed.name} to double") == 1 + assert sum("fptrunc" in line or "fpext" in line for line in lines) == 2 + assert_target_accepts(text, "gfx90a") + + +def test_each_conversion_lowers_to_one_hip_cast(): + b, value, narrowed, widened = _round_trip_kernel() + hip = lower_kernel_to_hip(b.kernel, arch="gfx90a") + lines = [line.strip() for line in hip.splitlines()] + value, narrowed, widened = ( + v.name.removeprefix("%") for v in (value, narrowed, widened) + ) + assert f"float {narrowed} = (float){value};" in lines + assert f"double {widened} = (double){narrowed};" in lines + + +def test_conversions_take_only_their_source_type(): + b = IRBuilder("f64_conversion_inputs") + x = b.param("X", F64) + y = b.param("Y", F32) + emitted = len(b.kernel.body.ops) + with pytest.raises(ValueError, match="^cvt_f32_to_f64 expects f32 input, got f64$"): + b.cvt_f32_to_f64(x) + with pytest.raises(ValueError, match="^cvt_f64_to_f32 expects f64 input, got f32$"): + b.cvt_f64_to_f32(y) + assert len(b.kernel.body.ops) == emitted + + +def test_the_f16_bf16_hub_casts_refuse_f64(): + # cast_to_f32 / cast_f32_to (and the helpers built on them) never convert a double; + # f64 goes through the explicit conversions above. + b = IRBuilder("f64_hub_casts") + x = b.param("X", F64) + y = b.param("Y", F32) + with pytest.raises(ValueError, match="^cast_to_f32 unsupported from f64$"): + b.cast_to_f32(x) + with pytest.raises(ValueError, match="^cast_f32_to unsupported to f64$"): + b.cast_f32_to(y, F64) + + +# ---- on GPU (only where the visible GPU supports these types) -------------------- + +# Every arch rocke lowers for executes f64 arithmetic. +_EVERY_ARCH = wired_arches() + +# The largest finite f32, and half an f32 ulp at that magnitude. +_F32_MAX = float(np.finfo(np.float32).max) +_F32_MAX_HALF_ULP = 2.0**103 + +# Values exactly halfway between two neighbouring f32 values. A tie is the only input on +# which round-to-nearest-even differs from round-half-away-from-zero, so these are what +# prove that narrowing picks the neighbour with the even significand. +_F32_TIES = [ + # Between 1 and 1 + 2**-23: rounds down to the even 1. + 1 + 2**-24, + # Between 1 + 2**-23 and 1 + 2**-22: rounds up to the even 1 + 2**-22. + 1 + 3 * 2**-24, + -(1 + 2**-24), + -(1 + 3 * 2**-24), + # Between 0 and the smallest f32 subnormal 2**-149: rounds to the even 0. + 2**-150, + # Between the subnormals 2**-149 and 2**-148: rounds to the even 2**-148. + 3 * 2**-150, + -3 * 2**-150, + # The overflow threshold: the tie rounds to inf, one f64 ulp below it to the f32 max. + _F32_MAX + _F32_MAX_HALF_ULP, + float(np.nextafter(_F32_MAX + _F32_MAX_HALF_ULP, 0.0)), +] + +# 62 ordinary values (most not representable in f32, so narrowing must round), the f32 +# ties above, then the IEEE edge values: a value beyond the f32 range, an f64 subnormal, +# a value that narrows to an f32 subnormal, both zeros, both infinities and a NaN. +_VALUES = np.array( + [0.1 * i - 3.05 for i in range(62)] + + _F32_TIES + + [1e300, 5e-324, -1e-40, 0.0, -0.0, np.inf, -np.inf, np.nan], + dtype=np.float64, +) + + +def _edge_value_kernel(): + """Elementwise copy, neg, abs, add, sub and mul of X and Y, and both conversions.""" + b = IRBuilder("f64_edge_values") + params = { + name: b.param(name, PtrType(dtype, "global")) + for name, dtype in ( + ("X", F64), + ("Y", F64), + ("XF32", F32), + ("COPY", F64), + ("NEG", F64), + ("ABS", F64), + ("ADD", F64), + ("SUB", F64), + ("MUL", F64), + ("NARROW", F32), + ("WIDE", F64), + ) + } + tid = b.thread_id_x() + x = b.global_load(params["X"], tid, F64, align=8) + y = b.global_load(params["Y"], tid, F64, align=8) + x_f32 = b.global_load(params["XF32"], tid, F32, align=4) + results = { + "COPY": x, + "NEG": b.fneg(x), + # |x| as the elementwise family computes it. + "ABS": b.fmax(x, b.fneg(x)), + "ADD": b.fadd(x, y), + "SUB": b.fsub(x, y), + "MUL": b.fmul(x, y), + "NARROW": b.cvt_f64_to_f32(x), + "WIDE": b.cvt_f32_to_f64(x_f32), + } + for name, value in results.items(): + b.global_store(params[name], tid, value, align=4 if name == "NARROW" else 8) + return b.kernel + + +def _assert_ieee_identical( + actual, expected, *, nan_bits_exact=False, zero_sign_exact=True +): + """Bit-exact, except where IEEE 754 leaves the result's bits open: the sign and + payload of a NaN an arithmetic op produces, and the sign of the zero maxNum returns + for operands that compare equal.""" + bits = np.dtype(f"u{actual.itemsize}") + is_nan = np.isnan(expected) + np.testing.assert_array_equal(np.isnan(actual), is_nan) + exact = ~is_nan if not nan_bits_exact else np.ones_like(is_nan) + if not zero_sign_exact: + exact &= expected != 0 + np.testing.assert_array_equal(actual[expected == 0], 0) + np.testing.assert_array_equal(actual.view(bits)[exact], expected.view(bits)[exact]) + + +def test_f64_arithmetic_and_conversions_are_exact_on_edge_values_on_gpu(): + x = _VALUES + y = x[::-1].copy() + # numpy warns about the overflow of 1e300 to f32; that overflow is a wanted case. + with np.errstate(all="ignore"): + x_f32 = x.astype(np.float32) + expected = { + "COPY": x, + "NEG": -x, + "ABS": np.abs(x), + "ADD": x + y, + "SUB": x - y, + "MUL": x * y, + "NARROW": x_f32, + "WIDE": x_f32.astype(np.float64), + } + buffers = {"X": x.copy(), "Y": y, "XF32": x_f32.copy()} + buffers.update({name: np.zeros_like(want) for name, want in expected.items()}) + run_on_device(_edge_value_kernel(), buffers, len(x), _EVERY_ARCH) + + for name, want in expected.items(): + _assert_ieee_identical( + buffers[name], + want, + # A copy and a sign flip move bits, so even a NaN comes back exactly. + nan_bits_exact=name in ("COPY", "NEG"), + zero_sign_exact=name != "ABS", + ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_f64_lowering.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_f64_lowering.cpp index c685a7cf37a6..7d510fb0d56d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_f64_lowering.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_f64_lowering.cpp @@ -3,9 +3,13 @@ /* * test_f64_lowering.cpp -- C++ engine coverage for the f64 op boundary. * Mirrors TestHipF64Rejection / TestF64VectorArith in - * tests/core/test_amdgpu_lowering_ops.py. + * tests/core/test_amdgpu_lowering_ops.py, and the C++ builder side of + * tests/core/test_f64_conversions.py (the explicit f32 <-> f64 conversions) + * and tests/core/test_zero_vec_and_lds_element_types.py (zero_vec of f64 and + * i16, LDS loads of i16). */ +#include "rocke/error.hpp" #include "rocke/ir.h" #include "rocke/lower_hip.h" #include "rocke/lower_llvm.h" @@ -166,10 +170,97 @@ static int test_vector_binops_use_float_ops() return 0; } +/* out[0] = f64(f32(out[0])): one conversion each way. */ +static int test_conversions_lower_to_one_instruction_each() +{ + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "f64_conversions") == ROCKE_OK); + rocke_value_t* out = rocke_b_param(&b, "out", rocke_ptr_type(&b, rocke_f64(), "global"), NULL); + rocke_value_t* idx = rocke_b_const_i32(&b, 0); + rocke_value_t* narrowed + = rocke_b_cvt_f64_to_f32(&b, rocke_b_global_load(&b, out, idx, rocke_f64(), 8)); + rocke_value_t* widened = rocke_b_cvt_f32_to_f64(&b, narrowed); + rocke_b_global_store(&b, out, idx, widened, 8); + rocke_b_ret(&b); + CHECK(rocke_ir_builder_ok(&b)); + + char* llvm_text = NULL; + CHECK(rocke_lower_kernel_to_llvm(b.kernel, ROCKE_LLVM_FLAVOR_LLVM20, "gfx942", &llvm_text) + == ROCKE_OK); + char expected[96]; + std::snprintf(expected, sizeof(expected), "%s = fptrunc double ", narrowed->name); + CHECK(std::strstr(llvm_text, expected) != NULL); + std::snprintf( + expected, sizeof(expected), "%s = fpext float %s to double", widened->name, narrowed->name); + CHECK(std::strstr(llvm_text, expected) != NULL); + std::free(llvm_text); + + 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); + CHECK(std::strstr(hip.data, " = (float)") != NULL); + CHECK(std::strstr(hip.data, " = (double)") != NULL); + rocke_strbuf_free(&hip); + rocke_ir_builder_free(&b); + return 0; +} + +/* Each conversion takes exactly its source type, with the Python builder's message, + * and emits no op when it refuses. */ +static int test_conversions_reject_the_wrong_source_type() +{ + const struct + { + bool from_f32; + const char* message; + } cases[] = {{true, "cvt_f32_to_f64 expects f32 input, got f64"}, + {false, "cvt_f64_to_f32 expects f64 input, got f32"}}; + for(const auto& c : cases) + { + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "f64_conversion_reject") == ROCKE_OK); + rocke_value_t* wrong = rocke_b_param(&b, "x", c.from_f32 ? rocke_f64() : rocke_f32(), NULL); + int before = b.kernel->body->num_ops; + bool rejected = false; + try + { + c.from_f32 ? rocke_b_cvt_f32_to_f64(&b, wrong) : rocke_b_cvt_f64_to_f32(&b, wrong); + } + catch(const ckc::Error& error) + { + rejected = error.code() == ROCKE_ERR_VALUE && std::strcmp(error.what(), c.message) == 0; + } + CHECK(rejected); + CHECK(b.kernel->body->num_ops == before); + rocke_ir_builder_free(&b); + } + return 0; +} + +static int test_zero_vec_and_lds_load_admit_f64_and_i16() +{ + rocke_ir_builder_t b; + CHECK(rocke_ir_builder_init(&b, "zero_vec_and_lds") == ROCKE_OK); + CHECK(rocke_b_zero_vec(&b, rocke_f64(), 2) != NULL); + CHECK(rocke_b_zero_vec(&b, rocke_i16(), 4) != NULL); + const int shape[] = {256}; + rocke_value_t* smem = rocke_b_smem_alloc(&b, rocke_i16(), shape, 1, "halves"); + rocke_value_t* index = rocke_b_thread_id_x(&b); + CHECK(rocke_b_smem_load_vN(&b, smem, &index, 1, rocke_i16(), 4) != NULL); + rocke_b_ret(&b); + CHECK(rocke_ir_builder_ok(&b)); + rocke_ir_builder_free(&b); + return 0; +} + int main() { if(test_unsupported_ops_rejected() || test_rcp_divides_in_double() - || test_vector_binops_use_float_ops()) + || test_vector_binops_use_float_ops() || test_conversions_lower_to_one_instruction_each() + || test_conversions_reject_the_wrong_source_type() + || test_zero_vec_and_lds_load_admit_f64_and_i16()) { std::fprintf(stderr, "f64 lowering contract failed\n"); return 1; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_opcode_values.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_opcode_values.cpp new file mode 100644 index 000000000000..eb3a692a27ab --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_opcode_values.cpp @@ -0,0 +1,334 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* + * test_opcode_values.cpp -- pins the numeric value of every opcode. + * + * rocke_opcode_t is part of the C ABI: a caller compiled against an older + * rocke/ir.h passes and receives opcodes as plain integers. A new opcode must + * therefore be appended at the end of the enum; inserting one anywhere else + * renumbers every opcode after it and silently changes what an old caller's + * integer means (see "when to bump" in rocke/abi.h). + * + * The table below lists every opcode with its value, its dotted name and its + * purity as they stood before any opcode was appended. It must never change: a + * failure here means an opcode was inserted, removed or reordered. Opcodes + * appended later are not listed; the second check covers them by requiring + * that every opcode, old or new, has a name that maps back to it. + */ + +#include "rocke/ir.h" +#include +#include + +namespace +{ + +struct PinnedOpcode +{ + rocke_opcode_t op; + int value; + const char* name; + bool pure; +}; + +constexpr PinnedOpcode kPinned[] = { + {ROCKE_OP_ARITH_CONSTANT, 1, "arith.constant", true}, + {ROCKE_OP_ARITH_CONSTANT_VEC, 2, "arith.constant_vec", true}, + {ROCKE_OP_ARITH_ADD, 3, "arith.add", true}, + {ROCKE_OP_ARITH_SUB, 4, "arith.sub", true}, + {ROCKE_OP_ARITH_MUL, 5, "arith.mul", true}, + {ROCKE_OP_ARITH_DIV, 6, "arith.div", true}, + {ROCKE_OP_ARITH_MOD, 7, "arith.mod", true}, + {ROCKE_OP_ARITH_FADD, 8, "arith.fadd", true}, + {ROCKE_OP_ARITH_FSUB, 9, "arith.fsub", true}, + {ROCKE_OP_ARITH_FMUL, 10, "arith.fmul", true}, + {ROCKE_OP_ARITH_FDIV, 11, "arith.fdiv", true}, + {ROCKE_OP_ARITH_FNEG, 12, "arith.fneg", true}, + {ROCKE_OP_ARITH_FABS, 13, "arith.fabs", false}, + {ROCKE_OP_ARITH_FMA, 14, "arith.fma", false}, + {ROCKE_OP_ARITH_FMAX3, 15, "arith.fmax3", false}, + {ROCKE_OP_ARITH_FMIN3, 16, "arith.fmin3", false}, + {ROCKE_OP_ARITH_CMP, 17, "arith.cmp", true}, + {ROCKE_OP_ARITH_FCMP, 18, "arith.fcmp", true}, + {ROCKE_OP_ARITH_FMAX, 19, "arith.fmax", true}, + {ROCKE_OP_ARITH_FMIN, 20, "arith.fmin", true}, + {ROCKE_OP_ARITH_AND, 21, "arith.and", true}, + {ROCKE_OP_ARITH_OR, 22, "arith.or", true}, + {ROCKE_OP_ARITH_NOT, 23, "arith.not", true}, + {ROCKE_OP_ARITH_SMAX, 24, "arith.smax", true}, + {ROCKE_OP_ARITH_SMIN, 25, "arith.smin", true}, + {ROCKE_OP_ARITH_XOR, 26, "arith.xor", true}, + {ROCKE_OP_ARITH_SHL, 27, "arith.shl", true}, + {ROCKE_OP_ARITH_LSHR, 28, "arith.lshr", true}, + {ROCKE_OP_ARITH_UMUL_HI_I32, 29, "arith.umul_hi_i32", true}, + {ROCKE_OP_ARITH_ZEXT, 30, "arith.zext", true}, + {ROCKE_OP_ARITH_SEXT, 31, "arith.sext", true}, + {ROCKE_OP_ARITH_TRUNC, 32, "arith.trunc", true}, + {ROCKE_OP_ARITH_SELECT, 33, "arith.select", true}, + {ROCKE_OP_ARITH_BITCAST, 34, "arith.bitcast", true}, + {ROCKE_OP_ARITH_TRUNC_F32_TO_F16, 35, "arith.trunc_f32_to_f16", true}, + {ROCKE_OP_ARITH_TRUNC_F32_TO_BF16, 36, "arith.trunc_f32_to_bf16", true}, + {ROCKE_OP_ARITH_RINT_F32, 37, "arith.rint_f32", true}, + {ROCKE_OP_ARITH_CAST_TO_F32, 38, "arith.cast_to_f32", true}, + {ROCKE_OP_ARITH_CAST_F32_TO, 39, "arith.cast_f32_to", true}, + {ROCKE_OP_ARITH_SITOFP_F32, 40, "arith.sitofp_f32", true}, + {ROCKE_OP_ARITH_CVT_FP8_TO_F32, 41, "arith.cvt_fp8_to_f32", true}, + {ROCKE_OP_ARITH_CVT_BF8_TO_F32, 42, "arith.cvt_bf8_to_f32", true}, + {ROCKE_OP_ARITH_CVT_PK_F32_FP8X4, 43, "arith.cvt_pk_f32_fp8x4", true}, + {ROCKE_OP_ARITH_CVT_PK_F32_BF8X4, 44, "arith.cvt_pk_f32_bf8x4", true}, + {ROCKE_OP_ARITH_CVT_PK_FP8_F32X4, 45, "arith.cvt_pk_fp8_f32x4", true}, + {ROCKE_OP_ARITH_CVT_PK_BF8_F32X4, 46, "arith.cvt_pk_bf8_f32x4", true}, + {ROCKE_OP_ARITH_CVT_PK_I8_F32X4, 47, "arith.cvt_pk_i8_f32x4", false}, + {ROCKE_OP_ARITH_CVT_F32_TO_FP8, 48, "arith.cvt_f32_to_fp8", true}, + {ROCKE_OP_ARITH_CVT_F32_TO_BF8, 49, "arith.cvt_f32_to_bf8", true}, + {ROCKE_OP_ARITH_CVT_F32_TO_I8_SAT, 50, "arith.cvt_f32_to_i8_sat", true}, + {ROCKE_OP_ARITH_CVT_SCALEF32_PK_F32_FP8, 51, "arith.cvt_scalef32_pk_f32_fp8", true}, + {ROCKE_OP_ARITH_CVT_SCALEF32_PK_F32_BF8, 52, "arith.cvt_scalef32_pk_f32_bf8", true}, + {ROCKE_OP_ARITH_CVT_SCALEF32_PK_FP8_F32, 53, "arith.cvt_scalef32_pk_fp8_f32", true}, + {ROCKE_OP_ARITH_CVT_SCALEF32_PK_BF8_F32, 54, "arith.cvt_scalef32_pk_bf8_f32", true}, + {ROCKE_OP_MATH_EXP2, 55, "math.exp2", true}, + {ROCKE_OP_MATH_EXP2_FAST, 56, "math.exp2_fast", true}, + {ROCKE_OP_MATH_LOG2, 57, "math.log2", true}, + {ROCKE_OP_MATH_RCP, 58, "math.rcp", true}, + {ROCKE_OP_MATH_RCP_FAST, 59, "math.rcp_fast", true}, + {ROCKE_OP_MATH_SQRT, 60, "math.sqrt", true}, + {ROCKE_OP_MATH_RSQRT, 61, "math.rsqrt", true}, + {ROCKE_OP_MATH_TANH, 62, "math.tanh", true}, + {ROCKE_OP_GPU_THREAD_ID, 63, "gpu.thread_id", true}, + {ROCKE_OP_GPU_BLOCK_ID, 64, "gpu.block_id", true}, + {ROCKE_OP_GPU_CLUSTER_ID, 65, "gpu.cluster_id", true}, + {ROCKE_OP_GPU_CLUSTER_WORKGROUP_ID, 66, "gpu.cluster_workgroup_id", true}, + {ROCKE_OP_GPU_CLUSTER_WORKGROUP_MAX_ID, 67, "gpu.cluster_workgroup_max_id", true}, + {ROCKE_OP_GPU_CLUSTER_WORKGROUP_FLAT_ID, 68, "gpu.cluster_workgroup_flat_id", true}, + {ROCKE_OP_GPU_CLUSTER_WORKGROUP_MAX_FLAT_ID, 69, "gpu.cluster_workgroup_max_flat_id", true}, + {ROCKE_OP_MEMREF_GLOBAL_LOAD, 70, "memref.global_load", false}, + {ROCKE_OP_MEMREF_GLOBAL_LOAD_TYPED, 71, "memref.global_load_typed", false}, + {ROCKE_OP_MEMREF_GLOBAL_LOAD_VN, 72, "memref.global_load_vN", false}, + {ROCKE_OP_MEMREF_GLOBAL_STORE, 73, "memref.global_store", false}, + {ROCKE_OP_MEMREF_GLOBAL_STORE_TYPED, 74, "memref.global_store_typed", false}, + {ROCKE_OP_MEMREF_GLOBAL_STORE_VN, 75, "memref.global_store_vN", false}, + {ROCKE_OP_MEMREF_GLOBAL_ATOMIC_ADD, 76, "memref.global_atomic_add", false}, + {ROCKE_OP_MEMREF_GLOBAL_ATOMIC_ADD_F32, 77, "memref.global_atomic_add_f32", false}, + {ROCKE_OP_MEMREF_GLOBAL_ATOMIC_ADD_PK_BF16, 78, "memref.global_atomic_add_pk_bf16", false}, + {ROCKE_OP_MEMREF_GLOBAL_ATOMIC_ADD_PK_F16, 79, "memref.global_atomic_add_pk_f16", false}, + {ROCKE_OP_MEMREF_COOPERATIVE_GLOBAL_STORE, 80, "memref.cooperative_global_store", false}, + {ROCKE_OP_VECTOR_ADD, 81, "vector.add", true}, + {ROCKE_OP_VECTOR_SUB, 82, "vector.sub", false}, + {ROCKE_OP_VECTOR_MUL, 83, "vector.mul", true}, + {ROCKE_OP_VECTOR_AND, 84, "vector.and", true}, + {ROCKE_OP_VECTOR_OR, 85, "vector.or", true}, + {ROCKE_OP_VECTOR_SHL, 86, "vector.shl", true}, + {ROCKE_OP_VECTOR_LSHR, 87, "vector.lshr", true}, + {ROCKE_OP_VECTOR_SMAX, 88, "vector.smax", true}, + {ROCKE_OP_VECTOR_SMIN, 89, "vector.smin", true}, + {ROCKE_OP_VECTOR_MAX, 90, "vector.max", true}, + {ROCKE_OP_VECTOR_FMA, 91, "vector.fma", false}, + {ROCKE_OP_VECTOR_SUM, 92, "vector.sum", true}, + {ROCKE_OP_VECTOR_REDUCE_MAX, 93, "vector.reduce_max", true}, + {ROCKE_OP_VECTOR_SPLAT, 94, "vector.splat", true}, + {ROCKE_OP_VECTOR_SELECT, 95, "vector.select", true}, + {ROCKE_OP_VECTOR_CMP, 96, "vector.cmp", true}, + {ROCKE_OP_VECTOR_TRUNC, 97, "vector.trunc", true}, + {ROCKE_OP_VECTOR_SEXT, 98, "vector.sext", true}, + {ROCKE_OP_VECTOR_TRUNC_F32_TO_F16, 99, "vector.trunc_f32_to_f16", true}, + {ROCKE_OP_VECTOR_TRUNC_F32_TO, 100, "vector.trunc_f32_to", true}, + {ROCKE_OP_VECTOR_BITCAST, 101, "vector.bitcast", true}, + {ROCKE_OP_VECTOR_EXTRACT, 102, "vector.extract", true}, + {ROCKE_OP_VECTOR_INSERT, 103, "vector.insert", true}, + {ROCKE_OP_VECTOR_PACK, 104, "vector.pack", true}, + {ROCKE_OP_VECTOR_CONCAT, 105, "vector.concat", true}, + {ROCKE_OP_TILE_SMEM_ALLOC, 106, "tile.smem_alloc", false}, + {ROCKE_OP_TILE_SMEM_STORE, 107, "tile.smem_store", false}, + {ROCKE_OP_TILE_SMEM_STORE_VN, 108, "tile.smem_store_vN", false}, + {ROCKE_OP_TILE_SMEM_STORE_VN_F32, 109, "tile.smem_store_vN_f32", false}, + {ROCKE_OP_TILE_SMEM_STORE_DISTRIBUTED, 110, "tile.smem_store_distributed", false}, + {ROCKE_OP_TILE_SMEM_LOAD_V4, 111, "tile.smem_load_v4", false}, + {ROCKE_OP_TILE_SMEM_LOAD_VN, 112, "tile.smem_load_vN", false}, + {ROCKE_OP_TILE_SMEM_LOAD_VN_F32, 113, "tile.smem_load_vN_f32", false}, + {ROCKE_OP_TILE_SMEM_ADDR_OF, 114, "tile.smem_addr_of", true}, + {ROCKE_OP_TILE_SMEM_PTR_ADD, 115, "tile.smem_ptr_add", true}, + {ROCKE_OP_TILE_LDS_ATOMIC_ADD, 116, "tile.lds_atomic_add", false}, + {ROCKE_OP_TILE_GLOBAL_PTR_ADD, 117, "tile.global_ptr_add", false}, + {ROCKE_OP_TILE_GLOBAL_LOAD_LDS, 118, "tile.global_load_lds", false}, + {ROCKE_OP_TILE_ASYNC_BUFFER_LOAD_LDS, 119, "tile.async_buffer_load_lds", false}, + {ROCKE_OP_TILE_ASYNC_BUFFER_LOAD_LDS_ADDR, 120, "tile.async_buffer_load_lds_addr", false}, + {ROCKE_OP_TILE_BUFFER_LOAD_LDS_ASYNC, 121, "tile.buffer_load_lds_async", false}, + {ROCKE_OP_TILE_GLOBAL_LOAD_ASYNC_TO_LDS, 122, "tile.global_load_async_to_lds", false}, + {ROCKE_OP_TILE_GLOBAL_STORE_ASYNC_FROM_LDS, 123, "tile.global_store_async_from_lds", false}, + {ROCKE_OP_TILE_GLOBAL_LOAD_TR16_B128, 124, "tile.global_load_tr16_b128", false}, + {ROCKE_OP_TILE_TENSOR_LOAD_TO_LDS, 125, "tile.tensor_load_to_lds", false}, + {ROCKE_OP_TILE_TENSOR_STORE_FROM_LDS, 126, "tile.tensor_store_from_lds", false}, + {ROCKE_OP_TILE_BUFFER_RSRC, 127, "tile.buffer_rsrc", false}, + {ROCKE_OP_TILE_BUFFER_LOAD_F16, 128, "tile.buffer_load_f16", false}, + {ROCKE_OP_TILE_BUFFER_LOAD_VN_F16, 129, "tile.buffer_load_vN_f16", false}, + {ROCKE_OP_TILE_BUFFER_LOAD_VN, 130, "tile.buffer_load_vN", false}, + {ROCKE_OP_TILE_BUFFER_STORE_F16, 131, "tile.buffer_store_f16", false}, + {ROCKE_OP_TILE_BUFFER_STORE_VN_F16, 132, "tile.buffer_store_vN_f16", false}, + {ROCKE_OP_TILE_BUFFER_LOAD_BF16, 133, "tile.buffer_load_bf16", false}, + {ROCKE_OP_TILE_BUFFER_LOAD_VN_BF16, 134, "tile.buffer_load_vN_bf16", false}, + {ROCKE_OP_TILE_BUFFER_STORE_BF16, 135, "tile.buffer_store_bf16", false}, + {ROCKE_OP_TILE_BUFFER_STORE_VN_BF16, 136, "tile.buffer_store_vN_bf16", false}, + {ROCKE_OP_TILE_BUFFER_STORE_F32, 137, "tile.buffer_store_f32", false}, + {ROCKE_OP_TILE_BUFFER_STORE_VN_F32, 138, "tile.buffer_store_vN_f32", false}, + {ROCKE_OP_TILE_MMA, 139, "tile.mma", false}, + {ROCKE_OP_TILE_REGISTER_P_FROM_QK_C, 140, "tile.register_p_from_qk_c", false}, + {ROCKE_OP_TILE_INLINE_ASM, 141, "tile.inline_asm", false}, + {ROCKE_OP_TILE_READFIRSTLANE, 142, "tile.readfirstlane", true}, + {ROCKE_OP_TILE_PIN_SGPR, 143, "tile.pin_sgpr", true}, + {ROCKE_OP_TILE_LANE_ID, 144, "tile.lane_id", true}, + {ROCKE_OP_TILE_WAVE_ALL, 145, "tile.wave_all", true}, + {ROCKE_OP_TILE_WAVE_ANY, 146, "tile.wave_any", true}, + {ROCKE_OP_TILE_WAVE_BALLOT, 147, "tile.wave_ballot", true}, + {ROCKE_OP_TILE_DS_BPERMUTE, 148, "tile.ds_bpermute", true}, + {ROCKE_OP_TILE_DS_BPERMUTE_B64, 149, "tile.ds_bpermute_b64", false}, + {ROCKE_OP_TILE_DS_SWIZZLE_XOR, 150, "tile.ds_swizzle_xor", true}, + {ROCKE_OP_TILE_DS_SWIZZLE, 151, "tile.ds_swizzle", true}, + {ROCKE_OP_TILE_MOV_DPP8, 152, "tile.mov_dpp8", true}, + {ROCKE_OP_TILE_QUAD_PERM, 153, "tile.quad_perm", true}, + {ROCKE_OP_TILE_DPP_XOR, 154, "tile.dpp_xor", true}, + {ROCKE_OP_TILE_WAVE_REDUCE, 155, "tile.wave_reduce", true}, + {ROCKE_OP_TILE_READLANE, 156, "tile.readlane", true}, + {ROCKE_OP_TILE_WRITELANE, 157, "tile.writelane", true}, + {ROCKE_OP_TILE_PERMLANE16, 158, "tile.permlane16", true}, + {ROCKE_OP_TILE_PERMLANE64, 159, "tile.permlane64", true}, + {ROCKE_OP_TILE_ALIGNBYTE, 160, "tile.alignbyte", true}, + {ROCKE_OP_TILE_S_WQM, 161, "tile.s_wqm", true}, + {ROCKE_OP_TILE_AV_LOAD_B128, 162, "tile.av_load_b128", true}, + {ROCKE_OP_TILE_AV_STORE_B128, 163, "tile.av_store_b128", false}, + {ROCKE_OP_TILE_S_ALLOC_VGPR, 164, "tile.s_alloc_vgpr", false}, + {ROCKE_OP_TILE_MOV_DPP, 165, "tile.mov_dpp", false}, + {ROCKE_OP_TILE_PERMLANE32_SWAP, 166, "tile.permlane32_swap", true}, + {ROCKE_OP_TILE_PERM_B32, 167, "tile.perm_b32", true}, + {ROCKE_OP_TILE_PERMLANEX16, 168, "tile.permlanex16", true}, + {ROCKE_OP_TILE_BYTE_PERM, 169, "tile.byte_perm", true}, + {ROCKE_OP_TILE_DS_READ_TR16_B64, 170, "tile.ds_read_tr16_b64", true}, + {ROCKE_OP_TILE_DS_READ_TR16_B128, 171, "tile.ds_read_tr16_b128", false}, + {ROCKE_OP_TILE_DS_READ_TR_B8, 172, "tile.ds_read_tr_b8", true}, + {ROCKE_OP_TILE_SYNC, 173, "tile.sync", false}, + {ROCKE_OP_TILE_SYNC_HALF_BLOCK, 174, "tile.sync_half_block", true}, + {ROCKE_OP_TILE_SYNC_LDS_ONLY, 175, "tile.sync_lds_only", false}, + {ROCKE_OP_TILE_S_BARRIER_BARE, 176, "tile.s_barrier_bare", false}, + {ROCKE_OP_TILE_S_WAITCNT, 177, "tile.s_waitcnt", false}, + {ROCKE_OP_TILE_S_WAIT_ASYNCCNT, 178, "tile.s_wait_asynccnt", false}, + {ROCKE_OP_TILE_S_WAIT_TENSORCNT, 179, "tile.s_wait_tensorcnt", false}, + {ROCKE_OP_TILE_S_BARRIER_SIGNAL, 180, "tile.s_barrier_signal", false}, + {ROCKE_OP_TILE_S_BARRIER_WAIT, 181, "tile.s_barrier_wait", false}, + {ROCKE_OP_TILE_S_BARRIER_INIT, 182, "tile.s_barrier_init", false}, + {ROCKE_OP_TILE_S_BARRIER_SIGNAL_VAR, 183, "tile.s_barrier_signal_var", false}, + {ROCKE_OP_TILE_S_BARRIER_JOIN, 184, "tile.s_barrier_join", false}, + {ROCKE_OP_TILE_S_WAKEUP_BARRIER, 185, "tile.s_wakeup_barrier", false}, + {ROCKE_OP_TILE_S_BARRIER_LEAVE, 186, "tile.s_barrier_leave", false}, + {ROCKE_OP_TILE_ASYNCMARK, 187, "tile.asyncmark", false}, + {ROCKE_OP_TILE_WAIT_ASYNCMARK, 188, "tile.wait_asyncmark", false}, + {ROCKE_OP_TILE_S_WAIT_EVENT, 189, "tile.s_wait_event", false}, + {ROCKE_OP_TILE_S_PREFETCH_INST, 190, "tile.s_prefetch_inst", false}, + {ROCKE_OP_TILE_S_SETREG, 191, "tile.s_setreg", false}, + {ROCKE_OP_TILE_S_PREFETCH_DATA, 192, "tile.s_prefetch_data", false}, + {ROCKE_OP_TILE_S_BUFFER_PREFETCH_DATA, 193, "tile.s_buffer_prefetch_data", false}, + {ROCKE_OP_TILE_GLOBAL_PREFETCH, 194, "tile.global_prefetch", false}, + {ROCKE_OP_TILE_FLAT_PREFETCH, 195, "tile.flat_prefetch", false}, + {ROCKE_OP_TILE_CLUSTER_BARRIER, 196, "tile.cluster_barrier", false}, + {ROCKE_OP_TILE_CLUSTER_LOAD, 197, "tile.cluster_load", false}, + {ROCKE_OP_TILE_CLUSTER_LOAD_ASYNC_TO_LDS, 198, "tile.cluster_load_async_to_lds", false}, + {ROCKE_OP_TILE_GLOBAL_PTR_TO_I64, 199, "tile.global_ptr_to_i64", true}, + {ROCKE_OP_TILE_S_SETPRIO, 200, "tile.s_setprio", false}, + {ROCKE_OP_TILE_IGLP_OPT, 201, "tile.iglp_opt", false}, + {ROCKE_OP_TILE_SCHED_BARRIER, 202, "tile.sched_barrier", false}, + {ROCKE_OP_TILE_SCHED_GROUP_BARRIER, 203, "tile.sched_group_barrier", false}, + {ROCKE_OP_TILE_EXEC_AND_SAVEEXEC, 204, "tile.exec_and_saveexec", false}, + {ROCKE_OP_TILE_EXEC_XOR, 205, "tile.exec_xor", false}, + {ROCKE_OP_TILE_EXEC_OR_SAVEEXEC, 206, "tile.exec_or_saveexec", false}, + {ROCKE_OP_TILE_EXEC_OR, 207, "tile.exec_or", false}, + {ROCKE_OP_SCF_FOR, 208, "scf.for", false}, + {ROCKE_OP_SCF_IF, 209, "scf.if", false}, + {ROCKE_OP_SCF_IF_ELSE, 210, "scf.if_else", false}, + {ROCKE_OP_SCF_YIELD, 211, "scf.yield", false}, + {ROCKE_OP_CF_RETURN, 212, "cf.return", false}, + {ROCKE_OP_MEMREF_FENCE, 213, "memref.fence", false}, + {ROCKE_OP_MEMREF_GLOBAL_FLAG_STORE, 214, "memref.global_flag_store", false}, + {ROCKE_OP_MEMREF_GLOBAL_FLAG_WAIT_EQ, 215, "memref.global_flag_wait_eq", false}, + {ROCKE_OP_VECTOR_UNDEF, 216, "vector.undef", false}, + {ROCKE_OP_TILE_BUFFER_LOAD_F16_D16, 217, "tile.buffer_load_f16_d16", false}, +}; + +/* Each pinned opcode keeps its value, its name and its purity. */ +int check_pinned_opcodes() +{ + int failures = 0; + for(const PinnedOpcode& p : kPinned) + { + if(static_cast(p.op) != p.value) + { + std::fprintf(stderr, + "opcode %s has value %d, but it must stay %d: append new opcodes at " + "the end of rocke_opcode_t instead of inserting them\n", + p.name, + static_cast(p.op), + p.value); + ++failures; + } + const char* name = rocke_opcode_name(p.op); + if(std::strcmp(name, p.name) != 0) + { + std::fprintf(stderr, + "opcode %d is named \"%s\" but should be \"%s\": the name table in " + "core_types.cpp is out of step with the enum\n", + p.value, + name, + p.name); + ++failures; + } + if(rocke_opcode_is_pure(p.op) != p.pure) + { + std::fprintf(stderr, + "opcode %s should be %s: the purity table in core_types.cpp is out of " + "step with the enum\n", + p.name, + p.pure ? "pure" : "impure"); + ++failures; + } + } + return failures; +} + +/* Every opcode, pinned or appended later, has a name that maps back to it. */ +int check_every_name_round_trips() +{ + int failures = 0; + for(int v = ROCKE_OP_INVALID + 1; v < ROCKE_OP__COUNT; ++v) + { + const rocke_opcode_t op = static_cast(v); + const char* name = rocke_opcode_name(op); + if(name == nullptr || name[0] == '\0') + { + std::fprintf(stderr, "opcode %d has no name in core_types.cpp\n", v); + ++failures; + } + else if(rocke_opcode_from_name(name) != op) + { + std::fprintf(stderr, + "opcode %d is named \"%s\", but that name maps to opcode %d\n", + v, + name, + static_cast(rocke_opcode_from_name(name))); + ++failures; + } + } + return failures; +} + +} // namespace + +int main() +{ + const int failures = check_pinned_opcodes() + check_every_name_round_trips(); + if(failures != 0) + { + std::fprintf(stderr, "%d opcode check(s) failed\n", failures); + return 1; + } + std::printf("opcode values: %zu pinned opcodes and all %d opcode names OK\n", + sizeof(kPinned) / sizeof(kPinned[0]), + ROCKE_OP__COUNT - 1); + return 0; +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_zero_vec_and_lds_element_types.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_zero_vec_and_lds_element_types.py new file mode 100644 index 000000000000..6a8e0f3e6455 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_zero_vec_and_lds_element_types.py @@ -0,0 +1,92 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""``zero_vec`` of f64 and i16, and LDS loads and stores of i16. + +``zero_vec`` is what ``fill_fragment`` builds on, so without f64 an f64 fragment could +not be zero-filled. Lowering goes through the active engine, so a strict both-engines run +(``ROCKE_BACKEND=both ROCKE_CPP_STRICT=1``) also proves the two engines agree. The numeric +test runs only on a visible GPU whose arch it names, and skips otherwise. +""" + +import numpy as np +import pytest +from rocke.core.ir import F64, I16, I64, TF32, IRBuilder, PtrType +from rocke.core.isa.backend import wired_arches +from rocke.core.lower_llvm import lower_kernel_to_llvm +from rocke.tests.gpu import run_on_device +from rocke.tests.ir_validity import assert_target_accepts + + +def test_zero_vec_accepts_f64_and_i16(): + b = IRBuilder("zero_vectors") + b.zero_vec(I16, 4) + b.zero_vec(F64, 2) + text = lower_kernel_to_llvm(b.kernel, arch="gfx90a") + assert "<4 x i16> zeroinitializer" in text + assert "<2 x double> zeroinitializer" in text + assert_target_accepts(text, "gfx90a") + with pytest.raises(ValueError, match="^zero_vec unsupported elem tf32$"): + IRBuilder("tf32_zero").zero_vec(TF32, 4) + + +def _i16_lds_round_trip(n: int) -> IRBuilder: + """DST[i*n : i*n+n] = SRC[i*n : i*n+n], staged through LDS as .""" + b = IRBuilder(f"lds_round_trip_i16_{n}") + src = b.param("SRC", PtrType(I16, "global")) + dst = b.param("DST", PtrType(I16, "global")) + tid = b.thread_id_x() + base = b.mul(tid, b.const_i32(n)) + smem = b.smem_alloc(I16, [64 * n]) + elems = [b.global_load(src, b.add(base, b.const_i32(j)), I16) for j in range(n)] + value = elems[0] if n == 1 else b.vec_pack(elems, I16) + b.smem_store_vN(smem, [base], value, n) + loaded = b.smem_load_vN(smem, base, dtype=I16, n=n) + for j in range(n): + b.global_store(dst, b.add(base, b.const_i32(j)), b.vec_extract(loaded, j)) + return b + + +@pytest.mark.parametrize("n", [1, 2, 4, 8]) +def test_lds_loads_and_stores_i16_at_natural_alignment(n): + text = lower_kernel_to_llvm(_i16_lds_round_trip(n).kernel, arch="gfx90a") + lds_lines = [line for line in text.splitlines() if "ptr addrspace(3)" in line] + loads = [line for line in lds_lines if " = load " in line] + stores = [line for line in lds_lines if line.strip().startswith("store ")] + assert loads and stores + for line in loads + stores: + assert "i16" in line and line.rstrip().endswith(f"align {2 * n}"), line + if n > 1: + assert any(f"<{n} x i16>" in line for line in loads) + assert_target_accepts(text, "gfx90a") + + +def test_lds_loads_outside_the_element_rules_still_raise(): + b = IRBuilder("lds_bad_loads") + smem = b.smem_alloc(I16, [64]) + idx = b.thread_id_x() + with pytest.raises( + ValueError, match="^unsupported vector width 3 for smem_load_vN" + ): + b.smem_load_vN(smem, idx, dtype=I16, n=3) + with pytest.raises( + ValueError, + match=( + r"^smem_load_vN supports f16 / bf16 / i16 / f32 / i32 / tf32 / f64 / " + r"fp8e4m3 / bf8e5m2 / i8, got i64$" + ), + ): + b.smem_load_vN(smem, idx, dtype=I64, n=1) + + +# Every arch rocke lowers for executes i16 LDS accesses. +_EVERY_ARCH = wired_arches() + + +def test_lds_round_trip_of_i16_on_gpu(): + n = 4 + # Span the whole i16 range, sign bit included, so a truncated or sign-extended + # lane shows up as a wrong value. + values = (np.arange(64 * n, dtype=np.int64) * 257 - 32768).astype(np.int16) + buffers = {"SRC": values.copy(), "DST": np.zeros_like(values)} + run_on_device(_i16_lds_round_trip(n).kernel, buffers, 64, _EVERY_ARCH) + np.testing.assert_array_equal(buffers["DST"], values) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/runtime/test_cluster_launch.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/runtime/test_cluster_launch.py index 0e51e18137ce..f6097fefc884 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/runtime/test_cluster_launch.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/runtime/test_cluster_launch.py @@ -77,7 +77,7 @@ def hip(): ), mock.patch.object( hip_module, "_hipStreamSynchronize", return_value=0 ), mock.patch.object( - hip_module, "_current_device", return_value=0 + hip_module, "current_device", return_value=0 ), mock.patch.object( hip_module, "device_supports_cluster_launch", return_value=True ): @@ -185,7 +185,7 @@ def test_bad_cluster_is_refused_before_hip(hip, grid, cluster, msg): def test_unsupported_device_is_refused(hip): with mock.patch.object( hip_module, "device_supports_cluster_launch", return_value=False - ), mock.patch.object(hip_module, "_current_device", return_value=3): + ), mock.patch.object(hip_module, "current_device", return_value=3): with pytest.raises(HipError, match="device 3 does not support cluster launch"): Runtime().launch( object(), (2, 1, 1), (64, 1, 1), b"", stream=_STREAM, cluster=(2, 1, 1) From fbc5e2f88a3b4dde30b8734cf3abf222008b5abd Mon Sep 17 00:00:00 2001 From: Chris Millette Date: Thu, 8 Oct 2026 18:08:19 -0400 Subject: [PATCH 04/11] feat(rocke): add the logical I4 type and name RDNA integer WMMA operands i8 and i4 I4 is a logical 4-bit integer type, exported next to I8. It is packed-only: its values only ever live packed in 32-bit words, so it has no scalar or vector form and no LLVM or HIP rendering. The 8- and 4-bit integer dtypes are i8 and i4. RDNA integer WMMA instructions say iu8/iu4 in their names, but signedness is an instruction operand, not a property of the data, so iu8 and iu4 are no longer dtypes: their dtype entries are removed and no alias replaces them. The gfx1151 and gfx11-generic integer WMMA rows now record i8 and i4 operands; their op ids are unchanged. A query for i8 or i4 on those targets finds those instructions; a query by iu8 or iu4 finds nothing. Both engines make both changes: the C++ engine gains rocke_i4(), its dtype table drops iu8/iu4, and its embedded catalog rows match the JSON ones. A C++ test pins the catalog lookups: i8/i4 find the ops, iu8/iu4 do not. The HIP lowerings refuse I4 the way the LLVM lowerings do. Every HIP type lookup goes through one function per engine (_hip_scalar_type in lower_hip.py, rocke_h_require_hip_scalar in the C++ lowerer), which refuses a type with no HIP spelling with NotImplementedError / ROCKE_ERR_NOTIMPL "no HIP mapping for type ", replacing a bare KeyError in Python and four differently worded C++ messages. The C++ HIP lowerer now records the reason on the builder's error channel, as its header documents. The IR builder of both engines enforces the packed-only rule: any op result, kernel parameter or LDS tile of i4 elements is refused with the TypeError ": i4 values travel packed in 32-bit words and have no scalar or vector form -- keep them in i32 storage and unpack explicitly". The check lives in one function per engine, _refuse_packed_only_value in core/ir.py and rocke_i_refuse_packed_only_value in the C++ builder, called for every op result and every parameter; zero_vec calls it before its own element list. The lowerers' scalar and LDS type tables carry a comment naming it as the reason they have no i4 entry, and their own refusal stays as a backstop. rocke.tests.engines lowers a kernel with each engine and requires the same IR, or the same error, whatever backend the run selects, for the LLVM and the HIP lowering; the rocke_engine binding gains lower_serialized_ir_hip for the second. test_packed_only.py pins the Python builder's refusal text for each builder that could make an i4 value, and uses rocke.tests.engines to check that both engines refuse an i4 parameter, assembled by hand past the builder, in both lowerings. dtype_support.md names the integer dtypes, describes the packed-only rule, and rescores i4. Co-Authored-By: Claude --- .../platform/cpp/bindings/rocke_engine.cpp | 102 +++++++++++++---- .../rocke/platform/cpp/core/arch/data.cpp | 16 +-- .../rocke/platform/cpp/core/dtypes.cpp | 76 +++++++++---- .../platform/cpp/core/ir/core_builder.cpp | 26 +++++ .../rocke/platform/cpp/core/ir/core_types.cpp | 4 + .../rocke/platform/cpp/core/ir/ir_arith.cpp | 1 + .../platform/cpp/core/lower_hip/core.cpp | 55 +++++---- .../cpp/core/lower_hip/lower_hip_arith.cpp | 8 +- .../cpp/core/lower_hip/lower_hip_mem.cpp | 11 +- .../cpp/core/lower_hip/lower_hip_mma.cpp | 31 +++--- .../platform/cpp/core/lower_llvm/core.cpp | 10 +- .../rocke/platform/cpp/include/rocke/ir.h | 2 + .../platform/cpp/include/rocke/ir_internal.h | 7 ++ .../cpp/include/rocke/lower_hip_internal.h | 18 ++- .../dsl_docs/reference/dtype_support.md | 22 +++- .../rocke/platform/python/rocke/__init__.py | 2 + .../platform/python/rocke/core/__init__.py | 2 + .../rocke/core/arch/data/arch_specs.json | 8 +- .../platform/python/rocke/core/dtypes.py | 11 +- .../rocke/platform/python/rocke/core/ir.py | 27 +++++ .../platform/python/rocke/core/lower_hip.py | 44 ++++++-- .../platform/python/rocke/core/lower_llvm.py | 6 + .../platform/python/rocke/tests/engines.py | 105 ++++++++++++++++++ .../platform/tests/core/test_arch_mma_ssot.py | 24 ++++ .../rocke/platform/tests/core/test_dtypes.cpp | 6 + .../rocke/platform/tests/core/test_dtypes.py | 18 ++- .../tests/core/test_mma_catalog_lookup.cpp | 29 ++++- .../platform/tests/core/test_packed_only.py | 72 ++++++++++++ .../platform/tests/core/test_storage.cpp | 3 +- 29 files changed, 602 insertions(+), 144 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/engines.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_packed_only.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 f804924d9d60..f9688a962779 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 @@ -61,6 +61,7 @@ extern "C" { #include "rocke/instance_tf32_mma_probe.h" #include "rocke/ir.h" #include "rocke/ir_serialize.h" +#include "rocke/lower_hip.h" #include "rocke/lower_llvm.h" #include "rocke/rocke_build_id.h" #include "rocke/verify.h" @@ -3974,17 +3975,40 @@ static std::string rocke_engine_flavor_list() return out; } -std::string lower_serialized_ir(const std::string& ir_text, - const std::string& arch, - const std::string& flavor) +/* Parse serialized ck.dsl.ir/v1 text into a kernel owned by `b`, which the + * caller has zero-initialised and frees. `entry` names the binding in the + * error messages. Throws std::runtime_error on any failure. */ +static const rocke_kernel_def_t* + parse_serialized_ir(const std::string& ir_text, const char* entry, rocke_ir_builder_t* b) { + const std::string prefix = std::string("rocke_engine.") + entry + ": "; if(ir_text.empty()) { - throw std::runtime_error("rocke_engine.lower_serialized_ir: empty IR input"); + throw std::runtime_error(prefix + "empty IR input"); } + if(rocke_ir_builder_init(b, entry) != ROCKE_OK) + { + throw std::runtime_error(prefix + "builder init failed (OOM)"); + } + rocke_kernel_def_t* kernel = nullptr; + const rocke_status_t st = rocke_ir_parse(ir_text.c_str(), b, &kernel); + if(st != ROCKE_OK || !kernel) + { + const char* m = rocke_ir_builder_error(b); + throw std::runtime_error(prefix + "parse failed (status " + std::to_string((int)st) + + "): " + ((m && *m) ? m : "unknown parse error")); + } + return kernel; +} + +std::string lower_serialized_ir(const std::string& ir_text, + const std::string& arch, + const std::string& flavor) +{ /* No empty-string fallback to gfx950: the target is required on both * engines or on neither. The core resolver's NULL check does not cover - * this -- an empty std::string is not NULL by the time it gets there. */ + * this -- an empty std::string is not NULL by the time it gets there. + * An empty IR text is refused by parse_serialized_ir. */ if(arch.empty()) { throw std::runtime_error("rocke_engine.lower_serialized_ir: lowering requires an " @@ -4013,22 +4037,7 @@ std::string lower_serialized_ir(const std::string& ir_text, rocke_ir_builder_t b = {}; std::unique_ptr builder( &b, rocke_ir_builder_free); - const auto init_status = rocke_ir_builder_init(&b, "lower_serialized_ir"); - if(init_status != ROCKE_OK) - { - throw std::runtime_error("rocke_engine.lower_serialized_ir: builder init failed (OOM)"); - } - - rocke_kernel_def_t* kernel = nullptr; - rocke_status_t st = rocke_ir_parse(ir_text.c_str(), &b, &kernel); - if(st != ROCKE_OK || !kernel) - { - const char* m = rocke_ir_builder_error(&b); - std::string msg = std::string("rocke_engine.lower_serialized_ir: parse failed (status ") - + std::to_string((int)st) - + "): " + ((m && *m) ? m : "unknown parse error"); - throw std::runtime_error(msg); - } + const rocke_kernel_def_t* kernel = parse_serialized_ir(ir_text, "lower_serialized_ir", &b); if(fl == ROCKE_LLVM_FLAVOR_AUTO) fl = resolve_python_llvm_flavor(); @@ -4036,7 +4045,8 @@ std::string lower_serialized_ir(const std::string& ir_text, char* out_ll = nullptr; char err[ROCKE_ERR_MSG_CAP]; err[0] = '\0'; - st = rocke_lower_kernel_to_llvm_ex(kernel, fl, a, &out_ll, err, sizeof(err)); + const rocke_status_t st + = rocke_lower_kernel_to_llvm_ex(kernel, fl, a, &out_ll, err, sizeof(err)); std::unique_ptr text(out_ll, std::free); if(st != ROCKE_OK || !out_ll) { @@ -4049,6 +4059,46 @@ std::string lower_serialized_ir(const std::string& ir_text, return std::string(text.get()); } +/* lower_serialized_ir_hip(ir_text, arch): the HIP-source twin of + * lower_serialized_ir -- parse, then rocke_lower_kernel_to_hip with its + * defaults (prologue included, launch bounds from the kernel). Lets a test + * compare the C++ HIP lowering of a Python-built kernel, and its refusals, + * with rocke.core.lower_hip.lower_kernel_to_hip. */ +std::string lower_serialized_ir_hip(const std::string& ir_text, const std::string& arch) +{ + /* The target is required, as for lower_serialized_ir. */ + if(arch.empty()) + { + throw std::runtime_error("rocke_engine.lower_serialized_ir_hip: lowering requires an " + "explicit gfx target; pass arch=... (there is no default)"); + } + const char* a = arch.c_str(); + rocke_ir_builder_t b = {}; + std::unique_ptr builder( + &b, rocke_ir_builder_free); + const rocke_kernel_def_t* kernel = parse_serialized_ir(ir_text, "lower_serialized_ir_hip", &b); + + rocke_strbuf_t text; + if(rocke_strbuf_init(&text, 0) != 0) + { + throw std::runtime_error("rocke_engine.lower_serialized_ir_hip: output allocation failed"); + } + std::unique_ptr owned_text(&text, + rocke_strbuf_free); + rocke_lower_hip_opts_t opts = {}; + opts.arch = a; + const rocke_status_t st = rocke_lower_kernel_to_hip(&b, kernel, &opts, &text); + if(st != ROCKE_OK) + { + const char* m = rocke_ir_builder_error(&b); + throw std::runtime_error(std::string("rocke_engine.lower_serialized_ir_hip: lower failed " + "for arch '") + + a + "' (status " + std::to_string((int)st) + + "): " + ((m && *m) ? m : "unknown lowering error")); + } + return std::string(rocke_strbuf_cstr(&text)); +} + } // namespace PYBIND11_MODULE(rocke_engine, m) @@ -4108,6 +4158,14 @@ PYBIND11_MODULE(rocke_engine, m) "the Python lowerer for the same serialized IR. flavor='' resolves " "the LLVM flavor automatically; a name from llvm_flavors() pins it."); + m.def("lower_serialized_ir_hip", + &lower_serialized_ir_hip, + py::arg("ir_text"), + py::arg("arch"), /* required: no default target */ + "Parse serialized ck.dsl.ir/v1 text and lower it to HIP source via the " + "C++ engine; the twin of rocke.core.lower_hip.lower_kernel_to_hip with " + "its default options."); + /* Exposes the engine's flavor ladder so a test can assert it still equals * rocke.core.lower_llvm.LLVM_FLAVORS. Without a way to read the C++ list * back, the two enumerations can drift into a flavor that lowers on one 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 4018f3547af5..3d81b1680dea 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 @@ -2228,8 +2228,8 @@ static const rocke_mma_op_t k_mma_gfx1151[] = { &lm_wmma_16x16x16_b, &lm_wmma_16x16x16_c}, {"wmma", - "iu8", - "iu8", + "i8", + "i8", "i32", 16, 16, @@ -2243,8 +2243,8 @@ static const rocke_mma_op_t k_mma_gfx1151[] = { &lm_wmma_iu8_b, &lm_wmma_iu8_c}, {"wmma", - "iu4", - "iu4", + "i4", + "i4", "i32", 16, 16, @@ -2326,8 +2326,8 @@ static const rocke_mma_op_t k_mma_gfx11_generic[] = { &lm_wmma_16x16x16_b, &lm_wmma_16x16x16_c}, {"wmma", - "iu8", - "iu8", + "i8", + "i8", "i32", 16, 16, @@ -2341,8 +2341,8 @@ static const rocke_mma_op_t k_mma_gfx11_generic[] = { &lm_wmma_iu8_b, &lm_wmma_iu8_c}, {"wmma", - "iu4", - "iu4", + "i4", + "i4", "i32", 16, 16, 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 d3035a346b12..06b38c2e96d0 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,17 +26,40 @@ 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"}, {"tf32", "tf32"}, - {"xf32", "tf32"}, {"f64", "fp64"}, {"double", "fp64"}, - {"fp64", "fp64"}, {"iu8", "iu8"}, {"iu4", "iu4"}, - {"i8", "i8"}, {"int8", "i8"}, {"i4", "i4"}, - {"int4", "i4"}, {"i32", ROCKE_DTYPE_I32}, {"int32", ROCKE_DTYPE_I32}, - {"fp8", "fp8e4m3"}, {"fp8e4m3", "fp8e4m3"}, {"e4m3", "fp8e4m3"}, - {"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"}, + {"f64", "fp64"}, + {"double", "fp64"}, + {"fp64", "fp64"}, + {"fp8", "fp8e4m3"}, + {"fp8e4m3", "fp8e4m3"}, + {"e4m3", "fp8e4m3"}, + {"bf8", "bf8e5m2"}, + {"bf8e5m2", "bf8e5m2"}, + {"fp6", "fp6e2m3"}, + {"fp6e2m3", "fp6e2m3"}, + {"bf6", "fp6e3m2"}, + {"fp6e3m2", "fp6e3m2"}, + {"fp4", "fp4e2m1"}, + {"fp4e2m1", "fp4e2m1"}, + /* Integers. RDNA integer WMMA instructions call their 8- and 4-bit operands + * "iu8"/"iu4" in the instruction name only; their operand dtypes are i8 / i4, + * because signedness is an instruction operand, not a property of the data. + * "i32" is the integer accumulator. */ + {"i8", "i8"}, + {"int8", "i8"}, + {"i4", "i4"}, + {"int4", "i4"}, + {"i32", ROCKE_DTYPE_I32}, + {"int32", ROCKE_DTYPE_I32}, }; #define K_NUM_DTYPE_ALIASES ((int)(sizeof(k_dtype_aliases) / sizeof(k_dtype_aliases[0]))) @@ -124,16 +147,27 @@ const char* rocke_normalize_dtype(const char* name, char* lowered, size_t cap) } static const rocke_dtype_info_t k_dtype_info[] = { - {"fp64", ROCKE_DTYPE_FLOAT, 64}, {"fp32", ROCKE_DTYPE_FLOAT, 32}, - {"tf32", ROCKE_DTYPE_FLOAT, 32}, {"fp16", ROCKE_DTYPE_FLOAT, 16}, - {"bf16", ROCKE_DTYPE_FLOAT, 16}, {"fp8e4m3", ROCKE_DTYPE_FLOAT, 8}, - {"bf8e5m2", ROCKE_DTYPE_FLOAT, 8}, {"fp6e2m3", ROCKE_DTYPE_FLOAT, 6}, - {"fp6e3m2", ROCKE_DTYPE_FLOAT, 6}, {"fp4e2m1", ROCKE_DTYPE_FLOAT, 4}, - {"e8m0", ROCKE_DTYPE_FLOAT, 8}, {"e5m3", ROCKE_DTYPE_FLOAT, 8}, - {"i4", ROCKE_DTYPE_INTEGER, 4}, {"iu4", ROCKE_DTYPE_INTEGER, 4}, - {"i8", ROCKE_DTYPE_INTEGER, 8}, {"iu8", ROCKE_DTYPE_INTEGER, 8}, - {"i16", ROCKE_DTYPE_INTEGER, 16}, {"i32", ROCKE_DTYPE_INTEGER, 32}, - {"i64", ROCKE_DTYPE_INTEGER, 64}, {"i1", ROCKE_DTYPE_PREDICATE, 1}, + /* Floating point, grouped as in _DTYPE_INFO. */ + {"fp64", ROCKE_DTYPE_FLOAT, 64}, + {"fp32", ROCKE_DTYPE_FLOAT, 32}, + {"tf32", ROCKE_DTYPE_FLOAT, 32}, + {"fp16", ROCKE_DTYPE_FLOAT, 16}, + {"bf16", ROCKE_DTYPE_FLOAT, 16}, + {"fp8e4m3", ROCKE_DTYPE_FLOAT, 8}, + {"bf8e5m2", ROCKE_DTYPE_FLOAT, 8}, + {"fp6e2m3", ROCKE_DTYPE_FLOAT, 6}, + {"fp6e3m2", ROCKE_DTYPE_FLOAT, 6}, + {"fp4e2m1", ROCKE_DTYPE_FLOAT, 4}, + {"e8m0", ROCKE_DTYPE_FLOAT, 8}, + {"e5m3", ROCKE_DTYPE_FLOAT, 8}, + /* Integers. */ + {"i4", ROCKE_DTYPE_INTEGER, 4}, + {"i8", ROCKE_DTYPE_INTEGER, 8}, + {"i16", ROCKE_DTYPE_INTEGER, 16}, + {"i32", ROCKE_DTYPE_INTEGER, 32}, + {"i64", ROCKE_DTYPE_INTEGER, 64}, + /* Predicates. */ + {"i1", ROCKE_DTYPE_PREDICATE, 1}, }; const rocke_dtype_info_t* rocke_dtype_info(const char* name) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_builder.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_builder.cpp index c49b01b7086d..a480ef1fea05 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_builder.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_builder.cpp @@ -385,6 +385,12 @@ rocke_op_t* rocke_i_op(rocke_ir_builder_t* b, if(num_regions < 0) num_regions = 0; + for(i = 0; i < num_results; i++) + { + rocke_i_refuse_packed_only_value(rocke_opcode_name(opcode), + result_types ? result_types[i] : NULL); + } + op = (rocke_op_t*)rocke_arena_calloc(&b->arena, sizeof(*op)); if(!op) { @@ -640,6 +646,24 @@ int rocke_i_count_of(const rocke_type_t* t) return 1; } +void rocke_i_refuse_packed_only_value(const char* op_name, const rocke_type_t* t) +{ + const rocke_type_t* elem + = (t != NULL && t->kind == ROCKE_TYPE_SMEM) ? t->elem : rocke_i_elem_of(t); + if(elem == NULL || elem->kind != ROCKE_TYPE_SCALAR) + { + return; + } + if(elem->scalar == ROCKE_SCALAR_I4) + { + throw ckc::TypeError( + ckc::format_error("%s: %s values travel packed in 32-bit words and have no scalar or " + "vector form -- keep them in i32 storage and unpack explicitly", + op_name, + elem->name)); + } +} + /* ============================== BUILDER ================================= */ rocke_status_t rocke_ir_builder_init(rocke_ir_builder_t* b, const char* kernel_name) @@ -820,6 +844,8 @@ rocke_value_t* rocke_b_param(rocke_ir_builder_t* b, } } + rocke_i_refuse_packed_only_value("param", t); + /* Value name carries the leading '%'. */ full_name = rocke_arena_printf(&b->arena, "%%%s", name); if(!full_name) 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 cf0557cff26a..8259a63f0196 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 @@ -46,6 +46,8 @@ } ROCKE_SCALAR_SINGLETON(rocke_i1, ROCKE_SCALAR_I1, "i1") +/* Logical only, like fp4e2m1: no LLVM rendering; its values travel packed. */ +ROCKE_SCALAR_SINGLETON(rocke_i4, ROCKE_SCALAR_I4, "i4") ROCKE_SCALAR_SINGLETON(rocke_i8, ROCKE_SCALAR_I8, "i8") ROCKE_SCALAR_SINGLETON(rocke_i16, ROCKE_SCALAR_I16, "i16") ROCKE_SCALAR_SINGLETON(rocke_i32, ROCKE_SCALAR_I32, "i32") @@ -72,6 +74,8 @@ const rocke_type_t* rocke_scalar_by_name(const char* name) return NULL; if(strcmp(name, "i1") == 0) return rocke_i1(); + if(strcmp(name, "i4") == 0) + return rocke_i4(); if(strcmp(name, "i8") == 0) return rocke_i8(); if(strcmp(name, "i16") == 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 1244fed1bda2..46a28e732a4b 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 @@ -121,6 +121,7 @@ rocke_value_t* rocke_b_zero_vec(rocke_ir_builder_t* b, const rocke_type_t* elem, return NULL; if(!elem) return (rocke_value_t*)rocke_i_set_err(b, ROCKE_ERR_VALUE, "zero_vec elem is NULL"); + rocke_i_refuse_packed_only_value("zero_vec", elem); if(rocke_type_eq(elem, rocke_f32())) return rocke_b_zero_vec_f32(b, n); if(rocke_type_eq(elem, rocke_f16()) || rocke_type_eq(elem, rocke_bf16()) 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 4f763b49b6d9..5e643761d02e 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 @@ -11,7 +11,7 @@ * / rocke_h_push_indent / rocke_h_pop_indent), * - exception-based errors and the NULL guard (rocke_h_fail / rocke_h_live), * - naming / type mapping (rocke_h_name / rocke_h_type_to_hip / rocke_h_hip_scalar / - * rocke_h_vec_prefix), + * rocke_h_require_hip_scalar / rocke_h_vec_prefix), * - float literal formatting (rocke_h_f32_literal / rocke_h_f64_literal), * - waitcnt encoders (rocke_h_encode_waitcnt + the two raw encoders), * - arch gates (rocke_h_require_ds_read_tr / rocke_h_require_wmma_arch), @@ -38,6 +38,7 @@ #include "rocke/arena.h" #include "rocke/error.hpp" /* ckc::Error boundary translation */ #include "rocke/ir.h" +#include "rocke/ir_internal.h" /* rocke_i_set_err_msg */ #include "rocke/lower_hip.h" #include "rocke/lower_hip_internal.h" #include "rocke/strbuf.h" @@ -278,7 +279,9 @@ const char* rocke_h_name(rocke_h_lowerer_t* lw, const rocke_value_t* v) } /* The raw _HIP_TYPE dict: scalar IR name -> HIP scalar spelling. NULL if - * unknown (caller decides whether that is an error). */ + * unknown (caller decides whether that is an error). + * No packed-only type has an entry: rocke_i_refuse_packed_only_value keeps + * its values out of the IR. */ const char* rocke_h_hip_scalar(const char* ir_scalar_name) { if(!ir_scalar_name) @@ -332,6 +335,21 @@ const char* rocke_h_hip_scalar(const char* ir_scalar_name) return NULL; } +const char* rocke_h_require_hip_scalar(rocke_h_lowerer_t* lw, const char* ir_type_name) +{ + const char* spelling = rocke_h_hip_scalar(ir_type_name); + if(!spelling) + { + rocke_h_fail(lw, + ROCKE_ERR_NOTIMPL, + "no HIP mapping for type %s", + ir_type_name ? ir_type_name : "(null)"); + } + return spelling; +} + +/* No packed-only type has an entry: rocke_i_refuse_packed_only_value keeps its + * values out of the IR. */ const char* rocke_h_vec_prefix(const char* ir_scalar_name, bool full_map) { if(ir_scalar_name) @@ -430,8 +448,8 @@ void rocke_h_reject_f64(rocke_h_lowerer_t* lw, const rocke_op_t* op, const char* } } -/* Python _type_to_hip(t). Returns an arena-owned string; throws on an unmappable - * HIP type (KeyError parity), including logical types without direct HIP lowering. */ +/* Python _type_to_hip(t). Returns an arena-owned string; a type it does not + * compose itself is refused by rocke_h_require_hip_scalar. */ const char* rocke_h_type_to_hip(rocke_h_lowerer_t* lw, const rocke_type_t* t) { if(!t) @@ -453,7 +471,7 @@ const char* rocke_h_type_to_hip(rocke_h_lowerer_t* lw, const rocke_type_t* t) out = rocke_arena_printf(&lw->b->arena, "%s*", pointee); return out ? out : ""; } - /* Other ptr spaces fall through to the scalar map (KeyError in Python). */ + /* Other ptr spaces fall through to the scalar map, which refuses them. */ } if(t->kind == ROCKE_TYPE_VECTOR) @@ -478,8 +496,9 @@ const char* rocke_h_type_to_hip(rocke_h_lowerer_t* lw, const rocke_type_t* t) && 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 ""; + /* Python falls through to the scalar map, which refuses the + * vector type by its own name. */ + return rocke_h_require_hip_scalar(lw, t->name); } } out = rocke_arena_printf(&lw->b->arena, "%s%d", pfx, t->count); @@ -498,19 +517,8 @@ const char* rocke_h_type_to_hip(rocke_h_lowerer_t* lw, const rocke_type_t* t) return out ? out : ""; } - /* scalar: _HIP_TYPE[t.name] */ - { - const char* s = rocke_h_hip_scalar(t->name); - if(!s) - { - rocke_h_fail(lw, - ROCKE_ERR_KEY, - "type_to_hip: unmappable type '%s'", - t->name ? t->name : "(null)"); - return ""; - } - return s; - } + /* scalar: _hip_scalar_type(t.name) */ + return rocke_h_require_hip_scalar(lw, t->name); } /* ============================== literals ============================= */ @@ -970,8 +978,9 @@ static void = 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* scalar = strcmp(prefix, "i8x") == 0 + ? "int8_t" + : ckc::rocke_h_require_hip_scalar(lw, t->elem->name); const char* macro = rocke_arena_printf(&lw->b->arena, "_ROCKE_VEC(%s, %s, %d)", scalar, prefix, t->count); const char* declaration @@ -1137,6 +1146,8 @@ rocke_status_t rocke_lower_kernel_to_hip(rocke_ir_builder_t* b, { rocke_strbuf_free(&sig); h_smem_release(&lw); + /* The header promises the reason on the builder's error channel. */ + rocke_i_set_err_msg(b, e.code(), e.what()); return e.code(); } catch(const std::bad_alloc&) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_arith.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_arith.cpp index 6afb2e8357aa..e7caab9e0684 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_arith.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_arith.cpp @@ -7,7 +7,7 @@ * * Each `_op_*` Python method becomes a static `rocke_h_op_*` handler with the * (lw, op) signature from lower_hip_internal.h. Shared helpers (rocke_h_emit / - * rocke_h_emitf / rocke_h_name / rocke_h_type_to_hip / rocke_h_hip_scalar / + * rocke_h_emitf / rocke_h_name / rocke_h_type_to_hip / rocke_h_require_hip_scalar / * rocke_h_f32_literal / rocke_attr_get_*) are DEFINED in lower_hip_core.c and only * called here. The registration table is exported via rocke_h_handlers_arith(), * which the core bucket stitches into the dispatch table. @@ -81,11 +81,7 @@ static rocke_status_t rocke_h_op_arith_constant(rocke_h_lowerer_t* lw, const roc { ity = "i32"; } - cpp_t = rocke_h_hip_scalar(ity); - if(!cpp_t) - { - return rocke_h_fail(lw, ROCKE_ERR_KEY, "arith.constant: unknown ity %s", ity); - } + cpp_t = rocke_h_require_hip_scalar(lw, ity); val = rocke_attr_get(&op->attrs, "value"); if(!val) { 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 8023901a6163..2becb116f6c5 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 @@ -12,7 +12,7 @@ * * Binds strictly to rocke/ir.h (frozen IR) and the shared helpers declared in * rocke/lower_hip_internal.h (defined in bucket 0). Shared helpers (rocke_h_emit*, - * rocke_h_name, rocke_h_type_to_hip, rocke_h_hip_scalar, rocke_h_vec_prefix, + * rocke_h_name, rocke_h_type_to_hip, rocke_h_require_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" @@ -126,14 +126,7 @@ static rocke_status_t _op_tile_smem_alloc(rocke_h_lowerer_t* lw, const rocke_op_ return rocke_h_fail(lw, ROCKE_ERR_TYPE, "tile.smem_alloc: result is not an smem type"); } - elem = rocke_h_hip_scalar(st->elem ? st->elem->name : ""); - if(!elem) - { - return rocke_h_fail(lw, - ROCKE_ERR_KEY, - "tile.smem_alloc: no HIP type for elem '%s'", - st->elem ? st->elem->name : "(null)"); - } + elem = rocke_h_require_hip_scalar(lw, st->elem ? st->elem->name : NULL); nice = rocke_h_name(lw, res); /* dims = "][".join(str(d) for d in st.shape) */ 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 a7aeb72e6f9e..f5a3af03bc7d 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 @@ -8,9 +8,10 @@ * Each `_op_*` Python method becomes a static `rocke_h_op_*` handler with the * (lw, op) signature from lower_hip_internal.h. Shared helpers * (rocke_h_emit / rocke_h_emitf / rocke_h_name / rocke_h_type_to_hip / - * rocke_h_hip_scalar / rocke_h_require_wmma_arch / rocke_h_require_ds_read_tr / - * rocke_h_lower_region / rocke_h_push_indent / rocke_h_pop_indent / rocke_attr_get_* / - * rocke_vector_type) are DEFINED elsewhere (bucket 0); only called here. + * rocke_h_hip_scalar / rocke_h_require_hip_scalar / rocke_h_require_wmma_arch / + * rocke_h_require_ds_read_tr / rocke_h_lower_region / rocke_h_push_indent / + * rocke_h_pop_indent / rocke_attr_get_* / rocke_vector_type) are DEFINED elsewhere + * (bucket 0); only called here. * * The simple handlers reproduce every _emit() format string byte-for-byte. * The complex matrix-engine handlers are fully ported too: @@ -55,11 +56,11 @@ static int h_vcount(const rocke_type_t* t) } /* Python idiom `t.elem.name if isinstance(t, VectorType) else t.name`, mapped - * to the HIP scalar spelling via _HIP_TYPE (rocke_h_hip_scalar). */ -static const char* h_elem_scalar(const rocke_type_t* t) + * to the HIP scalar spelling (rocke_h_require_hip_scalar). */ +static const char* h_elem_scalar(rocke_h_lowerer_t* lw, const rocke_type_t* t) { const char* nm = (t && t->kind == ROCKE_TYPE_VECTOR) ? t->elem->name : t->name; - return rocke_h_hip_scalar(nm); + return rocke_h_require_hip_scalar(lw, nm); } /* ============================== mma ====================================== */ @@ -858,7 +859,7 @@ static rocke_status_t rocke_h_op_vector_extract(rocke_h_lowerer_t* lw, const roc rocke_attr_get_int(&op->attrs, "index", &i); rocke_h_emitf(lw, "%s %s = %s[%d];", - h_elem_scalar(v->type), + h_elem_scalar(lw, v->type), rocke_h_name(lw, r), rocke_h_name(lw, v), (int)i); @@ -974,7 +975,7 @@ static rocke_status_t rocke_h_op_vector_sum(rocke_h_lowerer_t* lw, const rocke_o const char* nice = rocke_h_name(lw, r); const char* vn = rocke_h_name(lw, v); int n = h_vcount(v->type), i; - rocke_h_emitf(lw, "%s %s = %s[0];", h_elem_scalar(v->type), nice, vn); + rocke_h_emitf(lw, "%s %s = %s[0];", h_elem_scalar(lw, v->type), nice, vn); for(i = 1; i < n; i++) { rocke_h_emitf(lw, "%s = %s + %s[%d];", nice, nice, vn, i); @@ -990,7 +991,7 @@ static rocke_status_t rocke_h_op_vector_reduce_max(rocke_h_lowerer_t* lw, const const char* nice = rocke_h_name(lw, r); const char* vn = rocke_h_name(lw, v); int n = h_vcount(v->type), i; - rocke_h_emitf(lw, "%s %s = %s[0];", h_elem_scalar(v->type), nice, vn); + rocke_h_emitf(lw, "%s %s = %s[0];", h_elem_scalar(lw, v->type), nice, vn); for(i = 1; i < n; i++) { rocke_h_emitf(lw, "%s = (%s[%d] > %s) ? %s[%d] : %s;", nice, vn, i, nice, vn, i, nice); @@ -1235,7 +1236,7 @@ static rocke_status_t rocke_h_op_vector_trunc_f32_to(rocke_h_lowerer_t* lw, cons { target = "f16"; } - elem_cpp = rocke_h_hip_scalar(target); + elem_cpp = rocke_h_require_hip_scalar(lw, target); rocke_h_emitf(lw, "%s %s;", res_t, nice); for(i = 0; i < n; i++) { @@ -1344,15 +1345,11 @@ static rocke_status_t rocke_h_op_scf_for(rocke_h_lowerer_t* lw, const rocke_op_t } iv_name = (iv_full[0] == '%') ? iv_full + 1 : iv_full; /* [1:] */ iv_type = rocke_attr_get_str(&op->attrs, "iv_type"); - iv_ty = iv_type ? rocke_h_hip_scalar(iv_type) : NULL; - if(!iv_ty) + if(!iv_type) { - /* Python _HIP_TYPE[...] KeyError parity. */ - return rocke_h_fail(lw, - ROCKE_ERR_KEY, - "scf.for: no HIP type for iv_type '%s'", - iv_type ? iv_type : ""); + return rocke_h_fail(lw, ROCKE_ERR_KEY, "scf.for: missing 'iv_type' attr"); } + iv_ty = rocke_h_require_hip_scalar(lw, iv_type); /* Declare for-op results in the enclosing scope (zip stops at min length). */ for(i = 0; i < meta_n && i < op->num_results; i++) 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 a96bd600dd58..f2403bd0ea3a 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 @@ -614,7 +614,8 @@ const char* rocke_ll_llvm_type(rocke_lower_t* L, const rocke_type_t* t) { return "ptr addrspace(3)"; } - /* scalar */ + /* scalar. No packed-only type has an entry: rocke_i_refuse_packed_only_value + * keeps its values out of the IR. */ const char* n = t->name; if(n) { @@ -1224,7 +1225,9 @@ static void } } -/* Size of one smem allocation in bytes (Python _seg_size). */ +/* Size of one smem allocation in bytes (Python _seg_size). + * No packed-only type has an entry: rocke_i_refuse_packed_only_value refuses + * its LDS tiles. */ static int ll_smem_seg_size(const rocke_type_t* stype) { if(!stype || !stype->elem || !stype->elem->name) @@ -1247,7 +1250,8 @@ static int ll_smem_seg_size(const rocke_type_t* stype) return seg; } -/* Alignment for one smem allocation (Python _align). */ +/* Alignment for one smem allocation (Python _align). A packed-only type never + * reaches here: rocke_i_refuse_packed_only_value refuses its LDS tiles. */ static int ll_smem_align(const rocke_type_t* stype) { if(!stype || !stype->elem || !stype->elem->name) 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 7e781d9f76e0..fd3202dc01a3 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 @@ -112,6 +112,7 @@ typedef enum rocke_scalar_kind ROCKE_SCALAR_E5M3, ROCKE_SCALAR_TF32, ROCKE_SCALAR_F64, + ROCKE_SCALAR_I4, ROCKE_SCALAR__COUNT } rocke_scalar_kind_t; @@ -639,6 +640,7 @@ typedef struct rocke_mem_opts /* Interned scalar singletons (Python module-level I1, F32, ...). Always valid; * never NULL; never arena-owned (static storage). */ const rocke_type_t* rocke_i1(void); +const rocke_type_t* rocke_i4(void); const rocke_type_t* rocke_i8(void); const rocke_type_t* rocke_i16(void); const rocke_type_t* rocke_i32(void); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir_internal.h index 7ab6e5383614..de5e21c27f40 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir_internal.h @@ -136,6 +136,13 @@ bool rocke_i_is_vector(const rocke_type_t* t, const char* elem_name, int count); * `v.type.elem if isinstance(...,VectorType) else v.type`). */ const rocke_type_t* rocke_i_elem_of(const rocke_type_t* t); +/* Throw ckc::TypeError when `t` is a packed-only scalar (i4), a vector of one + * or an LDS tile of one: the twin of core/ir.py::_refuse_packed_only_value, + * which states the rule and the message. rocke_i_op calls it for every result + * and rocke_b_param for every parameter, so no builder makes such a value. + * NULL-safe. */ +void rocke_i_refuse_packed_only_value(const char* op_name, const rocke_type_t* t); + /* Lane count of a vector type, or 1 for a scalar. */ int rocke_i_count_of(const rocke_type_t* t); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/lower_hip_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/lower_hip_internal.h index 3192eabf6796..68bcd80b78f3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/lower_hip_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/lower_hip_internal.h @@ -146,16 +146,24 @@ const char* rocke_h_name(rocke_h_lowerer_t* lw, const rocke_value_t* v); * vec -> "f16xN"/"bf16xN"/"f32xN"/"i32xN"/"i16xN"/"i8xN" * (fp8e4m3/bf8e5m2 -> "i8xN"; i1 -> "boolxN") * smem -> "*" - * Returns an arena-owned string. Throws for an unmappable HIP type (Python - * KeyError parity), including logical types without direct HIP lowering. */ + * Returns an arena-owned string. A type it does not compose itself goes to + * rocke_h_require_hip_scalar, which refuses an unmappable one (logical types + * without direct HIP lowering included). */ const char* rocke_h_type_to_hip(rocke_h_lowerer_t* lw, const rocke_type_t* t); /* Map a scalar IR type *name* ("i32","f16",...) to its HIP scalar spelling - * (the raw _HIP_TYPE dict). Returns NULL on an unknown name (caller decides - * whether that is an error). Used where Python indexes _HIP_TYPE[name] directly - * (e.g. constants, vector.sum element type). */ + * (the raw _HIP_TYPE dict). Returns NULL on an unknown name. Only a caller with + * a fallback for a miss (scf.for iter-arg types use "auto") calls this + * directly; every other caller uses rocke_h_require_hip_scalar. */ const char* rocke_h_hip_scalar(const char* ir_scalar_name); +/* Python _hip_scalar_type: rocke_h_hip_scalar, or refuse the type with + * ROCKE_ERR_NOTIMPL "no HIP mapping for type " -- the HIP twin of the + * LLVM lowering's "no LLVM mapping for type ", so a type with no HIP + * spelling (a logical sub-byte type such as i4, or a vector of an unmapped + * element) is refused with one message wherever it is met. */ +const char* rocke_h_require_hip_scalar(rocke_h_lowerer_t* lw, const char* ir_type_name); + /* The "f16x"/"bf16x"/"f32x"/"i32x"/"i16x"/"i8x" vector-name prefix for an IR * scalar element name, used by the smem/global vN load/store handlers. Returns * the Python fallback "f16x" for names not in the small map (matching the 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 9166a6a2f499..a116b6c00d3b 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 @@ -125,6 +125,18 @@ from logical TF32, including gfx950's BF16x3-emulated atoms, which split each operand's payload into three BF16 products rather than rounding it to BF16; no production selector maps TF32 requests to BF16 compute. +**Integer types.** The 8- and 4-bit integer dtypes are `i8` and `i4`. RDNA integer +WMMA instructions say `iu8` / `iu4` in their names, but those are not dtypes: +signedness is an operand of those instructions, not a property of the data, so +their operands are `i8` / `i4`. `i4` is packed-only: its +values only ever live packed in 32-bit words, so it has no scalar or vector form +and no LLVM or HIP rendering. The IR builder of both engines refuses any op +result, kernel parameter or LDS tile of `i4` elements (`: i4 values travel +packed in 32-bit words and have no scalar or vector form -- keep them in i32 +storage and unpack explicitly`); a kernel keeps the words in `i32` storage and +unpacks them explicitly. Both lowerings still refuse an `i4` value that reaches +them (`no LLVM mapping for type i4`, `no HIP mapping for type i4`). + **Types that are entirely absent from the vocabulary:** unsigned `u8/u16/u32/u64`, and a distinct `bool` (predicates reuse `i1`). No engine references them. @@ -222,7 +234,7 @@ Gap categories: **[C]** conversion gap · **[P]** partial/plumbing gap · | `bf8e5m2` (OCP) | 3 | 3 | 1 | 1 | 1 | 1 | **56%** | [C][P] | | `fp8 fnuz` (gfx942 only) | 1 | 3 | 1 | 1 | 1 | 1 | **44%** | [P] | | **Sub-byte / block formats** | | | | | | | | | -| `i4` (packed) | 0 | 1 | 1 | 2 | 2 | 2 | **44%** | [C][P] | +| `i4` (packed) | 1 | 1 | 1 | 2 | 2 | 2 | **50%** | [C][P] | | `fp4` (E2M1, MX) | 0 | 2 | 1 | 1 | 0 | 2 | **33%** | [P][M] | | `fp6` (E2M3, MX) | 0 | 2 | 0 | 1 | 0 | 2 | **28%** | [C][P][M] | | `bf6` (E3M2, MX) | 0 | 2 | 0 | 1 | 0 | 2 | **28%** | [C][P][M] | @@ -293,12 +305,12 @@ mode. The missing selector and precision contract are described in §2.5. rather than 3 because the scale semantics are not part of the contract — the helpers take a scale operand but neither document nor enforce per-tensor vs per-channel, so two callers can disagree about what a scale means. -- **`i4` C1=0 / C3=1** — packed i4 is not a scalar IR `Type`. +- **`i4` C1=1 / C3=1** -- `I4` is a logical IR `Type` with no LLVM rendering; its values travel packed. [i4_dequant.py](../../python/rocke/helpers/i4_dequant.py) provides conversions from packed i4 to i32, f32, f16, fp8, and bf8, but no reverse f32→packed-i4 conversion. The quantization helpers accept only i8, fp8, and bf8 (`quant.py:58-70`). This is one-direction support under C3's round-trip - criterion. C2=1 reflects `iu4` WMMA on gfx1151/gfx11-generic; the MFMA targets + criterion. C2=1 reflects the i4 (`iu4`) WMMA on gfx1151/gfx11-generic; the MFMA targets and gfx1250 have no native INT4 matrix instruction (§4A). - **`fp4` / `fp6` / `bf6` C5=0** — the widest silicon-to-software gap in the matrix. The gfx950 atoms exist (`mfma_f32_16x16x128_fp4`, @@ -469,10 +481,10 @@ The original basis snapshot, with the gfx942 XF32 addition noted below: | `gfx90a` | `mfma_f32_{16x16x16,32x32x8}_{f16,bf16}` | f16, bf16 | | `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** | +| `gfx1151` | `wmma_f32_16x16x16_{f16,bf16}`, `wmma_i32_16x16x16_{iu8,iu4}` | f16, bf16, **i8, i4** | | `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) | -| `gfx11-generic` | same as `gfx1151` | f16, bf16, iu8, iu4 | +| `gfx11-generic` | same as `gfx1151` | f16, bf16, i8, i4 | | `gfx12-generic` | same as `gfx1201` | f16, bf16 | - The `gfx1250` catalog entry defines its target family, memory model, limits, 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 266997b8eb4b..e7a206eec32d 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 @@ -65,6 +65,7 @@ F64, FP8E4M3, I1, + I4, I8, I16, I32, @@ -172,6 +173,7 @@ "F64", "FP8E4M3", "I1", + "I4", "I8", "I16", "I32", 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 aa5916c050f5..125d4acd55c9 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 @@ E5M3, dtype_to_ir_type, I1, + I4, I8, I32, I64, @@ -125,6 +126,7 @@ "F64", "FP8E4M3", "I1", + "I4", "I8", "I32", "I64", 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 47510050f238..11e78fb12799 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 @@ -153,8 +153,8 @@ "mma": [ {"family": "wmma", "a": "fp16", "b": "fp16", "c": "fp32", "m": 16, "n": 16, "k": 16, "op_id": "wmma_f32_16x16x16_f16"}, {"family": "wmma", "a": "bf16", "b": "bf16", "c": "fp32", "m": 16, "n": 16, "k": 16, "op_id": "wmma_f32_16x16x16_bf16"}, - {"family": "wmma", "a": "iu8", "b": "iu8", "c": "i32", "m": 16, "n": 16, "k": 16, "op_id": "wmma_i32_16x16x16_iu8"}, - {"family": "wmma", "a": "iu4", "b": "iu4", "c": "i32", "m": 16, "n": 16, "k": 16, "op_id": "wmma_i32_16x16x16_iu4"} + {"family": "wmma", "a": "i8", "b": "i8", "c": "i32", "m": 16, "n": 16, "k": 16, "op_id": "wmma_i32_16x16x16_iu8"}, + {"family": "wmma", "a": "i4", "b": "i4", "c": "i32", "m": 16, "n": 16, "k": 16, "op_id": "wmma_i32_16x16x16_iu4"} ] }, "gfx1201": { @@ -272,8 +272,8 @@ "mma": [ {"family": "wmma", "a": "fp16", "b": "fp16", "c": "fp32", "m": 16, "n": 16, "k": 16, "op_id": "wmma_f32_16x16x16_f16"}, {"family": "wmma", "a": "bf16", "b": "bf16", "c": "fp32", "m": 16, "n": 16, "k": 16, "op_id": "wmma_f32_16x16x16_bf16"}, - {"family": "wmma", "a": "iu8", "b": "iu8", "c": "i32", "m": 16, "n": 16, "k": 16, "op_id": "wmma_i32_16x16x16_iu8"}, - {"family": "wmma", "a": "iu4", "b": "iu4", "c": "i32", "m": 16, "n": 16, "k": 16, "op_id": "wmma_i32_16x16x16_iu4"} + {"family": "wmma", "a": "i8", "b": "i8", "c": "i32", "m": 16, "n": 16, "k": 16, "op_id": "wmma_i32_16x16x16_iu8"}, + {"family": "wmma", "a": "i4", "b": "i4", "c": "i32", "m": 16, "n": 16, "k": 16, "op_id": "wmma_i32_16x16x16_iu4"} ] }, "gfx12-generic": { 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 86df9d5d5ef0..840bd920cbcf 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 @@ -39,11 +39,10 @@ "fp6e3m2": "fp6e3m2", "fp4": "fp4e2m1", "fp4e2m1": "fp4e2m1", - # Integer WMMA: "iu8"/"iu4" are the RDNA WMMA integer operand families - # (signedness is an instruction operand, not the dtype); "i32" is the - # integer accumulator. Scalar int spellings pass through for completeness. - "iu8": "iu8", - "iu4": "iu4", + # Integers. RDNA integer WMMA instructions call their 8- and 4-bit operands + # "iu8"/"iu4" in the instruction name only; their operand dtypes are i8 / i4, + # because signedness is an instruction operand, not a property of the data. + # "i32" is the integer accumulator. "i8": "i8", "int8": "i8", "i4": "i4", @@ -98,9 +97,7 @@ class DTypeInfo: DTypeCategory.INTEGER, ( ("i4", 4), - ("iu4", 4), ("i8", 8), - ("iu8", 8), ("i16", 16), ("i32", 32), ("i64", 64), 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 da1e9d1eab52..c72e7829b676 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 @@ -45,6 +45,7 @@ def __repr__(self) -> str: I1 = Type("i1") +I4 = Type("i4") I8 = Type("i8") I16 = Type("i16") I32 = Type("i32") @@ -71,6 +72,7 @@ def dtype_to_ir_type(dtype: str) -> Type: info = dtype_info(dtype) types = { "i1": I1, + "i4": I4, "i8": I8, "i16": I16, "i32": I32, @@ -262,6 +264,27 @@ def _check_flag_ptr(op_name: str, ptr: "Value") -> None: raise ValueError(f"{op_name} expects a global i32 pointer, got {ptr.type.name}") +# I4 is packed-only: its values only ever live packed in 32-bit words, so no IR +# value, scalar or vector, and no LDS tile has that element type. IRBuilder +# enforces it for every op result and kernel parameter through +# _refuse_packed_only_value, the one home of that rule; the lowerers' type +# tables have no entry for a packed-only type. +_PACKED_ONLY_SCALARS = ("i4",) + + +def _refuse_packed_only_value(op_name: str, t: "Type") -> None: + """Refuse an ``op_name`` result or parameter of type ``t`` whose elements + are packed-only: a scalar, a vector of one, or an LDS tile of one. An + untyped placeholder (``None``) has no elements and passes.""" + elem = t.elem if isinstance(t, (VectorType, SmemType)) else t + if elem is not None and elem.name in _PACKED_ONLY_SCALARS: + raise TypeError( + f"{op_name}: {elem.name} values travel packed in 32-bit words and have " + "no scalar or vector form -- keep them in i32 storage and unpack " + "explicitly" + ) + + def _sync_attr(op: "Op", key: str, allowed: Sequence[str]) -> str: """Scope / ordering attr of a fence or flag op, checked for the lowerers. @@ -831,6 +854,7 @@ def param(self, name: str, t: Type, **attrs: Any) -> Value: """ if name in self._param_values: raise ValueError(f"duplicate kernel parameter {name!r}") + _refuse_packed_only_value("param", t) v = Value(name=f"%{name}", type=t) self._param_values[name] = v self._params.append(Param(name=name, type=t, attrs=dict(attrs))) @@ -881,6 +905,8 @@ def _op( result_name_hint: str = "v", loc: Optional[str] = None, ) -> Op: + for t in result_types: + _refuse_packed_only_value(name, t) results = [Value(self._fresh(result_name_hint), t) for t in result_types] op = Op( name=name, @@ -1815,6 +1841,7 @@ def zero_vec_f32(self, n: int) -> Value: ).result def zero_vec(self, elem: Type, n: int) -> Value: + _refuse_packed_only_value("zero_vec", elem) if elem == F32: return self.zero_vec_f32(n) if elem in (F16, BF16, F64, FP8E4M3, BF8E5M2, I8, I16, I32): 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 8c10f09cebca..7eb1fa07c0c1 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 @@ -38,6 +38,8 @@ ) +# No packed-only type has an entry: core/ir.py _refuse_packed_only_value +# keeps its values out of the IR. _HIP_TYPE = { "i1": "bool", "i8": "int8_t", @@ -53,6 +55,24 @@ "bf8e5m2": "bf8e5m2", } + +def _hip_scalar_type(type_name: str) -> str: + """Return the HIP scalar spelling of the IR type named ``type_name``, or refuse it. + + Every handler that needs a scalar spelling asks here, and + :func:`_type_to_hip` sends here every type it does not compose itself, so + one message refuses them all (only the ``scf.for`` iter-arg declarations + fall back to ``auto``). A type with no HIP spelling -- a logical sub-byte + type such as ``i4`` whose values travel packed, or a vector of an unmapped + element -- is refused with the HIP twin of the LLVM lowering's ``no LLVM + mapping for type ``. + """ + try: + return _HIP_TYPE[type_name] + except KeyError: + raise NotImplementedError(f"no HIP mapping for type {type_name}") from None + + _HIP_UNSIGNED_INT_TYPE = { "i1": "bool", "i8": "uint8_t", @@ -69,6 +89,8 @@ def _type_to_hip(t) -> str: if isinstance(t, VectorType): # Naming convention matches the prologue's typedefs: # ``f16xN``, ``bf16xN``, ``f32xN``, ``i32xN``, ``i8xN`` for N>=1. + # No packed-only type has an entry: core/ir.py _refuse_packed_only_value + # keeps its values out of the IR. elem = t.elem.name if elem == "f16": return f"f16x{t.count}" @@ -98,7 +120,7 @@ def _type_to_hip(t) -> str: # smem allocations expand into __shared__ arrays at the top of # the kernel; the value at the use site is a typed pointer-ish. return f"{_type_to_hip(t.elem)}*" - return _HIP_TYPE[t.name] + return _hip_scalar_type(t.name) # Compilable-HIP prologue. Pasted at the top of every lowered source so @@ -394,7 +416,7 @@ def _op_arith_constant(self, op: Op) -> None: res = op.result ity = op.attrs.get("ity", "i32") val = op.attrs["value"] - cpp_t = _HIP_TYPE[ity] + cpp_t = _hip_scalar_type(ity) if ity in ("f16", "f32"): literal = _f32_literal(float(val)) if ity == "f16": @@ -531,7 +553,7 @@ def _op_tile_smem_alloc(self, op: Op) -> None: st = op.result.type assert isinstance(st, SmemType) dims = "][".join(str(d) for d in st.shape) - elem = _HIP_TYPE[st.elem.name] + elem = _hip_scalar_type(st.elem.name) nice = _name(op.result) decl = f" __shared__ {elem} {nice}_storage[{dims}];" self.smem_decls.append(decl) @@ -1073,7 +1095,7 @@ def _op_tile_buffer_load(self, op: Op) -> None: # 4-byte types (f32, i32) → __builtin_amdgcn_raw_buffer_load_b32. rsrc, voffset, soffset = op.operands elem_type = op.attrs["elem_type"] - hip_t = _HIP_TYPE[elem_type] + hip_t = _hip_scalar_type(elem_type) tmp = f"_bl_{_name(op.result).lstrip('%')}" if elem_type in ("f16", "bf16"): self._emit( @@ -1327,7 +1349,9 @@ def _op_vector_extract(self, op: Op) -> None: (v,) = op.operands i = op.attrs["index"] elem_t = v.type.elem if isinstance(v.type, VectorType) else v.type - self._emit(f"{_HIP_TYPE[elem_t.name]} {_name(op.result)} = {_name(v)}[{i}];") + self._emit( + f"{_hip_scalar_type(elem_t.name)} {_name(op.result)} = {_name(v)}[{i}];" + ) def _op_vector_trunc_f32_to_f16(self, op: Op) -> None: # Legacy op name; the post-merge IR emits ``vector.trunc_f32_to`` @@ -2402,7 +2426,7 @@ def _op_vector_sum(self, op: Op) -> None: (v,) = op.operands n = v.type.count if isinstance(v.type, VectorType) else 1 elem_t = v.type.elem.name if isinstance(v.type, VectorType) else v.type.name - cpp_t = _HIP_TYPE[elem_t] + cpp_t = _hip_scalar_type(elem_t) nice = _name(op.result) self._emit(f"{cpp_t} {nice} = {_name(v)}[0];") for i in range(1, n): @@ -2545,7 +2569,7 @@ def _op_vector_reduce_max(self, op: Op) -> None: (v,) = op.operands n = v.type.count if isinstance(v.type, VectorType) else 1 elem_t = v.type.elem.name if isinstance(v.type, VectorType) else v.type.name - cpp_t = _HIP_TYPE[elem_t] + cpp_t = _hip_scalar_type(elem_t) nice = _name(op.result) self._emit(f"{cpp_t} {nice} = {_name(v)}[0];") for i in range(1, n): @@ -2560,7 +2584,7 @@ def _op_vector_trunc_f32_to(self, op: Op) -> None: (v,) = op.operands n = v.type.count if isinstance(v.type, VectorType) else 1 res_t = _type_to_hip(op.result.type) - elem_cpp = _HIP_TYPE[op.attrs.get("target", "f16")] + elem_cpp = _hip_scalar_type(op.attrs.get("target", "f16")) nice = _name(op.result) self._emit(f"{res_t} {nice};") for i in range(n): @@ -2581,7 +2605,7 @@ def _op_scf_for(self, op: Op) -> None: iter_inits = op.operands[3 : 3 + num_iter] iter_meta = op.attrs.get("iter_args", []) iv_name = op.attrs["iv"][1:] - iv_ty = _HIP_TYPE[op.attrs["iv_type"]] + iv_ty = _hip_scalar_type(op.attrs["iv_type"]) def cpp_type_for(type_name: str) -> str: if type_name.startswith("vec None: 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] + scalar = "int8_t" if prefix == "i8x" else _hip_scalar_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})));" 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 877b82649784..a14916c04ec1 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 @@ -1229,6 +1229,8 @@ def _llvm_type(t: Type) -> str: # The value-level type of a smem allocation token is "the pointer # to its base", since uses GEP into it. return "ptr addrspace(3)" + # No packed-only type has an entry: core/ir.py _refuse_packed_only_value + # keeps its values out of the IR. if t.name == "i1": return "i1" if t.name == "i8": @@ -2089,6 +2091,8 @@ def _compute_smem_layout(self) -> None: Falls back to the original sequential packing when liveness analysis yields no intervals (e.g. zero smem allocations). """ + # No packed-only type has an entry: core/ir.py _refuse_packed_only_value + # refuses its LDS tiles. _elem_bytes = { "i8": 1, "fp8e4m3": 1, @@ -2118,6 +2122,8 @@ def _seg_size(stype: "SmemType") -> int: return seg def _align(stype: "SmemType") -> int: + # A packed-only type never reaches here: core/ir.py + # _refuse_packed_only_value refuses its LDS tiles. if stype.elem.name == "f64": return 8 return 16 if stype.elem.name in ("i8", "fp8e4m3", "bf8e5m2") else 4 diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/engines.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/engines.py new file mode 100644 index 000000000000..7d3736e2d4c4 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/engines.py @@ -0,0 +1,105 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Test helper: lower one kernel with both engines and require the same outcome. + +rocKE's Python and C++ engines must emit the same LLVM IR for the same kernel, +and refuse the same kernel with the same message. ``ROCKE_BACKEND=both`` checks +the first half for every ``lower_kernel_to_llvm`` call, but it cannot check the +second: the Python engine raises before the C++ engine runs. These helpers drive +each engine directly, so a test states the parity it relies on whatever backend +the run selects. + +- ``lower_with_both_engines`` returns the LLVM IR both engines emit, after + asserting that the texts are identical. +- ``assert_both_engines_reject`` asserts that both engines refuse the kernel + with exactly ``message``. The C++ engine reports through the ``rocke_engine`` + binding as a ``RuntimeError`` whose text ends with ``": " + message``. +- ``lower_hip_with_both_engines`` and ``assert_both_engines_reject_hip`` are + the same two checks for the HIP-source lowering. + +Both skip, naming the reason, when the ``rocke_engine`` binding is not built. +""" + +from __future__ import annotations + +import pytest + +from rocke.core.backend import _lower_via_cpp_engine +from rocke.core.ir import KernelDef +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 + +__all__ = [ + "assert_both_engines_reject", + "assert_both_engines_reject_hip", + "lower_hip_with_both_engines", + "lower_with_both_engines", +] + + +def lower_with_both_engines( + kernel: KernelDef, arch: str, llvm_flavor: str | None = None +) -> str: + """Lower ``kernel`` for ``arch`` with each engine; return the shared IR text.""" + pytest.importorskip("rocke_engine") + python_ir = _lower_kernel_to_llvm_python(kernel, arch=arch, llvm_flavor=llvm_flavor) + native_ir = _lower_via_cpp_engine(serialize(kernel), arch, llvm_flavor) + assert native_ir == python_ir, ( + f"the C++ engine lowered kernel '{kernel.name}' for {arch} differently " + "from the Python engine" + ) + return python_ir + + +def assert_both_engines_reject( + kernel: KernelDef, + arch: str, + error_type: type[Exception], + message: str, + llvm_flavor: str | None = None, +) -> None: + """Assert that both engines refuse ``kernel`` for ``arch`` with ``message``. + + The Python engine must raise ``error_type`` with exactly ``message``; the + C++ engine must fail with the same message after its binding prefix. + """ + pytest.importorskip("rocke_engine") + with pytest.raises(error_type) as python_error: + _lower_kernel_to_llvm_python(kernel, arch=arch, llvm_flavor=llvm_flavor) + assert str(python_error.value) == message + with pytest.raises(RuntimeError) as native_error: + _lower_via_cpp_engine(serialize(kernel), arch, llvm_flavor) + assert str(native_error.value).endswith(": " + message), str(native_error.value) + + +def lower_hip_with_both_engines(kernel: KernelDef, arch: str) -> str: + """Lower ``kernel`` to HIP source for ``arch`` with each engine; return the shared text.""" + engine = pytest.importorskip("rocke_engine") + python_hip = lower_kernel_to_hip(kernel, arch=arch) + native_hip = engine.lower_serialized_ir_hip(serialize(kernel), arch) + assert native_hip == python_hip, ( + f"the C++ engine lowered kernel '{kernel.name}' to HIP for {arch} " + "differently from the Python engine" + ) + return python_hip + + +def assert_both_engines_reject_hip( + kernel: KernelDef, + arch: str, + error_type: type[Exception], + message: str, +) -> None: + """Assert that both HIP lowerings refuse ``kernel`` for ``arch`` with ``message``. + + Same contract as :func:`assert_both_engines_reject`, for the HIP source. + """ + engine = pytest.importorskip("rocke_engine") + with pytest.raises(error_type) as python_error: + lower_kernel_to_hip(kernel, arch=arch) + assert str(python_error.value) == message + with pytest.raises(RuntimeError) as native_error: + engine.lower_serialized_ir_hip(serialize(kernel), arch) + assert str(native_error.value).endswith(": " + message), str(native_error.value) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_arch_mma_ssot.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_arch_mma_ssot.py index 59a8848bac3e..26cc5015abaf 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_arch_mma_ssot.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_arch_mma_ssot.py @@ -479,3 +479,27 @@ def test_mma_naming_without_target_lookup(): ) result = b.mma(arg, a, a, c, *extra) assert result.name.lstrip("%").startswith(hint) + + +def test_rdna_integer_wmma_rows_name_i8_and_i4(): + # The instruction names keep "iu8"/"iu4"; the operand dtypes are i8/i4, and only + # i8/i4 select the ops ("iu8"/"iu4" are not dtypes). + for gfx in ("gfx1151", "gfx11-generic"): + mma = ArchTarget.from_gfx(gfx).mma + for op_id, dtype, old in ( + ("wmma_i32_16x16x16_iu8", "i8", "iu8"), + ("wmma_i32_16x16x16_iu4", "i4", "iu4"), + ): + op = mma.by_op_id(op_id) + assert (op.a_dtype, op.b_dtype, op.c_dtype) == (dtype, dtype, "i32") + for spelling, want in ((dtype, op_id), (old, None)): + found = mma.op_for_shape( + family="wmma", + a_dtype=spelling, + b_dtype=spelling, + c_dtype="i32", + m=16, + n=16, + k=16, + ) + assert (found and found.op_id) == want diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_dtypes.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_dtypes.cpp index 70be3b998fd5..81d3694840d2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_dtypes.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_dtypes.cpp @@ -23,6 +23,7 @@ int main() {"int32", ROCKE_DTYPE_I32}, {"XF32", "tf32"}, {" TF32 ", "tf32"}, + {"INT8", "i8"}, {" Custom_Format ", "custom_format"}, {"", ""}}; for(const auto& entry : cases) @@ -39,5 +40,10 @@ int main() if(strcmp(rocke_normalize_dtype(NULL, scratch, sizeof(scratch)), "") != 0 || strcmp(rocke_normalize_dtype(" BF6 ", NULL, 0), "fp6e3m2") != 0) return 1; + /* "iu8"/"iu4" appear only in RDNA integer WMMA instruction names, so they are + * not dtypes: signedness is an instruction operand, not a property of the data. */ + const rocke_dtype_info_t* i4 = rocke_dtype_info("i4"); + if(!i4 || i4->encoded_bits != 4 || rocke_dtype_info("iu8") || rocke_dtype_info("iu4")) + return 1; return 0; } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_dtypes.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_dtypes.py index a27fa6fd5163..aab71559badb 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 @@ -60,11 +60,23 @@ def test_quant_type_recognition_does_not_enable_integer_conversion(dtype): dequantize_scalar_to_f32(b, value, scale=scale) -def test_unknown_dtype_and_unrepresented_integer_family(): +def test_unknown_dtype_is_rejected(): 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") + + +def test_the_integer_dtypes_are_i8_and_i4_only(): + # "iu8" / "iu4" appear only in RDNA integer WMMA instruction names: signedness is + # an instruction operand, not a property of the data, so they are not dtypes. + from rocke.core.ir import I4 + + assert dtype_to_ir_type("i8") is I8 + assert dtype_to_ir_type("i4") is I4 + assert dtype_info("i4").encoded_bits == 4 + assert storage_ir_type("i4") is I8 + for name in ("iu8", "iu4"): + with pytest.raises(ValueError, match="unknown dtype"): + dtype_info(name) @pytest.mark.parametrize( diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_catalog_lookup.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_catalog_lookup.cpp index fc94e77f4f1e..5d36f50d364d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_catalog_lookup.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_catalog_lookup.cpp @@ -690,13 +690,40 @@ static int test_fp32_never_reaches_an_xf32_atom() return 0; } +/* Mirrors test_rdna_integer_wmma_rows_name_i8_and_i4: the instruction names keep + * "iu8"/"iu4", the operand dtypes are i8/i4, and only i8/i4 select the op. */ +static int test_rdna_integer_wmma_rows_name_i8_and_i4() +{ + const char* rows[][3] + = {{"wmma_i32_16x16x16_iu8", "i8", "iu8"}, {"wmma_i32_16x16x16_iu4", "i4", "iu4"}}; + for(const char* gfx : {"gfx1151", "gfx11-generic"}) + { + const auto* arch = rocke_arch_target_from_gfx(gfx); + CHECK(arch); + for(const auto& row : rows) + { + const auto* op = rocke_mma_catalog_by_op_id(&arch->mma, row[0]); + CHECK(op && strcmp(op->a_dtype, row[1]) == 0 && strcmp(op->b_dtype, row[1]) == 0 + && strcmp(op->c_dtype, "i32") == 0); + CHECK(rocke_mma_catalog_op_for_shape( + &arch->mma, "wmma", row[1], row[1], "i32", 16, 16, 16) + == op); + CHECK(rocke_mma_catalog_op_for_shape( + &arch->mma, "wmma", row[2], row[2], "i32", 16, 16, 16) + == NULL); + } + } + return 0; +} + int main() { // Keep this first: no previous query may prime the shared family index. if(test_family_index_first_use() || test_family_index_duplicates() || test_mma_result_names() || test_scale_contracts() || test_scale_layouts_and_families() || test_scaled_wmma_supported_combinations() || test_emulated_rows_mirror_physical() - || test_pre_existing_selection_is_native_only() || test_fp32_never_reaches_an_xf32_atom()) + || test_pre_existing_selection_is_native_only() || test_fp32_never_reaches_an_xf32_atom() + || test_rdna_integer_wmma_rows_name_i8_and_i4()) return 1; int checked = 0; for(const char* gfx : {"gfx950", "gfx1250"}) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_packed_only.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_packed_only.py new file mode 100644 index 000000000000..5237f02cf9d0 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_packed_only.py @@ -0,0 +1,72 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Packed-only types have no IR value. + +I4 values only ever live packed in 32-bit words: a kernel keeps the words in +i32 storage and unpacks them explicitly. So the IR builder refuses every op +result, kernel parameter and LDS tile of i4 elements, and the lowerings, which +have no rendering for the type, still refuse one that reaches them. +""" + +import pytest + +from rocke.core.ir import I32, IRBuilder, Param, PtrType, VectorType, dtype_to_ir_type +from rocke.tests.engines import ( + assert_both_engines_reject, + assert_both_engines_reject_hip, +) + +PACKED_ONLY_DTYPES = ["i4"] + +# Builders that would make an IR value, or an LDS tile, of the element type +# they are given, with the IR op each refusal names. +_BUILDERS = { + "param": ("param", lambda b, t: b.param("X", t)), + "param_vector": ("param", lambda b, t: b.param("X", VectorType(t, 8))), + "zero_vec": ("zero_vec", lambda b, t: b.zero_vec(t, 8)), + "undef_vec": ("vector.undef", lambda b, t: b.undef_vec(t, 8)), + "trunc": ("arith.trunc", lambda b, t: b.trunc(b.const_i32(0), t)), + "vec_bitcast": ( + "vector.bitcast", + lambda b, t: b.vec_bitcast(b.zero_vec(I32, 1), VectorType(t, 8)), + ), + "global_load": ( + "memref.global_load_typed", + lambda b, t: b.global_load( + b.param("P", PtrType(t, "global")), b.const_i32(0), t + ), + ), + "smem_alloc": ("tile.smem_alloc", lambda b, t: b.smem_alloc(t, (64,))), +} + + +def _refusal(op: str, dtype: str) -> str: + return ( + f"{op}: {dtype} values travel packed in 32-bit words and have no scalar " + "or vector form -- keep them in i32 storage and unpack explicitly" + ) + + +@pytest.mark.parametrize("dtype", PACKED_ONLY_DTYPES) +@pytest.mark.parametrize("builder", sorted(_BUILDERS)) +def test_builder_refuses_a_packed_only_value(builder, dtype): + op, build = _BUILDERS[builder] + b = IRBuilder("packed_only") + with pytest.raises(TypeError) as error: + build(b, dtype_to_ir_type(dtype)) + assert str(error.value) == _refusal(op, dtype) + + +@pytest.mark.parametrize("dtype", PACKED_ONLY_DTYPES) +def test_lowerings_refuse_a_packed_only_value(dtype): + # The builder refuses the parameter, so the kernel is assembled by hand to + # reach the lowerings' own refusal: the type has no LLVM or HIP spelling. + b = IRBuilder(f"{dtype}_value") + b.ret() + b.kernel.params.append(Param(name="X", type=dtype_to_ir_type(dtype))) + assert_both_engines_reject( + b.kernel, "gfx90a", NotImplementedError, f"no LLVM mapping for type {dtype}" + ) + assert_both_engines_reject_hip( + b.kernel, "gfx90a", NotImplementedError, f"no HIP mapping for type {dtype}" + ) 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 c7f39ee4cb6e..64bc401dfefb 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 @@ -354,7 +354,8 @@ int main(int argc, char** argv) 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); + CHECK(rocke_lower_kernel_to_hip(&b, b.kernel, &opts, &text) == ROCKE_ERR_NOTIMPL); + CHECK(strstr(rocke_ir_builder_error(&b), "no HIP mapping for type vec<")); CHECK(text.len > 0); // The prologue is assembled before the vector walk. rocke_strbuf_free(&text); rocke_ir_builder_free(&b); From 6f747ce80ac0eee7f49cb2bbc7f8073895a847f3 Mon Sep 17 00:00:00 2001 From: Chris Millette Date: Thu, 8 Oct 2026 19:38:55 -0400 Subject: [PATCH 05/11] feat(rocke): add the unsigned U8 and U4 types U8 and U4 name unsigned 8- and 4-bit integer data (spellings u8/uint8 and u4/uint4), exported next to I8 and I4. They let a kernel say its data is unsigned so that matrix instructions with a per-operand signedness flag can be selected for it. U8 renders as LLVM i8, since LLVM integers carry no sign, and as HIP uint8_t. U4, like I4, is packed-only: the IR builder of both engines refuses a u4 value or LDS tile with the I4 refusal's text (_refuse_packed_only_value, rocke_i_refuse_packed_only_value), and it has no LLVM or HIP rendering, so no u4 value reaches an arithmetic op or a compare. U8 is an arithmetic type. One predicate per engine says whether a type is unsigned (_is_unsigned in core/ir.py, rocke_type_is_unsigned in the C++ engine), and every op whose result depends on signedness asks it which form to emit for u8: div and mod lower to udiv and urem, the ordered compares of arith.cmp and vector.cmp to icmp ult/ule/ugt/uge, and add, sub and mul drop the nsw flag, since unsigned values wrap. zext, lshr and the sign-agnostic ops lower as for i8. Three ops are new: umax and umin (llvm.umax.i8 / llvm.umin.i8, measured into the llvm20, llvm22 and llvm23 arch-domain artifacts with those toolchains), uitofp_f32, and cvt_f32_to_u8_sat, which rounds to nearest even and clamps to [0, 255] in f32 before fptoui, so a NaN gives 0. smax, smin, sext and their vector forms are signed by definition. Both builders refuse a u8 operand when one of them is built, vector_binary("smax" | "smin") included, naming the unsigned op to use where there is one. The LLVM lowerings of both engines, and their HIP lowerings of the scalar ops (HIP has no u8 vector type), refuse it again with the same text under the op's dotted name (_refuse_unsigned_operands, rocke_i_refuse_unsigned_operands), so IR that did not come through a builder cannot run a signed op on unsigned data. The new opcodes are appended to rocke_opcode_t, so no existing opcode value changes. In HIP every u8 value is a uint8_t, so C's own operators are the unsigned ones, and umax / umin share smax / smin's ternary without their refusal; zext casts a u8 source through uint8_t. dtype_support.md describes what u8 supports. A u8 LDS element takes one byte, like i8: the lowerers' element-size and alignment tables list u8 next to i8, so a u8 tile and its scalar stores lower exactly as the i8 ones do. smem_load_vN and smem_store_vN admit one shared list of vector element types (_SMEM_VECTOR_ELEMENTS in core/ir.py, ROCKE_SMEM_VECTOR_ELEMENTS in ir_tile.cpp), refused with one message generated from it. smem_store_vN had no element check and sized an unlisted element as two bytes; it now refuses a u8 vector as smem_load_vN and the global vector builders do. Both engines make every change. Tests lower a U8 copy and a U8 LDS store with both engines, to LLVM IR (the LDS store matches the I8 kernel) and to HIP source (compiled with hipcc when present). test_u8_arith.py lowers every u8 op with both engines and checks the unsigned LLVM forms and the HIP types, and on a gfx90a GPU checks division, remainder, the compares, umax/umin, lshr, wrapping add/sub/mul, uitofp, zext and the saturating conversion bit-exact against numpy uint8, on values above 127. Tests also check that vector_binary refuses smax and smin on u8, and that both engines refuse each signed op on u8 in deserialized IR, in the LLVM lowering and, for the scalar ops, the HIP lowering. test_packed_only.py checks the builders' refusal of u4 values, and that both engines refuse a u4 parameter in both lowerings. A new ctest, rocke_ir_types, builds LDS vector accesses, i4 and u4 values and the u8 ops and refusals with the C++ builders; tests compare its answers with the Python builders, and tests/run_all.py passes it in as ROCKE_IR_TYPES_TEST. Co-Authored-By: Claude --- ...ROCKE_ENGINE_test_categories_external.yaml | 2 + .../rocke/platform/cpp/core/dtypes.cpp | 7 + .../platform/cpp/core/ir/core_builder.cpp | 26 +- .../rocke/platform/cpp/core/ir/core_types.cpp | 33 +- .../rocke/platform/cpp/core/ir/ir_arith.cpp | 75 +++ .../rocke/platform/cpp/core/ir/ir_mem.cpp | 8 +- .../rocke/platform/cpp/core/ir/ir_tile.cpp | 49 +- .../platform/cpp/core/lower_hip/core.cpp | 4 + .../cpp/core/lower_hip/lower_hip_arith.cpp | 53 ++- .../cpp/core/lower_hip/lower_hip_cast.cpp | 37 +- .../platform/cpp/core/lower_llvm/arith.cpp | 94 ++-- .../platform/cpp/core/lower_llvm/control.cpp | 34 +- .../platform/cpp/core/lower_llvm/convert.cpp | 35 ++ .../platform/cpp/core/lower_llvm/core.cpp | 37 +- .../platform/cpp/core/lower_llvm/data.cpp | 3 + .../platform/cpp/core/lower_llvm/mem.cpp | 7 +- .../rocke/platform/cpp/include/rocke/ir.h | 23 + .../platform/cpp/include/rocke/ir_internal.h | 15 +- .../cpp/include/rocke/lower_llvm_internal.h | 9 + .../dsl_docs/reference/dtype_support.md | 31 +- .../rocke/platform/python/rocke/__init__.py | 4 + .../platform/python/rocke/core/__init__.py | 4 + .../data/intrinsic_arch_domain.llvm20.json | 72 +++ .../data/intrinsic_arch_domain.llvm22.json | 70 +++ .../data/intrinsic_arch_domain.llvm23.json | 70 +++ .../platform/python/rocke/core/dtypes.py | 7 + .../rocke/platform/python/rocke/core/ir.py | 155 ++++++- .../platform/python/rocke/core/lower_hip.py | 47 +- .../platform/python/rocke/core/lower_llvm.py | 133 ++++-- .../rocke/platform/python/rocke/tests/gpu.py | 1 + .../rocke/platform/tests/CMakeLists.txt | 9 + .../rocke/platform/tests/README.md | 12 +- .../rocke/platform/tests/core/test_dtypes.cpp | 10 + .../rocke/platform/tests/core/test_dtypes.py | 36 ++ .../platform/tests/core/test_ir_types.cpp | 433 ++++++++++++++++++ .../platform/tests/core/test_packed_only.py | 37 +- .../platform/tests/core/test_u8_arith.py | 385 ++++++++++++++++ .../platform/tests/core/test_u8_memory.py | 141 ++++++ .../rocke/platform/tests/run_all.py | 1 + .../platform/tests/tools/test_run_all.py | 7 +- 40 files changed, 2032 insertions(+), 184 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_u8_arith.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_u8_memory.py diff --git a/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml b/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml index 9a7d95abbfc9..69dd87872ba9 100644 --- a/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml +++ b/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml @@ -34,6 +34,7 @@ test_categories: - "rocke_wmma_fmha_decode_host" - "rocke_future_intrinsic_lowering" - "rocke_ir_types" + - "rocke_opcode_values" - "rocke_tiled_attention_2d_reentrancy" - "rocke_unified_attention_scalar_lower_to_llvm" - "rocke_python_import_smoke" @@ -66,6 +67,7 @@ test_categories: - "rocke_wmma_fmha_decode_host" - "rocke_future_intrinsic_lowering" - "rocke_ir_types" + - "rocke_opcode_values" - "rocke_tiled_attention_2d_reentrancy" - "rocke_unified_attention_scalar_lower_to_llvm" - "rocke_python_import_smoke" 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 06b38c2e96d0..8860bec148e1 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 @@ -60,6 +60,11 @@ static const rocke_dtype_alias_t k_dtype_aliases[] = { {"int4", "i4"}, {"i32", ROCKE_DTYPE_I32}, {"int32", ROCKE_DTYPE_I32}, + /* Unsigned integers (logical types; no unsigned arithmetic yet). */ + {"u8", "u8"}, + {"uint8", "u8"}, + {"u4", "u4"}, + {"uint4", "u4"}, }; #define K_NUM_DTYPE_ALIASES ((int)(sizeof(k_dtype_aliases) / sizeof(k_dtype_aliases[0]))) @@ -163,6 +168,8 @@ static const rocke_dtype_info_t k_dtype_info[] = { /* Integers. */ {"i4", ROCKE_DTYPE_INTEGER, 4}, {"i8", ROCKE_DTYPE_INTEGER, 8}, + {"u4", ROCKE_DTYPE_INTEGER, 4}, + {"u8", ROCKE_DTYPE_INTEGER, 8}, {"i16", ROCKE_DTYPE_INTEGER, 16}, {"i32", ROCKE_DTYPE_INTEGER, 32}, {"i64", ROCKE_DTYPE_INTEGER, 64}, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_builder.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_builder.cpp index a480ef1fea05..8252ffaeb55e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_builder.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_builder.cpp @@ -654,7 +654,7 @@ void rocke_i_refuse_packed_only_value(const char* op_name, const rocke_type_t* t { return; } - if(elem->scalar == ROCKE_SCALAR_I4) + if(elem->scalar == ROCKE_SCALAR_I4 || elem->scalar == ROCKE_SCALAR_U4) { throw ckc::TypeError( ckc::format_error("%s: %s values travel packed in 32-bit words and have no scalar or " @@ -664,6 +664,30 @@ void rocke_i_refuse_packed_only_value(const char* op_name, const rocke_type_t* t } } +void rocke_i_refuse_unsigned(const char* op_name, const rocke_type_t* t, const char* instead) +{ + if(!rocke_type_is_unsigned(t)) + { + return; + } + const rocke_type_t* elem = rocke_i_elem_of(t); + std::string message = ckc::format_error( + "%s treats its operands as signed, but %s is unsigned", op_name, elem->name); + if(instead != NULL) + { + message += std::string(" -- use ") + instead; + } + throw ckc::TypeError(message); +} + +void rocke_i_refuse_unsigned_operands(const rocke_op_t* op, const char* instead) +{ + for(int i = 0; i < op->num_operands; ++i) + { + rocke_i_refuse_unsigned(rocke_opcode_name(op->opcode), op->operands[i]->type, instead); + } +} + /* ============================== BUILDER ================================= */ rocke_status_t rocke_ir_builder_init(rocke_ir_builder_t* b, const char* kernel_name) 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 8259a63f0196..115ee05c25f5 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 @@ -48,6 +48,11 @@ ROCKE_SCALAR_SINGLETON(rocke_i1, ROCKE_SCALAR_I1, "i1") /* Logical only, like fp4e2m1: no LLVM rendering; its values travel packed. */ ROCKE_SCALAR_SINGLETON(rocke_i4, ROCKE_SCALAR_I4, "i4") +/* Unsigned data. u8 renders as LLVM i8 (LLVM integers carry no sign), so every op + * whose result depends on signedness asks rocke_type_is_unsigned which form to + * emit; u4, like i4, is logical only and travels packed. */ +ROCKE_SCALAR_SINGLETON(rocke_u4, ROCKE_SCALAR_U4, "u4") +ROCKE_SCALAR_SINGLETON(rocke_u8, ROCKE_SCALAR_U8, "u8") ROCKE_SCALAR_SINGLETON(rocke_i8, ROCKE_SCALAR_I8, "i8") ROCKE_SCALAR_SINGLETON(rocke_i16, ROCKE_SCALAR_I16, "i16") ROCKE_SCALAR_SINGLETON(rocke_i32, ROCKE_SCALAR_I32, "i32") @@ -76,6 +81,10 @@ const rocke_type_t* rocke_scalar_by_name(const char* name) return rocke_i1(); if(strcmp(name, "i4") == 0) return rocke_i4(); + if(strcmp(name, "u4") == 0) + return rocke_u4(); + if(strcmp(name, "u8") == 0) + return rocke_u8(); if(strcmp(name, "i8") == 0) return rocke_i8(); if(strcmp(name, "i16") == 0) @@ -255,6 +264,14 @@ bool rocke_type_eq(const rocke_type_t* a, const rocke_type_t* b) return strcmp(a->name, b->name) == 0; } +/* Python core/ir.py::_is_unsigned. */ +bool rocke_type_is_unsigned(const rocke_type_t* t) +{ + const rocke_type_t* elem = (t != NULL && t->kind == ROCKE_TYPE_VECTOR) ? t->elem : t; + return elem != NULL && elem->kind == ROCKE_TYPE_SCALAR + && (elem->scalar == ROCKE_SCALAR_U8 || elem->scalar == ROCKE_SCALAR_U4); +} + /* ============================== ATTR MAP ================================ */ void rocke_attr_map_init(rocke_attr_map_t* m) @@ -712,7 +729,13 @@ static const char* const rocke_opcode_names[ROCKE_OP__COUNT] = { /* arith.* f32 <-> f64 conversions */ "arith.cvt_f32_to_f64", - "arith.cvt_f64_to_f32"}; + "arith.cvt_f64_to_f32", + + /* arith.* unsigned (u8) ops */ + "arith.umax", + "arith.umin", + "arith.uitofp_f32", + "arith.cvt_f32_to_u8_sat"}; const char* rocke_opcode_name(rocke_opcode_t op) { @@ -985,7 +1008,13 @@ static const bool rocke_opcode_pure[ROCKE_OP__COUNT] = { /* arith.* f32 <-> f64 conversions */ /* arith.cvt_f32_to_f64 */ true, - /* arith.cvt_f64_to_f32 */ true}; + /* arith.cvt_f64_to_f32 */ true, + + /* arith.* unsigned (u8) ops */ + /* arith.umax */ true, + /* arith.umin */ true, + /* arith.uitofp_f32 */ true, + /* arith.cvt_f32_to_u8_sat */ true}; bool rocke_opcode_is_pure(rocke_opcode_t op) { 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 46a28e732a4b..cad405edaeea 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 @@ -182,14 +182,57 @@ rocke_value_t* rocke_b_lnot(rocke_ir_builder_t* b, rocke_value_t* a) rocke_value_t* rocke_b_smax(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_t* c) { + if(a && c) + { + rocke_i_refuse_unsigned("smax", a->type, "umax"); + rocke_i_refuse_unsigned("smax", c->type, "umax"); + } return rocke_i_binop(b, ROCKE_OP_ARITH_SMAX, a, c, "smax"); } rocke_value_t* rocke_b_smin(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_t* c) { + if(a && c) + { + rocke_i_refuse_unsigned("smin", a->type, "umin"); + rocke_i_refuse_unsigned("smin", c->type, "umin"); + } return rocke_i_binop(b, ROCKE_OP_ARITH_SMIN, a, c, "smin"); } +/* umax / umin take two scalar operands of one unsigned type (u8). */ +static rocke_value_t* unsigned_min_max(rocke_ir_builder_t* b, + rocke_opcode_t opcode, + const char* op_name, + rocke_value_t* a, + rocke_value_t* c) +{ + if(!rocke_i_live(b)) + return NULL; + if(!a || !c) + return (rocke_value_t*)rocke_i_set_err(b, ROCKE_ERR_VALUE, "%s NULL operand", op_name); + if(a->type->kind == ROCKE_TYPE_VECTOR || !rocke_type_is_unsigned(a->type) + || !rocke_type_eq(a->type, c->type)) + return (rocke_value_t*)rocke_i_set_err( + b, + ROCKE_ERR_VALUE, + "%s expects two operands of one unsigned integer type, got %s and %s", + op_name, + a->type->name, + c->type->name); + return rocke_i_binop(b, opcode, a, c, op_name); +} + +rocke_value_t* rocke_b_umax(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_t* c) +{ + return unsigned_min_max(b, ROCKE_OP_ARITH_UMAX, "umax", a, c); +} + +rocke_value_t* rocke_b_umin(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_t* c) +{ + return unsigned_min_max(b, ROCKE_OP_ARITH_UMIN, "umin", a, c); +} + rocke_value_t* rocke_b_xor(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_t* c) { return rocke_i_binop(b, ROCKE_OP_ARITH_XOR, a, c, "xor"); @@ -512,6 +555,7 @@ rocke_value_t* rocke_b_sext(rocke_ir_builder_t* b, rocke_value_t* v, const rocke return (rocke_value_t*)rocke_i_set_err(b, ROCKE_ERR_VALUE, "sext NULL value"); if(!target) return (rocke_value_t*)rocke_i_set_err(b, ROCKE_ERR_VALUE, "sext NULL target"); + rocke_i_refuse_unsigned("sext", v->type, "zext"); operands[0] = v; return rocke_i_op1(b, ROCKE_OP_ARITH_SEXT, operands, 1, target, NULL, "sx"); } @@ -702,6 +746,23 @@ rocke_value_t* rocke_b_sitofp_f32(rocke_ir_builder_t* b, rocke_value_t* v) return rocke_i_op1(b, ROCKE_OP_ARITH_SITOFP_F32, operands, 1, rocke_f32(), NULL, "sitof"); } +rocke_value_t* rocke_b_uitofp_f32(rocke_ir_builder_t* b, rocke_value_t* v) +{ + rocke_value_t* operands[1]; + if(!rocke_i_live(b)) + return NULL; + if(!v) + return (rocke_value_t*)rocke_i_set_err(b, ROCKE_ERR_VALUE, "uitofp_f32 NULL value"); + if(v->type->kind == ROCKE_TYPE_VECTOR || !rocke_type_is_unsigned(v->type)) + return (rocke_value_t*)rocke_i_set_err( + b, + ROCKE_ERR_VALUE, + "uitofp_f32 expects an unsigned integer input, got %s", + v->type->name); + operands[0] = v; + return rocke_i_op1(b, ROCKE_OP_ARITH_UITOFP_F32, operands, 1, rocke_f32(), NULL, "uitof"); +} + /* f32 <-> f64 has its own opcode per direction rather than riding on * cast_to_f32 / cast_f32_to: those two are the f16/bf16 <-> f32 hub and keep * rejecting f64, so no caller widens or narrows a double by accident. */ @@ -902,6 +963,20 @@ rocke_value_t* rocke_b_cvt_f32_to_i8_sat(rocke_ir_builder_t* b, rocke_value_t* v return rocke_i_op1(b, ROCKE_OP_ARITH_CVT_F32_TO_I8_SAT, operands, 1, rocke_i8(), NULL, "qi8"); } +rocke_value_t* rocke_b_cvt_f32_to_u8_sat(rocke_ir_builder_t* b, rocke_value_t* v) +{ + rocke_value_t* operands[1]; + if(!rocke_i_live(b)) + return NULL; + if(!v) + return (rocke_value_t*)rocke_i_set_err(b, ROCKE_ERR_VALUE, "cvt_f32_to_u8_sat NULL value"); + if(!rocke_i_type_is(v->type, "f32")) + return (rocke_value_t*)rocke_i_set_err( + b, ROCKE_ERR_VALUE, "cvt_f32_to_u8_sat expects f32 input, got %s", v->type->name); + operands[0] = v; + return rocke_i_op1(b, ROCKE_OP_ARITH_CVT_F32_TO_U8_SAT, operands, 1, rocke_u8(), NULL, "qu8"); +} + rocke_value_t* rocke_b_cvt_pk_fp8_f32x4(rocke_ir_builder_t* b, rocke_value_t* v) { rocke_value_t* operands[1]; 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 40a13f5fe1c0..346c1e1b1172 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 @@ -883,7 +883,8 @@ void rocke_b_global_atomic_add_f32(rocke_ir_builder_t* b, /* ===================== vector.* element-wise family ===================== */ -/* IRBuilder.vector_binary: matching vector operands -> a->type. */ +/* IRBuilder.vector_binary: matching vector operands -> a->type. The signed + * ops (smax, smin) refuse an unsigned element type. */ static rocke_value_t* rocke_i_vector_binary(rocke_ir_builder_t* b, rocke_opcode_t opcode, rocke_value_t* a, @@ -898,6 +899,10 @@ static rocke_value_t* rocke_i_vector_binary(rocke_ir_builder_t* b, if(!rocke_i_is_vector(a->type, NULL, -1) || !rocke_type_eq(a->type, c->type)) return (rocke_value_t*)rocke_i_set_err( b, ROCKE_ERR_VALUE, "vector_binary expects matching vector operands"); + if(opcode == ROCKE_OP_VECTOR_SMAX) + rocke_i_refuse_unsigned("vector_smax", a->type, NULL); + if(opcode == ROCKE_OP_VECTOR_SMIN) + rocke_i_refuse_unsigned("vector_smin", a->type, NULL); ops[0] = a; ops[1] = c; return rocke_i_op1(b, opcode, ops, 2, a->type, NULL, hint); @@ -1130,6 +1135,7 @@ rocke_value_t* if(!rocke_i_is_vector(v->type, NULL, -1)) return (rocke_value_t*)rocke_i_set_err( b, ROCKE_ERR_VALUE, "vector_sext expects vector input"); + rocke_i_refuse_unsigned("vector_sext", v->type, NULL); vt = rocke_vector_type(b, target, rocke_i_count_of(v->type)); if(!vt) return NULL; 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 6787e4e033b3..1b0a1da5a44e 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 @@ -112,6 +112,41 @@ static int rocke_elem_bytes_name(const char* elem_name) return 2; } +/* Element types an LDS vector access (smem_load_vN and smem_store_vN) admits, + * in the order its refusal lists them (Python core/ir.py + * _SMEM_VECTOR_ELEMENTS). One list for both directions keeps the pair + * symmetric: whatever one stores as a vector can be loaded back. */ +static const char* const ROCKE_SMEM_VECTOR_ELEMENTS[] + = {"f16", "bf16", "i16", "f32", "i32", "tf32", "f64", "fp8e4m3", "bf8e5m2", "i8"}; + +/* Python _require_smem_vector_element: refuse an LDS vector `access` of an + * element outside ROCKE_SMEM_VECTOR_ELEMENTS, listing the admitted types. */ +static void + rocke_require_smem_vector_elem(rocke_ir_builder_t* b, const char* access, const char* elem) +{ + char admitted[128]; + size_t used = 0; + for(const char* name : ROCKE_SMEM_VECTOR_ELEMENTS) + { + if(strcmp(elem, name) == 0) + { + return; + } + } + admitted[0] = '\0'; + for(const char* name : ROCKE_SMEM_VECTOR_ELEMENTS) + { + const int n + = snprintf(admitted + used, sizeof admitted - used, "%s%s", used ? " / " : "", name); + if(n < 0 || (size_t)n >= sizeof admitted - used) + { + break; /* cannot happen: the list fits; stop rather than overrun */ + } + used += (size_t)n; + } + rocke_i_set_err(b, ROCKE_ERR_VALUE, "%s supports %s, got %s", access, admitted, elem); +} + static bool rocke_n_in(int n, const int* allowed, int count) { int i; @@ -276,6 +311,7 @@ void rocke_b_smem_store_vN(rocke_ir_builder_t* b, return; } elem_name = value->type->elem->name; + rocke_require_smem_vector_elem(b, "smem_store_vN", elem_name); { const int* allowed = allowed_other; int acount = 3; @@ -380,18 +416,7 @@ rocke_value_t* rocke_b_smem_load_vN(rocke_ir_builder_t* b, return (rocke_value_t*)rocke_i_set_err(b, ROCKE_ERR_VALUE, "smem_load_vN dtype is NULL"); } dn = dtype->name; - if(!(strcmp(dn, "f16") == 0 || strcmp(dn, "bf16") == 0 || strcmp(dn, "i16") == 0 - || strcmp(dn, "f32") == 0 || strcmp(dn, "tf32") == 0 || strcmp(dn, "i32") == 0 - || strcmp(dn, "f64") == 0 || strcmp(dn, "fp8e4m3") == 0 || strcmp(dn, "bf8e5m2") == 0 - || strcmp(dn, "i8") == 0)) - { - return (rocke_value_t*)rocke_i_set_err( - b, - ROCKE_ERR_VALUE, - "smem_load_vN supports f16 / bf16 / i16 / f32 / i32 / tf32 / f64 / " - "fp8e4m3 / bf8e5m2 / i8, got %s", - dn); - } + rocke_require_smem_vector_elem(b, "smem_load_vN", dn); { const int* allowed = allowed_32bit; int acount = 5; 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 5e643761d02e..8d33c2e15930 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 @@ -296,6 +296,10 @@ const char* rocke_h_hip_scalar(const char* ir_scalar_name) { return "int8_t"; } + if(strcmp(ir_scalar_name, "u8") == 0) + { + return "uint8_t"; + } if(strcmp(ir_scalar_name, "i16") == 0) { return "int16_t"; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_arith.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_arith.cpp index e7caab9e0684..47f912b449f3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_arith.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_arith.cpp @@ -20,6 +20,7 @@ #include /* strcmp */ #include "rocke/ir.h" +#include "rocke/ir_internal.h" #include "rocke/lower_hip.h" #include "rocke/lower_hip_internal.h" @@ -358,11 +359,13 @@ static rocke_status_t rocke_h_op_arith_or(rocke_h_lowerer_t* lw, const rocke_op_ return lw->status; } -/* def _op_arith_smax(self, op): +/* def _select_by_compare(self, op, compare): * a, b = op.operands * self._emit(f"{_type_to_hip(op.result.type)} {_name(op.result)} = " - * f"({_name(a)} > {_name(b)} ? {_name(a)} : {_name(b)});") */ -static rocke_status_t rocke_h_op_arith_smax(rocke_h_lowerer_t* lw, const rocke_op_t* op) + * f"({_name(a)} {compare} {_name(b)} ? {_name(a)} : {_name(b)});") + * The maximum (">") or minimum ("<") of the two operands, as a C ternary. */ +static rocke_status_t + h_select_by_compare(rocke_h_lowerer_t* lw, const rocke_op_t* op, const char* compare) { const rocke_value_t* a = op->operands[0]; const rocke_value_t* b = op->operands[1]; @@ -370,36 +373,40 @@ static rocke_status_t rocke_h_op_arith_smax(rocke_h_lowerer_t* lw, const rocke_o const char* an = rocke_h_name(lw, a); const char* bn = rocke_h_name(lw, b); rocke_h_emitf(lw, - "%s %s = (%s > %s ? %s : %s);", + "%s %s = (%s %s %s ? %s : %s);", rocke_h_type_to_hip(lw, r->type), rocke_h_name(lw, r), an, + compare, bn, an, bn); return lw->status; } -/* def _op_arith_smin(self, op): - * a, b = op.operands - * self._emit(f"{_type_to_hip(op.result.type)} {_name(op.result)} = " - * f"({_name(a)} < {_name(b)} ? {_name(a)} : {_name(b)});") */ +/* smax / smin are signed, so they refuse a u8 operand; umax / umin take u8. + * uint8_t operands promote to int, so C's own comparison and the ternary are + * the unsigned maximum / minimum. */ +static rocke_status_t rocke_h_op_arith_smax(rocke_h_lowerer_t* lw, const rocke_op_t* op) +{ + rocke_i_refuse_unsigned_operands(op, "arith.umax"); + return h_select_by_compare(lw, op, ">"); +} + static rocke_status_t rocke_h_op_arith_smin(rocke_h_lowerer_t* lw, const rocke_op_t* op) { - const rocke_value_t* a = op->operands[0]; - const rocke_value_t* b = op->operands[1]; - const rocke_value_t* r = h_res(op); - const char* an = rocke_h_name(lw, a); - const char* bn = rocke_h_name(lw, b); - rocke_h_emitf(lw, - "%s %s = (%s < %s ? %s : %s);", - rocke_h_type_to_hip(lw, r->type), - rocke_h_name(lw, r), - an, - bn, - an, - bn); - return lw->status; + rocke_i_refuse_unsigned_operands(op, "arith.umin"); + return h_select_by_compare(lw, op, "<"); +} + +static rocke_status_t rocke_h_op_arith_umax(rocke_h_lowerer_t* lw, const rocke_op_t* op) +{ + return h_select_by_compare(lw, op, ">"); +} + +static rocke_status_t rocke_h_op_arith_umin(rocke_h_lowerer_t* lw, const rocke_op_t* op) +{ + return h_select_by_compare(lw, op, "<"); } /* def _op_arith_not(self, op): @@ -871,6 +878,8 @@ const rocke_h_handler_entry_t* rocke_h_handlers_arith(void) {ROCKE_OP_ARITH_OR, rocke_h_op_arith_or}, {ROCKE_OP_ARITH_SMAX, rocke_h_op_arith_smax}, {ROCKE_OP_ARITH_SMIN, rocke_h_op_arith_smin}, + {ROCKE_OP_ARITH_UMAX, rocke_h_op_arith_umax}, + {ROCKE_OP_ARITH_UMIN, rocke_h_op_arith_umin}, {ROCKE_OP_ARITH_NOT, rocke_h_op_arith_not}, {ROCKE_OP_ARITH_XOR, rocke_h_op_arith_xor}, {ROCKE_OP_ARITH_SHL, rocke_h_op_arith_shl}, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_cast.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_cast.cpp index 03150adab205..4d2e75746cce 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_cast.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_hip/lower_hip_cast.cpp @@ -16,6 +16,7 @@ * here rocke_h_emit/rocke_h_emitf does the same). */ #include "rocke/ir.h" +#include "rocke/ir_internal.h" #include "rocke/lower_hip.h" #include "rocke/lower_hip_internal.h" @@ -39,7 +40,7 @@ static const char* h_unsigned_int_type(const rocke_type_t* t) return NULL; if(__builtin_strcmp(name, "i1") == 0) return "bool"; - if(__builtin_strcmp(name, "i8") == 0) + if(__builtin_strcmp(name, "i8") == 0 || __builtin_strcmp(name, "u8") == 0) return "uint8_t"; if(__builtin_strcmp(name, "i16") == 0) return "uint16_t"; @@ -73,6 +74,7 @@ static rocke_status_t rocke_h_op_arith_zext(rocke_h_lowerer_t* lw, const rocke_o /* def _op_arith_sext(self, op): -- identical lowering to zext in HIP. */ static rocke_status_t rocke_h_op_arith_sext(rocke_h_lowerer_t* lw, const rocke_op_t* op) { + rocke_i_refuse_unsigned_operands(op, "arith.zext"); const rocke_value_t* v = op->operands[0]; const rocke_value_t* r = h_res(op); const char* t = rocke_h_type_to_hip(lw, r->type); @@ -204,6 +206,18 @@ static rocke_status_t rocke_h_op_arith_sitofp_f32(rocke_h_lowerer_t* lw, const r return lw->status; } +/* def _op_arith_uitofp_f32(self, op): + * # u8 -> f32: a uint8_t converts exactly. + * (v,) = op.operands + * self._emit(f"float {_name(op.result)} = (float){_name(v)};") */ +static rocke_status_t rocke_h_op_arith_uitofp_f32(rocke_h_lowerer_t* lw, const rocke_op_t* op) +{ + const rocke_value_t* v = op->operands[0]; + const rocke_value_t* r = h_res(op); + rocke_h_emitf(lw, "float %s = (float)%s;", rocke_h_name(lw, r), rocke_h_name(lw, v)); + return lw->status; +} + /* ----------------------------- arith: fp8/bf8 scalar cvt ------------------- */ /* def _op_arith_cvt_fp8_to_f32(self, op): @@ -296,6 +310,25 @@ static rocke_status_t rocke_h_op_arith_cvt_f32_to_i8_sat(rocke_h_lowerer_t* lw, return lw->status; } +/* def _op_arith_cvt_f32_to_u8_sat(self, op): + * (v,) = op.operands + * rounded = f"{_name(op.result)}_r" + * self._emit(f"float {rounded} = rintf({_name(v)});") + * self._emit(f"uint8_t {_name(op.result)} = " + * f"(uint8_t)fminf(fmaxf({rounded}, 0.0f), 255.0f);") + * The clamp into [0, 255] happens in f32 so the conversion is always in range, + * and fmaxf turns a NaN into 0, as the LLVM lowering's maxnum does. */ +static rocke_status_t rocke_h_op_arith_cvt_f32_to_u8_sat(rocke_h_lowerer_t* lw, + const rocke_op_t* op) +{ + const rocke_value_t* v = op->operands[0]; + const rocke_value_t* r = h_res(op); + const char* rn = rocke_h_name(lw, r); + rocke_h_emitf(lw, "float %s_r = rintf(%s);", rn, rocke_h_name(lw, v)); + rocke_h_emitf(lw, "uint8_t %s = (uint8_t)fminf(fmaxf(%s_r, 0.0f), 255.0f);", rn, rn); + return lw->status; +} + /* ----------------------------- arith: packed fp8/bf8 cvt ------------------- */ /* def _op_arith_cvt_pk_f32_fp8x4(self, op): @@ -480,11 +513,13 @@ const rocke_h_handler_entry_t* rocke_h_handlers_cast(void) {ROCKE_OP_ARITH_CVT_F32_TO_F64, rocke_h_op_arith_cvt_f32_to_f64}, {ROCKE_OP_ARITH_CVT_F64_TO_F32, rocke_h_op_arith_cvt_f64_to_f32}, {ROCKE_OP_ARITH_SITOFP_F32, rocke_h_op_arith_sitofp_f32}, + {ROCKE_OP_ARITH_UITOFP_F32, rocke_h_op_arith_uitofp_f32}, {ROCKE_OP_ARITH_CVT_FP8_TO_F32, rocke_h_op_arith_cvt_fp8_to_f32}, {ROCKE_OP_ARITH_CVT_BF8_TO_F32, rocke_h_op_arith_cvt_bf8_to_f32}, {ROCKE_OP_ARITH_CVT_F32_TO_FP8, rocke_h_op_arith_cvt_f32_to_fp8}, {ROCKE_OP_ARITH_CVT_F32_TO_BF8, rocke_h_op_arith_cvt_f32_to_bf8}, {ROCKE_OP_ARITH_CVT_F32_TO_I8_SAT, rocke_h_op_arith_cvt_f32_to_i8_sat}, + {ROCKE_OP_ARITH_CVT_F32_TO_U8_SAT, rocke_h_op_arith_cvt_f32_to_u8_sat}, {ROCKE_OP_ARITH_CVT_PK_F32_FP8X4, rocke_h_op_arith_cvt_pk_f32_fp8x4}, {ROCKE_OP_ARITH_CVT_PK_F32_BF8X4, rocke_h_op_arith_cvt_pk_f32_bf8x4}, {ROCKE_OP_ARITH_CVT_PK_FP8_F32X4, rocke_h_op_arith_cvt_pk_fp8_f32x4}, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/arith.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/arith.cpp index df6b60af0aa9..65e2ab4ad972 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/arith.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/arith.cpp @@ -9,7 +9,7 @@ * arith.add/sub/mul/div/mod (integer), * arith.fadd/fsub/fmul/fdiv/fneg/fabs/fma/fmax3/fmin3 (float), * arith.cmp/fcmp, arith.fmax/fmin, - * arith.select/and/or/not/smax/smin/xor/shl/lshr/umul_hi_i32, + * arith.select/and/or/not/smax/smin/umax/umin/xor/shl/lshr/umul_hi_i32, * math.exp2/log2/rcp/rcp_fast/sqrt/rsqrt/tanh, * gpu.thread_id / gpu.block_id. * @@ -17,6 +17,7 @@ * rocke_ll_need, rocke_ll_fresh, rocke_ll_fp32_hex, rocke_ll_binop, rocke_ll_fail, ...) * lives in BUCKET 0; this file only calls them through the internal header. */ +#include "rocke/ir_internal.h" #include "rocke/lower_llvm_internal.h" #include @@ -167,25 +168,38 @@ static void _op_arith_constant_vec(rocke_lower_t* L, const rocke_op_t* op) /* The same-type binary handlers all defer to the shared rocke_ll_binop helper * (Python self._binop(op, llvm_op)). */ + +/* Python _wrapping_binop: `nsw` makes a signed overflow poison, as C does for + * signed types; unsigned values wrap modulo 2**N, so they get no flag. */ +static void ll_wrapping_binop(rocke_lower_t* L, + const rocke_op_t* op, + const char* signed_op, + const char* unsigned_op) +{ + const rocke_value_t* res = ll_result(op); + rocke_ll_binop(L, op, (res && rocke_type_is_unsigned(res->type)) ? unsigned_op : signed_op); +} static void _op_arith_add(rocke_lower_t* L, const rocke_op_t* op) { - rocke_ll_binop(L, op, "add nsw"); + ll_wrapping_binop(L, op, "add nsw", "add"); } static void _op_arith_sub(rocke_lower_t* L, const rocke_op_t* op) { - rocke_ll_binop(L, op, "sub nsw"); + ll_wrapping_binop(L, op, "sub nsw", "sub"); } static void _op_arith_mul(rocke_lower_t* L, const rocke_op_t* op) { - rocke_ll_binop(L, op, "mul nsw"); + ll_wrapping_binop(L, op, "mul nsw", "mul"); } static void _op_arith_div(rocke_lower_t* L, const rocke_op_t* op) { - rocke_ll_binop(L, op, "sdiv"); + const rocke_value_t* res = ll_result(op); + rocke_ll_binop(L, op, (res && rocke_type_is_unsigned(res->type)) ? "udiv" : "sdiv"); } static void _op_arith_mod(rocke_lower_t* L, const rocke_op_t* op) { - rocke_ll_binop(L, op, "srem"); + const rocke_value_t* res = ll_result(op); + rocke_ll_binop(L, op, (res && rocke_type_is_unsigned(res->type)) ? "urem" : "srem"); } static void _op_arith_fadd(rocke_lower_t* L, const rocke_op_t* op) { @@ -354,37 +368,9 @@ static void _op_arith_cmp(rocke_lower_t* L, const rocke_op_t* op) { pred = "lt"; } - const char* llvm_pred; - if(strcmp(pred, "lt") == 0) - { - llvm_pred = "slt"; - } - else if(strcmp(pred, "le") == 0) - { - llvm_pred = "sle"; - } - else if(strcmp(pred, "gt") == 0) - { - llvm_pred = "sgt"; - } - else if(strcmp(pred, "ge") == 0) - { - llvm_pred = "sge"; - } - else if(strcmp(pred, "eq") == 0) - { - llvm_pred = "eq"; - } - else if(strcmp(pred, "ne") == 0) - { - llvm_pred = "ne"; - } - else - { - rocke_ll_fail(L, ROCKE_ERR_KEY, "arith.cmp: unknown pred %s", pred); - } const rocke_value_t* a = op->operands[0]; const rocke_value_t* b = op->operands[1]; + const char* llvm_pred = rocke_ll_icmp_predicate(L, "arith.cmp", pred, a->type); rocke_ll_emitf(L, " %s = icmp %s %s %s, %s", res->name, @@ -546,6 +532,7 @@ static void _op_arith_smax(rocke_lower_t* L, const rocke_op_t* op) { return; } + rocke_i_refuse_unsigned_operands(op, "arith.umax"); const rocke_value_t* a = op->operands[0]; const rocke_value_t* b = op->operands[1]; rocke_ll_need(L, "smax.i32"); @@ -564,6 +551,7 @@ static void _op_arith_smin(rocke_lower_t* L, const rocke_op_t* op) { return; } + rocke_i_refuse_unsigned_operands(op, "arith.umin"); const rocke_value_t* a = op->operands[0]; const rocke_value_t* b = op->operands[1]; rocke_ll_need(L, "smin.i32"); @@ -574,6 +562,40 @@ static void _op_arith_smin(rocke_lower_t* L, const rocke_op_t* op) rocke_ll_operand(L, b)); } +/* Python _unsigned_min_max -> llvm.umax. / llvm.umin.. */ +static void ll_unsigned_min_max(rocke_lower_t* L, const rocke_op_t* op, const char* which) +{ + const rocke_value_t* res = ll_result(op); + if(!rocke_ll_live(L) || !res) + { + return; + } + const rocke_value_t* a = op->operands[0]; + const rocke_value_t* b = op->operands[1]; + const char* ty = rocke_ll_llvm_type(L, a->type); + char key[32]; + snprintf(key, sizeof(key), "%s.%s", which, ty); + rocke_ll_need(L, key); + rocke_ll_emitf(L, + " %s = call %s @llvm.%s.%s(%s %s, %s %s)", + res->name, + ty, + which, + ty, + ty, + rocke_ll_operand(L, a), + ty, + rocke_ll_operand(L, b)); +} +static void _op_arith_umax(rocke_lower_t* L, const rocke_op_t* op) +{ + ll_unsigned_min_max(L, op, "umax"); +} +static void _op_arith_umin(rocke_lower_t* L, const rocke_op_t* op) +{ + ll_unsigned_min_max(L, op, "umin"); +} + /* Python _op_arith_xor (result-typed). */ static void _op_arith_xor(rocke_lower_t* L, const rocke_op_t* op) { @@ -891,6 +913,8 @@ void rocke_ll_register_arith(void) rocke_ll_set_handler(ROCKE_OP_ARITH_NOT, _op_arith_not); rocke_ll_set_handler(ROCKE_OP_ARITH_SMAX, _op_arith_smax); rocke_ll_set_handler(ROCKE_OP_ARITH_SMIN, _op_arith_smin); + rocke_ll_set_handler(ROCKE_OP_ARITH_UMAX, _op_arith_umax); + rocke_ll_set_handler(ROCKE_OP_ARITH_UMIN, _op_arith_umin); rocke_ll_set_handler(ROCKE_OP_ARITH_XOR, _op_arith_xor); rocke_ll_set_handler(ROCKE_OP_ARITH_SHL, _op_arith_shl); rocke_ll_set_handler(ROCKE_OP_ARITH_LSHR, _op_arith_lshr); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/control.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/control.cpp index cf22352e800f..9273d2e0b1af 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/control.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/control.cpp @@ -22,6 +22,7 @@ * rocke_ll_need, rocke_ll_fresh, rocke_ll_vector_binop, rocke_ll_fail, ...) lives in * BUCKET 0; this file only calls them through the internal header. */ +#include "rocke/ir_internal.h" #include "rocke/lower_llvm_internal.h" #include @@ -340,37 +341,9 @@ static void _op_vector_cmp(rocke_lower_t* L, const rocke_op_t* op) { pred = "lt"; } - const char* llvm_pred; - if(strcmp(pred, "lt") == 0) - { - llvm_pred = "slt"; - } - else if(strcmp(pred, "le") == 0) - { - llvm_pred = "sle"; - } - else if(strcmp(pred, "gt") == 0) - { - llvm_pred = "sgt"; - } - else if(strcmp(pred, "ge") == 0) - { - llvm_pred = "sge"; - } - else if(strcmp(pred, "eq") == 0) - { - llvm_pred = "eq"; - } - else if(strcmp(pred, "ne") == 0) - { - llvm_pred = "ne"; - } - else - { - rocke_ll_fail(L, ROCKE_ERR_KEY, "vector.cmp: unknown pred %s", pred); - } const rocke_value_t* a = op->operands[0]; const rocke_value_t* b = op->operands[1]; + const char* llvm_pred = rocke_ll_icmp_predicate(L, "vector.cmp", pred, a->type); rocke_ll_emitf(L, " %s = icmp %s %s %s, %s", res->name, @@ -388,6 +361,7 @@ static void _op_vector_smax(rocke_lower_t* L, const rocke_op_t* op) { return; } + rocke_i_refuse_unsigned_operands(op, NULL); const rocke_value_t* a = op->operands[0]; const rocke_value_t* b = op->operands[1]; const rocke_type_t* vec_ty = a->type; @@ -423,6 +397,7 @@ static void _op_vector_smin(rocke_lower_t* L, const rocke_op_t* op) { return; } + rocke_i_refuse_unsigned_operands(op, NULL); const rocke_value_t* a = op->operands[0]; const rocke_value_t* b = op->operands[1]; int count = (res->type && res->type->kind == ROCKE_TYPE_VECTOR) ? res->type->count : 0; @@ -469,6 +444,7 @@ static void _op_vector_sext(rocke_lower_t* L, const rocke_op_t* op) { return; } + rocke_i_refuse_unsigned_operands(op, NULL); const rocke_value_t* v = op->operands[0]; rocke_ll_emitf(L, " %s = sext %s %s to %s", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/convert.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/convert.cpp index a85cd62fe14b..86d104b77f9c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/convert.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/convert.cpp @@ -15,6 +15,7 @@ * reached through rocke/lower_llvm_internal.h. */ +#include "rocke/ir_internal.h" #include "rocke/lower_llvm_internal.h" #include @@ -64,6 +65,7 @@ static void _op_arith_zext(rocke_lower_t* L, const rocke_op_t* op) static void _op_arith_sext(rocke_lower_t* L, const rocke_op_t* op) { + rocke_i_refuse_unsigned_operands(op, "arith.zext"); const rocke_value_t* v = op->operands[0]; rocke_ll_emitf(L, " %s = sext %s %s to %s", @@ -163,6 +165,16 @@ static void _op_arith_sitofp_f32(rocke_lower_t* L, const rocke_op_t* op) L, " %s = sitofp i32 %s to float", rocke_ll_result_name(op), rocke_ll_operand(L, v)); } +static void _op_arith_uitofp_f32(rocke_lower_t* L, const rocke_op_t* op) +{ + const rocke_value_t* v = op->operands[0]; + rocke_ll_emitf(L, + " %s = uitofp %s %s to float", + rocke_ll_result_name(op), + rocke_ll_llvm_type(L, v->type), + rocke_ll_operand(L, v)); +} + static void _op_arith_rint_f32(rocke_lower_t* L, const rocke_op_t* op) { const rocke_value_t* v = op->operands[0]; @@ -442,6 +454,27 @@ static void _op_arith_cvt_f32_to_i8_sat(rocke_lower_t* L, const rocke_op_t* op) rocke_ll_emitf(L, " %s = trunc i32 %s to i8", res, smin_v); } +/* Python _op_arith_cvt_f32_to_u8_sat: clamping in f32 keeps fptoui in range, + * and maxnum turns a NaN into 0. */ +static void _op_arith_cvt_f32_to_u8_sat(rocke_lower_t* L, const rocke_op_t* op) +{ + const rocke_value_t* v = op->operands[0]; + const char* res = rocke_ll_result_name(op); + const char* rounded = rocke_ll_derived(L, op, "r"); + const char* low = rocke_ll_derived(L, op, "lo"); + const char* high = rocke_ll_derived(L, op, "hi"); + rocke_ll_need(L, "rint.f32"); + rocke_ll_need(L, "maxnum.f32"); + rocke_ll_need(L, "minnum.f32"); + rocke_ll_emitf( + L, " %s = call float @llvm.rint.f32(float %s)", rounded, rocke_ll_operand(L, v)); + rocke_ll_emitf( + L, " %s = call float @llvm.maxnum.f32(float %s, float 0.000000e+00)", low, rounded); + rocke_ll_emitf( + L, " %s = call float @llvm.minnum.f32(float %s, float 2.550000e+02)", high, low); + rocke_ll_emitf(L, " %s = fptoui float %s to i8", res, high); +} + static void _op_arith_cvt_pk_i8_f32x4(rocke_lower_t* L, const rocke_op_t* op) { const rocke_value_t* v = op->operands[0]; @@ -524,6 +557,7 @@ void rocke_ll_register_convert(void) rocke_ll_set_handler(ROCKE_OP_ARITH_CVT_F32_TO_F64, _op_arith_cvt_f32_to_f64); rocke_ll_set_handler(ROCKE_OP_ARITH_CVT_F64_TO_F32, _op_arith_cvt_f64_to_f32); rocke_ll_set_handler(ROCKE_OP_ARITH_SITOFP_F32, _op_arith_sitofp_f32); + rocke_ll_set_handler(ROCKE_OP_ARITH_UITOFP_F32, _op_arith_uitofp_f32); rocke_ll_set_handler(ROCKE_OP_ARITH_RINT_F32, _op_arith_rint_f32); rocke_ll_set_handler(ROCKE_OP_ARITH_CVT_FP8_TO_F32, _op_arith_cvt_fp8_to_f32); @@ -540,6 +574,7 @@ void rocke_ll_register_convert(void) rocke_ll_set_handler(ROCKE_OP_ARITH_CVT_PK_BF8_F32X4, _op_arith_cvt_pk_bf8_f32x4); rocke_ll_set_handler(ROCKE_OP_ARITH_CVT_PK_I8_F32X4, _op_arith_cvt_pk_i8_f32x4); rocke_ll_set_handler(ROCKE_OP_ARITH_CVT_F32_TO_I8_SAT, _op_arith_cvt_f32_to_i8_sat); + rocke_ll_set_handler(ROCKE_OP_ARITH_CVT_F32_TO_U8_SAT, _op_arith_cvt_f32_to_u8_sat); } } /* namespace ckc */ 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 f2403bd0ea3a..d85815ee3c57 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 @@ -621,7 +621,8 @@ const char* rocke_ll_llvm_type(rocke_lower_t* L, const rocke_type_t* t) { if(strcmp(n, "i1") == 0) return "i1"; - if(strcmp(n, "i8") == 0) + /* LLVM integers carry no sign; signedness lives in the operations. */ + if(strcmp(n, "i8") == 0 || strcmp(n, "u8") == 0) return "i8"; if(strcmp(n, "i16") == 0) return "i16"; @@ -994,6 +995,34 @@ void rocke_ll_vector_binop(rocke_lower_t* L, const rocke_op_t* op, const char* l rocke_ll_operand(L, b)); } +const char* rocke_ll_icmp_predicate(rocke_lower_t* L, + const char* op_name, + const char* pred, + const rocke_type_t* t) +{ + static const struct + { + const char* pred; + const char* signed_form; + const char* unsigned_form; + } k_icmp[] = { + {"lt", "slt", "ult"}, + {"le", "sle", "ule"}, + {"gt", "sgt", "ugt"}, + {"ge", "sge", "uge"}, + {"eq", "eq", "eq"}, + {"ne", "ne", "ne"}, + }; + for(const auto& row : k_icmp) + { + if(strcmp(pred, row.pred) == 0) + { + return rocke_type_is_unsigned(t) ? row.unsigned_form : row.signed_form; + } + } + rocke_ll_fail(L, ROCKE_ERR_KEY, "%s: unknown pred %s", op_name, pred); +} + /* ====================================================================== */ /* smem pre-pass (Python _collect_smem / _smem_global_name) */ /* ====================================================================== */ @@ -1234,7 +1263,8 @@ static int ll_smem_seg_size(const rocke_type_t* stype) return 0; const char* n = stype->elem->name; int eb; - if(strcmp(n, "i8") == 0 || strcmp(n, "fp8e4m3") == 0 || strcmp(n, "bf8e5m2") == 0) + if(strcmp(n, "i8") == 0 || strcmp(n, "u8") == 0 || strcmp(n, "fp8e4m3") == 0 + || strcmp(n, "bf8e5m2") == 0) eb = 1; else if(strcmp(n, "f16") == 0 || strcmp(n, "bf16") == 0) eb = 2; @@ -1259,7 +1289,8 @@ static int ll_smem_align(const rocke_type_t* stype) const char* n = stype->elem->name; if(strcmp(n, "f64") == 0) return 8; - if(strcmp(n, "i8") == 0 || strcmp(n, "fp8e4m3") == 0 || strcmp(n, "bf8e5m2") == 0) + if(strcmp(n, "i8") == 0 || strcmp(n, "u8") == 0 || strcmp(n, "fp8e4m3") == 0 + || strcmp(n, "bf8e5m2") == 0) return 16; return 4; } 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 694ea26888b7..93466dfe6b50 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 @@ -342,6 +342,9 @@ const rocke_ll_decl_t ROCKE_LL_INTRINSIC_DECLS[] = { {"rint.f32", "declare float @llvm.rint.f32(float)"}, {"smax.i32", "declare i32 @llvm.smax.i32(i32, i32)"}, {"smin.i32", "declare i32 @llvm.smin.i32(i32, i32)"}, + /* Unsigned u8 maximum / minimum (u8 renders as i8). */ + {"umax.i8", "declare i8 @llvm.umax.i8(i8, i8)"}, + {"umin.i8", "declare i8 @llvm.umin.i8(i8, i8)"}, {"amdgcn.perm", "declare i32 @llvm.amdgcn.perm(i32, i32, i32)"}, {"amdgcn.cvt.scalef32.pk.f32.fp8", "declare <2 x float> @llvm.amdgcn.cvt.scalef32.pk.f32.fp8(i32, float, i1)"}, 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 8c37f74c03b3..a82d8d17af71 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 @@ -37,7 +37,8 @@ static int ll_elem_bytes(const char* name) { return 2; } - if(strcmp(name, "i8") == 0 || strcmp(name, "fp8e4m3") == 0 || strcmp(name, "bf8e5m2") == 0) + if(strcmp(name, "i8") == 0 || strcmp(name, "u8") == 0 || strcmp(name, "fp8e4m3") == 0 + || strcmp(name, "bf8e5m2") == 0) { return 1; } @@ -697,12 +698,12 @@ static void op_tile_smem_load_vN(rocke_lower_t* L, const rocke_op_t* op) * 2 (not 1). Replicate that exact dict here rather than the shared * ll_elem_bytes (which maps fp8/bf8 -> 1), or the fp8 down-GEMM LDS reads * emit `align 16` instead of the Python `align 32`. - * {"i8":1,"f16":2,"bf16":2,"i32":4,"f32":4,"i64":8,"f64":8}.get(name, 2) */ + * {"i8":1,"u8":1,"f16":2,"bf16":2,"i32":4,"f32":4,"i64":8,"f64":8}.get(name, 2) */ const char* en = op->results[0]->type->elem->name; int elem_bytes = 2; if(en) { - if(strcmp(en, "i8") == 0) + if(strcmp(en, "i8") == 0 || strcmp(en, "u8") == 0) { elem_bytes = 1; } 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 fd3202dc01a3..dd72f3a15f45 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 @@ -113,6 +113,8 @@ typedef enum rocke_scalar_kind ROCKE_SCALAR_TF32, ROCKE_SCALAR_F64, ROCKE_SCALAR_I4, + ROCKE_SCALAR_U4, + ROCKE_SCALAR_U8, ROCKE_SCALAR__COUNT } rocke_scalar_kind_t; @@ -457,6 +459,12 @@ typedef enum rocke_opcode ROCKE_OP_ARITH_CVT_F32_TO_F64, ROCKE_OP_ARITH_CVT_F64_TO_F32, + /* arith.* unsigned (u8) ops. Appended so no existing opcode index shifts. */ + ROCKE_OP_ARITH_UMAX, + ROCKE_OP_ARITH_UMIN, + ROCKE_OP_ARITH_UITOFP_F32, + ROCKE_OP_ARITH_CVT_F32_TO_U8_SAT, + ROCKE_OP__COUNT } rocke_opcode_t; @@ -641,6 +649,8 @@ typedef struct rocke_mem_opts * never NULL; never arena-owned (static storage). */ const rocke_type_t* rocke_i1(void); const rocke_type_t* rocke_i4(void); +const rocke_type_t* rocke_u4(void); +const rocke_type_t* rocke_u8(void); const rocke_type_t* rocke_i8(void); const rocke_type_t* rocke_i16(void); const rocke_type_t* rocke_i32(void); @@ -678,6 +688,11 @@ const rocke_type_t* rocke_smem_type( * canonical name, which encodes kind + components). */ bool rocke_type_eq(const rocke_type_t* a, const rocke_type_t* b); +/* True when `t`, or the element of a vector `t`, is an unsigned integer (u8, + * u4). The twin of core/ir.py::_is_unsigned: every op whose result depends on + * signedness asks it which form to emit. NULL-safe. */ +bool rocke_type_is_unsigned(const rocke_type_t* t); + /* AMDGPU buffer-load AUX cache-coherency hints (Python module constants). */ typedef enum rocke_cache_policy { @@ -784,8 +799,12 @@ rocke_value_t* rocke_b_mod(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_ rocke_value_t* rocke_b_land(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_t* c); rocke_value_t* rocke_b_lor(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_t* c); rocke_value_t* rocke_b_lnot(rocke_ir_builder_t* b, rocke_value_t* a); +/* Signed max / min; an unsigned operand is refused (TypeError). */ rocke_value_t* rocke_b_smax(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_t* c); rocke_value_t* rocke_b_smin(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_t* c); +/* Unsigned max / min of two u8 values. */ +rocke_value_t* rocke_b_umax(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_t* c); +rocke_value_t* rocke_b_umin(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_t* c); rocke_value_t* rocke_b_xor(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_t* c); rocke_value_t* rocke_b_shl(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_t* c); rocke_value_t* rocke_b_lshr(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_t* c); @@ -852,6 +871,8 @@ rocke_value_t* rocke_b_cast_to_f32(rocke_ir_builder_t* b, rocke_value_t* v); rocke_value_t* rocke_b_cast_f32_to(rocke_ir_builder_t* b, rocke_value_t* v, const rocke_type_t* target); rocke_value_t* rocke_b_sitofp_f32(rocke_ir_builder_t* b, rocke_value_t* v); +/* u8 -> f32 (LLVM uitofp). Exact. */ +rocke_value_t* rocke_b_uitofp_f32(rocke_ir_builder_t* b, rocke_value_t* v); /* f32 <-> f64 have their own opcodes; cast_to_f32 / cast_f32_to reject f64. */ rocke_value_t* rocke_b_cvt_f32_to_f64(rocke_ir_builder_t* b, rocke_value_t* v); rocke_value_t* rocke_b_cvt_f64_to_f32(rocke_ir_builder_t* b, rocke_value_t* v); @@ -870,6 +891,8 @@ 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); +/* f32 -> u8, round to nearest even and saturate to [0, 255]; NaN gives 0. */ +rocke_value_t* rocke_b_cvt_f32_to_u8_sat(rocke_ir_builder_t* b, rocke_value_t* v); rocke_value_t* rocke_b_cvt_pk_fp8_f32x4(rocke_ir_builder_t* b, rocke_value_t* v); rocke_value_t* rocke_b_cvt_pk_bf8_f32x4(rocke_ir_builder_t* b, rocke_value_t* v); rocke_value_t* rocke_b_cvt_pk_i8_f32x4(rocke_ir_builder_t* b, rocke_value_t* v); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir_internal.h index de5e21c27f40..eed8c7d3f96a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir_internal.h @@ -136,13 +136,26 @@ bool rocke_i_is_vector(const rocke_type_t* t, const char* elem_name, int count); * `v.type.elem if isinstance(...,VectorType) else v.type`). */ const rocke_type_t* rocke_i_elem_of(const rocke_type_t* t); -/* Throw ckc::TypeError when `t` is a packed-only scalar (i4), a vector of one +/* Throw ckc::TypeError when `t` is a packed-only scalar (i4, u4), a vector of one * or an LDS tile of one: the twin of core/ir.py::_refuse_packed_only_value, * which states the rule and the message. rocke_i_op calls it for every result * and rocke_b_param for every parameter, so no builder makes such a value. * NULL-safe. */ void rocke_i_refuse_packed_only_value(const char* op_name, const rocke_type_t* t); +/* Throw ckc::TypeError when `t` is unsigned (rocke_type_is_unsigned) and + * `op_name` treats its operands as signed; `instead` names the unsigned op to + * use, or is NULL when there is none. The twin of core/ir.py::_refuse_unsigned, + * with the same message. NULL-safe. */ +void rocke_i_refuse_unsigned(const char* op_name, const rocke_type_t* t, const char* instead); + +/* rocke_i_refuse_unsigned for every operand of `op`, a signed op, under the + * op's dotted name. The builders refuse such an op when it is built; the + * lowerings call this for IR that did not come through them (hand-assembled + * or deserialized), so no signed op lowers on unsigned data. The twin of + * core/ir.py::_refuse_unsigned_operands. */ +void rocke_i_refuse_unsigned_operands(const rocke_op_t* op, const char* instead); + /* Lane count of a vector type, or 1 for a scalar. */ int rocke_i_count_of(const rocke_type_t* t); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/lower_llvm_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/lower_llvm_internal.h index b948c7c20c1f..a8582d87a131 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/lower_llvm_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/lower_llvm_internal.h @@ -660,6 +660,15 @@ void rocke_ll_binop(rocke_lower_t* L, const rocke_op_t* op, const char* llvm_op) /* Vector same-type binary op (Python _vector_binop). */ void rocke_ll_vector_binop(rocke_lower_t* L, const rocke_op_t* op, const char* llvm_op); +/* The icmp predicate for an arith.cmp / vector.cmp predicate ("lt", ...) on + * operands of type `t`: ordered compares of unsigned integers are unsigned + * (Python _icmp_predicate). An unknown predicate fails ": unknown pred + * ". */ +const char* rocke_ll_icmp_predicate(rocke_lower_t* L, + const char* op_name, + const char* pred, + const rocke_type_t* t); + /* Shared FP8/BF8 MFMA lowering body (Python _lower_mfma_fp8_bf8): bitcasts the * <8 x i8> A/B to the flavor-correct packed type (i64 on LLVM22, <2 x i32> on * LLVM20) and emits the call. */ 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 a116b6c00d3b..2414256fb598 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 @@ -137,8 +137,29 @@ storage and unpack explicitly`); a kernel keeps the words in `i32` storage and unpacks them explicitly. Both lowerings still refuse an `i4` value that reaches them (`no LLVM mapping for type i4`, `no HIP mapping for type i4`). +**Unsigned integers.** `u8` and `u4` name unsigned data, so integer matrix +instructions with a per-operand signedness flag can be selected for it. `u8` +renders as LLVM `i8` (LLVM integers carry no sign) and HIP `uint8_t`, and takes +one byte in LDS, like `i8`. `u8` is an arithmetic type: every op whose result +depends on signedness emits its unsigned form for it -- `udiv` / `urem` for +`div` / `mod`, `icmp ult` / `ule` / `ugt` / `uge` for the ordered compares, and +no `nsw` flag on `add` / `sub` / `mul`, since unsigned values wrap. `zext` and +`lshr` widen and shift it; `umax` / `umin` (`llvm.umax` / `llvm.umin`), +`uitofp_f32` and the saturating `cvt_f32_to_u8_sat` (round to nearest even, +then clamp to [0, 255] in f32 before `fptoui`; a NaN gives 0) are its own ops. +The ops that are signed by definition -- `smax`, `smin`, `sext` and their +vector forms -- refuse a `u8` operand (`smax treats its operands as signed, but +u8 is unsigned -- use umax`). Native `u16`/`u32` types remain absent. +A single `u8` can be loaded from and stored to global memory and stored to +LDS. The vector memory builders refuse `u8`: the global pair, and the LDS pair +`smem_load_vN` / `smem_store_vN`, which admit one shared element list (so an +LDS load of `u8`, which `smem_load_vN` also serves for one element, is refused +too). `u4`, like `i4`, is packed-only: the IR builder of both engines refuses +its values with the same text, and both lowerings refuse a `u4` value that +reaches them. + **Types that are entirely absent from the vocabulary:** unsigned -`u8/u16/u32/u64`, and a distinct `bool` (predicates reuse `i1`). No engine +`u16/u32/u64`, and a distinct `bool` (predicates reuse `i1`). No engine references them. TF32 source anchors: [Python type and builders](../../python/rocke/core/ir.py), @@ -245,7 +266,7 @@ Gap categories: **[C]** conversion gap · **[P]** partial/plumbing gap · | `i8` | 3 | 2 | 2 | 2 | 1 | 2 | **67%** | [P] | | `i64` | 3 | — | — | 1 | 1 | — | **56%** | [P] | | `i16` | 3 | — | — | 2 | 1 | — | **67%** | [P] | -| `u8` | 0 | 0 | 0 | 0 | 0 | — | **0%** | [M] | +| `u8` | 1 | 0 | 0 | 0 | 0 | -- | **7%** | [M] | | `u16/u32/u64` | 0 | 0 | 0 | 0 | 0 | — | **0%** | [M] | | **Compute mode (not a storage type)** | | | | | | | | | | `bf16-mode` (reduced-precision compute over `f32`) | 0 | 3 | 2 | 3 | 0 | 1 | **50%** | [P][M] | @@ -349,9 +370,11 @@ mode. The missing selector and precision contract are described in §2.5. (guarded), so vector-width predicate init must be open-coded. - **`i64` C4=1** — 8-byte LDS works and addressing is fine, but there is no `global_store_vN` and no zero-vec. It is an addressing type with no matrix atom. -- **`u8` — 0 at every layer.** `UINT8` appears in the hipDNN enum +- **`u8` C1=1, 0 elsewhere.** The row records the snapshot before `u8` gained its + unsigned arithmetic, compares and conversions, described in section 1; it still has no + vector memory helpers. `UINT8` appears in the hipDNN enum (`GpuReferenceValidationFactory.cpp:31`), but only in a reference-validation - switch; no rocKE surface accepts it. Note that the silicon is not the + switch. Note that the silicon is not the constraint: the RDNA/GFX12 WMMA integer atom is `IU8`, whose per-operand signedness selects make it an *unsigned* int8 matmul as well as a signed one (§4A). 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 e7a206eec32d..d8a4a2e1b005 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 @@ -66,6 +66,8 @@ FP8E4M3, I1, I4, + U4, + U8, I8, I16, I32, @@ -174,6 +176,8 @@ "FP8E4M3", "I1", "I4", + "U4", + "U8", "I8", "I16", "I32", 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 125d4acd55c9..e2fe75d654a0 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 @@ -62,6 +62,8 @@ dtype_to_ir_type, I1, I4, + U4, + U8, I8, I32, I64, @@ -127,6 +129,8 @@ "FP8E4M3", "I1", "I4", + "U4", + "U8", "I8", "I32", "I64", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm20.json b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm20.json index fadca37b7ea5..0cdedfc0f126 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm20.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm20.json @@ -121,6 +121,8 @@ "smax.i32": "llvm.smax.i32", "smin.i32": "llvm.smin.i32", "sqrt.f32": "llvm.sqrt.f32", + "umax.i8": "llvm.umax.i8", + "umin.i8": "llvm.umin.i8", "update.dpp.i32": "llvm.amdgcn.update.dpp.i32", "wmma.f32.16x16x16.bf16": "llvm.amdgcn.wmma.f32.16x16x16.bf16.v8f32.v16i16", "wmma.f32.16x16x16.f16": "llvm.amdgcn.wmma.f32.16x16x16.f16.v8f32.v16f16", @@ -6662,6 +6664,76 @@ "verified_on": "llvm20" } }, + "umax.i8": { + "gfx11-generic": { + "status": "ok", + "verified_on": "llvm20" + }, + "gfx1151": { + "status": "ok", + "verified_on": "llvm20" + }, + "gfx12-generic": { + "status": "ok", + "verified_on": "llvm20" + }, + "gfx1201": { + "status": "ok", + "verified_on": "llvm20" + }, + "gfx1250": { + "evidence": "invalid target ID", + "status": "target_unsupported", + "verified_on": "llvm20" + }, + "gfx90a": { + "status": "ok", + "verified_on": "llvm20" + }, + "gfx942": { + "status": "ok", + "verified_on": "llvm20" + }, + "gfx950": { + "status": "ok", + "verified_on": "llvm20" + } + }, + "umin.i8": { + "gfx11-generic": { + "status": "ok", + "verified_on": "llvm20" + }, + "gfx1151": { + "status": "ok", + "verified_on": "llvm20" + }, + "gfx12-generic": { + "status": "ok", + "verified_on": "llvm20" + }, + "gfx1201": { + "status": "ok", + "verified_on": "llvm20" + }, + "gfx1250": { + "evidence": "invalid target ID", + "status": "target_unsupported", + "verified_on": "llvm20" + }, + "gfx90a": { + "status": "ok", + "verified_on": "llvm20" + }, + "gfx942": { + "status": "ok", + "verified_on": "llvm20" + }, + "gfx950": { + "status": "ok", + "verified_on": "llvm20" + } + }, "update.dpp.i32": { "gfx11-generic": { "status": "ok", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm22.json b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm22.json index cf81bb653cb0..3e66e514758a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm22.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm22.json @@ -152,6 +152,8 @@ "smax.i32": "llvm.smax.i32", "smin.i32": "llvm.smin.i32", "sqrt.f32": "llvm.sqrt.f32", + "umax.i8": "llvm.umax.i8", + "umin.i8": "llvm.umin.i8", "update.dpp.i32": "llvm.amdgcn.update.dpp.i32", "wave.reduce.add.i32": "llvm.amdgcn.wave.reduce.add.i32", "wave.reduce.max.i32": "llvm.amdgcn.wave.reduce.max.i32", @@ -6509,6 +6511,74 @@ "verified_on": "llvm22" } }, + "umax.i8": { + "gfx11-generic": { + "status": "ok", + "verified_on": "llvm22" + }, + "gfx1151": { + "status": "ok", + "verified_on": "llvm22" + }, + "gfx12-generic": { + "status": "ok", + "verified_on": "llvm22" + }, + "gfx1201": { + "status": "ok", + "verified_on": "llvm22" + }, + "gfx1250": { + "status": "ok", + "verified_on": "llvm22" + }, + "gfx90a": { + "status": "ok", + "verified_on": "llvm22" + }, + "gfx942": { + "status": "ok", + "verified_on": "llvm22" + }, + "gfx950": { + "status": "ok", + "verified_on": "llvm22" + } + }, + "umin.i8": { + "gfx11-generic": { + "status": "ok", + "verified_on": "llvm22" + }, + "gfx1151": { + "status": "ok", + "verified_on": "llvm22" + }, + "gfx12-generic": { + "status": "ok", + "verified_on": "llvm22" + }, + "gfx1201": { + "status": "ok", + "verified_on": "llvm22" + }, + "gfx1250": { + "status": "ok", + "verified_on": "llvm22" + }, + "gfx90a": { + "status": "ok", + "verified_on": "llvm22" + }, + "gfx942": { + "status": "ok", + "verified_on": "llvm22" + }, + "gfx950": { + "status": "ok", + "verified_on": "llvm22" + } + }, "update.dpp.i32": { "gfx11-generic": { "status": "ok", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm23.json b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm23.json index bf7030268285..393ca7d26964 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm23.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/data/intrinsic_arch_domain.llvm23.json @@ -167,6 +167,8 @@ "sqrt.f32": "llvm.sqrt.f32", "tensor.load.to.lds": "llvm.amdgcn.tensor.load.to.lds", "tensor.store.from.lds": "llvm.amdgcn.tensor.store.from.lds", + "umax.i8": "llvm.umax.i8", + "umin.i8": "llvm.umin.i8", "update.dpp.i32": "llvm.amdgcn.update.dpp.i32", "wait.asyncmark": "llvm.amdgcn.wait.asyncmark", "wave.reduce.add.i32": "llvm.amdgcn.wave.reduce.add.i32", @@ -6468,6 +6470,74 @@ "verified_on": "llvm23" } }, + "umax.i8": { + "gfx11-generic": { + "status": "ok", + "verified_on": "llvm23" + }, + "gfx1151": { + "status": "ok", + "verified_on": "llvm23" + }, + "gfx12-generic": { + "status": "ok", + "verified_on": "llvm23" + }, + "gfx1201": { + "status": "ok", + "verified_on": "llvm23" + }, + "gfx1250": { + "status": "ok", + "verified_on": "llvm23" + }, + "gfx90a": { + "status": "ok", + "verified_on": "llvm23" + }, + "gfx942": { + "status": "ok", + "verified_on": "llvm23" + }, + "gfx950": { + "status": "ok", + "verified_on": "llvm23" + } + }, + "umin.i8": { + "gfx11-generic": { + "status": "ok", + "verified_on": "llvm23" + }, + "gfx1151": { + "status": "ok", + "verified_on": "llvm23" + }, + "gfx12-generic": { + "status": "ok", + "verified_on": "llvm23" + }, + "gfx1201": { + "status": "ok", + "verified_on": "llvm23" + }, + "gfx1250": { + "status": "ok", + "verified_on": "llvm23" + }, + "gfx90a": { + "status": "ok", + "verified_on": "llvm23" + }, + "gfx942": { + "status": "ok", + "verified_on": "llvm23" + }, + "gfx950": { + "status": "ok", + "verified_on": "llvm23" + } + }, "update.dpp.i32": { "gfx11-generic": { "status": "ok", 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 840bd920cbcf..fdffab0bc251 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 @@ -49,6 +49,11 @@ "int4": "i4", "i32": "i32", "int32": "i32", + # Unsigned integers (logical types; no unsigned arithmetic yet). + "u8": "u8", + "uint8": "u8", + "u4": "u4", + "uint4": "u4", } @@ -98,6 +103,8 @@ class DTypeInfo: ( ("i4", 4), ("i8", 8), + ("u4", 4), + ("u8", 8), ("i16", 16), ("i32", 32), ("i64", 64), 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 c72e7829b676..b43077944f60 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 @@ -46,6 +46,8 @@ def __repr__(self) -> str: I1 = Type("i1") I4 = Type("i4") +U4 = Type("u4") +U8 = Type("u8") I8 = Type("i8") I16 = Type("i16") I32 = Type("i32") @@ -73,6 +75,8 @@ def dtype_to_ir_type(dtype: str) -> Type: types = { "i1": I1, "i4": I4, + "u4": U4, + "u8": U8, "i8": I8, "i16": I16, "i32": I32, @@ -248,6 +252,30 @@ def _streaming(temporal_hint: TemporalHint) -> bool: _FENCE_ORDERINGS = ("acquire", "release", "acq_rel", "seq_cst") _FLAG_STORE_ORDERINGS = ("monotonic", "release", "seq_cst") +# Element types an LDS vector access (``smem_load_vN`` and ``smem_store_vN``) +# admits, in the order its refusal lists them. One list for both directions +# keeps the pair symmetric: whatever one stores as a vector can be loaded back. +_SMEM_VECTOR_ELEMENTS = ( + "f16", + "bf16", + "i16", + "f32", + "i32", + "tf32", + "f64", + "fp8e4m3", + "bf8e5m2", + "i8", +) + + +def _require_smem_vector_element(access: str, elem_name: str) -> None: + """Refuse an LDS vector ``access`` of an element outside ``_SMEM_VECTOR_ELEMENTS``.""" + if elem_name not in _SMEM_VECTOR_ELEMENTS: + raise ValueError( + f"{access} supports {' / '.join(_SMEM_VECTOR_ELEMENTS)}, got {elem_name}" + ) + def _is_flag_ptr(t: "Type") -> bool: """True for the ``ptr`` a cross-workgroup flag op addresses. @@ -264,12 +292,12 @@ def _check_flag_ptr(op_name: str, ptr: "Value") -> None: raise ValueError(f"{op_name} expects a global i32 pointer, got {ptr.type.name}") -# I4 is packed-only: its values only ever live packed in 32-bit words, so no IR -# value, scalar or vector, and no LDS tile has that element type. IRBuilder -# enforces it for every op result and kernel parameter through +# I4 and U4 are packed-only: their values only ever live packed in 32-bit +# words, so no IR value, scalar or vector, and no LDS tile has either element +# type. IRBuilder enforces it for every op result and kernel parameter through # _refuse_packed_only_value, the one home of that rule; the lowerers' type # tables have no entry for a packed-only type. -_PACKED_ONLY_SCALARS = ("i4",) +_PACKED_ONLY_SCALARS = ("i4", "u4") def _refuse_packed_only_value(op_name: str, t: "Type") -> None: @@ -285,6 +313,45 @@ def _refuse_packed_only_value(op_name: str, t: "Type") -> None: ) +# U8 and U4 hold unsigned integers. LLVM integers carry no sign, so u8 renders +# as i8 and every op whose result depends on signedness asks _is_unsigned which +# form to emit: unsigned division, remainder and ordered compares, zero rather +# than sign extension, and no signed-overflow flag on add / sub / mul. Ops that +# are signed by definition (smax, smin, sext and their vector forms) refuse an +# unsigned operand through _refuse_unsigned when they are built, and again +# through _refuse_unsigned_operands when they are lowered. U4 is packed-only +# (above), so only U8 values ever reach an op. +_UNSIGNED_SCALARS = ("u4", "u8") + + +def _is_unsigned(t: "Type") -> bool: + """True when ``t``, or the element of a vector ``t``, is an unsigned integer.""" + elem = t.elem if isinstance(t, VectorType) else t + return elem.name in _UNSIGNED_SCALARS + + +def _refuse_unsigned(op_name: str, t: "Type", instead: Optional[str] = None) -> None: + """Refuse an unsigned operand of type ``t`` to ``op_name``, an op that treats + its operands as signed; ``instead`` names the unsigned op to use, if any.""" + if _is_unsigned(t): + elem = t.elem if isinstance(t, VectorType) else t + message = ( + f"{op_name} treats its operands as signed, but {elem.name} is unsigned" + ) + if instead is not None: + message += f" -- use {instead}" + raise TypeError(message) + + +def _refuse_unsigned_operands(op: "Op", instead: Optional[str] = None) -> None: + """_refuse_unsigned for every operand of ``op``, a signed op, under the op's + dotted name. IRBuilder refuses such an op when it is built; the lowerings + call this for IR that did not come through it (hand-assembled or + deserialized), so no signed op lowers on unsigned data.""" + for operand in op.operands: + _refuse_unsigned(op.name, operand.type, instead) + + def _sync_attr(op: "Op", key: str, allowed: Sequence[str]) -> str: """Scope / ordering attr of a fence or flag op, checked for the lowerers. @@ -1159,15 +1226,41 @@ def lnot(self, a: Value) -> Value: return self._op("arith.not", [a], [a.type], result_name_hint="not").result def smax(self, a: Value, b: Value) -> Value: + """Signed maximum; an unsigned operand is refused (use :meth:`umax`).""" + _refuse_unsigned("smax", a.type, "umax") + _refuse_unsigned("smax", b.type, "umax") return self._op("arith.smax", [a, b], [a.type], result_name_hint="smax").result def smin(self, a: Value, b: Value) -> Value: + """Signed minimum; an unsigned operand is refused (use :meth:`umin`).""" + _refuse_unsigned("smin", a.type, "umin") + _refuse_unsigned("smin", b.type, "umin") return self._op("arith.smin", [a, b], [a.type], result_name_hint="smin").result + def _require_unsigned_pair(self, op_name: str, a: Value, b: Value) -> None: + scalar = not isinstance(a.type, VectorType) + if not (scalar and _is_unsigned(a.type) and a.type == b.type): + raise ValueError( + f"{op_name} expects two operands of one unsigned integer type, " + f"got {a.type.name} and {b.type.name}" + ) + + def umax(self, a: Value, b: Value) -> Value: + """Unsigned maximum of two u8 values. Lowers to ``llvm.umax``.""" + self._require_unsigned_pair("umax", a, b) + return self._op("arith.umax", [a, b], [a.type], result_name_hint="umax").result + + def umin(self, a: Value, b: Value) -> Value: + """Unsigned minimum of two u8 values. Lowers to ``llvm.umin``.""" + self._require_unsigned_pair("umin", a, b) + return self._op("arith.umin", [a, b], [a.type], result_name_hint="umin").result + def zext(self, v: Value, target: Type) -> Value: return self._op("arith.zext", [v], [target], result_name_hint="zx").result def sext(self, v: Value, target: Type) -> Value: + """Sign-extend; an unsigned source is refused (use :meth:`zext`).""" + _refuse_unsigned("sext", v.type, "zext") return self._op("arith.sext", [v], [target], result_name_hint="sx").result def trunc(self, v: Value, target: Type) -> Value: @@ -1278,6 +1371,14 @@ def sitofp_f32(self, v: Value) -> Value: raise ValueError(f"sitofp_f32 expects i32 input, got {v.type.name}") return self._op("arith.sitofp_f32", [v], [F32], result_name_hint="sitof").result + def uitofp_f32(self, v: Value) -> Value: + """Unsigned-integer to f32 conversion (LLVM uitofp) of a u8 value. Exact.""" + if isinstance(v.type, VectorType) or not _is_unsigned(v.type): + raise ValueError( + f"uitofp_f32 expects an unsigned integer input, got {v.type.name}" + ) + return self._op("arith.uitofp_f32", [v], [F32], result_name_hint="uitof").result + # f32 <-> f64 has its own opcode per direction rather than riding on # cast_to_f32 / cast_f32_to: those two are the f16/bf16 <-> f32 hub and # keep rejecting f64, so no caller widens or narrows a double by accident. @@ -1561,6 +1662,18 @@ def cvt_f32_to_i8_sat(self, v: Value) -> Value: "arith.cvt_f32_to_i8_sat", [v], [I8], result_name_hint="qi8" ).result + def cvt_f32_to_u8_sat(self, v: Value) -> Value: + """Round + saturate one f32 element to u8 (round-to-nearest-even). + + ``(uint8_t) min(255, max(0, rintf(v)))``; a NaN gives 0. The unsigned + twin of :meth:`cvt_f32_to_i8_sat`. + """ + if v.type.name != "f32": + raise ValueError(f"cvt_f32_to_u8_sat expects f32 input, got {v.type.name}") + return self._op( + "arith.cvt_f32_to_u8_sat", [v], [U8], result_name_hint="qu8" + ).result + def clamp_f32(self, v: Value, lo: Value, hi: Value) -> Value: """``min(hi, max(lo, v))`` for f32 -- saturating clamp every quantisation path needs before rounding. The two ``fmin`` / @@ -2146,8 +2259,12 @@ def global_load_vN( ).result def vector_binary(self, op_name: str, a: Value, b: Value) -> Value: + """Lane-wise ``vector.`` of two matching vectors. The signed ops + (``smax``, ``smin``) refuse an unsigned element type.""" if not isinstance(a.type, VectorType) or a.type != b.type: raise ValueError("vector_binary expects matching vector operands") + if op_name in ("smax", "smin"): + _refuse_unsigned(f"vector_{op_name}", a.type) return self._op( f"vector.{op_name}", [a, b], [a.type], result_name_hint=f"v{op_name}" ).result @@ -2268,6 +2385,7 @@ def vector_sext(self, v: Value, target: Type) -> Value: (``v_pk_max_i16`` etc.) can consume it.""" if not isinstance(v.type, VectorType): raise ValueError("vector_sext expects vector input") + _refuse_unsigned("vector_sext", v.type) return self._op( "vector.sext", [v], @@ -2300,9 +2418,10 @@ def smem_store_vN( ) -> None: """Vectorised LDS store of N consecutive values. - Supports scalar stores (`n=1`) plus vector stores for the element types - accepted by :meth:`smem_load_vN`. For 8-bit element types, ``n=16`` maps - to a 16-byte LDS transaction, which is the native-int WMMA staging shape. + Supports scalar stores (`n=1`) of any element type, plus vector stores + of the element types :meth:`smem_load_vN` accepts (the shared + ``_SMEM_VECTOR_ELEMENTS``). For 8-bit element types, ``n=16`` maps to a + 16-byte LDS transaction, which is the native-int WMMA staging shape. """ if n == 1: # Single-element store; route through the scalar `tile.smem_store`. @@ -2315,6 +2434,7 @@ def smem_store_vN( if not isinstance(value.type, VectorType): raise ValueError("smem_store_vN expects vector value for n > 1") elem_name = value.type.elem.name + _require_smem_vector_element("smem_store_vN", elem_name) if elem_name in ("i8", "fp8e4m3", "bf8e5m2"): allowed_n = (2, 4, 8, 16) elem_bytes = 1 @@ -2366,22 +2486,7 @@ def smem_load_vN(self, smem: Value, *indices, dtype: Type, n: int = 0) -> Value: or three words), using element alignment. The target and alignment determine instruction selection. """ - if dtype.name not in ( - "f16", - "bf16", - "i16", - "f32", - "i32", - "tf32", - "f64", - "fp8e4m3", - "bf8e5m2", - "i8", - ): - raise ValueError( - "smem_load_vN supports f16 / bf16 / i16 / f32 / i32 / tf32 / f64 / " - f"fp8e4m3 / bf8e5m2 / i8, got {dtype.name}" - ) + _require_smem_vector_element("smem_load_vN", dtype.name) if dtype.name in ("fp8e4m3", "bf8e5m2", "i8"): allowed_n = (1, 2, 4, 8, 12, 16) elif dtype.name in ("f16", "bf16", "i16"): @@ -5452,6 +5557,10 @@ def scf_if_else(self, cond: Value): "arith.cvt_f32_to_f64", "arith.cvt_f64_to_f32", "arith.cvt_fp8_to_f32", + "arith.umax", + "arith.umin", + "arith.uitofp_f32", + "arith.cvt_f32_to_u8_sat", "math.exp2", "math.exp2_fast", "math.log2", 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 7eb1fa07c0c1..e516ae443810 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 @@ -31,6 +31,7 @@ Type, Value, VectorType, + _refuse_unsigned_operands, _sync_attr, dpp_ctrl_attr, dpp_xor_ctrl, @@ -43,6 +44,7 @@ _HIP_TYPE = { "i1": "bool", "i8": "int8_t", + "u8": "uint8_t", "i16": "int16_t", "i32": "int", "tf32": "int", @@ -76,6 +78,7 @@ def _hip_scalar_type(type_name: str) -> str: _HIP_UNSIGNED_INT_TYPE = { "i1": "bool", "i8": "uint8_t", + "u8": "uint8_t", "i16": "uint16_t", "i32": "uint32_t", "i64": "uint64_t", @@ -490,19 +493,31 @@ def _op_arith_or(self, op: Op) -> None: f"{_type_to_hip(op.result.type)} {_name(op.result)} = {_name(a)} | {_name(b)};" ) - def _op_arith_smax(self, op: Op) -> None: + def _select_by_compare(self, op: Op, compare: str) -> None: + """The maximum (``>``) or minimum (``<``) of the two operands, as a C + ternary.""" a, b = op.operands self._emit( f"{_type_to_hip(op.result.type)} {_name(op.result)} = " - f"({_name(a)} > {_name(b)} ? {_name(a)} : {_name(b)});" + f"({_name(a)} {compare} {_name(b)} ? {_name(a)} : {_name(b)});" ) + # smax / smin are signed, so they refuse a u8 operand; umax / umin take u8. + # uint8_t operands promote to int, so C's own comparison and the ternary + # are the unsigned maximum / minimum. + def _op_arith_smax(self, op: Op) -> None: + _refuse_unsigned_operands(op, "arith.umax") + self._select_by_compare(op, ">") + def _op_arith_smin(self, op: Op) -> None: - a, b = op.operands - self._emit( - f"{_type_to_hip(op.result.type)} {_name(op.result)} = " - f"({_name(a)} < {_name(b)} ? {_name(a)} : {_name(b)});" - ) + _refuse_unsigned_operands(op, "arith.umin") + self._select_by_compare(op, "<") + + def _op_arith_umax(self, op: Op) -> None: + self._select_by_compare(op, ">") + + def _op_arith_umin(self, op: Op) -> None: + self._select_by_compare(op, "<") def _op_arith_zext(self, op: Op) -> None: (v,) = op.operands @@ -518,6 +533,7 @@ def _op_arith_zext(self, op: Op) -> None: ) def _op_arith_sext(self, op: Op) -> None: + _refuse_unsigned_operands(op, "arith.zext") (v,) = op.operands self._emit( f"{_type_to_hip(op.result.type)} {_name(op.result)} = ({_type_to_hip(op.result.type)}){_name(v)};" @@ -1614,6 +1630,11 @@ def _op_arith_sitofp_f32(self, op: Op) -> None: (v,) = op.operands self._emit(f"float {_name(op.result)} = (float){_name(v)};") + def _op_arith_uitofp_f32(self, op: Op) -> None: + # u8 -> f32: a uint8_t converts exactly. + (v,) = op.operands + self._emit(f"float {_name(op.result)} = (float){_name(v)};") + def _op_arith_cvt_fp8_to_f32(self, op: Op) -> None: # AMDGPU's per-byte fp8e4m3 -> f32 builtin. (v,) = op.operands @@ -1789,6 +1810,18 @@ def _op_arith_cvt_f32_to_i8_sat(self, op: Op) -> None: f"({as_i32} > 127 ? 127 : {as_i32}));" ) + def _op_arith_cvt_f32_to_u8_sat(self, op: Op) -> None: + # Saturating f32 -> u8, round-to-nearest-even. The clamp into + # [0, 255] happens in f32 so the conversion is always in range, and + # fmaxf turns a NaN into 0, as the LLVM lowering's maxnum does. + (v,) = op.operands + rounded = f"{_name(op.result)}_r" + self._emit(f"float {rounded} = rintf({_name(v)});") + self._emit( + f"uint8_t {_name(op.result)} = " + f"(uint8_t)fminf(fmaxf({rounded}, 0.0f), 255.0f);" + ) + def _op_arith_rint_f32(self, op: Op) -> None: (v,) = op.operands self._emit(f"float {_name(op.result)} = rintf({_name(v)});") 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 a14916c04ec1..fe979e660a2f 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 @@ -60,6 +60,8 @@ Type, Value, VectorType, + _is_unsigned, + _refuse_unsigned_operands, _sync_attr, check_cluster_dims, dpp_ctrl_attr, @@ -841,6 +843,9 @@ def _resolve_llvm_flavor() -> str: "rint.f32": "declare float @llvm.rint.f32(float)", "smax.i32": "declare i32 @llvm.smax.i32(i32, i32)", "smin.i32": "declare i32 @llvm.smin.i32(i32, i32)", + # Unsigned u8 maximum / minimum (u8 renders as i8). + "umax.i8": "declare i8 @llvm.umax.i8(i8, i8)", + "umin.i8": "declare i8 @llvm.umin.i8(i8, i8)", # AMDGPU ``v_perm_b32`` byte-select. Used by P09's packed i8 quant # to fuse four clamped i32s into one i32 in 2 perm calls (lo half + # hi half) followed by one perm to interleave. @@ -1233,7 +1238,8 @@ def _llvm_type(t: Type) -> str: # keeps its values out of the IR. if t.name == "i1": return "i1" - if t.name == "i8": + if t.name in ("i8", "u8"): + # LLVM integers carry no sign; signedness lives in the operations. return "i8" if t.name == "i16": return "i16" @@ -2095,6 +2101,7 @@ def _compute_smem_layout(self) -> None: # refuses its LDS tiles. _elem_bytes = { "i8": 1, + "u8": 1, "fp8e4m3": 1, "bf8e5m2": 1, "f16": 2, @@ -2126,7 +2133,7 @@ def _align(stype: "SmemType") -> int: # _refuse_packed_only_value refuses its LDS tiles. if stype.elem.name == "f64": return 8 - return 16 if stype.elem.name in ("i8", "fp8e4m3", "bf8e5m2") else 4 + return 16 if stype.elem.name in ("i8", "u8", "fp8e4m3", "bf8e5m2") else 4 # ---- live intervals from the kernel body ---- live, used = self._collect_smem_liveness(self.kernel.body) @@ -2342,20 +2349,26 @@ def _binop(self, op: Op, llvm_op: str) -> None: f"{self._operand(a)}, {self._operand(b)}" ) + def _wrapping_binop(self, op: Op, llvm_op: str) -> None: + # ``nsw`` makes a signed overflow poison, as C does for signed types; + # unsigned values wrap modulo 2**N, so they get no flag. + flag = "" if _is_unsigned(op.result.type) else " nsw" + self._binop(op, llvm_op + flag) + def _op_arith_add(self, op: Op) -> None: - self._binop(op, "add nsw") + self._wrapping_binop(op, "add") def _op_arith_sub(self, op: Op) -> None: - self._binop(op, "sub nsw") + self._wrapping_binop(op, "sub") def _op_arith_mul(self, op: Op) -> None: - self._binop(op, "mul nsw") + self._wrapping_binop(op, "mul") def _op_arith_div(self, op: Op) -> None: - self._binop(op, "sdiv") + self._binop(op, "udiv" if _is_unsigned(op.result.type) else "sdiv") def _op_arith_mod(self, op: Op) -> None: - self._binop(op, "srem") + self._binop(op, "urem" if _is_unsigned(op.result.type) else "srem") def _op_arith_fadd(self, op: Op) -> None: self._binop(op, "fadd") @@ -2465,18 +2478,10 @@ def _op_arith_fmin3(self, op: Op) -> None: ) def _op_arith_cmp(self, op: Op) -> None: - pred = op.attrs.get("pred", "lt") - pmap = { - "lt": "slt", - "le": "sle", - "gt": "sgt", - "ge": "sge", - "eq": "eq", - "ne": "ne", - } a, b = op.operands + pred = _icmp_predicate(op.attrs.get("pred", "lt"), a.type) self._current().emit( - f" {op.result.name} = icmp {pmap[pred]} {_llvm_type(a.type)} " + f" {op.result.name} = icmp {pred} {_llvm_type(a.type)} " f"{self._operand(a)}, {self._operand(b)}" ) @@ -2571,6 +2576,7 @@ def _op_arith_not(self, op: Op) -> None: ) def _op_arith_smax(self, op: Op) -> None: + _refuse_unsigned_operands(op, "arith.umax") a, b = op.operands self._need("smax.i32") self._current().emit( @@ -2579,6 +2585,7 @@ def _op_arith_smax(self, op: Op) -> None: ) def _op_arith_smin(self, op: Op) -> None: + _refuse_unsigned_operands(op, "arith.umin") a, b = op.operands self._need("smin.i32") self._current().emit( @@ -2586,6 +2593,21 @@ def _op_arith_smin(self, op: Op) -> None: f"i32 {self._operand(b)})" ) + def _unsigned_min_max(self, op: Op, which: str) -> None: + a, b = op.operands + ty = _llvm_type(a.type) + self._need(f"{which}.{ty}") + self._current().emit( + f" {op.result.name} = call {ty} @llvm.{which}.{ty}({ty} {self._operand(a)}, " + f"{ty} {self._operand(b)})" + ) + + def _op_arith_umax(self, op: Op) -> None: + self._unsigned_min_max(op, "umax") + + def _op_arith_umin(self, op: Op) -> None: + self._unsigned_min_max(op, "umin") + def _op_arith_zext(self, op: Op) -> None: (v,) = op.operands self._current().emit( @@ -2594,6 +2616,7 @@ def _op_arith_zext(self, op: Op) -> None: ) def _op_arith_sext(self, op: Op) -> None: + _refuse_unsigned_operands(op, "arith.zext") (v,) = op.operands self._current().emit( f" {op.result.name} = sext {_llvm_type(v.type)} {self._operand(v)} " @@ -2654,6 +2677,12 @@ def _op_arith_sitofp_f32(self, op: Op) -> None: f" {op.result.name} = sitofp i32 {self._operand(v)} to float" ) + def _op_arith_uitofp_f32(self, op: Op) -> None: + (v,) = op.operands + self._current().emit( + f" {op.result.name} = uitofp {_llvm_type(v.type)} {self._operand(v)} to float" + ) + def _op_arith_cvt_fp8_to_f32(self, op: Op) -> None: """Lower fp8e4m3->f32 conversion to AMDGPU intrinsic. @@ -2990,6 +3019,37 @@ def _op_arith_cvt_f32_to_i8_sat(self, op: Op) -> None: ) self._current().emit(f" {op.result.name} = trunc i32 {smin_v} to i8") + def _op_arith_cvt_f32_to_u8_sat(self, op: Op) -> None: + """Lower saturating f32->u8 (round-to-nearest-even). + + .. code-block:: text + + %r = call float @llvm.rint.f32(float %v) + %lo = call float @llvm.maxnum.f32(float %r, float 0.0) + %hi = call float @llvm.minnum.f32(float %lo, float 255.0) + %q = fptoui float %hi to i8 + + Clamping in f32 keeps ``fptoui`` in range; ``maxnum`` turns a NaN + into 0. + """ + (v,) = op.operands + self._need("rint.f32") + self._need("maxnum.f32") + self._need("minnum.f32") + rounded = f"{op.result.name}r" + low = f"{op.result.name}lo" + high = f"{op.result.name}hi" + self._current().emit( + f" {rounded} = call float @llvm.rint.f32(float {self._operand(v)})" + ) + self._current().emit( + f" {low} = call float @llvm.maxnum.f32(float {rounded}, float 0.000000e+00)" + ) + self._current().emit( + f" {high} = call float @llvm.minnum.f32(float {low}, float 2.550000e+02)" + ) + self._current().emit(f" {op.result.name} = fptoui float {high} to i8") + def _op_arith_rint_f32(self, op: Op) -> None: """Round an f32 to nearest integer (still f32), round-to-nearest-even.""" (v,) = op.operands @@ -3446,11 +3506,12 @@ def _op_tile_smem_store(self, op: Op) -> None: f" {gep} = getelementptr inbounds {agg_ty}, ptr addrspace(3) {base_ptr}, " f"{', '.join(gidx)}" ) - # Alignment is the element byte size: 1 for i8, 2 for f16/bf16, + # Alignment is the element byte size: 1 for i8/u8, 2 for f16/bf16, # 4 for f32/i32, 8 for i64/f64. The AMDGPU backend rejects loads / # stores with under-aligned addresses on ``addrspace(3)``. align = { "i8": 1, + "u8": 1, "fp8e4m3": 1, "bf8e5m2": 1, "f16": 2, @@ -3517,6 +3578,7 @@ def _op_tile_smem_store_vN(self, op: Op) -> None: elem_ty = _llvm_type(value.type.elem) # type: ignore[attr-defined] elem_bytes = { "i8": 1, + "u8": 1, "fp8e4m3": 1, "bf8e5m2": 1, "f16": 2, @@ -3604,6 +3666,7 @@ def _op_tile_smem_load_vN(self, op: Op) -> None: # (i64 / f64): 8 bytes. elem_bytes = { "i8": 1, + "u8": 1, "f16": 2, "bf16": 2, "i32": 4, @@ -6203,6 +6266,7 @@ def _op_memref_global_store_vN(self, op: Op) -> None: ) elem_bytes = { "i8": 1, + "u8": 1, "fp8e4m3": 1, "bf8e5m2": 1, "i16": 2, @@ -6370,22 +6434,15 @@ def _op_vector_lshr(self, op: Op) -> None: self._vector_binop(op, "lshr") def _op_vector_cmp(self, op: Op) -> None: - pred = op.attrs.get("pred", "lt") - pmap = { - "lt": "slt", - "le": "sle", - "gt": "sgt", - "ge": "sge", - "eq": "eq", - "ne": "ne", - } a, b = op.operands + pred = _icmp_predicate(op.attrs.get("pred", "lt"), a.type) self._current().emit( - f" {op.result.name} = icmp {pmap[pred]} {_llvm_type(a.type)} " + f" {op.result.name} = icmp {pred} {_llvm_type(a.type)} " f"{self._operand(a)}, {self._operand(b)}" ) def _op_vector_smax(self, op: Op) -> None: + _refuse_unsigned_operands(op) a, b = op.operands # Prefer the ``llvm.smax.vi`` intrinsic over icmp+select: the # AMDGPU backend reliably lowers the i16 vector form to packed @@ -6410,6 +6467,7 @@ def _op_vector_smax(self, op: Op) -> None: ) def _op_vector_smin(self, op: Op) -> None: + _refuse_unsigned_operands(op) a, b = op.operands cmp = self._fresh("vsmin.cmp") self._current().emit( @@ -6428,6 +6486,7 @@ def _op_vector_trunc(self, op: Op) -> None: ) def _op_vector_sext(self, op: Op) -> None: + _refuse_unsigned_operands(op) (v,) = op.operands self._current().emit( f" {op.result.name} = sext {_llvm_type(v.type)} {self._operand(v)} " @@ -7117,6 +7176,24 @@ def _encode_waitcnt_gfx11(vmcnt: int, expcnt: int, lgkmcnt: int) -> int: return (ec_b & 0x7) | ((lk_b & 0x3F) << 4) | ((vm_b & 0x3F) << 10) +# arith.cmp / vector.cmp predicate -> (signed icmp predicate, unsigned one). +_ICMP_PREDICATES = { + "lt": ("slt", "ult"), + "le": ("sle", "ule"), + "gt": ("sgt", "ugt"), + "ge": ("sge", "uge"), + "eq": ("eq", "eq"), + "ne": ("ne", "ne"), +} + + +def _icmp_predicate(pred: str, t: Type) -> str: + """The icmp predicate for ``pred`` on operands of type ``t``: ordered + compares of unsigned integers are unsigned.""" + signed, unsigned = _ICMP_PREDICATES[pred] + return unsigned if _is_unsigned(t) else signed + + def _fp64_hex(x: float) -> str: # LLVM textual IR spells `double` hex constants as the raw IEEE-754 bits. import struct diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/gpu.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/gpu.py index c6c221b50ffe..a3a401273c7a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/gpu.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/gpu.py @@ -31,6 +31,7 @@ np.int32: "i32", np.int16: "i16", np.int8: "i8", + np.uint8: "u8", } 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 615def4e59f7..bb6dc92e616a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/CMakeLists.txt @@ -46,6 +46,15 @@ if(BUILD_TESTING) 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_ir_types core/test_ir_types.cpp) + target_link_libraries(rocke_ir_types PRIVATE rocke_core) + add_test(NAME rocke_ir_types COMMAND rocke_ir_types) + set_tests_properties(rocke_ir_types PROPERTIES LABELS "ckc;unit;host") + rocke_cxx_test_installed_name(_rocke_installed_name rocke_ir_types) + set_target_properties(rocke_ir_types PROPERTIES OUTPUT_NAME "${_rocke_installed_name}") + install(TARGETS rocke_ir_types RUNTIME DESTINATION ${ROCKE_CXX_TEST_INSTALL_DIR}) + set_property(GLOBAL APPEND PROPERTY ROCKE_INSTALL_TEST_TARGETS rocke_ir_types) + 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/README.md b/dnn-providers/hip-kernel-provider/rocke/platform/tests/README.md index fd9a7268bb05..df7d6d02925a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/README.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/README.md @@ -130,14 +130,16 @@ divergence - the Python builder correctly rejects wave32 WMMA on gfx942.) `run_all.py --build-root ` builds all configured targets before pytest, then obtains each native parity executable from CTest, as listed in -`NATIVE_PARITY_TESTS`: `rocke_storage` (`ROCKE_STORAGE_TEST`) and -`rocke_nontemporal_hip` (`ROCKE_NONTEMPORAL_HIP_TEST`). Both pytest passes -receive those paths, so the storage IR/HIP parity, serialization and nontemporal -HIP-source parity tests run automatically. `--config` selects the native test +`NATIVE_PARITY_TESTS`: `rocke_storage` (`ROCKE_STORAGE_TEST`), +`rocke_ir_types` (`ROCKE_IR_TYPES_TEST`) and `rocke_nontemporal_hip` +(`ROCKE_NONTEMPORAL_HIP_TEST`). Both pytest passes receive those paths, so the +storage IR/HIP parity, serialization, C++ builder type-rule parity and +nontemporal HIP-source parity 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` / `ROCKE_NONTEMPORAL_HIP_TEST` overrides +An explicit `ROCKE_STORAGE_TEST` / `ROCKE_IR_TYPES_TEST` / +`ROCKE_NONTEMPORAL_HIP_TEST` overrides discovery for that binary and must name an existing executable; it does not skip the configured build. With no configured build or explicit override, the runner reports that parity as skipped; direct pytest invocations can use the same diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_dtypes.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_dtypes.cpp index 81d3694840d2..3ce6e2441f0d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_dtypes.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_dtypes.cpp @@ -7,6 +7,7 @@ #include #include +#include int main() { @@ -24,6 +25,8 @@ int main() {"XF32", "tf32"}, {" TF32 ", "tf32"}, {"INT8", "i8"}, + {"UINT8", "u8"}, + {"uint4", "u4"}, {" Custom_Format ", "custom_format"}, {"", ""}}; for(const auto& entry : cases) @@ -45,5 +48,12 @@ int main() const rocke_dtype_info_t* i4 = rocke_dtype_info("i4"); if(!i4 || i4->encoded_bits != 4 || rocke_dtype_info("iu8") || rocke_dtype_info("iu4")) return 1; + /* u8/u4 are logical unsigned integer types. */ + for(const auto& [name, bits] : {std::pair{"u8", 8}, std::pair{"u4", 4}}) + { + const rocke_dtype_info_t* info = rocke_dtype_info(name); + if(!info || info->encoded_bits != bits || info->category != ROCKE_DTYPE_INTEGER) + return 1; + } return 0; } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_dtypes.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_dtypes.py index aab71559badb..a21a7e31314c 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 @@ -182,3 +182,39 @@ def test_quant_dtype_annotations_match_scalar_and_logical_contracts(): 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 + + +def test_u8_and_u4_are_logical_unsigned_types(): + # Types only: they name unsigned data for matrix instructions with a signedness flag. + import rocke + from rocke.core.ir import U4, U8, IRBuilder, PtrType + from rocke.tests.engines import ( + lower_hip_with_both_engines, + lower_with_both_engines, + ) + from rocke.tests.ir_validity import assert_target_accepts + + for spelling, ir_type, bits in ( + ("u8", U8, 8), + ("uint8", U8, 8), + ("u4", U4, 4), + ("uint4", U4, 4), + ): + assert dtype_to_ir_type(spelling) is ir_type + assert dtype_info(spelling).encoded_bits == bits + assert dtype_info(spelling).category.value == "integer" + assert U8 != I8 + assert rocke.U8 is U8 and rocke.U4 is U4 + assert storage_ir_type("u4") is I8 + # U8 values are LLVM i8 (LLVM integers carry no sign) and HIP uint8_t. + b = IRBuilder("u8_pointer") + src = b.param("SRC", PtrType(U8, "global")) + dst = b.param("DST", PtrType(U8, "global")) + tid = b.thread_id_x() + loaded = b.global_load(src, tid, U8) + b.global_store(dst, tid, loaded) + ir_text = lower_with_both_engines(b.kernel, "gfx90a") + assert "load i8, ptr addrspace(1)" in ir_text + assert_target_accepts(ir_text, "gfx90a") + hip = lower_hip_with_both_engines(b.kernel, "gfx90a") + assert "uint8_t* __restrict__ SRC" in hip diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp new file mode 100644 index 000000000000..1d998c69a538 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp @@ -0,0 +1,433 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* The IR type rules the C++ builders apply, checked against the Python builders. + * + * The serialized-IR parity seam (rocke.tests.engines) lowers a kernel the + * PYTHON builder built, so it never runs a C++ rocke_b_* check. This test does. + * + * Without arguments it self-checks with the literals core/ir.py uses. With + * arguments it answers one question for tests/core/test_u8_memory.py, which + * compares the answer with the Python builder (found through + * ROCKE_IR_TYPES_TEST): + * + * --unsigned + * build the u8 op named (a row of k_unsigned_cases, mirrored in + * tests/core/test_u8_arith.py) on u8 or i32 operands; print "ok", or + * "error: " when the builder refuses it. + * + * --packed-only + * make one IR value, or LDS tile, of the packed-only elem with the + * builder named by a row of k_packed_only_builders (the rows of + * tests/core/test_packed_only.py); print "ok", or "error: " + * when the builder refuses it. + * + * --smem-vector + * build one LDS vector access of n elem values; print "ok", or + * "error: " when the builder refuses it. + */ +#include "rocke/ir.h" + +#include +#include +#include +#include + +#include "rocke/error.hpp" + +#define CHECK(x) \ + do \ + { \ + if(!(x)) \ + { \ + fprintf(stderr, "check failed at %d: %s\n", __LINE__, #x); \ + return 1; \ + } \ + } while(0) + +/* The outcome of one builder call: accepted, or refused with a message. */ +struct build_outcome +{ + bool accepted; + std::string error; + /* Ops the measured call appended to the kernel body, counted after a + * refusal too, so a refused call that left an op behind is caught; -1 when + * the setup was refused and the measured call never ran. */ + int ops_added; +}; + +/* Build an LDS allocation of `elem` and one vector access of `n` elements: + * smem_load_vN, or smem_store_vN of vector_splat(, n). Only the + * access is measured; the setup ops are counted out. */ +static build_outcome smem_vector_access(bool store, const char* elem_name, int n) +{ + build_outcome outcome = {false, "", 0}; + const rocke_type_t* elem = rocke_scalar_by_name(elem_name); + if(!elem) + { + outcome.error = std::string("no scalar type named '") + elem_name + "'"; + return outcome; + } + rocke_ir_builder_t b; + if(rocke_ir_builder_init(&b, "smem_vector_access") != ROCKE_OK) + { + outcome.error = "builder init failed"; + return outcome; + } + int ops_before = -1; /* set once the setup is built */ + try + { + const int shape[] = {64}; + rocke_value_t* smem = rocke_b_smem_alloc(&b, elem, shape, 1, "tile"); + rocke_value_t* zero = rocke_b_const_i32(&b, 0); + rocke_value_t* value = NULL; + if(store) + { + rocke_value_t* src = rocke_b_param(&b, "SRC", rocke_ptr_type(&b, elem, "global"), NULL); + value = rocke_b_vector_splat(&b, rocke_b_global_load(&b, src, zero, elem, 1), n); + } + ops_before = b.kernel->body->num_ops; + if(store) + rocke_b_smem_store_vN(&b, smem, &zero, 1, value, n); + else + rocke_b_smem_load_vN(&b, smem, &zero, 1, elem, n); + outcome.accepted = true; + } + catch(const ckc::Error& error) + { + outcome.error = error.what(); + } + outcome.ops_added = ops_before < 0 ? -1 : b.kernel->body->num_ops - ops_before; + rocke_ir_builder_free(&b); + return outcome; +} + +static int print_smem_vector_access(const char* direction, const char* elem, const char* count) +{ + const bool store = strcmp(direction, "store") == 0; + CHECK(store || strcmp(direction, "load") == 0); + const build_outcome outcome = smem_vector_access(store, elem, atoi(count)); + if(outcome.accepted) + printf("ok\n"); + else + printf("error: %s\n", outcome.error.c_str()); + return 0; +} + +/* A builder that would make an IR value or LDS tile of its element type, and + * the IR op its refusal names. */ +struct packed_only_builder +{ + const char* name; + const char* op; +}; + +static const packed_only_builder k_packed_only_builders[] = { + {"param", "param"}, + {"param_vector", "param"}, + {"zero_vec", "zero_vec"}, + {"undef_vec", "vector.undef"}, + {"trunc", "arith.trunc"}, + {"vec_bitcast", "vector.bitcast"}, + {"global_load", "memref.global_load_typed"}, + {"smem_alloc", "tile.smem_alloc"}, +}; + +/* Build the operands, then make one `elem_name` value or LDS tile with + * `builder`; only that call is measured. */ +static build_outcome + packed_only_value(const char* builder, const char* elem_name, rocke_status_t* code) +{ + build_outcome outcome = {false, "", 0}; + *code = ROCKE_OK; + const rocke_type_t* elem = rocke_scalar_by_name(elem_name); + if(!elem) + { + outcome.error = std::string("no scalar type named '") + elem_name + "'"; + return outcome; + } + rocke_ir_builder_t b; + if(rocke_ir_builder_init(&b, "packed_only_value") != ROCKE_OK) + { + outcome.error = "builder init failed"; + return outcome; + } + int ops_before = -1; /* set once the operands are built */ + try + { + rocke_value_t* zero = rocke_b_const_i32(&b, 0); + rocke_value_t* word = rocke_b_zero_vec(&b, rocke_i32(), 1); + rocke_value_t* ptr = rocke_b_param(&b, "P", rocke_ptr_type(&b, elem, "global"), NULL); + const rocke_type_t* vector = rocke_vector_type(&b, elem, 8); + const int shape[] = {64}; + ops_before = b.kernel->body->num_ops; + if(strcmp(builder, "param") == 0) + rocke_b_param(&b, "X", elem, NULL); + else if(strcmp(builder, "param_vector") == 0) + rocke_b_param(&b, "X", vector, NULL); + else if(strcmp(builder, "zero_vec") == 0) + rocke_b_zero_vec(&b, elem, 8); + else if(strcmp(builder, "undef_vec") == 0) + rocke_b_undef_vec(&b, elem, 8); + else if(strcmp(builder, "trunc") == 0) + rocke_b_trunc(&b, zero, elem); + else if(strcmp(builder, "vec_bitcast") == 0) + rocke_b_vec_bitcast(&b, word, vector); + else if(strcmp(builder, "global_load") == 0) + rocke_b_global_load(&b, ptr, zero, elem, 1); + else if(strcmp(builder, "smem_alloc") == 0) + rocke_b_smem_alloc(&b, elem, shape, 1, "tile"); + else + throw ckc::ValueError(std::string("no packed-only builder named '") + builder + "'"); + outcome.accepted = true; + } + catch(const ckc::Error& error) + { + outcome.error = error.what(); + *code = error.code(); + } + outcome.ops_added = ops_before < 0 ? -1 : b.kernel->body->num_ops - ops_before; + rocke_ir_builder_free(&b); + return outcome; +} + +static int print_packed_only_value(const char* builder, const char* elem) +{ + rocke_status_t code; + const build_outcome outcome = packed_only_value(builder, elem, &code); + if(outcome.accepted) + printf("ok\n"); + else + printf("error: %s\n", outcome.error.c_str()); + return 0; +} + +/* One builder call on u8 operands (or i32 ones, to show what the unsigned ops + * refuse). The rows match _UNSIGNED_CASES in tests/core/test_u8_arith.py. */ +struct unsigned_case +{ + const char* name; + rocke_value_t* (*build)(rocke_ir_builder_t* b, rocke_value_t* x, rocke_value_t* y); + /* The builder's refusal status and text; error is NULL when the call builds + * one op. */ + rocke_status_t code; + const char* error; +}; + +static const unsigned_case k_unsigned_cases[] = { + {"umax", + [](rocke_ir_builder_t* b, rocke_value_t* x, rocke_value_t* y) { + return rocke_b_umax(b, x, y); + }, + ROCKE_OK, + NULL}, + {"umin", + [](rocke_ir_builder_t* b, rocke_value_t* x, rocke_value_t* y) { + return rocke_b_umin(b, x, y); + }, + ROCKE_OK, + NULL}, + {"uitofp_f32", + [](rocke_ir_builder_t* b, rocke_value_t* x, rocke_value_t*) { + return rocke_b_uitofp_f32(b, x); + }, + ROCKE_OK, + NULL}, + {"zext", + [](rocke_ir_builder_t* b, rocke_value_t* x, rocke_value_t*) { + return rocke_b_zext(b, x, rocke_i32()); + }, + ROCKE_OK, + NULL}, + {"smax", + [](rocke_ir_builder_t* b, rocke_value_t* x, rocke_value_t* y) { + return rocke_b_smax(b, x, y); + }, + ROCKE_ERR_TYPE, + "smax treats its operands as signed, but u8 is unsigned -- use umax"}, + {"smin", + [](rocke_ir_builder_t* b, rocke_value_t* x, rocke_value_t* y) { + return rocke_b_smin(b, x, y); + }, + ROCKE_ERR_TYPE, + "smin treats its operands as signed, but u8 is unsigned -- use umin"}, + {"sext", + [](rocke_ir_builder_t* b, rocke_value_t* x, rocke_value_t*) { + return rocke_b_sext(b, x, rocke_i32()); + }, + ROCKE_ERR_TYPE, + "sext treats its operands as signed, but u8 is unsigned -- use zext"}, + {"vector_smax", + [](rocke_ir_builder_t* b, rocke_value_t* x, rocke_value_t* y) { + return rocke_b_vector_smax( + b, rocke_b_vector_splat(b, x, 4), rocke_b_vector_splat(b, y, 4)); + }, + ROCKE_ERR_TYPE, + "vector_smax treats its operands as signed, but u8 is unsigned"}, + {"vector_smin", + [](rocke_ir_builder_t* b, rocke_value_t* x, rocke_value_t* y) { + return rocke_b_vector_smin( + b, rocke_b_vector_splat(b, x, 4), rocke_b_vector_splat(b, y, 4)); + }, + ROCKE_ERR_TYPE, + "vector_smin treats its operands as signed, but u8 is unsigned"}, + {"vector_sext", + [](rocke_ir_builder_t* b, rocke_value_t* x, rocke_value_t*) { + return rocke_b_vector_sext(b, rocke_b_vector_splat(b, x, 4), rocke_i32()); + }, + ROCKE_ERR_TYPE, + "vector_sext treats its operands as signed, but u8 is unsigned"}, + {"umax_of_i32", + [](rocke_ir_builder_t* b, rocke_value_t*, rocke_value_t*) { + rocke_value_t* i = rocke_b_const_i32(b, 1); + return rocke_b_umax(b, i, i); + }, + ROCKE_ERR_VALUE, + "umax expects two operands of one unsigned integer type, got i32 and i32"}, + {"uitofp_f32_of_i32", + [](rocke_ir_builder_t* b, rocke_value_t*, rocke_value_t*) { + return rocke_b_uitofp_f32(b, rocke_b_const_i32(b, 1)); + }, + ROCKE_ERR_VALUE, + "uitofp_f32 expects an unsigned integer input, got i32"}, + {"cvt_f32_to_u8_sat_of_u8", + [](rocke_ir_builder_t* b, rocke_value_t* x, rocke_value_t*) { + return rocke_b_cvt_f32_to_u8_sat(b, x); + }, + ROCKE_ERR_VALUE, + "cvt_f32_to_u8_sat expects f32 input, got u8"}, +}; + +static const unsigned_case* find_unsigned_case(const char* name) +{ + for(const unsigned_case& c : k_unsigned_cases) + if(strcmp(c.name, name) == 0) + return &c; + return NULL; +} + +/* Declare two u8 parameters X and Y, then run `c`; ops_added counts every op + * the case adds, the operands it builds itself included. */ +static build_outcome build_unsigned_case(const unsigned_case& c, rocke_status_t* code) +{ + build_outcome outcome = {false, "", 0}; + *code = ROCKE_OK; + rocke_ir_builder_t b; + if(rocke_ir_builder_init(&b, "unsigned_case") != ROCKE_OK) + { + outcome.error = "builder init failed"; + return outcome; + } + int ops_before = -1; /* set once the parameters are declared */ + try + { + rocke_value_t* x = rocke_b_param(&b, "X", rocke_u8(), NULL); + rocke_value_t* y = rocke_b_param(&b, "Y", rocke_u8(), NULL); + ops_before = b.kernel->body->num_ops; + c.build(&b, x, y); + outcome.accepted = true; + } + catch(const ckc::Error& error) + { + outcome.error = error.what(); + *code = error.code(); + } + outcome.ops_added = ops_before < 0 ? -1 : b.kernel->body->num_ops - ops_before; + rocke_ir_builder_free(&b); + return outcome; +} + +static int print_unsigned_case(const char* name) +{ + const unsigned_case* c = find_unsigned_case(name); + CHECK(c != NULL); + rocke_status_t code; + const build_outcome outcome = build_unsigned_case(*c, &code); + if(outcome.accepted) + printf("ok\n"); + else + printf("error: %s\n", outcome.error.c_str()); + return 0; +} + +/* The unsigned ops build one op on u8; the signed-by-definition ops refuse u8, + * and the unsigned ops refuse a signed operand, with core/ir.py's text. */ +static int test_unsigned_ops() +{ + for(const unsigned_case& c : k_unsigned_cases) + { + rocke_status_t code; + const build_outcome outcome = build_unsigned_case(c, &code); + if(c.error == NULL) + { + CHECK(outcome.accepted); + CHECK(outcome.ops_added == 1); + continue; + } + CHECK(!outcome.accepted); + CHECK(code == c.code); + CHECK(outcome.error == c.error); + } + return 0; +} + +/* Every builder refuses an i4 or u4 value or LDS tile with core/ir.py's + * TypeError text, before any op is built. */ +static int test_packed_only_values() +{ + for(const char* elem : {"i4", "u4"}) + for(const packed_only_builder& builder : k_packed_only_builders) + { + rocke_status_t code; + const build_outcome outcome = packed_only_value(builder.name, elem, &code); + CHECK(!outcome.accepted); + CHECK(code == ROCKE_ERR_TYPE); + CHECK(outcome.error + == std::string(builder.op) + ": " + elem + + " values travel packed in 32-bit words and have no scalar or vector " + "form -- keep them in i32 storage and unpack explicitly"); + CHECK(outcome.ops_added == 0); + } + return 0; +} + +/* The LDS vector pair admits one element list (core/ir.py + * _SMEM_VECTOR_ELEMENTS), so u8 is refused in both directions with the same + * text, before any op is built; admitted elements build exactly one op. */ +static int test_smem_vector_elements() +{ + static const char* const admitted + = "f16 / bf16 / i16 / f32 / i32 / tf32 / f64 / fp8e4m3 / bf8e5m2 / i8"; + for(bool store : {false, true}) + { + const build_outcome refused = smem_vector_access(store, "u8", 4); + CHECK(!refused.accepted); + CHECK(refused.error + == std::string(store ? "smem_store_vN" : "smem_load_vN") + " supports " + admitted + + ", got u8"); + CHECK(refused.ops_added == 0); + for(const char* elem : {"i8", "i16", "f64"}) + { + const build_outcome built = smem_vector_access(store, elem, 2); + CHECK(built.accepted); + CHECK(built.ops_added == 1); + } + CHECK(smem_vector_access(store, "i8", 16).accepted); + } + return 0; +} + +int main(int argc, char** argv) +{ + if(argc == 3 && strcmp(argv[1], "--unsigned") == 0) + return print_unsigned_case(argv[2]); + if(argc == 4 && strcmp(argv[1], "--packed-only") == 0) + return print_packed_only_value(argv[2], argv[3]); + if(argc == 5 && strcmp(argv[1], "--smem-vector") == 0) + return print_smem_vector_access(argv[2], argv[3], argv[4]); + CHECK(argc == 1); + CHECK(test_unsigned_ops() == 0); + CHECK(test_packed_only_values() == 0); + CHECK(test_smem_vector_elements() == 0); + printf("rocke_ir_types: ok\n"); + return 0; +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_packed_only.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_packed_only.py index 5237f02cf9d0..424cea6cc3ca 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_packed_only.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_packed_only.py @@ -2,12 +2,20 @@ # SPDX-License-Identifier: MIT """Packed-only types have no IR value. -I4 values only ever live packed in 32-bit words: a kernel keeps the words in -i32 storage and unpacks them explicitly. So the IR builder refuses every op -result, kernel parameter and LDS tile of i4 elements, and the lowerings, which -have no rendering for the type, still refuse one that reaches them. +I4 and U4 values only ever live packed in 32-bit words: a kernel keeps the +words in i32 storage and unpacks them explicitly. So the IR builder of each +engine refuses every op result, kernel parameter and LDS tile of i4 or u4 +elements with the same text, and the lowerings, which have no rendering for +either type, still refuse one that reaches them. + +The serialized-IR seam (rocke.tests.engines) lowers kernels the Python builder +built, so the C++ builder's refusals are asked of the rocke_ir_types test +binary (ROCKE_IR_TYPES_TEST). """ +import os +import subprocess + import pytest from rocke.core.ir import I32, IRBuilder, Param, PtrType, VectorType, dtype_to_ir_type @@ -16,10 +24,11 @@ assert_both_engines_reject_hip, ) -PACKED_ONLY_DTYPES = ["i4"] +PACKED_ONLY_DTYPES = ["i4", "u4"] # Builders that would make an IR value, or an LDS tile, of the element type -# they are given, with the IR op each refusal names. +# they are given, with the IR op each refusal names. rocke_ir_types builds the +# same calls under the same names. _BUILDERS = { "param": ("param", lambda b, t: b.param("X", t)), "param_vector": ("param", lambda b, t: b.param("X", VectorType(t, 8))), @@ -57,6 +66,22 @@ def test_builder_refuses_a_packed_only_value(builder, dtype): assert str(error.value) == _refusal(op, dtype) +@pytest.mark.parametrize("dtype", PACKED_ONLY_DTYPES) +@pytest.mark.parametrize("builder", sorted(_BUILDERS)) +def test_native_builder_refuses_a_packed_only_value(builder, dtype): + executable = os.environ.get("ROCKE_IR_TYPES_TEST") + if not executable: + pytest.skip("set ROCKE_IR_TYPES_TEST to the built rocke_ir_types test") + native = subprocess.run( + [executable, "--packed-only", builder, dtype], + check=True, + capture_output=True, + text=True, + ).stdout.rstrip("\n") + op, _ = _BUILDERS[builder] + assert native == f"error: {_refusal(op, dtype)}" + + @pytest.mark.parametrize("dtype", PACKED_ONLY_DTYPES) def test_lowerings_refuse_a_packed_only_value(dtype): # The builder refuses the parameter, so the kernel is assembled by hand to diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_u8_arith.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_u8_arith.py new file mode 100644 index 000000000000..b85c53534aa6 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_u8_arith.py @@ -0,0 +1,385 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Unsigned arithmetic, compares and conversions of U8. + +U8 renders as LLVM i8, and LLVM integers carry no sign, so every op whose +result depends on signedness must emit its unsigned form for a u8 operand: +udiv / urem, icmp ult / ule / ugt / uge, no nsw flag on add / sub / mul, zext, +lshr, llvm.umax / llvm.umin, uitofp and fptoui. The ops that are signed by +definition (smax, smin, sext and their vector forms) refuse a u8 operand. The +C++ builders are reached through the ``rocke_ir_types`` test binary +(``ROCKE_IR_TYPES_TEST``), because the serialized-IR seam lowers kernels the +Python builder built. The numeric test runs on a visible gfx90a only, through +the active engine, so a strict both-engines run +(``ROCKE_BACKEND=both ROCKE_CPP_STRICT=1``) checks both engines on the GPU. +""" + +import os +import re +import shutil +import subprocess + +import numpy as np +import pytest +from rocke.core.ir import F32, I8, I32, U8, IRBuilder, PtrType +from rocke.core.ir_serialize import parse, serialize +from rocke.core.lower_hip import lower_kernel_to_hip +from rocke.tests.engines import ( + assert_both_engines_reject, + assert_both_engines_reject_hip, + lower_hip_with_both_engines, + lower_with_both_engines, +) +from rocke.tests.gpu import run_on_device +from rocke.tests.ir_validity import assert_target_accepts + +# Output buffer -> its element type, in kernel parameter order after the inputs. +_OUTPUTS = { + "DIV": U8, + "REM": U8, + "LT": U8, + "LE": U8, + "GT": U8, + "GE": U8, + "MAX": U8, + "MIN": U8, + "SHR": U8, + "ADD": U8, + "SUB": U8, + "MUL": U8, + "TO_F32": F32, + "WIDE": I32, + "SAT": U8, +} + +_ALIGN = {U8: 1, F32: 4, I32: 4} + + +def _u8_arith_kernel(): + """Per thread, one result of every u8 op on A[i] and B[i], and the saturating + f32 -> u8 conversion of F[i]. Compares store 0 or 1; SHR shifts by B & 7. + + REM divides by C[i], not B[i]: LLVM's AMDGPU backend at -O0 merges a division + and a remainder of the same i8 operands into one i16 divrem it cannot select, + in either signedness, and the IR-validity check compiles at -O0.""" + b = IRBuilder("u8_arith") + a_ptr = b.param("A", PtrType(U8, "global")) + b_ptr = b.param("B", PtrType(U8, "global")) + c_ptr = b.param("C", PtrType(U8, "global")) + f_ptr = b.param("F", PtrType(F32, "global")) + out = {name: b.param(name, PtrType(t, "global")) for name, t in _OUTPUTS.items()} + tid = b.thread_id_x() + x = b.global_load(a_ptr, tid, U8) + y = b.global_load(b_ptr, tid, U8) + z = b.global_load(c_ptr, tid, U8) + f = b.global_load(f_ptr, tid, F32, align=4) + seven = b.trunc(b.const_i32(7), U8) + results = { + "DIV": b.div(x, y), + "REM": b.mod(x, z), + "LT": b.zext(b.cmp_lt(x, y), U8), + "LE": b.zext(b.cmp_le(x, y), U8), + "GT": b.zext(b.cmp_gt(x, y), U8), + "GE": b.zext(b.cmp_ge(x, y), U8), + "MAX": b.umax(x, y), + "MIN": b.umin(x, y), + "SHR": b.lshr(x, b.land(y, seven)), + "ADD": b.add(x, y), + "SUB": b.sub(x, y), + "MUL": b.mul(x, y), + "TO_F32": b.uitofp_f32(x), + "WIDE": b.zext(x, I32), + "SAT": b.cvt_f32_to_u8_sat(f), + } + for name, value in results.items(): + b.global_store(out[name], tid, value, align=_ALIGN[value.type]) + return b.kernel + + +def test_u8_ops_lower_to_their_unsigned_llvm_forms(): + text = lower_with_both_engines(_u8_arith_kernel(), "gfx90a") + for form in ( + r"= udiv i8 ", + r"= urem i8 ", + r"= icmp ult i8 ", + r"= icmp ule i8 ", + r"= icmp ugt i8 ", + r"= icmp uge i8 ", + r"= call i8 @llvm\.umax\.i8\(i8 ", + r"= call i8 @llvm\.umin\.i8\(i8 ", + r"= lshr i8 ", + r"= add i8 ", + r"= sub i8 ", + r"= mul i8 ", + r"= uitofp i8 \S+ to float", + r"= zext i8 \S+ to i32", + r"= fptoui float \S+ to i8", + ): + assert len(re.findall(form, text)) == 1, form + # No signed form of any i8 op, and no signed-overflow flag on one. + assert not re.search( + r"\b(sdiv|srem|ashr|sext|sitofp) i8 |icmp s\w+ i8 |nsw i8 ", text + ) + assert "declare i8 @llvm.umax.i8(i8, i8)" in text + assert "declare i8 @llvm.umin.i8(i8, i8)" in text + assert_target_accepts(text, "gfx90a") + + +def test_u8_ops_lower_to_uint8_t_in_hip(): + hip = lower_hip_with_both_engines(_u8_arith_kernel(), "gfx90a") + # Every u8 result is a uint8_t, so C's own operators are the unsigned ones. + for hint in ("div", "mod", "umax", "umin", "add", "sub", "mul", "qu8"): + assert re.search(rf"\buint8_t \w*{hint}\w* = ", hip), hint + assert re.search(r"\bfloat \w*uitof\w* = \(float\)", hip) + + +def test_u8_hip_source_compiles(tmp_path): + if not shutil.which("hipcc"): + pytest.skip("hipcc not in PATH") + path = tmp_path / "kernel.hip" + path.write_text(lower_kernel_to_hip(_u8_arith_kernel(), arch="gfx90a")) + subprocess.run( + [ + "hipcc", + "--offload-arch=gfx90a", + "--cuda-device-only", + "-fsyntax-only", + str(path), + ], + check=True, + capture_output=True, + text=True, + ) + + +# One builder call on two u8 parameters X and Y, its expected refusal (None when +# it builds one op), as the rows of k_unsigned_cases in test_ir_types.cpp. +_UNSIGNED_CASES = { + "umax": (lambda b, x, y: b.umax(x, y), None), + "umin": (lambda b, x, y: b.umin(x, y), None), + "uitofp_f32": (lambda b, x, y: b.uitofp_f32(x), None), + "zext": (lambda b, x, y: b.zext(x, I32), None), + "smax": ( + lambda b, x, y: b.smax(x, y), + ( + TypeError, + "smax treats its operands as signed, but u8 is unsigned -- use umax", + ), + ), + "smin": ( + lambda b, x, y: b.smin(x, y), + ( + TypeError, + "smin treats its operands as signed, but u8 is unsigned -- use umin", + ), + ), + "sext": ( + lambda b, x, y: b.sext(x, I32), + ( + TypeError, + "sext treats its operands as signed, but u8 is unsigned -- use zext", + ), + ), + "vector_smax": ( + lambda b, x, y: b.vector_smax(b.vector_splat(x, 4), b.vector_splat(y, 4)), + (TypeError, "vector_smax treats its operands as signed, but u8 is unsigned"), + ), + "vector_smin": ( + lambda b, x, y: b.vector_smin(b.vector_splat(x, 4), b.vector_splat(y, 4)), + (TypeError, "vector_smin treats its operands as signed, but u8 is unsigned"), + ), + "vector_sext": ( + lambda b, x, y: b.vector_sext(b.vector_splat(x, 4), I32), + (TypeError, "vector_sext treats its operands as signed, but u8 is unsigned"), + ), + "umax_of_i32": ( + lambda b, x, y: b.umax(b.const_i32(1), b.const_i32(1)), + ( + ValueError, + "umax expects two operands of one unsigned integer type, got i32 and i32", + ), + ), + "uitofp_f32_of_i32": ( + lambda b, x, y: b.uitofp_f32(b.const_i32(1)), + (ValueError, "uitofp_f32 expects an unsigned integer input, got i32"), + ), + "cvt_f32_to_u8_sat_of_u8": ( + lambda b, x, y: b.cvt_f32_to_u8_sat(x), + (ValueError, "cvt_f32_to_u8_sat expects f32 input, got u8"), + ), +} + + +def _python_unsigned_case(name: str) -> str: + """Run one case with the Python builder; ``ok`` or ``error: ``.""" + build, _ = _UNSIGNED_CASES[name] + b = IRBuilder("unsigned_case") + x = b.param("X", U8) + y = b.param("Y", U8) + try: + build(b, x, y) + except (TypeError, ValueError) as error: + return f"error: {error}" + return "ok" + + +@pytest.mark.parametrize("name", sorted(_UNSIGNED_CASES)) +def test_unsigned_builders(name): + build, refusal = _UNSIGNED_CASES[name] + b = IRBuilder("unsigned_case") + x = b.param("X", U8) + y = b.param("Y", U8) + if refusal is None: + ops_before = len(b.kernel.body.ops) + build(b, x, y) + assert len(b.kernel.body.ops) == ops_before + 1 + return + error_type, message = refusal + with pytest.raises(error_type) as error: + build(b, x, y) + assert str(error.value) == message + + +@pytest.mark.parametrize("name", sorted(_UNSIGNED_CASES)) +def test_native_unsigned_builders_match_the_python_builder(name): + executable = os.environ.get("ROCKE_IR_TYPES_TEST") + if not executable: + pytest.skip("set ROCKE_IR_TYPES_TEST to the built rocke_ir_types test") + native = subprocess.run( + [executable, "--unsigned", name], + check=True, + capture_output=True, + text=True, + ).stdout.rstrip("\n") + assert native == _python_unsigned_case(name) + + +@pytest.mark.parametrize("op_name", ["smax", "smin"]) +def test_vector_binary_refuses_a_signed_op_on_u8(op_name): + # vector_binary builds every vector., so it refuses a signed op itself + # rather than leaving the refusal to the vector_smax / vector_smin wrappers. + b = IRBuilder("vector_binary_case") + v = b.vector_splat(b.param("X", U8), 4) + with pytest.raises(TypeError) as error: + b.vector_binary(op_name, v, v) + assert str(error.value) == ( + f"vector_{op_name} treats its operands as signed, but u8 is unsigned" + ) + + +# Signed op -> (builds it on an i8 value x, the unsigned op it names, if any). +_SIGNED_OPS = { + "arith.smax": (lambda b, x: b.smax(x, x), "arith.umax"), + "arith.smin": (lambda b, x: b.smin(x, x), "arith.umin"), + "arith.sext": (lambda b, x: b.sext(x, I32), "arith.zext"), + "vector.smax": ( + lambda b, x: b.vector_smax(b.vector_splat(x, 4), b.vector_splat(x, 4)), + None, + ), + "vector.smin": ( + lambda b, x: b.vector_smin(b.vector_splat(x, 4), b.vector_splat(x, 4)), + None, + ), + "vector.sext": (lambda b, x: b.vector_sext(b.vector_splat(x, 4), I32), None), +} + + +@pytest.mark.parametrize("op_name", sorted(_SIGNED_OPS)) +def test_lowerings_refuse_a_signed_op_on_u8(op_name): + # The builders refuse the op on u8, so the kernel is built on i8 and its + # serialized text retyped to u8: IR that did not come through a builder, + # which each lowering must refuse with the builders' text under the op's + # name. HIP has no u8 vector type, so the HIP lowerings refuse a u8 vector + # before they reach a vector op. + build, instead = _SIGNED_OPS[op_name] + b = IRBuilder("signed_op") + build(b, b.param("X", I8)) + b.ret() + kernel = parse(serialize(b.kernel).replace("i8", "u8")) + message = f"{op_name} treats its operands as signed, but u8 is unsigned" + if instead is not None: + message += f" -- use {instead}" + assert_both_engines_reject(kernel, "gfx90a", TypeError, message) + if op_name.startswith("vector."): + assert_both_engines_reject_hip( + kernel, "gfx90a", NotImplementedError, "no HIP mapping for type vec" + ) + else: + assert_both_engines_reject_hip(kernel, "gfx90a", TypeError, message) + + +# ---- on GPU (gfx90a) --------------------------------------------------------- + +_LANES = 64 + + +def _inputs(): + """A, B and C cover both halves of the u8 range (B and C are never 0, so DIV + and REM are defined); F holds the rounding ties, the clamp edges, the infinities and + a NaN, then ordinary values.""" + rng = np.random.default_rng(0) + a = np.array( + [0, 1, 2, 127, 128, 129, 200, 254, 255, 255, 128, 3, 250, 17, 255, 0] + + list(rng.integers(0, 256, _LANES - 16)), + dtype=np.uint8, + ) + b = np.array( + [1, 255, 2, 128, 127, 129, 3, 255, 1, 2, 255, 200, 7, 17, 254, 9] + + list(rng.integers(1, 256, _LANES - 16)), + dtype=np.uint8, + ) + c = np.array( + [255, 1, 128, 129, 127, 2, 200, 3, 254, 255, 5, 128, 251, 1, 255, 10] + + list(rng.integers(1, 256, _LANES - 16)), + dtype=np.uint8, + ) + f_edges = [ + -3.5, -0.4, 0.5, 1.5, 2.5, 127.5, 128.5, 254.5, + 255.4, 255.6, 300.0, np.inf, -np.inf, np.nan, 0.0, -0.0, + ] # fmt: skip + f = np.array( + f_edges + list(rng.uniform(-20.0, 280.0, _LANES - len(f_edges))), + dtype=np.float32, + ) + return a, b, c, f + + +def _expected(a, b, c, f): + """numpy's uint8 results, which wrap modulo 256 like the u8 ops.""" + with np.errstate(over="ignore", invalid="ignore"): + saturated = np.nan_to_num(np.clip(np.rint(f), 0.0, 255.0), nan=0.0) + return { + "DIV": a // b, + "REM": a % c, + "LT": (a < b).astype(np.uint8), + "LE": (a <= b).astype(np.uint8), + "GT": (a > b).astype(np.uint8), + "GE": (a >= b).astype(np.uint8), + "MAX": np.maximum(a, b), + "MIN": np.minimum(a, b), + "SHR": a >> (b & np.uint8(7)), + "ADD": a + b, + "SUB": a - b, + "MUL": a * b, + "TO_F32": a.astype(np.float32), + "WIDE": a.astype(np.int32), + "SAT": saturated.astype(np.uint8), + } + + +def test_u8_ops_on_gpu_match_numpy_uint8(): + a, b, c, f = _inputs() + host_dtype = {U8: np.uint8, F32: np.float32, I32: np.int32} + outputs = { + name: np.zeros(_LANES, dtype=host_dtype[t]) for name, t in _OUTPUTS.items() + } + run_on_device( + _u8_arith_kernel(), + {"A": a, "B": b, "C": c, "F": f, **outputs}, + threads=_LANES, + arches={"gfx90a"}, + ) + for name, expected in _expected(a, b, c, f).items(): + np.testing.assert_array_equal( + outputs[name].view(np.uint8), expected.view(np.uint8), err_msg=name + ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_u8_memory.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_u8_memory.py new file mode 100644 index 000000000000..757d1122e567 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_u8_memory.py @@ -0,0 +1,141 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""U8 in memory, and the element list the LDS vector builders admit. + +A u8 element takes one byte, like i8, in both engines' LLVM and HIP lowerings. +``smem_load_vN`` and ``smem_store_vN`` admit one shared element list, which +leaves u8 out, so u8 vectors are refused by the LDS pair as by the global pair. +The C++ builders are reached through the ``rocke_ir_types`` test binary +(``ROCKE_IR_TYPES_TEST``), because the serialized-IR seam lowers kernels the +Python builder built. +""" + +import os +import re +import shutil +import subprocess + +import pytest +from rocke.core.ir import ( + I8, + U8, + IRBuilder, + KernelDef, + PtrType, + Type, + dtype_to_ir_type, +) +from rocke.core.lower_hip import lower_kernel_to_hip +from rocke.tests.engines import lower_hip_with_both_engines, lower_with_both_engines +from rocke.tests.ir_validity import assert_target_accepts + +_ADMITTED = "f16 / bf16 / i16 / f32 / i32 / tf32 / f64 / fp8e4m3 / bf8e5m2 / i8" + + +def _byte_copy_kernel(dtype: Type) -> KernelDef: + """Copy one ``dtype`` element per thread from SRC to DST.""" + b = IRBuilder("byte_copy") + src = b.param("SRC", PtrType(dtype, "global")) + dst = b.param("DST", PtrType(dtype, "global")) + tid = b.thread_id_x() + b.global_store(dst, tid, b.global_load(src, tid, dtype)) + return b.kernel + + +def _lds_byte_store_kernel(dtype: Type) -> KernelDef: + """Store one ``dtype`` element per thread from SRC into a 64-element LDS tile.""" + b = IRBuilder("lds_byte_store") + src = b.param("SRC", PtrType(dtype, "global")) + tid = b.thread_id_x() + tile = b.smem_alloc(dtype, [64], name_hint="tile") + b.smem_store_vN(tile, [tid], b.global_load(src, tid, dtype), 1) + return b.kernel + + +def _python_smem_vector_access(direction: str, elem: str, n: int) -> str: + """Build one LDS vector access with the Python builder; ``ok`` or ``error: ``. + + The twin of ``rocke_ir_types --smem-vector``: an LDS tile of ``elem``, then + ``smem_load_vN``, or ``smem_store_vN`` of ``vector_splat(, n)``. + A refusal must leave the kernel without a new op. + """ + dtype = dtype_to_ir_type(elem) + b = IRBuilder("smem_vector_access") + tile = b.smem_alloc(dtype, [64], name_hint="tile") + zero = b.const_i32(0) + value = None + if direction == "store": + src = b.param("SRC", PtrType(dtype, "global")) + value = b.vector_splat(b.global_load(src, zero, dtype), n) + ops_before = len(b.kernel.body.ops) + try: + if direction == "store": + b.smem_store_vN(tile, [zero], value, n) + else: + b.smem_load_vN(tile, zero, dtype=dtype, n=n) + except ValueError as error: + assert len(b.kernel.body.ops) == ops_before + return f"error: {error}" + return "ok" + + +def test_u8_lds_store_is_laid_out_like_i8(): + # A u8 LDS element is one byte, so the tile and the store match the i8 kernel. + u8_ir = lower_with_both_engines(_lds_byte_store_kernel(U8), "gfx90a") + assert u8_ir == lower_with_both_engines(_lds_byte_store_kernel(I8), "gfx90a") + assert "[64 x i8]" in u8_ir + assert re.search(r"store i8 \S+, ptr addrspace\(3\) \S+, align 1\n", u8_ir) + assert_target_accepts(u8_ir, "gfx90a") + hip = lower_hip_with_both_engines(_lds_byte_store_kernel(U8), "gfx90a") + assert "__shared__ uint8_t" in hip + + +@pytest.mark.parametrize("direction", ["load", "store"]) +def test_lds_vector_access_refuses_u8(direction): + assert _python_smem_vector_access(direction, "u8", 4) == ( + f"error: smem_{direction}_vN supports {_ADMITTED}, got u8" + ) + assert _python_smem_vector_access(direction, "i8", 16) == "ok" + + +@pytest.mark.parametrize("direction", ["load", "store"]) +@pytest.mark.parametrize("elem", ["u8", "i8", "i16"]) +@pytest.mark.parametrize("n", [4, 16]) +def test_native_lds_vector_access_matches_the_python_builder(direction, elem, n): + executable = os.environ.get("ROCKE_IR_TYPES_TEST") + if not executable: + pytest.skip("set ROCKE_IR_TYPES_TEST to the built rocke_ir_types test") + native = subprocess.run( + [executable, "--smem-vector", direction, elem, str(n)], + check=True, + capture_output=True, + text=True, + ).stdout.rstrip("\n") + python = _python_smem_vector_access(direction, elem, n) + if elem == "u8": + # The element refusal is one text in both builders. + assert native == python + else: + # The width refusals word the allowed widths differently in the two + # builders; compare the outcome only. + assert native.split(":")[0] == python.split(":")[0] + + +@pytest.mark.parametrize("kernel_builder", [_byte_copy_kernel, _lds_byte_store_kernel]) +def test_u8_hip_source_compiles(kernel_builder, tmp_path): + if not shutil.which("hipcc"): + pytest.skip("hipcc not in PATH") + path = tmp_path / "kernel.hip" + path.write_text(lower_kernel_to_hip(kernel_builder(U8), arch="gfx90a")) + subprocess.run( + [ + "hipcc", + "--offload-arch=gfx90a", + "--cuda-device-only", + "-fsyntax-only", + str(path), + ], + check=True, + capture_output=True, + text=True, + ) 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 649d25168ffd..de60c458e46f 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 @@ -133,6 +133,7 @@ def ctest_ready(build_root: Path, config: str) -> bool: # instead of skipping: (env var, CTest name, label). NATIVE_PARITY_TESTS = ( ("ROCKE_STORAGE_TEST", "rocke_storage", "native storage parity"), + ("ROCKE_IR_TYPES_TEST", "rocke_ir_types", "native IR type rule parity"), ( "ROCKE_NONTEMPORAL_HIP_TEST", "rocke_nontemporal_hip", 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 581926d726d9..794a194fd161 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 @@ -42,6 +42,8 @@ def test_runner_passes_registered_fixture_to_both_backends( executable = tmp_path / "Debug" / filename executable.parent.mkdir() executable.touch() + ir_types = tmp_path / "Debug" / "rocke_ir_types" + ir_types.touch() nontemporal = tmp_path / "Debug" / "rocke_nontemporal_hip" nontemporal.touch() selected = tmp_path / "override" if override else executable @@ -55,6 +57,7 @@ def run(command, **kwargs): listing = { "tests": [ {"name": "rocke_storage", "command": [str(executable)]}, + {"name": "rocke_ir_types", "command": [str(ir_types)]}, {"name": "rocke_nontemporal_hip", "command": [str(nontemporal)]}, ] } @@ -82,6 +85,7 @@ def run(command, **kwargs): ] assert len(children) == 3 # default pytest, extension import probe, both pytest assert all(env["ROCKE_STORAGE_TEST"] == str(selected) for _, env in children) + assert all(env["ROCKE_IR_TYPES_TEST"] == str(ir_types) for _, env in children) assert all( env["ROCKE_NONTEMPORAL_HIP_TEST"] == str(nontemporal) for _, env in children ) @@ -268,7 +272,7 @@ def test_fresh_build_prepares_entire_ctest_suite( "enable_testing()\n" 'add_custom_target(rocke_core COMMAND "${CMAKE_COMMAND}" -E touch core.ready)\n' 'get_filename_component(tool_suffix "${CMAKE_COMMAND}" LAST_EXT)\n' - "foreach(name rocke_storage rocke_nontemporal_hip rocke_dtypes)\n" + "foreach(name rocke_storage rocke_ir_types rocke_nontemporal_hip rocke_dtypes)\n" ' set(output "${CMAKE_CURRENT_BINARY_DIR}/fixtures/${name}${tool_suffix}")\n' ' add_custom_command(OUTPUT "${output}"\n' ' COMMAND "${CMAKE_COMMAND}" -E make_directory "${CMAKE_CURRENT_BINARY_DIR}/fixtures"\n' @@ -293,6 +297,7 @@ def test_fresh_build_prepares_entire_ctest_suite( monkeypatch.setenv("ROCKE_STORAGE_TEST", str(executable)) env = runner.native_pytest_env(build, "Release") assert Path(env["ROCKE_STORAGE_TEST"]).is_file() + assert Path(env["ROCKE_IR_TYPES_TEST"]).is_file() assert Path(env["ROCKE_NONTEMPORAL_HIP_TEST"]).is_file() if override: assert env["ROCKE_STORAGE_TEST"] == str(executable) From 11e977ab2d989c3fcf71dfad7325243efb81a7c1 Mon Sep 17 00:00:00 2001 From: Chris Millette Date: Fri, 9 Oct 2026 15:55:33 -0400 Subject: [PATCH 06/11] fix(rocke): refuse a bitcast that changes the bit width A bitcast keeps every bit of its operand, so a source and target of different widths is always a mistake; until now it was caught only by clang ("invalid cast opcode"). type_bits(t) in core/ir.py is the one place that answers how many bits a value of an IR type occupies: a scalar's width comes from its format, a vector's is its count times its element width, and a pointer's comes from its address space (32 bits for LDS, 64 otherwise). An LDS allocation is memory, not a value, and raises TypeError. IRBuilder.bitcast and IRBuilder.vec_bitcast raise ValueError naming both types and widths when the widths differ, and TypeError for a pointer operand or target, which LLVM does not bitcast. Equal-width casts lower exactly as before, so no existing kernel's IR changes. The C++ builders refuse the same casts with the same text: rocke_b_bitcast and rocke_b_vec_bitcast call rocke_i_check_bitcast_width before building the op. rocke_i_type_bits is the C++ twin of type_bits, and rocke_require_dtype_info (rocke/dtypes_internal.h) is the raising twin of rocke_dtype_info, so "unknown dtype ''" has one home in each engine. The rocke_ir_types test binary gains a --bitcast mode; the tests require its answers to equal the Python builder's, and lower the equal-width casts with both engines. Co-Authored-By: Claude --- .../rocke/platform/cpp/core/dtypes.cpp | 17 +- .../rocke/platform/cpp/core/ir/core_types.cpp | 28 +++ .../rocke/platform/cpp/core/ir/ir_arith.cpp | 31 +++ .../rocke/platform/cpp/core/ir/ir_mem.cpp | 2 + .../cpp/include/rocke/dtypes_internal.h | 20 ++ .../rocke/platform/cpp/include/rocke/ir.h | 4 + .../platform/cpp/include/rocke/ir_internal.h | 19 ++ .../rocke/platform/python/rocke/core/ir.py | 47 ++++- .../platform/tests/core/test_bitcast_width.py | 182 ++++++++++++++++ .../platform/tests/core/test_ir_types.cpp | 197 ++++++++++++++++++ 10 files changed, 543 insertions(+), 4 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/dtypes_internal.h create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_bitcast_width.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 8860bec148e1..9242b7d2f865 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 @@ -1,13 +1,18 @@ // Copyright (c) Advanced Micro Devices, Inc., or its affiliates. // SPDX-License-Identifier: MIT /* Target-independent dtype names. Mirrors rocke.core.dtypes: - * normalize_dtype -> rocke_normalize_dtype. - * No architecture, IR builder, or lowering dependencies. + * normalize_dtype -> rocke_normalize_dtype, dtype_info -> rocke_dtype_info + * (NULL on a miss) and rocke_require_dtype_info (raises). + * No architecture, IR builder, or lowering dependencies; the engine exception + * types (rocke/error.hpp) carry the raising lookup's refusal. */ #include "rocke/dtypes.h" #include +#include "rocke/dtypes_internal.h" +#include "rocke/error.hpp" + /* ========================================================================= * target-independent dtype normalization * ========================================================================= @@ -188,3 +193,11 @@ const rocke_dtype_info_t* rocke_dtype_info(const char* name) return &info; return NULL; } + +const rocke_dtype_info_t* rocke_require_dtype_info(const char* name) +{ + const rocke_dtype_info_t* info = rocke_dtype_info(name); + if(!info) + throw ckc::ValueError(ckc::format_error("unknown dtype '%s'", name ? name : "(null)")); + return info; +} 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 115ee05c25f5..f2515c9999e4 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 @@ -22,6 +22,8 @@ #include "rocke/arena.h" #include "rocke/dtypes.h" +#include "rocke/dtypes_internal.h" +#include "rocke/error.hpp" #include "rocke/ir.h" #include "rocke/ir_internal.h" @@ -251,6 +253,32 @@ const rocke_type_t* rocke_smem_type( return t; } +/* Pointer widths by address space, as the AMDGPU data layout states them (p3 is + * 32 bits; global, constant and flat pointers are 64). Mirrors _POINTER_BITS. */ +static int rocke_i_pointer_bits(const char* space) +{ + return space && strcmp(space, "lds") == 0 ? 32 : 64; +} + +/* Mirrors core/ir.py::type_bits. A scalar's width comes from its format + * (k_dtype_info in dtypes.cpp, the one home of format widths). */ +int rocke_i_type_bits(const rocke_type_t* t) +{ + switch(t->kind) + { + case ROCKE_TYPE_VECTOR: + return t->count * rocke_i_type_bits(t->elem); + case ROCKE_TYPE_PTR: + return rocke_i_pointer_bits(t->space); + case ROCKE_TYPE_SMEM: + throw ckc::TypeError( + ckc::format_error("%s is an LDS allocation, not a value with a bit width", t->name)); + case ROCKE_TYPE_SCALAR: + break; + } + return rocke_require_dtype_info(t->name)->encoded_bits; +} + /* Structural equality == Python frozen-dataclass __eq__: canonical name encodes * kind + all components, so a name compare is sufficient. NULL-safe. */ bool rocke_type_eq(const rocke_type_t* a, const rocke_type_t* b) 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 cad405edaeea..e1d1e5a25268 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 @@ -14,6 +14,7 @@ */ #include +#include "rocke/error.hpp" #include "rocke/ir_internal.h" /* ------------------------------------------------------------ arith constants */ @@ -573,6 +574,34 @@ rocke_value_t* rocke_b_trunc(rocke_ir_builder_t* b, rocke_value_t* v, const rock return rocke_i_op1(b, ROCKE_OP_ARITH_TRUNC, operands, 1, target, NULL, "tr"); } +/* A bitcast keeps every bit of its operand, so a source and target of different + * widths is always a mistake; refusing it here names both types instead of + * leaving clang to reject the lowered IR ("invalid cast opcode"). */ +void rocke_i_check_bitcast_width(const char* verb, + const rocke_type_t* source, + const rocke_type_t* target, + const char* hint) +{ + if(source->kind == ROCKE_TYPE_PTR || target->kind == ROCKE_TYPE_PTR) + throw ckc::TypeError( + ckc::format_error("%s of %s to %s: a bitcast does not convert pointers", + verb, + source->name, + target->name)); + const int source_bits = rocke_i_type_bits(source); + const int target_bits = rocke_i_type_bits(target); + if(source_bits != target_bits) + throw ckc::ValueError( + ckc::format_error("%s of %s (%d bits) to %s (%d bits) changes the bit " + "width -- a bitcast keeps every bit; %s", + verb, + source->name, + source_bits, + target->name, + target_bits, + hint)); +} + rocke_value_t* rocke_b_bitcast(rocke_ir_builder_t* b, rocke_value_t* v, const rocke_type_t* target) { rocke_value_t* operands[1]; @@ -583,6 +612,8 @@ rocke_value_t* rocke_b_bitcast(rocke_ir_builder_t* b, rocke_value_t* v, const ro return (rocke_value_t*)rocke_i_set_err(b, ROCKE_ERR_VALUE, "bitcast NULL value"); if(!target) return (rocke_value_t*)rocke_i_set_err(b, ROCKE_ERR_VALUE, "bitcast NULL target"); + rocke_i_check_bitcast_width( + "bitcast", v->type, target, "convert the value to change its width"); operands[0] = v; at = rocke_i_attrs(b); rocke_attr_set_str(b, &at, "target", target->name); 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 346c1e1b1172..21629130da5f 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 @@ -1284,6 +1284,8 @@ rocke_value_t* if(!v || !target) return (rocke_value_t*)rocke_i_set_err( b, ROCKE_ERR_VALUE, "vec_bitcast: null operand/target"); + rocke_i_check_bitcast_width( + "vec_bitcast", v->type, target, "slice or concatenate the vector first"); ops[0] = v; a = rocke_i_attrs(b); rocke_attr_set_str(b, &a, "target", target->name); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/dtypes_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/dtypes_internal.h new file mode 100644 index 000000000000..df5150dfd0bb --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/dtypes_internal.h @@ -0,0 +1,20 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* Engine-internal dtype lookups that raise. NOT a public API. + * + * rocke/dtypes.h is the C contract: rocke_dtype_info returns NULL for an unknown + * name. Engine code that mirrors a Python `raise` calls the twin below instead, + * so the refusal text has one home in the C++ engine as it has in Python + * (core/dtypes.py::dtype_info). + */ +#ifndef ROCKE_DTYPES_INTERNAL_H +#define ROCKE_DTYPES_INTERNAL_H + +#include "rocke/dtypes.h" + +/* The encoding of `name` (any accepted spelling), as rocke_dtype_info returns + * it. Throws ckc::ValueError "unknown dtype ''", naming the spelling the + * caller passed, when no encoding has that name. */ +const rocke_dtype_info_t* rocke_require_dtype_info(const char* name); + +#endif /* ROCKE_DTYPES_INTERNAL_H */ 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 dd72f3a15f45..47cba1adc986 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 @@ -855,6 +855,8 @@ rocke_value_t* rocke_b_tanh(rocke_ir_builder_t* b, rocke_value_t* a); rocke_value_t* rocke_b_zext(rocke_ir_builder_t* b, rocke_value_t* v, const rocke_type_t* target); rocke_value_t* rocke_b_sext(rocke_ir_builder_t* b, rocke_value_t* v, const rocke_type_t* target); rocke_value_t* rocke_b_trunc(rocke_ir_builder_t* b, rocke_value_t* v, const rocke_type_t* target); +/* Bitcast to a type of the same bit width; another width (ValueError) or a + * pointer (TypeError) is refused before any op is built. */ rocke_value_t* rocke_b_bitcast(rocke_ir_builder_t* b, rocke_value_t* v, const rocke_type_t* target); rocke_value_t* rocke_b_select(rocke_ir_builder_t* b, rocke_value_t* cond, @@ -1372,6 +1374,8 @@ rocke_value_t* rocke_b_ds_read_tr_b8(rocke_ir_builder_t* b, const rocke_type_t* dtype /* NULL=>fp8e4m3 */); /* ----- vector bitcast / packed f32->f16 conversion ----- */ +/* Reinterpret a vector as another vector type of the same bit width; another + * width (ValueError) or a pointer (TypeError) is refused before any op is built. */ rocke_value_t* rocke_b_vec_bitcast(rocke_ir_builder_t* b, rocke_value_t* v, const rocke_type_t* target); rocke_value_t* rocke_b_vec_trunc_f32_to_f16(rocke_ir_builder_t* b, rocke_value_t* v); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir_internal.h index eed8c7d3f96a..283abc505f36 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir_internal.h @@ -159,6 +159,25 @@ void rocke_i_refuse_unsigned_operands(const rocke_op_t* op, const char* instead) /* Lane count of a vector type, or 1 for a scalar. */ int rocke_i_count_of(const rocke_type_t* t); +/* The number of bits a value of type `t` occupies (Python core/ir.py::type_bits): + * a scalar's format width, a vector's element width times its count, a + * pointer's address-space width. Throws ckc::TypeError for an LDS allocation + * (memory, not a value) and ckc::ValueError "unknown dtype ''" for a + * scalar with no format. */ +int rocke_i_type_bits(const rocke_type_t* t); + +/* Refuse a bitcast of a `source` value to `target` unless both are values of + * the same bit width (Python core/ir.py::_check_bitcast_width). `verb` names + * the builder ("bitcast", "vec_bitcast") and `hint` says what to do instead of + * changing the width. Throws ckc::TypeError when either type is a pointer + * (LLVM does not bitcast pointers) or an LDS allocation, and ckc::ValueError + * naming both types and widths when the widths differ. Builders call it before + * creating the op, so a refused cast leaves the kernel unchanged. */ +void rocke_i_check_bitcast_width(const char* verb, + const rocke_type_t* source, + const rocke_type_t* target, + const char* hint); + /* ------------------------------------------------------------- attr helpers */ /* Build a small attr map IN the arena and return it by value (the map's 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 b43077944f60..c9a29bcd2540 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 @@ -575,6 +575,43 @@ def __init__( object.__setattr__(self, "exclusive", exclusive) +# Pointer widths by address space, as the AMDGPU data layout states them +# (p3 is 32 bits; global, constant and flat pointers are 64). +_POINTER_BITS = {"lds": 32} +_DEFAULT_POINTER_BITS = 64 + + +def type_bits(t: Type) -> int: + """The number of bits a value of IR type ``t`` occupies. + + A scalar's width comes from its format (``dtype_info``); a vector's is its + element width times its count; a pointer's comes from its address space. + An LDS allocation is memory, not a value, and raises ``TypeError``. + """ + if isinstance(t, VectorType): + return t.count * type_bits(t.elem) + if isinstance(t, PtrType): + return _POINTER_BITS.get(t.space, _DEFAULT_POINTER_BITS) + if isinstance(t, SmemType): + raise TypeError(f"{t.name} is an LDS allocation, not a value with a bit width") + return dtype_info(t.name).encoded_bits + + +def _check_bitcast_width(verb: str, source: Type, target: Type, hint: str) -> None: + if isinstance(source, PtrType) or isinstance(target, PtrType): + raise TypeError( + f"{verb} of {source.name} to {target.name}: a bitcast does not convert " + "pointers" + ) + source_bits, target_bits = type_bits(source), type_bits(target) + if source_bits != target_bits: + raise ValueError( + f"{verb} of {source.name} ({source_bits} bits) to {target.name} " + f"({target_bits} bits) changes the bit width -- a bitcast keeps every " + f"bit; {hint}" + ) + + # ----------------------------- Values / Ops ------------------------------ @@ -3055,8 +3092,11 @@ def vec_bitcast(self, v: Value, target: Type) -> Value: `buffer_load_dwordxN`) and `<2N x f16>` (the form the MFMA wants), or `` (raw 32-bit dwords) and ``. The LLVM lowering emits an `addrspacecast`-less `bitcast`; the gfx - backend folds it away. + backend folds it away. A target of a different bit width raises. """ + _check_bitcast_width( + "vec_bitcast", v.type, target, "slice or concatenate the vector first" + ) return self._op( "vector.bitcast", [v], @@ -3861,7 +3901,10 @@ def lane_id(self) -> Value: return self._op("tile.lane_id", [], [I32], result_name_hint="lane").result def bitcast(self, v: Value, target: Type) -> Value: - """Bitcast a scalar to another type of the same size.""" + """Bitcast a scalar to a type of the same bit width; another width raises.""" + _check_bitcast_width( + "bitcast", v.type, target, "convert the value to change its width" + ) return self._op( "arith.bitcast", [v], diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_bitcast_width.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_bitcast_width.py new file mode 100644 index 000000000000..64aa4f69bec8 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_bitcast_width.py @@ -0,0 +1,182 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""A bitcast keeps every bit, so ``bitcast`` and ``vec_bitcast`` refuse a width change. + +``type_bits`` is the one place that answers how many bits a value of an IR type +occupies. Equal-width casts lower to one LLVM ``bitcast``; a cast that changes the +width fails when the kernel is built instead of inside clang. Both engines lower +the casts identically, and the C++ builders refuse the same casts with the same +text: they are reached through the ``rocke_ir_types`` test binary +(``ROCKE_IR_TYPES_TEST``), because the serialized-IR seam lowers kernels the +Python builder built. +""" + +import os +import re +import subprocess + +import pytest +from rocke.core import ir +from rocke.core.dtypes import dtype_info +from rocke.core.ir import ( + F16, + F32, + F64, + I32, + I64, + IRBuilder, + PtrType, + SmemType, + VectorType, + type_bits, +) +from rocke.tests.engines import lower_with_both_engines +from rocke.tests.ir_validity import assert_target_accepts + + +def _lower(b: IRBuilder) -> str: + return lower_with_both_engines(b.kernel, "gfx90a") + + +def _scalar_types() -> list: + """Every scalar IR type constant the module defines, so later types are covered.""" + return [t for t in vars(ir).values() if type(t) is ir.Type] + + +@pytest.mark.parametrize("t", _scalar_types(), ids=lambda t: t.name) +def test_a_scalar_occupies_its_format_width(t): + assert type_bits(t) == dtype_info(t.name).encoded_bits + + +def test_a_vector_occupies_count_times_its_element_width(): + assert type_bits(VectorType(F16, 8)) == 128 + assert type_bits(VectorType(I32, 4)) == 128 + assert type_bits(VectorType(ir.FP4E2M1, 32)) == 128 + assert type_bits(VectorType(ir.FP6E2M3, 32)) == 192 + + +@pytest.mark.parametrize( + ("space", "bits"), [("global", 64), ("constant", 64), ("lds", 32)] +) +def test_a_pointer_occupies_its_address_space_width(space, bits): + assert type_bits(PtrType(F32, space)) == bits + + +def test_an_lds_allocation_is_not_a_value(): + t = SmemType(F32, (16, 16)) + message = f"{t.name} is an LDS allocation, not a value with a bit width" + with pytest.raises(TypeError, match=f"^{re.escape(message)}$"): + type_bits(t) + + +def test_a_bitcast_refuses_pointers(): + b = IRBuilder("pointer_bitcast") + p = b.param("P", PtrType(F32, "global")) + message = "bitcast of ptr to i64: a bitcast does not convert pointers" + with pytest.raises(TypeError, match=f"^{re.escape(message)}$"): + b.bitcast(p, I64) + + +def test_equal_width_casts_lower_to_one_bitcast_each(): + b = IRBuilder("equal_width_bitcasts") + p = b.param("P", PtrType(I32, "global")) + q = b.param("Q", PtrType(F32, "global")) + tid = b.thread_id_x() + words = b.global_load_vN(p, tid, I32, 4) + halves = b.vec_bitcast(words, VectorType(F16, 8)) + b.global_store_vN(p, tid, b.vec_bitcast(halves, VectorType(I32, 4)), 4) + single = b.vec_pack([b.global_load(q, tid, F32, align=4)], F32) + b.global_store(q, tid, b.bitcast(single, F32), align=4) + text = _lower(b) + assert re.search(r"bitcast <4 x i32> %\S+ to <8 x half>", text) + assert re.search(r"bitcast <8 x half> %\S+ to <4 x i32>", text) + assert re.search(r"bitcast <1 x float> %\S+ to float", text) + assert_target_accepts(text, "gfx90a") + + +def test_a_vector_bitcast_that_changes_the_width_raises(): + b = IRBuilder("widening_vec_bitcast") + halves = b.zero_vec(F16, 8) + message = ( + "vec_bitcast of vec (128 bits) to vec (256 bits) changes the bit " + "width -- a bitcast keeps every bit; slice or concatenate the vector first" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + b.vec_bitcast(halves, VectorType(I32, 8)) + + +def test_a_scalar_bitcast_that_changes_the_width_raises(): + b = IRBuilder("widening_bitcast") + message = ( + "bitcast of f32 (32 bits) to i64 (64 bits) changes the bit width -- a bitcast " + "keeps every bit; convert the value to change its width" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + b.bitcast(b.const_f32(1.0), I64) + + +def test_bitcast_then_slice_then_bitcast_back_links(): + """Reading and replacing one f16 inside raw 32-bit words is done with core ops.""" + b = IRBuilder("bitcast_and_slice") + p = b.param("P", PtrType(I32, "global")) + tid = b.thread_id_x() + halves = b.vec_bitcast(b.global_load_vN(p, tid, I32, 4), VectorType(F16, 8)) + third = b.vec_extract(halves, 3) + replaced = b.vec_insert(halves, third, 0) + b.global_store_vN(p, tid, b.vec_bitcast(replaced, VectorType(I32, 4)), 4) + text = _lower(b) + assert "extractelement <8 x half>" in text + assert "insertelement <8 x half>" in text + assert_target_accepts(text, "gfx90a") + + +# One bitcast per row: (builder, source type, target type). The rows match +# k_bitcast_cases in tests/core/test_ir_types.cpp. An SmemType source is an LDS +# tile from smem_alloc; any other source is a kernel parameter. +_BITCAST_CASES = { + "f32_as_i32": ("bitcast", F32, I32), + "f64_as_i32x2": ("bitcast", F64, VectorType(I32, 2)), + "i32x4_as_f16x8": ("vec_bitcast", VectorType(I32, 4), VectorType(F16, 8)), + "f32_as_i64": ("bitcast", F32, I64), + "f16x8_as_i32x8": ("vec_bitcast", VectorType(F16, 8), VectorType(I32, 8)), + "pointer_as_i64": ("bitcast", PtrType(F32, "global"), I64), + "i64_as_pointer": ("bitcast", I64, PtrType(F32, "global")), + "lds_tile_as_i32": ("bitcast", SmemType(F32, (16, 16)), I32), +} + + +def _python_bitcast(case: str) -> str: + """Build one bitcast with the Python builder; ``ok`` or ``error: : ``. + + The twin of ``rocke_ir_types --bitcast``. A refusal must leave the kernel + without a new op. + """ + verb, source_type, target = _BITCAST_CASES[case] + b = IRBuilder("bitcast_case") + if isinstance(source_type, SmemType): + source = b.smem_alloc( + source_type.elem, list(source_type.shape), name_hint="tile" + ) + else: + source = b.param("X", source_type) + ops_before = len(b.kernel.body.ops) + try: + getattr(b, verb)(source, target) + except (TypeError, ValueError) as error: + assert len(b.kernel.body.ops) == ops_before + return f"error: {type(error).__name__}: {error}" + return "ok" + + +@pytest.mark.parametrize("case", sorted(_BITCAST_CASES)) +def test_native_bitcast_matches_the_python_builder(case): + executable = os.environ.get("ROCKE_IR_TYPES_TEST") + if not executable: + pytest.skip("set ROCKE_IR_TYPES_TEST to the built rocke_ir_types test") + native = subprocess.run( + [executable, "--bitcast", case], + check=True, + capture_output=True, + text=True, + ).stdout.rstrip("\n") + assert native == _python_bitcast(case) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp index 1d998c69a538..7e067ffc75a9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp @@ -24,6 +24,12 @@ * --smem-vector * build one LDS vector access of n elem values; print "ok", or * "error: " when the builder refuses it. + * + * --bitcast + * build the bitcast named (a row of k_bitcast_cases, mirrored in + * tests/core/test_bitcast_width.py); print "ok", or + * "error: : " when the builder refuses + * it, since a refusal is a TypeError or a ValueError. */ #include "rocke/ir.h" @@ -32,7 +38,9 @@ #include #include +#include "rocke/dtypes_internal.h" #include "rocke/error.hpp" +#include "rocke/ir_internal.h" #define CHECK(x) \ do \ @@ -416,6 +424,191 @@ static int test_smem_vector_elements() return 0; } +/* The Python exception class a refusal raises, for the --bitcast report. */ +static const char* python_exception_class(rocke_status_t code) +{ + switch(code) + { + case ROCKE_ERR_TYPE: + return "TypeError"; + case ROCKE_ERR_VALUE: + return "ValueError"; + default: + return "unexpected status"; + } +} + +/* One bitcast: the builder (`vec_bitcast` or `bitcast`), the source value and + * the target type. The rows match _BITCAST_CASES in test_bitcast_width.py. */ +struct bitcast_case +{ + const char* name; + bool vector_builder; + rocke_value_t* (*source)(rocke_ir_builder_t* b); + const rocke_type_t* (*target)(rocke_ir_builder_t* b); +}; + +static rocke_value_t* param_of(rocke_ir_builder_t* b, const rocke_type_t* type) +{ + return rocke_b_param(b, "X", type, NULL); +} + +static const bitcast_case k_bitcast_cases[] = { + {"f32_as_i32", + false, + [](rocke_ir_builder_t* b) { return param_of(b, rocke_f32()); }, + [](rocke_ir_builder_t*) { return rocke_i32(); }}, + {"f64_as_i32x2", + false, + [](rocke_ir_builder_t* b) { return param_of(b, rocke_f64()); }, + [](rocke_ir_builder_t* b) { return rocke_vector_type(b, rocke_i32(), 2); }}, + {"i32x4_as_f16x8", + true, + [](rocke_ir_builder_t* b) { return param_of(b, rocke_vector_type(b, rocke_i32(), 4)); }, + [](rocke_ir_builder_t* b) { return rocke_vector_type(b, rocke_f16(), 8); }}, + {"f32_as_i64", + false, + [](rocke_ir_builder_t* b) { return param_of(b, rocke_f32()); }, + [](rocke_ir_builder_t*) { return rocke_i64(); }}, + {"f16x8_as_i32x8", + true, + [](rocke_ir_builder_t* b) { return param_of(b, rocke_vector_type(b, rocke_f16(), 8)); }, + [](rocke_ir_builder_t* b) { return rocke_vector_type(b, rocke_i32(), 8); }}, + {"pointer_as_i64", + false, + [](rocke_ir_builder_t* b) { return param_of(b, rocke_ptr_type(b, rocke_f32(), "global")); }, + [](rocke_ir_builder_t*) { return rocke_i64(); }}, + {"i64_as_pointer", + false, + [](rocke_ir_builder_t* b) { return param_of(b, rocke_i64()); }, + [](rocke_ir_builder_t* b) { return rocke_ptr_type(b, rocke_f32(), "global"); }}, + {"lds_tile_as_i32", + false, + [](rocke_ir_builder_t* b) { + const int shape[] = {16, 16}; + return rocke_b_smem_alloc(b, rocke_f32(), shape, 2, "tile"); + }, + [](rocke_ir_builder_t*) { return rocke_i32(); }}, +}; + +static const bitcast_case* find_bitcast_case(const char* name) +{ + for(const bitcast_case& c : k_bitcast_cases) + if(strcmp(c.name, name) == 0) + return &c; + return NULL; +} + +/* Build the source and target of `c`, then the bitcast itself; only the cast + * is measured. */ +static build_outcome build_bitcast(const bitcast_case& c, rocke_status_t* code) +{ + build_outcome outcome = {false, "", 0}; + *code = ROCKE_OK; + rocke_ir_builder_t b; + if(rocke_ir_builder_init(&b, "bitcast_case") != ROCKE_OK) + { + outcome.error = "builder init failed"; + return outcome; + } + int ops_before = -1; /* set once the source and target are built */ + try + { + rocke_value_t* source = c.source(&b); + const rocke_type_t* target = c.target(&b); + ops_before = b.kernel->body->num_ops; + if(c.vector_builder) + rocke_b_vec_bitcast(&b, source, target); + else + rocke_b_bitcast(&b, source, target); + outcome.accepted = true; + } + catch(const ckc::Error& error) + { + outcome.error = error.what(); + *code = error.code(); + } + outcome.ops_added = ops_before < 0 ? -1 : b.kernel->body->num_ops - ops_before; + rocke_ir_builder_free(&b); + return outcome; +} + +static int print_bitcast(const char* name) +{ + const bitcast_case* c = find_bitcast_case(name); + CHECK(c != NULL); + rocke_status_t code; + const build_outcome outcome = build_bitcast(*c, &code); + if(outcome.accepted) + printf("ok\n"); + else + printf("error: %s: %s\n", python_exception_class(code), outcome.error.c_str()); + return 0; +} + +/* Equal-width casts build one op; a width change, a pointer or an LDS + * allocation is refused with core/ir.py's text before any op is built. */ +static int test_bitcast_width() +{ + for(const char* name : {"f32_as_i32", "f64_as_i32x2", "i32x4_as_f16x8"}) + { + rocke_status_t code; + const build_outcome built = build_bitcast(*find_bitcast_case(name), &code); + CHECK(built.accepted); + CHECK(built.ops_added == 1); + } + const struct + { + const char* name; + rocke_status_t code; + const char* error; + } refused[] = { + {"f32_as_i64", + ROCKE_ERR_VALUE, + "bitcast of f32 (32 bits) to i64 (64 bits) changes the bit width -- a bitcast " + "keeps every bit; convert the value to change its width"}, + {"f16x8_as_i32x8", + ROCKE_ERR_VALUE, + "vec_bitcast of vec (128 bits) to vec (256 bits) changes the bit " + "width -- a bitcast keeps every bit; slice or concatenate the vector first"}, + {"pointer_as_i64", + ROCKE_ERR_TYPE, + "bitcast of ptr to i64: a bitcast does not convert pointers"}, + {"lds_tile_as_i32", + ROCKE_ERR_TYPE, + "smem is an LDS allocation, not a value with a bit width"}, + }; + for(const auto& r : refused) + { + rocke_status_t code; + const build_outcome outcome = build_bitcast(*find_bitcast_case(r.name), &code); + CHECK(!outcome.accepted); + CHECK(code == r.code); + CHECK(outcome.error == r.error); + CHECK(outcome.ops_added == 0); + } + return 0; +} + +/* A scalar's width is its format's width; an unknown format is refused with + * core/dtypes.py's text, naming the spelling that was passed. */ +static int test_scalar_widths() +{ + CHECK(rocke_i_type_bits(rocke_f64()) == 64); + CHECK(rocke_i_type_bits(rocke_fp6e2m3()) == 6); + CHECK(rocke_require_dtype_info("half") == rocke_dtype_info("f16")); + try + { + rocke_require_dtype_info("custom_format"); + CHECK(false); + } + catch(const ckc::ValueError& error) + { + CHECK(strcmp(error.what(), "unknown dtype 'custom_format'") == 0); + } + return 0; +} + int main(int argc, char** argv) { if(argc == 3 && strcmp(argv[1], "--unsigned") == 0) @@ -424,10 +617,14 @@ int main(int argc, char** argv) return print_packed_only_value(argv[2], argv[3]); if(argc == 5 && strcmp(argv[1], "--smem-vector") == 0) return print_smem_vector_access(argv[2], argv[3], argv[4]); + if(argc == 3 && strcmp(argv[1], "--bitcast") == 0) + return print_bitcast(argv[2]); CHECK(argc == 1); CHECK(test_unsigned_ops() == 0); CHECK(test_packed_only_values() == 0); CHECK(test_smem_vector_elements() == 0); + CHECK(test_bitcast_width() == 0); + CHECK(test_scalar_widths() == 0); printf("rocke_ir_types: ok\n"); return 0; } From e04f6f3e0790c86a7b0b0723b02bf32ca60230be Mon Sep 17 00:00:00 2001 From: Chris Millette Date: Fri, 9 Oct 2026 15:57:19 -0400 Subject: [PATCH 07/11] feat(rocke): define how sub-byte values pack into 32-bit registers register_packing(dtype, count) in core/storage.py returns the dense FragmentPacking for count elements of a sub-byte dtype (fp4, fp6, i4, u4) held in whole 32-bit registers: elements pack low-bit-first with no padding, so 32 fp4 values take 4 registers and 32 fp6 values take 6. A payload that leaves a register partly filled raises ValueError, a predicate raises, and a dtype of 8 or more bits returns None: it has an LLVM element type, and LLVM packs its vectors into registers itself. packed_word_range(packing, start, length) returns the range of registers holding elements [start, start + length), and raises when the range does not start and end on a register boundary (naming the element step that does) or runs past the fragment. Both reject a non-integer count, start or length with TypeError. rocke_register_packing and rocke_packed_word_range in rocke/storage.h are the C++ twins. They keep the header's promise that no exception crosses the C API: they return false and write Python's refusal text into a caller-owned buffer. The rocke_storage test binary answers both questions on the command line, and test_storage.py requires its answers and texts to equal core/storage.py's. Both are pure arithmetic with no caller yet. dtype_storage_packing.md describes them. Co-Authored-By: Claude --- .../rocke/platform/cpp/core/storage.cpp | 148 ++++++++++++++ .../platform/cpp/include/rocke/storage.h | 35 +++- .../architecture/dtype_storage_packing.md | 12 ++ .../platform/python/rocke/core/storage.py | 65 ++++++- .../platform/tests/core/test_storage.cpp | 125 +++++++++++- .../rocke/platform/tests/core/test_storage.py | 181 +++++++++++++++++- 6 files changed, 560 insertions(+), 6 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 c951bf74096f..0f43394d7cf0 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 @@ -5,8 +5,14 @@ */ #include "rocke/storage.h" +#include +#include +#include #include +#include "rocke/dtypes_internal.h" +#include "rocke/error.hpp" + 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; @@ -213,3 +219,145 @@ bool rocke_matrix_fragment_coord( && add(chunk, lane / p->lanes_per_group, &chunk) && mul(chunk, p->chunk_elements, &chunk) && add(chunk, slot % p->chunk_elements, k); } + +/* ======================= sub-byte values in registers ======================= + * + * Mirrors core/storage.py register_packing and packed_word_range. The bodies + * throw a ckc::Error where Python raises, with Python's text; the extern "C" + * entry points turn it into `false` plus the message, so no exception crosses + * the C API. + */ + +/* Width of the registers packed operands occupy (Python _REGISTER_BITS). */ +static const int k_register_bits = 32; + +/* Python _checked's refusal of a size or offset past the uint64 range. */ +[[noreturn]] static void raise_storage_overflow() +{ + throw ckc::ValueError("storage size or offset exceeds uint64 range"); +} + +/* Python _require_int_at_least, less the integer-type check C makes at compile + * time. */ +static void require_at_least(const char* name, int64_t value, int64_t minimum) +{ + if(value < minimum) + throw ckc::ValueError(ckc::format_error( + "%s must be at least %" PRId64 ", got %" PRId64, name, minimum, value)); +} + +/* Run `body` and report its outcome through the C API: true, or false with the + * exception's message copied into `err`. */ +template +static bool report_through_c_api(char* err, size_t err_cap, Body&& body) +{ + try + { + body(); + } + catch(const std::exception& error) + { + if(err && err_cap) + snprintf(err, err_cap, "%s", error.what()); + return false; + } + if(err && err_cap) + err[0] = '\0'; + return true; +} + +bool rocke_register_packing(const char* dtype, + int64_t count, + bool* packed, + rocke_fragment_packing_t* out, + char* err, + size_t err_cap) +{ + return report_through_c_api(err, err_cap, [&] { + if(!packed || !out) + throw ckc::ValueError("rocke_register_packing: packed and out must not be NULL"); + const rocke_dtype_info_t* info = rocke_require_dtype_info(dtype); + require_at_least("element count", count, 1); + if(info->category == ROCKE_DTYPE_PREDICATE) + throw ckc::ValueError( + ckc::format_error("%s is a predicate, not register data", info->name)); + if(info->encoded_bits >= 8) + { + *packed = false; + return; + } + uint64_t bits; + if(!mul(uint64_t(count), uint64_t(info->encoded_bits), &bits)) + raise_storage_overflow(); + if(bits % k_register_bits) + throw ckc::ValueError(ckc::format_error( + "%" PRId64 " %s elements occupy %" PRIu64 " bits, not a whole number of %d-bit " + "registers -- packed operands must fill whole registers", + count, + info->name, + bits, + k_register_bits)); + rocke_bit_packing_t element; + if(!rocke_bit_packing_init(&element, info->encoded_bits, 0) + || !rocke_fragment_packing_init( + out, &element, uint64_t(count), k_register_bits, bits / k_register_bits)) + raise_storage_overflow(); + *packed = true; + }); +} + +bool rocke_packed_word_range(const rocke_fragment_packing_t* packing, + int64_t start, + int64_t length, + uint64_t* first_word, + uint64_t* end_word, + char* err, + size_t err_cap) +{ + return report_through_c_api(err, err_cap, [&] { + // Python cannot hold an invalid FragmentPacking; a C descriptor can be + // filled in by hand, so check it the way rocke_fragment_pack does. + rocke_fragment_packing_t checked; + if(!packing || !first_word || !end_word + || !rocke_fragment_packing_init(&checked, + &packing->packing, + packing->count, + packing->carrier_bits, + packing->carrier_count)) + throw ckc::ValueError("rocke_packed_word_range: needs a valid fragment packing " + "(rocke_fragment_packing_init) and two output words"); + require_at_least("start", start, 0); + require_at_least("length", length, 1); + // Both are below 2^63, so the sum fits. + const uint64_t first = uint64_t(start); + const uint64_t end = first + uint64_t(length); + if(end > packing->count) + throw ckc::ValueError(ckc::format_error("elements [%" PRIu64 ", %" PRIu64 + ") run past the fragment's " + "%" PRIu64 " elements", + first, + end, + packing->count)); + uint64_t first_bit, end_bit; + if(!rocke_bit_packing_offset(&packing->packing, first, &first_bit) + || !rocke_bit_packing_offset(&packing->packing, end, &end_bit)) + raise_storage_overflow(); + const uint64_t word = uint64_t(packing->carrier_bits); + if(first_bit % word || end_bit % word) + { + uint64_t step, carriers; + rocke_bit_packing_group(&packing->packing, packing->carrier_bits, &step, &carriers); + throw ckc::ValueError(ckc::format_error( + "elements [%" PRIu64 ", %" PRIu64 ") of a fragment packed at %d bits per element " + "do not start and end on a register boundary (%d bits) -- start and length " + "must be multiples of %" PRIu64 " elements", + first, + end, + packing->packing.slot_bits, + packing->carrier_bits, + step)); + } + *first_word = first_bit / word; + *end_word = end_bit / word; + }); +} 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 37a3ee9be0c9..ac0707c69396 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,8 +3,13 @@ /* Shared bit storage; mirrors core/storage.py. * BitPacking -> rocke_bit_packing_t * FragmentPacking -> rocke_fragment_packing_t + * register_packing -> rocke_register_packing + * packed_word_range -> rocke_packed_word_range * All descriptors are caller-owned. Functions return false for invalid inputs, - * overflow or insufficient buffers; no exceptions cross this C API. + * overflow or insufficient buffers; no exceptions cross this C API. The two + * register functions also write the reason into a caller-owned `err` buffer + * (truncated to `err_cap`, may be NULL), worded as core/storage.py words the + * ValueError it raises. */ #ifndef ROCKE_STORAGE_H #define ROCKE_STORAGE_H @@ -83,6 +88,34 @@ bool rocke_fragment_pack(const rocke_fragment_packing_t* p, size_t count, uint64_t* carriers, size_t carrier_count); + +/* How `count` elements of `dtype` (any accepted spelling) sit densely in 32-bit + * registers: low-bit-first, no padding between elements, whole registers only. + * On success *packed is false for a dtype of 8 or more bits (Python's None: it + * has an LLVM element type, and LLVM packs its vectors into registers itself), + * and true with the dense packing in *out for a sub-byte dtype. Fails for an + * unknown dtype, a count below 1, a predicate, or a payload that leaves a + * register partly filled. Python's TypeError for a non-integer count has no C + * twin. */ +bool rocke_register_packing(const char* dtype, + int64_t count, + bool* packed, + rocke_fragment_packing_t* out, + char* err, + size_t err_cap); + +/* The registers [*first_word, *end_word) that hold elements + * [start, start + length) of `packing`. Fails for a start below 0, a length + * below 1, a range that runs past the fragment, or one that does not start and + * end on a register boundary (the message names the element step that does). + * Python's TypeError for a non-integer start or length has no C twin. */ +bool rocke_packed_word_range(const rocke_fragment_packing_t* packing, + int64_t start, + int64_t length, + uint64_t* first_word, + uint64_t* end_word, + char* err, + size_t err_cap); #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 c6c4788f951b..094cd92bd58a 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 @@ -69,6 +69,18 @@ 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. +`register_packing(dtype, count)` builds the dense `FragmentPacking` for `count` +sub-byte elements in 32-bit registers and refuses a payload that leaves a register +partly filled. A dtype of 8 or more bits returns `None`: it has an LLVM element +type, and LLVM packs its vectors into registers itself (a `<4 x half>` occupies +two registers), so rocke never packs it by hand. A predicate (`i1`) is refused. + +`packed_word_range(packing, start, length)` names the registers that hold an +element range; the range must start and end on a register boundary, so its start +and length are multiples of 8 elements for 4-bit types and 16 for 6-bit types. +The C++ engine's twins are `rocke_register_packing` and `rocke_packed_word_range` +in `cpp/include/rocke/storage.h`, with the same answers and refusal texts. + 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. 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 e0770a939c1c..e17702939ee2 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 @@ -4,11 +4,14 @@ from __future__ import annotations +from collections.abc import Sequence from dataclasses import dataclass from math import gcd -from collections.abc import Sequence + +from .dtypes import DTypeCategory, dtype_info _MAX_BITS = (1 << 64) - 1 +_REGISTER_BITS = 32 def _checked(value: int) -> int: @@ -122,6 +125,66 @@ def pack(self, patterns: Sequence[int]) -> tuple[int, ...]: ) +def _require_int_at_least(name: str, value: int, minimum: int) -> int: + """Return ``value`` if it is an int (not a bool) of at least ``minimum``, else raise.""" + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an integer, got {value!r}") + if value < minimum: + raise ValueError(f"{name} must be at least {minimum}, got {value}") + return value + + +def register_packing(dtype: str, count: int) -> FragmentPacking | None: + """How ``count`` elements of a sub-byte ``dtype`` sit densely in 32-bit registers. + + Elements pack low-bit-first with no padding between them, and must fill + whole registers. A dtype of 8 or more bits returns None: it has an LLVM + element type, and LLVM packs its vectors into registers itself (a + ``<4 x half>`` occupies two registers), so rocke never packs it by hand. A + predicate raises. + """ + info = dtype_info(dtype) + _require_int_at_least("element count", count, 1) + if info.category is DTypeCategory.PREDICATE: + raise ValueError(f"{info.name} is a predicate, not register data") + if info.encoded_bits >= 8: + return None + bits = count * info.encoded_bits + if bits % _REGISTER_BITS: + raise ValueError( + f"{count} {info.name} elements occupy {bits} bits, not a whole number of " + f"{_REGISTER_BITS}-bit registers -- packed operands must fill whole " + "registers" + ) + return FragmentPacking( + BitPacking(info.encoded_bits), count, _REGISTER_BITS, bits // _REGISTER_BITS + ) + + +def packed_word_range(packing: FragmentPacking, start: int, length: int) -> range: + """The registers that hold elements ``[start, start + length)`` of ``packing``.""" + end = _require_int_at_least("start", start, 0) + _require_int_at_least( + "length", length, 1 + ) + if end > packing.count: + raise ValueError( + f"elements [{start}, {end}) run past the fragment's {packing.count} " + "elements" + ) + first_bit = packing.packing.bit_offset(start) + end_bit = packing.packing.bit_offset(end) + word = packing.carrier_bits + if first_bit % word or end_bit % word: + step, _ = packing.packing.group(word) + raise ValueError( + f"elements [{start}, {end}) of a fragment packed at " + f"{packing.packing.slot_bits} bits per element do not start and end on " + f"a register boundary ({word} bits) -- start and length must be " + f"multiples of {step} elements" + ) + return range(first_bit // word, end_bit // word) + + @dataclass(frozen=True) class MatrixFragmentLayout: """Contiguous K chunks interleaved between lane groups. 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 64bc401dfefb..51c66f2811d6 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 @@ -1,9 +1,22 @@ // 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. */ +/* Host packing fixtures and native authoring emission for test_mma_io.py. + * + * It also answers register-packing questions for test_storage.py, which + * compares each answer with core/storage.py (found through ROCKE_STORAGE_TEST): + * + * --register-packing + * print "none" (not packed by hand), the packing's fields, or + * "error: ". + * --packed-word-range + * pack elements, then print the registers holding + * elements [start, start + length) as " ", or + * "error: ". + */ #include "rocke/storage.h" #include +#include #include #include #include @@ -158,6 +171,111 @@ static int emit(const char* dtype, bool hip, bool include_prologue = true) return 0; } +/* Parse a whole decimal argument; false when any of it is not a number. */ +static bool parse_int64(const char* text, int64_t* value) +{ + char* end = NULL; + *value = strtoll(text, &end, 10); + return end != text && *end == '\0'; +} + +static int print_register_packing(const char* dtype, const char* count_text) +{ + int64_t count; + CHECK(parse_int64(count_text, &count)); + bool packed = false; + rocke_fragment_packing_t packing; + char err[ROCKE_ERR_MSG_CAP]; + if(!rocke_register_packing(dtype, count, &packed, &packing, err, sizeof err)) + printf("error: %s\n", err); + else if(!packed) + printf("none\n"); + else + printf("element_bits=%d slot_bits=%d count=%" PRIu64 + " carrier_bits=%d carrier_count=%" PRIu64 "\n", + packing.packing.element_bits, + packing.packing.slot_bits, + packing.count, + packing.carrier_bits, + packing.carrier_count); + return 0; +} + +static int print_packed_word_range(const char* dtype, + const char* count_text, + const char* start_text, + const char* length_text) +{ + int64_t count, start, length; + CHECK(parse_int64(count_text, &count) && parse_int64(start_text, &start) + && parse_int64(length_text, &length)); + bool packed = false; + rocke_fragment_packing_t packing; + char err[ROCKE_ERR_MSG_CAP]; + CHECK(rocke_register_packing(dtype, count, &packed, &packing, err, sizeof err) && packed); + uint64_t first, end; + if(rocke_packed_word_range(&packing, start, length, &first, &end, err, sizeof err)) + printf("%" PRIu64 " %" PRIu64 "\n", first, end); + else + printf("error: %s\n", err); + return 0; +} + +/* The register-packing answers and refusals, with core/storage.py's text. */ +static int test_register_packing() +{ + bool packed = false; + rocke_fragment_packing_t packing; + char err[ROCKE_ERR_MSG_CAP]; + CHECK(rocke_register_packing("fp6e2m3", 32, &packed, &packing, err, sizeof err) && packed); + CHECK(packing.packing.element_bits == 6 && packing.count == 32); + CHECK(packing.carrier_bits == 32 && packing.carrier_count == 6); + CHECK(err[0] == '\0'); + CHECK(rocke_register_packing("fp16", 8, &packed, &packing, err, sizeof err) && !packed); + const struct + { + const char* dtype; + int64_t count; + const char* error; + } refused[] = { + {"fp6e2m3", + 3, + "3 fp6e2m3 elements occupy 18 bits, not a whole number of 32-bit registers -- packed " + "operands must fill whole registers"}, + {"fp4e2m1", -4, "element count must be at least 1, got -4"}, + {"i1", 32, "i1 is a predicate, not register data"}, + {"custom_format", 8, "unknown dtype 'custom_format'"}, + }; + for(const auto& r : refused) + { + CHECK(!rocke_register_packing(r.dtype, r.count, &packed, &packing, err, sizeof err)); + CHECK(strcmp(err, r.error) == 0); + } + + CHECK(rocke_register_packing("fp4e2m1", 64, &packed, &packing, err, sizeof err) && packed); + uint64_t first, end; + CHECK(rocke_packed_word_range(&packing, 32, 32, &first, &end, err, sizeof err)); + CHECK(first == 4 && end == 8); + CHECK(!rocke_packed_word_range(&packing, 3, 32, &first, &end, err, sizeof err)); + CHECK(strcmp(err, + "elements [3, 35) of a fragment packed at 4 bits per element do not start and " + "end on a register boundary (32 bits) -- start and length must be multiples of " + "8 elements") + == 0); + CHECK(!rocke_packed_word_range(&packing, 56, 16, &first, &end, err, sizeof err)); + CHECK(strcmp(err, "elements [56, 72) run past the fragment's 64 elements") == 0); + // A message longer than the caller's buffer is truncated, never overrun. + char small[8]; + memset(small, 0xA5, sizeof small); + CHECK(!rocke_packed_word_range(&packing, -8, 8, &first, &end, small, 6)); + CHECK(strcmp(small, "start") == 0 && (unsigned char)small[6] == 0xA5); + // A hand-filled descriptor that no init() would accept is refused. + rocke_fragment_packing_t invalid = packing; + invalid.carrier_count = 1; + CHECK(!rocke_packed_word_range(&invalid, 0, 8, &first, &end, err, sizeof err)); + return 0; +} + static int test_fragment_inputs() { const struct @@ -247,6 +365,10 @@ int main(int argc, char** argv) return emit(argv[2], true); if(argc == 3 && strcmp(argv[1], "--hip-bare") == 0) return emit(argv[2], true, false); + if(argc == 4 && strcmp(argv[1], "--register-packing") == 0) + return print_register_packing(argv[2], argv[3]); + if(argc == 6 && strcmp(argv[1], "--packed-word-range") == 0) + return print_packed_word_range(argv[2], argv[3], argv[4], argv[5]); if(argc == 3 && strcmp(argv[1], "--parse") == 0) { rocke_ir_builder_t b; @@ -293,6 +415,7 @@ int main(int argc, char** argv) } CHECK(test_fragment_inputs() == 0); CHECK(test_fragment_offset_bounds() == 0); + CHECK(test_register_packing() == 0); for(int count : {1, 2, 4, 8}) { const rocke_scale_packing_t scales = {count, 32}; 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 d97df3bedf81..d71e27c67c4b 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,24 @@ # Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -"""Independent bit-stream fixtures for matrices and scales.""" +"""Independent bit-stream fixtures for matrices and scales, and register packing. -import pytest +The C++ engine's register packing answers through the ``rocke_storage`` test +binary (``ROCKE_STORAGE_TEST``); the native tests require its answers and its +refusal texts to equal ``core/storage.py``'s. +""" + +import os +import re +import subprocess +import pytest from rocke.core.arch.wmma_scale import ScalePacking, scaled_matrix_layout -from rocke.core.storage import BitPacking, FragmentPacking +from rocke.core.storage import ( + BitPacking, + FragmentPacking, + packed_word_range, + register_packing, +) @pytest.mark.parametrize("bits", [4, 6, 8, 16, 32, 64]) @@ -84,3 +97,165 @@ def test_invalid_packing_and_overflow(): BitPacking(6).byte_size(1 << 63) with pytest.raises(ValueError, match="fit"): BitPacking(6).pack([64]) + + +_PACKED = [ + ("fp4e2m1", 32, 4), + ("fp6e2m3", 32, 6), + ("fp6e3m2", 16, 3), + ("i4", 16, 2), + ("u4", 8, 1), +] +_NOT_PACKED = ["fp8e4m3", "i8", "fp16", "bf16", "fp32", "fp64"] +_WORD_RANGES = [ + ("fp4e2m1", 32, 32, range(4, 8)), + ("fp6e2m3", 32, 32, range(6, 12)), + ("fp6e2m3", 16, 16, range(3, 6)), + ("i4", 0, 8, range(1)), +] +_FRAGMENT_OVERRUNS = [ + (-8, 8, "^start must be at least 0, got -8$"), + (8, 0, "^length must be at least 1, got 0$"), + (56, 16, "^elements \\[56, 72\\) run past the fragment's 64 elements$"), +] + + +@pytest.mark.parametrize(("dtype", "count", "registers"), _PACKED) +def test_sub_byte_elements_pack_densely_into_whole_registers(dtype, count, registers): + packing = register_packing(dtype, count) + assert packing.carrier_bits == 32 + assert packing.carrier_count == registers + assert packing.padding_bits == 0 + + +@pytest.mark.parametrize("dtype", _NOT_PACKED) +def test_byte_or_wider_elements_are_not_packed(dtype): + assert register_packing(dtype, 8) is None + + +def test_register_packing_lays_out_bits_like_a_hand_built_fragment(): + patterns = [(3 * i + 1) % 64 for i in range(16)] + by_hand = FragmentPacking(BitPacking(6), 16, 32, 3) + assert register_packing("fp6e2m3", 16).pack(patterns) == by_hand.pack(patterns) + + +def test_a_payload_that_leaves_a_register_partly_filled_raises(): + message = ( + "3 fp6e2m3 elements occupy 18 bits, not a whole number of 32-bit registers " + "-- packed operands must fill whole registers" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + register_packing("fp6e2m3", 3) + + +@pytest.mark.parametrize("count", [2.0, True]) +def test_register_packing_needs_an_integer_count(count): + with pytest.raises(TypeError, match="^element count must be an integer"): + register_packing("fp4e2m1", count) + + +@pytest.mark.parametrize("count", [0, -4]) +def test_register_packing_needs_a_positive_count(count): + with pytest.raises(ValueError, match="^element count must be at least 1"): + register_packing("fp4e2m1", count) + + +def test_a_predicate_is_not_register_data(): + with pytest.raises(ValueError, match="^i1 is a predicate, not register data$"): + register_packing("i1", 32) + + +@pytest.mark.parametrize(("dtype", "start", "length", "words"), _WORD_RANGES) +def test_an_element_range_maps_to_the_registers_that_hold_it( + dtype, start, length, words +): + assert packed_word_range(register_packing(dtype, 64), start, length) == words + + +def test_an_element_range_off_a_register_boundary_raises(): + message = ( + "elements [3, 35) of a fragment packed at 4 bits per element do not start " + "and end on a register boundary (32 bits) -- start and length must be " + "multiples of 8 elements" + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + packed_word_range(register_packing("fp4e2m1", 64), 3, 32) + + +@pytest.mark.parametrize(("start", "length", "message"), _FRAGMENT_OVERRUNS) +def test_an_element_range_must_lie_inside_the_fragment(start, length, message): + with pytest.raises(ValueError, match=message): + packed_word_range(register_packing("fp4e2m1", 64), start, length) + + +@pytest.mark.parametrize(("start", "length"), [(True, 8), (8.0, 8), (0, 8.0)]) +def test_an_element_range_needs_integer_ends(start, length): + with pytest.raises(TypeError, match="must be an integer"): + packed_word_range(register_packing("fp4e2m1", 64), start, length) + + +def test_the_boundary_error_names_the_step_for_fp6(): + with pytest.raises(ValueError, match="multiples of 16 elements$"): + packed_word_range(register_packing("fp6e2m3", 64), 8, 16) + + +def _native(*arguments: str) -> str: + """One answer of the ``rocke_storage`` binary, or skip when it is not built.""" + executable = os.environ.get("ROCKE_STORAGE_TEST") + if not executable: + pytest.skip("set ROCKE_STORAGE_TEST to the built native storage test") + return subprocess.run( + [executable, *arguments], check=True, capture_output=True, text=True + ).stdout.rstrip("\n") + + +def _python_register_packing(dtype: str, count: int) -> str: + """The twin of ``rocke_storage --register-packing``.""" + try: + packing = register_packing(dtype, count) + except ValueError as error: + return f"error: {error}" + if packing is None: + return "none" + return ( + f"element_bits={packing.packing.element_bits} " + f"slot_bits={packing.packing.slot_bits} count={packing.count} " + f"carrier_bits={packing.carrier_bits} carrier_count={packing.carrier_count}" + ) + + +def _python_packed_word_range(dtype: str, count: int, start: int, length: int) -> str: + """The twin of ``rocke_storage --packed-word-range``.""" + try: + words = packed_word_range(register_packing(dtype, count), start, length) + except ValueError as error: + return f"error: {error}" + return f"{words.start} {words.stop}" + + +@pytest.mark.parametrize( + ("dtype", "count"), + [(dtype, count) for dtype, count, _ in _PACKED] + + [(dtype, 8) for dtype in _NOT_PACKED] + + [ + ("fp6e2m3", 3), + ("fp4e2m1", 0), + ("fp4e2m1", -4), + ("i1", 32), + ("custom_format", 8), + ], +) +def test_native_register_packing_matches(dtype, count): + native = _native("--register-packing", dtype, str(count)) + assert native == _python_register_packing(dtype, count) + + +@pytest.mark.parametrize( + ("dtype", "start", "length"), + [(dtype, start, length) for dtype, start, length, _ in _WORD_RANGES] + + [("fp4e2m1", start, length) for start, length, _ in _FRAGMENT_OVERRUNS] + + [("fp4e2m1", 3, 32), ("fp6e2m3", 8, 16)], +) +def test_native_packed_word_range_matches(dtype, start, length): + native = _native("--packed-word-range", dtype, "64", str(start), str(length)) + assert native == _python_packed_word_range(dtype, 64, start, length) From 575e09e512c6a30fe15146a314da156d61d81b53 Mon Sep 17 00:00:00 2001 From: Chris Millette Date: Fri, 9 Oct 2026 16:45:57 -0400 Subject: [PATCH 08/11] test(rocke): pin every IR scalar type's width, and document adding a dtype test_dtypes.py gains _IR_SCALAR_BITS, the width of every IR scalar type written out by hand. The format table in core/dtypes.py and type_bits are both checked against it, and a completeness test fails when an IR scalar type is added without its width. The scalar width check in test_bitcast_width.py moves here. The same map checks the C++ engine: rocke_ir_types --scalar-bits prints every IR scalar type with the width rocke_i_type_bits gives it. Its name list holds one name per scalar kind, which a static_assert and a self-check enforce, so a new kind cannot go unchecked. extending.md section 5.4 becomes the one recipe for adding a dtype: format table, IR type, reference width, register form, memory, conversions, MMA catalog, manifest, the C++ engine mirror, docs and tests. backend_details.md points to it instead of keeping its own checklist, and its list of scalar LLVM mappings gains i16, f64 and u8. Co-Authored-By: Claude --- .../dsl_docs/development/extending.md | 44 ++++++++++--- .../dsl_docs/ir_lowering/backend_details.md | 12 +--- .../platform/tests/core/test_bitcast_width.py | 23 ++----- .../rocke/platform/tests/core/test_dtypes.py | 61 ++++++++++++++++++- .../platform/tests/core/test_ir_types.cpp | 38 ++++++++++++ 5 files changed, 142 insertions(+), 36 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/development/extending.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/development/extending.md index 4910b3f3b711..f4bafccf1d84 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/development/extending.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/development/extending.md @@ -404,14 +404,42 @@ Asserts: - a small kernel using the atom lowers to the expected intrinsic call; - `MfmaAtom.lane_to_output` for a few specific lanes returns the expected `(row, col)` SSA values when materialized (compare LLVM text against expected arithmetic). -For a new dtype: - -- Add the `Type` singleton in `core/ir.py`. -- Add LLVM `_llvm_type` mapping. -- Add load / store / cast support paths. -- Add MFMA atom variants (`mfma_f32_xx_
`). -- Add manifest dtype string canonicalization (`helpers/spec.py::ptr_type_str`). -- Verify with both `lower_kernel_to_llvm` and an end-to-end `compile_kernel + KernelLauncher` test. +### 5.4 Adding a dtype + +One fact, one home. Work through the list in order; a step that does not apply +(for example, a 4-bit format has no LLVM register type) is skipped, not faked. + +1. **Format** -- `core/dtypes.py`: a `(name, bits)` row under its category in + `_DTYPE_INFO`, and every accepted spelling in `_DTYPE_ALIASES`. This table is the + home of a format's bit width; `type_bits` and the packing helpers read it. +2. **IR type** -- `core/ir.py`: the `Type` constant and its `dtype_to_ir_type` + entry; export it from `core/__init__.py` and `rocke/__init__.py` if kernels name it. +3. **Reference width** -- `tests/core/test_dtypes.py`: add the width to + `_IR_SCALAR_BITS`, written by hand. Its completeness test fails until you do, + so the format table is always checked against an independent source. The + same map checks the C++ engine's widths (`rocke_ir_types --scalar-bits`), whose + name list in `tests/core/test_ir_types.cpp` fails to compile until the new + scalar kind is named there. +4. **Register form** -- `core/lower_llvm.py::_llvm_type` if values of the type live + in registers as an LLVM type. Sub-byte formats have none: they travel packed in + 32-bit registers (`core/storage.py::register_packing`). +5. **Memory and constants** -- loads, stores, `zero_vec`, and the LDS element-byte + tables, only for the widths the type needs. +6. **Conversions** -- explicit `cvt__to_` primitives, each one + instruction; no existing cast widens to the new type implicitly. +7. **MMA** -- catalog rows in `core/arch/data/arch_specs.json` if an instruction + takes the type as an operand or accumulator. +8. **Manifest** -- `helpers/spec.py::ptr_type_str` canonicalization. +9. **C++ engine** -- mirror the change in the same change (`cpp/core/dtypes.cpp`, + `cpp/core/ir/core_types.cpp`, `cpp/core/lower_llvm/`, `cpp/core/lower_hip/`): + both engines emit the same IR and refuse the same input with the same text. +10. **Docs** -- the type's row in [dtype support](../reference/dtype_support.md). +11. **Tests** -- a lowering test that passes + `rocke.tests.ir_validity.assert_target_accepts`, and a numeric test that names + the architectures it supports (`rocke.tests.gpu.run_on_device`). Lower through + `lower_kernel_to_llvm`, so a `ROCKE_BACKEND=both ROCKE_CPP_STRICT=1` run + compares the engines, and state a refusal with + `rocke.tests.engines.assert_both_engines_reject`. ## 6. Adding A New Optimization Lever To The Autotuner diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/ir_lowering/backend_details.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/ir_lowering/backend_details.md index 84216ec469d8..e2e7c0d0ebaa 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/ir_lowering/backend_details.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/ir_lowering/backend_details.md @@ -77,17 +77,9 @@ Verified mappings: - pointers: `ptr addrspace(1)` (`global`), `ptr addrspace(3)` (`lds`), `ptr` otherwise. - vectors: ``. - SmemType: lowers to `ptr addrspace(3)` for value-level usage. -- scalars: `i1`, `i8`, `i32`, `i64`, `f16`, `bf16`, `f32`. (FP8E4M3 / BF8E5M2 lower through the convert intrinsics rather than as first-class load/store types.) +- scalars: `i1`, `i8` (also `u8`), `i16`, `i32`, `i64`, `f16`, `bf16`, `f32`, `f64`. (FP8E4M3 / BF8E5M2 lower through the convert intrinsics rather than as first-class load/store types.) -If you add a new dtype, check: - -- scalar LLVM type mapping in `_llvm_type`; -- vector mapping if `` is needed; -- pointer / address-space handling; -- load and store ops; -- cast / pack / unpack; -- MFMA atom support (and `MFMA_F16_ATOMS` catalog or new catalog); -- manifest dtype parsing and reference verification. +To add a dtype, follow [Adding a dtype](../development/extending.md#54-adding-a-dtype). ## Address Spaces diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_bitcast_width.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_bitcast_width.py index 64aa4f69bec8..568411cd4dff 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_bitcast_width.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_bitcast_width.py @@ -3,12 +3,12 @@ """A bitcast keeps every bit, so ``bitcast`` and ``vec_bitcast`` refuse a width change. ``type_bits`` is the one place that answers how many bits a value of an IR type -occupies. Equal-width casts lower to one LLVM ``bitcast``; a cast that changes the -width fails when the kernel is built instead of inside clang. Both engines lower -the casts identically, and the C++ builders refuse the same casts with the same -text: they are reached through the ``rocke_ir_types`` test binary -(``ROCKE_IR_TYPES_TEST``), because the serialized-IR seam lowers kernels the -Python builder built. +occupies; the scalar widths are pinned in test_dtypes.py. Equal-width casts lower to +one LLVM ``bitcast``; a cast that changes the width fails when the kernel is built +instead of inside clang. Both engines lower the casts identically, and the C++ +builders refuse the same casts with the same text: they are reached through the +``rocke_ir_types`` test binary (``ROCKE_IR_TYPES_TEST``), because the serialized-IR +seam lowers kernels the Python builder built. """ import os @@ -17,7 +17,6 @@ import pytest from rocke.core import ir -from rocke.core.dtypes import dtype_info from rocke.core.ir import ( F16, F32, @@ -38,16 +37,6 @@ def _lower(b: IRBuilder) -> str: return lower_with_both_engines(b.kernel, "gfx90a") -def _scalar_types() -> list: - """Every scalar IR type constant the module defines, so later types are covered.""" - return [t for t in vars(ir).values() if type(t) is ir.Type] - - -@pytest.mark.parametrize("t", _scalar_types(), ids=lambda t: t.name) -def test_a_scalar_occupies_its_format_width(t): - assert type_bits(t) == dtype_info(t.name).encoded_bits - - def test_a_vector_occupies_count_times_its_element_width(): assert type_bits(VectorType(F16, 8)) == 128 assert type_bits(VectorType(I32, 4)) == 128 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 a21a7e31314c..e6b5f13aa423 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,8 @@ """Core dtype names are independent of target support and retain public aliases.""" import ast +import os +import subprocess from typing import get_args, get_type_hints import pytest @@ -10,7 +12,8 @@ from rocke.core.dtypes import normalize_dtype 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 import ir +from rocke.core.ir import FP8E4M3, I8, IRBuilder, Type, dtype_to_ir_type, type_bits from rocke.helpers.mma_io import storage_ir_type from rocke.core.ir_serialize import parse, serialize from rocke.helpers.quant import ( @@ -48,6 +51,62 @@ def test_logical_dtype_and_serialization(dtype, bits): assert serialize(parse(serialize(b.kernel))) == serialize(b.kernel) +# The width of every IR scalar type, written out by hand so that the format table +# in core/dtypes.py and type_bits are checked against an independent source. +# Adding an IR scalar type means adding its width here; see section 5.4, +# "Adding a dtype", in dsl_docs/development/extending.md. +_IR_SCALAR_BITS = { + "i1": 1, + "i4": 4, + "u4": 4, + "i8": 8, + "u8": 8, + "i16": 16, + "i32": 32, + "i64": 64, + "f16": 16, + "bf16": 16, + "f32": 32, + "tf32": 32, + "f64": 64, + "fp8e4m3": 8, + "bf8e5m2": 8, + "fp6e2m3": 6, + "fp6e3m2": 6, + "fp4e2m1": 4, + "e8m0": 8, + "e5m3": 8, +} + + +def test_every_ir_scalar_type_has_a_reference_width(): + defined = {t.name for t in vars(ir).values() if type(t) is Type} + assert not defined - set(_IR_SCALAR_BITS), ( + f"add the width of {sorted(defined - set(_IR_SCALAR_BITS))} to " + "_IR_SCALAR_BITS" + ) + assert not set(_IR_SCALAR_BITS) - defined + + +@pytest.mark.parametrize(("name", "bits"), sorted(_IR_SCALAR_BITS.items())) +def test_an_ir_scalar_type_has_its_reference_width(name, bits): + assert dtype_info(name).encoded_bits == bits + assert type_bits(Type(name)) == bits + + +def test_native_scalar_widths_match_the_reference(): + # The C++ engine lists every scalar kind it defines (a static_assert keeps + # the list complete) with the width rocke_i_type_bits gives it. + executable = os.environ.get("ROCKE_IR_TYPES_TEST") + if not executable: + pytest.skip("set ROCKE_IR_TYPES_TEST to the built rocke_ir_types test") + lines = subprocess.run( + [executable, "--scalar-bits"], check=True, capture_output=True, text=True + ).stdout.splitlines() + native = {name: int(bits) for name, bits in (line.split() for line in lines)} + assert native == _IR_SCALAR_BITS + + @pytest.mark.parametrize("dtype", ["fp4", "fp6", "bf6"]) def test_quant_type_recognition_does_not_enable_integer_conversion(dtype): logical = quant_ir_type(dtype) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp index 7e067ffc75a9..f822b523395a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp @@ -25,6 +25,10 @@ * build one LDS vector access of n elem values; print "ok", or * "error: " when the builder refuses it. * + * --scalar-bits + * print " " for every IR scalar type, the width the engine + * gives it (test_dtypes.py compares the map with _IR_SCALAR_BITS). + * * --bitcast * build the bitcast named (a row of k_bitcast_cases, mirrored in * tests/core/test_bitcast_width.py); print "ok", or @@ -609,6 +613,37 @@ static int test_scalar_widths() return 0; } +/* Every IR scalar type by name: one per rocke_scalar_kind_t, which the + * static_assert and test_scalar_kinds_are_covered hold to. Widths are not + * listed: they come from the engine. */ +static const char* const k_scalar_names[] + = {"i1", "i4", "u4", "i8", "u8", "i16", "i32", "i64", "f16", "bf16", + "f32", "tf32", "f64", "fp8e4m3", "bf8e5m2", "fp6e2m3", "fp6e3m2", "fp4e2m1", "e8m0", "e5m3"}; +static_assert(sizeof k_scalar_names / sizeof k_scalar_names[0] == ROCKE_SCALAR__COUNT, + "add the new scalar type's name to k_scalar_names"); + +/* Each name resolves to the singleton of that name, and no two names share a + * scalar kind, so the list holds every kind exactly once. */ +static int test_scalar_kinds_are_covered() +{ + bool seen[ROCKE_SCALAR__COUNT] = {}; + for(const char* name : k_scalar_names) + { + const rocke_type_t* t = rocke_scalar_by_name(name); + CHECK(t != NULL && strcmp(t->name, name) == 0); + CHECK(!seen[t->scalar]); + seen[t->scalar] = true; + } + return 0; +} + +static int print_scalar_bits() +{ + for(const char* name : k_scalar_names) + printf("%s %d\n", name, rocke_i_type_bits(rocke_scalar_by_name(name))); + return 0; +} + int main(int argc, char** argv) { if(argc == 3 && strcmp(argv[1], "--unsigned") == 0) @@ -617,6 +652,8 @@ int main(int argc, char** argv) return print_packed_only_value(argv[2], argv[3]); if(argc == 5 && strcmp(argv[1], "--smem-vector") == 0) return print_smem_vector_access(argv[2], argv[3], argv[4]); + if(argc == 2 && strcmp(argv[1], "--scalar-bits") == 0) + return print_scalar_bits(); if(argc == 3 && strcmp(argv[1], "--bitcast") == 0) return print_bitcast(argv[2]); CHECK(argc == 1); @@ -625,6 +662,7 @@ int main(int argc, char** argv) CHECK(test_smem_vector_elements() == 0); CHECK(test_bitcast_width() == 0); CHECK(test_scalar_widths() == 0); + CHECK(test_scalar_kinds_are_covered() == 0); printf("rocke_ir_types: ok\n"); return 0; } From 04d966a0545fa9fc0f562f3b05fe5ef47b13195c Mon Sep 17 00:00:00 2001 From: Chris Millette Date: Fri, 9 Oct 2026 18:48:34 -0400 Subject: [PATCH 09/11] fix(rocke): lower MMA operands as the instruction's register type Four MMA lowerings printed a fixed register type around whatever value they received: the two f32 MFMAs print `float`, the gfx11 integer WMMA prints `` and the gfx1250 fp8 WMMA prints `<8 x i32>`. A value of the instruction's own format in another shape, such as a one-element `<1 x float>` vector or packed `<16 x i8>` / `<32 x i8>` bytes, produced IR that LLVM rejects. The HIP lowering passed the same values to the builtins unchanged: clang rejects the one-element vector for an f32 MFMA, and accepts the WMMA bytes only through an implicit vector conversion. core/mma_operand.py holds the rule, with no output syntax. mma_operand_needs_reinterpret(op_id, role, operand_type, register_type, arch) returns False for an operand of the register type, without consulting the MMA catalog, and True when one reinterpretation of the same bits gives the register type: the operand holds the register's element type or the element type the target's MMA catalog records for that operand, and has the register's bit width. Any other number format, and any other bit width, raises ValueError naming the operand, the instruction and the register type. mma_operand_value(op, role) picks operand a or b. Each lowerer only renders the answer. In the LLVM lowering, _Lowerer._mma_operand(op, role, register_type) emits one `bitcast` through _Lowerer._reinterpret; _mma_op_id reads the op_id from an MMA op's `tile.` name. In the HIP lowering, _Lowerer._mma_operand(op, op_id, role, register_type) writes `__builtin_bit_cast(, )`. The four lowerings pass their operands through them in both lowerings. Operands that already have the register type lower exactly as before, so no existing kernel's IR or HIP source changes. Tests cover each lowering with different A and B values, the number-format and width errors for both operands with the same text from the LLVM and HIP lowerings, pin the gfx11 integer WMMA signedness and clamp flags, and compile the HIP source with hipcc when it is present. The C++ engine applies the same rule with the same fresh names and error text. rocke_mma_operand_needs_reinterpret and rocke_mma_operand_value (cpp/core/mma_operand.cpp, rocke/mma_operand_internal.h) are the twins of the rule module; mma_operand in lower_llvm/mma.cpp renders through rocke_ll_reinterpret, and the HIP lowering of the two f32 MFMAs renders through h_mma_operand in lower_hip_mma.cpp. rocke_require_ir_type is the raising twin of rocke_dtype_to_ir_type, with dtype_to_ir_type's texts, and rocke_i_vector_type_in_arena builds the register type in the lowerer's arena. Every LLVM test lowers with both engines and requires the same IR, or the same error; the f32 MFMA HIP tests do the same for the HIP source. Co-Authored-By: Claude --- .../platform/cpp/core/ir/core_builder.cpp | 21 +- .../rocke/platform/cpp/core/ir/core_types.cpp | 36 +- .../rocke/platform/cpp/core/isa/backend.cpp | 3 +- .../cpp/core/lower_hip/lower_hip_mma.cpp | 33 +- .../platform/cpp/core/lower_llvm/core.cpp | 14 + .../platform/cpp/core/lower_llvm/mma.cpp | 82 +++- .../rocke/platform/cpp/core/mma_operand.cpp | 108 +++++ .../platform/cpp/include/rocke/ir_internal.h | 21 +- .../platform/cpp/include/rocke/isa_backend.h | 4 +- .../cpp/include/rocke/lower_llvm_internal.h | 7 + .../cpp/include/rocke/mma_operand_internal.h | 47 ++ .../rocke/platform/python/rocke/core/ir.py | 15 +- .../platform/python/rocke/core/isa/backend.py | 33 +- .../platform/python/rocke/core/lower_hip.py | 71 ++- .../platform/python/rocke/core/lower_llvm.py | 55 ++- .../platform/python/rocke/core/mma_operand.py | 101 +++++ .../platform/tests/core/test_ir_types.cpp | 20 + .../core/test_mma_operand_register_type.py | 429 ++++++++++++++++++ 18 files changed, 1020 insertions(+), 80 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/mma_operand.cpp create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/mma_operand_internal.h create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/mma_operand.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_operand_register_type.py diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_builder.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_builder.cpp index 8252ffaeb55e..a636b4bd307b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_builder.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_builder.cpp @@ -646,16 +646,23 @@ int rocke_i_count_of(const rocke_type_t* t) return 1; } +static const rocke_type_t* packed_only_candidate(const rocke_type_t* t) +{ + return (t != NULL && t->kind == ROCKE_TYPE_SMEM) ? t->elem : rocke_i_elem_of(t); +} + +bool rocke_i_is_packed_only(const rocke_type_t* t) +{ + const rocke_type_t* elem = packed_only_candidate(t); + return elem != NULL && elem->kind == ROCKE_TYPE_SCALAR + && (elem->scalar == ROCKE_SCALAR_I4 || elem->scalar == ROCKE_SCALAR_U4); +} + void rocke_i_refuse_packed_only_value(const char* op_name, const rocke_type_t* t) { - const rocke_type_t* elem - = (t != NULL && t->kind == ROCKE_TYPE_SMEM) ? t->elem : rocke_i_elem_of(t); - if(elem == NULL || elem->kind != ROCKE_TYPE_SCALAR) - { - return; - } - if(elem->scalar == ROCKE_SCALAR_I4 || elem->scalar == ROCKE_SCALAR_U4) + if(rocke_i_is_packed_only(t)) { + const rocke_type_t* elem = packed_only_candidate(t); throw ckc::TypeError( ckc::format_error("%s: %s values travel packed in 32-bit words and have no scalar or " "vector form -- keep them in i32 storage and unpack explicitly", 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 f2515c9999e4..2ae36769cf13 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 @@ -136,24 +136,42 @@ const rocke_type_t* rocke_dtype_to_ir_type(const char* dtype) return rocke_scalar_by_name(info->name); } -/* VectorType(elem, count) -> "vec<{elem}x{count}>" */ +/* Mirrors core/ir.py::dtype_to_ir_type, raising where rocke_dtype_to_ir_type + * returns NULL. */ +const rocke_type_t* rocke_require_ir_type(const char* dtype) +{ + const rocke_dtype_info_t* info = rocke_require_dtype_info(dtype); + const rocke_type_t* t = rocke_dtype_to_ir_type(info->name); + if(!t) + throw ckc::ValueError(ckc::format_error("no logical IR type for dtype '%s'", info->name)); + return t; +} + +/* VectorType(elem, count) -> "vec<{elem}x{count}>", allocated in `arena`. */ +const rocke_type_t* + rocke_i_vector_type_in_arena(rocke_arena_t* arena, const rocke_type_t* elem, int count) +{ + rocke_type_t* t = (rocke_type_t*)rocke_arena_calloc(arena, sizeof(*t)); + if(!t) + return NULL; + t->kind = ROCKE_TYPE_VECTOR; + t->elem = elem; + t->count = count; + t->name = rocke_arena_printf(arena, "vec<%sx%d>", elem->name, count); + return t->name ? t : NULL; +} + const rocke_type_t* rocke_vector_type(rocke_ir_builder_t* b, const rocke_type_t* elem, int count) { - rocke_type_t* t; + const rocke_type_t* t; if(!rocke_i_live(b)) return NULL; if(!elem) return (const rocke_type_t*)rocke_i_set_err( b, ROCKE_ERR_VALUE, "vector_type: NULL element type"); - t = (rocke_type_t*)rocke_arena_calloc(&b->arena, sizeof(*t)); + t = rocke_i_vector_type_in_arena(&b->arena, elem, count); if(!t) return (const rocke_type_t*)rocke_i_set_err(b, ROCKE_ERR_OOM, "vector_type: OOM"); - t->kind = ROCKE_TYPE_VECTOR; - t->elem = elem; - t->count = count; - t->name = rocke_arena_printf(&b->arena, "vec<%sx%d>", elem->name, count); - if(!t->name) - return (const rocke_type_t*)rocke_i_set_err(b, ROCKE_ERR_OOM, "vector_type: OOM"); return t; } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/isa/backend.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/isa/backend.cpp index accace5420d6..ea40c71f087b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/isa/backend.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/isa/backend.cpp @@ -596,7 +596,8 @@ int rocke_isa_emit_wmma_int_call(rocke_strbuf_t* out, return -1; } /* signedA=1, signedB=1 (our quantized data is signed), clamp=0 (values stay - * within i32 range). Operands arrive as in SSA, no bitcast. */ + * within i32 range). a_name and b_name already name values; + * the lowerer reinterprets an operand of another shape before the call. */ if(rocke_strbuf_appendf(out, " %s = call <%d x i32> @%s(" "i1 1, <%d x i32> %s, " 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 f5a3af03bc7d..391560c9ff57 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 @@ -34,6 +34,7 @@ #include "rocke/lower_hip.h" #include "rocke/lower_hip_internal.h" #include "rocke/mfma_scale_internal.h" +#include "rocke/mma_operand_internal.h" #include "rocke/tf32_internal.h" #include "rocke/wmma_scale_internal.h" @@ -65,6 +66,29 @@ static const char* h_elem_scalar(rocke_h_lowerer_t* lw, const rocke_type_t* t) /* ============================== mma ====================================== */ +/* MMA operand `role` (A or B) of `op` as the instruction's register type + * (Python _Lowerer._mma_operand): unchanged, or reinterpreted with one + * __builtin_bit_cast, as rocke_mma_operand_needs_reinterpret decides; that rule + * also throws for an operand it refuses. The LLVM lowering renders the same + * decision as a `bitcast`. */ +static const char* h_mma_operand(rocke_h_lowerer_t* lw, + const rocke_op_t* op, + const char* op_id, + rocke_mma_role_t role, + const rocke_type_t* register_type) +{ + const rocke_value_t* value = rocke_mma_operand_value(op, role); + if(!rocke_mma_operand_needs_reinterpret(op_id, role, value->type, register_type, lw->arch.gfx)) + return rocke_h_name(lw, value); + const char* reinterpreted = rocke_arena_printf(&lw->b->arena, + "__builtin_bit_cast(%s, %s)", + rocke_h_type_to_hip(lw, register_type), + rocke_h_name(lw, value)); + if(!reinterpreted) + rocke_h_fail(lw, ROCKE_ERR_OOM, "out of memory formatting an MMA operand"); + return reinterpreted ? reinterpreted : ""; +} + /* def _op_tile_mma(self, op): * op_id = op.attrs["op_id"] * legacy = Op(name=f"tile.{op_id}", operands=list(op.operands), @@ -195,17 +219,20 @@ static rocke_status_t rocke_h_op_tile_mma(rocke_h_lowerer_t* lw, const rocke_op_ rocke_arch_mma_op_physical_op_id(atom)); } } - /* Ordinary FP32 controls use catalog IDs that have no legacy opcode. */ + /* Ordinary FP32 controls use catalog IDs that have no legacy opcode. A and + * B are passed as f32 (h_mma_operand). */ 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; + const char* a_arg = h_mma_operand(lw, op, op_id, ROCKE_MMA_ROLE_A, rocke_f32()); + const char* b_arg = h_mma_operand(lw, op, op_id, ROCKE_MMA_ROLE_B, rocke_f32()); 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]), + a_arg, + b_arg, rocke_h_name(lw, op->operands[2])); return ROCKE_OK; } 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 d85815ee3c57..11a61c9819dc 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 @@ -955,6 +955,20 @@ const char* rocke_ll_operand_with_type(rocke_lower_t* L, const rocke_value_t* v) return rocke_arena_printf(&L->arena, "%s %s", ty, op); } +const char* rocke_ll_reinterpret(rocke_lower_t* L, + const rocke_value_t* v, + const rocke_type_t* register_type) +{ + /* The fresh name is taken before the operand is rendered, as in Python. */ + const char* reinterpreted = rocke_ll_fresh(L, "reg"); + rocke_ll_emitf(L, + " %s = bitcast %s to %s", + reinterpreted, + rocke_ll_operand_with_type(L, v), + rocke_ll_llvm_type(L, register_type)); + return reinterpreted; +} + /* ====================================================================== */ /* Shared binary-op helpers (Python _binop / _vector_binop) */ /* ====================================================================== */ 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 2335c01e4ed5..b4e09d1ff1b6 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 @@ -26,8 +26,10 @@ #include "rocke/arch_target.h" #include "rocke/ir.h" +#include "rocke/ir_internal.h" #include "rocke/lower_llvm.h" #include "rocke/lower_llvm_internal.h" +#include "rocke/mma_operand_internal.h" #include "rocke/wmma_scale_internal.h" namespace ckc @@ -72,6 +74,33 @@ static const char* mma_result_name(rocke_lower_t* L, const rocke_op_t* op) return op->results[0]->name; } +/* The register type `vec`, built in the lowerer's arena. */ +static const rocke_type_t* + mma_register_vector(rocke_lower_t* L, const rocke_type_t* elem, int count) +{ + const rocke_type_t* t = rocke_i_vector_type_in_arena(&L->arena, elem, count); + if(!t) + rocke_ll_fail(L, ROCKE_ERR_OOM, "MMA register type: OOM"); + return t; +} + +/* MMA operand `role` (A or B) of `op` as the instruction's register type + * (Python _Lowerer._mma_operand): unchanged, or reinterpreted with one bitcast, + * as rocke_mma_operand_needs_reinterpret decides; that rule also throws for an + * operand it refuses. */ +static const char* mma_operand(rocke_lower_t* L, + const rocke_op_t* op, + const char* op_id, + rocke_mma_role_t role, + const rocke_type_t* register_type) +{ + const rocke_value_t* value = rocke_mma_operand_value(op, role); + if(rocke_mma_operand_needs_reinterpret( + op_id, role, value->type, register_type, L->backend->gfx)) + return rocke_ll_reinterpret(L, value, register_type); + return rocke_ll_operand(L, value); +} + /* ====================================================================== */ /* Emulated atoms: BF16x3 expansion (Python _expand_emulated_mma) */ /* ====================================================================== */ @@ -766,7 +795,7 @@ static const _wmma_spec_t WMMA_SPECS[] = { "bfloat", 16, WMMA_SHAPE_GFX1250}, - /* _GFX1250_WMMA_FP8 (K=64). A/B arrive in SSA already as <8 x i32>. */ + /* _GFX1250_WMMA_FP8 (K=64). A/B are passed as <8 x i32> (mma_operand). */ {"wmma_gfx1250_f32_16x16x64_fp8_fp8", "wmma.gfx1250.f32.16x16x64.fp8.fp8", "llvm.amdgcn.wmma.f32.16x16x64.fp8.fp8.v8f32.v8i32", @@ -801,8 +830,9 @@ static const int WMMA_SPECS_N = (int)(sizeof(WMMA_SPECS) / sizeof(WMMA_SPECS[0]) /* Integer WMMA spec table (Python _RDNA_WMMA_INT). Integer WMMA differs from * the float path: operands/accumulator are i32 vectors (A/B packed, C/D the i32 * accumulator), and the intrinsic signature carries i1 signedness flags before - * each matrix operand and a trailing i1 clamp. Operands arrive in SSA already - * as so no bitcast is needed. Quantized data is signed and within i32 + * each matrix operand and a trailing i1 clamp. A and B are passed as ; + * packed bytes of the same width (for example <16 x i8>) are reinterpreted with + * one bitcast (mma_operand). Quantized data is signed and within i32 * range, so the flags are (signedA=1, signedB=1, clamp=0). */ typedef struct _wmma_int_spec { @@ -830,14 +860,14 @@ static const int WMMA_INT_SPECS_N = (int)(sizeof(WMMA_INT_SPECS) / sizeof(WMMA_I /* Emit an integer WMMA (iu8/iu4) call (Python _emit_wmma_int). The signature is * (i1 signedA, A, i1 signedB, B, <8 x i32> C, i1 clamp) with * an <8 x i32> result. Both signedness flags are 1 (signed quant data); clamp - * is 0 (values stay within i32 range). No operand bitcast (already ). */ + * is 0 (values stay within i32 range). */ static void _emit_wmma_int(rocke_lower_t* L, const rocke_op_t* op, const _wmma_int_spec_t* spec) { - const rocke_value_t *a, *b, *c; - a = op->operands[0]; - b = op->operands[1]; - c = op->operands[2]; + const rocke_value_t* c = op->operands[2]; rocke_ll_need(L, spec->decl_key); + const rocke_type_t* register_type = mma_register_vector(L, rocke_i32(), spec->op_vec); + const char* a_arg = mma_operand(L, op, spec->op_id, ROCKE_MMA_ROLE_A, register_type); + const char* b_arg = mma_operand(L, op, spec->op_id, ROCKE_MMA_ROLE_B, register_type); rocke_ll_emitf(L, " %s = call <%d x i32> @%s(" "i1 1, <%d x i32> %s, " @@ -847,9 +877,9 @@ static void _emit_wmma_int(rocke_lower_t* L, const rocke_op_t* op, const _wmma_i spec->acc_vec, spec->intrinsic, spec->op_vec, - rocke_ll_operand(L, a), + a_arg, spec->op_vec, - rocke_ll_operand(L, b), + b_arg, spec->acc_vec, rocke_ll_operand(L, c)); } @@ -909,8 +939,19 @@ static void _emit_wmma(rocke_lower_t* L, const rocke_op_t* op, const char* op_id w = spec->frag_width; rocke_ll_need(L, spec->decl_key); - a_arg = rocke_ll_operand(L, a); - b_arg = rocke_ll_operand(L, b); + if(spec->shape == WMMA_SHAPE_GFX1250_FP8) + { + /* A and B are passed as the intrinsic's <8 x i32> (32 low-bit bytes per + * lane); a <32 x i8> fragment is reinterpreted with one bitcast. */ + const rocke_type_t* register_type = mma_register_vector(L, rocke_i32(), w); + a_arg = mma_operand(L, op, op_id, ROCKE_MMA_ROLE_A, register_type); + b_arg = mma_operand(L, op, op_id, ROCKE_MMA_ROLE_B, register_type); + } + else + { + a_arg = rocke_ll_operand(L, a); + b_arg = rocke_ll_operand(L, b); + } if(strcmp(spec->call_elt, spec->ssa_elt) != 0) { @@ -1013,6 +1054,10 @@ typedef struct _mfma_spec const char* ab_ty; /* A/B SSA operand type spelling */ const char* acc_ty; /* accumulator/result vector spelling */ const char* bitcast_to; /* operand bitcast target (NULL = no bitcast) */ + /* The A/B register type when A and B pass through mma_operand (a value of + * another shape is reinterpreted); NULL prints them as is, as Python does + * for every atom but the two f32 ones. */ + const rocke_type_t* (*ab_register)(void); } _mfma_spec_t; static const _mfma_spec_t MFMA_SPECS[] = { @@ -1072,13 +1117,15 @@ static const _mfma_spec_t MFMA_SPECS[] = { "llvm.amdgcn.mfma.f32.16x16x4f32", "float", "<4 x float>", - NULL}, + NULL, + rocke_f32}, {"mfma_f32_32x32x2_f32", "mfma.f32.32x32x2f32", "llvm.amdgcn.mfma.f32.32x32x2f32", "float", "<16 x float>", - NULL}, + NULL, + rocke_f32}, {"mfma_f32_32x32x16_f16", "mfma.f32.32x32x16.f16", "llvm.amdgcn.mfma.f32.32x32x16.f16", @@ -1147,6 +1194,13 @@ static void _emit_mfma(rocke_lower_t* L, const rocke_op_t* op, const _mfma_spec_ b_arg = b_cast; call_ty = spec->bitcast_to; } + else if(spec->ab_register != NULL) + { + const rocke_type_t* register_type = spec->ab_register(); + a_arg = mma_operand(L, op, spec->op_id, ROCKE_MMA_ROLE_A, register_type); + b_arg = mma_operand(L, op, spec->op_id, ROCKE_MMA_ROLE_B, register_type); + call_ty = spec->ab_ty; + } else { a_arg = rocke_ll_operand(L, a); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/mma_operand.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/mma_operand.cpp new file mode 100644 index 000000000000..00eed55ec35d --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/mma_operand.cpp @@ -0,0 +1,108 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* Mirrors core/mma_operand.py: the rule for passing an MMA operand as the + * instruction's register type. The lowerers only render its answer. */ +#include "rocke/mma_operand_internal.h" + +#include +#include + +#include "rocke/error.hpp" +#include "rocke/ir_internal.h" + +/* The multiplicand roles: A and B. The operand rule applies to no other. */ +static bool is_multiplicand(rocke_mma_role_t role) +{ + return role == ROCKE_MMA_ROLE_A || role == ROCKE_MMA_ROLE_B; +} + +const rocke_value_t* rocke_mma_operand_value(const rocke_op_t* op, rocke_mma_role_t role) +{ + if(!is_multiplicand(role)) + throw ckc::ValueError("an MMA operand role is A or B"); + /* A is operand 0 and B operand 1; C, the accumulator, follows them. */ + const int index = role == ROCKE_MMA_ROLE_A ? 0 : 1; + if(op->num_operands <= index) + throw ckc::ValueError(ckc::format_error("%s has %d operands; an MMA op takes A, B and C", + op->name ? op->name : "(unnamed op)", + op->num_operands)); + return op->operands[index]; +} + +/* The element types MMA operand `role` may hold in another shape (Python + * _admitted_elements), sorted by name as Python joins them: the register's own + * element type (a reshape, such as vec for f32) and the element type the + * MMA catalog of `gfx` records for that operand (the instruction's packed + * values, such as vec for vec) unless it is packed-only. Writes + * one or two types; returns how many. */ +static int admitted_elements(const char* op_id, + rocke_mma_role_t role, + const rocke_type_t* register_type, + const char* gfx, + const rocke_type_t* elements[2]) +{ + int count = 0; + elements[count++] = rocke_i_elem_of(register_type); + const rocke_arch_target_t* arch = rocke_arch_target_from_gfx(gfx); + if(!arch) + throw ckc::ValueError( + ckc::format_error("no MMA catalog for target '%s'", gfx ? gfx : "(null)")); + const rocke_mma_op_t* catalog_op = rocke_mma_catalog_by_op_id(&arch->mma, op_id); + if(!catalog_op) + return count; + const rocke_type_t* catalog_element = rocke_require_ir_type( + role == ROCKE_MMA_ROLE_A ? catalog_op->a_dtype : catalog_op->b_dtype); + /* A packed-only catalog element (i4 for the iu4 instructions) is never a + * value's element type, so those operands arrive in the register type. */ + if(rocke_i_is_packed_only(catalog_element)) + return count; + if(!rocke_type_eq(catalog_element, elements[0])) + elements[count++] = catalog_element; + if(count == 2 && strcmp(elements[0]->name, elements[1]->name) > 0) + std::swap(elements[0], elements[1]); + return count; +} + +bool rocke_mma_operand_needs_reinterpret(const char* op_id, + rocke_mma_role_t role, + const rocke_type_t* operand_type, + const rocke_type_t* register_type, + const char* gfx) +{ + if(!is_multiplicand(role)) + throw ckc::ValueError("an MMA operand role is A or B"); + if(rocke_type_eq(operand_type, register_type)) + return false; + /* Python spells the roles "a" and "b". */ + const char* role_name = role == ROCKE_MMA_ROLE_A ? "a" : "b"; + const rocke_type_t* element = rocke_i_elem_of(operand_type); + const rocke_type_t* admitted[2]; + const int num_admitted = admitted_elements(op_id, role, register_type, gfx, admitted); + bool is_admitted = false; + for(int i = 0; i < num_admitted; ++i) + is_admitted = is_admitted || rocke_type_eq(element, admitted[i]); + if(!is_admitted) + throw ckc::ValueError(ckc::format_error( + "MMA operand %s of '%s' holds %s values; the instruction takes %s%s%s -- convert " + "the values first", + role_name, + op_id, + element->name, + admitted[0]->name, + num_admitted == 2 ? " or " : "", + num_admitted == 2 ? admitted[1]->name : "")); + const int operand_bits = rocke_i_type_bits(operand_type); + const int register_bits = rocke_i_type_bits(register_type); + if(operand_bits != register_bits) + throw ckc::ValueError(ckc::format_error( + "MMA operand %s of '%s' is %s (%d bits); the instruction takes %s (%d bits) -- " + "pass a %d-bit value", + role_name, + op_id, + operand_type->name, + operand_bits, + register_type->name, + register_bits, + register_bits)); + return true; +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir_internal.h index 283abc505f36..f56f4876e9b8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/ir_internal.h @@ -136,8 +136,12 @@ bool rocke_i_is_vector(const rocke_type_t* t, const char* elem_name, int count); * `v.type.elem if isinstance(...,VectorType) else v.type`). */ const rocke_type_t* rocke_i_elem_of(const rocke_type_t* t); -/* Throw ckc::TypeError when `t` is a packed-only scalar (i4, u4), a vector of one - * or an LDS tile of one: the twin of core/ir.py::_refuse_packed_only_value, +/* Whether `t`'s elements are packed-only (i4, u4): `t` is such a scalar, a vector + * of one or an LDS tile of one. The twin of core/ir.py::is_packed_only. NULL is + * not. */ +bool rocke_i_is_packed_only(const rocke_type_t* t); + +/* Throw ckc::TypeError when rocke_i_is_packed_only(t): the twin of core/ir.py::_refuse_packed_only_value, * which states the rule and the message. rocke_i_op calls it for every result * and rocke_b_param for every parameter, so no builder makes such a value. * NULL-safe. */ @@ -159,6 +163,19 @@ void rocke_i_refuse_unsigned_operands(const rocke_op_t* op, const char* instead) /* Lane count of a vector type, or 1 for a scalar. */ int rocke_i_count_of(const rocke_type_t* t); +/* VectorType(elem, count) allocated in `arena`, for callers that hold an arena + * but no builder (the lowerers). `elem` must be non-NULL; NULL on OOM. The + * name is the one rocke_vector_type gives the same type. */ +const rocke_type_t* + rocke_i_vector_type_in_arena(rocke_arena_t* arena, const rocke_type_t* elem, int count); + +/* The logical IR type of `dtype` (any accepted spelling), as + * rocke_dtype_to_ir_type returns it (Python core/ir.py::dtype_to_ir_type). + * Throws ckc::ValueError "unknown dtype ''" for an unknown name and + * "no logical IR type for dtype ''" for an encoding with no + * IR type. */ +const rocke_type_t* rocke_require_ir_type(const char* dtype); + /* The number of bits a value of type `t` occupies (Python core/ir.py::type_bits): * a scalar's format width, a vector's element width times its count, a * pointer's address-space width. Throws ckc::TypeError for an LDS allocation diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/isa_backend.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/isa_backend.h index 2a99099acced..a7d4263590ac 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/isa_backend.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/isa_backend.h @@ -230,7 +230,9 @@ int rocke_isa_emit_wmma_call(rocke_strbuf_t* out, const char* b_cast_name); /* Emit the byte-identical text of an integer WMMA (iu8/iu4) call. Mirrors - * Gfx11Backend._emit_wmma_int: signedA=1, signedB=1, clamp=0, no bitcast. + * Gfx11Backend._emit_wmma_int: signedA=1, signedB=1, clamp=0. `a_name` and + * `b_name` must already name values; reinterpreting an operand + * of another shape is the lowerer's job (mma_operand in lower_llvm/mma.cpp). * % = call @( * i1 1, a, i1 1, b, * c, i1 0) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/lower_llvm_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/lower_llvm_internal.h index a8582d87a131..cc439a518758 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/lower_llvm_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/lower_llvm_internal.h @@ -532,6 +532,13 @@ const char* rocke_ll_operand(rocke_lower_t* L, const rocke_value_t* v); /* " " (Python _operand_with_type). */ const char* rocke_ll_operand_with_type(rocke_lower_t* L, const rocke_value_t* v); +/* Emit one `bitcast` of `v` to `register_type` and return the new name + * (Python _Lowerer._reinterpret). Only renders: the caller has established + * that the bit widths match (rocke_mma_operand_needs_reinterpret). */ +const char* rocke_ll_reinterpret(rocke_lower_t* L, + const rocke_value_t* v, + const rocke_type_t* register_type); + /* Map an IR Type to its LLVM textual form (Python _llvm_type). Sets NOTIMPL on * an unmapped type and returns "" . */ const char* rocke_ll_llvm_type(rocke_lower_t* L, const rocke_type_t* t); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/mma_operand_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/mma_operand_internal.h new file mode 100644 index 000000000000..70ccde1fbdb1 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/mma_operand_internal.h @@ -0,0 +1,47 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* How an MMA operand reaches the instruction's register type. NOT a public API. + * + * Mirrors core/mma_operand.py. An MMA instruction takes its A and B operands as + * one fixed register type (f32 for the f32 MFMAs, vec for the gfx11 iu8 + * WMMA). A kernel may hold the same bits in another shape: a one-element + * vec, or the instruction's packed values as vec. Reading such an + * operand as the register type reinterprets the same bits; it never converts. + * + * rocke_mma_operand_needs_reinterpret is the rule, with no output syntax. Each + * lowerer renders its answer in its own language (LLVM `bitcast`, HIP + * `__builtin_bit_cast`), so both lowerers of both engines accept and refuse the + * same operands with the same text. + */ +#ifndef ROCKE_MMA_OPERAND_INTERNAL_H +#define ROCKE_MMA_OPERAND_INTERNAL_H + +#include "rocke/arch_target.h" +#include "rocke/ir.h" + +/* MMA operand `role` of `op` (Python mma_operand_value): operand 0 for + * ROCKE_MMA_ROLE_A, operand 1 for ROCKE_MMA_ROLE_B. Throws ckc::ValueError for + * any other role, or when `op` has fewer than two operands. */ +const rocke_value_t* rocke_mma_operand_value(const rocke_op_t* op, rocke_mma_role_t role); + +/* Whether MMA operand `role` (A or B) of `op_id`, of type `operand_type`, must be + * reinterpreted to be passed as `register_type` on target `gfx` (Python + * mma_operand_needs_reinterpret). + * + * False when the operand already has the register type; the target's MMA + * catalog is not consulted then, so a target without a catalog still lowers + * operands of the register type. True when one reinterpretation of the same + * bits gives the register type: the operand holds the register's element type + * or the element type the target's MMA catalog records for that operand, and + * has the register's bit width. + * + * Throws ckc::ValueError, with the Python text, for an operand holding another + * number format and for an operand of another bit width; and for a mismatched + * operand on a target with no MMA catalog. */ +bool rocke_mma_operand_needs_reinterpret(const char* op_id, + rocke_mma_role_t role, + const rocke_type_t* operand_type, + const rocke_type_t* register_type, + const char* gfx); + +#endif /* ROCKE_MMA_OPERAND_INTERNAL_H */ 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 c9a29bcd2540..17864d41f23a 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 @@ -300,12 +300,19 @@ def _check_flag_ptr(op_name: str, ptr: "Value") -> None: _PACKED_ONLY_SCALARS = ("i4", "u4") +def is_packed_only(t: "Type") -> bool: + """Whether ``t``'s elements are packed-only (i4, u4): ``t`` is such a + scalar, a vector of one, or an LDS tile of one. An untyped placeholder + (``None``) has no elements and is not.""" + elem = t.elem if isinstance(t, (VectorType, SmemType)) else t + return elem is not None and elem.name in _PACKED_ONLY_SCALARS + + def _refuse_packed_only_value(op_name: str, t: "Type") -> None: """Refuse an ``op_name`` result or parameter of type ``t`` whose elements - are packed-only: a scalar, a vector of one, or an LDS tile of one. An - untyped placeholder (``None``) has no elements and passes.""" - elem = t.elem if isinstance(t, (VectorType, SmemType)) else t - if elem is not None and elem.name in _PACKED_ONLY_SCALARS: + are packed-only (:func:`is_packed_only`).""" + if is_packed_only(t): + elem = t.elem if isinstance(t, (VectorType, SmemType)) else t raise TypeError( f"{op_name}: {elem.name} values travel packed in 32-bit words and have " "no scalar or vector form -- keep them in i32 storage and unpack " diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/isa/backend.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/isa/backend.py index 40066fc6bdff..34d68dc6fff4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/isa/backend.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/isa/backend.py @@ -32,6 +32,7 @@ from ..arch import ArchTarget from ..arch.wmma_scale import gfx1250_scaled_wmma +from ..ir import I32, Op, VectorType from .wmma_scale import ScaledWmmaLLVM @@ -133,8 +134,6 @@ def emit_mma(self, lowerer, op) -> None: which calls :meth:`emit_wmma` and correctly raises ``NotImplementedError`` (WMMA is an RDNA-only instruction). """ - from ..ir import Op - op_id = op.attrs["op_id"] from ..arch.mfma_scale import gfx950_scaled_mfma @@ -364,8 +363,6 @@ def emit_mma(self, lowerer, op) -> None: name-keyed lookup keeps working, then emits the WMMA call. This is the RDNA half of the neutral-MMA contract; the CDNA base emits MFMA. """ - from ..ir import Op - op_id = op.attrs["op_id"] from ..arch.mfma_scale import gfx950_scaled_mfma @@ -426,15 +423,16 @@ def _emit_wmma_int(self, lowerer, op, spec) -> None: verified empirically on gfx11-generic (iu8 GEMM probe): passing ``0`` made the unit compute the **unsigned** dot product (all-positive results matching ``A.view(uint8) @ B.view(uint8).T``). Our quantized - data is signed, so both flags are ``1``. Operands arrive as - ```` in SSA (int8/int4 packed into i32), so no bitcast is - needed; values stay within i32 range -> ``clamp = 0`` (exact wrap). + data is signed, so both flags are ``1``. A and B are passed as the + intrinsic's ````; packed bytes of the same width (for example + ``<16 x i8>``) are reinterpreted with one bitcast. Values stay within + i32 range -> ``clamp = 0`` (exact wrap). """ decl_key, intrinsic, op_vec, acc_vec = spec - a, b, c = op.operands + c = op.operands[2] lowerer._need(decl_key) - a_arg = lowerer._operand(a) - b_arg = lowerer._operand(b) + a_arg = lowerer._mma_operand(op, "a", VectorType(I32, op_vec)) + b_arg = lowerer._mma_operand(op, "b", VectorType(I32, op_vec)) c_arg = lowerer._operand(c) lowerer._current().emit( f" {op.result.name} = call <{acc_vec} x i32> @{intrinsic}(" @@ -607,17 +605,20 @@ def emit_wmma(self, lowerer, op) -> None: def _emit_wmma_fp8(self, lowerer, op, spec) -> None: """Emit a gfx1250 K=64 FP8/BF8 WMMA call. - The A/B fragments arrive in SSA as ``<8 x i32>`` (32 low-bit bytes per - lane). The 6-operand form is ``(A, B, i16 fmt, C, i1, i1)`` with the - format / reuse immediates pinned to 0 (plain unscaled MMA). + A and B are passed as the intrinsic's ``<8 x i32>`` (32 low-bit bytes + per lane); a ``<32 x i8>`` fragment is reinterpreted with one bitcast. + The 6-operand form is ``(A, B, i16 fmt, C, i1, i1)`` with the format / + reuse immediates pinned to 0 (plain unscaled MMA). """ decl_key, intrinsic = spec - a, b, c = op.operands + c = op.operands[2] lowerer._need(decl_key) + a_arg = lowerer._mma_operand(op, "a", VectorType(I32, 8)) + b_arg = lowerer._mma_operand(op, "b", VectorType(I32, 8)) lowerer._current().emit( f" {op.result.name} = call <8 x float> @{intrinsic}(" - f"<8 x i32> {lowerer._operand(a)}, " - f"<8 x i32> {lowerer._operand(b)}, " + f"<8 x i32> {a_arg}, " + f"<8 x i32> {b_arg}, " f"i16 0, <8 x float> {lowerer._operand(c)}, " f"i1 false, i1 false)" ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_hip.py index e516ae443810..47ed1c553c24 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 @@ -23,6 +23,8 @@ _FENCE_ORDERINGS, _FENCE_SCOPES, _FLAG_STORE_ORDERINGS, + F32, + I32, KernelDef, Op, PtrType, @@ -37,6 +39,11 @@ dpp_xor_ctrl, require_streaming_arch, ) +from .mma_operand import ( + MmaOperandRole, + mma_operand_needs_reinterpret, + mma_operand_value, +) # No packed-only type has an entry: core/ir.py _refuse_packed_only_value @@ -720,18 +727,40 @@ def _op_tile_mma(self, op: Op) -> None: ) self.lower_op(legacy) + def _mma_operand( + self, op: Op, op_id: str, role: MmaOperandRole, register_type: Type + ) -> str: + """Return MMA operand ``role`` of ``op`` as the instruction's register type. + + The operand passes through unchanged, or is reinterpreted with one + ``__builtin_bit_cast``, as + :func:`~rocke.core.mma_operand.mma_operand_needs_reinterpret` decides; + that rule also raises for an operand it refuses. The LLVM lowering + renders the same decision as a ``bitcast``. + """ + value = mma_operand_value(op, role) + if mma_operand_needs_reinterpret( + op_id, role, value.type, register_type, self.arch + ): + return f"__builtin_bit_cast({_type_to_hip(register_type)}, {_name(value)})" + return _name(value) + def _op_tile_mfma_f32_16x16x4_f32(self, op: Op) -> None: - a, b, c = op.operands + op_id = "mfma_f32_16x16x4_f32" + a = self._mma_operand(op, op_id, "a", F32) + b = self._mma_operand(op, op_id, "b", F32) self._emit( f"f32x4 {_name(op.result)} = __builtin_amdgcn_mfma_f32_16x16x4f32(" - f"{_name(a)}, {_name(b)}, {_name(c)}, 0, 0, 0);" + f"{a}, {b}, {_name(op.operands[2])}, 0, 0, 0);" ) def _op_tile_mfma_f32_32x32x2_f32(self, op: Op) -> None: - a, b, c = op.operands + op_id = "mfma_f32_32x32x2_f32" + a = self._mma_operand(op, op_id, "a", F32) + b = self._mma_operand(op, op_id, "b", F32) self._emit( f"f32x16 {_name(op.result)} = __builtin_amdgcn_mfma_f32_32x32x2f32(" - f"{_name(a)}, {_name(b)}, {_name(c)}, 0, 0, 0);" + f"{a}, {b}, {_name(op.operands[2])}, 0, 0, 0);" ) def _op_tile_mfma_f32_16x16x16_f16(self, op: Op) -> None: @@ -833,14 +862,16 @@ def _op_tile_wmma_gfx1250_f32_16x16x64_bf8_bf8(self, op: Op) -> None: self._emit_wmma_gfx1250_fp8(op, "bf8_bf8") def _emit_wmma_gfx1250_fp8(self, op: Op, ab: str) -> None: - # gfx1250 K=64 FP8/BF8 builtin: A/B are <8 x i32> (32 low-bit + # gfx1250 K=64 FP8/BF8 builtin: A/B are passed as i32x8 (32 low-bit # bytes per lane), 6-operand form (A, B, fmt, C, reuseA, reuseB). - self._require_wmma_arch(f"wmma_gfx1250_f32_16x16x64_{ab}") - a, b, c = op.operands + op_id = f"wmma_gfx1250_f32_16x16x64_{ab}" + self._require_wmma_arch(op_id) + a = self._mma_operand(op, op_id, "a", VectorType(I32, 8)) + b = self._mma_operand(op, op_id, "b", VectorType(I32, 8)) self._emit( f"f32x8 {_name(op.result)} = " f"__builtin_amdgcn_wmma_f32_16x16x64_{ab}(" - f"{_name(a)}, {_name(b)}, (int16_t)0, {_name(c)}, false, false);" + f"{a}, {b}, (int16_t)0, {_name(op.operands[2])}, false, false);" ) def _emit_wmma_gfx1250_scaled(self, op: Op) -> None: @@ -878,19 +909,23 @@ def _op_tile_wmma_gfx1250_f32_16x16x32_bf16(self, op: Op) -> None: # # The integer twin of the f16 WMMA handler. Per-lane A/B operands are # ``i32x4`` (16 int8 packed 4-per-i32) and the accumulator / result are - # ``i32x8``. The builtin's two leading bool flags select each operand's - # *signedness* (``true`` => signed; verified on gfx11 — passing ``false`` - # makes the unit compute the *unsigned* dot product), and the trailing bool + # ``i32x8``; packed bytes of the same width (``i8x16``) are passed through + # one ``__builtin_bit_cast`` (``_mma_operand``). The builtin's two leading + # bool flags select each operand's *signedness* (``true`` => signed; + # verified on gfx11 -- passing ``false`` makes the unit compute the + # *unsigned* dot product), and the trailing bool # is the clamp flag (``false`` => exact i32 wrap). This mirrors the verified # LLVM-IR emission in ``core/isa/backend.py`` (``_emit_wmma_int``): # ``llvm.amdgcn.wmma.i32.16x16x16.iu8(i1 1, A, i1 1, B, C, i1 0)``. def _op_tile_wmma_i32_16x16x16_iu8(self, op: Op) -> None: - self._require_wmma_arch("wmma_i32_16x16x16_iu8") - a, b, c = op.operands + op_id = "wmma_i32_16x16x16_iu8" + self._require_wmma_arch(op_id) + a = self._mma_operand(op, op_id, "a", VectorType(I32, 4)) + b = self._mma_operand(op, op_id, "b", VectorType(I32, 4)) self._emit( f"i32x8 {_name(op.result)} = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(" - f"true, {_name(a)}, true, {_name(b)}, {_name(c)}, false);" + f"true, {a}, true, {b}, {_name(op.operands[2])}, false);" ) # ---- WMMA iu4 (RDNA3/3.5, gfx11, wave32) — int4 matrix engine ---- @@ -901,11 +936,13 @@ def _op_tile_wmma_i32_16x16x16_iu8(self, op: Op) -> None: # trailing bool is the (off) clamp. Mirrors the LLVM emission # ``llvm.amdgcn.wmma.i32.16x16x16.iu4(i1 1, A, i1 1, B, C, i1 0)``. def _op_tile_wmma_i32_16x16x16_iu4(self, op: Op) -> None: - self._require_wmma_arch("wmma_i32_16x16x16_iu4") - a, b, c = op.operands + op_id = "wmma_i32_16x16x16_iu4" + self._require_wmma_arch(op_id) + a = self._mma_operand(op, op_id, "a", VectorType(I32, 2)) + b = self._mma_operand(op, op_id, "b", VectorType(I32, 2)) self._emit( f"i32x8 {_name(op.result)} = __builtin_amdgcn_wmma_i32_16x16x16_iu4_w32(" - f"true, {_name(a)}, true, {_name(b)}, {_name(c)}, false);" + f"true, {a}, true, {b}, {_name(op.operands[2])}, false);" ) def _require_wmma_arch(self, op_id: str) -> None: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py index fe979e660a2f..50c79ba27b40 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 @@ -69,6 +69,11 @@ require_streaming_arch, split_loc, ) +from .mma_operand import ( + MmaOperandRole, + mma_operand_needs_reinterpret, + mma_operand_value, +) # Datalayout / triple. Copied verbatim from clang's output for the same @@ -1690,6 +1695,13 @@ def render(self) -> List[str]: _BF16_SIGN_BIT = -0x8000 +def _mma_op_id(op: Op) -> str: + """The op_id of an MMA op; MMA ops reach their handlers as ``tile.``.""" + if not op.name.startswith("tile."): + raise ValueError(f"{op.name!r} is not an MMA op named 'tile.'") + return op.name.removeprefix("tile.") + + class _Lowerer: def __init__( self, @@ -1790,6 +1802,33 @@ def _operand(self, v: Value) -> str: def _operand_with_type(self, v: Value) -> str: return f"{_llvm_type(v.type)} {self._operand(v)}" + def _reinterpret(self, v: Value, register_type: Type) -> str: + """Emit one ``bitcast`` of ``v`` to ``register_type``; return its name. + + Only renders: the caller has established that the bit widths match + (:func:`~rocke.core.mma_operand.mma_operand_needs_reinterpret`). + """ + reinterpreted = self._fresh("reg") + self._current().emit( + f" {reinterpreted} = bitcast {self._operand_with_type(v)} " + f"to {_llvm_type(register_type)}" + ) + return reinterpreted + + def _mma_operand(self, op: Op, role: MmaOperandRole, register_type: Type) -> str: + """Return MMA operand ``role`` of ``op`` as the instruction's register type. + + The operand passes through unchanged, or is reinterpreted with one + ``bitcast``, as :func:`~rocke.core.mma_operand.mma_operand_needs_reinterpret` + decides; that rule also raises for an operand it refuses. + """ + value = mma_operand_value(op, role) + if mma_operand_needs_reinterpret( + _mma_op_id(op), role, value.type, register_type, self._backend.arch + ): + return self._reinterpret(value, register_type) + return self._operand(value) + @staticmethod def _composed_constant(value: int | float, ty: Type, ity: str) -> Value: """Make an inline constant for a lowering-time operation expansion.""" @@ -3973,23 +4012,27 @@ def _emit_xf32(self, op: Op, intrinsic: str, count: int) -> None: ) def _op_tile_mfma_f32_16x16x4_f32(self, op: Op) -> None: - a, b, c = op.operands + c = op.operands[2] self._need("mfma.f32.16x16x4f32") + a_arg = self._mma_operand(op, "a", F32) + b_arg = self._mma_operand(op, "b", F32) self._current().emit( f" {op.result.name} = call <4 x float> @llvm.amdgcn.mfma.f32.16x16x4f32(" - f"float {self._operand(a)}, " - f"float {self._operand(b)}, " + f"float {a_arg}, " + f"float {b_arg}, " f"<4 x float> {self._operand(c)}, " f"i32 0, i32 0, i32 0)" ) def _op_tile_mfma_f32_32x32x2_f32(self, op: Op) -> None: - a, b, c = op.operands + c = op.operands[2] self._need("mfma.f32.32x32x2f32") + a_arg = self._mma_operand(op, "a", F32) + b_arg = self._mma_operand(op, "b", F32) self._current().emit( f" {op.result.name} = call <16 x float> @llvm.amdgcn.mfma.f32.32x32x2f32(" - f"float {self._operand(a)}, " - f"float {self._operand(b)}, " + f"float {a_arg}, " + f"float {b_arg}, " f"<16 x float> {self._operand(c)}, " f"i32 0, i32 0, i32 0)" ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/mma_operand.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/mma_operand.py new file mode 100644 index 000000000000..3d7b55a900a4 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/mma_operand.py @@ -0,0 +1,101 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""How an MMA operand reaches the instruction's register type. + +An MMA instruction takes its A and B operands as one fixed register type, such +as ``f32`` for the f32 MFMAs or ``vec`` for the gfx11 iu8 WMMA. A kernel +may hold the same bits in another shape: a one-element ``vec``, or the +instruction's packed values as ``vec``. Reading such an operand as the +register type is a reinterpretation of the same bits, never a conversion. + +:func:`mma_operand_needs_reinterpret` is the rule; it has no output syntax. +Each lowerer renders its answer in its own language (LLVM ``bitcast``, HIP +``__builtin_bit_cast``), so both lowerers of both engines accept and refuse the +same operands with the same text. The C++ engine's twin is +``rocke_mma_operand_needs_reinterpret`` in ``cpp/core/mma_operand.cpp``. +""" + +from __future__ import annotations + +from typing import Literal + +from .arch import ArchTarget +from .ir import Op, Type, Value, VectorType, dtype_to_ir_type, is_packed_only, type_bits + +MmaOperandRole = Literal["a", "b"] + +# Position of each multiplicand among an MMA op's operands (A, B, C). +_OPERAND_INDEX = {"a": 0, "b": 1} + + +def mma_operand_value(op: Op, role: MmaOperandRole) -> Value: + """MMA operand ``role`` ("a" or "b") of ``op``.""" + return op.operands[_OPERAND_INDEX[role]] + + +def _element_type(t: Type) -> Type: + """The element type of a vector type, or the type itself for a scalar.""" + return t.elem if isinstance(t, VectorType) else t + + +def _admitted_elements( + op_id: str, role: MmaOperandRole, register_type: Type, arch: ArchTarget +) -> frozenset[Type]: + """The element types MMA operand ``role`` may hold in another shape. + + They are the register's own element type (a reshape, such as ``vec`` + for ``f32``) and the element type ``arch``'s MMA catalog records for that + operand (the instruction's packed values, such as ``vec`` for + ``vec``). An op_id the catalog does not list admits only the first, + and so does a catalog element that is packed-only (``i4`` for the iu4 + instructions): no value has that element type, so its operands arrive in + the register type. + """ + elements = {_element_type(register_type)} + catalog_op = arch.mma.by_op_id(op_id) + if catalog_op is not None: + dtype = catalog_op.a_dtype if role == "a" else catalog_op.b_dtype + catalog_element = dtype_to_ir_type(dtype) + if not is_packed_only(catalog_element): + elements.add(catalog_element) + return frozenset(elements) + + +def mma_operand_needs_reinterpret( + op_id: str, + role: MmaOperandRole, + operand_type: Type, + register_type: Type, + arch: ArchTarget, +) -> bool: + """Whether MMA operand ``role`` of ``op_id`` must be reinterpreted to be passed. + + False when ``operand_type`` already is ``register_type``; the catalog is not + consulted then, so a target without an MMA catalog still lowers operands of + the register type. True when one reinterpretation of the same bits gives the + register type: the operand holds an admitted element type (see + :func:`_admitted_elements`) and has the register's bit width. + + Raises ``ValueError`` for an operand holding another number format (its + values would have to be converted, which is the kernel's decision, not the + lowerer's) and for an operand of another bit width. + """ + if operand_type == register_type: + return False + label = f"MMA operand {role} of {op_id!r}" + element = _element_type(operand_type) + admitted = _admitted_elements(op_id, role, register_type, arch) + if element not in admitted: + names = " or ".join(sorted(t.name for t in admitted)) + raise ValueError( + f"{label} holds {element.name} values; the instruction takes {names} " + "-- convert the values first" + ) + operand_bits, register_bits = type_bits(operand_type), type_bits(register_type) + if operand_bits != register_bits: + raise ValueError( + f"{label} is {operand_type.name} ({operand_bits} bits); the instruction " + f"takes {register_type.name} ({register_bits} bits) -- pass a " + f"{register_bits}-bit value" + ) + return True diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp index f822b523395a..4ceadaa6b62b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp @@ -613,6 +613,25 @@ static int test_scalar_widths() return 0; } +/* A dtype's logical IR type, as core/ir.py::dtype_to_ir_type gives it: any + * accepted spelling resolves, and an unknown name is refused with + * core/dtypes.py's text. */ +static int test_ir_type_of_dtype() +{ + CHECK(rocke_require_ir_type("half") == rocke_f16()); + CHECK(rocke_require_ir_type("fp8") == rocke_fp8e4m3()); + try + { + rocke_require_ir_type("custom_format"); + CHECK(false); + } + catch(const ckc::ValueError& error) + { + CHECK(strcmp(error.what(), "unknown dtype 'custom_format'") == 0); + } + return 0; +} + /* Every IR scalar type by name: one per rocke_scalar_kind_t, which the * static_assert and test_scalar_kinds_are_covered hold to. Widths are not * listed: they come from the engine. */ @@ -662,6 +681,7 @@ int main(int argc, char** argv) CHECK(test_smem_vector_elements() == 0); CHECK(test_bitcast_width() == 0); CHECK(test_scalar_widths() == 0); + CHECK(test_ir_type_of_dtype() == 0); CHECK(test_scalar_kinds_are_covered() == 0); printf("rocke_ir_types: ok\n"); return 0; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_operand_register_type.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_operand_register_type.py new file mode 100644 index 000000000000..cc6d1a4582c9 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_operand_register_type.py @@ -0,0 +1,429 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""MMA operands are passed as the register type the instruction takes. + +An operand of the register type passes through unchanged. An operand in another +shape is reinterpreted once -- an LLVM ``bitcast``, a HIP ``__builtin_bit_cast`` -- +when it holds the register's own element type (a one-element vector) or the +operand's element type in the target's MMA catalog (the instruction's packed +bytes). An operand holding any other number format, or of another bit width, +fails while lowering, naming the operand and the instruction, with the same text +in the LLVM and HIP lowerings. + +Every LLVM case lowers with both engines and requires the same IR, or the same +error. The HIP cases do the same for the f32 MFMAs; the C++ HIP lowerer has no +integer or gfx1250 fp8 WMMA lowering, so those cases run the Python HIP lowering +and require the LLVM lowering's error text. +""" + +import re +import shutil +import subprocess + +import pytest +from rocke.core.arch import ArchTarget +from rocke.core.ir import ( + BF8E5M2, + F16, + F32, + FP8E4M3, + I8, + I32, + IRBuilder, + KernelDef, + PtrType, +) +from rocke.core.lower_hip import lower_kernel_to_hip +from rocke.tests.engines import ( + assert_both_engines_reject, + assert_both_engines_reject_hip, + lower_hip_with_both_engines, + lower_with_both_engines, +) +from rocke.tests.ir_validity import assert_target_accepts + + +def _mma_kernel( + arch, op_id, a_operand, b_operand, accumulator, result_type +) -> KernelDef: + """D = A x B + C for one MMA; each operand is built by its callable.""" + op = ArchTarget.from_gfx(arch).mma.by_op_id(op_id) + b = IRBuilder(f"mma_operand_{op_id}") + out = b.param("D", PtrType(result_type, "global")) + result = b.mma(op, a_operand(b), b_operand(b), accumulator(b)) + b.global_store(out, b.thread_id_x(), b.vec_extract(result, 0)) + return b.kernel + + +def _lower_mma(arch, op_id, a_operand, b_operand, accumulator, result_type) -> str: + """The LLVM IR both engines emit for :func:`_mma_kernel`.""" + kernel = _mma_kernel(arch, op_id, a_operand, b_operand, accumulator, result_type) + return lower_with_both_engines(kernel, arch) + + +def _mma_call(text: str) -> str: + return next( + line for line in text.splitlines() if "= call" in line and "amdgcn" in line + ) + + +def _f32_scalar(b): + return b.const_f32(1.0) + + +def _f32_one_element_vector(b): + return b.vec_pack([b.const_f32(2.0)], F32) + + +def _zeros(elem, count): + return lambda b: b.zero_vec(elem, count) + + +_F32_MFMAS = [("mfma_f32_16x16x4_f32", 4), ("mfma_f32_32x32x2_f32", 16)] + + +@pytest.mark.parametrize(("op_id", "accumulator_count"), _F32_MFMAS) +@pytest.mark.parametrize("arch", ["gfx942", "gfx950"]) +def test_an_f32_mfma_takes_scalars_unchanged(arch, op_id, accumulator_count): + text = _lower_mma( + arch, + op_id, + _f32_scalar, + _f32_scalar, + _zeros(F32, accumulator_count), + F32, + ) + assert "bitcast" not in text + assert_target_accepts(text, arch) + + +@pytest.mark.parametrize(("op_id", "accumulator_count"), _F32_MFMAS) +@pytest.mark.parametrize("arch", ["gfx942", "gfx950"]) +def test_an_f32_mfma_reinterprets_only_the_one_element_vector( + arch, op_id, accumulator_count +): + text = _lower_mma( + arch, + op_id, + _f32_one_element_vector, + _f32_scalar, + _zeros(F32, accumulator_count), + F32, + ) + (reinterpreted,) = re.findall( + r"(%reg\.\d+) = bitcast <1 x float> \S+ to float", text + ) + assert re.search( + rf"\(float {re.escape(reinterpreted)}, float 0x3FF0000000000000,", + _mma_call(text), + ) + assert_target_accepts(text, arch) + + +def test_integer_wmma_takes_packed_words_unchanged(): + text = _lower_mma( + "gfx1151", + "wmma_i32_16x16x16_iu8", + _zeros(I32, 4), + _zeros(I32, 4), + _zeros(I32, 8), + I32, + ) + assert "bitcast" not in text + assert_target_accepts(text, "gfx1151") + + +def test_integer_wmma_reinterprets_only_the_packed_byte_operand(): + text = _lower_mma( + "gfx1151", + "wmma_i32_16x16x16_iu8", + _zeros(I8, 16), + _zeros(I32, 4), + _zeros(I32, 8), + I32, + ) + (reinterpreted,) = re.findall( + r"(%reg\.\d+) = bitcast <16 x i8> \S+ to <4 x i32>", text + ) + assert re.search( + rf"\(i1 1, <4 x i32> {re.escape(reinterpreted)}, i1 1, <4 x i32> %cz\d+,", + _mma_call(text), + ) + assert_target_accepts(text, "gfx1151") + + +@pytest.mark.parametrize( + ("op_id", "a_operand"), + [ + ("wmma_i32_16x16x16_iu8", _zeros(I32, 4)), + ("wmma_i32_16x16x16_iu8", _zeros(I8, 16)), + ("wmma_i32_16x16x16_iu4", _zeros(I32, 2)), + ], +) +def test_integer_wmma_marks_both_operands_signed_without_clamping(op_id, a_operand): + words = 4 if op_id.endswith("iu8") else 2 + text = _lower_mma( + "gfx1151", op_id, a_operand, _zeros(I32, words), _zeros(I32, 8), I32 + ) + assert re.search( + rf"\(i1 1, <{words} x i32> \S+, i1 1, <{words} x i32> \S+, " + r"<8 x i32> \S+, i1 0\)", + _mma_call(text), + ) + + +def test_gfx1250_fp8_wmma_reinterprets_each_operand_as_its_catalog_format(): + text = _lower_mma( + "gfx1250", + "wmma_gfx1250_f32_16x16x64_fp8_bf8", + _zeros(FP8E4M3, 32), + _zeros(BF8E5M2, 32), + _zeros(F32, 8), + F32, + ) + assert len(re.findall(r"bitcast <32 x i8> \S+ to <8 x i32>", text)) == 2 + assert_target_accepts(text, "gfx1250") + + +@pytest.mark.parametrize( + ("arch", "op_id", "a_operand", "b_operand", "accumulator", "message"), + [ + ( + "gfx942", + "mfma_f32_16x16x4_f32", + lambda b: b.const_i32(7), + _f32_scalar, + _zeros(F32, 4), + ( + "MMA operand a of 'mfma_f32_16x16x4_f32' holds i32 values; the " + "instruction takes f32 -- convert the values first" + ), + ), + ( + "gfx942", + "mfma_f32_16x16x4_f32", + _f32_scalar, + _zeros(F16, 2), + _zeros(F32, 4), + ( + "MMA operand b of 'mfma_f32_16x16x4_f32' holds f16 values; the " + "instruction takes f32 -- convert the values first" + ), + ), + ( + "gfx1151", + "wmma_i32_16x16x16_iu8", + _zeros(F16, 8), + _zeros(I32, 4), + _zeros(I32, 8), + ( + "MMA operand a of 'wmma_i32_16x16x16_iu8' holds f16 values; the " + "instruction takes i32 or i8 -- convert the values first" + ), + ), + ( + "gfx1151", + "wmma_i32_16x16x16_iu4", + _zeros(I8, 8), + _zeros(I32, 2), + _zeros(I32, 8), + ( + "MMA operand a of 'wmma_i32_16x16x16_iu4' holds i8 values; the " + "instruction takes i32 -- convert the values first" + ), + ), + ( + "gfx1250", + "wmma_gfx1250_f32_16x16x64_fp8_bf8", + _zeros(FP8E4M3, 32), + _zeros(FP8E4M3, 32), + _zeros(F32, 8), + ( + "MMA operand b of 'wmma_gfx1250_f32_16x16x64_fp8_bf8' holds fp8e4m3 " + "values; the instruction takes bf8e5m2 or i32 -- convert the values first" + ), + ), + ], + ids=["i32-into-f32", "f16-into-f32", "f16-into-iu8", "i8-into-iu4", "fp8-into-bf8"], +) +def test_an_operand_of_another_number_format_fails( + arch, op_id, a_operand, b_operand, accumulator, message +): + kernel = _mma_kernel(arch, op_id, a_operand, b_operand, accumulator, F32) + assert_both_engines_reject(kernel, arch, ValueError, message) + _assert_hip_rejects(kernel, arch, op_id, message) + + +@pytest.mark.parametrize("role", ["a", "b"]) +def test_an_operand_of_another_width_fails_naming_its_role(role): + narrow, words = _zeros(I8, 8), _zeros(I32, 4) + message = ( + f"MMA operand {role} of 'wmma_i32_16x16x16_iu8' is vec (64 bits); the " + "instruction takes vec (128 bits) -- pass a 128-bit value" + ) + kernel = _mma_kernel( + "gfx1151", + "wmma_i32_16x16x16_iu8", + narrow if role == "a" else words, + narrow if role == "b" else words, + _zeros(I32, 8), + I32, + ) + assert_both_engines_reject(kernel, "gfx1151", ValueError, message) + _assert_hip_rejects(kernel, "gfx1151", "wmma_i32_16x16x16_iu8", message) + + +def test_an_f32_operand_of_another_width_fails_in_every_lowering(): + message = ( + "MMA operand a of 'mfma_f32_16x16x4_f32' is vec (64 bits); the " + "instruction takes f32 (32 bits) -- pass a 32-bit value" + ) + kernel = _mma_kernel( + "gfx942", + "mfma_f32_16x16x4_f32", + _zeros(F32, 2), + _f32_scalar, + _zeros(F32, 4), + F32, + ) + assert_both_engines_reject(kernel, "gfx942", ValueError, message) + _assert_hip_rejects(kernel, "gfx942", "mfma_f32_16x16x4_f32", message) + + +# ------------------------------------------------------------------- HIP source + +# The MMAs the C++ HIP lowerer lowers among those this file covers. +_NATIVE_HIP_MMAS = frozenset(op_id for op_id, _ in _F32_MFMAS) + + +def _assert_hip_rejects(kernel: KernelDef, arch: str, op_id: str, message: str) -> None: + """The HIP lowering refuses ``kernel`` with the LLVM lowering's ``message``.""" + if op_id in _NATIVE_HIP_MMAS: + assert_both_engines_reject_hip(kernel, arch, ValueError, message) + return + with pytest.raises(ValueError) as error: + lower_kernel_to_hip(kernel, arch=arch) + assert str(error.value) == message + + +def _loaded(elem, count=None): + """An operand built from one value loaded from its own global buffer. + + A scalar when ``count`` is None, else that value in every one of ``count`` + lanes. The C++ HIP lowering spells float literals differently from the + Python one, so the HIP parity kernels take every value from memory instead. + """ + + def build(b): + source = b.param(f"S{len(b.kernel.params)}", PtrType(elem, "global")) + value = b.global_load(source, b.thread_id_x(), elem) + return value if count is None else b.vector_splat(value, count) + + return build + + +def _mma_builtin_call(source: str) -> str: + return next( + line + for line in source.splitlines() + if "__builtin_amdgcn_" in line and ("mfma" in line or "wmma" in line) + ) + + +def _assert_hipcc_accepts(source: str, arch: str, tmp_path) -> None: + """Compile ``source`` for ``arch`` with hipcc, when hipcc is present.""" + if not shutil.which("hipcc"): + pytest.skip("hipcc not in PATH") + path = tmp_path / "kernel.hip" + path.write_text(source) + subprocess.run( + [ + "hipcc", + f"--offload-arch={arch}", + "--cuda-device-only", + "-fsyntax-only", + str(path), + ], + check=True, + capture_output=True, + text=True, + ) + + +@pytest.mark.parametrize(("op_id", "accumulator_count"), _F32_MFMAS) +@pytest.mark.parametrize("arch", ["gfx942", "gfx950"]) +def test_hip_f32_mfma_takes_scalars_unchanged(arch, op_id, accumulator_count, tmp_path): + kernel = _mma_kernel( + arch, + op_id, + _loaded(F32), + _loaded(F32), + _loaded(F32, accumulator_count), + F32, + ) + source = lower_hip_with_both_engines(kernel, arch) + assert "__builtin_bit_cast" not in _mma_builtin_call(source) + _assert_hipcc_accepts(source, arch, tmp_path) + + +@pytest.mark.parametrize(("op_id", "accumulator_count"), _F32_MFMAS) +@pytest.mark.parametrize("arch", ["gfx942", "gfx950"]) +def test_hip_f32_mfma_reinterprets_only_the_one_element_vector( + arch, op_id, accumulator_count, tmp_path +): + kernel = _mma_kernel( + arch, + op_id, + _loaded(F32), + _loaded(F32, 1), + _loaded(F32, accumulator_count), + F32, + ) + source = lower_hip_with_both_engines(kernel, arch) + call = _mma_builtin_call(source) + assert re.search(r"\(\w+, __builtin_bit_cast\(float, \w+\), \w+, 0, 0, 0\);", call) + assert call.count("__builtin_bit_cast") == 1 + _assert_hipcc_accepts(source, arch, tmp_path) + + +def test_hip_integer_wmma_reinterprets_only_the_packed_byte_operand(tmp_path): + kernel = _mma_kernel( + "gfx1151", + "wmma_i32_16x16x16_iu8", + _zeros(I8, 16), + _zeros(I32, 4), + _zeros(I32, 8), + I32, + ) + source = lower_kernel_to_hip(kernel, arch="gfx1151") + assert re.search( + r"\(true, __builtin_bit_cast\(i32x4, \w+\), true, \w+, \w+, false\);", + _mma_builtin_call(source), + ) + _assert_hipcc_accepts(source, "gfx1151", tmp_path) + + +def test_hip_integer_wmma_takes_packed_words_unchanged(): + kernel = _mma_kernel( + "gfx1151", + "wmma_i32_16x16x16_iu4", + _zeros(I32, 2), + _zeros(I32, 2), + _zeros(I32, 8), + I32, + ) + call = _mma_builtin_call(lower_kernel_to_hip(kernel, arch="gfx1151")) + assert "__builtin_bit_cast" not in call + + +def test_hip_gfx1250_fp8_wmma_reinterprets_each_operand_as_its_catalog_format(): + kernel = _mma_kernel( + "gfx1250", + "wmma_gfx1250_f32_16x16x64_fp8_bf8", + _zeros(FP8E4M3, 32), + _zeros(BF8E5M2, 32), + _zeros(F32, 8), + F32, + ) + call = _mma_builtin_call(lower_kernel_to_hip(kernel, arch="gfx1250")) + assert len(re.findall(r"__builtin_bit_cast\(i32x8, \w+\)", call)) == 2 From 962dab5adb071ee53b3bbb169efd6f0ebb1e8fc8 Mon Sep 17 00:00:00 2001 From: Chris Millette Date: Fri, 9 Oct 2026 18:48:53 -0400 Subject: [PATCH 10/11] fix(rocke): type MMA results by the instruction's accumulator dtype IRBuilder.mma chose the result element as i32 for integer accumulators and f32 for everything else, so an instruction with any other accumulator type would have been given the wrong result type. _mma_c_dtype(op_id) replaces _mma_c_is_int and returns the accumulator dtype the arch catalog records for an op_id (fp32 for op_ids the catalog does not list). IRBuilder.mma types its result as dtype_to_ir_type of that dtype, or of the MmaOp's own c_dtype when it is given one. Every catalog instruction today accumulates in fp32 or i32, so no kernel's IR changes. Tests check that every catalog accumulator has an IR type, that an op_id without a catalog row accumulates in fp32, and that IRBuilder.mma results take the accumulator dtype for an fp32 and an i32 instruction. The IRBuilder.mma docstring describes the result type. The C++ rocke_b_mma types its result the same way: rocke_mma_c_dtype replaces rocke_mma_is_int_acc, and rocke_require_ir_type resolves the dtype, so a refusal has core/ir.py's text. It now checks the operands before the op_id, in IRBuilder.mma's order (scaled MFMA, then XF32), so both builders refuse the same call with the same text. The scaled MFMA check gives one refusal for a wrong operand count or type, as mfma_scale.py does, instead of a separate "requires five operands". The rocke_ir_types test gains --mma , which builds one MMA and prints its result type or the refusal; test_mma_frag_tables.py compares it with the Python builder for every op_id in the catalog and the fragment table, and checks that both refuse an unknown op_id. Its self-check pins that a refusal builds no op. The rocke_mma_frag_ssot test adds an op_id with no catalog row and checks every catalog accumulator. Co-Authored-By: Claude --- .../rocke/platform/cpp/core/ir/ir_tile.cpp | 56 ++++++------ .../cpp/include/rocke/mfma_scale_internal.h | 9 +- .../rocke/platform/python/rocke/core/ir.py | 27 +++--- .../platform/tests/core/mma_frag_ssot.cpp | 31 ++++++- .../platform/tests/core/test_ir_types.cpp | 66 ++++++++++++++ .../tests/core/test_mma_frag_tables.py | 90 +++++++++++++++++-- 6 files changed, 223 insertions(+), 56 deletions(-) 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 1b0a1da5a44e..c32c1839f5f8 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 @@ -61,12 +61,13 @@ static int rocke_mma_c_frag_len(const char* op_id) return rocke_arch_mma_c_frag_len(op_id); } -/* True when op_id accumulates in i32 (integer WMMA), from the arch catalog SSOT - * (target._op_id_c_dtype()[op_id] == "i32"). */ -static bool rocke_mma_is_int_acc(const char* op_id) +/* The accumulator dtype of op_id, from the arch catalog SSOT + * (target._op_id_c_dtype). Op_ids absent from the catalog accumulate in fp32. + * Mirrors Python _mma_c_dtype. */ +static const char* rocke_mma_c_dtype(const char* op_id) { const char* c_dtype = rocke_arch_mma_op_id_c_dtype(op_id); - return c_dtype != NULL && strcmp(c_dtype, ROCKE_DTYPE_I32) == 0; + return c_dtype != NULL ? c_dtype : "fp32"; } static const char* rocke_mma_result_hint(const char* op_id) @@ -611,8 +612,6 @@ rocke_value_t* rocke_b_mma(rocke_ir_builder_t* b, int num_extra) { int c_frag_len; - bool is_int_acc; - const rocke_type_t* c_elem; const rocke_type_t* vt; const char* hint; rocke_attr_map_t attrs; @@ -627,25 +626,6 @@ rocke_value_t* rocke_b_mma(rocke_ir_builder_t* b, { return (rocke_value_t*)rocke_i_set_err(b, ROCKE_ERR_VALUE, "mma op_id is NULL"); } - /* C has no MmaOp object; op_id is always a bare string, so the frag length - * and accumulator element are resolved from the arch SSOT (the Python - * bare-string code path). A <= 0 frag length means the op_id is unknown to - * the SSOT (the zero-length _frag_info fallback). */ - c_frag_len = rocke_mma_c_frag_len(op_id); - if(c_frag_len <= 0) - { - return (rocke_value_t*)rocke_i_set_err( - b, ROCKE_ERR_VALUE, "unknown MMA op_id '%s'; pass a known mfma_*/wmma_* op_id", op_id); - } - is_int_acc = rocke_mma_is_int_acc(op_id); - c_elem = is_int_acc ? rocke_i32() : rocke_f32(); - vt = rocke_vector_type(b, c_elem, c_frag_len); - if(!vt) - { - return NULL; - } - hint = rocke_mma_result_hint(op_id); - if(num_extra < 0) { num_extra = 0; @@ -663,12 +643,32 @@ 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); + /* Operands first, then the op_id, in IRBuilder.mma's order, so both engines + * report the same refusal for the same call. */ const char* mx_error = rocke_scaled_mfma_error(op_id, ops, nops); if(mx_error) return (rocke_value_t*)rocke_i_set_err(b, ROCKE_ERR_VALUE, "%s", mx_error); + 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); + /* C has no MmaOp object; op_id is always a bare string, so the frag length + * and accumulator element are resolved from the arch SSOT (the Python + * bare-string code path). A <= 0 frag length means the op_id is unknown to + * the SSOT (the zero-length _frag_info fallback). */ + c_frag_len = rocke_mma_c_frag_len(op_id); + if(c_frag_len <= 0) + { + return (rocke_value_t*)rocke_i_set_err( + b, ROCKE_ERR_VALUE, "unknown MMA op_id '%s'; pass a known mfma_*/wmma_* op_id", op_id); + } + /* The result element is the instruction's accumulator dtype. C has no MmaOp + * argument, so Python's "the MmaOp's own c_dtype" case has no twin here. */ + vt = rocke_vector_type(b, rocke_require_ir_type(rocke_mma_c_dtype(op_id)), c_frag_len); + if(!vt) + { + return NULL; + } + hint = rocke_mma_result_hint(op_id); 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/include/rocke/mfma_scale_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/mfma_scale_internal.h index 9cf5ef0f2940..b0e877f3e2b2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/mfma_scale_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/mfma_scale_internal.h @@ -39,13 +39,16 @@ static inline const char* rocke_scaled_mfma_error(const char* id, rocke_value_t* if(!id || !rocke_arch_mma_op_id_family(id) || strcmp(rocke_arch_mma_op_id_family(id), "mfma_scaled") != 0) return NULL; + /* One refusal for a wrong count or a wrong type, as mfma_scale.py compares + * the whole operand type list. */ + const char* error = "scaled MFMA requires two vec payloads, vec " + "accumulator and two i32 E8M0 scales"; const char* types[] = {"vec", "vec", "vec", "i32", "i32"}; if(n != 5) - return "scaled MFMA requires five operands"; + return error; for(int i = 0; i < 5; ++i) if(!args[i] || strcmp(args[i]->type->name, types[i]) != 0) - return "scaled MFMA requires two vec payloads, vec " - "accumulator and two i32 E8M0 scales"; + return error; return NULL; } #endif 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 17864d41f23a..73f7dfe36fb0 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 @@ -528,14 +528,14 @@ def _mma_c_frag_len(op_id: str) -> int: return frag_len -def _mma_c_is_int(op_id: str) -> bool: - """True when ``op_id`` accumulates in i32 (integer WMMA). +def _mma_c_dtype(op_id: str) -> str: + """The accumulator dtype of ``op_id``, from the arch catalog. - Sourced from the arch catalog's accumulator dtype - (``core/arch/data/arch_specs.json`` via ``target._op_id_c_dtype``). - Op_ids absent from the catalog default to the f32 accumulator. + Sourced from ``core/arch/data/arch_specs.json`` via + ``target._op_id_c_dtype``. Op_ids absent from the catalog accumulate in + fp32. """ - return _arch._op_id_c_dtype().get(op_id) == "i32" + return _arch._op_id_c_dtype().get(op_id, "fp32") @dataclass(frozen=True) @@ -2574,8 +2574,9 @@ def mma( matching MFMA call on CDNA or the WMMA call on RDNA. **One kernel body, two ISAs.** - The result vector type is ```` (the per-lane - accumulator length the atom produces). When ``op`` is an ``op_id`` + The result vector type is ````: the per-lane + accumulator length the atom produces, with C the instruction's + accumulator dtype from the arch catalog. When ``op`` is an ``op_id`` string the frag length is resolved from the static MMA fragment table; an emulated atom's op_id resolves through the physical atom it expands into. @@ -2603,12 +2604,10 @@ def mma( if hasattr(op, "c_frag_len") and op.c_frag_len else _mma_c_frag_len(op_id) ) - # Accumulator element type: integer WMMA atoms (iu8/iu4) accumulate in - # i32; everything else in f32. Prefer the atom's own c_dtype when ``op`` - # is an MmaOp, else resolve from the arch SSOT via op_id. - c_dtype = getattr(op, "c_dtype", None) - is_int_acc = c_dtype == "i32" if c_dtype is not None else _mma_c_is_int(op_id) - c_elem = I32 if is_int_acc else F32 + # Accumulator element type: the atom's own c_dtype when ``op`` is an + # MmaOp, else the catalog's accumulator dtype for the op_id. + c_dtype = getattr(op, "c_dtype", None) or _mma_c_dtype(op_id) + c_elem = dtype_to_ir_type(c_dtype) hint = ( "mxacc" if _arch._op_id_family().get(op_id) in ("wmma_scaled", "mfma_scaled") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/mma_frag_ssot.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/mma_frag_ssot.cpp index 0629f45a6727..0a29e41b96ce 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/mma_frag_ssot.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/mma_frag_ssot.cpp @@ -5,10 +5,12 @@ * frag-length / accumulator-dtype tables consulted by rocke_b_mma. * * rocke_b_mma sizes its tile.mma result vector as , - * where c_frag_len comes from the op_id frag-length table and acc_elem is i32 - * for integer WMMA atoms (else f32). This test pins that mapping for a - * representative set of atoms and checks the unknown-op_id error path, so a - * table edit that changes a result width/dtype is caught here. + * where c_frag_len comes from the op_id frag-length table and acc_elem is the + * IR type of the instruction's accumulator dtype in the arch catalog (fp32 for + * an op_id the catalog does not list). This test pins that mapping for a + * representative set of atoms, checks that every catalog accumulator has an IR + * type, and checks the unknown-op_id error path, so a table edit that changes + * a result width/dtype is caught here. Mirrors tests/core/test_mma_frag_tables.py. * * Plain executable: returns non-zero on the first failed check (a clean run is * the pass criterion). Registered via tests/CMakeLists.txt so it is installed @@ -123,6 +125,27 @@ int main(void) check_emulated_atom(&b, "mfma_f32_16x16x32_xf32", 8, 4); check_emulated_atom(&b, "mfma_f32_32x32x16_xf32", 8, 16); + /* No catalog row (fragment information only): accumulates in fp32. */ + CHECK(rocke_arch_mma_op_id_c_dtype("mfma_f32_4x4x4_f16") == NULL, + "mfma_f32_4x4x4_f16 has no catalog row"); + check_atom(&b, "mfma_f32_4x4x4_f16", 4, false); + + /* Every catalog accumulator has an IR type, and the op_id lookup rocke_b_mma + * uses returns the row's accumulator dtype. */ + int num_arches = 0; + const char* const* arches = rocke_known_arches(&num_arches); + for(int i = 0; i < num_arches; ++i) + { + const rocke_arch_target_t* arch = rocke_arch_target_from_gfx(arches[i]); + for(int j = 0; arch && j < arch->mma.num_ops; ++j) + { + const rocke_mma_op_t* op = &arch->mma.ops[j]; + const char* c_dtype = rocke_arch_mma_op_id_c_dtype(op->op_id); + CHECK(rocke_dtype_to_ir_type(op->c_dtype) != NULL, op->op_id); + CHECK(c_dtype != NULL && strcmp(c_dtype, op->c_dtype) == 0, op->op_id); + } + } + /* Keep the C catalog in parity with Python _MMA_FRAGMENT_INFO. The scaled * matrix operands are <16 x i32>, so fragment length is 16 elements, not * the equivalent 64-byte storage width. */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp index 4ceadaa6b62b..05e845897ada 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_ir_types.cpp @@ -34,6 +34,11 @@ * tests/core/test_bitcast_width.py); print "ok", or * "error: : " when the builder refuses * it, since a refusal is a TypeError or a ValueError. + * + * --mma + * build one target-neutral MMA of op_id on i32 zero operands; print its + * result type, or "error: " when the builder refuses the call + * (tests/core/test_mma_frag_tables.py runs it for every op_id). */ #include "rocke/ir.h" @@ -632,6 +637,64 @@ static int test_ir_type_of_dtype() return 0; } +/* One target-neutral MMA of `op_id` on i32 zero operands (the operands of + * test_mma_frag_tables.py's twin). The builder types the result from the op_id + * alone, so ordinary op_ids build; XF32 and scaled MFMA op_ids refuse these + * operands, and an unknown op_id is refused. */ +static build_outcome build_mma(const char* op_id, std::string* result_type) +{ + build_outcome outcome = {false, "", 0}; + rocke_ir_builder_t b; + if(rocke_ir_builder_init(&b, "mma_case") != ROCKE_OK) + { + outcome.error = "builder init failed"; + return outcome; + } + rocke_value_t* zero = rocke_b_const_i32(&b, 0); + const int ops_before = b.kernel->body->num_ops; + try + { + const rocke_value_t* result = rocke_b_mma(&b, op_id, zero, zero, zero, NULL, 0); + *result_type = result->type->name; + outcome.accepted = true; + } + catch(const ckc::Error& error) + { + outcome.error = error.what(); + } + /* Counted after a refusal too: a refused call must leave no op behind. */ + outcome.ops_added = b.kernel->body->num_ops - ops_before; + rocke_ir_builder_free(&b); + return outcome; +} + +static int print_mma(const char* op_id) +{ + std::string result_type; + const build_outcome outcome = build_mma(op_id, &result_type); + if(outcome.accepted) + printf("%s\n", result_type.c_str()); + else + printf("error: %s\n", outcome.error.c_str()); + return 0; +} + +/* The operands are checked before the op_id, as in IRBuilder.mma: an XF32 op_id + * on the wrong operands names the operands, and a refusal builds no op. */ +static int test_mma_refusals() +{ + std::string result_type; + const build_outcome built = build_mma("wmma_i32_16x16x16_iu8", &result_type); + CHECK(built.accepted && built.ops_added == 1 && result_type == "vec"); + const build_outcome xf32 = build_mma("mfma_f32_16x16x8_xf32", &result_type); + CHECK(!xf32.accepted && xf32.ops_added == 0); + CHECK(xf32.error == "XF32 MMA requires two vec operands and its FP32 accumulator"); + const build_outcome unknown = build_mma("not_an_op_id", &result_type); + CHECK(!unknown.accepted && unknown.ops_added == 0); + CHECK(unknown.error.rfind("unknown MMA op_id 'not_an_op_id'; ", 0) == 0); + return 0; +} + /* Every IR scalar type by name: one per rocke_scalar_kind_t, which the * static_assert and test_scalar_kinds_are_covered hold to. Widths are not * listed: they come from the engine. */ @@ -675,6 +738,8 @@ int main(int argc, char** argv) return print_scalar_bits(); if(argc == 3 && strcmp(argv[1], "--bitcast") == 0) return print_bitcast(argv[2]); + if(argc == 3 && strcmp(argv[1], "--mma") == 0) + return print_mma(argv[2]); CHECK(argc == 1); CHECK(test_unsigned_ops() == 0); CHECK(test_packed_only_values() == 0); @@ -682,6 +747,7 @@ int main(int argc, char** argv) CHECK(test_bitcast_width() == 0); CHECK(test_scalar_widths() == 0); CHECK(test_ir_type_of_dtype() == 0); + CHECK(test_mma_refusals() == 0); CHECK(test_scalar_kinds_are_covered() == 0); printf("rocke_ir_types: ok\n"); return 0; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_frag_tables.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_frag_tables.py index 99745c8c995d..7af80e6450d1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_frag_tables.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_frag_tables.py @@ -9,8 +9,8 @@ * ``_mma_c_frag_len`` — accumulator fragment length, resolved from the arch SSOT ``core.arch.target._MMA_FRAGMENT_INFO`` (no ir-side copy); -* ``_mma_c_is_int`` — whether the atom accumulates in i32, resolved from the JSON - catalog accumulator dtype (``target._op_id_c_dtype``); +* ``_mma_c_dtype`` -- the atom's accumulator dtype, resolved from the JSON catalog + (``target._op_id_c_dtype``); every catalog accumulator has an IR type; * ``_MMA_RESULT_HINT`` — the ir-side SSA result-name hints (naming only, kept in ir.py to preserve byte-identical value numbering). @@ -18,14 +18,21 @@ i32, or that carries a result-name hint, must also have a fragment length in the SSOT, or ``IRBuilder.mma`` would raise on an otherwise-valid atom. +The C++ ``rocke_b_mma`` is reached through the ``rocke_ir_types`` test binary +(``ROCKE_IR_TYPES_TEST``): for every op_id both builders give the same result +type, or refuse the call with the same text. + CPU-only: imports the accessors and the arch SSOT directly, no GPU / compile / launch required. """ from __future__ import annotations +import os +import subprocess import unittest +from rocke.core.arch import ArchTarget, known_arches from rocke.core.arch.target import ( _MMA_FRAGMENT_INFO, _frag_info, @@ -35,8 +42,13 @@ ) from rocke.core.ir import ( _MMA_RESULT_HINT, + F16, + F32, + I32, + IRBuilder, + _mma_c_dtype, _mma_c_frag_len, - _mma_c_is_int, + dtype_to_ir_type, ) @@ -49,6 +61,30 @@ def _sizable_op_ids(): } +def _python_mma(op_id: str) -> str: + """Build one MMA of ``op_id`` on i32 zero operands with the Python builder. + + The twin of ``rocke_ir_types --mma``: the result type, or ``error: `` + when the builder refuses the call. A refusal must leave the kernel without a + new op. + """ + b = IRBuilder("mma_case") + zero = b.const_i32(0) + ops_before = len(b.kernel.body.ops) + try: + return b.mma(op_id, zero, zero, zero).type.name + except ValueError as error: + assert len(b.kernel.body.ops) == ops_before + return f"error: {error}" + + +def _native_mma(executable: str, op_id: str) -> str: + """The C++ builder's answer for ``op_id``, from ``rocke_ir_types --mma``.""" + return subprocess.run( + [executable, "--mma", op_id], check=True, capture_output=True, text=True + ).stdout.rstrip("\n") + + class TestMmaFragTables(unittest.TestCase): def test_every_scaled_catalog_atom_has_the_same_padded_abi(self): rows = [ @@ -90,10 +126,7 @@ def test_int_acc_op_ids_have_frag_lengths(self): int_op_ids, "expected at least one i32-accumulator atom in the catalog" ) for op_id in int_op_ids: - self.assertTrue( - _mma_c_is_int(op_id), - msg=f"{op_id!r} is i32 in the catalog but _mma_c_is_int disagrees", - ) + self.assertEqual(_mma_c_dtype(op_id), "i32") self.assertGreater( _mma_c_frag_len(op_id), 0, @@ -115,6 +148,49 @@ def test_emulated_op_ids_size_through_their_physical_atom(self): _mma_c_frag_len(op_id), _MMA_FRAGMENT_INFO[physical].c_frag_len ) + def test_an_op_id_outside_the_catalog_accumulates_in_fp32(self): + # mfma_f32_4x4x4_f16 has fragment information but no catalog row. + self.assertNotIn("mfma_f32_4x4x4_f16", _op_id_c_dtype()) + self.assertEqual(_mma_c_dtype("mfma_f32_4x4x4_f16"), "fp32") + + def test_mma_results_take_the_accumulator_dtype(self): + for gfx, op_id, c, a_type, result in ( + ("gfx942", "mfma_f32_16x16x16_f16", (F32, 4), (F16, 4), "vec"), + ("gfx1151", "wmma_i32_16x16x16_iu8", (I32, 8), (I32, 4), "vec"), + ): + with self.subTest(op_id=op_id): + b = IRBuilder(f"result_{op_id}") + a = b.zero_vec(*a_type) + op = ArchTarget.from_gfx(gfx).mma.by_op_id(op_id) + self.assertEqual(b.mma(op, a, a, b.zero_vec(*c)).type.name, result) + self.assertEqual(b.mma(op_id, a, a, b.zero_vec(*c)).type.name, result) + + def test_every_catalog_accumulator_has_an_ir_type(self): + for gfx in known_arches(): + for op in ArchTarget.from_gfx(gfx).mma.ops: + with self.subTest(gfx=gfx, op_id=op.op_id): + dtype_to_ir_type(op.c_dtype) + self.assertEqual(_mma_c_dtype(op.op_id), op.c_dtype) + + def test_native_mma_results_match_the_python_builder(self): + executable = os.environ.get("ROCKE_IR_TYPES_TEST") + if not executable: + self.skipTest("set ROCKE_IR_TYPES_TEST to the built rocke_ir_types test") + catalog_op_ids = { + op.op_id + for gfx in known_arches() + for op in ArchTarget.from_gfx(gfx).mma.ops + } + for op_id in sorted(catalog_op_ids | set(_MMA_FRAGMENT_INFO)): + with self.subTest(op_id=op_id): + self.assertEqual(_native_mma(executable, op_id), _python_mma(op_id)) + # An unknown op_id is refused by both. The hint after the op_id differs: + # Python can take an MmaOp and lists the known op_ids, C cannot. + native = _native_mma(executable, "not_a_real_op_id") + refusal = "error: unknown MMA op_id 'not_a_real_op_id'; " + self.assertTrue(native.startswith(refusal), native) + self.assertTrue(_python_mma("not_a_real_op_id").startswith(refusal)) + def test_result_hint_op_ids_have_frag_lengths(self): for op_id in _MMA_RESULT_HINT: self.assertGreater( From 618356336723639bdfd466d5147e28c83ae13733 Mon Sep 17 00:00:00 2001 From: Chris Millette Date: Fri, 9 Oct 2026 18:49:05 -0400 Subject: [PATCH 11/11] fix(rocke): decide tf32 MMA availability from the target's catalog The LLVM and HIP lowerings refused a TF32 MMA on any target other than gfx942 by comparing the target name, with the message "XF32 MMA requires gfx942". The arch catalog already records which targets list each XF32 instruction. tf32_mma_target_error(op_id, arch) in core/tf32.py returns the error for a TF32 op_id that the target's MMA catalog does not list: "no MMA instruction found for '' on '' -- the target's MMA catalog does not list it; it is listed for ". The three call sites (the neutral MMA and XF32 LLVM lowerings, and the HIP MMA lowering) raise it as before, as a ValueError; the lowerers import core/tf32.py at module level, as it imports nothing that imports them. gfx942 lowers exactly as before, so no kernel's IR changes, and a target that gains the instruction in its catalog needs no code change. The C++ LLVM and HIP lowerers decide the same way through rocke_tf32_mma_target_error, with the same text. It is declared in tf32_internal.h and defined in the new cpp/core/tf32.cpp, so its catalog lookup is compiled once instead of inlined into every file that includes the header. The tests require the exact message from both LLVM lowerers and from the C++ HIP lowerer. dtype_support.md describes the rule. Co-Authored-By: Claude --- .../cpp/core/lower_hip/lower_hip_mma.cpp | 5 ++- .../platform/cpp/core/lower_llvm/mma.cpp | 5 ++- .../rocke/platform/cpp/core/tf32.cpp | 45 +++++++++++++++++++ .../cpp/include/rocke/tf32_internal.h | 5 +++ .../dsl_docs/reference/dtype_support.md | 5 ++- .../platform/python/rocke/core/lower_hip.py | 10 ++--- .../platform/python/rocke/core/lower_llvm.py | 15 +++---- .../rocke/platform/python/rocke/core/tf32.py | 23 ++++++++++ .../rocke/platform/tests/core/test_tf32.cpp | 36 +++++++++++++++ .../rocke/platform/tests/core/test_tf32.py | 15 ++++--- 10 files changed, 139 insertions(+), 25 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/tf32.cpp 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 391560c9ff57..084b988cc70e 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 @@ -187,8 +187,9 @@ static rocke_status_t rocke_h_op_tile_mma(rocke_h_lowerer_t* lw, const rocke_op_ 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"); + char tf32_error[ROCKE_ERR_MSG_CAP]; + if(rocke_tf32_mma_target_error(op_id, lw->arch.gfx, tf32_error, sizeof tf32_error)) + return rocke_h_fail(lw, ROCKE_ERR_VALUE, "%s", tf32_error); rocke_h_emitf(lw, "f32x%d %s = __builtin_amdgcn_%s(__builtin_bit_cast(f32x2, %s), " "__builtin_bit_cast(f32x2, %s), %s, 0, 0, 0);", 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 b4e09d1ff1b6..39e4947741a6 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 @@ -500,8 +500,9 @@ static void _op_tile_mma(rocke_lower_t* L, const rocke_op_t* op) rocke_ll_operand(L, op->operands[4])); return; } - if(rocke_tf32_mma_count(op_id) && strcmp(L->backend->gfx, "gfx942") != 0) - rocke_ll_fail(L, ROCKE_ERR_VALUE, "XF32 MMA requires gfx942"); + char tf32_error[ROCKE_ERR_MSG_CAP]; + if(rocke_tf32_mma_target_error(op_id, L->backend->gfx, tf32_error, sizeof tf32_error)) + rocke_ll_fail(L, ROCKE_ERR_VALUE, "%s", tf32_error); /* An emulated atom is expanded here, before ISA dispatch, so the expansion * is written once rather than per ISA family. Native atoms -- every atom diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/tf32.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/tf32.cpp new file mode 100644 index 000000000000..b385acbdea99 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/tf32.cpp @@ -0,0 +1,45 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* Mirrors the target check of core/tf32.py: whether a target's MMA catalog lists + * a TF32 MMA. The operand and op checks need no catalog and stay inline in + * tf32_internal.h. */ +#include "rocke/tf32_internal.h" + +#include + +#include "rocke/arch_target.h" + +/* Whether the MMA catalog of target `gfx` lists `op_id` (an unknown target lists + * nothing). */ +static bool target_lists(const char* gfx, const char* op_id) +{ + const rocke_arch_target_t* target = rocke_arch_target_from_gfx(gfx); + return target && rocke_mma_catalog_by_op_id(&target->mma, op_id); +} + +/* The targets whose catalogs list the op are named in known_arches() order. */ +bool rocke_tf32_mma_target_error(const char* op_id, const char* gfx, char* out, size_t cap) +{ + if(!rocke_tf32_mma_count(op_id) || target_lists(gfx, op_id)) + return false; + int used = snprintf(out, + cap, + "no MMA instruction found for '%s' on '%s' -- the target's MMA catalog " + "does not list it; ", + op_id, + gfx); + int listed = 0; + int num_arches = 0; + const char* const* arches = rocke_known_arches(&num_arches); + for(int i = 0; i < num_arches; ++i) + { + if(!target_lists(arches[i], op_id)) + continue; + const char* lead = listed++ ? ", " : "it is listed for "; + if(used >= 0 && (size_t)used < cap) + used += snprintf(out + used, cap - (size_t)used, "%s%s", lead, arches[i]); + } + if(!listed && used >= 0 && (size_t)used < cap) + snprintf(out + used, cap - (size_t)used, "no target lists it"); + return true; +} 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 21222b142e47..8c0df8655767 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 @@ -17,6 +17,11 @@ static inline int rocke_tf32_mma_count(const char* id) return 16; return 0; } +/* Mirrors tf32_mma_target_error: for a TF32 MMA `op_id` that the MMA catalog of + * target `gfx` does not list, writes the error into `out` and returns true. Any + * other op_id, and a TF32 op_id the catalog lists, returns false. Defined in + * cpp/core/tf32.cpp, the one TU that needs the arch catalog for it. */ +bool rocke_tf32_mma_target_error(const char* op_id, const char* gfx, char* out, size_t cap); /* The operand contract of a TF32 atom the hardware does not encode (mirrors * core/tf32.py::_emulated_tf32_mma_error). Same ABI as the native atoms -- logical * TF32 A/B carried as I32, FP32 accumulator -- with fragment lengths from the 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 2414256fb598..c57433b1cdc4 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 @@ -573,7 +573,10 @@ gfx1250 has no native F64 matrix instruction and is excluded from this gap. **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; gfx950's catalog +other targets, deciding from the target's MMA catalog: a target whose catalog +does not list the atom raises "no MMA instruction found for +'' on ''", so a target that gains the atom in its catalog needs no +code change. gfx950/gfx1250 have no native XF32 instruction; gfx950's catalog carries two emulated atoms, `mfma_f32_16x16x32_xf32` and `mfma_f32_32x32x16_xf32`, which take the same logical TF32 operands and expand into three BF16 MFMAs. This atom support does not add production GEMM/convolution 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 47ed1c553c24..da2962732a64 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 @@ -44,6 +44,7 @@ mma_operand_needs_reinterpret, mma_operand_value, ) +from .tf32 import TF32_MMA, tf32_mma_target_error, tf32_op_error # No packed-only type has an entry: core/ir.py _refuse_packed_only_value @@ -403,8 +404,6 @@ 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) @@ -675,8 +674,6 @@ 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"] from .arch.mfma_scale import ( gfx950_scaled_mfma, @@ -699,8 +696,9 @@ def _op_tile_mma(self, op: Op) -> None: ) return if op_id in TF32_MMA: - if self.arch.gfx != "gfx942": - raise ValueError("XF32 MMA requires gfx942") + error = tf32_mma_target_error(op_id, self.arch) + if error: + raise ValueError(error) a, b, c = op.operands self._emit( f"f32x{TF32_MMA[op_id]} {_name(op.result)} = __builtin_amdgcn_{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 50c79ba27b40..0f8996a930f5 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 @@ -74,6 +74,7 @@ mma_operand_needs_reinterpret, mma_operand_value, ) +from .tf32 import tf32_mma_target_error, tf32_op_error # Datalayout / triple. Copied verbatim from clang's output for the same @@ -2333,8 +2334,6 @@ 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) @@ -3756,10 +3755,9 @@ 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") + error = tf32_mma_target_error(op.attrs.get("op_id"), self._backend.arch) + if error: + raise ValueError(error) # 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 @@ -3993,8 +3991,9 @@ 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") + error = tf32_mma_target_error(_mma_op_id(op), self._backend.arch) + if error: + raise ValueError(error) a, b, c = op.operands self._need(intrinsic) a_cast = self._fresh("mfma_a_i16") 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 10eaf63df226..5876c07c2919 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 @@ -2,9 +2,32 @@ # SPDX-License-Identifier: MIT """Logical TF32 operand contracts, independent of the physical i32 carrier.""" +from __future__ import annotations + +from .arch import ArchTarget, known_arches + TF32_MMA = {"mfma_f32_16x16x8_xf32": 4, "mfma_f32_32x32x4_xf32": 16} +def tf32_mma_target_error(op_id: str | None, arch: ArchTarget) -> str | None: + """The error for a TF32 MMA ``op_id`` that ``arch``'s MMA catalog does not list. + + Returns None for any other op_id, and for a TF32 op_id the catalog lists. + """ + if op_id not in TF32_MMA or arch.mma.by_op_id(op_id) is not None: + return None + listing = [ + gfx for gfx in known_arches() if ArchTarget.from_gfx(gfx).mma.by_op_id(op_id) + ] + where = ( + f"it is listed for {', '.join(listing)}" if listing else "no target lists it" + ) + return ( + f"no MMA instruction found for {op_id!r} on {arch.gfx!r} -- the target's " + f"MMA catalog does not list it; {where}" + ) + + def _emulated_tf32_mma_error(op_id, lens, operands, results): """The operand contract of a TF32 atom the hardware does not encode. 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 33b64dd7b735..ef3dc803b43f 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 @@ -198,11 +198,47 @@ static int test_invalid_store_alignment() return 0; } +/* Mirrors test_rejects_unsupported_target: a target whose MMA catalog does not + * list the XF32 instruction refuses it in both C++ lowerers, with the text + * Python raises. The HIP lowerer reports it on the builder's error channel. */ +static int test_rejects_unsupported_target() +{ + for(const char* gfx : {"gfx950", "gfx1151", "gfx1250"}) + { + char message[256]; + std::snprintf(message, + sizeof(message), + "no MMA instruction found for 'mfma_f32_16x16x8_xf32' on '%s' -- the " + "target's MMA catalog does not list it; it is listed for gfx942", + gfx); + rocke_ir_builder_t b; + auto* kernel = rocke_build_tf32_mma_probe(&b, 16, "raw"); + CHECK(kernel); + char* ll = nullptr; + char err[ROCKE_ERR_MSG_CAP] = {}; + CHECK(rocke_lower_kernel_to_llvm_ex( + kernel, ROCKE_LLVM_FLAVOR_LLVM20, gfx, &ll, err, sizeof(err)) + == ROCKE_ERR_VALUE); + CHECK(std::strcmp(err, message) == 0); + std::free(ll); + rocke_strbuf_t hip; + CHECK(rocke_strbuf_init(&hip, 0) == 0); + rocke_lower_hip_opts_t opts = {}; + opts.arch = gfx; + CHECK(rocke_lower_kernel_to_hip(&b, kernel, &opts, &hip) == ROCKE_ERR_VALUE); + CHECK(std::strcmp(rocke_ir_builder_error(&b), message) == 0); + rocke_strbuf_free(&hip); + 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(test_rejects_unsupported_target() == 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 5df782de5905..e8824da649d8 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 @@ -2,6 +2,7 @@ # SPDX-License-Identifier: MIT """Logical TF32 contracts and strict native parity, independent of GPU access.""" import os +import re import subprocess from pathlib import Path @@ -25,6 +26,7 @@ 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.tests.engines import assert_both_engines_reject from rocke.instances.gfx942.tf32_mma_probe import ( PREPARATIONS, Tf32MmaProbeSpec, @@ -128,13 +130,14 @@ def test_mma_requires_logical_tf32(dtype, n): @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"): + message = ( + f"no MMA instruction found for 'mfma_f32_16x16x8_xf32' on '{gfx}' -- the " + "target's MMA catalog does not list it; it is listed for gfx942" + ) + assert_both_engines_reject(kernel, gfx, ValueError, message) + # The C++ HIP lowerer is checked against the same text by the rocke_tf32 ctest. + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): 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():