From eaba867ac016035662b4883f2c96b70d3be57a34 Mon Sep 17 00:00:00 2001 From: lohiaj Date: Mon, 4 May 2026 11:05:30 +0000 Subject: [PATCH 1/4] feat(amdgpu): wave-scope subgroupOr / subgroupAnd / subgroupBarrier Adds wave-scope subgroup boolean reductions and barrier to the AMDGPU LLVM codegen path. `qd.simt.subgroup.reduce_or_i32(value)` and `reduce_and_i32(value)` now produce single-instruction-equivalent wave reductions on AMDGPU (amdgcn.icmp + ballot + s_or_b64 / s_and_b64 with EXEC after LLVM lowering), matching what the SPIRV codegen emits via OpGroupNonUniformBitwiseOr / And. `qd.simt.subgroup.barrier()` is also wired up on AMDGPU (was failing at codegen with `LLVMRuntime function subgroupBarrier not found`). Implementation follows the existing AMDGPU patch_intrinsic pattern in llvm_context.cpp (e.g. block_barrier -> amdgcn_s_barrier, cuda_shfl_down_sync_i32 -> amdgcn_ds_bpermute custom impl): 1. C++ stubs in runtime.cpp so the LLVM linker resolves the symbols 2. AMDGPU patch_intrinsic rewrites the stub bodies at codegen time using amdgcn.icmp.i32 (ballot of arg != 0), amdgcn.ballot(true) (EXEC mask), and icmp against 0 (OR) or against EXEC (AND) 3. Functions marked always_inline so the body inlines at the call site Generated AMDGCN at the call site is the expected 4-instruction sequence: v_cmp_ne_u32_e32 vcc, 0, v4 s_cmp_lg_u64 vcc, 0 s_cselect_b64 s[0:1], -1, 0 v_cndmask_b32_e64 v4, 0, 1, s[0:1] v1 limitations (acknowledged): - Boolean-only (i32 input treated as 0 vs non-zero). Generalising to bitwise OR/AND of arbitrary i32 values needs DPP-based tree reduction (~50 lines, clean follow-up). - AMDGPU-only. CUDA path falls back to no-op stubs (CUDA users already have cuda_ballot_i32 for this purpose). - OR / AND only. Add / Mul / Min / Max / Xor variants are straightforward extensions but kept out of v1 to keep PR scope small. Tests: 3 new AMDGPU-gated tests in tests/python/test_simt.py: - test_subgroup_reduce_or_i32_single_wave (block_dim=64, 1 wave) - test_subgroup_reduce_and_i32_single_wave (block_dim=64, 1 wave) - test_subgroup_reduce_or_i32_multi_wave (block_dim=128, 2 waves) explicitly verifies WAVE-SCOPE (not block-scope) semantics: a non-zero lane in wave 1 must NOT propagate to wave 0's outputs. All 3 pass. Targeted regression sweep (test_simt + test_basics + test_internal_func + test_shared_array_codegen_bug): 39 passed, 21 skipped, 0 failed. Zero behavior change for code that doesn't use the new primitives. --- python/quadrants/lang/simt/subgroup.py | 22 +++++ quadrants/inc/internal_ops.inc.h | 7 ++ quadrants/ir/type_system.cpp | 6 ++ quadrants/runtime/llvm/llvm_context.cpp | 58 ++++++++++++ .../runtime/llvm/runtime_module/runtime.cpp | 22 +++++ tests/python/test_simt.py | 89 +++++++++++++++++++ 6 files changed, 204 insertions(+) diff --git a/python/quadrants/lang/simt/subgroup.py b/python/quadrants/lang/simt/subgroup.py index edec8978d8..601362be0b 100644 --- a/python/quadrants/lang/simt/subgroup.py +++ b/python/quadrants/lang/simt/subgroup.py @@ -25,6 +25,28 @@ def any_true(cond): pass +def reduce_or_i32(value): + """Wave-scope OR reduction (boolean): returns 1 if any active lane + in the wave has value != 0, else 0. + + AMDGPU implementation: ``s_or_b64``-equivalent via ``amdgcn.icmp.i32`` + ballot, single-instruction at the AMDGCN level after LLVM lowering. + See ``quadrants/runtime/llvm/llvm_context.cpp`` for the patcher. + + Differs from the polymorphic ``reduce_or(value)`` (SPIRV-only at the + moment): this variant is i32-typed and treats input as boolean + (0 vs non-zero), sufficient for wave-vote convergence checks. + """ + return impl.call_internal("subgroupOr_i32", value, with_runtime_context=False) + + +def reduce_and_i32(value): + """Wave-scope AND reduction (boolean): returns 1 if all active lanes + in the wave have value != 0, else 0. See ``reduce_or_i32`` for context. + """ + return impl.call_internal("subgroupAnd_i32", value, with_runtime_context=False) + + def all_equal(value): # TODO pass diff --git a/quadrants/inc/internal_ops.inc.h b/quadrants/inc/internal_ops.inc.h index bef62c0d05..58296efe7a 100644 --- a/quadrants/inc/internal_ops.inc.h +++ b/quadrants/inc/internal_ops.inc.h @@ -72,5 +72,12 @@ PER_INTERNAL_OP(warp_barrier) // AMDGPU PER_INTERNAL_OP(amdgpu_clock_i64) +// Wave-scope subgroup primitives (LLVM-runtime path; AMDGPU-implemented in +// llvm_context.cpp via amdgcn.icmp/ballot intrinsics, no-op stubs on other +// LLVM archs for now). i32-typed boolean variants distinct from the +// polymorphic SPIRV `subgroupOr` / `subgroupAnd`. +PER_INTERNAL_OP(subgroupOr_i32) +PER_INTERNAL_OP(subgroupAnd_i32) + // CPU PER_INTERNAL_OP(cpu_clock_i64) diff --git a/quadrants/ir/type_system.cpp b/quadrants/ir/type_system.cpp index 97a6a33704..8e5225757a 100644 --- a/quadrants/ir/type_system.cpp +++ b/quadrants/ir/type_system.cpp @@ -339,6 +339,12 @@ void Operations::init_internals() { PLAIN_OP(block_barrier_or_i32, i32, false, i32); PLAIN_OP(block_barrier_count_i32, i32, false, i32); PLAIN_OP(grid_memfence, i32_void, false); + + // Wave-scope OR/AND boolean reductions. AMDGPU-implemented (see + // llvm_context.cpp); no-op stubs on other LLVM archs. + PLAIN_OP(subgroupOr_i32, i32, false, i32); + PLAIN_OP(subgroupAnd_i32, i32, false, i32); + CUDA_VOTE_SYNC(all); CUDA_VOTE_SYNC(any); CUDA_VOTE_SYNC(uni); diff --git a/quadrants/runtime/llvm/llvm_context.cpp b/quadrants/runtime/llvm/llvm_context.cpp index 2e03bc104e..2bebbe4674 100644 --- a/quadrants/runtime/llvm/llvm_context.cpp +++ b/quadrants/runtime/llvm/llvm_context.cpp @@ -582,8 +582,66 @@ std::unique_ptr QuadrantsLLVMContext::module_from_file( patch_intrinsic("block_idx", llvm::Intrinsic::amdgcn_workgroup_id_x); patch_intrinsic("block_barrier", llvm::Intrinsic::amdgcn_s_barrier, false); + // Wave-scope barrier (synonym for block_barrier on a wave-sized + // workgroup; on multi-wave workgroups it's an alias to s_barrier + // because AMDGPU doesn't expose a finer-grained intrinsic — the EXEC + // mask handles per-lane divergence within a wave automatically). + patch_intrinsic("subgroupBarrier", llvm::Intrinsic::amdgcn_s_barrier, + false); patch_intrinsic("amdgpu_clock_i64", llvm::Intrinsic::amdgcn_s_memtime); + // Wave-scope OR/AND boolean reduction. Treats the input i32 as boolean + // (0 vs non-zero) and returns 1 if any/all active lanes have val!=0, + // else 0. Implemented via amdgcn.icmp.i32 which produces a wave-wide + // ballot mask (i64 on wave64 / i32 on wave32). Compares the ballot + // result against zero (for OR / "any") or against EXEC (for AND / + // "all"). Single-instruction-equivalent at the AMDGCN level + // (s_or_b64 / s_and_b64 with EXEC) once LLVM lowers it. + // + // v1 limitation: only the boolean case (sufficient for wave-vote + // convergence checks on the CG iter loop). Generalising to bitwise + // OR/AND of arbitrary i32 values needs DPP-based tree reduction. + auto patch_subgroup_bool_reduce = [&](std::string name, bool is_and) { + auto func = module->getFunction(name); + if (!func) { + return; + } + func->deleteBody(); + auto bb = llvm::BasicBlock::Create(*ctx, "entry", func); + IRBuilder<> builder(*ctx); + builder.SetInsertPoint(bb); + auto *arg = &*func->arg_begin(); + auto *zero32 = llvm::ConstantInt::get(builder.getInt32Ty(), 0); + // amdgcn.icmp.i32(a, b, predicate) → i64 (on wave64) ballot mask. + // Predicate IDs match LLVM CmpInst::Predicate ordering: + // CmpInst::ICMP_NE = 33 (used here to ballot lanes where arg!=0) + auto *icmp_ne = llvm::ConstantInt::get(builder.getInt32Ty(), + llvm::CmpInst::ICMP_NE); + llvm::Value *ballot = builder.CreateIntrinsic( + llvm::Intrinsic::amdgcn_icmp, + {builder.getInt64Ty(), builder.getInt32Ty()}, + {arg, zero32, icmp_ne}); + // For OR: result = (ballot != 0). + // For AND: result = (ballot == EXEC). Read EXEC via the readlane + // intrinsic isn't quite right; use amdgcn.read.exec instead. + llvm::Value *result_i1; + if (is_and) { + // EXEC mask of currently-active lanes (i64 on wave64). + auto *exec = builder.CreateIntrinsic( + llvm::Intrinsic::amdgcn_ballot, + {builder.getInt64Ty()}, + {llvm::ConstantInt::getTrue(builder.getInt1Ty())}); + result_i1 = builder.CreateICmpEQ(ballot, exec); + } else { + auto *zero64 = llvm::ConstantInt::get(builder.getInt64Ty(), 0); + result_i1 = builder.CreateICmpNE(ballot, zero64); + } + builder.CreateRet(builder.CreateZExt(result_i1, builder.getInt32Ty())); + QuadrantsLLVMContext::mark_inline(func); + }; + patch_subgroup_bool_reduce("subgroupOr_i32", /*is_and=*/false); + patch_subgroup_bool_reduce("subgroupAnd_i32", /*is_and=*/true); + auto patch_amdgpu_shfl_down = [&](std::string name, bool is_float, bool has_mask) { auto func = module->getFunction(name); diff --git a/quadrants/runtime/llvm/runtime_module/runtime.cpp b/quadrants/runtime/llvm/runtime_module/runtime.cpp index a96aa8796e..e61c423917 100644 --- a/quadrants/runtime/llvm/runtime_module/runtime.cpp +++ b/quadrants/runtime/llvm/runtime_module/runtime.cpp @@ -1284,6 +1284,28 @@ int32 block_barrier_count_i32(int32 predicate) { return 0; } +// Wave-scope subgroup primitives. Stubs here; rewritten at LLVM IR codegen +// time by patch_intrinsic in llvm_context.cpp for each backend (CUDA/AMDGPU). +// The Python frontend (qd.simt.subgroup.reduce_or / reduce_and) targets +// these. Until 2026-05-04 these existed only on the SPIRV path (Vulkan/Metal) +// — the LLVM-runtime archs (AMDGPU/CUDA/CPU) had no implementation. +// +// Semantics (v1, boolean-only): subgroupOr_i32(val) returns 1 if any active +// lane in the wave has val != 0, else 0. subgroupAnd_i32 is the dual. +// Generalising to bitwise OR/AND of arbitrary i32 values would need DPP-based +// tree reduction; the boolean case is sufficient for wave-vote convergence +// checks on the CG iter loop. +int32 subgroupOr_i32(int32 val) { + return 0; +} + +int32 subgroupAnd_i32(int32 val) { + return 0; +} + +void subgroupBarrier() { +} + void warp_barrier(uint32 mask) { } diff --git a/tests/python/test_simt.py b/tests/python/test_simt.py index 948a21e58a..46f1493379 100644 --- a/tests/python/test_simt.py +++ b/tests/python/test_simt.py @@ -470,6 +470,95 @@ def foo(): assert a[i] == N - 1 +# ===================================================================== +# AMDGPU-specific wave-scope subgroup primitives. Implemented via the +# patch_intrinsic mechanism in quadrants/runtime/llvm/llvm_context.cpp, +# emitting amdgcn.icmp / amdgcn.ballot for single-instruction-equivalent +# wave reductions (s_or_b64 / s_and_b64 with EXEC after LLVM lowering). +# ===================================================================== + + +@test_utils.test(arch=qd.amdgpu) +def test_subgroup_reduce_or_i32_single_wave(): + N = 64 # one wave on AMDGPU + inputs = qd.field(qd.i32, shape=(N,)) + outputs = qd.field(qd.i32, shape=(N,)) + + @qd.kernel + def k(): + qd.loop_config(block_dim=N) + for i in range(N): + outputs[i] = qd.simt.subgroup.reduce_or_i32(inputs[i]) + + # All zero → OR is 0 in every lane. + for i in range(N): + inputs[i] = 0 + k() + for i in range(N): + assert outputs[i] == 0 + + # One lane non-zero → OR is 1 in every lane. + inputs[7] = 1 + k() + for i in range(N): + assert outputs[i] == 1 + + +@test_utils.test(arch=qd.amdgpu) +def test_subgroup_reduce_and_i32_single_wave(): + N = 64 + inputs = qd.field(qd.i32, shape=(N,)) + outputs = qd.field(qd.i32, shape=(N,)) + + @qd.kernel + def k(): + qd.loop_config(block_dim=N) + for i in range(N): + outputs[i] = qd.simt.subgroup.reduce_and_i32(inputs[i]) + + # All non-zero → AND is 1. + for i in range(N): + inputs[i] = 1 + k() + for i in range(N): + assert outputs[i] == 1 + + # One lane zero → AND is 0. + inputs[31] = 0 + k() + for i in range(N): + assert outputs[i] == 0 + + +@test_utils.test(arch=qd.amdgpu) +def test_subgroup_reduce_or_i32_multi_wave(): + """Verifies wave-scope semantics: with block_dim=128, two waves of 64 + each must reduce independently. A single non-zero lane in wave 1 must + NOT propagate to wave 0's outputs. + """ + N = 128 + inputs = qd.field(qd.i32, shape=(N,)) + outputs = qd.field(qd.i32, shape=(N,)) + + @qd.kernel + def k(): + qd.loop_config(block_dim=N) + for i in range(N): + outputs[i] = qd.simt.subgroup.reduce_or_i32(inputs[i]) + + # Set lane 64 (wave 1's first lane) to 1, all others 0. + for i in range(N): + inputs[i] = 0 + inputs[64] = 1 + k() + # Wave 0 (lanes 0-63) sees no non-zero → OR=0. + for i in range(64): + assert outputs[i] == 0 + # Wave 1 (lanes 64-127) sees the non-zero → OR=1 in every lane. + for i in range(64, N): + assert outputs[i] == 1 + + # TODO: replace this with a stronger test case @test_utils.test(arch=qd.cuda) def test_grid_memfence(): From e46931bd7d949941697d18bf5c6c8281f51be86c Mon Sep 17 00:00:00 2001 From: lohiaj Date: Wed, 6 May 2026 16:54:05 +0000 Subject: [PATCH 2/4] fix(amdgpu): use wave_barrier (not s_barrier) for subgroupBarrier s_barrier is the workgroup-scope barrier and would deadlock when subgroupBarrier() is called from non-uniform control flow with block_dim > wave_size -- non-participating waves never arrive at the barrier (per Dipto's review). subgroupBarrier is meant to match SPIR-V OpControlBarrier with Subgroup execution scope. Switch the lowering to llvm.amdgcn.wave.barrier: a discardable compiler + memory-ordering barrier scoped to the current wave. It emits no hardware sync (the 64 lanes of an AMDGPU wave already execute in lockstep) and carries IntrConvergent + IntrHasSideEffects to prevent illegal CFG transforms and memory-op reordering. Add two AMDGPU regression tests in tests/python/test_simt.py: - test_subgroup_barrier_partial_participation: exercises the exact deadlock pattern from the review thread (block_dim=128, only lanes with `if i < 64:` call the barrier). With s_barrier this hung; with wave.barrier it completes. - test_subgroup_barrier_in_uniform_path: read-after-write across the barrier, verifies compiler ordering is preserved. Also addresses vecheruk-amd's review nits: - Add reduce_or_i32 / reduce_and_i32 to subgroup.py __all__. - Document that non-AMDGPU LLVM-runtime backends (CPU, CUDA today) fall through to the runtime stubs (return 0 / no-op) and that Vulkan/Metal users should keep using the polymorphic reduce_or / reduce_and. Tests: AMDGPU subgroup tests 5/5 pass. Targeted regression sweep (test_simt + test_basics + test_internal_func + test_shared_array_codegen_bug) clean: 41 passed, 21 skipped, 0 failed. --- python/quadrants/lang/simt/subgroup.py | 31 ++++++++++ quadrants/runtime/llvm/llvm_context.cpp | 28 +++++++-- .../runtime/llvm/runtime_module/runtime.cpp | 29 ++++++--- tests/python/test_simt.py | 62 +++++++++++++++++++ 4 files changed, 134 insertions(+), 16 deletions(-) diff --git a/python/quadrants/lang/simt/subgroup.py b/python/quadrants/lang/simt/subgroup.py index 601362be0b..e6df5151f7 100644 --- a/python/quadrants/lang/simt/subgroup.py +++ b/python/quadrants/lang/simt/subgroup.py @@ -1,9 +1,38 @@ # type: ignore +"""Wave-scope (a.k.a. SPIR-V "Subgroup") primitives. + +Backend availability for the wave-scope reductions added in this module: + +- AMDGPU: full implementation via ``patch_intrinsic`` (see + ``quadrants/runtime/llvm/llvm_context.cpp``); ``reduce_or_i32`` / + ``reduce_and_i32`` lower to ``amdgcn.icmp`` + ``amdgcn.ballot`` + (single-instruction-equivalent at AMDGCN level), and ``barrier`` lowers + to ``llvm.amdgcn.wave.barrier`` (compiler + memory-ordering barrier + scoped to the current wave; no hardware sync, no deadlock hazard on + partial-wave participation). +- Other LLVM-runtime backends (CPU, CUDA today): no patcher yet — + ``reduce_or_i32`` / ``reduce_and_i32`` fall through to the runtime stubs + in ``quadrants/runtime/llvm/runtime_module/runtime.cpp`` which return 0. + This is intentional placeholder behaviour for backends where wave-scope + semantics aren't wired up; callers that need cross-backend portability + should gate on ``qd.cfg.arch == qd.amdgpu`` for now. +- SPIR-V backends (Vulkan, Metal): the polymorphic ``reduce_or`` / + ``reduce_and`` (no ``_i32`` suffix) cover those targets. +""" from quadrants.lang import impl def barrier(): + """Wave-scope (subgroup) execution + memory-ordering barrier. + + On AMDGPU lowers to ``llvm.amdgcn.wave.barrier``: no hardware wait + (all 64 lanes of an AMDGPU wave already execute in lockstep) but + prevents the compiler from reordering memory ops or moving + convergent ops across this point. Safe to call from non-uniform + control flow — does NOT synchronize across waves in a workgroup. + For workgroup-scope synchronization use ``qd.simt.block.sync()``. + """ return impl.call_internal("subgroupBarrier", with_runtime_context=False) @@ -190,7 +219,9 @@ def shuffle_down(value, offset): "reduce_min", "reduce_max", "reduce_and", + "reduce_and_i32", "reduce_or", + "reduce_or_i32", "reduce_xor", "inclusive_add", "inclusive_mul", diff --git a/quadrants/runtime/llvm/llvm_context.cpp b/quadrants/runtime/llvm/llvm_context.cpp index 2bebbe4674..5c5b3d6a47 100644 --- a/quadrants/runtime/llvm/llvm_context.cpp +++ b/quadrants/runtime/llvm/llvm_context.cpp @@ -582,12 +582,28 @@ std::unique_ptr QuadrantsLLVMContext::module_from_file( patch_intrinsic("block_idx", llvm::Intrinsic::amdgcn_workgroup_id_x); patch_intrinsic("block_barrier", llvm::Intrinsic::amdgcn_s_barrier, false); - // Wave-scope barrier (synonym for block_barrier on a wave-sized - // workgroup; on multi-wave workgroups it's an alias to s_barrier - // because AMDGPU doesn't expose a finer-grained intrinsic — the EXEC - // mask handles per-lane divergence within a wave automatically). - patch_intrinsic("subgroupBarrier", llvm::Intrinsic::amdgcn_s_barrier, - false); + // Wave-scope subgroup barrier — matches SPIR-V OpControlBarrier with + // Subgroup execution scope. Lowers to llvm.amdgcn.wave.barrier, a + // discardable barrier that carries IntrConvergent + IntrHasSideEffects: + // - IntrConvergent prevents the compiler from hoisting/sinking + // convergent ops (v_readfirstlane, ds_swizzle, ballots, the + // subgroupOr/subgroupAnd primitives below) across this point + // or moving them in/out of divergent control flow. + // - IntrHasSideEffects prevents memory-op reordering across it. + // No hardware barrier instruction is emitted: the 64 lanes of an + // AMDGPU wave already execute in lockstep, so there is nothing for + // the hardware to wait on. See LLVM IntrinsicsAMDGPU.td: + // def int_amdgcn_wave_barrier ... + // + // We deliberately do NOT use amdgcn_s_barrier here. s_barrier is the + // workgroup-scope barrier (waits for every wave in the workgroup to + // arrive) and would deadlock on partial-wave participation, e.g. + // calling qd.simt.subgroup.barrier() inside `if lane < 64:` with + // block_dim=128 — wave 1 never reaches the barrier, so wave 0 hangs + // forever waiting on it. wave_barrier has no such hazard because it + // emits no actual hardware wait. + patch_intrinsic("subgroupBarrier", + llvm::Intrinsic::amdgcn_wave_barrier, false); patch_intrinsic("amdgpu_clock_i64", llvm::Intrinsic::amdgcn_s_memtime); // Wave-scope OR/AND boolean reduction. Treats the input i32 as boolean diff --git a/quadrants/runtime/llvm/runtime_module/runtime.cpp b/quadrants/runtime/llvm/runtime_module/runtime.cpp index e61c423917..56a877528b 100644 --- a/quadrants/runtime/llvm/runtime_module/runtime.cpp +++ b/quadrants/runtime/llvm/runtime_module/runtime.cpp @@ -1284,17 +1284,26 @@ int32 block_barrier_count_i32(int32 predicate) { return 0; } -// Wave-scope subgroup primitives. Stubs here; rewritten at LLVM IR codegen -// time by patch_intrinsic in llvm_context.cpp for each backend (CUDA/AMDGPU). -// The Python frontend (qd.simt.subgroup.reduce_or / reduce_and) targets -// these. Until 2026-05-04 these existed only on the SPIRV path (Vulkan/Metal) -// — the LLVM-runtime archs (AMDGPU/CUDA/CPU) had no implementation. +// Wave-scope subgroup primitives. These bodies are placeholder stubs +// (return 0 / no-op) for LLVM-runtime backends that don't wire up a real +// implementation. AMDGPU rewrites them at LLVM-IR codegen time via +// patch_intrinsic in quadrants/runtime/llvm/llvm_context.cpp: +// - subgroupOr_i32 / subgroupAnd_i32 → amdgcn.icmp + amdgcn.ballot +// (single-instruction-equivalent at AMDGCN level after lowering) +// - subgroupBarrier → llvm.amdgcn.wave.barrier (discardable barrier; +// compiler + memory-ordering only, emits no hardware instruction) // -// Semantics (v1, boolean-only): subgroupOr_i32(val) returns 1 if any active -// lane in the wave has val != 0, else 0. subgroupAnd_i32 is the dual. -// Generalising to bitwise OR/AND of arbitrary i32 values would need DPP-based -// tree reduction; the boolean case is sufficient for wave-vote convergence -// checks on the CG iter loop. +// CUDA / CPU LLVM-runtime backends currently fall through to the stubs +// below and get backend-limited behaviour (return 0 / no sync). Wiring +// them up for those backends is left for follow-up PRs. Vulkan / Metal +// users should keep using the existing polymorphic subgroupOr / +// subgroupAnd in subgroup.py. +// +// Semantics (v1, boolean-only): subgroupOr_i32(val) returns 1 if any +// active lane in the wave has val != 0, else 0. subgroupAnd_i32 is the +// dual. Generalising to bitwise OR/AND of arbitrary i32 values would +// need DPP-based tree reduction; the boolean case is sufficient for +// wave-vote convergence checks. int32 subgroupOr_i32(int32 val) { return 0; } diff --git a/tests/python/test_simt.py b/tests/python/test_simt.py index 46f1493379..a0f1503a9d 100644 --- a/tests/python/test_simt.py +++ b/tests/python/test_simt.py @@ -559,6 +559,68 @@ def k(): assert outputs[i] == 1 +@test_utils.test(arch=qd.amdgpu) +def test_subgroup_barrier_partial_participation(): + """Regression test for the s_barrier-vs-wave_barrier fix. + + On AMDGPU, subgroupBarrier() must lower to llvm.amdgcn.wave.barrier + (wave-scope, no hardware sync) and NOT to llvm.amdgcn.s.barrier + (workgroup-scope, waits for every wave to arrive). Calling a + workgroup-scope barrier from non-uniform control flow with + block_dim > wave_size deadlocks the workgroup, because non- + participating waves never reach the barrier. + + The kernel below uses block_dim=128 (two wave64 waves) and only + calls subgroupBarrier() on the first half of lanes (i < 64, i.e. + only wave 0). With s_barrier this hangs forever; with wave.barrier + it completes and the assertions verify that the kernel produced + the correct result. + """ + N = 128 # 2 waves of 64 + out = qd.field(qd.i32, shape=(N,)) + + @qd.kernel + def k(): + qd.loop_config(block_dim=N) + for i in range(N): + if i < 64: + # Only wave 0 hits this barrier. Workgroup-scope + # s_barrier would deadlock here; wave.barrier is safe. + qd.simt.subgroup.barrier() + out[i] = 1 + else: + out[i] = 2 + + k() + for i in range(64): + assert out[i] == 1 + for i in range(64, N): + assert out[i] == 2 + + +@test_utils.test(arch=qd.amdgpu) +def test_subgroup_barrier_in_uniform_path(): + """subgroupBarrier called by every lane in every wave — should + behave the same as the partial-participation variant (no hardware + sync, just compiler + memory ordering).""" + N = 128 + out = qd.field(qd.i32, shape=(N,)) + + @qd.kernel + def k(): + qd.loop_config(block_dim=N) + for i in range(N): + out[i] = i + qd.simt.subgroup.barrier() + # Read-after-write within the same lane: barrier prevents + # the compiler from reordering the store and load across it. + out[i] = out[i] * 2 + + k() + for i in range(N): + assert out[i] == 2 * i + + # TODO: replace this with a stronger test case @test_utils.test(arch=qd.cuda) def test_grid_memfence(): From 2c2300259bf92be567e2a6079c12374476619aac Mon Sep 17 00:00:00 2001 From: lohiaj Date: Wed, 6 May 2026 17:10:52 +0000 Subject: [PATCH 3/4] style(amdgpu): clang-format the new subgroupBarrier patch_intrinsic call Pre-commit `clang-format` v19.1.7 wants patch_intrinsic("subgroupBarrier", llvm::Intrinsic::amdgcn_wave_barrier, false); instead of the wrap I had. Pure whitespace; semantically identical. --- quadrants/runtime/llvm/llvm_context.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/quadrants/runtime/llvm/llvm_context.cpp b/quadrants/runtime/llvm/llvm_context.cpp index 5c5b3d6a47..11cdb7ada7 100644 --- a/quadrants/runtime/llvm/llvm_context.cpp +++ b/quadrants/runtime/llvm/llvm_context.cpp @@ -602,8 +602,8 @@ std::unique_ptr QuadrantsLLVMContext::module_from_file( // block_dim=128 — wave 1 never reaches the barrier, so wave 0 hangs // forever waiting on it. wave_barrier has no such hazard because it // emits no actual hardware wait. - patch_intrinsic("subgroupBarrier", - llvm::Intrinsic::amdgcn_wave_barrier, false); + patch_intrinsic("subgroupBarrier", llvm::Intrinsic::amdgcn_wave_barrier, + false); patch_intrinsic("amdgpu_clock_i64", llvm::Intrinsic::amdgcn_s_memtime); // Wave-scope OR/AND boolean reduction. Treats the input i32 as boolean From 427919ea3f1c1058c78b98322e14784a21acd8ea Mon Sep 17 00:00:00 2001 From: lohiaj Date: Thu, 7 May 2026 06:57:11 +0000 Subject: [PATCH 4/4] feat(amdgpu): predicate subgroupOr/subgroupAnd lowering on wave size Per Dipto's review on PR #26: keep the implementation flexible enough to be predicated on wave size so it is upstream-friendly (the community cares about RDNA, all wave32) and forward-compatible with MI450 (expected wave32 per AMD). The wave-vote ballot intrinsics (`llvm.amdgcn.icmp` and `llvm.amdgcn.ballot`) are overloaded by their return-type argument: i64 on wave64 targets (CDNA1/2/3, gfx9xx) and i32 on wave32 targets (RDNA, expected MI450). The previous lowering hard-coded i64, which would either fail to lower or produce wrong code on a wave32 target. Resolution at codegen time: unsigned get_amdgpu_wave_size(llvm::Function *F): 1. If F has a `+wavefrontsize32` / `+wavefrontsize64` feature in its target-features attribute, honor that. 2. Else fall back to the mcpu-based default: gfx9xx -> 64, gfx10xx+ -> 32. Quadrants's AMDGPU codegen path currently sets target-features to "" (see llvm_context_pass.h:88), so step 2 is what's actually used today. The fallback list assumes `gfx9xx` is always wave64 -- true for every shipped MI300/MI325/MI355 part. If MI450 ships under a `gfx9...` mcpu prefix yet defaults to wave32, Quadrants's AMDGPU codegen path should also be updated to set "+wavefrontsize32" explicitly so step 1 picks it up; that's outside this PR's scope. The mask type is then used uniformly: - llvm.amdgcn.icmp.{mask_ty}.i32(arg, 0, ICMP_NE) -> ballot - For OR: ballot != 0 - For AND: ballot == amdgcn.ballot.{mask_ty}(true) (active-lane mask) No semantic change on wave64 hardware -- the i64 path is identical to what landed in eaba867ac. The new branch is purely additive for wave32. Tests - All 5 AMDGPU subgroup tests pass on MI300X/wave64: test_subgroup_reduce_or_i32_single_wave test_subgroup_reduce_and_i32_single_wave test_subgroup_reduce_or_i32_multi_wave test_subgroup_barrier_partial_participation test_subgroup_barrier_in_uniform_path - All 5 Vulkan subgroup tests pass. - Targeted regression sweep (test_simt + test_basics + test_internal_func + test_shared_array_codegen_bug): 41 passed, 21 skipped, 0 failed. - Wave32 path is correct-by-construction (symmetric to wave64) but untested locally -- no RDNA hardware in the bench setup. --- quadrants/runtime/llvm/llvm_context.cpp | 66 ++++++++++++++++++++----- 1 file changed, 53 insertions(+), 13 deletions(-) diff --git a/quadrants/runtime/llvm/llvm_context.cpp b/quadrants/runtime/llvm/llvm_context.cpp index 11cdb7ada7..1501938ff7 100644 --- a/quadrants/runtime/llvm/llvm_context.cpp +++ b/quadrants/runtime/llvm/llvm_context.cpp @@ -606,13 +606,45 @@ std::unique_ptr QuadrantsLLVMContext::module_from_file( false); patch_intrinsic("amdgpu_clock_i64", llvm::Intrinsic::amdgcn_s_memtime); + // Determine the AMDGPU wave (subgroup) size for the current target. + // The wave-vote ballot intrinsics return an i32 mask on wave32 targets + // (RDNA, expected MI450) and an i64 mask on wave64 targets (CDNA1/2/3 + // — gfx9xx). Pick the type from the function's `target-features` + // attribute first; if not set, fall back to the mcpu-based default. + // + // Per Dipto's review on PR #26: keep this predicated on wavesize so + // the lowering is upstream-friendly (the community cares about RDNA, + // all wave32) and forward-compatible with MI450. + auto get_amdgpu_wave_size = [](llvm::Function *F) -> unsigned { + auto fn_attrs = F->getAttributes().getFnAttrs(); + if (fn_attrs.hasAttribute("target-features")) { + auto feats = + fn_attrs.getAttribute("target-features").getValueAsString(); + if (feats.contains("+wavefrontsize32")) + return 32; + if (feats.contains("+wavefrontsize64")) + return 64; + } +#if defined(QD_WITH_AMDGPU) + // Quadrants's AMDGPU codegen path sets target-features to "" and + // relies on the mcpu default. gfx9xx (CDNA1/2/3) is wave64 by + // default; gfx10xx+ (RDNA, future CDNA-Next) is wave32 by default. + const std::string &mcpu = AMDGPUContext::get_instance().get_mcpu(); + if (mcpu.size() >= 4 && mcpu[3] == '9') + return 64; + return 32; +#else + return 64; +#endif + }; + // Wave-scope OR/AND boolean reduction. Treats the input i32 as boolean // (0 vs non-zero) and returns 1 if any/all active lanes have val!=0, // else 0. Implemented via amdgcn.icmp.i32 which produces a wave-wide // ballot mask (i64 on wave64 / i32 on wave32). Compares the ballot // result against zero (for OR / "any") or against EXEC (for AND / // "all"). Single-instruction-equivalent at the AMDGCN level - // (s_or_b64 / s_and_b64 with EXEC) once LLVM lowers it. + // (s_or_b{32,64} / s_and_b{32,64} with EXEC) once LLVM lowers it. // // v1 limitation: only the boolean case (sufficient for wave-vote // convergence checks on the CG iter loop). Generalising to bitwise @@ -628,29 +660,37 @@ std::unique_ptr QuadrantsLLVMContext::module_from_file( builder.SetInsertPoint(bb); auto *arg = &*func->arg_begin(); auto *zero32 = llvm::ConstantInt::get(builder.getInt32Ty(), 0); - // amdgcn.icmp.i32(a, b, predicate) → i64 (on wave64) ballot mask. + + // Pick the ballot mask type to match the wave size of the target. + // amdgcn.icmp / amdgcn.ballot are overloaded by their return type; + // the type MUST match the subtarget's wave size or codegen will + // either fail to lower or produce wrong results. + unsigned wave_size = get_amdgpu_wave_size(func); + llvm::Type *mask_ty = + (wave_size == 64) ? builder.getInt64Ty() : builder.getInt32Ty(); + + // amdgcn.icmp.i32(a, b, predicate) → mask_ty ballot mask. // Predicate IDs match LLVM CmpInst::Predicate ordering: - // CmpInst::ICMP_NE = 33 (used here to ballot lanes where arg!=0) + // CmpInst::ICMP_NE (used here to ballot lanes where arg!=0) auto *icmp_ne = llvm::ConstantInt::get(builder.getInt32Ty(), llvm::CmpInst::ICMP_NE); llvm::Value *ballot = builder.CreateIntrinsic( - llvm::Intrinsic::amdgcn_icmp, - {builder.getInt64Ty(), builder.getInt32Ty()}, + llvm::Intrinsic::amdgcn_icmp, {mask_ty, builder.getInt32Ty()}, {arg, zero32, icmp_ne}); - // For OR: result = (ballot != 0). - // For AND: result = (ballot == EXEC). Read EXEC via the readlane - // intrinsic isn't quite right; use amdgcn.read.exec instead. + + // For OR ("any"): result = (ballot != 0). + // For AND ("all"): result = (ballot == EXEC), where EXEC is read + // via amdgcn.ballot of `true` (the canonical idiom + // for "currently-active lane mask"). llvm::Value *result_i1; if (is_and) { - // EXEC mask of currently-active lanes (i64 on wave64). auto *exec = builder.CreateIntrinsic( - llvm::Intrinsic::amdgcn_ballot, - {builder.getInt64Ty()}, + llvm::Intrinsic::amdgcn_ballot, {mask_ty}, {llvm::ConstantInt::getTrue(builder.getInt1Ty())}); result_i1 = builder.CreateICmpEQ(ballot, exec); } else { - auto *zero64 = llvm::ConstantInt::get(builder.getInt64Ty(), 0); - result_i1 = builder.CreateICmpNE(ballot, zero64); + auto *mask_zero = llvm::ConstantInt::get(mask_ty, 0); + result_i1 = builder.CreateICmpNE(ballot, mask_zero); } builder.CreateRet(builder.CreateZExt(result_i1, builder.getInt32Ty())); QuadrantsLLVMContext::mark_inline(func);