diff --git a/python/quadrants/lang/simt/block.py b/python/quadrants/lang/simt/block.py index 96c94537b7..e65b83f4cd 100644 --- a/python/quadrants/lang/simt/block.py +++ b/python/quadrants/lang/simt/block.py @@ -23,7 +23,6 @@ from quadrants.lang.util import quadrants_scope from quadrants.types.annotations import template from quadrants.types.primitive_types import i32 as _i32 -from quadrants.types.primitive_types import i64 as _i64 from quadrants.types.primitive_types import u32 as _u32 from quadrants.types.primitive_types import u64 as _u64 @@ -263,6 +262,49 @@ def reduce_all_max(value, block_dim: template(), dtype: template()): # input - we recover the inclusive total with one extra `op` on the publish path. +@_func +def _block_scan_lds(value, block_dim: template(), op: template(), identity, dtype: template(), exclusive: template()): + """Permlane-free block-scope scan: double-buffered Hillis-Steele over a ``2 * block_dim`` LDS scratch. + + AMDGPU-only path. The default subgroup-shuffle + cross-subgroup-fold scan (see ``inclusive_scan`` / + ``exclusive_scan`` below) relies on the wave64 cross-half shuffle, whose ``ds_bpermute`` + ``permlane64`` + LDS-roundtrip emulation races under heavy-kernel instruction scheduling at high wave counts (observable from + ``BLOCK_DIM == 512`` / 8 waves, where the tiled scan that feeds the fold returns nondeterministic garbage; the + native ``v_permlane64_b32`` that would avoid the emulation crashes the bundled LLVM 22.1.0 AMDGPU backend). This + pure-LDS scan issues no cross-lane shuffle at all -- only ``block.sync()`` barriers and double-buffered LDS + reads/writes -- so it is correct and deterministic at every ``block_dim`` and every wave count. + + Algorithm: stage ``value`` into ``buf[0 : block_dim]``; then for ``offset = 1, 2, 4, ... >= block_dim`` read from + the current half and write ``op(buf[tid - offset], buf[tid])`` (or the passthrough for ``tid < offset``) into the + *other* half, swapping halves each step so there is never an in-place read/write hazard. ``exclusive`` shifts the + final inclusive result right by one lane (lane 0 := ``identity``). ``ceil(log2(block_dim))`` steps; works for any + ``block_dim`` (need not be a power of two). + """ + NSTEPS = impl.static((block_dim - 1).bit_length()) + tid = thread_idx() + buf = SharedArray(impl.static((2 * block_dim,)), dtype) + buf[tid] = value + sync() + # Statically unrolled Hillis-Steele. ``cur``/``oth``/``off`` are derived from the (compile-time) step index so + # they stay Python ints -- a runtime ``for`` loop would turn them into ``Expr`` and break ``impl.static``. + for s in impl.static(range(NSTEPS)): + cur_base = impl.static((s % 2) * block_dim) + oth_base = impl.static(((s + 1) % 2) * block_dim) + off = impl.static(1 << s) + x = buf[cur_base + tid] + if tid >= off: + x = op(buf[cur_base + tid - off], x) + buf[oth_base + tid] = x + sync() + final_base = impl.static((NSTEPS % 2) * block_dim) + if impl.static(exclusive): + result = identity + if tid > 0: + result = buf[final_base + tid - 1] + return result + return buf[final_base + tid] + + @_func def inclusive_scan(value, block_dim: template(), op: template(), dtype: template()): """Block-scope inclusive scan under a generic associative ``op``. Every thread receives a valid result. @@ -288,6 +330,11 @@ def inclusive_scan(value, block_dim: template(), op: template(), dtype: template ) NUM_SUBGROUPS = impl.static(block_dim // SUBGROUP_SIZE) + # AMDGPU: use the permlane-free pure-LDS scan -- the wave64 cross-half shuffle emulation races at high wave counts + # (see ``_block_scan_lds``). ``identity`` is unused on the inclusive path; ``value`` is just a typed placeholder. + if impl.static(impl.get_runtime().prog.config().arch == _qd_core.amdgpu): + return _block_scan_lds(value, block_dim, op, value, dtype, False) + inclusive = _reductions._inclusive_scan_tiled(value, op, log2_subgroup) if impl.static(NUM_SUBGROUPS == 1): @@ -336,6 +383,11 @@ def exclusive_scan(value, block_dim: template(), op: template(), identity, dtype ) NUM_SUBGROUPS = impl.static(block_dim // SUBGROUP_SIZE) + # AMDGPU: use the permlane-free pure-LDS scan -- the wave64 cross-half shuffle emulation races at high wave counts + # (see ``_block_scan_lds``). + if impl.static(impl.get_runtime().prog.config().arch == _qd_core.amdgpu): + return _block_scan_lds(value, block_dim, op, identity, dtype, True) + exclusive = _reductions._exclusive_scan_tiled(value, op, identity, log2_subgroup) if impl.static(NUM_SUBGROUPS == 1): @@ -602,18 +654,20 @@ def _radix_rank_match_atomic_or_wave64( NUM_BITS_MASK = impl.static((1 << num_bits) - 1) BINS_PER_LANE = impl.static(RADIX_DIGITS // SUBGROUP_THREADS) + # No ``smem_match`` region on wave64: the per-digit peer mask is derived from subgroup ballots (step 5), not from + # an i64 LDS atomic_or cell. The old atomic_or-on-i64-LDS + u64-popcnt match miscompiles on the wave64 AMDGPU + # backend (garbage ranks -> OOB scatter -> GPU fault in device_radix_sort n=1M); the ballot form sidesteps LDS + # atomics entirely. This also frees the 8 KiB i64 LDS region the match cell used to need at radix_bits=8. smem_offsets = SharedArray(impl.static((BLOCK_SUBGROUPS * RADIX_DIGITS,)), _i32) - smem_match = SharedArray(impl.static((BLOCK_SUBGROUPS * RADIX_DIGITS,)), _i64) tid = thread_idx() subgroup_idx = tid // SUBGROUP_THREADS lane = _ops.cast(_subgroup.invocation_id(), _i32) - # Step 1: zero per-subgroup histograms and match_masks. + # Step 1: zero per-subgroup histograms. for b in impl.static(range(BINS_PER_LANE)): bin_idx = lane + impl.static(b * SUBGROUP_THREADS) smem_offsets[subgroup_idx * RADIX_DIGITS + bin_idx] = _i32(0) - smem_match[subgroup_idx * RADIX_DIGITS + bin_idx] = _i64(0) _subgroup_sync_fence() digit = _ops.cast(_ops.bit_and(_ops.bit_shr(key, _u32(bit_start)), _u32(NUM_BITS_MASK)), _i32) @@ -629,6 +683,15 @@ def _radix_rank_match_atomic_or_wave64( exclusive_digit_prefix = exclusive_add(bin_count, block_dim, _i32) + # Publish the per-digit total + exclusive prefix to the caller's outparams *now*, while ``bin_count`` / + # ``exclusive_digit_prefix`` are still live and known-correct. Deferring these writes until after steps 4-5 (as + # the original did) makes their source registers span the downsweep + match phase; under the wave64 AMDGPU + # backend that long live range gets clobbered (observed as garbage ``bins`` / ``excl_prefix`` -> wrong scatter + # offsets). ``bins`` / ``excl_prefix`` are caller-owned and untouched by steps 4-5, so an early write is safe and + # the trailing ``sync()`` still guarantees visibility. + bins[tid] = bin_count + excl_prefix[tid] = exclusive_digit_prefix + for j_subgroup in impl.static(range(BLOCK_SUBGROUPS)): smem_offsets[impl.static(j_subgroup * RADIX_DIGITS) + tid] = ( smem_offsets[impl.static(j_subgroup * RADIX_DIGITS) + tid] + exclusive_digit_prefix @@ -636,37 +699,40 @@ def _radix_rank_match_atomic_or_wave64( sync() - # Step 5 - wave64 specifics: u64 ballot mask via inline ``one_at_lane | (one_at_lane - 1)`` (avoids UB on - # lane=63), atomic_or on the i64 match cell, clz / popcnt on u64. Leader formula is ``63 - clz(u64)``. + # Step 5 - wave64 match-and-rank via subgroup ballots (CUB ``MatchAny``). Instead of OR-ing each lane's bit into + # a shared i64 cell and reading it back, every lane reconstructs its per-digit peer mask -- the set of lanes in its + # subgroup that share ``digit`` -- with ``num_bits`` subgroup ballots: + # + # peers = &_{b in [0,num_bits)} ( bit_b(digit) ? ballot(bit_b==1) : ~ballot(bit_b==1) ) + # + # ``subgroup.ballot`` lowers to ``llvm.amdgcn.ballot.i64`` (full 64-lane mask, isel-clean on wave64 -- unlike the + # i32 ballot and unlike the i64 LDS atomic_or this replaces). No LDS atomics and no cross-lane shuffle, so neither + # the i64-atomic_or miscompile nor the divergent ``ds_bpermute`` emulation bug can be hit. + # + # The branch on ``bit_b(digit)`` is done branchlessly: ``xor_mask = u64(bit) - 1`` is ``0`` when ``bit==1`` and + # all-ones (0xFFFF...F) when ``bit==0`` (unsigned wraparound), so ``ballot_b ^ xor_mask`` yields ``ballot_b`` or + # its complement without divergent control flow. + digit_u32 = _ops.cast(digit, _u32) + peers = _u64(0xFFFFFFFFFFFFFFFF) + for bit_b in impl.static(range(num_bits)): + bit_val = _ops.bit_and(_ops.bit_shr(digit_u32, _u32(impl.static(bit_b))), _u32(1)) + ballot_b = _ops.cast(_subgroup.ballot(_ops.cast(bit_val, _i32)), _u64) + xor_mask = _ops.cast(bit_val, _u64) - _u64(1) + peers = _ops.bit_and(peers, _ops.bit_xor(ballot_b, xor_mask)) + + # ``popc`` = 1-based position of this lane among its peers at indices ``<= lane`` (lowest matching lane -> 1). lane_u64 = _ops.cast(lane, _u64) - lane_mask = _u64(1) << lane_u64 - lane_mask_le_v = lane_mask | (lane_mask - _u64(1)) - - match_idx = subgroup_idx * RADIX_DIGITS + digit - - _ops.atomic_or(smem_match[match_idx], _ops.cast(lane_mask, _i64)) - _subgroup_sync_fence() - - # u64 clz via FindUMsb-equivalent on every backend; the wave32 path's caveat about FindSMsb vs FindUMsb on i64 - # would apply on SPIR-V wave64 devices if those existed (today wave64 = AMDGPU only). - bin_mask = _ops.cast(smem_match[match_idx], _u64) - leader = _i32(63) - _ops.cast(_ops.clz(bin_mask), _i32) - popc = _ops.popcnt(_ops.bit_and(bin_mask, lane_mask_le_v)) - - subgroup_offset = _i32(0) - if lane == leader: - subgroup_offset = _ops.atomic_add(smem_offsets[subgroup_idx * RADIX_DIGITS + digit], _ops.cast(popc, _i32)) + one_at_lane = _u64(1) << lane_u64 + lane_mask_le_v = one_at_lane | (one_at_lane - _u64(1)) + popc = _ops.popcnt(_ops.bit_and(peers, lane_mask_le_v)) - subgroup_offset = _subgroup.shuffle(subgroup_offset, _ops.cast(leader, _u32)) + # Every lane of a (subgroup, digit) group shares one offset cell, ``smem_offsets[sg*RD + digit]``, which already + # holds that group's block-wide base offset after the step-4 downsweep + the ``sync()`` above. Each lane reads it + # directly and derives its rank as ``base + popc - 1`` -- so the lowest matching lane gets ``base + 0``, the next + # ``base + 1``, ... ``items_per_thread == 1`` so the cell needs no post-increment. + base_offset = smem_offsets[subgroup_idx * RADIX_DIGITS + digit] + rank = base_offset + _ops.cast(popc, _i32) - _i32(1) - if lane == leader: - smem_match[match_idx] = _i64(0) - _subgroup_sync_fence() - - rank = subgroup_offset + _ops.cast(popc, _i32) - _i32(1) - - bins[tid] = bin_count - excl_prefix[tid] = exclusive_digit_prefix sync() return rank diff --git a/quadrants/codegen/amdgpu/codegen_amdgpu.cpp b/quadrants/codegen/amdgpu/codegen_amdgpu.cpp index 4945f3481e..fab5515b22 100644 --- a/quadrants/codegen/amdgpu/codegen_amdgpu.cpp +++ b/quadrants/codegen/amdgpu/codegen_amdgpu.cpp @@ -24,6 +24,59 @@ namespace lang { using namespace llvm; +namespace { +constexpr unsigned kAmdgpuLdsAddrSpace = 3; + +// True if ``v`` is, or (for ``block.SharedArray``) transitively references, a +// value living in AMDGPU LDS (addrspace(3)). Shared arrays are emitted as a +// module-level ``addrspace(3)`` global immediately ``addrspacecast``-ed to +// addrspace(0); because the global is constant, that cast folds into a +// ``ConstantExpr`` rather than an instruction, so the addrspace(3) marker only +// survives *nested inside* the constant operand of the loads/stores. We +// therefore have to descend through constant expressions, not just look at +// top-level operand types. ``seen`` guards against cyclic constant graphs. +bool value_references_lds(llvm::Value *v, llvm::SmallPtrSetImpl &seen) { + if (!v || !seen.insert(v).second) { + return false; + } + if (auto *pty = llvm::dyn_cast(v->getType())) { + if (pty->getAddressSpace() == kAmdgpuLdsAddrSpace) { + return true; + } + } + if (auto *ce = llvm::dyn_cast(v)) { + for (llvm::Value *op : ce->operands()) { + if (value_references_lds(op, seen)) { + return true; + } + } + } + return false; +} + +// Returns true if ``func`` reads or writes LDS anywhere in its body. Used to +// force-inline LDS-touching range_for bodies so all threads of the offloaded +// task share one LDS allocation (see the call site). +bool function_uses_lds(llvm::Function *func) { + llvm::SmallPtrSet seen; + for (llvm::BasicBlock &bb : *func) { + for (llvm::Instruction &inst : bb) { + seen.clear(); + if (value_references_lds(&inst, seen)) { + return true; + } + for (llvm::Value *op : inst.operands()) { + seen.clear(); + if (value_references_lds(op, seen)) { + return true; + } + } + } + } + return false; +} +} // namespace + class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { public: using IRVisitor::visit; @@ -305,6 +358,22 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { should_inline = true; } } + // CORRECTNESS OVERRIDE (must win over the size gate above): a range_for + // body that touches workgroup-shared LDS (addrspace(3) -- ``block.SharedArray``) + // MUST be inlined into the offloaded task, regardless of size or the + // ``force_inline`` request. When such a body is left out-of-line, it is + // emitted as a separate ``amdgpu_kernel``-attributed function whose + // module-level LDS globals are allocated independently of the caller's; + // the per-thread ``s_swappc_b64`` call then operates on a different LDS + // frame than the surrounding block, so cross-thread publish/observe via + // SharedArray (e.g. radix-rank histogram bins, block reductions) reads + // stale / uninitialised LDS and the result is nondeterministically wrong. + // Inlining keeps the body in the task function so all threads share the + // one LDS allocation. This only force-inlines bodies that actually use + // LDS, so the perf-motivated size gate still applies to LDS-free bodies. + if (!should_inline && function_uses_lds(body)) { + should_inline = true; + } if (should_inline) { tlctx->mark_inline(body); } diff --git a/quadrants/runtime/llvm/llvm_context.cpp b/quadrants/runtime/llvm/llvm_context.cpp index 870136330c..71df0a963e 100644 --- a/quadrants/runtime/llvm/llvm_context.cpp +++ b/quadrants/runtime/llvm/llvm_context.cpp @@ -692,11 +692,20 @@ std::unique_ptr QuadrantsLLVMContext::module_from_file(const std:: auto store_ptr = builder.CreateInBoundsGEP(buf_ty, lds_global, llvm::ArrayRef(store_idxs)); builder.CreateStore(value_arg, store_ptr); - // Wavefront-scope acquire-release fence -> ``s_waitcnt lgkmcnt(0)`` on AMDGPU; orders LDS writes against - // subsequent LDS reads within the wave without touching cross-wave state (avoids the ``s_barrier`` that a - // workgroup-scope fence would emit, which deadlocks if only some waves in the workgroup reach this point). - llvm::SyncScope::ID wave_scope = ctx->getOrInsertSyncScopeID("wavefront"); - builder.CreateFence(llvm::AtomicOrdering::AcquireRelease, wave_scope); + // Acquire-release fence that orders the LDS store above against the cross-lane LDS load below. This MUST be + // ``workgroup`` scope, not ``wavefront``: LDS is workgroup-shared address space, and the AMDGPU + // ``SIMemoryLegalizer`` only emits the ``s_waitcnt lgkmcnt(0)`` that orders LDS traffic for a fence whose scope + // is at least as wide as the memory's scope. A ``wavefront``-scope fence is *narrower* than workgroup, so the + // legalizer treats it as not governing LDS at all and emits no waitcnt -- which left the ``ds_write`` -> + // partner + // ``ds_read`` pair with no synchronization, so each lane could read its partner's stale / in-flight slot. That + // is a cross-lane Read-After-Write hazard and showed up as a nondeterministic, wrong cross-half shuffle (the + // ``permlane64`` swap silently returned garbage for some lanes). A workgroup-scope *fence* lowers to + // ``s_waitcnt`` + cache ops, NOT an ``s_barrier`` (only the ``barrier()`` control intrinsic emits + // ``s_barrier``), so there is no divergent-wave deadlock risk -- this is the same scope ``block_mem_fence`` + // already uses safely. + llvm::SyncScope::ID lds_scope = ctx->getOrInsertSyncScopeID("workgroup"); + builder.CreateFence(llvm::AtomicOrdering::AcquireRelease, lds_scope); // partner_lane = lane ^ 32; result = lds[wave_base + partner_lane] auto partner_lane = builder.CreateXor(lane, llvm::ConstantInt::get(i32_ty, 32)); @@ -704,6 +713,15 @@ std::unique_ptr QuadrantsLLVMContext::module_from_file(const std:: llvm::Value *load_idxs[] = {zero32, partner_slot}; auto load_ptr = builder.CreateInBoundsGEP(buf_ty, lds_global, llvm::ArrayRef(load_idxs)); auto result = builder.CreateAlignedLoad(i32_ty, load_ptr, llvm::Align(4)); + // Second workgroup-scope fence -> a trailing ``s_waitcnt lgkmcnt(0)`` that forces this call's LDS *read* to + // retire before any later code (in particular, the next inlined ``amdgpu_permlane64`` call's LDS *store* to the + // same single shared ``__amdgpu_permlane64_lds`` buffer) can issue. Without it, back-to-back permlane64 calls + // -- e.g. ``_exclusive_scan_tiled``'s leading ``shuffle_up(value, 1)`` immediately followed by the inclusive + // scan's own ``shuffle_up`` chain -- form a Write-After-Read hazard: call N+1's store can overtake call N's + // still-outstanding read of the same slot, so a lane reads its partner's *next* value instead of the current + // one. The leading fence above closes the within-call Read-After-Write; this one closes the across-call + // Write-After-Read. Both must be workgroup scope for the LDS waitcnt to be emitted (see above). + builder.CreateFence(llvm::AtomicOrdering::AcquireRelease, lds_scope); builder.CreateRet(result); QuadrantsLLVMContext::mark_inline(permlane64_func); diff --git a/quadrants/runtime/llvm/runtime_module/runtime.cpp b/quadrants/runtime/llvm/runtime_module/runtime.cpp index fbc4c23eed..1792444a57 100644 --- a/quadrants/runtime/llvm/runtime_module/runtime.cpp +++ b/quadrants/runtime/llvm/runtime_module/runtime.cpp @@ -1056,12 +1056,10 @@ i32 amdgpu_lane_id() { i32 amdgpu_cross_half_shuffle_i32(i32 target_lane, i32 value) { // Two parallel reads, then a per-lane select. ``permlane64`` is convergent and must execute uniformly across the // wave -- lifting it above the select keeps the AMDGPU backend happy and lets it issue exactly one - // ``v_permlane64_b32``. ``ds_bpermute`` on RDNA wave64 is SIMD32-scoped with a 5-bit address (top half of the wave - // is unreachable directly), so ``from_self_half`` handles the same-SIMD case and ``from_other_half`` handles the - // cross-SIMD case via the ``swapped`` payload. On CDNA the wave is one SIMD64 so both reads return the same value - // and the select is a no-op; we don't try to optimize that out because the dead read is cheap (LLVM CSE may fold - // it anyway). - i32 self_lane = amdgpu_lane_id(); + // ``v_permlane64_b32``. ``ds_bpermute`` has a 5-bit lane address (see the block comment above), so it always reads + // a bottom-half lane: ``from_self_half`` is ``value[target_lane & 31]`` and ``from_other_half`` (the + // ``permlane64``-swapped read) is ``value[(target_lane & 31) + 32]``. The select below picks between them purely + // on whether the *target* lane lives in the top half; the issuing lane is irrelevant. i32 swapped = amdgpu_permlane64(value); i32 byte = (target_lane & 31) * 4; // ``llvm.amdgcn.ds.bpermute`` is the real hardware ``ds_bpermute_b32`` -- but if LLVM's uniformity analysis decides @@ -1089,7 +1087,15 @@ i32 amdgpu_cross_half_shuffle_i32(i32 target_lane, i32 value) { #endif i32 from_self_half = amdgpu_ds_bpermute(byte, value); i32 from_other_half = amdgpu_ds_bpermute(byte, swapped); - return ((target_lane ^ self_lane) & 32) ? from_other_half : from_self_half; + // ``ds_bpermute`` always reads a *bottom-half* lane (``target_lane & 31``), so ``from_self_half`` holds + // ``value[target_lane & 31]`` and ``from_other_half`` (the ``permlane64``-swapped read) holds + // ``value[(target_lane & 31) + 32]``. The desired ``value[target_lane]`` therefore lives in ``from_other_half`` + // exactly when the *target* lane is in the top half -- ``target_lane & 32`` -- independent of which lane issues the + // read. (The previous ``(target_lane ^ self_lane) & 32`` form only coincides with this for bottom-half readers; + // for top-half readers it inverts the selection, corrupting every cross-half read on lanes 32..63 -- e.g. block + // scans via ``shuffle_up`` first diverge at thread 32. Reductions masked the bug because their result is consumed + // only on lane 0.) + return (target_lane & 32) ? from_other_half : from_self_half; } i32 amdgpu_shuffle_i32(i32 index, i32 value) {