Add storage class memory semantics bit to LLVM IR atomic translation (#3826)

Deduce storage class memory semantic from the pointer accessed by atomic
intrinsics

Inspired by llvm/llvm-project#193696,
llvm/llvm-project#195049

Author: aobolensk

lib/SPIRV/OCLToSPIRV.cpp

@@ -69,87 +69,6 @@ using namespace OCLUtil;
 
 namespace SPIRV {
 
-static unsigned getAddressSpaceFromType(const Type *Ty) {
-  assert(Ty && "Can't deduce pointer AS");
-  if (auto *TypedPtr = dyn_cast<TypedPointerType>(Ty))
-    return TypedPtr->getAddressSpace();
-  if (auto *Ptr = dyn_cast<PointerType>(Ty))
-    return Ptr->getAddressSpace();
-  llvm_unreachable("Can't deduce pointer AS");
-}
-
-// Performs an address space inference analysis.
-static unsigned getAddressSpaceFromValue(const Value *Ptr) {
-  assert(Ptr && "Can't deduce pointer AS");
-
-  SmallPtrSet<const Value *, 8> Visited;
-  SmallVector<const Value *, 8> Worklist;
-  Worklist.push_back(Ptr);
-  unsigned AS = SPIRAS_Generic;
-
-  while (!Worklist.empty()) {
-    const Value *Current = Worklist.pop_back_val();
-    if (!Visited.insert(Current).second)
-      continue;
-
-    unsigned DeducedAS = getAddressSpaceFromType(Current->getType());
-    if (DeducedAS != SPIRAS_Generic)
-      return DeducedAS;
-    AS = DeducedAS;
-
-    // Find origins of the pointer and add to the worklist.
-    if (auto *Op = dyn_cast<Operator>(Current)) {
-      switch (Op->getOpcode()) {
-      case Instruction::AddrSpaceCast:
-      case Instruction::BitCast:
-      case Instruction::GetElementPtr:
-        Worklist.push_back(Op->getOperand(0));
-        break;
-      case Instruction::Select:
-        Worklist.push_back(Op->getOperand(1));
-        Worklist.push_back(Op->getOperand(2));
-        break;
-      case Instruction::PHI: {
-        auto *Phi = cast<PHINode>(Op);
-        for (Value *Incoming : Phi->incoming_values())
-          Worklist.push_back(Incoming);
-        break;
-      }
-      default:
-        break;
-      }
-    }
-  }
-
-  return AS;
-}
-
-// Sets memory semantic mask of an atomic depending on a pointer argument
-// address space.
-static unsigned
-getAtomicPointerMemorySemanticsMemoryMask(const Value *Ptr,
-                                          const Type *RecordedType) {
-  assert((Ptr && RecordedType) &&
-         "Can't evaluate atomic builtin's memory semantic");
-  unsigned AddrSpace = getAddressSpaceFromType(RecordedType);
-  if (AddrSpace == SPIRAS_Generic)
-    AddrSpace = getAddressSpaceFromValue(Ptr);
-
-  switch (AddrSpace) {
-  case SPIRAS_Global:
-  case SPIRAS_GlobalDevice:
-  case SPIRAS_GlobalHost:
-    return MemorySemanticsCrossWorkgroupMemoryMask;
-  case SPIRAS_Local:
-    return MemorySemanticsWorkgroupMemoryMask;
-  case SPIRAS_Generic:
-    return MemorySemanticsCrossWorkgroupMemoryMask |
-           MemorySemanticsWorkgroupMemoryMask;
-  default:
-    return MemorySemanticsMaskNone;
-  }
-}
-
 static size_t getOCLCpp11AtomicMaxNumOps(StringRef Name) {
   return StringSwitch<size_t>(Name)
       .Cases({"load", "flag_test_and_set", "flag_clear"}, 3)
@@ -790,8 +709,8 @@ void OCLToSPIRVBase::transAtomicBuiltin(CallInst *CI,
 
   unsigned PtrMemSemantics = MemorySemanticsMaskNone;
   if (Mutator.arg_size() > 0)
-    PtrMemSemantics = getAtomicPointerMemorySemanticsMemoryMask(
-        Mutator.getArg(0), Mutator.getType(0));
+    PtrMemSemantics = getAtomicPointerMemorySemanticsMask(Mutator.getArg(0),
+                                                          Mutator.getType(0));
 
   if (NeedsNegate) {
     Mutator.mapArg(1, [=](Value *V) {

lib/SPIRV/OCLUtil.cpp

@@ -41,10 +41,13 @@
 #include "SPIRVFunction.h"
 #include "SPIRVInstruction.h"
 #include "SPIRVInternal.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/StringSwitch.h"
+#include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstVisitor.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/Operator.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
@@ -670,6 +673,83 @@ template <> void LLVMSPIRVAtomicRmwOpCodeMap::init() {
 
 namespace OCLUtil {
 
+static unsigned getAddressSpaceFromType(const Type *Ty) {
+  assert(Ty && "Can't deduce pointer AS");
+  if (auto *TypedPtr = dyn_cast<TypedPointerType>(Ty))
+    return TypedPtr->getAddressSpace();
+  if (auto *Ptr = dyn_cast<PointerType>(Ty))
+    return Ptr->getAddressSpace();
+  llvm_unreachable("Can't deduce pointer AS");
+}
+
+// Performs an address space inference analysis.
+static unsigned getAddressSpaceFromValue(const Value *Ptr) {
+  assert(Ptr && "Can't deduce pointer AS");
+
+  SmallPtrSet<const Value *, 8> Visited;
+  SmallVector<const Value *, 8> Worklist;
+  Worklist.push_back(Ptr);
+  unsigned AS = SPIRAS_Generic;
+
+  while (!Worklist.empty()) {
+    const Value *Current = Worklist.pop_back_val();
+    if (!Visited.insert(Current).second)
+      continue;
+
+    unsigned DeducedAS = getAddressSpaceFromType(Current->getType());
+    if (DeducedAS != SPIRAS_Generic)
+      return DeducedAS;
+
+    // Find origins of the pointer and add to the worklist.
+    if (auto *Op = dyn_cast<Operator>(Current)) {
+      switch (Op->getOpcode()) {
+      case Instruction::AddrSpaceCast:
+      case Instruction::BitCast:
+      case Instruction::GetElementPtr:
+        Worklist.push_back(Op->getOperand(0));
+        break;
+      case Instruction::Select:
+        Worklist.push_back(Op->getOperand(1));
+        Worklist.push_back(Op->getOperand(2));
+        break;
+      case Instruction::PHI: {
+        auto *Phi = cast<PHINode>(Op);
+        for (Value *Incoming : Phi->incoming_values())
+          Worklist.push_back(Incoming);
+        break;
+      }
+      default:
+        break;
+      }
+    }
+  }
+
+  return AS;
+}
+
+unsigned getAtomicPointerMemorySemanticsMask(const Value *Ptr,
+                                             const Type *RecordedType) {
+  assert((Ptr && RecordedType) &&
+         "Can't evaluate atomic builtin's memory semantic");
+  unsigned AddrSpace = getAddressSpaceFromType(RecordedType);
+  if (AddrSpace == SPIRAS_Generic)
+    AddrSpace = getAddressSpaceFromValue(Ptr);
+
+  switch (AddrSpace) {
+  case SPIRAS_Global:
+  case SPIRAS_GlobalDevice:
+  case SPIRAS_GlobalHost:
+    return MemorySemanticsCrossWorkgroupMemoryMask;
+  case SPIRAS_Local:
+    return MemorySemanticsWorkgroupMemoryMask;
+  case SPIRAS_Generic:
+    return MemorySemanticsCrossWorkgroupMemoryMask |
+           MemorySemanticsWorkgroupMemoryMask;
+  default:
+    return MemorySemanticsMaskNone;
+  }
+}
+
 AtomicWorkItemFenceLiterals getAtomicWorkItemFenceLiterals(CallInst *CI) {
   return std::make_tuple(getArgAsInt(CI, 0),
                          static_cast<OCLMemOrderKind>(getArgAsInt(CI, 1)),

lib/SPIRV/OCLUtil.h

@@ -508,6 +508,12 @@ bool isPipeOrAddressSpaceCastBI(const StringRef MangledName);
 bool isEnqueueKernelBI(const StringRef MangledName);
 bool isKernelQueryBI(const StringRef MangledName);
 
+// Returns the storage-class memory-semantics bit mask derived from the pointer
+// address space. RecordedType is checked first; if it resolves to Generic, Ptr
+// is analyzed via use-def walk.
+unsigned getAtomicPointerMemorySemanticsMask(const Value *Ptr,
+                                             const Type *RecordedType);
+
 /// Check that the type is the sampler_t
 bool isSamplerTy(Type *Ty);
 

lib/SPIRV/SPIRVRegularizeLLVM.cpp

@@ -771,8 +771,12 @@ bool SPIRVRegularizeLLVMBase::regularize() {
               llvm::toCABI(Cmpxchg->getSuccessOrdering()));
           auto FailureOrder = static_cast<OCLMemOrderKind>(
               llvm::toCABI(Cmpxchg->getFailureOrdering()));
-          Value *EqualSem = getInt32(M, OCLMemOrderMap::map(SuccessOrder));
-          Value *UnequalSem = getInt32(M, OCLMemOrderMap::map(FailureOrder));
+          unsigned SCMask =
+              getAtomicPointerMemorySemanticsMask(Ptr, Ptr->getType());
+          Value *EqualSem =
+              getInt32(M, OCLMemOrderMap::map(SuccessOrder) | SCMask);
+          Value *UnequalSem =
+              getInt32(M, OCLMemOrderMap::map(FailureOrder) | SCMask);
           Value *Val = Cmpxchg->getNewValOperand();
           Value *Comparator = Cmpxchg->getCompareOperand();
 

lib/SPIRV/SPIRVWriter.cpp

@@ -2028,8 +2028,10 @@ SPIRVValue *LLVMToSPIRVBase::transAtomicStore(StoreInst *ST,
                                               SPIRVBasicBlock *BB) {
   spv::Scope S = toSPIRVScope(ST->getContext(), ST->getSyncScopeID());
 
-  std::vector<Value *> Ops{ST->getPointerOperand(), getUInt32(M, S),
-                           getUInt32(M, transAtomicOrdering(ST->getOrdering())),
+  Value *Ptr = ST->getPointerOperand();
+  auto MemSem = transAtomicOrdering(ST->getOrdering()) |
+                getAtomicPointerMemorySemanticsMask(Ptr, Ptr->getType());
+  std::vector<Value *> Ops{Ptr, getUInt32(M, S), getUInt32(M, MemSem),
                            ST->getValueOperand()};
   std::vector<SPIRVValue *> SPIRVOps = transValue(Ops, BB);
 
@@ -2041,9 +2043,10 @@ SPIRVValue *LLVMToSPIRVBase::transAtomicLoad(LoadInst *LD,
                                              SPIRVBasicBlock *BB) {
   spv::Scope S = toSPIRVScope(LD->getContext(), LD->getSyncScopeID());
 
-  std::vector<Value *> Ops{
-      LD->getPointerOperand(), getUInt32(M, S),
-      getUInt32(M, transAtomicOrdering(LD->getOrdering()))};
+  Value *Ptr = LD->getPointerOperand();
+  auto MemSem = transAtomicOrdering(LD->getOrdering()) |
+                getAtomicPointerMemorySemanticsMask(Ptr, Ptr->getType());
+  std::vector<Value *> Ops{Ptr, getUInt32(M, S), getUInt32(M, MemSem)};
   std::vector<SPIRVValue *> SPIRVOps = transValue(Ops, BB);
 
   return mapValue(LD, BM->addInstTemplate(OpAtomicLoad, BM->getIds(SPIRVOps),
@@ -2807,9 +2810,11 @@ LLVMToSPIRVBase::transValueWithoutDecoration(Value *V, SPIRVBasicBlock *BB,
       return nullptr;
 
     AtomicOrderingCABI Ordering = llvm::toCABI(ARMW->getOrdering());
-    auto MemSem = OCLMemOrderMap::map(static_cast<OCLMemOrderKind>(Ordering));
+    Value *Ptr = ARMW->getPointerOperand();
+    auto MemSem = OCLMemOrderMap::map(static_cast<OCLMemOrderKind>(Ordering)) |
+                  getAtomicPointerMemorySemanticsMask(Ptr, Ptr->getType());
     std::vector<Value *> Operands(4);
-    Operands[0] = ARMW->getPointerOperand();
+    Operands[0] = Ptr;
     spv::Scope S = toSPIRVScope(ARMW->getContext(), ARMW->getSyncScopeID());
     Operands[1] = getUInt32(M, S);
     Operands[2] = getUInt32(M, MemSem);

test/atomicrmw.ll

@@ -8,58 +8,82 @@
 ; RUN: llvm-spirv -to-text %t.spv -o - | FileCheck %s
 
 ; CHECK: TypeInt [[Int:[0-9]+]] 32 0
-; CHECK-DAG: Constant [[Int]] [[MemSem_Relaxed:[0-9]+]] 0
-; CHECK-DAG: Constant [[Int]] [[MemSem_Acquire:[0-9]+]] 2
-; CHECK-DAG: Constant [[Int]] [[MemSem_Release:[0-9]+]] 4 {{$}}
-; CHECK-DAG: Constant [[Int]] [[MemSem_AcquireRelease:[0-9]+]] 8
-; CHECK-DAG: Constant [[Int]] [[MemSem_SequentiallyConsistent:[0-9]+]] 16
+; CHECK-DAG: Constant [[Int]] [[MemSem_CW_Relaxed:[0-9]+]] 512
+; CHECK-DAG: Constant [[Int]] [[MemSem_CW_Acquire:[0-9]+]] 514
+; CHECK-DAG: Constant [[Int]] [[MemSem_CW_Release:[0-9]+]] 516
+; CHECK-DAG: Constant [[Int]] [[MemSem_CW_AcquireRelease:[0-9]+]] 520
+; CHECK-DAG: Constant [[Int]] [[MemSem_CW_SequentiallyConsistent:[0-9]+]] 528
+; CHECK-DAG: Constant [[Int]] [[MemSem_WG_Relaxed:[0-9]+]] 256
+; CHECK-DAG: Constant [[Int]] [[MemSem_WG_AcquireRelease:[0-9]+]] 264
+; CHECK-DAG: Constant [[Int]] [[MemSem_WG_SequentiallyConsistent:[0-9]+]] 272
 ; CHECK-DAG: Constant [[Int]] [[Value:[0-9]+]] 42
 ; CHECK: TypeFloat [[Float:[0-9]+]] 32
 ; CHECK: {{(Variable|UntypedVariableKHR)}} {{[0-9]+}} [[Pointer:[0-9]+]]
 ; CHECK: {{(Variable|UntypedVariableKHR)}} {{[0-9]+}} [[FPPointer:[0-9]+]]
+; CHECK: {{(Variable|UntypedVariableKHR)}} {{[0-9]+}} [[LocalPointer:[0-9]+]]
 ; CHECK: Constant [[Float]] [[FPValue:[0-9]+]] 1109917696
 
 target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024"
 target triple = "spir64"
 
 @ui = common dso_local addrspace(1) global i32 0, align 4
 @f = common dso_local local_unnamed_addr addrspace(1) global float 0.000000e+00, align 4
+@li = common dso_local addrspace(3) global i32 0, align 4
 
 ; Function Attrs: nounwind
 define dso_local spir_func void @test_atomicrmw() local_unnamed_addr #0 {
 entry:
   %0 = atomicrmw xchg ptr addrspace(1) @ui, i32 42 acq_rel
-; CHECK: AtomicExchange [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_AcquireRelease]] [[Value]]
+; CHECK: AtomicExchange [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_CW_AcquireRelease]] [[Value]]
 
   %1 = atomicrmw xchg ptr addrspace(1) @f, float 42.000000e+00 seq_cst
-; CHECK: AtomicExchange [[Float]] {{[0-9]+}} [[FPPointer]] {{.+}} [[MemSem_SequentiallyConsistent]] [[FPValue]]
+; CHECK: AtomicExchange [[Float]] {{[0-9]+}} [[FPPointer]] {{.+}} [[MemSem_CW_SequentiallyConsistent]] [[FPValue]]
 
   %2 = atomicrmw add ptr addrspace(1) @ui, i32 42 monotonic
-; CHECK: AtomicIAdd [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_Relaxed]] [[Value]]
+; CHECK: AtomicIAdd [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_CW_Relaxed]] [[Value]]
 
   %3 = atomicrmw sub ptr addrspace(1) @ui, i32 42 acquire
-; CHECK: AtomicISub [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_Acquire]] [[Value]]
+; CHECK: AtomicISub [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_CW_Acquire]] [[Value]]
 
   %4 = atomicrmw or ptr addrspace(1) @ui, i32 42 release
-; CHECK: AtomicOr [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_Release]] [[Value]]
+; CHECK: AtomicOr [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_CW_Release]] [[Value]]
 
   %5 = atomicrmw xor ptr addrspace(1) @ui, i32 42 acq_rel
-; CHECK: AtomicXor [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_AcquireRelease]] [[Value]]
+; CHECK: AtomicXor [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_CW_AcquireRelease]] [[Value]]
 
   %6 = atomicrmw and ptr addrspace(1) @ui, i32 42 seq_cst
-; CHECK: AtomicAnd [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_SequentiallyConsistent]] [[Value]]
+; CHECK: AtomicAnd [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_CW_SequentiallyConsistent]] [[Value]]
 
   %7 = atomicrmw max ptr addrspace(1) @ui, i32 42 monotonic
-; CHECK: AtomicSMax [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_Relaxed]] [[Value]]
+; CHECK: AtomicSMax [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_CW_Relaxed]] [[Value]]
 
   %8 = atomicrmw min ptr addrspace(1) @ui, i32 42 acquire
-; CHECK: AtomicSMin [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_Acquire]] [[Value]]
+; CHECK: AtomicSMin [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_CW_Acquire]] [[Value]]
 
   %9 = atomicrmw umax ptr addrspace(1) @ui, i32 42 release
-; CHECK: AtomicUMax [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_Release]] [[Value]]
+; CHECK: AtomicUMax [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_CW_Release]] [[Value]]
 
   %10 = atomicrmw umin ptr addrspace(1) @ui, i32 42 acq_rel
