Translate LLVM's cmpxchg instruction to SPIR-V

Signed-off-by: Alexey Sotkin <[email protected]>

Author: AlexeySotkin

llvm-spirv/lib/SPIRV/SPIRVRegularizeLLVM.cpp

@@ -49,6 +49,7 @@
 #include "llvm/Support/Debug.h"
 
 #include <set>
+#include <vector>
 
 using namespace llvm;
 using namespace SPIRV;
@@ -112,17 +113,18 @@ bool SPIRVRegularizeLLVM::regularize() {
       continue;
     }
 
-    for (auto BI = F->begin(), BE = F->end(); BI != BE; ++BI) {
-      for (auto II = BI->begin(), IE = BI->end(); II != IE; ++II) {
-        if (auto Call = dyn_cast<CallInst>(II)) {
+    std::vector<Instruction *> ToErase;
+    for (BasicBlock &BB : *F) {
+      for (Instruction &II : BB) {
+        if (auto Call = dyn_cast<CallInst>(&II)) {
           Call->setTailCall(false);
           Function *CF = Call->getCalledFunction();
           if (CF && CF->isIntrinsic())
             removeFnAttr(Call, Attribute::NoUnwind);
         }
 
         // Remove optimization info not supported by SPIRV
-        if (auto BO = dyn_cast<BinaryOperator>(II)) {
+        if (auto BO = dyn_cast<BinaryOperator>(&II)) {
           if (isa<PossiblyExactOperator>(BO) && BO->isExact())
             BO->setIsExact(false);
         }
@@ -133,12 +135,68 @@ bool SPIRVRegularizeLLVM::regularize() {
             "range",
         };
         for (auto &MDName : MDs) {
-          if (II->getMetadata(MDName)) {
-            II->setMetadata(MDName, nullptr);
+          if (II.getMetadata(MDName)) {
+            II.setMetadata(MDName, nullptr);
           }
         }
+        if (auto Cmpxchg = dyn_cast<AtomicCmpXchgInst>(&II)) {
+          Value *Ptr = Cmpxchg->getPointerOperand();
+          // To get memory scope argument we might use Cmpxchg->getSyncScopeID()
+          // but LLVM's cmpxchg instruction is not aware of OpenCL(or SPIR-V)
+          // memory scope enumeration. And assuming the produced SPIR-V module
+          // will be consumed in an OpenCL environment, we can use the same
+          // memory scope as OpenCL atomic functions that do not have
+          // memory_scope argument, i.e. memory_scope_device. See the OpenCL C
+          // specification p6.13.11. Atomic Functions
+          Value *MemoryScope = getInt32(M, spv::ScopeDevice);
+          auto SuccessOrder = static_cast<OCLMemOrderKind>(
+              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));
+          Value *Val = Cmpxchg->getNewValOperand();
+          Value *Comparator = Cmpxchg->getCompareOperand();
+
+          llvm::Value *Args[] = {Ptr,        MemoryScope, EqualSem,
+                                 UnequalSem, Val,         Comparator};
+          auto *Res = addCallInstSPIRV(M, "__spirv_AtomicCompareExchange",
+                                       Cmpxchg->getCompareOperand()->getType(),
+                                       Args, nullptr, &II, "cmpxchg.res");
+          // cmpxchg LLVM instruction returns a pair: the original value and
+          // a flag indicating success (true) or failure (false).
+          // OpAtomicCompareExchange SPIR-V instruction returns only the
+          // original value. So we replace all uses of the original value
+          // extracted from the pair with the result of OpAtomicCompareExchange
+          // instruction. And we replace all uses of the flag with result of an
+          // OpIEqual instruction. The OpIEqual instruction returns true if the
+          // original value equals to the comparator which matches with
+          // semantics of cmpxchg.
+          for (User *U : Cmpxchg->users()) {
+            if (auto *Extract = dyn_cast<ExtractValueInst>(U)) {
+              if (Extract->getIndices()[0] == 0) {
+                Extract->replaceAllUsesWith(Res);
+              } else if (Extract->getIndices()[0] == 1) {
+                auto *Cmp = new ICmpInst(Extract, CmpInst::ICMP_EQ, Res,
+                                         Comparator, "cmpxchg.success");
+                Extract->replaceAllUsesWith(Cmp);
+              } else {
+                llvm_unreachable("Unxpected cmpxchg pattern");
+              }
+              assert(Extract->user_empty());
+              Extract->dropAllReferences();
+              ToErase.push_back(Extract);
+            }
+          }
+          if (Cmpxchg->user_empty())
+            ToErase.push_back(Cmpxchg);
+        }
       }
     }
+    for (Instruction *V : ToErase) {
+      assert(V->user_empty());
+      V->eraseFromParent();
+    }
   }
 
   std::string Err;

llvm-spirv/test/AtomicCompareExchange.ll

@@ -0,0 +1,50 @@
+; RUN: llvm-as < %s -o %t.bc
+; RUN: llvm-spirv %t.bc -o %t.spv
+; RUN: llvm-spirv -to-text %t.spv -o - | FileCheck %s --check-prefix=CHECK-SPIRV
+; RUN: spirv-val %t.spv
+
+; CHECK-SPIRV: TypeInt [[Int:[0-9]+]] 32 0
+; CHECK-SPIRV: Constant [[Int]] [[MemScope_Device:[0-9]+]] 1
+; CHECK-SPIRV: Constant [[Int]] [[MemSemEqual_SeqCst:[0-9]+]] 16
+; CHECK-SPIRV: Constant [[Int]] [[MemSemUnequal_Acquire:[0-9]+]] 2
+
+; CHECK-SPIRV: FunctionParameter {{[0-9]+}} [[Pointer:[0-9]+]]
+; CHECK-SPIRV: FunctionParameter {{[0-9]+}} [[Value_ptr:[0-9]+]]
+; CHECK-SPIRV: FunctionParameter {{[0-9]+}} [[Comparator:[0-9]+]]
+
+; CHECK-SPIRV: Load [[Int]] [[Value:[0-9]+]] [[Value_ptr]]
+; CHECK-SPIRV: AtomicCompareExchange [[Int]] [[Res:[0-9]+]] [[Pointer]] [[MemScope_Device]]
+; CHECK-SPIRV-SAME:                  [[MemSemEqual_SeqCst]] [[MemSemUnequal_Acquire]] [[Value]] [[Comparator]]
+; CHECK-SPIRV: IEqual {{[0-9]+}} [[Success:[0-9]+]] [[Res]] [[Comparator]]
+; CHECK-SPIRV: BranchConditional [[Success]]
+
+; CHECK-SPIRV: Store [[Value_ptr]] [[Res]]
+
+target datalayout = "e-p:32:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024"
+target triple = "spir"
+
+; Function Attrs: nounwind
+define dso_local spir_func void @test(i32* %ptr, i32* %value_ptr, i32 %comparator) local_unnamed_addr #0 {
+entry:
+  %0 = load i32, i32* %value_ptr, align 4
+  %1 = cmpxchg i32* %ptr, i32 %comparator, i32 %0 seq_cst acquire
+  %2 = extractvalue { i32, i1 } %1, 1
+  br i1 %2, label %cmpxchg.continue, label %cmpxchg.store_expected
+
+cmpxchg.store_expected:                           ; preds = %entry
+  %3 = extractvalue { i32, i1 } %1, 0
+  store i32 %3, i32* %value_ptr, align 4
+  br label %cmpxchg.continue
+
+cmpxchg.continue:                                 ; preds = %cmpxchg.store_expected, %entry
+  ret void
+}
+
+attributes #0 = { nounwind "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "frame-pointer"="all" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="false" "stack-protector-buffer-size"="8" "unsafe-fp-math"="false" "use-soft-float"="false" }
+attributes #1 = { nounwind }
+
+!llvm.module.flags = !{!0}
+!llvm.ident = !{!1}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{!"clang version 11.0.0 (https://github.com/llvm/llvm-project.git cfebd7774229885e7ec88ae9ef1f4ae819cce1d2)"}