[Backport to 9] [NFC] Refactor logic of memory scopes, orders and flags translation

AlexeySachkov · stanleygambarin · commit f5bc1ff7ac7d · 2023-03-01T09:55:38.000-08:00
There are several OpenCL built-ins which accept cl_mem_fence_flags,
memory_scope and memory_order arguments (like barriers, fences or atomics) and
representation of that info is different between OpenCL and SPIR-V:
different values for enumerations and bitmasks, different arguments for
OpenCL built-ins and SPIR-V instructions.

To avoid code duplication, refactored all common code for translation
aforementioned entities into helper functions.

Simplified interface of getOrCreateSwithFunc helper: `llvm::Module`
argument is now infered from `InsertPoint` argument.
diff --git a/lib/SPIRV/OCL20ToSPIRV.cpp b/lib/SPIRV/OCL20ToSPIRV.cpp
@@ -823,29 +823,13 @@ void OCL20ToSPIRV::transAtomicBuiltin(CallInst *CI, OCLBuiltinTransInfo &Info) {
         const size_t ArgsCount = Args.size();
         const size_t ScopeIdx = ArgsCount - 1;
         const size_t OrderIdx = ScopeIdx - NumOrder;
-        if (auto ScopeInt = dyn_cast_or_null<ConstantInt>(Args[ScopeIdx])) {
-          Args[ScopeIdx] = mapUInt(M, ScopeInt, [](unsigned I) {
-            return map<Scope>(static_cast<OCLScopeKind>(I));
-          });
-        } else {
-          Args[ScopeIdx] =
-              getOrCreateSwitchFunc(kSPIRVName::TranslateOCLMemScope,
-                                    Args[ScopeIdx], OCLMemScopeMap::getMap(),
-                                    false /*IsReverse*/, OCLMS_device, CI, M);
-        }
+
+        Args[ScopeIdx] =
+            transOCLMemScopeIntoSPIRVScope(Args[ScopeIdx], OCLMS_device, CI);
+
         for (size_t I = 0; I < NumOrder; ++I) {
-          if (auto OrderInt =
-                  dyn_cast_or_null<ConstantInt>(Args[OrderIdx + I])) {
-            Args[OrderIdx + I] = mapUInt(M, OrderInt, [](unsigned Ord) {
-              return mapOCLMemSemanticToSPIRV(
-                  0, static_cast<OCLMemOrderKind>(Ord));
-            });
-          } else {
-            Args[OrderIdx + I] = getOrCreateSwitchFunc(
-                kSPIRVName::TranslateOCLMemOrder, Args[OrderIdx + I],
-                OCLMemOrderMap::getMap(), false /*IsReverse*/, OCLMO_seq_cst,
-                CI, M);
-          }
+          Args[OrderIdx + I] = transOCLMemOrderIntoSPIRVMemorySemantics(
+              Args[OrderIdx + I], OCLMO_seq_cst, CI);
         }
         // Order of args in SPIR-V:
         // object, scope, 1-2 order, 0-2 other args
diff --git a/lib/SPIRV/OCLUtil.cpp b/lib/SPIRV/OCLUtil.cpp
@@ -952,6 +952,105 @@ std::string getIntelSubgroupBlockDataPostfix(unsigned ElementBitSize,
 }
 } // namespace OCLUtil
 
