[AMDGPU] Check atomics aliasing in the clobbering annotation

[rampitec]

Summary

MemorySSA considers any atomic a def to any operation it dominates
just like a barrier or fence. That is correct from memory state
perspective, but not required for the no-clobber metadata since
we are not using it for reordering. Skip such atomics during the
scan just like a barrier if it does not alias with the load.

[arsenm]

llvm/test/CodeGen/AMDGPU/noclobber-barrier.ll

452

volatile makes this test stricter, should probably not have it

505

Ditto

601

Is there already a test that verifies load atomic is not optimizable?

[rampitec]

Removed volatile from atomics in tests.

[rampitec]

llvm/test/CodeGen/AMDGPU/noclobber-barrier.ll

601

Yes, clobber_by_atomic_load in this file.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUAnnotateUniformValues.cpp

141

Typo unversal

[rampitec]

Fixed typo in the comment.

[rampitec]

Fixed not calling getClobberingMemoryAccess, it was gathering more defs for inspection than needed. Added test for this case.

[arsenm]

llvm/test/CodeGen/AMDGPU/noclobber-barrier.ll

623

No volatile

[foad]

llvm/lib/Target/AMDGPU/AMDGPUAnnotateUniformValues.cpp

52

Should add an addUsedIfAvailable call for AA here. But I think MemorySSA already requires it, so you could make it a hard requirement here too?

[foad]

MemorySSA considers any atomic a def to any operation it dominates

I guess this is because of the problem described by the TODO in MemorySSA::createNewAccess:

// The isOrdered check is used to ensure that volatiles end up as defs
// (atomics end up as ModRef right now anyway).  Until we separate the
// ordering chain from the memory chain, this enables people to see at least
// some relative ordering to volatiles.  Note that getClobberingMemoryAccess
// will still give an answer that bypasses other volatile loads.  TODO:
// Separate memory aliasing and ordering into two different chains so that
// we can precisely represent both "what memory will this read/write/is
// clobbered by" and "what instructions can I move this past".

?

Maybe a cleaner way to handle this would be to implement a new MemorySSAWalker subclass where getClobberingMemoryAccess(I) ignores ordering, and only returns a def that mayAlias I. But I suppose that's more work.

[rampitec]

In D118661#3287271, @foad wrote:

MemorySSA considers any atomic a def to any operation it dominates

I guess this is because of the problem described by the TODO in MemorySSA::createNewAccess:

// The isOrdered check is used to ensure that volatiles end up as defs
// (atomics end up as ModRef right now anyway).  Until we separate the
// ordering chain from the memory chain, this enables people to see at least
// some relative ordering to volatiles.  Note that getClobberingMemoryAccess
// will still give an answer that bypasses other volatile loads.  TODO:
// Separate memory aliasing and ordering into two different chains so that
// we can precisely represent both "what memory will this read/write/is
// clobbered by" and "what instructions can I move this past".

?

Maybe a cleaner way to handle this would be to implement a new MemorySSAWalker subclass where getClobberingMemoryAccess(I) ignores ordering, and only returns a def that mayAlias I. But I suppose that's more work.

Yes, that's the issue. This is how ordering is modeled.

[rampitec]

• Added AA as required analysis.
• Dropped volatile in the test.

[rampitec]

Fixed rebase artifact.

[vpykhtin]

LGTM

llvm/lib/Target/AMDGPU/AMDGPUAnnotateUniformValues.cpp

 /// This pass adds amdgpu.uniform metadata to IR values so this information
 /// can be used during instruction selection.
 //
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPU.h"
 #include "Utils/AMDGPUBaseInfo.h"
 #include "llvm/ADT/SmallSet.h"