-; CHECK: AtomicUMin [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_AcquireRelease]] [[Value]]
+; CHECK: AtomicUMin [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_CW_AcquireRelease]] [[Value]]
+
+  %11 = atomicrmw add ptr addrspace(3) @li, i32 42 acq_rel
+; CHECK: AtomicIAdd [[Int]] {{[0-9]+}} [[LocalPointer]] {{.+}} [[MemSem_WG_AcquireRelease]] [[Value]]
+
+  ret void
+}
+
+; Function Attrs: nounwind
+define dso_local spir_func void @test_atomic_load_store() local_unnamed_addr #0 {
+entry:
+  %0 = load atomic i32, ptr addrspace(1) @ui seq_cst, align 4
+; CHECK: AtomicLoad [[Int]] {{[0-9]+}} [[Pointer]] {{.+}} [[MemSem_CW_SequentiallyConsistent]]
+  store atomic i32 42, ptr addrspace(1) @ui monotonic, align 4
+; CHECK: AtomicStore [[Pointer]] {{.+}} [[MemSem_CW_Relaxed]] [[Value]]
+
+  %1 = load atomic i32, ptr addrspace(3) @li seq_cst, align 4
+; CHECK: AtomicLoad [[Int]] {{[0-9]+}} [[LocalPointer]] {{.+}} [[MemSem_WG_SequentiallyConsistent]]
+  store atomic i32 42, ptr addrspace(3) @li monotonic, align 4
+; CHECK: AtomicStore [[LocalPointer]] {{.+}} [[MemSem_WG_Relaxed]] [[Value]]
 
   ret void
 }

