Relaxed atomic ordering and AMDGPU-scoped atomics

docs/source/user_guide/index.md

@@ -38,6 +38,7 @@ interop
 
 cuda_graph
 perf_dispatch
+amdgpu_atomics
 ```
 
 ```{toctree}

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,6 +161,37 @@ class TaskCodeGenAMDGPU : public TaskCodeGenLLVM {
 #undef UNARY_STD
   }
 
+  // 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
   // reduce_* helpers. The runtime helpers expect addrspace(0) pointers,
   // but SNode destinations arrive in addrspace(1). Calling the helpers
@@ -184,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::SequentiallyConsistent);
+            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::SequentiallyConsistent);
+            default_atomic_ordering(), default_atomic_scope());
       } else if (op == AtomicOpType::min) {
         return atomic_op_using_cas(
             dest, val,
@@ -594,8 +630,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,

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<PrimitiveTypeID,
                      std::unordered_map<AtomicOpType, std::string>>
       atomics;

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;

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;