diff --git a/python/quadrants/lang/simt/subgroup.py b/python/quadrants/lang/simt/subgroup.py index edec8978d8..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) @@ -25,6 +54,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 @@ -168,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/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..1501938ff7 100644 --- a/quadrants/runtime/llvm/llvm_context.cpp +++ b/quadrants/runtime/llvm/llvm_context.cpp @@ -582,8 +582,122 @@ 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 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); + // 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_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 + // 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); + + // 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 (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, {mask_ty, builder.getInt32Ty()}, + {arg, zero32, icmp_ne}); + + // 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) { + auto *exec = builder.CreateIntrinsic( + llvm::Intrinsic::amdgcn_ballot, {mask_ty}, + {llvm::ConstantInt::getTrue(builder.getInt1Ty())}); + result_i1 = builder.CreateICmpEQ(ballot, exec); + } else { + 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); + }; + 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..56a877528b 100644 --- a/quadrants/runtime/llvm/runtime_module/runtime.cpp +++ b/quadrants/runtime/llvm/runtime_module/runtime.cpp @@ -1284,6 +1284,37 @@ int32 block_barrier_count_i32(int32 predicate) { return 0; } +// 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) +// +// 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; +} + +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..a0f1503a9d 100644 --- a/tests/python/test_simt.py +++ b/tests/python/test_simt.py @@ -470,6 +470,157 @@ 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 + + +@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():