From d2b8b4acb58656bd12996235a5149f69a93c0c64 Mon Sep 17 00:00:00 2001 From: David Peabody Date: Fri, 16 Oct 2026 22:20:24 -0600 Subject: [PATCH 1/3] fix(rocke): wave_all is never true on wave32 ballot.i64 sets only the low wave_size bits, so comparing it against -1 can only succeed on wave64. Compare against the full-wave mask (0xffffffff on wave32, -1 on wave64) in both the Python and C LLVM lowerings; wave64 output is unchanged. On gfx11 this re-enables the lazy-rescale skip in the transposed-QK attention path (its 16 parity anchors plus 2 replays are re-blessed). Replaces the IR-text test with an on-GPU wave_all/wave_any test that runs a full wave on the detected device. --- .../cpp/core/lower_llvm/crosslane.cpp | 15 ++- .../primitives/wave_and_cross_lane.md | 2 +- .../platform/python/rocke/core/lower_llvm.py | 16 ++- .../differential/golden/llvm_gfx_all.json | 36 +++--- .../tests/numeric/test_wave_vote_numeric.py | 111 ++++++++++++++++++ .../rocke/platform/tests/test_rocke.py | 15 --- 6 files changed, 150 insertions(+), 45 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/numeric/test_wave_vote_numeric.py diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/crosslane.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/crosslane.cpp index 709260d0e9b1..2c7b20c18f04 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/crosslane.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/lower_llvm/crosslane.cpp @@ -165,15 +165,26 @@ static void _op_tile_wave_ballot(rocke_lower_t* L, const rocke_op_t* op) static void _op_tile_wave_all(rocke_lower_t* L, const rocke_op_t* op) { - /* wave_all(p) = (ballot(p) == -1). */ + /* wave_all(p) = (ballot(p) == full-wave mask): -1 on wave64, 0xffffffff on + * wave32 (ballot.i64 sets only the low wave_size bits). Mirrors Python. */ const rocke_value_t* pred = op->operands[0]; + const rocke_arch_target_t* target; const char* b_name; const char* eq_name; rocke_ll_need(L, "ballot.i64"); b_name = rocke_ll_fresh(L, "ballot"); rocke_ll_emit_wave_ballot(L, pred, b_name); + target = rocke_arch_target_from_gfx(L->backend->gfx); + if(target == NULL) + { + rocke_ll_fail(L, ROCKE_ERR_KEY, "arch metadata not present for %s", L->backend->gfx); + } eq_name = rocke_ll_fresh(L, "all_eq"); - rocke_ll_emitf(L, " %s = icmp eq i64 %s, -1", eq_name, b_name); + rocke_ll_emitf(L, + " %s = icmp eq i64 %s, %s", + eq_name, + b_name, + target->wave_size == 64 ? "-1" : "4294967295"); rocke_ll_emitf(L, " %s = zext i1 %s to i32", ll_result_name(op), eq_name); } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/wave_and_cross_lane.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/wave_and_cross_lane.md index 750382fc29d6..6a4a66df6011 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/wave_and_cross_lane.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/wave_and_cross_lane.md @@ -57,7 +57,7 @@ wave_any(pred) # i32, 1 iff any lane's pred != 0 wave_ballot(pred) # i64 bitmask ``` -LLVM emission uses `llvm.amdgcn.ballot.i64`. `wave_all` / `wave_any` compare the ballot mask to specific values (`-1` for all, `0` for none). +LLVM emission uses `llvm.amdgcn.ballot.i64`, which sets only the low `wave_size` bits. `wave_all` compares the mask to the full-wave value (`-1` on wave64, `0xffffffff` on wave32); `wave_any` compares it to `0`. Typical use: adaptive online softmax. When `wave_all(max_diff < THRESHOLD)`, the workgroup can skip the rescale path entirely. 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..2f23c5c149b5 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 @@ -4399,20 +4399,18 @@ def _op_tile_wave_ballot(self, op: Op) -> None: self._emit_wave_ballot(pred, op.result.name) def _op_tile_wave_all(self, op: Op) -> None: - # wave_all(p) = (ballot(p) == -1). - # On wave64 with all lanes active, ``ballot.i64`` returns - # ``-1 == 0xffffffffffffffff`` iff every lane voted true. - # Matches HIP's ``__all(p)`` which uses the same shorthand. - # This form is robust under wave64 with full EXEC (our standard - # workgroup launch profile); under a partially-predicated EXEC - # mask the result counts only the active lanes (still the - # semantics we want for "did every *active* lane vote true"). + # wave_all(p) = (ballot(p) == full-wave mask). ``ballot.i64`` sets + # only the low ``wave_size`` bits, so the full mask is -1 on wave64 + # but 0xffffffff on wave32 (comparing against -1 there is never + # true). Like HIP's ``__all(p)`` this assumes the full-EXEC launch + # profile every rocKE kernel uses. (pred,) = op.operands self._need("ballot.i64") b_name = self._fresh("ballot") self._emit_wave_ballot(pred, b_name) + full = -1 if self._backend.arch.wave_size == 64 else (1 << 32) - 1 eq_name = self._fresh("all_eq") - self._current().emit(f" {eq_name} = icmp eq i64 {b_name}, -1") + self._current().emit(f" {eq_name} = icmp eq i64 {b_name}, {full}") self._current().emit(f" {op.result.name} = zext i1 {eq_name} to i32") def _op_tile_wave_any(self, op: Op) -> None: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json index 7c846e35d95e..a79346a2254e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json @@ -1522,8 +1522,8 @@ "134": "497a69c14a7db7199ca1ceeabdc5a477831b30d40e5a09f0caf32d9ff24ce0ab", "135": "58e12c23677155dceda8e80caa6a8e065f11d596cd3108fa83543c4b7b2c4cc6", "136": "9245080477727e4e4e7840b2d75694447ac15b77d0a7a54af819b50daf4b0826", - "137": "c8d51fc6cfd635cb09fe41c74a50a46e1bff906326cab0b23a4a20e4532dccf9", - "138": "fcd33f301315d1c64a2787e46308e40404b0c22ed74160727d74e6674d4d2f5f", + "137": "c5dd8d55db78eb544d6a17958e1d5b0bb7cc6037318f7d013b25d0cbf5ebe7d9", + "138": "6e158220bafb4c020c8fb74832801f706cc3b7c5e7fddbe3038a17ec26b276b1", "139": "fc52889792813615dc2171ea788eb854e966e75dac2c6c01da7933aac15f64d2", "14": "9a1b1005ea08665c7c187bbf2f036b1c5c1ec97fc1298c02f95240d4c61eaba1", "140": "69582e941f2e52a2c14ebd8c68b059eed50e2bb9107d240065c359820ff8d93f", @@ -1604,23 +1604,23 @@ "68": "c77b6e0e9852d5134d8927e3ae15f245e152e287edd2a4562fdf417caa66533a", "69": "45f51c24447c4941cf9b466b5ecd25da5a2acb25dc911753fbbe77818e424f8d", "7": "0ebed7c112418c6d60eee90a944349ccd2e51bd611232a412e2108966cb07a47", - "70": "c8d51fc6cfd635cb09fe41c74a50a46e1bff906326cab0b23a4a20e4532dccf9", - "71": "12b342be7ca100f260b1a513910a692e8eae81c86b9bb8fd40424d8f91db260e", - "72": "d542789ef4f7a6aad385d04b81d3557eec1a5945a72f25278abea85fdfabf729", - "73": "36aa7528fff44acef01925531e9d5041b1e2315e60ea0e9593a4b2e4dbe8a734", - "74": "6a0b8ac5b7d781f9c730f7381be739efb4d1890a44bac9dcc93c27bc8433a281", - "75": "8045fda7812d9fa89b400d6ac736d2f0a04fccd5d9be7ea6bbe5fdde4e583d97", - "76": "8cd952f9e566b0b0d9b8db313e208b327f4249cdac7acfb2def5862b31c50655", - "77": "7c5e36ba5e28893d0c219a943499361baa3e4fa22d54f65d373a807b6a3cc8cd", - "78": "593fd300a180cf4ce01f717669679864becdd412204baa6101c6268a5502577c", - "79": "5e90dc8b294b173a45a1aae688fb35529fc22f61cfdd792af19e5e50090396b9", + "70": "c5dd8d55db78eb544d6a17958e1d5b0bb7cc6037318f7d013b25d0cbf5ebe7d9", + "71": "9e8ba381439edbec0b63b0f71010835625706ccd101c82d582d4e86ea7ce9c90", + "72": "c69c6ad07c44d093c304fdf822f2a135245bcb726f8cbb8fe1e1f546969cf12b", + "73": "df2b1756e57ede9a4bddf1f1c0d0c23f2c4ea027f92bc1d04f58df16690728ad", + "74": "bd1c4634694808a33b587fffedc85eefb7b7a07ada4309a3edf836047efa1243", + "75": "6df6c7f070a7cbee0d2a0d410fcbb21b9f7cf9d88cdb25a133b937b4b3664692", + "76": "8cef92b5e8d5a29630c58079f4a4cc1182cbb821dc3c7878c3c81b53cceeb91a", + "77": "abecc12660f6c6f2f7ba4e92b1db163270080024760310b74681dc330caef02d", + "78": "f3b14ceaca3a130fad33e38f22a153279b7f8a932611ae7b7da4da2874f52c37", + "79": "28978cf5a8f94f31f00719b7dc577c1f864d10afa31d17570768bfb08bd905a8", "8": "09eea85e2bf58c29ae59f436906e8a35484efbab438bfbaec34038622c0e9ecf", - "80": "b69bb77bf983bbd047ecf874d8bb1cb40914a680cddfcedd62b6d0b8ee662955", - "81": "68af93cb3b81effcc29c3752eadb7e500354803042b0932e305f0de98e09bc47", - "82": "884d722dd9e41361e5079dcb2366fdfd1e484d4e95c557993d0dc021173a28a2", - "83": "c2110a939d008c15ad3e75f4821d4e5814304ee56e0d423034dfdd541dcca9bf", - "84": "4d4eabd49e71e76f1d1e2f55c3f13bc7311ef9e639a774592bc1946e6b9db997", - "85": "fcd33f301315d1c64a2787e46308e40404b0c22ed74160727d74e6674d4d2f5f", + "80": "f171dc8cb46717873504573a83262b528d37462ab75b1711bb2b44f7e5f49cb9", + "81": "bb401137f92fc912513113e01f8630a2a254c4bd2a60cbca51417f888aeaf277", + "82": "222cf5cb29f3472432e338211d5d7ef249bc1a3fa2b2847f92cf2ccf92b79b9d", + "83": "17213fa5b6865e24eee09eb06f7dbbb93403080190b51025e42b0c1691dc6582", + "84": "f7fab1c4dda69fc784753fba697dd8da45450d0f8141f1e37103670c8e7bda56", + "85": "6e158220bafb4c020c8fb74832801f706cc3b7c5e7fddbe3038a17ec26b276b1", "86": "e79f5e017dfa950e85d2d07334888811cf67073d25811cd6b845873f41cd4f4e", "87": "447f370aad0812bfe5eecace568a40f4ee6adc506e9fe3395e231d9ea86a89df", "88": "9e627a8a11115ac399a3250fda20a15eacd0245b015898428ca88a4d87e28486", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/numeric/test_wave_vote_numeric.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/numeric/test_wave_vote_numeric.py new file mode 100644 index 000000000000..da3275c95ba8 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/numeric/test_wave_vote_numeric.py @@ -0,0 +1,111 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Real-GPU checks for ``IRBuilder.wave_all`` / ``wave_any`` on one full wave +of the detected device (wave32 on RDNA, wave64 on CDNA). + +Each lane votes its input; every lane must receive the wave-wide AND/OR. +The cases where every lane votes true pin the wave32 full-mask comparison +(``ballot.i64`` sets only the low 32 bits there). Both engines lower the +kernel; ``cpp`` needs the binding. Machines without a HIP device skip. +""" + +import ctypes +import importlib.util +import struct + +import numpy as np +import pytest + +from rocke.core.arch import ArchTarget +from rocke.core.ir import I32, IRBuilder, PtrType +from rocke.helpers import compile_kernel +from rocke.runtime.hip_module import Runtime, get_device_arch + + +@pytest.fixture(scope="module") +def arch(): + try: + target = get_device_arch(0) + except Exception as exc: # noqa: BLE001 - no HIP runtime or device + pytest.skip(f"requires a HIP device: {exc}") + try: + ArchTarget.from_gfx(target) + except KeyError: + pytest.skip(f"no rocKE arch target for {target}") + return target + + +@pytest.fixture(params=("python", "cpp")) +def backend(request): + if request.param == "cpp" and importlib.util.find_spec("rocke_engine") is None: + pytest.skip("cpp engine needs the rocke_engine binding") + return request.param + + +def _vote_kernel(vote): + b = IRBuilder(f"wave_vote_{vote}") + src = b.param("S", PtrType(I32, "global"), noalias=True, readonly=True) + dst = b.param("D", PtrType(I32, "global"), noalias=True) + tid = b.thread_id_x() + pred = b.global_load(src, tid, I32, align=4) + b.global_store(dst, tid, getattr(b, vote)(pred), align=4) + b.ret() + return b.kernel + + +def _run(vote, values, arch, backend): + wave = ArchTarget.from_gfx(arch).wave_size + art = compile_kernel(_vote_kernel(vote), arch=arch, backend=backend) + rt = Runtime() + module = rt.load_module(art.hsaco) + nbytes = values.nbytes + src = rt.alloc(nbytes) + dst = rt.alloc(nbytes) + try: + host = (ctypes.c_uint8 * nbytes).from_buffer_copy(values.tobytes()) + rt.memcpy_h2d(src, host, nbytes) + out = (ctypes.c_uint8 * nbytes)() + rt.launch( + module.get_function(art.kernel_name), + (1, 1, 1), + (wave, 1, 1), + struct.pack(" Date: Sun, 18 Oct 2026 00:07:32 -0600 Subject: [PATCH 2/3] test(rocke): guard wave_all full-wave mask on wave32 and fix stale docs Add a host-only lowering test and a C++ ctest case that pin the wave_all full-wave mask (-1 on wave64, 4294967295 on wave32), and correct the wave_all/wave_ballot docstrings and helpers README. --- .../rocke/platform/python/rocke/core/ir.py | 8 ++--- .../platform/python/rocke/helpers/README.md | 3 +- .../tests/core/future_intrinsic_lowering.cpp | 33 +++++++++++++++++++ .../rocke/platform/tests/test_rocke.py | 32 ++++++++++++++++++ 4 files changed, 71 insertions(+), 5 deletions(-) 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 d692a1c3969f..6df869ad5067 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 @@ -2968,8 +2968,8 @@ def wave_all(self, predicate: Value) -> Value: Returns a wave-uniform i32 that is 1 iff *every* lane's ``predicate`` was non-zero, 0 otherwise. Lowered to - ``__builtin_amdgcn_read_exec()`` AND/compare on AMDGPU - (a single ``s_or_b64`` / ``s_cmp_eq`` pair) — no ds_bpermute, + ``@llvm.amdgcn.ballot.i64`` compared against the full-wave mask + (``-1`` on wave64, ``0xffffffff`` on wave32) — no ds_bpermute, no LDS round-trip. Pairs naturally with adaptive online-softmax rescaling: when @@ -3000,8 +3000,8 @@ def wave_ballot(self, predicate: Value) -> Value: """Wave-wide ballot: returns a 64-bit mask of which lanes satisfied the predicate. - Lowered to ``__builtin_amdgcn_ballot_w64`` (or the wave32 - equivalent on gfx12). For boolean reductions prefer the + Lowered to ``@llvm.amdgcn.ballot.i64`` (the mask has only the low + ``wave_size`` bits set on wave32). For boolean reductions prefer the higher-level :meth:`wave_all` / :meth:`wave_any` which avoid forcing the mask materialisation on the consumer side. """ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/README.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/README.md index 4d7af56c32b1..9ff0249a98cc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/README.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/README.md @@ -856,7 +856,8 @@ if_branch = b.cmp_eq(all_below, b.const_i32(1)) # ... skip the per-row rescale when every lane agrees ... ``` -Lowered to `llvm.amdgcn.ballot.i64` + `icmp eq i64 ... -1`. Pairs +Lowered to `llvm.amdgcn.ballot.i64` + `icmp eq i64 ...` against the +full-wave mask (`-1` on wave64, `0xffffffff` on wave32). Pairs naturally with adaptive online-softmax rescaling: skip the rescale pass entirely when every lane's `max_diff` stayed below threshold. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp index 003d51ba5522..a94caef5e239 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/future_intrinsic_lowering.cpp @@ -3188,6 +3188,38 @@ void case_lowbit_dtype_aliases() } } +/* wave_all compares ballot.i64 against the full-wave mask: -1 on wave64 but + * 0xffffffff on wave32 (ballot.i64 only sets the low wave_size bits). */ +void case_wave_all_full_wave_mask() +{ + struct ArchMask + { + const char* arch; + const char* mask; + }; + const ArchMask archs[] = {{"gfx950", ", -1\n"}, + {"gfx942", ", -1\n"}, + {"gfx1151", ", 4294967295\n"}, + {"gfx11-generic", ", 4294967295\n"}, + {"gfx12-generic", ", 4294967295\n"}, + {"gfx1201", ", 4294967295\n"}}; + for(const ArchMask& a : archs) + { + const std::string ir = lower_one( + "wave_all_mask", + [](rocke_ir_builder_t* b) { rocke_b_wave_all(b, rocke_b_const_i32(b, 1)); }, + a.arch); + EXPECT_IR(ir, "@llvm.amdgcn.ballot.i64"); + const size_t at = ir.find("icmp eq i64 "); + const size_t eol = ir.find('\n', at); + const std::string line = at == std::string::npos ? "" : ir.substr(at, eol - at); + const std::string want(a.mask, strlen(a.mask) - 1); + if(line.size() < want.size() + || line.compare(line.size() - want.size(), want.size(), want) != 0) + fail(a.arch, __LINE__); + } +} + struct TestCase { const char* name; @@ -3273,6 +3305,7 @@ const TestCase k_cases[] = { {"flag_ops_verify_and_llvm_arity", case_flag_ops_verify_and_llvm_arity}, {"flag_ops_reject_non_i32_ptr", case_flag_ops_reject_non_i32_ptr}, {"opcode_names_are_aligned", case_opcode_names_are_aligned}, + {"wave_all_full_wave_mask", case_wave_all_full_wave_mask}, }; } // namespace diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py index e3f7882eec2b..7b09550f5a54 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py @@ -2736,6 +2736,38 @@ def test_to_sgpr_u32_emits_readfirstlane_and_pin(self): # and silently breaks downstream uniform-value selection. self.assertNotIn("asm sideeffect", ll) + # ---- wave_all / wave_any ---- + def _lower_wave_votes(self, arch): + from rocke.core.ir import IRBuilder, PtrType, I32 + from rocke.core.lower_llvm import lower_kernel_to_llvm + + b = IRBuilder("wave_vote_smoke") + b.kernel.attrs["max_workgroup_size"] = 64 + Out = b.param("Out", PtrType(I32, "global")) + pred = b.const_i32(1) + b.global_store(Out, b.const_i32(0), b.wave_all(pred)) + b.global_store(Out, b.const_i32(1), b.wave_any(pred)) + return lower_kernel_to_llvm(b.kernel, arch=arch) + + def test_wave_all_full_mask_matches_wave_size(self): + import re + + # ballot.i64 only sets the low wave_size bits: the full-wave mask + # is -1 on wave64 but 0xffffffff on wave32. + for arch, full in ( + ("gfx950", "-1"), + ("gfx942", "-1"), + ("gfx1151", "4294967295"), + ("gfx11-generic", "4294967295"), + ("gfx12-generic", "4294967295"), + ("gfx1201", "4294967295"), + ): + with self.subTest(arch=arch): + ll = self._lower_wave_votes(arch) + self.assertIn("@llvm.amdgcn.ballot.i64", ll) + self.assertRegex(ll, rf"icmp eq i64 %\S+, {re.escape(full)}\n") + self.assertRegex(ll, r"icmp ne i64 %\S+, 0\n") + # ---- coherency hints ---- def test_async_buffer_load_lds_addr_propagates_aux(self): from rocke.core.ir import ( From 1d602ff49a0f673f093d8bc91a2a04dc72d5732a Mon Sep 17 00:00:00 2001 From: David Peabody Date: Mon, 19 Oct 2026 13:14:11 -0600 Subject: [PATCH 3/3] docs(rocke): state wave_all's full-EXEC precondition --- .../platform/dsl_docs/primitives/wave_and_cross_lane.md | 2 ++ .../rocke/platform/python/rocke/core/ir.py | 4 ++++ .../rocke/platform/python/rocke/core/lower_llvm.py | 6 ++++-- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/wave_and_cross_lane.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/wave_and_cross_lane.md index 6a4a66df6011..99bd0799dd98 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/wave_and_cross_lane.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/wave_and_cross_lane.md @@ -59,6 +59,8 @@ wave_ballot(pred) # i64 bitmask LLVM emission uses `llvm.amdgcn.ballot.i64`, which sets only the low `wave_size` bits. `wave_all` compares the mask to the full-wave value (`-1` on wave64, `0xffffffff` on wave32); `wave_any` compares it to `0`. +`wave_all` requires full EXEC (every lane of the wave active). Unlike HIP's `__all`, which counts only active lanes, it returns 0 under partial EXEC (a divergent region or a sub-wave workgroup). + Typical use: adaptive online softmax. When `wave_all(max_diff < THRESHOLD)`, the workgroup can skip the rescale path entirely. ## `ds_bpermute` 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 6df869ad5067..df3a991d8920 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 @@ -2972,6 +2972,10 @@ def wave_all(self, predicate: Value) -> Value: (``-1`` on wave64, ``0xffffffff`` on wave32) — no ds_bpermute, no LDS round-trip. + Precondition: every lane of the wave is active (full EXEC). Unlike + HIP's ``__all``, which counts only active lanes, this returns 0 + under partial EXEC (a divergent region or a sub-wave workgroup). + Pairs naturally with adaptive online-softmax rescaling: when ``wave_all(max_diff < THRESHOLD)`` is 1, the workgroup can skip the rescale path entirely. 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 2f23c5c149b5..49a7aa0f0660 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 @@ -4402,8 +4402,10 @@ def _op_tile_wave_all(self, op: Op) -> None: # wave_all(p) = (ballot(p) == full-wave mask). ``ballot.i64`` sets # only the low ``wave_size`` bits, so the full mask is -1 on wave64 # but 0xffffffff on wave32 (comparing against -1 there is never - # true). Like HIP's ``__all(p)`` this assumes the full-EXEC launch - # profile every rocKE kernel uses. + # true). Unlike HIP's ``__all(p)``, which counts only active lanes, + # this requires full EXEC: under partial EXEC (a divergent region or + # a sub-wave workgroup) it returns 0. Every rocKE kernel launches + # with full EXEC. (pred,) = op.operands self._need("ballot.i64") b_name = self._fresh("ballot")