Prepare SPIRVWriter for type conversion without opaque pointers. (#1499)

This changeset is, by itself, not yet enough to get most SPIR-V files to be
emitted when the input module is in opaque pointer mode. However, this does
remove all of the calls to `getPointerElementType` that SPIRVWriter makes
(directly or indirectly), except for the ones that directly correspond to
translating a pointer type.

A later changeset will add a type scavenger that will be used to find the
pointee type of a pointer. All calls to `getPointerElementType` that remain
after this one will be instead shifted to query the type scavenger instead. To
facilitate this change, several methods are added to avoid querying pointer
element types, and they have been added in several places where their need is
known.

The most basic of basic kernels, those that do not use pointer types (other than
declaring global values and functions) will work in opaque pointer mode with
this changeset.

Author: jcranmer-intel

lib/SPIRV/OCLToSPIRV.cpp

@@ -976,13 +976,18 @@ void OCLToSPIRVBase::visitCallReadImageWithSampler(CallInst *CI,
                                                    StringRef MangledName) {
   assert(MangledName.find(kMangledName::Sampler) != StringRef::npos);
   assert(CI->getCalledFunction() && "Unexpected indirect call");
-  AttributeList Attrs = CI->getCalledFunction()->getAttributes();
+  Function *Func = CI->getCalledFunction();
+  AttributeList Attrs = Func->getAttributes();
   bool IsRetScalar = !CI->getType()->isVectorTy();
+  SmallVector<StructType *, 3> ArgStructTys;
+  getParameterTypes(CI, ArgStructTys);
   mutateCallInstSPIRV(
       M, CI,
       [=](CallInst *, std::vector<Value *> &Args, Type *&Ret) {
         auto *ImageTy =
-            OCLTypeToSPIRVPtr->getAdaptedType(Args[0])->getPointerElementType();
+            OCLTypeToSPIRVPtr->getAdaptedArgumentType(Func, 0).second;
+        if (!ImageTy)
+          ImageTy = ArgStructTys[0];
         ImageTy = adaptSPIRVImageType(M, ImageTy);
         auto SampledImgTy = getSPIRVTypeByChangeBaseTypeName(
             M, ImageTy, kSPIRVTypeName::Image, kSPIRVTypeName::SampledImg);
@@ -1696,8 +1701,12 @@ void OCLToSPIRVBase::visitSubgroupAVCBuiltinCallWithSampler(
           if (!isOCLImageStructType(ParamTys[I]))
             continue;
 
-          auto *ImageTy = OCLTypeToSPIRVPtr->getAdaptedType(Args[I])
-                              ->getPointerElementType();
+          auto *ImageTy =
+              OCLTypeToSPIRVPtr
+                  ->getAdaptedArgumentType(CI->getCalledFunction(), I)
+                  .second;
+          if (!ImageTy)
+            ImageTy = ParamTys[I];
           ImageTy = adaptSPIRVImageType(M, ImageTy);
           auto *SampledImgTy = getSPIRVTypeByChangeBaseTypeName(
               M, ImageTy, kSPIRVTypeName::Image, kSPIRVTypeName::VmeImageINTEL);

lib/SPIRV/OCLTypeToSPIRV.cpp

@@ -70,8 +70,8 @@ bool OCLTypeToSPIRVLegacy::runOnModule(Module &M) {
   return runOCLTypeToSPIRV(M);
 }
 
-OCLTypeToSPIRVBase OCLTypeToSPIRVPass::run(llvm::Module &M,
-                                           llvm::ModuleAnalysisManager &MAM) {
+OCLTypeToSPIRVBase &OCLTypeToSPIRVPass::run(llvm::Module &M,
+                                            llvm::ModuleAnalysisManager &MAM) {
   runOCLTypeToSPIRV(M);
   return *this;
 }
@@ -104,11 +104,12 @@ bool OCLTypeToSPIRVBase::runOCLTypeToSPIRV(Module &Module) {
   return false;
 }
 
-void OCLTypeToSPIRVBase::addAdaptedType(Value *V, Type *T) {
+void OCLTypeToSPIRVBase::addAdaptedType(Value *V, Type *Ty,
+                                        unsigned AddrSpace) {
   LLVM_DEBUG(dbgs() << "[add adapted type] ";
              V->printAsOperand(dbgs(), true, M);
-             dbgs() << " => " << *T << '\n');
-  AdaptedTy[V] = T;
+             dbgs() << " => " << *Ty << '\n');
+  AdaptedTy[V] = {Ty, AddrSpace};
 }
 
 void OCLTypeToSPIRVBase::addWork(Function *F) {
@@ -117,34 +118,6 @@ void OCLTypeToSPIRVBase::addWork(Function *F) {
   WorkSet.insert(F);
 }
 
-/// Find index of \param V as argument of function call \param CI.
-static unsigned getArgIndex(CallInst *CI, Value *V) {
-  for (unsigned AI = 0, AE = CI->arg_size(); AI != AE; ++AI) {
-    if (CI->getArgOperand(AI) == V)
-      return AI;
-  }
-  llvm_unreachable("Not argument of function call");
-  return ~0U;
-}
-
-/// Find index of \param V as argument of function call \param CI.
-static unsigned getArgIndex(Function *F, Value *V) {
-  auto A = F->arg_begin(), E = F->arg_end();
-  for (unsigned I = 0; A != E; ++I, ++A) {
-    if (&(*A) == V)
-      return I;
-  }
-  llvm_unreachable("Not argument of function");
-  return ~0U;
-}
-
-/// Get i-th argument of a function.
-static Argument *getArg(Function *F, unsigned I) {
-  auto AI = F->arg_begin();
-  std::advance(AI, I);
-  return &(*AI);
-}
-
 /// Create a new function type if \param F has arguments in AdaptedTy, and
 /// propagates the adapted arguments to functions called by \param F.
 void OCLTypeToSPIRVBase::adaptFunction(Function *F) {
@@ -158,15 +131,17 @@ void OCLTypeToSPIRVBase::adaptFunction(Function *F) {
     auto Loc = AdaptedTy.find(&I);
     auto Found = (Loc != AdaptedTy.end());
     Changed |= Found;
-    ArgTys.push_back(Found ? Loc->second : I.getType());
+    ArgTys.push_back(Found ? Loc->second.first : I.getType());
 
     if (Found) {
-      for (auto U : I.users()) {
-        if (auto CI = dyn_cast<CallInst>(U)) {
-          auto ArgIndex = getArgIndex(CI, &I);
+      auto *Ty = Loc->second.first;
+      unsigned AddrSpace = Loc->second.second;
+      for (auto &U : I.uses()) {
+        if (auto *CI = dyn_cast<CallInst>(U.getUser())) {
+          auto ArgIndex = CI->getArgOperandNo(&U);
           auto CF = CI->getCalledFunction();
           if (AdaptedTy.count(CF) == 0) {
-            addAdaptedType(getArg(CF, ArgIndex), Loc->second);
+            addAdaptedType(CF->getArg(ArgIndex), Ty, AddrSpace);
             addWork(CF);
           }
         }
@@ -179,7 +154,7 @@ void OCLTypeToSPIRVBase::adaptFunction(Function *F) {
 
   auto FT = F->getFunctionType();
   FT = FunctionType::get(FT->getReturnType(), ArgTys, FT->isVarArg());
-  addAdaptedType(F, FT);
+  addAdaptedType(F, FT, 0);
 }
 
 // Handle functions with sampler arguments that don't get called by
@@ -204,15 +179,10 @@ void OCLTypeToSPIRVBase::adaptArgumentsBySamplerUse(Module &M) {
           AdaptedTy.count(SamplerArg) != 0) // Already traced this, move on.
         continue;
 
-      if (SamplerArg->getType()->isPointerTy() &&
-          isSPIRVStructType(SamplerArg->getType()->getPointerElementType(),
-                            kSPIRVTypeName::Sampler))
-        return;
-
-      addAdaptedType(SamplerArg, getSamplerType(&M));
+      addAdaptedType(SamplerArg, getSamplerStructType(&M), SPIRAS_Constant);
       auto Caller = cast<Argument>(SamplerArg)->getParent();
       addWork(Caller);
-      TraceArg(Caller, getArgIndex(Caller, SamplerArg));
+      TraceArg(Caller, Idx);
     }
   };
 
@@ -235,20 +205,28 @@ void OCLTypeToSPIRVBase::adaptFunctionArguments(Function *F) {
   if (TypeMD)
     return;
   bool Changed = false;
-  auto FT = F->getFunctionType();
-  auto PI = FT->param_begin();
   auto Arg = F->arg_begin();
-  for (unsigned I = 0; I < F->arg_size(); ++I, ++PI, ++Arg) {
-    auto NewTy = *PI;
-    if (isPointerToOpaqueStructType(NewTy)) {
-      auto STName = NewTy->getPointerElementType()->getStructName();
+  SmallVector<StructType *, 4> ParamTys;
+  getParameterTypes(F, ParamTys);
+
+  // If we couldn't get any information from demangling, there is nothing that
+  // can be done.
+  if (ParamTys.empty())
+    return;
+
+  for (unsigned I = 0; I < F->arg_size(); ++I, ++Arg) {
+    StructType *NewTy = ParamTys[I];
+    if (NewTy && NewTy->isOpaque()) {
+      auto STName = NewTy->getStructName();
       if (!hasAccessQualifiedName(STName))
         continue;
       if (STName.startswith(kSPR2TypeName::ImagePrefix)) {
         auto Ty = STName.str();
         auto AccStr = getAccessQualifierFullName(Ty);
-        addAdaptedType(&*Arg, getOrCreateOpaquePtrType(
-                                  M, mapOCLTypeNameToSPIRV(Ty, AccStr)));
+        addAdaptedType(
+            &*Arg,
+            getOrCreateOpaqueStructType(M, mapOCLTypeNameToSPIRV(Ty, AccStr)),
+            SPIRAS_Global);
         Changed = true;
       }
     }
@@ -269,16 +247,18 @@ void OCLTypeToSPIRVBase::adaptArgumentsByMetadata(Function *F) {
   for (unsigned I = 0, E = TypeMD->getNumOperands(); I != E; ++I, ++Arg) {
     auto OCLTyStr = getMDOperandAsString(TypeMD, I);
     if (OCLTyStr == OCL_TYPE_NAME_SAMPLER_T) {
-      addAdaptedType(&(*Arg), getSamplerType(M));
+      addAdaptedType(&(*Arg), getSamplerStructType(M), SPIRAS_Constant);
       Changed = true;
     } else if (OCLTyStr.startswith("image") && OCLTyStr.endswith("_t")) {
       auto Ty = (Twine("opencl.") + OCLTyStr).str();
       if (StructType::getTypeByName(F->getContext(), Ty)) {
         auto AccMD = F->getMetadata(SPIR_MD_KERNEL_ARG_ACCESS_QUAL);
         assert(AccMD && "Invalid access qualifier metadata");
         auto AccStr = getMDOperandAsString(AccMD, I);
-        addAdaptedType(&(*Arg), getOrCreateOpaquePtrType(
-                                    M, mapOCLTypeNameToSPIRV(Ty, AccStr)));
+        addAdaptedType(
+            &(*Arg),
+            getOrCreateOpaqueStructType(M, mapOCLTypeNameToSPIRV(Ty, AccStr)),
+            SPIRAS_Global);
         Changed = true;
       }
     }
@@ -315,14 +295,15 @@ void OCLTypeToSPIRVBase::adaptArgumentsByMetadata(Function *F) {
 // opencl data type x and access qualifier y, and use opencl.image_x.y to
 // represent image_x type with access qualifier y.
 //
-Type *OCLTypeToSPIRVBase::getAdaptedType(Value *V) {
-  auto Loc = AdaptedTy.find(V);
-  if (Loc != AdaptedTy.end())
-    return Loc->second;
-
-  if (auto F = dyn_cast<Function>(V))
-    return F->getFunctionType();
-  return V->getType();
+std::pair<Type *, Type *>
+OCLTypeToSPIRVBase::getAdaptedArgumentType(Function *F, unsigned ArgNo) {
+  Value *Arg = F->getArg(ArgNo);
+  auto Loc = AdaptedTy.find(Arg);
+  if (Loc == AdaptedTy.end())
+    return {nullptr, nullptr};
+  Type *PointeeTy = Loc->second.first;
+  Type *PointerTy = PointerType::get(PointeeTy, Loc->second.second);
+  return {PointerTy, PointeeTy};
 }
 
 } // namespace SPIRV

lib/SPIRV/OCLTypeToSPIRV.h

@@ -46,6 +46,7 @@
 #include "llvm/IR/Function.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/IR/ValueMap.h"
 #include "llvm/Pass.h"
 
 #include <map>
@@ -58,23 +59,26 @@ class OCLTypeToSPIRVBase {
   OCLTypeToSPIRVBase();
 
   bool runOCLTypeToSPIRV(llvm::Module &M);
-  /// \return Adapted type based on kernel argument metadata. If \p V is
-  ///   a function, returns function type.
-  /// E.g. for a function with argument of read only opencl.image_2d_t* type
-  /// returns a function with argument of type opencl.image2d_t.read_only*.
-  llvm::Type *getAdaptedType(llvm::Value *V);
+
+  /// Returns the adapted type of the corresponding argument for a function.
+  /// The first value of the returned pair is the LLVM type of the argument.
+  /// The second value of the returned pair is the pointer element type of the
+  /// argument, if the type is a pointer.
+  std::pair<llvm::Type *, llvm::Type *>
+  getAdaptedArgumentType(llvm::Function *F, unsigned ArgNo);
 
 private:
   llvm::Module *M;
   llvm::LLVMContext *Ctx;
-  std::map<llvm::Value *, llvm::Type *> AdaptedTy; // Adapted types for values
-  std::set<llvm::Function *> WorkSet;              // Functions to be adapted
+  // Map of argument/Function -> {pointee type, address space}
+  llvm::ValueMap<llvm::Value *, std::pair<llvm::Type *, unsigned>> AdaptedTy;
+  std::set<llvm::Function *> WorkSet; // Functions to be adapted
 
   void adaptFunctionArguments(llvm::Function *F);
   void adaptArgumentsByMetadata(llvm::Function *F);
   void adaptArgumentsBySamplerUse(llvm::Module &M);
   void adaptFunction(llvm::Function *F);
-  void addAdaptedType(llvm::Value *V, llvm::Type *T);
+  void addAdaptedType(llvm::Value *V, llvm::Type *PointeeTy, unsigned AS);
   void addWork(llvm::Function *F);
 };
 
@@ -92,7 +96,7 @@ class OCLTypeToSPIRVPass : public OCLTypeToSPIRVBase,
 public:
   using Result = OCLTypeToSPIRVBase;
   static llvm::AnalysisKey Key;
-  OCLTypeToSPIRVBase run(llvm::Module &F, llvm::ModuleAnalysisManager &MAM);
+  OCLTypeToSPIRVBase &run(llvm::Module &F, llvm::ModuleAnalysisManager &MAM);
 };
 
 } // namespace SPIRV

lib/SPIRV/OCLUtil.cpp

@@ -1326,47 +1326,30 @@ Instruction *mutateCallInstOCL(
                         TakeFuncName);
 }
 
-static std::pair<StringRef, StringRef>
-getSrcAndDstElememntTypeName(BitCastInst *BIC) {
-  if (!BIC)
-    return std::pair<StringRef, StringRef>("", "");
-
-  Type *SrcTy = BIC->getSrcTy();
-  Type *DstTy = BIC->getDestTy();
-  if (SrcTy->isPointerTy())
-    SrcTy = SrcTy->getPointerElementType();
-  if (DstTy->isPointerTy())
-    DstTy = DstTy->getPointerElementType();
-  auto SrcST = dyn_cast<StructType>(SrcTy);
-  auto DstST = dyn_cast<StructType>(DstTy);
-  if (!DstST || !DstST->hasName() || !SrcST || !SrcST->hasName())
-    return std::pair<StringRef, StringRef>("", "");
-
-  return std::make_pair(SrcST->getName(), DstST->getName());
+static StringRef getStructName(Type *Ty) {
+  if (auto *STy = dyn_cast<StructType>(Ty))
+    return STy->isLiteral() ? "" : Ty->getStructName();
+  return "";
 }
 
-bool isSamplerInitializer(Instruction *Inst) {
-  BitCastInst *BIC = dyn_cast<BitCastInst>(Inst);
-  auto Names = getSrcAndDstElememntTypeName(BIC);
-  if (Names.second == getSPIRVTypeName(kSPIRVTypeName::Sampler) &&
-      Names.first == getSPIRVTypeName(kSPIRVTypeName::ConstantSampler))
-    return true;
-
-  return false;
-}
-
-bool isPipeStorageInitializer(Instruction *Inst) {
-  BitCastInst *BIC = dyn_cast<BitCastInst>(Inst);
-  auto Names = getSrcAndDstElememntTypeName(BIC);
-  if (Names.second == getSPIRVTypeName(kSPIRVTypeName::PipeStorage) &&
-      Names.first == getSPIRVTypeName(kSPIRVTypeName::ConstantPipeStorage))
-    return true;
-
-  return false;
-}
-
-bool isSpecialTypeInitializer(Instruction *Inst) {
-  return isSamplerInitializer(Inst) || isPipeStorageInitializer(Inst);
+Value *unwrapSpecialTypeInitializer(Value *V) {
+  if (auto *BC = dyn_cast<BitCastOperator>(V)) {
+    Type *DestTy = BC->getDestTy();
+    Type *SrcTy = BC->getSrcTy();
+    if (SrcTy->isPointerTy() && !SrcTy->isOpaquePointerTy()) {
+      StringRef SrcName =
+          getStructName(SrcTy->getNonOpaquePointerElementType());
+      StringRef DestName =
+          getStructName(DestTy->getNonOpaquePointerElementType());
+      if (DestName == getSPIRVTypeName(kSPIRVTypeName::PipeStorage) &&
+          SrcName == getSPIRVTypeName(kSPIRVTypeName::ConstantPipeStorage))
+        return BC->getOperand(0);
+      if (DestName == getSPIRVTypeName(kSPIRVTypeName::Sampler) &&
+          SrcName == getSPIRVTypeName(kSPIRVTypeName::ConstantSampler))
+        return BC->getOperand(0);
+    }
+  }
+  return nullptr;
 }
 
 bool isSamplerStructTy(StructType *STy) {

lib/SPIRV/OCLUtil.h

@@ -495,15 +495,10 @@ Instruction *mutateCallInstOCL(
     std::function<Instruction *(CallInst *)> RetMutate,
     AttributeList *Attrs = nullptr, bool TakeFuncName = false);
 
-/// Check if instruction is bitcast from spirv.ConstantSampler to spirv.Sampler
-bool isSamplerInitializer(Instruction *Inst);
-
-/// Check if instruction is bitcast from spirv.ConstantPipeStorage
-/// to spirv.PipeStorage
-bool isPipeStorageInitializer(Instruction *Inst);
-
-/// Check (isSamplerInitializer || isPipeStorageInitializer)
-bool isSpecialTypeInitializer(Instruction *Inst);
+/// If the value is a special type initializer (something that bitcasts from
+/// spirv.ConstantSampler to spirv.Sampler or likewise for PipeStorage), get the
+/// original type initializer, unwrap the bitcast. Otherwise, return nullptr.
+Value *unwrapSpecialTypeInitializer(Value *V);
 
 bool isPipeOrAddressSpaceCastBI(const StringRef MangledName);
 bool isEnqueueKernelBI(const StringRef MangledName);