Rework the demangler to support recovering more pointer element types. (#1556)

Original commit:
KhronosGroup/SPIRV-LLVM-Translator@b412ae5

Author: jcranmer-intel

llvm-spirv/lib/SPIRV/OCLToSPIRV.cpp

@@ -1010,7 +1010,7 @@ void OCLToSPIRVBase::visitCallReadImageWithSampler(CallInst *CI,
   Function *Func = CI->getCalledFunction();
   AttributeList Attrs = Func->getAttributes();
   bool IsRetScalar = !CI->getType()->isVectorTy();
-  SmallVector<StructType *, 3> ArgStructTys;
+  SmallVector<Type *, 3> ArgStructTys;
   getParameterTypes(CI, ArgStructTys);
   mutateCallInstSPIRV(
       M, CI,
@@ -1073,7 +1073,7 @@ void OCLToSPIRVBase::visitCallGetImageSize(CallInst *CI,
   AttributeList Attrs = CI->getCalledFunction()->getAttributes();
   StringRef TyName;
   SmallVector<StringRef, 4> SubStrs;
-  SmallVector<StructType *, 4> ParamTys;
+  SmallVector<Type *, 4> ParamTys;
   getParameterTypes(CI, ParamTys);
   auto IsImg = isOCLImageStructType(ParamTys[0], &TyName);
   (void)IsImg;
@@ -1632,7 +1632,7 @@ static void processSubgroupBlockReadWriteINTEL(CallInst *CI,
 // reads and vector block reads.
 void OCLToSPIRVBase::visitSubgroupBlockReadINTEL(CallInst *CI) {
   OCLBuiltinTransInfo Info;
-  SmallVector<StructType *, 2> ParamTys;
+  SmallVector<Type *, 2> ParamTys;
   getParameterTypes(CI, ParamTys);
   if (isOCLImageStructType(ParamTys[0]))
     Info.UniqName = getSPIRVFuncName(spv::OpSubgroupImageBlockReadINTEL);
@@ -1647,7 +1647,7 @@ void OCLToSPIRVBase::visitSubgroupBlockReadINTEL(CallInst *CI) {
 // instructions.
 void OCLToSPIRVBase::visitSubgroupBlockWriteINTEL(CallInst *CI) {
   OCLBuiltinTransInfo Info;
-  SmallVector<StructType *, 3> ParamTys;
+  SmallVector<Type *, 3> ParamTys;
   getParameterTypes(CI, ParamTys);
   if (isOCLImageStructType(ParamTys[0]))
     Info.UniqName = getSPIRVFuncName(spv::OpSubgroupImageBlockWriteINTEL);
@@ -1768,7 +1768,7 @@ void OCLToSPIRVBase::visitSubgroupAVCWrapperBuiltinCall(
   OCLSPIRVSubgroupAVCIntelBuiltinMap::find(ToMCEFName, &ToMCEOC);
   assert(ToMCEOC != OpNop && "Invalid Subgroup AVC Intel built-in call");
 
-  SmallVector<StructType *, 2> ParamTys;
+  SmallVector<Type *, 2> ParamTys;
   getParameterTypes(CI, ParamTys);
 
   if (std::strcmp(TyKind, "payload") == 0) {
@@ -1837,7 +1837,7 @@ void OCLToSPIRVBase::visitSubgroupAVCBuiltinCallWithSampler(
   mutateCallInstSPIRV(
       M, CI,
       [=](CallInst *, std::vector<Value *> &Args) {
-        SmallVector<StructType *, 4> ParamTys;
+        SmallVector<Type *, 4> ParamTys;
         getParameterTypes(CI, ParamTys);
         auto *TyIt =
             std::find_if(ParamTys.begin(), ParamTys.end(), isSamplerStructTy);

llvm-spirv/lib/SPIRV/OCLTypeToSPIRV.cpp

@@ -208,7 +208,7 @@ void OCLTypeToSPIRVBase::adaptFunctionArguments(Function *F) {
     return;
   bool Changed = false;
   auto Arg = F->arg_begin();
-  SmallVector<StructType *, 4> ParamTys;
+  SmallVector<Type *, 4> ParamTys;
   getParameterTypes(F, ParamTys);
 
   // If we couldn't get any information from demangling, there is nothing that
@@ -217,7 +217,7 @@ void OCLTypeToSPIRVBase::adaptFunctionArguments(Function *F) {
     return;
 
   for (unsigned I = 0; I < F->arg_size(); ++I, ++Arg) {
-    StructType *NewTy = ParamTys[I];
+    StructType *NewTy = dyn_cast_or_null<StructType>(ParamTys[I]);
     if (NewTy && NewTy->isOpaque()) {
       auto STName = NewTy->getStructName();
       if (!hasAccessQualifiedName(STName))

llvm-spirv/lib/SPIRV/OCLUtil.cpp

@@ -1356,7 +1356,8 @@ Value *unwrapSpecialTypeInitializer(Value *V) {
   return nullptr;
 }
 
-bool isSamplerStructTy(StructType *STy) {
+bool isSamplerStructTy(Type *Ty) {
+  auto *STy = dyn_cast_or_null<StructType>(Ty);
   return STy && STy->hasName() && STy->getName() == kSPR2TypeName::Sampler;
 }
 

llvm-spirv/lib/SPIRV/OCLUtil.h

@@ -507,7 +507,7 @@ bool isEnqueueKernelBI(const StringRef MangledName);
 bool isKernelQueryBI(const StringRef MangledName);
 
 /// Check that the type is the sampler_t
-bool isSamplerStructTy(StructType *Ty);
+bool isSamplerStructTy(Type *Ty);
 
 // Checks if the binary operator is an unfused fmul + fadd instruction.
 bool isUnfusedMulAdd(BinaryOperator *B);

llvm-spirv/lib/SPIRV/SPIRVInternal.h

@@ -50,6 +50,7 @@
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/TypedPointerType.h"
 
 #include <functional>
 #include <utility>
@@ -969,11 +970,15 @@ std::string mangleBuiltin(StringRef UniqName, ArrayRef<Type *> ArgTypes,
 /// Extract the pointee types of arguments from a mangled function name. If the
 /// corresponding type is not a pointer to a struct type, its value will be a
 /// nullptr instead.
-void getParameterTypes(Function *F, SmallVectorImpl<StructType *> &ArgTys);
-inline void getParameterTypes(CallInst *CI,
-                              SmallVectorImpl<StructType *> &ArgTys) {
+void getParameterTypes(
+    Function *F, SmallVectorImpl<Type *> &ArgTys,
+    std::function<std::string(StringRef)> StructNameMapFn = nullptr);
+inline void getParameterTypes(CallInst *CI, SmallVectorImpl<Type *> &ArgTys) {
   return getParameterTypes(CI->getCalledFunction(), ArgTys);
 }
+void getParameterTypes(
+    Function *F, SmallVectorImpl<TypedPointerType *> &ArgTys,
+    std::function<std::string(StringRef)> StructNameMapFn = nullptr);
 
 /// Mangle a function from OpenCL extended instruction set in SPIR-V friendly IR
 /// manner

llvm-spirv/lib/SPIRV/SPIRVToOCL.cpp

@@ -256,9 +256,9 @@ void SPIRVToOCLBase::visitCastInst(CastInst &Cast) {
 
 void SPIRVToOCLBase::visitCallSPIRVImageQuerySize(CallInst *CI) {
   // Get image type
-  SmallVector<StructType *, 4> ParamTys;
+  SmallVector<Type *, 4> ParamTys;
   getParameterTypes(CI, ParamTys);
-  StructType *ImgTy = ParamTys[0];
+  StructType *ImgTy = cast<StructType>(ParamTys[0]);
   assert(ImgTy && ImgTy->isOpaque() &&
          "image type must be an opaque structure");
   StringRef ImgTyName = ImgTy->getName();
@@ -812,7 +812,7 @@ void SPIRVToOCLBase::visitCallSPIRVImageSampleExplicitLodBuiltIn(CallInst *CI,
   assert(CI->getCalledFunction() && "Unexpected indirect call");
   AttributeList Attrs = CI->getCalledFunction()->getAttributes();
   CallInst *CallSampledImg = cast<CallInst>(CI->getArgOperand(0));
-  SmallVector<StructType *, 6> ParamTys;
+  SmallVector<Type *, 6> ParamTys;
   getParameterTypes(CallSampledImg, ParamTys);
   StringRef ImageTypeName;
   bool IsDepthImage = false;
@@ -1322,22 +1322,8 @@ std::string SPIRVToOCLBase::translateOpaqueType(StringRef STName) {
 }
 
 void SPIRVToOCLBase::getParameterTypes(CallInst *CI,
-                                       SmallVectorImpl<StructType *> &Tys) {
-  ::getParameterTypes(CI, Tys);
-  for (auto &Ty : Tys) {
-    if (!Ty)
-      continue;
-    StringRef STName = Ty->getStructName();
-    bool IsSPIRVOpaque =
-        Ty->isOpaque() && STName.startswith(kSPIRVTypeName::PrefixAndDelim);
-
-    if (!IsSPIRVOpaque)
-      continue;
-
-    std::string NewName = translateOpaqueType(STName);
-    if (NewName != STName)
-      Ty = getOrCreateOpaqueStructType(M, NewName);
-  };
+                                       SmallVectorImpl<Type *> &Tys) {
+  ::getParameterTypes(CI->getCalledFunction(), Tys, translateOpaqueType);
 }
 
 void addSPIRVBIsLoweringPass(ModulePassManager &PassMgr,

llvm-spirv/lib/SPIRV/SPIRVToOCL.h

@@ -269,14 +269,16 @@ class SPIRVToOCLBase : public InstVisitor<SPIRVToOCLBase> {
   std::string groupOCToOCLBuiltinName(CallInst *CI, Op OC);
   /// Transform SPV-IR image opaque type into OpenCL representation,
   /// example: spirv.Image._void_1_0_0_0_0_0_1 => opencl.image2d_wo_t
-  std::string getOCLImageOpaqueType(SmallVector<std::string, 8> &Postfixes);
+  static std::string
+  getOCLImageOpaqueType(SmallVector<std::string, 8> &Postfixes);
   /// Transform SPV-IR pipe opaque type into OpenCL representation,
   /// example: spirv.Pipe._0 => opencl.pipe_ro_t
-  std::string getOCLPipeOpaqueType(SmallVector<std::string, 8> &Postfixes);
+  static std::string
+  getOCLPipeOpaqueType(SmallVector<std::string, 8> &Postfixes);
 
-  void getParameterTypes(CallInst *CI, SmallVectorImpl<StructType *> &Tys);
+  void getParameterTypes(CallInst *CI, SmallVectorImpl<Type *> &Tys);
 
-  std::string translateOpaqueType(StringRef STName);
+  static std::string translateOpaqueType(StringRef STName);
 
   /// Mutate the argument list based on (optional) image operands at position
   /// ImOpArgIndex.  Set IsSigned according to any SignExtend/ZeroExtend Image

llvm-spirv/lib/SPIRV/SPIRVTypeScavenger.cpp

@@ -229,7 +229,7 @@ void SPIRVTypeScavenger::deduceFunctionType(Function &F) {
   // If the function is a mangled name, try to recover types from the Itanium
   // name mangling.
   if (F.getName().startswith("_Z")) {
-    SmallVector<StructType *, 8> ParameterTypes;
+    SmallVector<Type *, 8> ParameterTypes;
     getParameterTypes(&F, ParameterTypes);
     for (Argument *Arg : PointerArgs) {
       if (auto *Ty = ParameterTypes[Arg->getArgNo()]) {