+Value *SPIRV::transOCLMemScopeIntoSPIRVScope(Value *MemScope,
+                                             Optional<int> DefaultCase,
+                                             Instruction *InsertBefore) {
+  if (auto *C = dyn_cast<ConstantInt>(MemScope)) {
+    return ConstantInt::get(
+        C->getType(), map<Scope>(static_cast<OCLScopeKind>(C->getZExtValue())));
+  }
+
+  // If memory_scope is not a constant, then we have to insert dynamic mapping:
+  return getOrCreateSwitchFunc(kSPIRVName::TranslateOCLMemScope, MemScope,
+                               OCLMemScopeMap::getMap(), /* IsReverse */ false,
+                               DefaultCase, InsertBefore);
+}
+
+Value *SPIRV::transOCLMemOrderIntoSPIRVMemorySemantics(
+    Value *MemOrder, Optional<int> DefaultCase, Instruction *InsertBefore) {
+  if (auto *C = dyn_cast<ConstantInt>(MemOrder)) {
+    return ConstantInt::get(
+        C->getType(), mapOCLMemSemanticToSPIRV(
+                          0, static_cast<OCLMemOrderKind>(C->getZExtValue())));
+  }
+
+  return getOrCreateSwitchFunc(kSPIRVName::TranslateOCLMemOrder, MemOrder,
+                               OCLMemOrderMap::getMap(), /* IsReverse */ false,
+                               DefaultCase, InsertBefore);
+}
+
+Value *
+SPIRV::transSPIRVMemoryScopeIntoOCLMemoryScope(Value *MemScope,
+                                               Instruction *InsertBefore) {
+  if (auto *C = dyn_cast<ConstantInt>(MemScope)) {
+    return ConstantInt::get(C->getType(), rmap<OCLScopeKind>(static_cast<Scope>(
+                                              C->getZExtValue())));
+  }
+
+  if (auto *CI = dyn_cast<CallInst>(MemScope)) {
+    Function *F = CI->getCalledFunction();
+    if (F && F->getName().equals(kSPIRVName::TranslateOCLMemScope)) {
+      // In case the SPIR-V module was created from an OpenCL program by
+      // *this* SPIR-V generator, we know that the value passed to
+      // __translate_ocl_memory_scope is what we should pass to the
+      // OpenCL builtin now.
+      return CI->getArgOperand(0);
+    }
+  }
+
+  return getOrCreateSwitchFunc(kSPIRVName::TranslateSPIRVMemScope, MemScope,
+                               OCLMemScopeMap::getRMap(),
+                               /* IsReverse */ true, None, InsertBefore);
+}
+
+Value *
+SPIRV::transSPIRVMemorySemanticsIntoOCLMemoryOrder(Value *MemorySemantics,
+                                                   Instruction *InsertBefore) {
+  if (auto *C = dyn_cast<ConstantInt>(MemorySemantics)) {
+    return ConstantInt::get(C->getType(),
+                            mapSPIRVMemSemanticToOCL(C->getZExtValue()).second);
+  }
+
+  if (auto *CI = dyn_cast<CallInst>(MemorySemantics)) {
+    Function *F = CI->getCalledFunction();
+    if (F && F->getName().equals(kSPIRVName::TranslateOCLMemOrder)) {
+      // In case the SPIR-V module was created from an OpenCL program by
+      // *this* SPIR-V generator, we know that the value passed to
+      // __translate_ocl_memory_order is what we should pass to the
+      // OpenCL builtin now.
+      return CI->getArgOperand(0);
+    }
+  }
+
+  // SPIR-V MemorySemantics contains both OCL mem_fence_flags and mem_order and
+  // therefore, we need to apply mask
+  int Mask = MemorySemanticsMaskNone | MemorySemanticsAcquireMask |
+             MemorySemanticsReleaseMask | MemorySemanticsAcquireReleaseMask |
+             MemorySemanticsSequentiallyConsistentMask;
+  return getOrCreateSwitchFunc(kSPIRVName::TranslateSPIRVMemOrder,
+                               MemorySemantics, OCLMemOrderMap::getRMap(),
+                               /* IsReverse */ true, None, InsertBefore, Mask);
+}
+
+Value *SPIRV::transSPIRVMemorySemanticsIntoOCLMemFenceFlags(
+    Value *MemorySemantics, Instruction *InsertBefore) {
+  if (auto *C = dyn_cast<ConstantInt>(MemorySemantics)) {
+    return ConstantInt::get(C->getType(),
+                            mapSPIRVMemSemanticToOCL(C->getZExtValue()).first);
+  }
+
+  // TODO: any possible optimizations?
+  // SPIR-V MemorySemantics contains both OCL mem_fence_flags and mem_order and
+  // therefore, we need to apply mask
+  int Mask = MemorySemanticsWorkgroupMemoryMask |
+             MemorySemanticsCrossWorkgroupMemoryMask |
+             MemorySemanticsImageMemoryMask;
+  return getOrCreateSwitchFunc(kSPIRVName::TranslateSPIRVMemFence,
+                               MemorySemantics,
+                               OCLMemFenceExtendedMap::getRMap(),
+                               /* IsReverse */ true, None, InsertBefore, Mask);
+}
+
 void llvm::mangleOpenClBuiltin(const std::string &UniqName,
                                ArrayRef<Type *> ArgTypes,
                                std::string &MangledName) {
diff --git a/lib/SPIRV/OCLUtil.h b/lib/SPIRV/OCLUtil.h
@@ -491,12 +491,13 @@ Instruction *
 getOrCreateSwitchFunc(StringRef MapName, Value *V,
                       const SPIRVMap<KeyTy, ValTy, Identifier> &Map,
                       bool IsReverse, Optional<int> DefaultCase,
-                      Instruction *InsertPoint, Module *M, int KeyMask = 0) {
+                      Instruction *InsertPoint, int KeyMask = 0) {
   static_assert(std::is_convertible<KeyTy, int>::value &&
                     std::is_convertible<ValTy, int>::value,
                 "Can map only integer values");
   Type *Ty = V->getType();
   assert(Ty && Ty->isIntegerTy() && "Can't map non-integer types");
+  Module *M = InsertPoint->getModule();
   Function *F = getOrCreateFunction(M, Ty, Ty, MapName);
   if (!F->empty()) // The switch function already exists. just call it.
     return addCallInst(M, MapName, Ty, V, nullptr, InsertPoint);
@@ -538,6 +539,85 @@ getOrCreateSwitchFunc(StringRef MapName, Value *V,
   return addCallInst(M, MapName, Ty, V, nullptr, InsertPoint);
 }
 
+/// Performs conversion from OpenCL memory_scope into SPIR-V Scope.
+///
+/// Supports both constant and non-constant values. To handle the latter case,
+/// function with switch..case statement will be inserted into module which
+/// \arg InsertBefore belongs to (in order to perform mapping at runtime)
+///
+/// \param [in] MemScope memory_scope value which needs to be translated
+/// \param [in] DefaultCase default value for switch..case construct if
+///             dynamic mapping is used
+/// \param [in] InsertBefore insertion point for call into conversion function
+///             which is generated if \arg MemScope is not a constant
+/// \returns \c Value corresponding to SPIR-V Scope equivalent to OpenCL
+///          memory_scope passed in \arg MemScope
+Value *transOCLMemScopeIntoSPIRVScope(Value *MemScope,
+                                      Optional<int> DefaultCase,
+                                      Instruction *InsertBefore);
+
+/// Performs conversion from OpenCL memory_order into SPIR-V Memory Semantics.
+///
+/// Supports both constant and non-constant values. To handle the latter case,
+/// function with switch..case statement will be inserted into module which
+/// \arg InsertBefore belongs to (in order to perform mapping at runtime)
+///
+/// \param [in] MemOrder memory_scope value which needs to be translated
+/// \param [in] DefaultCase default value for switch..case construct if
+///             dynamic mapping is used
+/// \param [in] InsertBefore insertion point for call into conversion function
+///             which is generated if \arg MemOrder is not a constant
+/// \returns \c Value corresponding to SPIR-V Memory Semantics equivalent to
+///          OpenCL memory_order passed in \arg MemOrder
+Value *transOCLMemOrderIntoSPIRVMemorySemantics(Value *MemOrder,
+                                                Optional<int> DefaultCase,
+                                                Instruction *InsertBefore);
+
+/// Performs conversion from SPIR-V Scope into OpenCL memory_scope.
+///
+/// Supports both constant and non-constant values. To handle the latter case,
+/// function with switch..case statement will be inserted into module which
+/// \arg InsertBefore belongs to (in order to perform mapping at runtime)
+///
+/// \param [in] MemScope Scope value which needs to be translated
+/// \param [in] InsertBefore insertion point for call into conversion function
+///             which is generated if \arg MemScope is not a constant
+/// \returns \c Value corresponding to  OpenCL memory_scope equivalent to SPIR-V
+///          Scope passed in \arg MemScope
+Value *transSPIRVMemoryScopeIntoOCLMemoryScope(Value *MemScope,
+                                               Instruction *InsertBefore);
+
+/// Performs conversion from SPIR-V Memory Semantics into OpenCL memory_order.
+///
+/// Supports both constant and non-constant values. To handle the latter case,
+/// function with switch..case statement will be inserted into module which
+/// \arg InsertBefore belongs to (in order to perform mapping at runtime)
+///
+/// \param [in] MemorySemantics Memory Semantics value which needs to be
+///             translated
+/// \param [in] InsertBefore insertion point for call into conversion function
+///             which is generated if \arg MemorySemantics is not a constant
+/// \returns \c Value corresponding to  OpenCL memory_order equivalent to SPIR-V
+///          Memory Semantics passed in \arg MemorySemantics
+Value *transSPIRVMemorySemanticsIntoOCLMemoryOrder(Value *MemorySemantics,
+                                                   Instruction *InsertBefore);
+
+/// Performs conversion from SPIR-V Memory Semantics into OpenCL
+/// mem_fence_flags.
+///
+/// Supports both constant and non-constant values. To handle the latter case,
+/// function with switch..case statement will be inserted into module which
+/// \arg InsertBefore belongs to (in order to perform mapping at runtime)
+///
+/// \param [in] MemorySemantics Memory Semantics value which needs to be
+///             translated
+/// \param [in] InsertBefore insertion point for call into conversion function
+///             which is generated if \arg MemorySemantics is not a constant
+/// \returns \c Value corresponding to  OpenCL mem_fence_flags equivalent to
+///          SPIR-V Memory Semantics passed in \arg MemorySemantics
+Value *transSPIRVMemorySemanticsIntoOCLMemFenceFlags(Value *MemorySemantics,
+                                                     Instruction *InsertBefore);
+
 template <> inline void SPIRVMap<std::string, SPIRVGroupOperationKind>::init() {
   add("reduce", GroupOperationReduce);
   add("scan_inclusive", GroupOperationInclusiveScan);
diff --git a/lib/SPIRV/SPIRVToOCL12.cpp b/lib/SPIRV/SPIRVToOCL12.cpp
@@ -131,26 +131,9 @@ void SPIRVToOCL12::visitCallSPIRVMemoryBarrier(CallInst *CI) {
   mutateCallInstOCL(
       M, CI,
       [=](CallInst *, std::vector<Value *> &Args) {
-        if (auto Arg = dyn_cast<ConstantInt>(Args[1])) {
-          auto Sema = mapSPIRVMemSemanticToOCL(Arg->getZExtValue());
-          Args.resize(1);
-          Args[0] = getInt32(M, Sema.first);
-        } else {
-          CallInst *TransCall = dyn_cast<CallInst>(Args[1]);
-          Function *F = TransCall ? TransCall->getCalledFunction() : nullptr;
-          if (F && F->getName().equals(kSPIRVName::TranslateOCLMemScope)) {
-            Args[0] = TransCall->getArgOperand(0);
-          } else {
-            int ClMemFenceMask = MemorySemanticsWorkgroupMemoryMask |
-                                 MemorySemanticsCrossWorkgroupMemoryMask |
-                                 MemorySemanticsImageMemoryMask;
-            Args[0] = getOrCreateSwitchFunc(
-                kSPIRVName::TranslateSPIRVMemFence, Args[1],
-                OCLMemFenceExtendedMap::getRMap(), true /*IsReverse*/, None, CI,
-                M, ClMemFenceMask);
-          }
-          Args.resize(1);
-        }
+        Value *MemFenceFlags =
+            transSPIRVMemorySemanticsIntoOCLMemFenceFlags(Args[1], CI);
+        Args.assign(1, MemFenceFlags);
         return kOCLBuiltinName::MemFence;
       },
       &Attrs);
@@ -163,26 +146,9 @@ void SPIRVToOCL12::visitCallSPIRVControlBarrier(CallInst *CI) {
   mutateCallInstOCL(
       M, CI,
       [=](CallInst *, std::vector<Value *> &Args) {
-        if (auto Arg = dyn_cast<ConstantInt>(Args[2])) {
-          auto Sema = mapSPIRVMemSemanticToOCL(Arg->getZExtValue());
-          Args.resize(1);
-          Args[0] = getInt32(M, Sema.first);
-        } else {
-          CallInst *TransCall = dyn_cast<CallInst>(Args[2]);
-          Function *F = TransCall ? TransCall->getCalledFunction() : nullptr;
-          if (F && F->getName().equals(kSPIRVName::TranslateOCLMemScope)) {
-            Args[0] = TransCall->getArgOperand(0);
-          } else {
-            int ClMemFenceMask = MemorySemanticsWorkgroupMemoryMask |
-                                 MemorySemanticsCrossWorkgroupMemoryMask |
-                                 MemorySemanticsImageMemoryMask;
-            Args[0] = getOrCreateSwitchFunc(
-                kSPIRVName::TranslateSPIRVMemFence, Args[2],
-                OCLMemFenceExtendedMap::getRMap(), true /*IsReverse*/, None, CI,
-                M, ClMemFenceMask);
-          }
-          Args.resize(1);
-        }
+        auto *MemFenceFlags =
+            transSPIRVMemorySemanticsIntoOCLMemFenceFlags(Args[2], CI);
+        Args.assign(1, MemFenceFlags);
         return kOCLBuiltinName::Barrier;
       },
       &Attrs);
diff --git a/lib/SPIRV/SPIRVToOCL20.cpp b/lib/SPIRV/SPIRVToOCL20.cpp
@@ -37,6 +37,7 @@
 //===----------------------------------------------------------------------===//
 #define DEBUG_TYPE "spvtocl20"
 
+#include "OCLUtil.h"
 #include "SPIRVToOCL.h"
 #include "llvm/IR/Verifier.h"
 
@@ -136,11 +137,15 @@ void SPIRVToOCL20::visitCallSPIRVControlBarrier(CallInst *CI) {
           return cast<ConstantInt>(Args[I])->getZExtValue();
         };
         auto ExecScope = static_cast<Scope>(GetArg(0));
-        auto MScope = static_cast<Scope>(GetArg(1));
-        auto Sema = mapSPIRVMemSemanticToOCL(GetArg(2));
+        Value *MemScope =
+            getInt32(M, rmap<OCLScopeKind>(static_cast<Scope>(GetArg(1))));
+        Value *MemFenceFlags =
+            SPIRV::transSPIRVMemorySemanticsIntoOCLMemFenceFlags(Args[2], CI);
+
         Args.resize(2);
-        Args[0] = getInt32(M, Sema.first);
-        Args[1] = getInt32(M, rmap<OCLScopeKind>(MScope));
+        Args[0] = MemFenceFlags;
+        Args[1] = MemScope;
+
         return (ExecScope == ScopeWorkgroup) ? kOCLBuiltinName::WorkGroupBarrier
                                              : kOCLBuiltinName::SubGroupBarrier;
       },
@@ -222,46 +227,13 @@ CallInst *SPIRVToOCL20::mutateCommonAtomicArguments(CallInst *CI, Op OC) {
         auto NumOrder = getSPIRVAtomicBuiltinNumMemoryOrderArgs(OC);
         auto ScopeIdx = Ptr + 1;
         auto OrderIdx = Ptr + 2;
-        if (auto *ScopeInt = dyn_cast_or_null<ConstantInt>(Args[ScopeIdx])) {
-          Args[ScopeIdx] = mapUInt(M, ScopeInt, [](unsigned I) {
-            return rmap<OCLScopeKind>(static_cast<Scope>(I));
-          });
-        } else {
-          CallInst *TransCall = dyn_cast<CallInst>(Args[ScopeIdx]);
-          Function *F = TransCall ? TransCall->getCalledFunction() : nullptr;
-          if (F && F->getName().equals(kSPIRVName::TranslateOCLMemScope)) {
-            // In case the SPIR-V module was created from an OpenCL program by
-            // *this* SPIR-V generator, we know that the value passed to
-            // __translate_ocl_memory_scope is what we should pass to the OpenCL
-            // builtin now.
-            Args[ScopeIdx] = TransCall->getArgOperand(0);
-          } else {
-            Args[ScopeIdx] = getOrCreateSwitchFunc(
-                kSPIRVName::TranslateSPIRVMemScope, Args[ScopeIdx],
-                OCLMemScopeMap::getRMap(), true /*IsReverse*/, None, CI, M);
-          }
-        }
+
+        Args[ScopeIdx] =
+            SPIRV::transSPIRVMemoryScopeIntoOCLMemoryScope(Args[ScopeIdx], CI);
         for (size_t I = 0; I < NumOrder; ++I) {
-          if (auto OrderInt =
-                  dyn_cast_or_null<ConstantInt>(Args[OrderIdx + I])) {
-            Args[OrderIdx + I] = mapUInt(M, OrderInt, [](unsigned Ord) {
-              return mapSPIRVMemOrderToOCL(Ord);
-            });
-          } else {
-            CallInst *TransCall = dyn_cast<CallInst>(Args[OrderIdx + I]);
-            Function *F = TransCall ? TransCall->getCalledFunction() : nullptr;
-            if (F && F->getName().equals(kSPIRVName::TranslateOCLMemOrder)) {
-              // In case the SPIR-V module was created from an OpenCL program by
-              // *this* SPIR-V generator, we know that the value passed to
-              // __translate_ocl_memory_order is what we should pass to the
-              // OpenCL builtin now.
-              Args[OrderIdx + I] = TransCall->getArgOperand(0);
-            } else {
-              Args[OrderIdx + I] = getOrCreateSwitchFunc(
-                  kSPIRVName::TranslateSPIRVMemOrder, Args[OrderIdx + I],
-                  OCLMemOrderMap::getRMap(), true /*IsReverse*/, None, CI, M);
-            }
-          }
+          Args[OrderIdx + I] =
+              SPIRV::transSPIRVMemorySemanticsIntoOCLMemoryOrder(
+                  Args[OrderIdx + I], CI);
         }
         std::swap(Args[ScopeIdx], Args.back());
         return Name;
diff --git a/test/atomic_explicit_arguments.spt b/test/atomic_explicit_arguments.spt
@@ -67,7 +67,8 @@
 
 ; CHECK: define private spir_func i32 @__translate_spirv_memory_order(i32 %key) {
 ; CHECK: entry:
-; CHECK:   switch i32 %key, label %default [
+; CHECK:   %key.masked = and i32 30, %key
+; CHECK:   switch i32 %key.masked, label %default [
 ; CHECK:     i32 0, label %case.0
 ; CHECK:     i32 2, label %case.2
 ; CHECK:     i32 4, label %case.4