+#include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/LegacyDivergenceAnalysis.h"
 #include "llvm/Analysis/MemorySSA.h"
 #include "llvm/IR/InstVisitor.h"
 #include "llvm/IR/IntrinsicsAMDGPU.h"
 #include "llvm/InitializePasses.h"
 
 #define DEBUG_TYPE "amdgpu-annotate-uniform"
 
 using namespace llvm;
 
 namespace {
 
 class AMDGPUAnnotateUniformValues : public FunctionPass,
                        public InstVisitor<AMDGPUAnnotateUniformValues> {
   LegacyDivergenceAnalysis *DA;
   MemorySSA *MSSA;
+  AliasAnalysis *AA;
   DenseMap<Value*, GetElementPtrInst*> noClobberClones;
   bool isEntryFunc;
 
 public:
   static char ID;
   AMDGPUAnnotateUniformValues() :
     FunctionPass(ID) { }
   bool doInitialization(Module &M) override;
   bool runOnFunction(Function &F) override;
   StringRef getPassName() const override {
     return "AMDGPU Annotate Uniform Values";
   }
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<LegacyDivergenceAnalysis>();
     AU.addRequired<MemorySSAWrapperPass>();
+    AU.addRequired<AAResultsWrapperPass>();
     AU.setPreservesAll();
  }
 
   void visitBranchInst(BranchInst &I);
   void visitLoadInst(LoadInst &I);
   bool isClobberedInFunction(LoadInst * Load);
 };
 
 } // End anonymous namespace
 
 INITIALIZE_PASS_BEGIN(AMDGPUAnnotateUniformValues, DEBUG_TYPE,
                       "Add AMDGPU uniform metadata", false, false)
 INITIALIZE_PASS_DEPENDENCY(LegacyDivergenceAnalysis)
 INITIALIZE_PASS_DEPENDENCY(MemorySSAWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
 INITIALIZE_PASS_END(AMDGPUAnnotateUniformValues, DEBUG_TYPE,
                     "Add AMDGPU uniform metadata", false, false)
 
 char AMDGPUAnnotateUniformValues::ID = 0;
 
 static void setUniformMetadata(Instruction *I) {
   I->setMetadata("amdgpu.uniform", MDNode::get(I->getContext(), {}));
 }
 static void setNoClobberMetadata(Instruction *I) {
   I->setMetadata("amdgpu.noclobber", MDNode::get(I->getContext(), {}));
 }
 
 bool AMDGPUAnnotateUniformValues::isClobberedInFunction(LoadInst *Load) {
   MemorySSAWalker *Walker = MSSA->getWalker();
   SmallVector<MemoryAccess *> WorkList{Walker->getClobberingMemoryAccess(Load)};
   SmallSet<MemoryAccess *, 8> Visited;
   MemoryLocation Loc(MemoryLocation::get(Load));
 
-  const auto isReallyAClobber = [](MemoryDef *Def) -> bool {
+  const auto isReallyAClobber = [this, Load](MemoryDef *Def) -> bool {
     Instruction *DefInst = Def->getMemoryInst();
     LLVM_DEBUG(dbgs() << "  Def: " << *DefInst << '\n');
 
     if (isa<FenceInst>(DefInst))
       return false;
 
     if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(DefInst)) {
       switch (II->getIntrinsicID()) {
       case Intrinsic::amdgcn_s_barrier:
       case Intrinsic::amdgcn_wave_barrier:
         return false;
       default:
         break;
       }
     }
 
+    // Ignore atomics not aliasing with the original load, any atomic is a
+    // universal MemoryDef from MSSA's point of view too, just like a fence.
+    const auto checkNoAlias = [this, Load](auto I) -> bool {
+      return I && AA->isNoAlias(I->getPointerOperand(),
+                                Load->getPointerOperand());
+    };
+
+    if (checkNoAlias(dyn_cast<AtomicCmpXchgInst>(DefInst)) ||
+        checkNoAlias(dyn_cast<AtomicRMWInst>(DefInst)))
+      return false;
+
     return true;
   };
 
   LLVM_DEBUG(dbgs() << "Checking clobbering of: " << *Load << '\n');
 
   // Start with a nearest dominating clobbering access, it will be either
   // live on entry (nothing to do, load is not clobbered), MemoryDef, or
   // MemoryPhi if several MemoryDefs can define this memory state. In that
while (!WorkList.empty()) {
 
     if (MemoryDef *Def = dyn_cast<MemoryDef>(MA)) {
       if (isReallyAClobber(Def)) {
         LLVM_DEBUG(dbgs() << "      -> load is clobbered\n");
         return true;
       }
 
       WorkList.push_back(
           Walker->getClobberingMemoryAccess(Def->getDefiningAccess(), Loc));
       continue;
     }
 
     const MemoryPhi *Phi = cast<MemoryPhi>(MA);
     for (auto &Use : Phi->incoming_values())
       WorkList.push_back(cast<MemoryAccess>(&Use));
   }
 

 }
 
 bool AMDGPUAnnotateUniformValues::runOnFunction(Function &F) {
   if (skipFunction(F))
     return false;
 
   DA = &getAnalysis<LegacyDivergenceAnalysis>();
   MSSA = &getAnalysis<MemorySSAWrapperPass>().getMSSA();
+  AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
   isEntryFunc = AMDGPU::isEntryFunctionCC(F.getCallingConv());
 
   visit(F);
   noClobberClones.clear();
   return true;
 }
 
 FunctionPass *
 llvm::createAMDGPUAnnotateUniformValues() {
   return new AMDGPUAnnotateUniformValues();
 }

llvm/test/CodeGen/AMDGPU/noclobber-barrier.ll

entry:
   tail call void @llvm.amdgcn.s.barrier()
   fence syncscope("workgroup") acquire
   %gep = getelementptr inbounds i32, i32 addrspace(1)* %in, i64 0
   %ld = load i32, i32 addrspace(1)* %gep, align 4
   store i32 %ld, i32 addrspace(1)* %out, align 4
   ret void
 }
 