test/extensions/EXT/SPV_EXT_shader_atomic_float_/atomicrmw_fadd_double.ll

@@ -7,7 +7,7 @@
 ; CHECK-DAG: Capability AtomicFloat64AddEXT
 ; CHECK: TypeInt [[Int:[0-9]+]] 32 0
 ; CHECK-DAG: Constant [[Int]] [[Scope_CrossDevice:[0-9]+]] 0 {{$}}
-; CHECK-DAG: Constant [[Int]] [[MemSem_SequentiallyConsistent:[0-9]+]] 16
+; CHECK-DAG: Constant [[Int]] [[MemSem_SequentiallyConsistent:[0-9]+]] 528
 ; CHECK: TypeFloat [[Double:[0-9]+]] 64
 ; CHECK: Variable {{[0-9]+}} [[DoublePointer:[0-9]+]]
 ; CHECK: Constant [[Double]] [[DoubleValue:[0-9]+]] 0 1078263808

test/extensions/EXT/SPV_EXT_shader_atomic_float_/atomicrmw_fadd_float.ll

@@ -7,7 +7,7 @@
 ; CHECK-DAG: Capability AtomicFloat32AddEXT
 ; CHECK: TypeInt [[Int:[0-9]+]] 32 0
 ; CHECK-DAG: Constant [[Int]] [[Scope_CrossDevice:[0-9]+]] 0 {{$}}
