Added option to choose programm address space in emitted LLVM IR

SPIRV does not specify storage class for functions, although llvm does. This patch enables choise of address space for function in data layout of emitted LLVM module.
KhronosGroup · Feb 2, 2022 · 7c720df · 7c720df
1 parent 984ce6e
commit 7c720df
Show file tree

Hide file tree

Showing 5 changed files with 74 additions and 3 deletions.
diff --git a/include/LLVMSPIRVOpts.h b/include/LLVMSPIRVOpts.h
@@ -181,6 +181,14 @@ class TranslatorOpts {
     PreserveOCLKernelArgTypeMetadataThroughString = Value;
   }
 
+  void setProgramAddressSpace(unsigned Value) { ProgramAddressSpace = Value; }
+
+  bool hasProgramAddressSpace() const { return ProgramAddressSpace.hasValue(); }
+
+  unsigned getProgramAddressSpace() const {
+    return ProgramAddressSpace.getValue();
+  }
+
 private:
   // Common translation options
   VersionNumber MaxVersion = VersionNumber::MaximumVersion;
@@ -208,6 +216,9 @@ class TranslatorOpts {
   // SPIR-V
   llvm::Optional<ArgList> SPIRVAllowUnknownIntrinsics{};
 
+  // Override program address space in module data layout
+  llvm::Optional<unsigned> ProgramAddressSpace;
+
   // Enable support for extra DIExpression opcodes not listed in the SPIR-V
   // DebugInfo specification.
   bool AllowExtraDIExpressions = false;

diff --git a/lib/SPIRV/SPIRVReader.cpp b/lib/SPIRV/SPIRVReader.cpp
@@ -355,6 +355,18 @@ std::string SPIRVToLLVM::transVCTypeName(SPIRVTypeBufferSurfaceINTEL *PST) {
   return VectorComputeUtil::getVCBufferSurfaceName();
 }
 
+static unsigned transPointerStorageClass(SPIRVTypePointer const *PtrType) {
+  assert(PtrType && "Must be valid pointer type.");
+
+  if (PtrType->getPointerElementType()->getOpCode() == OpTypeFunction) {
+    auto *M = PtrType->getModule();
+    if (M->hasProgramAddressSpace())
+      return M->getProgramAddressSpace();
+  }
+
+  return SPIRSPIRVAddrSpaceMap::rmap(PtrType->getPointerStorageClass());
+}
+
 Type *SPIRVToLLVM::transType(SPIRVType *T, bool IsClassMember) {
   auto Loc = TypeMap.find(T);
   if (Loc != TypeMap.end())
@@ -385,7 +397,7 @@ Type *SPIRVToLLVM::transType(SPIRVType *T, bool IsClassMember) {
     return mapType(
         T, PointerType::get(
                transType(T->getPointerElementType(), IsClassMember),
-               SPIRSPIRVAddrSpaceMap::rmap(T->getPointerStorageClass())));
+               transPointerStorageClass(static_cast<SPIRVTypePointer *>(T))));
   case OpTypeVector:
     return mapType(T,
                    FixedVectorType::get(transType(T->getVectorComponentType()),
@@ -3220,15 +3232,23 @@ bool SPIRVToLLVM::translate() {
   return true;
 }
 
+static std::string transDataLayout(SPIRVModule *M,
+                                   std::string const &GeneralLayout) {
+  return M->hasProgramAddressSpace()
+             ? GeneralLayout + "-P" +
+                   std::to_string(M->getProgramAddressSpace())
+             : GeneralLayout;
+}
+
 bool SPIRVToLLVM::transAddressingModel() {
   switch (BM->getAddressingModel()) {
   case AddressingModelPhysical64:
     M->setTargetTriple(SPIR_TARGETTRIPLE64);
-    M->setDataLayout(SPIR_DATALAYOUT64);
+    M->setDataLayout(transDataLayout(BM, SPIR_DATALAYOUT64));
     break;
   case AddressingModelPhysical32:
     M->setTargetTriple(SPIR_TARGETTRIPLE32);
-    M->setDataLayout(SPIR_DATALAYOUT32);
+    M->setDataLayout(transDataLayout(BM, SPIR_DATALAYOUT32));
     break;
   case AddressingModelLogical:
     // Do not set target triple and data layout

diff --git a/lib/SPIRV/libSPIRV/SPIRVModule.h b/lib/SPIRV/libSPIRV/SPIRVModule.h
@@ -488,6 +488,14 @@ class SPIRVModule {
     return TranslationOpts.isAllowedToUseExtension(RequestedExtension);
   }
 
+  bool hasProgramAddressSpace() const {
+    return TranslationOpts.hasProgramAddressSpace();
+  }
+
+  unsigned getProgramAddressSpace() const {
+    return TranslationOpts.getProgramAddressSpace();
+  }
+
   virtual bool isGenArgNameMDEnabled() const final {
     return TranslationOpts.isGenArgNameMDEnabled();
   }

diff --git a/test/program-addrspace-option.ll b/test/program-addrspace-option.ll
@@ -0,0 +1,25 @@
+; RUN: llvm-as %s -o %t.bc
+; RUN: llvm-spirv %t.bc -o %t.spv --spirv-ext=+SPV_INTEL_function_pointers
+; RUN: llvm-spirv -r %t.spv -o %t.bc --override-program-address-space=3
+; RUN: llvm-dis %t.bc -o %t.ll
+; RUN: FileCheck %s --input-file %t.ll
+
+; ModuleID = 'tmp_back.bc'
+; CHECK: target datalayout
+; CHECK-SAME: P3
+
+target datalayout = "e-p:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v24:32:32-v32:32:32-v48:64:64-v64:64:64-v96:128:128-v128:128:128-v192:256:256-v256:256:256-v512:512:512-v1024:1024:1024"
+target triple = "spir-unknown-unknown"
+
+; CHECK: @test
+; CHECK-SAME: addrspace(3)
+; CHECK: %p
+; CHECK-SAME: addrspace(3)
+
+define spir_kernel void @test() {
+entry:
+  %p = alloca void () *, align 8
+  store void () * @test, void () ** %p, align 8
+  ret void
+}
+
diff --git a/tools/llvm-spirv/llvm-spirv.cpp b/tools/llvm-spirv/llvm-spirv.cpp
@@ -152,6 +152,9 @@ static cl::opt<bool>
     SPIRVToolsDis("spirv-tools-dis", cl::init(false),
                   cl::desc("Emit textual assembly using SPIRV-Tools"));
 
+static cl::opt<unsigned> SPIRVProgramAddressSpace(
+    "override-program-address-space", cl::init(0),
+    cl::desc("Override program address space when reading from SPIRV"));
 using SPIRV::ExtensionID;
 
 #ifdef _SPIRV_SUPPORT_TEXT_FMT
@@ -708,6 +711,10 @@ int main(int Ac, char **Av) {
     }
   }
 
+  if (SPIRVProgramAddressSpace.getNumOccurrences() != 0) {
+    Opts.setProgramAddressSpace(SPIRVProgramAddressSpace);
+  }
+
   if (PreserveOCLKernelArgTypeMetadataThroughString.getNumOccurrences() != 0)
     Opts.setPreserveOCLKernelArgTypeMetadataThroughString(true);