Fix SPV-IR for OpBuildNDRange (#1145)

It is translated to a function with unmangled name
__spirv_BuildNDRange_{1|2|3}D with struct return parameter and array
arguments, since translator only translates it properly to SPIR-V with
this signature. _ND postfix is requred because array arguments are mangled
in the same way, so if there was no postfix, translator would produce
functions with same name for different dimensions.

Author: Fznamznon

docs/SPIRVRepresentationInLLVM.rst

@@ -150,8 +150,17 @@ The unmangled names of SPIR-V builtin GenericCastToPtrExplicit function follow t
 .. code-block:: c
 
   __spirv_GenericCastToPtrExplicit_To{Global|Local|Private}
-  
-SPIR-V 1.1 Builtin CreatePipeFromPipeStorage Function Name 
+
+SPIR-V Builtin BuildNDRange Function Name
+----------------------------------------
+
+The unmangled names of SPIR-V builtin BuildNDRange functions follow the convention:
+
+.. code-block:: c
+
+  __spirv_{BuildNDRange}_{1|2|3}D
+
+SPIR-V 1.1 Builtin CreatePipeFromPipeStorage Function Name
 ----------------------------------------
 
 The unmangled names of SPIR-V builtin CreatePipeFromPipeStorage function follow the convention:

lib/SPIRV/SPIRVInternal.h

@@ -1037,6 +1037,23 @@ bool lowerBuiltinVariableToCall(GlobalVariable *GV,
 // Transform all builtin variables into calls
 bool lowerBuiltinVariablesToCalls(Module *M);
 
+/// \brief Post-process OpenCL or SPIRV builtin function returning struct type.
+///
+/// Some builtin functions are translated to SPIR-V instructions with
+/// struct type result, e.g. NDRange creation functions. Such functions
+/// need to be post-processed to return the struct through sret argument.
+bool postProcessBuiltinReturningStruct(Function *F);
+
+/// \brief Post-process OpenCL or SPIRV builtin function having array argument.
+///
+/// These functions are translated to functions with array type argument
+/// first, then post-processed to have pointer arguments.
+bool postProcessBuiltinWithArrayArguments(Function *F, StringRef DemangledName);
+
+bool postProcessBuiltinsReturningStruct(Module *M, bool IsCpp = false);
+
+bool postProcessBuiltinsWithArrayArguments(Module *M, bool IsCpp = false);
+
 } // namespace SPIRV
 
 #endif // SPIRV_SPIRVINTERNAL_H

lib/SPIRV/SPIRVReader.cpp

@@ -562,41 +562,6 @@ bool SPIRVToLLVM::isSPIRVCmpInstTransToLLVMInst(SPIRVInstruction *BI) const {
   return isCmpOpCode(OC) && !(OC >= OpLessOrGreater && OC <= OpUnordered);
 }
 
-// TODO: Instead of direct translation to OCL we should always produce SPIR-V
-// friendly IR and apply lowering later if needed
-bool SPIRVToLLVM::isDirectlyTranslatedToOCL(Op OpCode) const {
-  if (isSubgroupAvcINTELInstructionOpCode(OpCode))
-    return false;
-  if (isIntelSubgroupOpCode(OpCode))
-    return false;
-  if (OpCode == OpImageSampleExplicitLod || OpCode == OpSampledImage)
-    return false;
-  if (OpCode == OpImageWrite || OpCode == OpImageRead ||
-      OpCode == OpImageQueryOrder || OpCode == OpImageQueryFormat ||
-      OpCode == OpImageQueryLevels)
-    return false;
-  if (OpCode == OpGenericCastToPtrExplicit)
-    return false;
-  if (isEventOpCode(OpCode))
-    return false;
-  if (OpBitFieldInsert <= OpCode && OpCode <= OpBitReverse)
-    return false;
-  if (OpCode == OpEnqueueMarker || OpCode == OpGetDefaultQueue)
-    return false;
-  if (OCLSPIRVBuiltinMap::rfind(OpCode, nullptr)) {
-    // Not every spirv opcode which is placed in OCLSPIRVBuiltinMap is
-    // translated directly to OCL builtin. Some of them are translated
-    // to LLVM representation without any modifications (SPIRV format of
-    // instruction is represented in LLVM) and then its translated to
-    // clang-consistent format in SPIRVToOCL pass.
-    return !(isAtomicOpCode(OpCode) || isGroupOpCode(OpCode) ||
-             isGroupNonUniformOpcode(OpCode) || isPipeOpCode(OpCode) ||
-             isMediaBlockINTELOpcode(OpCode) || OpCode == OpGroupAsyncCopy ||
-             OpCode == OpGroupWaitEvents);
-  }
-  return false;
-}
-
 void SPIRVToLLVM::setName(llvm::Value *V, SPIRVValue *BV) {
   auto Name = BV->getName();
   if (!Name.empty() && (!V->hasName() || Name != V->getName()))
@@ -1117,102 +1082,6 @@ Value *SPIRVToLLVM::transCmpInst(SPIRVValue *BV, BasicBlock *BB, Function *F) {
   return Inst;
 }
 
-bool SPIRVToLLVM::postProcessOCL() {
-  StringRef DemangledName;
-  SPIRVWord SrcLangVer = 0;
-  BM->getSourceLanguage(&SrcLangVer);
-  bool IsCpp = SrcLangVer == kOCLVer::CL21;
-  for (auto I = M->begin(), E = M->end(); I != E;) {
-    auto F = I++;
-    if (F->hasName() && F->isDeclaration()) {
-      LLVM_DEBUG(dbgs() << "[postProcessOCL sret] " << *F << '\n');
-      if (F->getReturnType()->isStructTy() &&
-          oclIsBuiltin(F->getName(), DemangledName, IsCpp)) {
-        if (!postProcessOCLBuiltinReturnStruct(&(*F)))
-          return false;
-      }
-    }
-  }
-  for (auto I = M->begin(), E = M->end(); I != E;) {
-    auto F = I++;
-    if (F->hasName() && F->isDeclaration()) {
-      LLVM_DEBUG(dbgs() << "[postProcessOCL array arg] " << *F << '\n');
-      if (hasArrayArg(&(*F)) &&
-          oclIsBuiltin(F->getName(), DemangledName, IsCpp))
-        if (!postProcessOCLBuiltinWithArrayArguments(&(*F), DemangledName))
-          return false;
-    }
-  }
-  return true;
-}
-
-bool SPIRVToLLVM::postProcessOCLBuiltinReturnStruct(Function *F) {
-  std::string Name = F->getName().str();
-  F->setName(Name + ".old");
-  for (auto I = F->user_begin(), E = F->user_end(); I != E;) {
-    if (auto CI = dyn_cast<CallInst>(*I++)) {
-      auto ST = dyn_cast<StoreInst>(*(CI->user_begin()));
-      assert(ST);
-      std::vector<Type *> ArgTys;
-      getFunctionTypeParameterTypes(F->getFunctionType(), ArgTys);
-      ArgTys.insert(ArgTys.begin(),
-                    PointerType::get(F->getReturnType(), SPIRAS_Private));
-      auto NewF =
-          getOrCreateFunction(M, Type::getVoidTy(*Context), ArgTys, Name);
-      NewF->setCallingConv(F->getCallingConv());
-      auto Args = getArguments(CI);
-      Args.insert(Args.begin(), ST->getPointerOperand());
-      auto NewCI = CallInst::Create(NewF, Args, CI->getName(), CI);
-      NewCI->setCallingConv(CI->getCallingConv());
-      ST->eraseFromParent();
-      CI->eraseFromParent();
-    }
-  }
-  F->eraseFromParent();
-  return true;
-}
-
-bool SPIRVToLLVM::postProcessOCLBuiltinWithArrayArguments(
-    Function *F, StringRef DemangledName) {
-  LLVM_DEBUG(dbgs() << "[postProcessOCLBuiltinWithArrayArguments] " << *F
-                    << '\n');
-  auto Attrs = F->getAttributes();
-  auto Name = F->getName();
-  mutateFunction(
-      F,
-      [=](CallInst *CI, std::vector<Value *> &Args) {
-        auto FBegin =
-            CI->getParent()->getParent()->begin()->getFirstInsertionPt();
-        for (auto &I : Args) {
-          auto T = I->getType();
-          if (!T->isArrayTy())
-            continue;
-          auto Alloca = new AllocaInst(T, 0, "", &(*FBegin));
-          new StoreInst(I, Alloca, false, CI);
-          auto Zero =
-              ConstantInt::getNullValue(Type::getInt32Ty(T->getContext()));
-          Value *Index[] = {Zero, Zero};
-          I = GetElementPtrInst::CreateInBounds(T, Alloca, Index, "", CI);
-        }
-        return Name.str();
-      },
-      nullptr, &Attrs);
-  return true;
-}
-
-CallInst *SPIRVToLLVM::postProcessOCLBuildNDRange(SPIRVInstruction *BI,
-                                                  CallInst *CI,
-                                                  const std::string &FuncName) {
-  assert(CI->getNumArgOperands() == 3);
-  auto GWS = CI->getArgOperand(0);
-  auto LWS = CI->getArgOperand(1);
-  auto GWO = CI->getArgOperand(2);
-  CI->setArgOperand(0, GWO);
-  CI->setArgOperand(1, GWS);
-  CI->setArgOperand(2, LWS);
-  return CI;
-}
-
 Type *SPIRVToLLVM::mapType(SPIRVType *BT, Type *T) {
   SPIRVDBG(dbgs() << *T << '\n';)
   TypeMap[BT] = T;
@@ -2569,14 +2438,17 @@ Value *SPIRVToLLVM::transValueWithoutDecoration(SPIRVValue *BV, Function *F,
 
   default: {
     auto OC = BV->getOpCode();
-    if (isSPIRVCmpInstTransToLLVMInst(static_cast<SPIRVInstruction *>(BV))) {
+    if (isSPIRVCmpInstTransToLLVMInst(static_cast<SPIRVInstruction *>(BV)))
       return mapValue(BV, transCmpInst(BV, BB, F));
-    } else if (isDirectlyTranslatedToOCL(OC)) {
-      return mapValue(
-          BV, transOCLBuiltinFromInst(static_cast<SPIRVInstruction *>(BV), BB));
-    } else if (isBinaryShiftLogicalBitwiseOpCode(OC) || isLogicalOpCode(OC)) {
+
+    if (OCLSPIRVBuiltinMap::rfind(OC, nullptr))
+      return mapValue(BV, transSPIRVBuiltinFromInst(
+                              static_cast<SPIRVInstruction *>(BV), BB));
+
+    if (isBinaryShiftLogicalBitwiseOpCode(OC) || isLogicalOpCode(OC))
       return mapValue(BV, transShiftLogicalBitwiseInst(BV, BB, F));
-    } else if (isCvtOpCode(OC) && OC != OpGenericCastToPtrExplicit) {
+
+    if (isCvtOpCode(OC) && OC != OpGenericCastToPtrExplicit) {
       auto BI = static_cast<SPIRVInstruction *>(BV);
       Value *Inst = nullptr;
       if (BI->hasFPRoundingMode() || BI->isSaturatedConversion())
@@ -2929,8 +2801,6 @@ SPIRVToLLVM::transOCLBuiltinPostproc(SPIRVInstruction *BI, CallInst *CI,
   }
   if (OC == OpGenericPtrMemSemantics)
     return BinaryOperator::CreateShl(CI, getInt32(M, 8), "", BB);
-  if (OC == OpBuildNDRange)
-    return postProcessOCLBuildNDRange(BI, CI, DemangledName);
   if (SPIRVEnableStepExpansion &&
       (DemangledName == "smoothstep" || DemangledName == "step"))
     return expandOCLBuiltinWithScalarArg(CI, DemangledName);
@@ -3161,30 +3031,6 @@ SPIRVToLLVM::SPIRVToLLVM(Module *LLVMModule, SPIRVModule *TheSPIRVModule)
   DbgTran.reset(new SPIRVToLLVMDbgTran(TheSPIRVModule, LLVMModule, this));
 }
 
-std::string SPIRVToLLVM::getOCLBuiltinName(SPIRVInstruction *BI) {
-  auto OC = BI->getOpCode();
-  if (OC == OpBuildNDRange) {
-    auto NDRangeInst = static_cast<SPIRVBuildNDRange *>(BI);
-    auto EleTy = ((NDRangeInst->getOperands())[0])->getType();
-    int Dim = EleTy->isTypeArray() ? EleTy->getArrayLength() : 1;
-    // cygwin does not have std::to_string
-    ostringstream OS;
-    OS << Dim;
-    assert((EleTy->isTypeInt() && Dim == 1) ||
-           (EleTy->isTypeArray() && Dim >= 2 && Dim <= 3));
-    return std::string(kOCLBuiltinName::NDRangePrefix) + OS.str() + "D";
-  }
-
-  return OCLSPIRVBuiltinMap::rmap(OC);
-}
-
-Instruction *SPIRVToLLVM::transOCLBuiltinFromInst(SPIRVInstruction *BI,
-                                                  BasicBlock *BB) {
-  assert(BB && "Invalid BB");
-  auto FuncName = getOCLBuiltinName(BI);
-  return transBuiltinFromInst(FuncName, BI, BB);
-}
-
 std::string getSPIRVFuncSuffix(SPIRVInstruction *BI) {
   string Suffix = "";
   if (BI->getOpCode() == OpCreatePipeFromPipeStorage) {
@@ -3231,6 +3077,17 @@ std::string getSPIRVFuncSuffix(SPIRVInstruction *BI) {
       llvm_unreachable("Invalid address space");
     }
   }
+  if (BI->getOpCode() == OpBuildNDRange) {
+    Suffix += kSPIRVPostfix::Divider;
+    auto *NDRangeInst = static_cast<SPIRVBuildNDRange *>(BI);
+    auto *EleTy = ((NDRangeInst->getOperands())[0])->getType();
+    int Dim = EleTy->isTypeArray() ? EleTy->getArrayLength() : 1;
+    assert((EleTy->isTypeInt() && Dim == 1) ||
+           (EleTy->isTypeArray() && Dim >= 2 && Dim <= 3));
+    ostringstream OS;
+    OS << Dim;
+    Suffix += OS.str() + "D";
+  }
   return Suffix;
 }
 
@@ -3319,8 +3176,13 @@ bool SPIRVToLLVM::translate() {
   // as calls.
   if (!lowerBuiltinVariablesToCalls(M))
     return false;
-  if (!postProcessOCL())
-    return false;
+  if (BM->getDesiredBIsRepresentation() == BIsRepresentation::SPIRVFriendlyIR) {
+    SPIRVWord SrcLangVer = 0;
+    BM->getSourceLanguage(&SrcLangVer);
+    bool IsCpp = SrcLangVer == kOCLVer::CL21;
+    if (!postProcessBuiltinsReturningStruct(M, IsCpp))
+      return false;
+  }
   eraseUselessFunctions(M);
 
   DbgTran->addDbgInfoVersion();

lib/SPIRV/SPIRVReader.h

@@ -78,7 +78,6 @@ class SPIRVToLLVM {
   SPIRVToLLVM(Module *LLVMModule, SPIRVModule *TheSPIRVModule);
 
   static const StringSet<> BuiltInConstFunc;
-  std::string getOCLBuiltinName(SPIRVInstruction *BI);
 
   Type *transType(SPIRVType *BT, bool IsClassMember = false);
   std::string transTypeToOCLTypeName(SPIRVType *BT, bool IsSigned = true);
@@ -123,34 +122,8 @@ class SPIRVToLLVM {
   Value *transConvertInst(SPIRVValue *BV, Function *F, BasicBlock *BB);
   Instruction *transBuiltinFromInst(const std::string &FuncName,
                                     SPIRVInstruction *BI, BasicBlock *BB);
-  Instruction *transOCLBuiltinFromInst(SPIRVInstruction *BI, BasicBlock *BB);
   Instruction *transSPIRVBuiltinFromInst(SPIRVInstruction *BI, BasicBlock *BB);
 
-  /// Post-process translated LLVM module for OpenCL.
-  bool postProcessOCL();
-
-  /// \brief Post-process OpenCL builtin functions returning struct type.
-  ///
-  /// Some OpenCL builtin functions are translated to SPIR-V instructions with
-  /// struct type result, e.g. NDRange creation functions. Such functions
-  /// need to be post-processed to return the struct through sret argument.
-  bool postProcessOCLBuiltinReturnStruct(Function *F);
-
-  /// \brief Post-process OpenCL builtin functions having array argument.
-  ///
-  /// These functions are translated to functions with array type argument
-  /// first, then post-processed to have pointer arguments.
-  bool postProcessOCLBuiltinWithArrayArguments(Function *F,
-                                               StringRef DemangledName);
-
-  /// \brief Post-process OpBuildNDRange.
-  ///   OpBuildNDRange GlobalWorkSize, LocalWorkSize, GlobalWorkOffset
-  /// =>
-  ///   call ndrange_XD(GlobalWorkOffset, GlobalWorkSize, LocalWorkSize)
-  /// \return transformed call instruction.
-  CallInst *postProcessOCLBuildNDRange(SPIRVInstruction *BI, CallInst *CI,
-                                       const std::string &DemangledName);
-
   /// \brief Expand OCL builtin functions with scalar argument, e.g.
   /// step, smoothstep.
   /// gentype func (fp edge, gentype x)

lib/SPIRV/SPIRVToOCL.cpp

@@ -135,6 +135,10 @@ void SPIRVToOCLBase::visitCallInst(CallInst &CI) {
     visitCallSPIRVAvcINTELInstructionBuiltin(&CI, OC);
     return;
   }
+  if (OC == OpBuildNDRange) {
+    visitCallBuildNDRangeBuiltIn(&CI, OC, DemangledName);
+    return;
+  }
   if (OC == OpGenericCastToPtrExplicit) {
     visitCallGenericCastToPtrExplicitBuiltIn(&CI, OC);
     return;
@@ -576,6 +580,34 @@ void SPIRVToOCLBase::visitCallSPIRVImageMediaBlockBuiltin(CallInst *CI, Op OC) {
       },
       &Attrs);
 }
+void SPIRVToOCLBase::visitCallBuildNDRangeBuiltIn(CallInst *CI, Op OC,
+                                                  StringRef DemangledName) {
+  AttributeList Attrs = CI->getCalledFunction()->getAttributes();
+  mutateCallInstOCL(
+      M, CI,
+      [=](CallInst *Call, std::vector<Value *> &Args) {
+        assert(Args.size() == 3);
+        // OpenCL built-in has another order of parameters.
+        auto *GlobalWorkSize = Args[0];
+        auto *LocalWorkSize = Args[1];
+        auto *GlobalWorkOffset = Args[2];
+        Args[0] = GlobalWorkOffset;
+        Args[1] = GlobalWorkSize;
+        Args[2] = LocalWorkSize;
+        // __spirv_BuildNDRange_nD, drop __spirv_
+        StringRef S = DemangledName;
+        S = S.drop_front(strlen(kSPIRVName::Prefix));
+        SmallVector<StringRef, 8> Split;
+        // BuildNDRange_nD
+        S.split(Split, kSPIRVPostfix::Divider,
+                /*MaxSplit=*/-1, /*KeepEmpty=*/false);
+        assert(Split.size() >= 2 && "Invalid SPIRV function name");
+        // Cut _nD and add it to function name.
+        return std::string(kOCLBuiltinName::NDRangePrefix) +
+               Split[1].substr(0, 3).str();
+      },
+      &Attrs);
+}
 
 void SPIRVToOCLBase::visitCallGenericCastToPtrExplicitBuiltIn(CallInst *CI,
                                                               Op OC) {

lib/SPIRV/SPIRVToOCL.h

@@ -107,6 +107,11 @@ class SPIRVToOCLBase : public InstVisitor<SPIRVToOCLBase> {
   /// to_{global|local|private} OCL builtin.
   void visitCallGenericCastToPtrExplicitBuiltIn(CallInst *CI, Op OC);
 
+  /// Transform __spirv_OpBuildINDRange_{1|2|3}D to
+  /// ndrange_{1|2|3}D OCL builtin.
+  void visitCallBuildNDRangeBuiltIn(CallInst *CI, Op OC,
+                                    StringRef DemangledName);
+
   /// Transform __spirv_*Convert_R{ReturnType}{_sat}{_rtp|_rtn|_rtz|_rte} to
   /// convert_{ReturnType}_{sat}{_rtp|_rtn|_rtz|_rte}
   /// example:  <2 x i8> __spirv_SatConvertUToS(<2 x i32>) =>

lib/SPIRV/SPIRVToOCL12.cpp

@@ -59,6 +59,9 @@ bool SPIRVToOCL12Base::runSPIRVToOCL(Module &Module) {
 
   visit(*M);
 
+  postProcessBuiltinsReturningStruct(M);
+  postProcessBuiltinsWithArrayArguments(M);
+
   eraseUselessFunctions(&Module);
 
   LLVM_DEBUG(dbgs() << "After SPIRVToOCL12:\n" << *M);