-; CHECK-DAG: Constant [[Int]] [[MemSem_SequentiallyConsistent:[0-9]+]] 16
+; CHECK-DAG: Constant [[Int]] [[MemSem_SequentiallyConsistent:[0-9]+]] 528
 ; CHECK: TypeFloat [[Float:[0-9]+]] 32
 ; CHECK: Variable {{[0-9]+}} [[FPPointer:[0-9]+]]
 ; CHECK: Constant [[Float]] [[FPValue:[0-9]+]] 1109917696

test/extensions/EXT/SPV_EXT_shader_atomic_float_/atomicrmw_fadd_half.ll

@@ -8,7 +8,7 @@
 ; CHECK-DAG: Capability AtomicFloat16AddEXT
 ; CHECK: TypeInt [[TypeIntID:[0-9]+]] 32 0
 ; CHECK-DAG: Constant [[TypeIntID]] [[ScopeCrossDevice:[0-9]+]] 0 {{$}}
-; CHECK-DAG: Constant [[TypeIntID]] [[MemSem_SequentiallyConsistent:[0-9]+]] 16
+; CHECK-DAG: Constant [[TypeIntID]] [[MemSem_SequentiallyConsistent:[0-9]+]] 528
 ; CHECK: TypeFloat [[TypeFloatHalfID:[0-9]+]] 16
 ; CHECK: Variable {{[0-9]+}} [[HalfPointer:[0-9]+]]
 ; CHECK: Constant [[TypeFloatHalfID]] [[HalfValue:[0-9]+]] 20800