SPIRVReader: make image and sampler types match clang output (#3293)

For e.g. a `foo(image2d_t %img)` function type, clang produces

    @foo(target("spirv.Image", void, 1, 0, 0, 0, 0, 0, 0) %img)

but after OpenCL -> SPIR-V translation, the SPIR-V reader would
translate it into

    @foo(ptr addrspace(1) %img)

Preserve image argument types as `TargetExtType`s when translating to an
OpenCL target environment. `TargetExtType`s are already emitted for
SPIR-V friendly IR.

Make sure not to mangle the `TargetExtType` argument in the usual way,
but continue using the special type mangling emitted by clang (e.g.
`ocl_image1d_rw`).

Contributes to
#3287

Author: svenvh

lib/SPIRV/OCLUtil.h

@@ -523,6 +523,7 @@ class OCLBuiltinFuncMangleInfo : public SPIRV::BuiltinFuncMangleInfo {
 public:
   OCLBuiltinFuncMangleInfo(Function *F) : F(F) {}
   OCLBuiltinFuncMangleInfo() = default;
+  bool isOpenCL() const override { return true; }
   void init(StringRef UniqName) override;
   // Auxiliary information, it is expected that it is relevant at the moment
   // the init method is called.

lib/SPIRV/SPIRVBuiltinHelper.cpp

@@ -318,11 +318,23 @@ Type *BuiltinCallHelper::getSPIRVType(
                       UseRealType);
 }
 
+/// Return true if we always need to emit a TargetExtType for the given type.
+static bool needsTargetExtTy(spv::Op TypeOpcode) {
+  switch (TypeOpcode) {
+  default:
+    return false;
+  case OpTypeImage:
+  case OpTypeSampledImage:
+  case OpTypeSampler:
+    return true;
+  }
+}
+
 Type *BuiltinCallHelper::getSPIRVType(spv::Op TypeOpcode,
                                       StringRef InnerTypeName,
                                       ArrayRef<unsigned> Parameters,
                                       bool UseRealType) {
-  if (UseTargetTypes) {
+  if (UseTargetTypes || needsTargetExtTy(TypeOpcode)) {
     std::string BaseName = (Twine(kSPIRVTypeName::PrefixAndDelim) +
                             SPIRVOpaqueTypeOpCodeMap::rmap(TypeOpcode))
                                .str();

lib/SPIRV/SPIRVInternal.h

@@ -512,6 +512,7 @@ class BuiltinFuncMangleInfo {
   virtual void init(StringRef UniqUnmangledName) {
     UnmangledName = UniqUnmangledName.str();
   }
+  virtual bool isOpenCL() const { return false; }
 
 protected:
   std::string UnmangledName;

lib/SPIRV/SPIRVToOCL.h

@@ -262,6 +262,8 @@ class SPIRVToOCLBase : public InstVisitor<SPIRVToOCLBase>,
   /// examples: __spirv_SDotKHR => dot, __spirv_SDotAccSatKHR => dot_acc_sat
   void visitCallSPIRVDot(CallInst *CI, Op OC, StringRef DemangledName);
 
+  static std::string translateOpaqueType(StringRef STName);
+
   void translateOpaqueTypes();
 
 private:
@@ -291,8 +293,6 @@ class SPIRVToOCLBase : public InstVisitor<SPIRVToOCLBase>,
   static std::string
   getOCLPipeOpaqueType(SmallVector<std::string, 8> &Postfixes);
 
-  static std::string translateOpaqueType(StringRef STName);
-
   /// Mutate the call instruction based on (optional) image operands at position
   /// ImOpArgIndex. The new function name will be based on NewFuncName, and the
   /// type suffix based on Ty and whether the image operand was signed.

lib/SPIRV/SPIRVUtil.cpp

@@ -49,6 +49,7 @@
 #include "ParameterType.h"
 #include "SPIRVInternal.h"
 #include "SPIRVMDWalker.h"
+#include "SPIRVToOCL.h"
 #include "libSPIRV/SPIRVDecorate.h"
 #include "libSPIRV/SPIRVValue.h"
 
@@ -1288,11 +1289,33 @@ SPIR::TypePrimitiveEnum getOCLTypePrimitiveEnum(StringRef TyName) {
             SPIR::PRIMITIVE_SUB_GROUP_AVC_IME_DUAL_REF_STREAMIN_T)
       .Default(SPIR::PRIMITIVE_NONE);
 }
+
+/// Get TypePrimitiveEnum for special OpenCL type from TargetExtType.
+static SPIR::TypePrimitiveEnum getOCLTypePrimitiveEnum(TargetExtType *Ty) {
+  if (Ty->getName() == "spirv.Image") {
+    std::string SPIRVName;
+    {
+      Type *InnerType = Ty->getTypeParameter(0);
+      raw_string_ostream OS(SPIRVName);
+      OS << kSPIRVTypeName::PrefixAndDelim
+         << SPIRVOpaqueTypeOpCodeMap::rmap(OpTypeImage) << "._"
+         << (InnerType ? convertTypeToPostfix(InnerType) : "void");
+      for (unsigned IntParam : Ty->int_params())
+        OS << kSPIRVTypeName::PostfixDelim << IntParam;
+    }
+    std::string OpenCLName = SPIRVToOCLBase::translateOpaqueType(SPIRVName);
+    return getOCLTypePrimitiveEnum(OpenCLName);
+  }
+
+  return SPIR::PRIMITIVE_NONE;
+}
+
 /// Translates LLVM type to descriptor for mangler.
 /// \param Signed indicates integer type should be translated as signed.
 /// \param VoidPtr indicates i8* should be translated as void*.
 static SPIR::RefParamType transTypeDesc(Type *Ty,
                                         const BuiltinArgTypeMangleInfo &Info,
+                                        bool IsOpenCL,
                                         StringRef InstName = "") {
   bool Signed = Info.IsSigned;
   unsigned Attr = Info.Attr;
@@ -1308,7 +1331,8 @@ static SPIR::RefParamType transTypeDesc(Type *Ty,
   if (Info.IsAtomic && !isa<TypedPointerType>(Ty)) {
     BuiltinArgTypeMangleInfo DTInfo = Info;
     DTInfo.IsAtomic = false;
-    return SPIR::RefParamType(new SPIR::AtomicType(transTypeDesc(Ty, DTInfo)));
+    return SPIR::RefParamType(
+        new SPIR::AtomicType(transTypeDesc(Ty, DTInfo, IsOpenCL)));
   }
   if (auto *IntTy = dyn_cast<IntegerType>(Ty)) {
     switch (IntTy->getBitWidth()) {
@@ -1342,11 +1366,12 @@ static SPIR::RefParamType transTypeDesc(Type *Ty,
     return SPIR::RefParamType(new SPIR::PrimitiveType(SPIR::PRIMITIVE_BFLOAT));
   if (auto *VecTy = dyn_cast<FixedVectorType>(Ty)) {
     return SPIR::RefParamType(new SPIR::VectorType(
-        transTypeDesc(VecTy->getElementType(), Info), VecTy->getNumElements()));
+        transTypeDesc(VecTy->getElementType(), Info, IsOpenCL),
+        VecTy->getNumElements()));
   }
   if (Ty->isArrayTy()) {
     return transTypeDesc(TypedPointerType::get(Ty->getArrayElementType(), 0),
-                         Info);
+                         Info, IsOpenCL);
   }
   if (Ty->isStructTy()) {
     auto Name = Ty->getStructName();
@@ -1373,6 +1398,13 @@ static SPIR::RefParamType transTypeDesc(Type *Ty,
     return SPIR::RefParamType(new SPIR::UserDefinedType(Name.str()));
   }
   if (auto *TargetTy = dyn_cast<TargetExtType>(Ty)) {
+    if (IsOpenCL) {
+      SPIR::TypePrimitiveEnum TPE = getOCLTypePrimitiveEnum(TargetTy);
+      if (TPE != SPIR::PRIMITIVE_NONE) {
+        return SPIR::RefParamType(
+            new SPIR::PrimitiveType(getOCLTypePrimitiveEnum(TargetTy)));
+      }
+    }
     std::string FullName;
     unsigned AS = 0;
     {
@@ -1407,7 +1439,8 @@ static SPIR::RefParamType transTypeDesc(Type *Ty,
       FunctionType *FT = cast<FunctionType>(ET);
       if (InstName.consume_front(kSPIRVName::Prefix) &&
           InstName.starts_with("TaskSequence")) {
-        EPT = new SPIR::PointerType(transTypeDesc(FT->getReturnType(), Info));
+        EPT = new SPIR::PointerType(
+            transTypeDesc(FT->getReturnType(), Info, IsOpenCL));
       } else {
         assert((isVoidFuncTy(FT)) && "Not supported");
         EPT = new SPIR::BlockType;
@@ -1461,7 +1494,7 @@ static SPIR::RefParamType transTypeDesc(Type *Ty,
 
     if (VoidPtr && ET->isIntegerTy(8))
       ET = Type::getVoidTy(ET->getContext());
-    auto *PT = new SPIR::PointerType(transTypeDesc(ET, Info));
+    auto *PT = new SPIR::PointerType(transTypeDesc(ET, Info, IsOpenCL));
     PT->setAddressSpace(static_cast<SPIR::TypeAttributeEnum>(
         TPT->getAddressSpace() + (unsigned)SPIR::ATTR_ADDR_SPACE_FIRST));
     for (unsigned I = SPIR::ATTR_QUALIFIER_FIRST, E = SPIR::ATTR_QUALIFIER_LAST;
@@ -1743,8 +1776,8 @@ std::string mangleBuiltin(StringRef UniqName, ArrayRef<Type *> ArgTypes,
       if (MangleInfo.PointerTy && T->isPointerTy()) {
         T = MangleInfo.PointerTy;
       }
-      FD.Parameters.emplace_back(
-          transTypeDesc(T, BtnInfo->getTypeMangleInfo(I), UniqName));
+      FD.Parameters.emplace_back(transTypeDesc(T, BtnInfo->getTypeMangleInfo(I),
+                                               BtnInfo->isOpenCL(), UniqName));
     }
   }
   // Ellipsis must be the last argument of any function

test/GroupAndSubgroupInstructions.spvasm

@@ -104,8 +104,8 @@
 ; CHECK-CL-DAG: declare spir_func i32 @_Z27intel_sub_group_shuffle_xorij(i32, i32) #[[#Attrs]]
 ; CHECK-CL-DAG: declare spir_func i32 @_Z26intel_sub_group_block_readPU3AS1Kj(ptr addrspace(1)) #[[#Attrs]]
 ; CHECK-CL-DAG: declare spir_func void @_Z27intel_sub_group_block_writePU3AS1jj(ptr addrspace(1), i32) #[[#Attrs]]
-; CHECK-CL-DAG: declare spir_func i32 @_Z26intel_sub_group_block_read14ocl_image2d_rwDv2_i(ptr addrspace(1), <2 x i32>) #[[#Attrs]]
-; CHECK-CL-DAG: declare spir_func void @_Z27intel_sub_group_block_write14ocl_image2d_rwDv2_ij(ptr addrspace(1), <2 x i32>, i32) #[[#Attrs]]
+; CHECK-CL-DAG: declare spir_func i32 @_Z26intel_sub_group_block_read14ocl_image2d_rwDv2_i(target("spirv.Image", void, 1, 0, 0, 0, 0, 0, 2), <2 x i32>) #[[#Attrs]]
+; CHECK-CL-DAG: declare spir_func void @_Z27intel_sub_group_block_write14ocl_image2d_rwDv2_ij(target("spirv.Image", void, 1, 0, 0, 0, 0, 0, 2), <2 x i32>, i32) #[[#Attrs]]
 ; CHECK-CL-DAG: declare spir_func i32 @_Z15sub_group_electv() #[[#Attrs]]
 ; CHECK-CL-DAG: declare spir_func i32 @_Z25sub_group_non_uniform_alli(i32) #[[#Attrs]]
 ; CHECK-CL-DAG: declare spir_func i32 @_Z25sub_group_non_uniform_anyi(i32) #[[#Attrs]]