clang/AMDGPU: Set noalias.addrspace metadata on atomicrmw (llvm#102462)

Author: arsenm

clang/include/clang/AST/Expr.h

@@ -6777,6 +6777,17 @@ class AtomicExpr : public Expr {
            getOp() <= AO__opencl_atomic_store;
   }
 
+  bool isHIP() const {
+    return Op >= AO__hip_atomic_compare_exchange_strong &&
+           Op <= AO__hip_atomic_store;
+  }
+
+  /// Return true if atomics operations targeting allocations in private memory
+  /// are undefined.
+  bool threadPrivateMemoryAtomicsAreUndefined() const {
+    return isOpenCL() || isHIP();
+  }
+
   SourceLocation getBuiltinLoc() const { return BuiltinLoc; }
   SourceLocation getRParenLoc() const { return RParenLoc; }
 

clang/include/clang/Basic/LangOptions.h

@@ -698,6 +698,14 @@ class LangOptions : public LangOptionsBase {
     return ConvergentFunctions;
   }
 
+  /// Return true if atomicrmw operations targeting allocations in private
+  /// memory are undefined.
+  bool threadPrivateMemoryAtomicsAreUndefined() const {
+    // Should be false for OpenMP.
+    // TODO: Should this be true for SYCL?
+    return OpenCL || CUDA;
+  }
+
   /// Return the OpenCL C or C++ version as a VersionTuple.
   VersionTuple getOpenCLVersionTuple() const;
 

clang/lib/CodeGen/CGAtomic.cpp

@@ -389,6 +389,7 @@ static void emitAtomicCmpXchg(CodeGenFunction &CGF, AtomicExpr *E, bool IsWeak,
       Ptr, Expected, Desired, SuccessOrder, FailureOrder, Scope);
   Pair->setVolatile(E->isVolatile());
   Pair->setWeak(IsWeak);
+  CGF.getTargetHooks().setTargetAtomicMetadata(CGF, *Pair, E);
 
   // Cmp holds the result of the compare-exchange operation: true on success,
   // false on failure.
@@ -727,7 +728,7 @@ static void EmitAtomicOp(CodeGenFunction &CGF, AtomicExpr *E, Address Dest,
 
   llvm::Value *LoadVal1 = CGF.Builder.CreateLoad(Val1);
   llvm::AtomicRMWInst *RMWI =
-      CGF.emitAtomicRMWInst(Op, Ptr, LoadVal1, Order, Scope);
+      CGF.emitAtomicRMWInst(Op, Ptr, LoadVal1, Order, Scope, E);
   RMWI->setVolatile(E->isVolatile());
 
   // For __atomic_*_fetch operations, perform the operation again to
@@ -2048,11 +2049,11 @@ std::pair<RValue, llvm::Value *> CodeGenFunction::EmitAtomicCompareExchange(
 llvm::AtomicRMWInst *
 CodeGenFunction::emitAtomicRMWInst(llvm::AtomicRMWInst::BinOp Op, Address Addr,
                                    llvm::Value *Val, llvm::AtomicOrdering Order,
-                                   llvm::SyncScope::ID SSID) {
-
+                                   llvm::SyncScope::ID SSID,
+                                   const AtomicExpr *AE) {
   llvm::AtomicRMWInst *RMW =
       Builder.CreateAtomicRMW(Op, Addr, Val, Order, SSID);
-  getTargetHooks().setTargetAtomicMetadata(*this, *RMW);
+  getTargetHooks().setTargetAtomicMetadata(*this, *RMW, AE);
   return RMW;
 }
 

clang/lib/CodeGen/CodeGenFunction.h

@@ -4166,7 +4166,8 @@ class CodeGenFunction : public CodeGenTypeCache {
   llvm::AtomicRMWInst *emitAtomicRMWInst(
       llvm::AtomicRMWInst::BinOp Op, Address Addr, llvm::Value *Val,
       llvm::AtomicOrdering Order = llvm::AtomicOrdering::SequentiallyConsistent,
-      llvm::SyncScope::ID SSID = llvm::SyncScope::System);
+      llvm::SyncScope::ID SSID = llvm::SyncScope::System,
+      const AtomicExpr *AE = nullptr);
 
   void EmitAtomicUpdate(LValue LVal, llvm::AtomicOrdering AO,
                         const llvm::function_ref<RValue(RValue)> &UpdateOp,

clang/lib/CodeGen/TargetInfo.h

@@ -336,7 +336,9 @@ class TargetCodeGenInfo {
 
   /// Allow the target to apply other metadata to an atomic instruction
   virtual void setTargetAtomicMetadata(CodeGenFunction &CGF,
-                                       llvm::AtomicRMWInst &RMW) const {}
+                                       llvm::Instruction &AtomicInst,
+                                       const AtomicExpr *Expr = nullptr) const {
+  }
 
   /// Interface class for filling custom fields of a block literal for OpenCL.
   class TargetOpenCLBlockHelper {

clang/lib/CodeGen/Targets/AMDGPU.cpp

@@ -9,6 +9,7 @@
 #include "ABIInfoImpl.h"
 #include "TargetInfo.h"
 #include "clang/Basic/TargetOptions.h"
+#include "llvm/Support/AMDGPUAddrSpace.h"
 
 using namespace clang;
 using namespace clang::CodeGen;
@@ -312,7 +313,8 @@ class AMDGPUTargetCodeGenInfo : public TargetCodeGenInfo {
                                          llvm::AtomicOrdering Ordering,
                                          llvm::LLVMContext &Ctx) const override;
   void setTargetAtomicMetadata(CodeGenFunction &CGF,
-                               llvm::AtomicRMWInst &RMW) const override;
+                               llvm::Instruction &AtomicInst,
+                               const AtomicExpr *Expr = nullptr) const override;
   llvm::Value *createEnqueuedBlockKernel(CodeGenFunction &CGF,
                                          llvm::Function *BlockInvokeFunc,
                                          llvm::Type *BlockTy) const override;
@@ -546,19 +548,39 @@ AMDGPUTargetCodeGenInfo::getLLVMSyncScopeID(const LangOptions &LangOpts,
 }
 
 void AMDGPUTargetCodeGenInfo::setTargetAtomicMetadata(
-    CodeGenFunction &CGF, llvm::AtomicRMWInst &RMW) const {
-  if (!CGF.getTarget().allowAMDGPUUnsafeFPAtomics())
+    CodeGenFunction &CGF, llvm::Instruction &AtomicInst,
+    const AtomicExpr *AE) const {
+  auto *RMW = dyn_cast<llvm::AtomicRMWInst>(&AtomicInst);
+  auto *CmpX = dyn_cast<llvm::AtomicCmpXchgInst>(&AtomicInst);
+
+  // OpenCL and old style HIP atomics consider atomics targeting thread private
+  // memory to be undefined.
+  //
+  // TODO: This is probably undefined for atomic load/store, but there's not
+  // much direct codegen benefit to knowing this.
+  if (((RMW && RMW->getPointerAddressSpace() == llvm::AMDGPUAS::FLAT_ADDRESS) ||
+       (CmpX &&
+        CmpX->getPointerAddressSpace() == llvm::AMDGPUAS::FLAT_ADDRESS)) &&
+      AE && AE->threadPrivateMemoryAtomicsAreUndefined()) {
+    llvm::MDBuilder MDHelper(CGF.getLLVMContext());
+    llvm::MDNode *ASRange = MDHelper.createRange(
+        llvm::APInt(32, llvm::AMDGPUAS::PRIVATE_ADDRESS),
+        llvm::APInt(32, llvm::AMDGPUAS::PRIVATE_ADDRESS + 1));
+    AtomicInst.setMetadata(llvm::LLVMContext::MD_noalias_addrspace, ASRange);
+  }
+
+  if (!RMW || !CGF.getTarget().allowAMDGPUUnsafeFPAtomics())
     return;
 
   // TODO: Introduce new, more controlled options that also work for integers,
   // and deprecate allowAMDGPUUnsafeFPAtomics.
-  llvm::AtomicRMWInst::BinOp RMWOp = RMW.getOperation();
+  llvm::AtomicRMWInst::BinOp RMWOp = RMW->getOperation();
   if (llvm::AtomicRMWInst::isFPOperation(RMWOp)) {
     llvm::MDNode *Empty = llvm::MDNode::get(CGF.getLLVMContext(), {});
-    RMW.setMetadata("amdgpu.no.fine.grained.memory", Empty);
+    RMW->setMetadata("amdgpu.no.fine.grained.memory", Empty);
 
-    if (RMWOp == llvm::AtomicRMWInst::FAdd && RMW.getType()->isFloatTy())
-      RMW.setMetadata("amdgpu.ignore.denormal.mode", Empty);
+    if (RMWOp == llvm::AtomicRMWInst::FAdd && RMW->getType()->isFloatTy())
+      RMW->setMetadata("amdgpu.ignore.denormal.mode", Empty);
   }
 }