+; GCN-LABEL: {{^}}no_alias_atomic_rmw_relaxed:
+; GCN: ds_add_u32
+; GCN: s_load_dword s
+; GCN-NOT: global_load_dword
+; GCN: global_store_dword
+define protected amdgpu_kernel void @no_alias_atomic_rmw_relaxed(i32 addrspace(1)* %in, i32 addrspace(1)* %out) {
+; CHECK-LABEL: @no_alias_atomic_rmw_relaxed(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[UNUSED:%.*]] = atomicrmw add i32 addrspace(3)* @LDS, i32 5 monotonic, align 4
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i32, i32 addrspace(1)* [[IN:%.*]], i64 0, !amdgpu.uniform !0, !amdgpu.noclobber !0
+; CHECK-NEXT:    [[LD:%.*]] = load i32, i32 addrspace(1)* [[GEP]], align 4
+; CHECK-NEXT:    store i32 [[LD]], i32 addrspace(1)* [[OUT:%.*]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %unused = atomicrmw add i32 addrspace(3)* @LDS, i32 5 monotonic
+  %gep = getelementptr inbounds i32, i32 addrspace(1)* %in, i64 0
+  %ld = load i32, i32 addrspace(1)* %gep, align 4
+  store i32 %ld, i32 addrspace(1)* %out, align 4
+  ret void
+}
+
+; GCN-LABEL: {{^}}no_alias_atomic_cmpxchg:
+; GCN: ds_cmpst_b32
+; GCN: s_load_dword s
+; GCN-NOT: global_load_dword
+; GCN: global_store_dword
+define protected amdgpu_kernel void @no_alias_atomic_cmpxchg(i32 addrspace(1)* %in, i32 addrspace(1)* %out, i32 %swap) {
+; CHECK-LABEL: @no_alias_atomic_cmpxchg(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[UNUSED:%.*]] = cmpxchg i32 addrspace(3)* @LDS, i32 7, i32 [[SWAP:%.*]] seq_cst monotonic, align 4
+; CHECK-NEXT:    fence syncscope("workgroup") release
+; CHECK-NEXT:    tail call void @llvm.amdgcn.s.barrier()
+; CHECK-NEXT:    fence syncscope("workgroup") acquire
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i32, i32 addrspace(1)* [[IN:%.*]], i64 0, !amdgpu.uniform !0, !amdgpu.noclobber !0
+; CHECK-NEXT:    [[LD:%.*]] = load i32, i32 addrspace(1)* [[GEP]], align 4
+; CHECK-NEXT:    store i32 [[LD]], i32 addrspace(1)* [[OUT:%.*]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %unused = cmpxchg i32 addrspace(3)* @LDS, i32 7, i32 %swap seq_cst monotonic
+  fence syncscope("workgroup") release
+  tail call void @llvm.amdgcn.s.barrier()
+  fence syncscope("workgroup") acquire
+  %gep = getelementptr inbounds i32, i32 addrspace(1)* %in, i64 0
+  %ld = load i32, i32 addrspace(1)* %gep, align 4
+  store i32 %ld, i32 addrspace(1)* %out, align 4
+  ret void
+}
+
+; GCN-LABEL: {{^}}no_alias_atomic_rmw:
+; GCN: ds_add_u32
+; GCN: s_load_dword s
+; GCN-NOT: global_load_dword
+; GCN: global_store_dword
+define protected amdgpu_kernel void @no_alias_atomic_rmw(i32 addrspace(1)* %in, i32 addrspace(1)* %out) {
+; CHECK-LABEL: @no_alias_atomic_rmw(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[UNUSED:%.*]] = atomicrmw add i32 addrspace(3)* @LDS, i32 5 seq_cst, align 4
+; CHECK-NEXT:    fence syncscope("workgroup") release
+; CHECK-NEXT:    tail call void @llvm.amdgcn.s.barrier()
+; CHECK-NEXT:    fence syncscope("workgroup") acquire
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i32, i32 addrspace(1)* [[IN:%.*]], i64 0, !amdgpu.uniform !0, !amdgpu.noclobber !0
+; CHECK-NEXT:    [[LD:%.*]] = load i32, i32 addrspace(1)* [[GEP]], align 4
+; CHECK-NEXT:    store i32 [[LD]], i32 addrspace(1)* [[OUT:%.*]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %unused = atomicrmw add i32 addrspace(3)* @LDS, i32 5 seq_cst
+  fence syncscope("workgroup") release
+  tail call void @llvm.amdgcn.s.barrier()
+  fence syncscope("workgroup") acquire
+  %gep = getelementptr inbounds i32, i32 addrspace(1)* %in, i64 0
+  %ld = load i32, i32 addrspace(1)* %gep, align 4
+  store i32 %ld, i32 addrspace(1)* %out, align 4
+  ret void
+}
+
+; GCN-LABEL: {{^}}may_alias_atomic_cmpxchg:
+; GCN: global_atomic_cmpswap
+; GCN: global_load_dword
+; GCN: global_store_dword
+define protected amdgpu_kernel void @may_alias_atomic_cmpxchg(i32 addrspace(1)* %in, i32 addrspace(1)* %out, i32 %swap) {
+; CHECK-LABEL: @may_alias_atomic_cmpxchg(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[UNUSED:%.*]] = cmpxchg i32 addrspace(1)* [[OUT:%.*]], i32 7, i32 [[SWAP:%.*]] seq_cst monotonic, align 4
+; CHECK-NEXT:    fence syncscope("workgroup") release
+; CHECK-NEXT:    tail call void @llvm.amdgcn.s.barrier()
+; CHECK-NEXT:    fence syncscope("workgroup") acquire
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i32, i32 addrspace(1)* [[IN:%.*]], i64 0, !amdgpu.uniform !0
+; CHECK-NEXT:    [[LD:%.*]] = load i32, i32 addrspace(1)* [[GEP]], align 4
+; CHECK-NEXT:    store i32 [[LD]], i32 addrspace(1)* [[OUT]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %unused = cmpxchg i32 addrspace(1)* %out, i32 7, i32 %swap seq_cst monotonic
+  fence syncscope("workgroup") release
+  tail call void @llvm.amdgcn.s.barrier()
+  fence syncscope("workgroup") acquire
+  %gep = getelementptr inbounds i32, i32 addrspace(1)* %in, i64 0
+  %ld = load i32, i32 addrspace(1)* %gep, align 4
+  store i32 %ld, i32 addrspace(1)* %out, align 4
+  ret void
+}
+
+; GCN-LABEL: {{^}}may_alias_atomic_rmw:
+; GCN: global_atomic_add
+; GCN: global_load_dword
+; GCN: global_store_dword
+define protected amdgpu_kernel void @may_alias_atomic_rmw(i32 addrspace(1)* %in, i32 addrspace(1)* %out) {
+; CHECK-LABEL: @may_alias_atomic_rmw(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[UNUSED:%.*]] = atomicrmw add i32 addrspace(1)* [[OUT:%.*]], i32 5 seq_cst, align 4
+; CHECK-NEXT:    fence syncscope("workgroup") release
+; CHECK-NEXT:    tail call void @llvm.amdgcn.s.barrier()
+; CHECK-NEXT:    fence syncscope("workgroup") acquire
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i32, i32 addrspace(1)* [[IN:%.*]], i64 0, !amdgpu.uniform !0
+; CHECK-NEXT:    [[LD:%.*]] = load i32, i32 addrspace(1)* [[GEP]], align 4
+; CHECK-NEXT:    store i32 [[LD]], i32 addrspace(1)* [[OUT]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %unused = atomicrmw add i32 addrspace(1)* %out, i32 5 seq_cst
+  fence syncscope("workgroup") release
+  tail call void @llvm.amdgcn.s.barrier()
+  fence syncscope("workgroup") acquire
+  %gep = getelementptr inbounds i32, i32 addrspace(1)* %in, i64 0
+  %ld = load i32, i32 addrspace(1)* %gep, align 4
+  store i32 %ld, i32 addrspace(1)* %out, align 4
+  ret void
+}
+
+; GCN-LABEL: {{^}}no_alias_atomic_rmw_then_clobber:
+; CGN: global_store_dword
+; CGN: global_store_dword
+; GCN: ds_add_u32
+; GCN: global_load_dword
+; GCN: global_store_dword
+define protected amdgpu_kernel void @no_alias_atomic_rmw_then_clobber(i32 addrspace(1)* %in, i32 addrspace(1)* %out, i32 addrspace(1)* noalias %noalias) {
+; CHECK-LABEL: @no_alias_atomic_rmw_then_clobber(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    store i32 1, i32 addrspace(1)* [[OUT:%.*]], align 4
+; CHECK-NEXT:    store i32 2, i32 addrspace(1)* [[NOALIAS:%.*]], align 4
+; CHECK-NEXT:    [[UNUSED:%.*]] = atomicrmw add i32 addrspace(3)* @LDS, i32 5 seq_cst, align 4
+; CHECK-NEXT:    fence syncscope("workgroup") release
+; CHECK-NEXT:    tail call void @llvm.amdgcn.s.barrier()
+; CHECK-NEXT:    fence syncscope("workgroup") acquire
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i32, i32 addrspace(1)* [[IN:%.*]], i64 0, !amdgpu.uniform !0
+; CHECK-NEXT:    [[LD:%.*]] = load i32, i32 addrspace(1)* [[GEP]], align 4
+; CHECK-NEXT:    store i32 [[LD]], i32 addrspace(1)* [[OUT]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  store i32 1, i32 addrspace(1)* %out, align 4
+  store i32 2, i32 addrspace(1)* %noalias, align 4
+  %unused = atomicrmw add i32 addrspace(3)* @LDS, i32 5 seq_cst
+  fence syncscope("workgroup") release
+  tail call void @llvm.amdgcn.s.barrier()
+  fence syncscope("workgroup") acquire
+  %gep = getelementptr inbounds i32, i32 addrspace(1)* %in, i64 0
+  %ld = load i32, i32 addrspace(1)* %gep, align 4
+  store i32 %ld, i32 addrspace(1)* %out, align 4
+  ret void
+}
+
+; GCN-LABEL: {{^}}no_alias_atomic_rmw_then_no_alias_store:
+; CGN: global_store_dword
+; GCN: ds_add_u32
+; GCN: s_load_dword s
+; GCN-NOT: global_load_dword
+; GCN: global_store_dword
+define protected amdgpu_kernel void @no_alias_atomic_rmw_then_no_alias_store(i32 addrspace(1)* %in, i32 addrspace(1)* %out, i32 addrspace(1)* noalias %noalias) {
+; CHECK-LABEL: @no_alias_atomic_rmw_then_no_alias_store(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    store i32 2, i32 addrspace(1)* [[NOALIAS:%.*]], align 4
+; CHECK-NEXT:    [[UNUSED:%.*]] = atomicrmw add i32 addrspace(3)* @LDS, i32 5 seq_cst, align 4
+; CHECK-NEXT:    fence syncscope("workgroup") release
+; CHECK-NEXT:    tail call void @llvm.amdgcn.s.barrier()
+; CHECK-NEXT:    fence syncscope("workgroup") acquire
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i32, i32 addrspace(1)* [[IN:%.*]], i64 0, !amdgpu.uniform !0, !amdgpu.noclobber !0
+; CHECK-NEXT:    [[LD:%.*]] = load i32, i32 addrspace(1)* [[GEP]], align 4
+; CHECK-NEXT:    store i32 [[LD]], i32 addrspace(1)* [[OUT:%.*]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  store i32 2, i32 addrspace(1)* %noalias, align 4
+  %unused = atomicrmw add i32 addrspace(3)* @LDS, i32 5 seq_cst
+  fence syncscope("workgroup") release
+  tail call void @llvm.amdgcn.s.barrier()
+  fence syncscope("workgroup") acquire
+  %gep = getelementptr inbounds i32, i32 addrspace(1)* %in, i64 0
+  %ld = load i32, i32 addrspace(1)* %gep, align 4
+  store i32 %ld, i32 addrspace(1)* %out, align 4
+  ret void
+}
+
 declare void @llvm.amdgcn.s.barrier()
 declare void @llvm.amdgcn.wave.barrier()