[AMDGPU] Add threshold for DPP atomic optimizer on LDS atomics (#186762)

Add amdgpu.expected.active.lane metadata which can be applied to operations by a 
compiler front-end.  Do not apply the DPP atomic optimizer to LDS atomics where
less than five active lanes are expected.
This is an empirically derived threshold based on GFX11 and GFX12 testing.

---------

Co-authored-by: YaFan <YaFan.Tao@amd.com>

llvm/docs/AMDGPUUsage.rst

@@ -2440,6 +2440,28 @@ and
 
   !0 = !{}
 
+.. _amdgpu_expected_active_lanes:
+
+'``amdgpu.expected.active.lanes``' Metadata
+-------------------------------------------------
+
+A profiling-derived hint describing how many lanes of the wavefront are
+expected to be active. The metadata has a single ``i32`` constant operand
+giving the expected number of active lanes as an absolute count (the same
+value is used for both wave32 and wave64 targets). This structure is
+enforced by the verifier.
+
+The AMDGPU atomic optimizer uses this metadata on LDS atomics to decide
+whether to skip the DPP optimization.
+
+.. code-block:: llvm
+
+  ; Few expected active lanes: the atomic optimizer may skip the DPP optimization.
+  %old0 = atomicrmw add ptr addrspace(3) @lds, i32 %val acq_rel, !amdgpu.expected.active.lanes !{i32 4}
+
+  ; Many expected active lanes: the DPP optimization is not skipped.
+  %old1 = atomicrmw add ptr addrspace(3) @lds, i32 %val acq_rel, !amdgpu.expected.active.lanes !{i32 32}
+
 
 LLVM IR Attributes
 ==================

llvm/lib/IR/Verifier.cpp

@@ -5907,6 +5907,17 @@ void Verifier::visitInstruction(Instruction &I) {
   if (MDNode *MD = I.getMetadata(LLVMContext::MD_mem_cache_hint))
     visitMemCacheHintMetadata(I, MD);
 
+  if (MDNode *MD = I.getMetadata("amdgpu.expected.active.lanes")) {
+    Check(MD->getNumOperands() == 1,
+          "!amdgpu.expected.active.lanes must have exactly one operand", &I,
+          MD);
+    ConstantInt *CI =
+        mdconst::dyn_extract_or_null<ConstantInt>(MD->getOperand(0));
+    Check(CI && CI->getType()->isIntegerTy(32),
+          "!amdgpu.expected.active.lanes operand must be an i32 constant", &I,
+          MD);
+  }
+
   if (MDNode *N = I.getDebugLoc().getAsMDNode()) {
     CheckDI(isa<DILocation>(N), "invalid !dbg metadata attachment", &I, N);
     visitMDNode(*N, AreDebugLocsAllowed::Yes);

llvm/lib/Target/AMDGPU/AMDGPUAtomicOptimizer.cpp

@@ -45,6 +45,7 @@ struct ReplacementInfo {
   AtomicRMWInst::BinOp Op;
   unsigned ValIdx;
   bool ValDivergent;
+  bool IsLDS;
 };
 
 class AMDGPUAtomicOptimizer : public FunctionPass {
@@ -87,7 +88,7 @@ private:
                        BasicBlock *ComputeLoop, BasicBlock *ComputeEnd) const;
 
   void optimizeAtomic(Instruction &I, AtomicRMWInst::BinOp Op, unsigned ValIdx,
-                      bool ValDivergent) const;
+                      bool ValDivergent, bool IsLDS) const;
 
 public:
   AMDGPUAtomicOptimizerImpl() = delete;
@@ -161,8 +162,8 @@ bool AMDGPUAtomicOptimizerImpl::run() {
   if (ToReplace.empty())
     return false;
 
-  for (auto &[I, Op, ValIdx, ValDivergent] : ToReplace)
-    optimizeAtomic(*I, Op, ValIdx, ValDivergent);
+  for (auto &[I, Op, ValIdx, ValDivergent, IsLDS] : ToReplace)
+    optimizeAtomic(*I, Op, ValIdx, ValDivergent, IsLDS);
   ToReplace.clear();
   return true;
 }
@@ -241,10 +242,12 @@ void AMDGPUAtomicOptimizerImpl::visitAtomicRMWInst(AtomicRMWInst &I) {
       return;
   }
 
+  const bool IsLDS = I.getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS;
+
   // If we get here, we can optimize the atomic using a single wavefront-wide
   // atomic operation to do the calculation for the entire wavefront, so
   // remember the instruction so we can come back to it.
-  ToReplace.push_back({&I, Op, ValIdx, ValDivergent});
+  ToReplace.push_back({&I, Op, ValIdx, ValDivergent, IsLDS});
 }
 
 void AMDGPUAtomicOptimizerImpl::visitIntrinsicInst(IntrinsicInst &I) {
@@ -335,7 +338,8 @@ void AMDGPUAtomicOptimizerImpl::visitIntrinsicInst(IntrinsicInst &I) {
   // If we get here, we can optimize the atomic using a single wavefront-wide
   // atomic operation to do the calculation for the entire wavefront, so
   // remember the instruction so we can come back to it.
-  ToReplace.push_back({&I, Op, ValIdx, ValDivergent});
+  // Buffer atomics are never LDS.
+  ToReplace.push_back({&I, Op, ValIdx, ValDivergent, /*IsLDS=*/false});
 }
 
 // Use the builder to create the non-atomic counterpart of the specified
@@ -646,7 +650,23 @@ static Value *buildMul(IRBuilder<> &B, Value *LHS, Value *RHS) {
 void AMDGPUAtomicOptimizerImpl::optimizeAtomic(Instruction &I,
                                                AtomicRMWInst::BinOp Op,
                                                unsigned ValIdx,
-                                               bool ValDivergent) const {
+                                               bool ValDivergent,
+                                               bool IsLDS) const {
+  // Don't generate a DPP scan if !amdgpu.expected.active.lane hint indicates
+  // insufficient lanes to offset fixed overhead.
+
+  // FIXME: The threshold was tuned empirically on gfx11 and gfx12. The DPP scan
+  // overhead differs across subtargets, so the break-even point may differ too;
+  // this may need to become subtarget-dependent.
+  if (IsLDS && ValDivergent && ScanImpl == ScanOptions::DPP) {
+    if (MDNode *MD = I.getMetadata("amdgpu.expected.active.lanes")) {
+      auto *CI = mdconst::extract<ConstantInt>(MD->getOperand(0));
+      constexpr unsigned ActiveLanesThreshold = 5;
+      if (CI->getValue().ule(ActiveLanesThreshold))
+        return;
+    }
+  }
+
   // Start building just before the instruction.
   IRBuilder<> B(&I);
 

llvm/test/Verifier/amdgpu-expected-active-lanes.ll

@@ -0,0 +1,30 @@
+; RUN: not llvm-as -disable-output %s 2>&1 | FileCheck %s
+
+; CHECK-DAG: !amdgpu.expected.active.lanes must have exactly one operand
+define i32 @no_operands(ptr %p, i32 %v) {
+  %old = atomicrmw add ptr %p, i32 %v seq_cst, !amdgpu.expected.active.lanes !0
+  ret i32 %old
+}
+
+; CHECK-DAG: !amdgpu.expected.active.lanes must have exactly one operand
+define i32 @too_many_operands(ptr %p, i32 %v) {
+  %old = atomicrmw add ptr %p, i32 %v seq_cst, !amdgpu.expected.active.lanes !1
+  ret i32 %old
+}
+
+; CHECK-DAG: !amdgpu.expected.active.lanes operand must be an i32 constant
+define i32 @operand_not_a_constant(ptr %p, i32 %v) {
+  %old = atomicrmw add ptr %p, i32 %v seq_cst, !amdgpu.expected.active.lanes !2
+  ret i32 %old
+}
+
+; CHECK-DAG: !amdgpu.expected.active.lanes operand must be an i32 constant
+define i32 @operand_wrong_width(ptr %p, i32 %v) {
+  %old = atomicrmw add ptr %p, i32 %v seq_cst, !amdgpu.expected.active.lanes !3
+  ret i32 %old
+}
+
+!0 = !{}
+!1 = !{i32 4, i32 5}
+!2 = !{!"not a constant"}
+!3 = !{i64 4}