fix(rocke): wave_all is never true on wave32

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);
 }
 

dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/primitives/wave_and_cross_lane.md

@@ -57,7 +57,9 @@ 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`.
+
+`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.
 

dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py

@@ -2968,10 +2968,14 @@ 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.
 
+        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.
@@ -3000,8 +3004,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.
         """

dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py

@@ -4399,20 +4399,20 @@ 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). 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")
         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:

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.
 

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

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",

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("<QQ", src, dst),
+        )
+        rt.sync()
+        rt.memcpy_d2h(out, dst, nbytes)
+        return np.frombuffer(bytes(out), dtype=values.dtype)
+    finally:
+        rt.free(src)
+        rt.free(dst)
+        module.unload()
+
+
+def _votes(arch, case):
+    wave = ArchTarget.from_gfx(arch).wave_size
+    values = np.ones(wave, dtype=np.int32)
+    if case == "one_false_first":
+        values[0] = 0
+    elif case == "one_false_last":
+        values[-1] = 0
+    elif case == "all_false":
+        values[:] = 0
+    elif case == "nonzero_true":
+        values = np.arange(1, wave + 1, dtype=np.int32) * -3
+    return values
+
+
+_CASES = ("all_true", "nonzero_true", "one_false_first", "one_false_last", "all_false")
+
+
+@pytest.mark.parametrize("case", _CASES)
+def test_wave_all_is_wave_wide_and(arch, backend, case):
+    values = _votes(arch, case)
+    expected = np.full(values.shape, int(bool(np.all(values != 0))), dtype=np.int32)
+    np.testing.assert_array_equal(_run("wave_all", values, arch, backend), expected)
+
+
+@pytest.mark.parametrize("case", _CASES)
+def test_wave_any_is_wave_wide_or(arch, backend, case):
+    values = _votes(arch, case)
+    expected = np.full(values.shape, int(bool(np.any(values != 0))), dtype=np.int32)
+    np.testing.assert_array_equal(_run("wave_any", values, arch, backend), expected)

dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py

@@ -2737,19 +2737,36 @@ def test_to_sgpr_u32_emits_readfirstlane_and_pin(self):
         self.assertNotIn("asm sideeffect", ll)
 
     # ---- wave_all / wave_any ----
-    def test_wave_all_lowers_to_ballot_eq_minus_one(self):
+    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_all_smoke")
+        b = IRBuilder("wave_vote_smoke")
         b.kernel.attrs["max_workgroup_size"] = 64
         Out = b.param("Out", PtrType(I32, "global"))
-        all_t = b.wave_all(b.const_i32(1))
-        b.global_store(Out, b.const_i32(0), all_t)
-        ll = lower_kernel_to_llvm(b.kernel, arch="gfx950")
-        self.assertIn("@llvm.amdgcn.ballot.i64", ll)
-        # Compares ballot to -1 (i.e. all wave64 lanes voted true).
-        self.assertIn("icmp eq i64", ll)
+        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):