From 9b568f7bb05e0b4cc7d06355f66f9d2c2749c607 Mon Sep 17 00:00:00 2001 From: rtmadduri Date: Tue, 12 May 2026 00:06:47 -0400 Subject: [PATCH 1/4] implement atomic optimizations --- quadrants/codegen/amdgpu/codegen_amdgpu.cpp | 177 +++++++++++++++++++- 1 file changed, 175 insertions(+), 2 deletions(-) diff --git a/quadrants/codegen/amdgpu/codegen_amdgpu.cpp b/quadrants/codegen/amdgpu/codegen_amdgpu.cpp index a910c1ff69..649d724609 100644 --- a/quadrants/codegen/amdgpu/codegen_amdgpu.cpp +++ b/quadrants/codegen/amdgpu/codegen_amdgpu.cpp @@ -156,6 +156,17 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { #undef UNARY_STD } + // AMDGPU sync scope helper. Genesis atomics are producer-consumer + // across kernel launches; the launch boundary supplies release/acquire + // ordering, so the in-kernel atomic only needs Monotonic ordering at + // agent (device) scope. This skips the s_waitcnt vmcnt(0) lgkmcnt(0) + // and buffer_inv / buffer_wbinvl1_vol cache-invalidate pair that + // SequentiallyConsistent system-scope atomics emit on AMDGPU. Atomics + // are still emitted as single-instruction global_atomic_* on gfx942. + llvm::SyncScope::ID amdgpu_atomic_scope() { + return llvm_context->getOrInsertSyncScopeID("agent"); + } + // Emit reductions as direct LLVM atomics instead of calling runtime // reduce_* helpers. The runtime helpers expect addrspace(0) pointers, // but SNode destinations arrive in addrspace(1). Calling the helpers @@ -184,13 +195,13 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { if (i32_ops.find(op) != i32_ops.end()) { return builder->CreateAtomicRMW( i32_ops.at(op), dest, val, llvm::MaybeAlign(0), - llvm::AtomicOrdering::SequentiallyConsistent); + llvm::AtomicOrdering::Monotonic, amdgpu_atomic_scope()); } } else if (prim_type == PrimitiveTypeID::f32) { if (op == AtomicOpType::add) { return builder->CreateAtomicRMW( llvm::AtomicRMWInst::FAdd, dest, val, llvm::MaybeAlign(0), - llvm::AtomicOrdering::SequentiallyConsistent); + llvm::AtomicOrdering::Monotonic, amdgpu_atomic_scope()); } else if (op == AtomicOpType::min) { return atomic_op_using_cas( dest, val, @@ -206,6 +217,149 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { return nullptr; } + // Override the non-reduction integer atomic path so user-written + // qd.atomic_add(...) on integer SNode/Ndarray destinations (e.g. the + // work-queue counters in Genesis's narrowphase / broadphase / inequality + // constraint kernels, which dominate ~14% of the run after the tiled-wc + // CG body) emits Monotonic+agent on AMDGPU instead of SeqCst+system. + // Mirrors TaskCodeGenLLVM::integral_type_atomic; only the AtomicRMW + // ordering + sync scope differ. + llvm::Value *integral_type_atomic(AtomicOpStmt *stmt) override { + if (!is_integral(stmt->val->ret_type)) { + return nullptr; + } + if (stmt->op_type == AtomicOpType::mul) { + return atomic_op_using_cas( + llvm_val[stmt->dest], llvm_val[stmt->val], + [&](auto v1, auto v2) { return builder->CreateMul(v1, v2); }, + stmt->val->ret_type); + } + std::unordered_map bin_op; + bin_op[AtomicOpType::add] = llvm::AtomicRMWInst::BinOp::Add; + if (is_signed(stmt->val->ret_type)) { + bin_op[AtomicOpType::min] = llvm::AtomicRMWInst::BinOp::Min; + bin_op[AtomicOpType::max] = llvm::AtomicRMWInst::BinOp::Max; + } else { + bin_op[AtomicOpType::min] = llvm::AtomicRMWInst::BinOp::UMin; + bin_op[AtomicOpType::max] = llvm::AtomicRMWInst::BinOp::UMax; + } + bin_op[AtomicOpType::bit_and] = llvm::AtomicRMWInst::BinOp::And; + bin_op[AtomicOpType::bit_or] = llvm::AtomicRMWInst::BinOp::Or; + bin_op[AtomicOpType::bit_xor] = llvm::AtomicRMWInst::BinOp::Xor; + QD_ASSERT(bin_op.find(stmt->op_type) != bin_op.end()); + return builder->CreateAtomicRMW( + bin_op.at(stmt->op_type), llvm_val[stmt->dest], llvm_val[stmt->val], + llvm::MaybeAlign(0), llvm::AtomicOrdering::Monotonic, + amdgpu_atomic_scope()); + } + + // Override the non-reduction real atomic path. The f32 add fast path + // is what Genesis hits whenever it accumulates a float into a counter + // outside a TLS-promoted reduction. f16 and f64 paths fall through to + // the CAS-based atomic_op_using_cas below, which is also overridden. + llvm::Value *real_type_atomic(AtomicOpStmt *stmt) override { + if (!is_real(stmt->val->ret_type)) { + return nullptr; + } + PrimitiveTypeID prim_type = + stmt->val->ret_type->cast()->type; + AtomicOpType op = stmt->op_type; + if (prim_type == PrimitiveTypeID::f16) { + switch (op) { + case AtomicOpType::add: + return atomic_op_using_cas( + llvm_val[stmt->dest], llvm_val[stmt->val], + [&](auto v1, auto v2) { return builder->CreateFAdd(v1, v2); }, + stmt->val->ret_type); + case AtomicOpType::max: + return atomic_op_using_cas( + llvm_val[stmt->dest], llvm_val[stmt->val], + [&](auto v1, auto v2) { return builder->CreateMaxNum(v1, v2); }, + stmt->val->ret_type); + case AtomicOpType::min: + return atomic_op_using_cas( + llvm_val[stmt->dest], llvm_val[stmt->val], + [&](auto v1, auto v2) { return builder->CreateMinNum(v1, v2); }, + stmt->val->ret_type); + default: + break; + } + } + switch (op) { + case AtomicOpType::add: + return builder->CreateAtomicRMW( + llvm::AtomicRMWInst::FAdd, llvm_val[stmt->dest], + llvm_val[stmt->val], llvm::MaybeAlign(0), + llvm::AtomicOrdering::Monotonic, amdgpu_atomic_scope()); + case AtomicOpType::mul: + return atomic_op_using_cas( + llvm_val[stmt->dest], llvm_val[stmt->val], + [&](auto v1, auto v2) { return builder->CreateFMul(v1, v2); }, + stmt->val->ret_type); + default: + break; + } + // f32/f64 min/max fall through to the runtime atomic_min/max helpers + // declared in runtime.cpp; the base implementation is intentionally + // reused here (they are CAS-based regardless of ordering, and Genesis + // does not hit this path in the benchmark). + std::unordered_map> + atomics; + atomics[PrimitiveTypeID::f32][AtomicOpType::min] = "atomic_min_f32"; + atomics[PrimitiveTypeID::f64][AtomicOpType::min] = "atomic_min_f64"; + atomics[PrimitiveTypeID::f32][AtomicOpType::max] = "atomic_max_f32"; + atomics[PrimitiveTypeID::f64][AtomicOpType::max] = "atomic_max_f64"; + QD_ASSERT(atomics.find(prim_type) != atomics.end()); + QD_ASSERT(atomics.at(prim_type).find(op) != atomics.at(prim_type).end()); + return call(atomics.at(prim_type).at(op), llvm_val[stmt->dest], + llvm_val[stmt->val]); + } + + // CAS-loop fallback path on AMDGPU. Used for atomic_mul on integers / + // floats and for f32 atomic min/max in the reduction path. Identical + // to the base implementation except both AtomicOrdering args on the + // cmpxchg are relaxed to Monotonic at agent scope. + llvm::Value *atomic_op_using_cas( + llvm::Value *dest, + llvm::Value *val, + std::function op, + const DataType &type) override { + using namespace llvm; + BasicBlock *body = + BasicBlock::Create(*llvm_context, "while_loop_body", func); + BasicBlock *after_loop = + BasicBlock::Create(*llvm_context, "after_while", func); + + builder->CreateBr(body); + builder->SetInsertPoint(body); + + llvm::Value *old_val; + { + int bits = data_type_bits(type); + unsigned dest_as = dest->getType()->isPointerTy() + ? dest->getType()->getPointerAddressSpace() + : 0; + llvm::PointerType *typeIntPtr = + llvm::PointerType::get(*llvm_context, dest_as); + llvm::IntegerType *typeIntTy = get_integer_type(bits); + + old_val = builder->CreateLoad(val->getType(), dest); + auto new_val = op(old_val, val); + dest = builder->CreateBitCast(dest, typeIntPtr); + auto atomicCmpXchg = builder->CreateAtomicCmpXchg( + dest, builder->CreateBitCast(old_val, typeIntTy), + builder->CreateBitCast(new_val, typeIntTy), llvm::MaybeAlign(0), + AtomicOrdering::Monotonic, AtomicOrdering::Monotonic, + amdgpu_atomic_scope()); + auto ok = builder->CreateExtractValue(atomicCmpXchg, 1); + builder->CreateCondBr(builder->CreateNot(ok), body, after_loop); + } + + builder->SetInsertPoint(after_loop); + return old_val; + } + void visit(RangeForStmt *for_stmt) override { create_naive_range_for(for_stmt); } @@ -261,6 +415,21 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { QD_NOT_IMPLEMENTED } + // Pin amdgpu-flat-work-group-size on the current kernel function (the + // `func` member set by init_offloaded_task_function) to the exact + // dispatch shape. The base codegen leaves this attribute unset, so the + // inheritance loop in jit_amdgpu.cpp falls back to the conservative + // "1,128" bound for every body function. Setting it here lets the + // AMDGPU backend fold __ockl_get_local_size to a constant, tightens + // VGPR/SGPR allocator decisions, and propagates the tight bound to all + // body functions via the inheritance loop. + void set_amdgpu_flat_work_group_size(int block_dim) { + if (func && block_dim > 0) { + auto wgs = std::to_string(block_dim); + func->addFnAttr("amdgpu-flat-work-group-size", wgs + "," + wgs); + } + } + void emit_amdgpu_gc(OffloadedStmt *stmt) { auto snode_id = tlctx->get_constant(stmt->snode->id); { @@ -269,6 +438,7 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { finalize_offloaded_task_function(); current_task->grid_dim = compile_config.saturating_grid_dim; current_task->block_dim = 64; + set_amdgpu_flat_work_group_size(current_task->block_dim); offloaded_tasks.push_back(*current_task); current_task = nullptr; } @@ -278,6 +448,7 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { finalize_offloaded_task_function(); current_task->grid_dim = 1; current_task->block_dim = 1; + set_amdgpu_flat_work_group_size(current_task->block_dim); offloaded_tasks.push_back(*current_task); current_task = nullptr; } @@ -287,6 +458,7 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { finalize_offloaded_task_function(); current_task->grid_dim = compile_config.saturating_grid_dim; current_task->block_dim = 64; + set_amdgpu_flat_work_group_size(current_task->block_dim); offloaded_tasks.push_back(*current_task); current_task = nullptr; } @@ -525,6 +697,7 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { current_task->dynamic_shared_array_bytes = dynamic_shared_array_bytes; QD_ASSERT(current_task->grid_dim != 0); QD_ASSERT(current_task->block_dim != 0); + set_amdgpu_flat_work_group_size(current_task->block_dim); offloaded_tasks.push_back(*current_task); current_task = nullptr; dynamic_shared_array_bytes = 0; From f4886cdebc8ca082fb55f41e58db282926f55c16 Mon Sep 17 00:00:00 2001 From: rtmadduri Date: Thu, 14 May 2026 14:37:33 -0400 Subject: [PATCH 2/4] remove amdgpu-flat-work-group-size changes --- quadrants/codegen/amdgpu/codegen_amdgpu.cpp | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/quadrants/codegen/amdgpu/codegen_amdgpu.cpp b/quadrants/codegen/amdgpu/codegen_amdgpu.cpp index 649d724609..8f81081007 100644 --- a/quadrants/codegen/amdgpu/codegen_amdgpu.cpp +++ b/quadrants/codegen/amdgpu/codegen_amdgpu.cpp @@ -415,21 +415,6 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { QD_NOT_IMPLEMENTED } - // Pin amdgpu-flat-work-group-size on the current kernel function (the - // `func` member set by init_offloaded_task_function) to the exact - // dispatch shape. The base codegen leaves this attribute unset, so the - // inheritance loop in jit_amdgpu.cpp falls back to the conservative - // "1,128" bound for every body function. Setting it here lets the - // AMDGPU backend fold __ockl_get_local_size to a constant, tightens - // VGPR/SGPR allocator decisions, and propagates the tight bound to all - // body functions via the inheritance loop. - void set_amdgpu_flat_work_group_size(int block_dim) { - if (func && block_dim > 0) { - auto wgs = std::to_string(block_dim); - func->addFnAttr("amdgpu-flat-work-group-size", wgs + "," + wgs); - } - } - void emit_amdgpu_gc(OffloadedStmt *stmt) { auto snode_id = tlctx->get_constant(stmt->snode->id); { @@ -438,7 +423,6 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { finalize_offloaded_task_function(); current_task->grid_dim = compile_config.saturating_grid_dim; current_task->block_dim = 64; - set_amdgpu_flat_work_group_size(current_task->block_dim); offloaded_tasks.push_back(*current_task); current_task = nullptr; } @@ -448,7 +432,6 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { finalize_offloaded_task_function(); current_task->grid_dim = 1; current_task->block_dim = 1; - set_amdgpu_flat_work_group_size(current_task->block_dim); offloaded_tasks.push_back(*current_task); current_task = nullptr; } @@ -458,7 +441,6 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { finalize_offloaded_task_function(); current_task->grid_dim = compile_config.saturating_grid_dim; current_task->block_dim = 64; - set_amdgpu_flat_work_group_size(current_task->block_dim); offloaded_tasks.push_back(*current_task); current_task = nullptr; } @@ -697,7 +679,6 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { current_task->dynamic_shared_array_bytes = dynamic_shared_array_bytes; QD_ASSERT(current_task->grid_dim != 0); QD_ASSERT(current_task->block_dim != 0); - set_amdgpu_flat_work_group_size(current_task->block_dim); offloaded_tasks.push_back(*current_task); current_task = nullptr; dynamic_shared_array_bytes = 0; From fd0df6d816980f2575c819aa8048704cc43e4667 Mon Sep 17 00:00:00 2001 From: rtmadduri Date: Fri, 15 May 2026 14:43:43 -0500 Subject: [PATCH 3/4] fix the security issues --- .../scripts_new/linux/1_prerequisites.sh | 0 docs/source/user_guide/index.md | 1 + quadrants/codegen/amdgpu/codegen_amdgpu.cpp | 52 ++- quadrants/codegen/llvm/codegen_llvm.cpp | 37 +- quadrants/codegen/llvm/codegen_llvm.h | 27 ++ .../llvm/runtime_module/node_dynamic.h | 10 + tests/python/test_atomic_amdgpu.py | 437 ++++++++++++++++++ tests/python/test_atomic_amdgpu_ir.py | 374 +++++++++++++++ 8 files changed, 921 insertions(+), 17 deletions(-) mode change 100644 => 100755 .github/workflows/scripts_new/linux/1_prerequisites.sh create mode 100644 tests/python/test_atomic_amdgpu.py create mode 100644 tests/python/test_atomic_amdgpu_ir.py diff --git a/.github/workflows/scripts_new/linux/1_prerequisites.sh b/.github/workflows/scripts_new/linux/1_prerequisites.sh old mode 100644 new mode 100755 diff --git a/docs/source/user_guide/index.md b/docs/source/user_guide/index.md index e7f8c46b30..cdff37b8ea 100644 --- a/docs/source/user_guide/index.md +++ b/docs/source/user_guide/index.md @@ -38,6 +38,7 @@ interop cuda_graph perf_dispatch +amdgpu_atomics ``` ```{toctree} diff --git a/quadrants/codegen/amdgpu/codegen_amdgpu.cpp b/quadrants/codegen/amdgpu/codegen_amdgpu.cpp index 8f81081007..cdb83f8df3 100644 --- a/quadrants/codegen/amdgpu/codegen_amdgpu.cpp +++ b/quadrants/codegen/amdgpu/codegen_amdgpu.cpp @@ -35,6 +35,11 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { const Kernel *kernel, IRNode *ir = nullptr) : TaskCodeGenLLVM(id, config, tlctx, kernel, ir) { + // ``llvm_context`` is assigned by TaskCodeGenLLVM::initialize_context + // in the base ctor body, so it is non-null here. Cache the agent + // sync-scope ID once per codegen instance to avoid the + // string-hash lookup on every atomic emission. + amdgpu_agent_scope_ = llvm_context->getOrInsertSyncScopeID("agent"); } llvm::Value *create_print(std::string tag, @@ -156,15 +161,35 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { #undef UNARY_STD } - // AMDGPU sync scope helper. Genesis atomics are producer-consumer - // across kernel launches; the launch boundary supplies release/acquire - // ordering, so the in-kernel atomic only needs Monotonic ordering at - // agent (device) scope. This skips the s_waitcnt vmcnt(0) lgkmcnt(0) - // and buffer_inv / buffer_wbinvl1_vol cache-invalidate pair that - // SequentiallyConsistent system-scope atomics emit on AMDGPU. Atomics - // are still emitted as single-instruction global_atomic_* on gfx942. - llvm::SyncScope::ID amdgpu_atomic_scope() { - return llvm_context->getOrInsertSyncScopeID("agent"); + // AMDGPU atomic memory-model contract (validated on gfx942): + // - Monotonic ordering at "agent" sync scope is sufficient for every + // atomic Genesis emits. Producer-consumer ordering between kernels + // is supplied by the HIP stream / kernel-launch boundary, not by + // the atomic itself. + // - This avoids the s_waitcnt vmcnt(0) lgkmcnt(0) plus buffer_inv / + // buffer_wbinvl1_vol cache-invalidate sequence that + // SequentiallyConsistent + system scope would otherwise emit + // around every atomic on AMDGPU. + // - User-visible contract: see docs/source/user_guide/amdgpu_atomics.md. + // + // The actual relaxation is implemented by overriding two virtuals on + // the base class; all atomic-emission logic (integral_type_atomic, + // real_type_atomic, atomic_op_using_cas) lives in TaskCodeGenLLVM and + // picks these up automatically. + llvm::AtomicOrdering default_atomic_ordering() const override { + return llvm::AtomicOrdering::Monotonic; + } + llvm::SyncScope::ID default_atomic_scope() const override { + return amdgpu_agent_scope_; + } + + // Required because default_atomic_scope() is not ``System``. Without + // this, f32/f64 atomic_min / atomic_max would fall through to the + // SeqCst+system runtime helpers in atomic.h and any user kernel that + // mixes ``qd.atomic_add`` and ``qd.atomic_min/max`` on the same f32 / + // f64 field would have undefined behavior per the LLVM memory model. + bool prefer_cas_for_fp_minmax() const override { + return true; } // Emit reductions as direct LLVM atomics instead of calling runtime @@ -195,13 +220,13 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { if (i32_ops.find(op) != i32_ops.end()) { return builder->CreateAtomicRMW( i32_ops.at(op), dest, val, llvm::MaybeAlign(0), - llvm::AtomicOrdering::Monotonic, amdgpu_atomic_scope()); + default_atomic_ordering(), default_atomic_scope()); } } else if (prim_type == PrimitiveTypeID::f32) { if (op == AtomicOpType::add) { return builder->CreateAtomicRMW( llvm::AtomicRMWInst::FAdd, dest, val, llvm::MaybeAlign(0), - llvm::AtomicOrdering::Monotonic, amdgpu_atomic_scope()); + default_atomic_ordering(), default_atomic_scope()); } else if (op == AtomicOpType::min) { return atomic_op_using_cas( dest, val, @@ -748,8 +773,11 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { workgroup_dim_, llvm::Type::getInt32Ty(*llvm_context)); return std::make_tuple(thread_idx, block_dim); } + // Cached LLVM sync-scope ID for "agent". Populated in the ctor and + // returned by default_atomic_scope() so that every atomic emission + // skips the per-call string-hash lookup into LLVMContext. + llvm::SyncScope::ID amdgpu_agent_scope_{llvm::SyncScope::System}; }; - LLVMCompiledTask KernelCodeGenAMDGPU::compile_task( int task_codegen_id, const CompileConfig &config, diff --git a/quadrants/codegen/llvm/codegen_llvm.cpp b/quadrants/codegen/llvm/codegen_llvm.cpp index 50e8a1599b..b4de7d783e 100644 --- a/quadrants/codegen/llvm/codegen_llvm.cpp +++ b/quadrants/codegen/llvm/codegen_llvm.cpp @@ -1455,7 +1455,7 @@ llvm::Value *TaskCodeGenLLVM::integral_type_atomic(AtomicOpStmt *stmt) { QD_ASSERT(bin_op.find(stmt->op_type) != bin_op.end()); return builder->CreateAtomicRMW( bin_op.at(stmt->op_type), llvm_val[stmt->dest], llvm_val[stmt->val], - llvm::MaybeAlign(0), llvm::AtomicOrdering::SequentiallyConsistent); + llvm::MaybeAlign(0), default_atomic_ordering(), default_atomic_scope()); } llvm::Value *TaskCodeGenLLVM::atomic_op_using_cas( @@ -1484,14 +1484,20 @@ llvm::Value *TaskCodeGenLLVM::atomic_op_using_cas( llvm::PointerType::get(*llvm_context, dest_as); llvm::IntegerType *typeIntTy = get_integer_type(bits); - old_val = builder->CreateLoad(val->getType(), dest); + // Use an atomic load (matching the cmpxchg's ordering and scope) + // so that LLVM cannot legally hoist this load out of the CAS-retry + // loop. An atomic load also requires an explicit natural alignment. + auto *initial_load = builder->CreateLoad(val->getType(), dest); + initial_load->setAlignment(llvm::Align(std::max(1, bits / 8))); + initial_load->setAtomic(default_atomic_ordering(), default_atomic_scope()); + old_val = initial_load; auto new_val = op(old_val, val); dest = builder->CreateBitCast(dest, typeIntPtr); auto atomicCmpXchg = builder->CreateAtomicCmpXchg( dest, builder->CreateBitCast(old_val, typeIntTy), builder->CreateBitCast(new_val, typeIntTy), llvm::MaybeAlign(0), - AtomicOrdering::SequentiallyConsistent, - AtomicOrdering::SequentiallyConsistent); + default_atomic_ordering(), default_atomic_ordering(), + default_atomic_scope()); // Check whether CAS was succussful auto ok = builder->CreateExtractValue(atomicCmpXchg, 1); builder->CreateCondBr(builder->CreateNot(ok), body, after_loop); @@ -1535,7 +1541,8 @@ llvm::Value *TaskCodeGenLLVM::real_type_atomic(AtomicOpStmt *stmt) { case AtomicOpType::add: return builder->CreateAtomicRMW( llvm::AtomicRMWInst::FAdd, llvm_val[stmt->dest], llvm_val[stmt->val], - llvm::MaybeAlign(0), llvm::AtomicOrdering::SequentiallyConsistent); + llvm::MaybeAlign(0), default_atomic_ordering(), + default_atomic_scope()); case AtomicOpType::mul: return atomic_op_using_cas( llvm_val[stmt->dest], llvm_val[stmt->val], @@ -1545,6 +1552,26 @@ llvm::Value *TaskCodeGenLLVM::real_type_atomic(AtomicOpStmt *stmt) { break; } + // Backends that emit atomics at a non-system scope (e.g. AMDGPU at + // ``agent`` scope) must avoid the SeqCst+system runtime helpers below + // for f32/f64 min/max -- mixing scopes on the same address is UB per + // the LLVM / C++ memory model. Route through the CAS path instead, + // which honors default_atomic_ordering() / default_atomic_scope(). + if (prefer_cas_for_fp_minmax()) { + if (op == AtomicOpType::min) { + return atomic_op_using_cas( + llvm_val[stmt->dest], llvm_val[stmt->val], + [&](auto v1, auto v2) { return builder->CreateMinNum(v1, v2); }, + stmt->val->ret_type); + } + if (op == AtomicOpType::max) { + return atomic_op_using_cas( + llvm_val[stmt->dest], llvm_val[stmt->val], + [&](auto v1, auto v2) { return builder->CreateMaxNum(v1, v2); }, + stmt->val->ret_type); + } + } + std::unordered_map> atomics; diff --git a/quadrants/codegen/llvm/codegen_llvm.h b/quadrants/codegen/llvm/codegen_llvm.h index 5f072e6b59..b17265e3e5 100644 --- a/quadrants/codegen/llvm/codegen_llvm.h +++ b/quadrants/codegen/llvm/codegen_llvm.h @@ -253,6 +253,33 @@ class TaskCodeGenLLVM : public IRVisitor, public LLVMModuleBuilder { virtual llvm::Value *real_type_atomic(AtomicOpStmt *stmt); + // Default memory ordering and sync scope used by the atomic-emission + // helpers (integral_type_atomic, real_type_atomic, atomic_op_using_cas). + // The base class targets the strongest C++ contract + // (SequentiallyConsistent + system scope), matching what was emitted + // before PR #38. Backend codegens (e.g. TaskCodeGenAMDGPU) override + // these to relax atomics to a backend-specific ordering without having + // to duplicate the entire atomic-emission logic. + virtual llvm::AtomicOrdering default_atomic_ordering() const { + return llvm::AtomicOrdering::SequentiallyConsistent; + } + virtual llvm::SyncScope::ID default_atomic_scope() const { + return llvm::SyncScope::System; + } + + // Controls whether f32 / f64 ``atomic_min`` / ``atomic_max`` should be + // lowered through ``atomic_op_using_cas`` (which honors + // ``default_atomic_ordering()`` / ``default_atomic_scope()``) instead of + // falling through to the SeqCst+system runtime helpers + // ``atomic_{min,max}_{f32,f64}`` declared in + // ``runtime/llvm/runtime_module/atomic.h``. Backends whose + // ``default_atomic_scope()`` is not ``System`` MUST override this to + // return ``true``; otherwise the same memory location would be touched + // at two different sync scopes (UB per the LLVM memory model). + virtual bool prefer_cas_for_fp_minmax() const { + return false; + } + void visit(AtomicOpStmt *stmt) override; void visit(GlobalPtrStmt *stmt) override; diff --git a/quadrants/runtime/llvm/runtime_module/node_dynamic.h b/quadrants/runtime/llvm/runtime_module/node_dynamic.h index a68b330764..0ade6580ff 100644 --- a/quadrants/runtime/llvm/runtime_module/node_dynamic.h +++ b/quadrants/runtime/llvm/runtime_module/node_dynamic.h @@ -18,6 +18,16 @@ void Dynamic_activate(Ptr meta_, Ptr node_, int i) { auto node = (DynamicNode *)(node_); // We need to not only update node->n, but also make sure the chunk containing // element i is allocated. + // + // ``atomic_max_i32`` is defined in + // ``quadrants/runtime/llvm/runtime_module/atomic.h`` and uses + // ``memory_order_seq_cst`` at the default (system) sync scope. On + // AMDGPU, codegen-emitted ``qd.atomic_*`` operations use Monotonic + // ordering at ``"agent"`` sync scope. The two never touch the same + // address: ``node->n`` is runtime-internal SNode metadata that user + // kernels reach only through the ``dynamic.append`` / + // ``dynamic.length`` abstractions, never via ``qd.atomic_*``. So the + // scope mismatch is intentional and safe. atomic_max_i32(&node->n, i + 1); int chunk_start = 0; auto p_chunk_ptr = &node->ptr; diff --git a/tests/python/test_atomic_amdgpu.py b/tests/python/test_atomic_amdgpu.py new file mode 100644 index 0000000000..5c89616fe2 --- /dev/null +++ b/tests/python/test_atomic_amdgpu.py @@ -0,0 +1,437 @@ +"""AMDGPU-targeted atomic correctness tests. + +These tests exist as a regression net for the atomic-ordering relaxation +landed in PR #38 (Monotonic + ``syncscope("agent")`` instead of +SequentiallyConsistent + system scope) and the follow-up correctness fixes +planned in Phase 3 (mixed-scope cleanup, CAS-loop atomic load). + +Coverage is intentionally broader than ``test_atomic.py`` for the AMDGPU +arch specifically, because several existing tests are explicitly gated to +``[qd.cpu, qd.cuda]`` and therefore would not catch AMDGPU-side +regressions. Anything that exists for CPU/CUDA and could plausibly behave +differently on AMDGPU under relaxed atomics gets mirrored here. + +Each test runs *only* on AMDGPU and is skipped (rather than xfailed) on +other arches so we never accidentally validate the AMDGPU contract using +a CUDA / CPU run. +""" + +import pytest + +import quadrants as qd + +from tests import test_utils + +# Keep this small enough to run quickly on CI but large enough that the +# atomics actually contend (i.e. multiple threads hit the same counter). +N = 256 + + +# --------------------------------------------------------------------------- +# 1. End-to-end correctness for every (type, op) the AMDGPU codegen +# rewrites with the new ordering. Mirrors the CPU/CUDA-only tests in +# test_atomic.py. +# --------------------------------------------------------------------------- + + +@test_utils.test(arch=[qd.amdgpu]) +def test_atomic_add_i32_amdgpu(): + c = qd.field(qd.i32, shape=()) + c[None] = 0 + + @qd.kernel + def k(): + for _ in range(N): + qd.atomic_add(c[None], 1) + + k() + assert c[None] == N + + +@test_utils.test(arch=[qd.amdgpu]) +def test_atomic_add_i64_amdgpu(): + c = qd.field(qd.i64, shape=()) + c[None] = 0 + + @qd.kernel + def k(): + for _ in range(N): + qd.atomic_add(c[None], qd.cast(1, qd.i64)) + + k() + assert c[None] == N + + +@test_utils.test(arch=[qd.amdgpu]) +def test_atomic_add_u32_amdgpu(): + c = qd.field(qd.u32, shape=()) + c[None] = 0 + + @qd.kernel + def k(): + for _ in range(N): + qd.atomic_add(c[None], qd.cast(1, qd.u32)) + + k() + assert c[None] == N + + +@test_utils.test(arch=[qd.amdgpu]) +def test_atomic_add_u64_amdgpu(): + c = qd.field(qd.u64, shape=()) + c[None] = 0 + + @qd.kernel + def k(): + for _ in range(N): + qd.atomic_add(c[None], qd.cast(1, qd.u64)) + + k() + assert c[None] == N + + +@test_utils.test(arch=[qd.amdgpu]) +def test_atomic_add_f32_amdgpu(): + c = qd.field(qd.f32, shape=()) + c[None] = 0.0 + + @qd.kernel + def k(): + for _ in range(N): + qd.atomic_add(c[None], 1.0) + + k() + # Reduction over many threads; allow a small relative tolerance for + # ordering-dependent FP non-associativity. + assert c[None] == test_utils.approx(float(N), rel=1e-5) + + +@test_utils.test(arch=[qd.amdgpu]) +def test_atomic_add_f64_amdgpu(): + c = qd.field(qd.f64, shape=()) + c[None] = 0.0 + + @qd.kernel + def k(): + for _ in range(N): + qd.atomic_add(c[None], 1.0) + + k() + assert c[None] == test_utils.approx(float(N), rel=1e-10) + + +# --------------------------------------------------------------------------- +# Integer min/max — this is the gap left by test_atomic_min_max_uint, which +# is gated to [cpu, cuda]. AMDGPU now exercises the same path under +# Monotonic+agent. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("dtype", [qd.i32, qd.i64, qd.u32, qd.u64]) +@test_utils.test(arch=[qd.amdgpu]) +def test_atomic_min_int_amdgpu(dtype): + x = qd.field(dtype, shape=N) + + @qd.kernel + def init(): + # x[1..N-1] = 1, 2, ..., N-1. Then x[0] is overwritten to a + # sentinel so the reduction below has to actually update it. + # The minimum value present in the indices the reducer scans + # (1..N-1) is 1, which is what the assertion below checks. + for i in x: + x[i] = qd.cast(i, dtype) + x[0] = qd.cast(10**6, dtype) + + @qd.kernel + def reduce_min(): + for i in range(1, N): + qd.atomic_min(x[0], x[i]) + + init() + reduce_min() + assert int(x[0]) == 1 + + +@pytest.mark.parametrize("dtype", [qd.i32, qd.i64, qd.u32, qd.u64]) +@test_utils.test(arch=[qd.amdgpu]) +def test_atomic_max_int_amdgpu(dtype): + x = qd.field(dtype, shape=N) + + @qd.kernel + def init(): + for i in x: + x[i] = qd.cast(i + 1, dtype) + + @qd.kernel + def reduce_max(): + for i in range(N): + qd.atomic_max(x[0], x[i]) + + init() + reduce_max() + assert int(x[0]) == N + + +@pytest.mark.parametrize("dtype", [qd.i32, qd.i64, qd.u32, qd.u64]) +@test_utils.test(arch=[qd.amdgpu]) +def test_atomic_bitops_int_amdgpu(dtype): + # and / or / xor on the same destination, mirroring the + # *_expr_evaled tests in test_atomic.py but on AMDGPU specifically. + # + # Host-scope assignments (``val[None] = ...``) take a Python int and + # let the field setter do the dtype conversion. Calling ``qd.cast`` + # from host scope returns an ``Expr`` whose construction reads + # ``src_info_stack[-1]``, which is empty outside a kernel. + n_bits = 16 + val = qd.field(dtype, shape=()) + val[None] = (1 << n_bits) - 1 + + @qd.kernel + def do_and(): + for i in range(n_bits): + # Clear bit i. + qd.atomic_and(val[None], qd.cast(~(1 << i), dtype)) + + do_and() + assert int(val[None]) == 0 + + val[None] = 0 + + @qd.kernel + def do_or(): + for i in range(n_bits): + qd.atomic_or(val[None], qd.cast(1 << i, dtype)) + + do_or() + assert int(val[None]) == (1 << n_bits) - 1 + + val[None] = 0 + + @qd.kernel + def do_xor(): + for i in range(n_bits): + qd.atomic_xor(val[None], qd.cast(1 << i, dtype)) + + do_xor() + assert int(val[None]) == (1 << n_bits) - 1 + + +# --------------------------------------------------------------------------- +# Float min/max — exercises the runtime-helper fall-through in +# real_type_atomic for f32/f64. This is the call site that is currently +# (pre-Phase-3) mixed-scope: the runtime helper uses SeqCst+system while +# atomic_add on the same dest uses Monotonic+agent. The end-to-end +# correctness still holds for these isolated reductions, so the tests +# below are expected to pass today; the mixed-scope IR-level concern is +# asserted in test_atomic_amdgpu_ir.py. +# --------------------------------------------------------------------------- + + +@test_utils.test(arch=[qd.amdgpu]) +def test_atomic_min_f32_amdgpu(): + x = qd.field(qd.f32, shape=N) + + @qd.kernel + def init(): + # See test_atomic_min_int_amdgpu for why we use ``i`` (not + # ``i + 1``): the reducer scans x[1..N-1], whose smallest value + # must equal the assertion below. + for i in x: + x[i] = qd.cast(i, qd.f32) + x[0] = 1e9 + + @qd.kernel + def reduce_min(): + for i in range(1, N): + qd.atomic_min(x[0], x[i]) + + init() + reduce_min() + assert x[0] == test_utils.approx(1.0) + + +@test_utils.test(arch=[qd.amdgpu]) +def test_atomic_max_f32_amdgpu(): + x = qd.field(qd.f32, shape=N) + + @qd.kernel + def init(): + for i in x: + x[i] = qd.cast(i + 1, qd.f32) + + @qd.kernel + def reduce_max(): + for i in range(N): + qd.atomic_max(x[0], x[i]) + + init() + reduce_max() + assert x[0] == test_utils.approx(float(N)) + + +@test_utils.test(arch=[qd.amdgpu]) +def test_atomic_min_f64_amdgpu(): + x = qd.field(qd.f64, shape=N) + + @qd.kernel + def init(): + # See test_atomic_min_int_amdgpu for why we use ``i`` (not + # ``i + 1``): the reducer scans x[1..N-1], whose smallest value + # must equal the assertion below. + for i in x: + x[i] = qd.cast(i, qd.f64) + x[0] = 1e18 + + @qd.kernel + def reduce_min(): + for i in range(1, N): + qd.atomic_min(x[0], x[i]) + + init() + reduce_min() + assert x[0] == test_utils.approx(1.0, rel=1e-12) + + +@test_utils.test(arch=[qd.amdgpu]) +def test_atomic_max_f64_amdgpu(): + x = qd.field(qd.f64, shape=N) + + @qd.kernel + def init(): + for i in x: + x[i] = qd.cast(i + 1, qd.f64) + + @qd.kernel + def reduce_max(): + for i in range(N): + qd.atomic_max(x[0], x[i]) + + init() + reduce_max() + assert x[0] == test_utils.approx(float(N), rel=1e-12) + + +# --------------------------------------------------------------------------- +# atomic_mul — both integer and float fall through to +# atomic_op_using_cas, which the PR explicitly overrides. Make sure the +# CAS loop is still correct under Monotonic+agent. +# --------------------------------------------------------------------------- + + +@test_utils.test(arch=[qd.amdgpu]) +def test_atomic_mul_i32_amdgpu(): + # Single-thread mul; multi-thread mul is sensitive to non-associativity + # in floats, but for ints it should produce a deterministic result. + val = qd.field(qd.i32, shape=()) + val[None] = 1 + + @qd.kernel + def k(): + for i in range(1, 8): # serial loop -> single thread / deterministic + qd.atomic_mul(val[None], i) + + k() + assert int(val[None]) == 5040 + + +@test_utils.test(arch=[qd.amdgpu]) +def test_atomic_mul_f32_amdgpu(): + val = qd.field(qd.f32, shape=()) + val[None] = 1.0 + + @qd.kernel + def k(): + for i in range(1, 8): + qd.atomic_mul(val[None], qd.cast(i, qd.f32)) + + k() + assert val[None] == test_utils.approx(5040.0) + + +# --------------------------------------------------------------------------- +# 2. Index-allocator pattern (the central "is the relaxation safe?" test). +# +# Pattern: many threads atomically increment a counter; each thread writes +# its value to data[old_counter]. A *second* kernel reads counter and +# data[]. The kernel-launch boundary supplies the acquire/release ordering +# this PR's design depends on, so this should pass under Monotonic+agent. +# --------------------------------------------------------------------------- + + +@test_utils.test(arch=[qd.amdgpu]) +def test_atomic_index_allocator_inter_kernel_amdgpu(): + counter = qd.field(qd.i32, shape=()) + data = qd.field(qd.i32, shape=N) + counter[None] = 0 + + @qd.kernel + def producer(): + for i in range(N): + idx = qd.atomic_add(counter[None], 1) + data[idx] = i + 1 # tag with a non-zero value + + @qd.kernel + def consumer(out: qd.types.ndarray(qd.i32, 1)): + for i in range(N): + out[i] = data[i] + + producer() + import numpy as np + + out = np.zeros(N, dtype=np.int32) + consumer(out) + + assert int(counter[None]) == N + # Every slot should have been written by exactly one producer thread, + # so the set of values written must be exactly {1, ..., N}. + assert sorted(out.tolist()) == list(range(1, N + 1)) + + +# --------------------------------------------------------------------------- +# 3. Intra-kernel publish/subscribe — DOCUMENTED LIMITATION. +# +# Under the PR's Monotonic+agent ordering, a non-atomic write that follows +# an atomic on a different address has NO ordering guarantee to other +# threads in the same kernel. We do not assert correctness for this case; +# the test exists to mark the contract change explicitly and will be +# wired up if/when Phase 3.x adds an opt-in stronger fence primitive. +# --------------------------------------------------------------------------- + + +@pytest.mark.skip( + reason=( + "Intra-kernel atomic-publish-then-non-atomic-write is not " + "supported on AMDGPU under the Monotonic+agent ordering chosen " + "in PR #38. See docs/source/user_guide/amdgpu_atomics.md." + ) +) +@test_utils.test(arch=[qd.amdgpu]) +def test_atomic_intra_kernel_publish_unsupported_amdgpu(): + # Left here as a placeholder so the contract change is greppable. + pass + + +# --------------------------------------------------------------------------- +# 4. atomic_add on a global counter under heavy contention. +# +# This is the dominant pattern named in PR #38 (Genesis work-queue +# counters in narrowphase / broadphase / inequality constraint kernels). +# Test it explicitly so that any future regression in the contended path +# shows up loudly. +# --------------------------------------------------------------------------- + + +@test_utils.test(arch=[qd.amdgpu]) +def test_atomic_add_high_contention_amdgpu(): + counter = qd.field(qd.i32, shape=()) + counter[None] = 0 + + big_n = 4096 + + @qd.kernel + def k(): + for _ in range(big_n): + qd.atomic_add(counter[None], 1) + + k() + assert int(counter[None]) == big_n diff --git a/tests/python/test_atomic_amdgpu_ir.py b/tests/python/test_atomic_amdgpu_ir.py new file mode 100644 index 0000000000..3065d1fc6e --- /dev/null +++ b/tests/python/test_atomic_amdgpu_ir.py @@ -0,0 +1,374 @@ +"""IR-level assertions for AMDGPU atomic codegen (PR #38 follow-up). + +These tests are a tripwire for the codegen contract introduced in PR #38: +every atomic on AMDGPU should lower to an LLVM ``atomicrmw`` or +``cmpxchg`` carrying ``syncscope("agent")`` and ``monotonic`` ordering. +The C++/Python correctness tests in ``test_atomic_amdgpu.py`` only catch +*observable* regressions; if a future LLVM upgrade or codegen change +silently reverts the ordering / scope to ``seq_cst`` / system, the +program will still produce correct results but will pay the cache-flush +cost the PR was designed to eliminate. These tests catch that case. + +Implementation follows the pattern used in ``test_fn_attrs.py``: spawn a +child process with ``print_kernel_llvm_ir=True``, run a minimal kernel +that exercises one atomic shape, then parse the dumped ``.ll`` file and +assert on the IR text. + +Each test isolates a single atomic shape so that when the assertion +fails, the diagnostic points at exactly which lowering path regressed. + +Post-Phase-3 status: + * Integer / FP-add / FP-mul CAS paths emit ``syncscope("agent") + monotonic`` (asserted strictly). + * f32/f64 ``atomic_min``/``atomic_max`` now route through the + agent-scope CAS path (Phase 3.1), so they are also asserted + strictly. No more mixed-scope: ``atomic_add`` and ``atomic_min`` on + the same f32 / f64 field both use agent+monotonic. + * The CAS-loop's initial load is now an atomic monotonic load (Phase + 3.2), which prevents the optimizer from legally hoisting it out of + the retry block on future LLVM upgrades. +""" + +import os +import pathlib +import re +import subprocess +import sys + +import quadrants as qd + +from tests import test_utils + +RET_SUCCESS = 42 + +# Regexes that match the relevant atomic instructions. We intentionally do +# not pin the exact destination type / addrspace because those depend on +# SNode layout and may change innocuously; we only assert the +# ordering+scope tokens that this PR is responsible for. +_AGENT_MONO_RMW_RE = re.compile(r"atomicrmw\s+[^\n]*syncscope\(\"agent\"\)\s+monotonic") +_AGENT_MONO_CMPXCHG_RE = re.compile(r"cmpxchg\s+[^\n]*syncscope\(\"agent\"\)\s+monotonic(?:\s+monotonic)?") +_ANY_RMW_RE = re.compile(r"atomicrmw\s") +_ANY_CMPXCHG_RE = re.compile(r"cmpxchg\s") +_SYSTEM_SEQCST_RMW_RE = re.compile(r"atomicrmw\s+[^\n]*\sseq_cst\s*(?:,|$)") +_SYSTEM_SEQCST_CMPXCHG_RE = re.compile(r"cmpxchg\s+[^\n]*\sseq_cst\s+seq_cst") + + +# --------------------------------------------------------------------------- +# Child-process scaffolding. +# --------------------------------------------------------------------------- + + +_RUNTIME_BITCODE_MODULE_ID = "ModuleID = 'runtime_bitcode'" + + +def _read_all_ll(dump_dir: pathlib.Path) -> str: + """Concatenate every user-kernel IR dump in ``dump_dir`` into a + single string for grepping. We don't care which task the atomic + ended up in; we only care that the overall translation unit contains + the expected ops. + + The JIT writes one ``quadrants_kernel_amdgpu_llvm_ir_*.ll`` per + compiled module. That includes the runtime-bitcode side-module + (``atomic_exchange_i32`` / ``atomic_add_i32`` and friends in + ``runtime/llvm/runtime_module/atomic.h``), which is built from C++ + ``__atomic_*`` intrinsics under ``memory_order_seq_cst``. Those + helpers are not user-kernel atomics and are not reachable from any + of the kernels these tests compile, so a bare ``seq_cst atomicrmw`` + inside them must not be confused with a regression in the AMDGPU + codegen contract. We identify the runtime-bitcode dump by its + LLVM ``ModuleID`` header and skip it.""" + ll_files = sorted(dump_dir.glob("quadrants_kernel_amdgpu_llvm_ir_*.ll")) + assert ll_files, f"no LLVM IR dumps produced in {dump_dir}" + kernel_irs = [] + for p in ll_files: + text = p.read_text() + # The LLVM ``ModuleID`` header is always the first line of an + # IR dump, e.g. ``; ModuleID = 'kernel'`` or + # ``; ModuleID = 'runtime_bitcode'``. + first_line = text.split("\n", 1)[0] + if _RUNTIME_BITCODE_MODULE_ID in first_line: + continue + kernel_irs.append(text) + assert kernel_irs, f"no user-kernel LLVM IR dumps in {dump_dir} " f"(only runtime-bitcode dumps were produced)" + return "\n".join(kernel_irs) + + +def _run_kernel_dump(tmp_path: pathlib.Path, child_name: str) -> str: + """Run the named child entry point in a subprocess with IR dumping + enabled and return the concatenated IR text.""" + cmd = [sys.executable, __file__, child_name, str(tmp_path)] + env = dict(os.environ) + env["PYTHONPATH"] = env.get("PYTHONPATH", "") + os.pathsep + "." + proc = subprocess.run(cmd, capture_output=True, text=True, env=env) + if proc.returncode != RET_SUCCESS: + print(proc.stdout) + print("-" * 80) + print(proc.stderr) + assert proc.returncode == RET_SUCCESS, f"child '{child_name}' failed with exit code {proc.returncode}" + return _read_all_ll(tmp_path) + + +def _init_amdgpu_with_dump(dump_dir: str) -> None: + os.chdir(dump_dir) # print_kernel_llvm_ir writes to CWD + qd.init(arch=qd.amdgpu, print_kernel_llvm_ir=True, offline_cache=False) + + +# --------------------------------------------------------------------------- +# Per-op child entry points. Kept tiny so the IR dump is small and the +# atomicrmw / cmpxchg of interest is easy to grep. +# --------------------------------------------------------------------------- + + +def _child_atomic_add_i32(args): + _init_amdgpu_with_dump(args[0]) + c = qd.field(qd.i32, shape=()) + + @qd.kernel + def k(): + for _ in range(16): + qd.atomic_add(c[None], 1) + + k() + sys.exit(RET_SUCCESS) + + +def _child_atomic_add_f32(args): + _init_amdgpu_with_dump(args[0]) + c = qd.field(qd.f32, shape=()) + + @qd.kernel + def k(): + for _ in range(16): + qd.atomic_add(c[None], 1.0) + + k() + sys.exit(RET_SUCCESS) + + +def _child_atomic_or_i32(args): + _init_amdgpu_with_dump(args[0]) + flags = qd.field(qd.i32, shape=()) + + @qd.kernel + def k(): + for i in range(16): + qd.atomic_or(flags[None], 1 << (i % 8)) + + k() + sys.exit(RET_SUCCESS) + + +def _child_atomic_mul_i32(args): + _init_amdgpu_with_dump(args[0]) + v = qd.field(qd.i32, shape=()) + v[None] = 1 + + @qd.kernel + def k(): + for i in range(1, 4): + qd.atomic_mul(v[None], i) + + k() + sys.exit(RET_SUCCESS) + + +def _child_atomic_mul_f32(args): + _init_amdgpu_with_dump(args[0]) + v = qd.field(qd.f32, shape=()) + v[None] = 1.0 + + @qd.kernel + def k(): + for i in range(1, 4): + qd.atomic_mul(v[None], qd.cast(i, qd.f32)) + + k() + sys.exit(RET_SUCCESS) + + +def _child_atomic_min_f32(args): + _init_amdgpu_with_dump(args[0]) + v = qd.field(qd.f32, shape=()) + v[None] = 1e9 + + @qd.kernel + def k(): + for i in range(1, 16): + qd.atomic_min(v[None], qd.cast(i, qd.f32)) + + k() + sys.exit(RET_SUCCESS) + + +def _child_atomic_max_f64(args): + _init_amdgpu_with_dump(args[0]) + v = qd.field(qd.f64, shape=()) + v[None] = -1e18 + + @qd.kernel + def k(): + for i in range(1, 16): + qd.atomic_max(v[None], qd.cast(i, qd.f64)) + + k() + sys.exit(RET_SUCCESS) + + +_CHILDREN = { + fn.__name__: fn + for fn in [ + _child_atomic_add_i32, + _child_atomic_add_f32, + _child_atomic_or_i32, + _child_atomic_mul_i32, + _child_atomic_mul_f32, + _child_atomic_min_f32, + _child_atomic_max_f64, + ] +} + + +# --------------------------------------------------------------------------- +# Assertion helpers. +# --------------------------------------------------------------------------- + + +def _assert_no_system_seqcst_atomic(ir: str) -> None: + """All atomic instructions on AMDGPU should be agent+monotonic. Any + occurrence of bare ``seq_cst`` on an ``atomicrmw`` / ``cmpxchg`` is a + regression.""" + bad = _SYSTEM_SEQCST_RMW_RE.findall(ir) + assert not bad, f"unexpected seq_cst atomicrmw in AMDGPU IR: {bad[:3]}" + bad = _SYSTEM_SEQCST_CMPXCHG_RE.findall(ir) + assert not bad, f"unexpected seq_cst cmpxchg in AMDGPU IR: {bad[:3]}" + + +def _assert_has_agent_monotonic_rmw(ir: str) -> None: + assert _ANY_RMW_RE.search(ir), "expected at least one atomicrmw in IR" + assert _AGENT_MONO_RMW_RE.search(ir), ( + 'no atomicrmw carries syncscope("agent") monotonic; ' + "atomic was emitted with stronger ordering or different scope" + ) + + +def _assert_has_agent_monotonic_cmpxchg(ir: str) -> None: + assert _ANY_CMPXCHG_RE.search(ir), "expected at least one cmpxchg in IR" + assert _AGENT_MONO_CMPXCHG_RE.search(ir), ( + 'no cmpxchg carries syncscope("agent") monotonic; ' + "CAS-loop was emitted with stronger ordering or different scope" + ) + + +# --------------------------------------------------------------------------- +# Tests. +# --------------------------------------------------------------------------- + + +@test_utils.test(arch=[qd.amdgpu]) +def test_ir_atomic_add_i32_is_agent_monotonic(tmp_path: pathlib.Path): + ir = _run_kernel_dump(tmp_path, _child_atomic_add_i32.__name__) + _assert_has_agent_monotonic_rmw(ir) + _assert_no_system_seqcst_atomic(ir) + + +@test_utils.test(arch=[qd.amdgpu]) +def test_ir_atomic_add_f32_is_agent_monotonic(tmp_path: pathlib.Path): + ir = _run_kernel_dump(tmp_path, _child_atomic_add_f32.__name__) + _assert_has_agent_monotonic_rmw(ir) + _assert_no_system_seqcst_atomic(ir) + + +@test_utils.test(arch=[qd.amdgpu]) +def test_ir_atomic_or_i32_is_agent_monotonic(tmp_path: pathlib.Path): + ir = _run_kernel_dump(tmp_path, _child_atomic_or_i32.__name__) + _assert_has_agent_monotonic_rmw(ir) + _assert_no_system_seqcst_atomic(ir) + + +@test_utils.test(arch=[qd.amdgpu]) +def test_ir_atomic_mul_i32_is_agent_monotonic_cmpxchg(tmp_path: pathlib.Path): + # atomic_mul on int has no native AtomicRMW op -> CAS loop. + ir = _run_kernel_dump(tmp_path, _child_atomic_mul_i32.__name__) + _assert_has_agent_monotonic_cmpxchg(ir) + _assert_no_system_seqcst_atomic(ir) + + +@test_utils.test(arch=[qd.amdgpu]) +def test_ir_atomic_mul_f32_is_agent_monotonic_cmpxchg(tmp_path: pathlib.Path): + ir = _run_kernel_dump(tmp_path, _child_atomic_mul_f32.__name__) + _assert_has_agent_monotonic_cmpxchg(ir) + _assert_no_system_seqcst_atomic(ir) + + +# --------------------------------------------------------------------------- +# f32/f64 min/max — Phase 3.1 routes these through the overridden +# atomic_op_using_cas path so they pick up agent+monotonic from the same +# default_atomic_ordering() / default_atomic_scope() virtuals as the +# integer paths. Asserted strictly. +# --------------------------------------------------------------------------- + + +@test_utils.test(arch=[qd.amdgpu]) +def test_ir_atomic_min_f32_is_agent_monotonic(tmp_path: pathlib.Path): + ir = _run_kernel_dump(tmp_path, _child_atomic_min_f32.__name__) + _assert_no_system_seqcst_atomic(ir) + _assert_has_agent_monotonic_cmpxchg(ir) + + +@test_utils.test(arch=[qd.amdgpu]) +def test_ir_atomic_max_f64_is_agent_monotonic(tmp_path: pathlib.Path): + ir = _run_kernel_dump(tmp_path, _child_atomic_max_f64.__name__) + _assert_no_system_seqcst_atomic(ir) + _assert_has_agent_monotonic_cmpxchg(ir) + + +# --------------------------------------------------------------------------- +# Cross-cutting: a kernel that exercises both atomic_add and atomic_min +# on the same f32 field. Pre-Phase-3.1 this produced mixed-scope IR +# (agent+monotonic atomicrmw plus SeqCst+system runtime helper call). +# Post-Phase-3.1 both routes through default_atomic_*() and the IR +# contains only agent+monotonic atomics. +# --------------------------------------------------------------------------- + + +def _child_atomic_add_then_min_f32(args): + _init_amdgpu_with_dump(args[0]) + v = qd.field(qd.f32, shape=()) + v[None] = 0.0 + + @qd.kernel + def k_add(): + for _ in range(16): + qd.atomic_add(v[None], 1.0) + + @qd.kernel + def k_min(): + for i in range(1, 16): + qd.atomic_min(v[None], qd.cast(i, qd.f32)) + + k_add() + k_min() + sys.exit(RET_SUCCESS) + + +_CHILDREN[_child_atomic_add_then_min_f32.__name__] = _child_atomic_add_then_min_f32 + + +@test_utils.test(arch=[qd.amdgpu]) +def test_ir_no_mixed_scope_on_same_field(tmp_path: pathlib.Path): + ir = _run_kernel_dump(tmp_path, _child_atomic_add_then_min_f32.__name__) + _assert_no_system_seqcst_atomic(ir) + # Both atomic shapes should appear (atomicrmw for the add, cmpxchg + # for the min CAS loop) and both should carry agent+monotonic. + _assert_has_agent_monotonic_rmw(ir) + _assert_has_agent_monotonic_cmpxchg(ir) + + +# --------------------------------------------------------------------------- +# Subprocess dispatch (mirrors test_fn_attrs.py). +# --------------------------------------------------------------------------- + + +if __name__ == "__main__": + name, *rest = sys.argv[1:] + _CHILDREN[name](rest) From 6546712acf1a9b803e4a5450226aa0690a6daba0 Mon Sep 17 00:00:00 2001 From: rtmadduri Date: Fri, 15 May 2026 15:05:48 -0500 Subject: [PATCH 4/4] fix a broken codegen_amd --- quadrants/codegen/amdgpu/codegen_amdgpu.cpp | 143 -------------------- 1 file changed, 143 deletions(-) diff --git a/quadrants/codegen/amdgpu/codegen_amdgpu.cpp b/quadrants/codegen/amdgpu/codegen_amdgpu.cpp index cdb83f8df3..eef0ec9049 100644 --- a/quadrants/codegen/amdgpu/codegen_amdgpu.cpp +++ b/quadrants/codegen/amdgpu/codegen_amdgpu.cpp @@ -242,149 +242,6 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM { return nullptr; } - // Override the non-reduction integer atomic path so user-written - // qd.atomic_add(...) on integer SNode/Ndarray destinations (e.g. the - // work-queue counters in Genesis's narrowphase / broadphase / inequality - // constraint kernels, which dominate ~14% of the run after the tiled-wc - // CG body) emits Monotonic+agent on AMDGPU instead of SeqCst+system. - // Mirrors TaskCodeGenLLVM::integral_type_atomic; only the AtomicRMW - // ordering + sync scope differ. - llvm::Value *integral_type_atomic(AtomicOpStmt *stmt) override { - if (!is_integral(stmt->val->ret_type)) { - return nullptr; - } - if (stmt->op_type == AtomicOpType::mul) { - return atomic_op_using_cas( - llvm_val[stmt->dest], llvm_val[stmt->val], - [&](auto v1, auto v2) { return builder->CreateMul(v1, v2); }, - stmt->val->ret_type); - } - std::unordered_map bin_op; - bin_op[AtomicOpType::add] = llvm::AtomicRMWInst::BinOp::Add; - if (is_signed(stmt->val->ret_type)) { - bin_op[AtomicOpType::min] = llvm::AtomicRMWInst::BinOp::Min; - bin_op[AtomicOpType::max] = llvm::AtomicRMWInst::BinOp::Max; - } else { - bin_op[AtomicOpType::min] = llvm::AtomicRMWInst::BinOp::UMin; - bin_op[AtomicOpType::max] = llvm::AtomicRMWInst::BinOp::UMax; - } - bin_op[AtomicOpType::bit_and] = llvm::AtomicRMWInst::BinOp::And; - bin_op[AtomicOpType::bit_or] = llvm::AtomicRMWInst::BinOp::Or; - bin_op[AtomicOpType::bit_xor] = llvm::AtomicRMWInst::BinOp::Xor; - QD_ASSERT(bin_op.find(stmt->op_type) != bin_op.end()); - return builder->CreateAtomicRMW( - bin_op.at(stmt->op_type), llvm_val[stmt->dest], llvm_val[stmt->val], - llvm::MaybeAlign(0), llvm::AtomicOrdering::Monotonic, - amdgpu_atomic_scope()); - } - - // Override the non-reduction real atomic path. The f32 add fast path - // is what Genesis hits whenever it accumulates a float into a counter - // outside a TLS-promoted reduction. f16 and f64 paths fall through to - // the CAS-based atomic_op_using_cas below, which is also overridden. - llvm::Value *real_type_atomic(AtomicOpStmt *stmt) override { - if (!is_real(stmt->val->ret_type)) { - return nullptr; - } - PrimitiveTypeID prim_type = - stmt->val->ret_type->cast()->type; - AtomicOpType op = stmt->op_type; - if (prim_type == PrimitiveTypeID::f16) { - switch (op) { - case AtomicOpType::add: - return atomic_op_using_cas( - llvm_val[stmt->dest], llvm_val[stmt->val], - [&](auto v1, auto v2) { return builder->CreateFAdd(v1, v2); }, - stmt->val->ret_type); - case AtomicOpType::max: - return atomic_op_using_cas( - llvm_val[stmt->dest], llvm_val[stmt->val], - [&](auto v1, auto v2) { return builder->CreateMaxNum(v1, v2); }, - stmt->val->ret_type); - case AtomicOpType::min: - return atomic_op_using_cas( - llvm_val[stmt->dest], llvm_val[stmt->val], - [&](auto v1, auto v2) { return builder->CreateMinNum(v1, v2); }, - stmt->val->ret_type); - default: - break; - } - } - switch (op) { - case AtomicOpType::add: - return builder->CreateAtomicRMW( - llvm::AtomicRMWInst::FAdd, llvm_val[stmt->dest], - llvm_val[stmt->val], llvm::MaybeAlign(0), - llvm::AtomicOrdering::Monotonic, amdgpu_atomic_scope()); - case AtomicOpType::mul: - return atomic_op_using_cas( - llvm_val[stmt->dest], llvm_val[stmt->val], - [&](auto v1, auto v2) { return builder->CreateFMul(v1, v2); }, - stmt->val->ret_type); - default: - break; - } - // f32/f64 min/max fall through to the runtime atomic_min/max helpers - // declared in runtime.cpp; the base implementation is intentionally - // reused here (they are CAS-based regardless of ordering, and Genesis - // does not hit this path in the benchmark). - std::unordered_map> - atomics; - atomics[PrimitiveTypeID::f32][AtomicOpType::min] = "atomic_min_f32"; - atomics[PrimitiveTypeID::f64][AtomicOpType::min] = "atomic_min_f64"; - atomics[PrimitiveTypeID::f32][AtomicOpType::max] = "atomic_max_f32"; - atomics[PrimitiveTypeID::f64][AtomicOpType::max] = "atomic_max_f64"; - QD_ASSERT(atomics.find(prim_type) != atomics.end()); - QD_ASSERT(atomics.at(prim_type).find(op) != atomics.at(prim_type).end()); - return call(atomics.at(prim_type).at(op), llvm_val[stmt->dest], - llvm_val[stmt->val]); - } - - // CAS-loop fallback path on AMDGPU. Used for atomic_mul on integers / - // floats and for f32 atomic min/max in the reduction path. Identical - // to the base implementation except both AtomicOrdering args on the - // cmpxchg are relaxed to Monotonic at agent scope. - llvm::Value *atomic_op_using_cas( - llvm::Value *dest, - llvm::Value *val, - std::function op, - const DataType &type) override { - using namespace llvm; - BasicBlock *body = - BasicBlock::Create(*llvm_context, "while_loop_body", func); - BasicBlock *after_loop = - BasicBlock::Create(*llvm_context, "after_while", func); - - builder->CreateBr(body); - builder->SetInsertPoint(body); - - llvm::Value *old_val; - { - int bits = data_type_bits(type); - unsigned dest_as = dest->getType()->isPointerTy() - ? dest->getType()->getPointerAddressSpace() - : 0; - llvm::PointerType *typeIntPtr = - llvm::PointerType::get(*llvm_context, dest_as); - llvm::IntegerType *typeIntTy = get_integer_type(bits); - - old_val = builder->CreateLoad(val->getType(), dest); - auto new_val = op(old_val, val); - dest = builder->CreateBitCast(dest, typeIntPtr); - auto atomicCmpXchg = builder->CreateAtomicCmpXchg( - dest, builder->CreateBitCast(old_val, typeIntTy), - builder->CreateBitCast(new_val, typeIntTy), llvm::MaybeAlign(0), - AtomicOrdering::Monotonic, AtomicOrdering::Monotonic, - amdgpu_atomic_scope()); - auto ok = builder->CreateExtractValue(atomicCmpXchg, 1); - builder->CreateCondBr(builder->CreateNot(ok), body, after_loop); - } - - builder->SetInsertPoint(after_loop); - return old_val; - } - void visit(RangeForStmt *for_stmt) override { create_naive_range_for(for_stmt); }