[SPIRV] Add reads from image buffer for shaders.

llvm/docs/SPIRVUsage.rst

@@ -392,6 +392,14 @@ SPIR-V backend, along with their descriptions and argument details.
        If `arraySize > 1`, then the binding represents an array of resources\
        of the given size, and the handle for the resource at the given index is returned.\
        If the index is possibly non-uniform, then `isUniformIndex` must get set to true.
+   * - `int_spv_typeBufferLoad`
+     - Scalar or vector
+     - `[spirv.Image ImageBuffer, 32-bit Integer coordinate]`
+     - Loads a value from a Vulkan image buffer at the given coordinate. The \
+       image buffer data is assumed to be stored as a 4-element vector. If the \
+       return type is a scalar, then the first element of the vector is \
+       returned. If the return type is an n-element vector, then the first \
+       n-elements of the 4-element vector are returned.
 
 .. _spirv-builtin-functions:
 

llvm/include/llvm/IR/IntrinsicsSPIRV.td

@@ -102,4 +102,10 @@ let TargetPrefix = "spv" in {
             [IntrNoMem]>;
   def int_spv_firstbituhigh : DefaultAttrsIntrinsic<[LLVMScalarOrSameVectorWidth<0, llvm_i32_ty>], [llvm_anyint_ty], [IntrNoMem]>;
   def int_spv_firstbitshigh : DefaultAttrsIntrinsic<[LLVMScalarOrSameVectorWidth<0, llvm_i32_ty>], [llvm_anyint_ty], [IntrNoMem]>;
+
+  // Read a value from the image buffer. It does not translate directly to a
+  // single OpImageRead because the result type is not necessarily a 4 element
+  // vector.
+  def int_spv_typedBufferLoad
+      : DefaultAttrsIntrinsic<[llvm_any_ty], [llvm_any_ty, llvm_i32_ty]>;
 }

llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.cpp

@@ -1056,6 +1056,11 @@ SPIRVGlobalRegistry::getSPIRVTypeForVReg(Register VReg,
   return nullptr;
 }
 
+SPIRVType *SPIRVGlobalRegistry::getResultType(Register VReg) {
+  MachineInstr *Instr = CurMF->getRegInfo().getVRegDef(VReg);
+  return getSPIRVTypeForVReg(Instr->getOperand(1).getReg());
+}
+
 SPIRVType *SPIRVGlobalRegistry::getOrCreateSPIRVType(
     const Type *Ty, MachineIRBuilder &MIRBuilder,
     SPIRV::AccessQualifier::AccessQualifier AccessQual, bool EmitIR) {
@@ -1126,6 +1131,21 @@ SPIRVGlobalRegistry::getScalarOrVectorComponentCount(SPIRVType *Type) const {
              : 1;
 }
 
+SPIRVType *
+SPIRVGlobalRegistry::getScalarOrVectorComponentType(Register VReg) const {
+  return getScalarOrVectorComponentType(getSPIRVTypeForVReg(VReg));
+}
+
+SPIRVType *
+SPIRVGlobalRegistry::getScalarOrVectorComponentType(SPIRVType *Type) const {
+  if (!Type)
+    return nullptr;
+  Register ScalarReg = Type->getOpcode() == SPIRV::OpTypeVector
+                           ? Type->getOperand(1).getReg()
+                           : Type->getOperand(0).getReg();
+  return getSPIRVTypeForVReg(ScalarReg);
+}
+
 unsigned
 SPIRVGlobalRegistry::getScalarOrVectorBitWidth(const SPIRVType *Type) const {
   assert(Type && "Invalid Type pointer");

llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.h

@@ -353,6 +353,9 @@ class SPIRVGlobalRegistry {
   SPIRVType *getSPIRVTypeForVReg(Register VReg,
                                  const MachineFunction *MF = nullptr) const;
 
+  // Return the result type of the instruction defining the register.
+  SPIRVType *getResultType(Register VReg);
+
   // Whether the given VReg has a SPIR-V type mapped to it yet.
   bool hasSPIRVTypeForVReg(Register VReg) const {
     return getSPIRVTypeForVReg(VReg) != nullptr;
@@ -388,6 +391,12 @@ class SPIRVGlobalRegistry {
   unsigned getScalarOrVectorComponentCount(Register VReg) const;
   unsigned getScalarOrVectorComponentCount(SPIRVType *Type) const;
 
+  // Return the component type in a vector if the argument is associated with
+  // a vector type. Returns the argument itself for a scalar type, and nullptr
+  // for a missing type.
+  SPIRVType *getScalarOrVectorComponentType(Register VReg) const;
+  SPIRVType *getScalarOrVectorComponentType(SPIRVType *Type) const;
+
   // For vectors or scalars of booleans, integers and floats, return the scalar
   // type's bitwidth. Otherwise calls llvm_unreachable().
   unsigned getScalarOrVectorBitWidth(const SPIRVType *Type) const;

llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp

@@ -264,6 +264,9 @@ class SPIRVInstructionSelector : public InstructionSelector {
   void selectHandleFromBinding(Register &ResVReg, const SPIRVType *ResType,
                                MachineInstr &I) const;
 
+  void selectReadImageIntrinsic(Register &ResVReg, const SPIRVType *ResType,
+                                MachineInstr &I) const;
+
   // Utilities
   Register buildI32Constant(uint32_t Val, MachineInstr &I,
                             const SPIRVType *ResType = nullptr) const;
@@ -288,6 +291,12 @@ class SPIRVInstructionSelector : public InstructionSelector {
                                   uint32_t Binding, uint32_t ArraySize,
                                   Register IndexReg, bool IsNonUniform,
                                   MachineIRBuilder MIRBuilder) const;
+  SPIRVType *getCorrespondingVec4Type(const SPIRVType *Type,
+                                      MachineInstr &I) const;
+  void extractScalarOrVectorFromVector(Register &ResultReg,
+                                       const SPIRVType *ResType,
+                                       Register &InputReg,
+                                       MachineInstr &InsertionPoint) const;
 };
 
 } // end anonymous namespace
@@ -2767,6 +2776,10 @@ bool SPIRVInstructionSelector::selectIntrinsic(Register ResVReg,
     selectHandleFromBinding(ResVReg, ResType, I);
     return true;
   }
+  case Intrinsic::spv_typedBufferLoad: {
+    selectReadImageIntrinsic(ResVReg, ResType, I);
+    return true;
+  }
   default: {
     std::string DiagMsg;
     raw_string_ostream OS(DiagMsg);
@@ -2803,6 +2816,83 @@ void SPIRVInstructionSelector::selectHandleFromBinding(Register &ResVReg,
       .addUse(VarReg);
 }
 
+void SPIRVInstructionSelector::selectReadImageIntrinsic(
+    Register &ResVReg, const SPIRVType *ResType, MachineInstr &I) const {
+
+  // If the load of the image is in a different basic block, then
+  // this will generate invalid code. A proper solution is to move
+  // the OpLoad from selectHandleFromBinding here. However, to do
+  // that we will need to change the return type of the intrinsic.
+  // We will do that when we can, but for now trying to move forward with other
+  // issues.
+  Register ImageReg = I.getOperand(2).getReg();
+
+  SPIRVType *ReadType = getCorrespondingVec4Type(ResType, I);
+  Register ReadReg = MRI->createVirtualRegister(GR.getRegClass(ReadType));
+  BuildMI(*I.getParent(), I, I.getDebugLoc(), TII.get(SPIRV::OpImageRead))
+      .addDef(ReadReg)
+      .addUse(GR.getSPIRVTypeID(ReadType))
+      .addUse(ImageReg)
+      .addUse(I.getOperand(3).getReg());
+
+  extractScalarOrVectorFromVector(ResVReg, ResType, ReadReg, I);
+}
+
+void SPIRVInstructionSelector::extractScalarOrVectorFromVector(
+    Register &ResultReg, const SPIRVType *ResType, Register &InputReg,
+    MachineInstr &InsertionPoint) const {
+  SPIRVType *InputType = GR.getResultType(InputReg);
+  assert(InputType->getOpcode() == SPIRV::OpTypeVector);
+
+  if (ResType->getOpcode() != SPIRV::OpTypeVector) {
+    assert(ResType == GR.getScalarOrVectorComponentType(InputType));
+    BuildMI(*InsertionPoint.getParent(), InsertionPoint,
+            InsertionPoint.getDebugLoc(), TII.get(SPIRV::OpCompositeExtract))
+        .addDef(ResultReg)
+        .addUse(GR.getSPIRVTypeID(ResType))
+        .addUse(InputReg)
+        .addImm(0);
+    return;
+  }
+
+  uint64_t InputSize = GR.getScalarOrVectorComponentCount(InputType);
+  uint64_t VectorSize = GR.getScalarOrVectorComponentCount(ResType);
+  if (VectorSize == InputSize) {
+    BuildMI(*InsertionPoint.getParent(), InsertionPoint,
+            InsertionPoint.getDebugLoc(), TII.get(SPIRV::OpCopyObject))
+        .addDef(ResultReg)
+        .addUse(GR.getSPIRVTypeID(ResType))
+        .addUse(InputReg);
+    return;
+  }
+
+  assert(VectorSize < InputSize &&
+         "Cannot extract more element than there are in the input.");
+  SmallVector<Register> ComponentRegisters;
+  SPIRVType *ScalarType = GR.getScalarOrVectorComponentType(ResType);
+  const TargetRegisterClass *ScalarRegClass = GR.getRegClass(ScalarType);
+  for (uint64_t i = 0; i < VectorSize; i++) {
+    Register ComponentReg = MRI->createVirtualRegister(ScalarRegClass);
+    BuildMI(*InsertionPoint.getParent(), InsertionPoint,
+            InsertionPoint.getDebugLoc(), TII.get(SPIRV::OpCompositeExtract))
+        .addDef(ComponentReg)
+        .addUse(ScalarType->getOperand(0).getReg())
+        .addUse(InputReg)
+        .addImm(i);
+    ComponentRegisters.emplace_back(ComponentReg);
+  }
+
+  MachineInstrBuilder MIB = BuildMI(*InsertionPoint.getParent(), InsertionPoint,
+                                    InsertionPoint.getDebugLoc(),
+                                    TII.get(SPIRV::OpCompositeConstruct))
+                                .addDef(ResultReg)
+                                .addUse(GR.getSPIRVTypeID(ResType));
+
+  for (Register ComponentReg : ComponentRegisters) {
+    MIB.addUse(ComponentReg);
+  }
+}
+
 Register SPIRVInstructionSelector::buildPointerToResource(
     const SPIRVType *ResType, uint32_t Set, uint32_t Binding,
     uint32_t ArraySize, Register IndexReg, bool IsNonUniform,
@@ -3305,6 +3395,24 @@ bool SPIRVInstructionSelector::selectSpvThreadId(Register ResVReg,
   return MIB.constrainAllUses(TII, TRI, RBI);
 }
 
+SPIRVType *
+SPIRVInstructionSelector::getCorrespondingVec4Type(const SPIRVType *Type,
+                                                   MachineInstr &I) const {
+  MachineIRBuilder MIRBuilder(I);
+  if (Type->getOpcode() != SPIRV::OpTypeVector) {
+    return GR.getOrCreateSPIRVVectorType(Type, 4, MIRBuilder);
+  }
+
+  uint64_t VectorSize = Type->getOperand(2).getImm();
+  if (VectorSize == 4) {
+    return Type;
+  }
+
+  Register ScalarTypeReg = Type->getOperand(1).getReg();
+  const SPIRVType *ScalarType = GR.getSPIRVTypeForVReg(ScalarTypeReg);
+  return GR.getOrCreateSPIRVVectorType(ScalarType, 4, MIRBuilder);
+}
+
 namespace llvm {
 InstructionSelector *
 createSPIRVInstructionSelector(const SPIRVTargetMachine &TM,

llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp

@@ -660,7 +660,8 @@ void RequirementHandler::initAvailableCapabilitiesForOpenCL(
   // Add the min requirements for different OpenCL and SPIR-V versions.
   addAvailableCaps({Capability::Addresses, Capability::Float16Buffer,
                     Capability::Kernel, Capability::Vector16,
-                    Capability::Groups, Capability::GenericPointer});
+                    Capability::Groups, Capability::GenericPointer,
+                    Capability::StorageImageReadWithoutFormat});
   if (ST.hasOpenCLFullProfile())
     addAvailableCaps({Capability::Int64, Capability::Int64Atomics});
   if (ST.hasOpenCLImageSupport()) {
@@ -719,6 +720,11 @@ void RequirementHandler::initAvailableCapabilitiesForVulkan(
          Capability::UniformTexelBufferArrayNonUniformIndexingEXT,
          Capability::StorageTexelBufferArrayNonUniformIndexingEXT});
   }
+
+  // Became core in Vulkan 1.3
+  if (ST.isAtLeastSPIRVVer(VersionTuple(1, 6))) {
+    addAvailableCaps({Capability::StorageImageReadWithoutFormat});
+  }
 }
 
 } // namespace SPIRV
@@ -1005,6 +1011,13 @@ void addOpAccessChainReqs(const MachineInstr &Instr,
   }
 }
 
+static bool imageTypeHasUnknownFormat(SPIRVType *TypeInst) {
+  if (TypeInst->getOpcode() != SPIRV::OpTypeImage)
+    return false;
+  assert(TypeInst->getOperand(7).isImm() && "The image format must be an imm.");
+  return TypeInst->getOperand(7).getImm() == 0;
+}
+
 static void AddDotProductRequirements(const MachineInstr &MI,
                                       SPIRV::RequirementHandler &Reqs,
                                       const SPIRVSubtarget &ST) {
@@ -1411,6 +1424,14 @@ void addInstrRequirements(const MachineInstr &MI,
   case SPIRV::OpUDot:
     AddDotProductRequirements(MI, Reqs, ST);
     break;
+  case SPIRV::OpImageRead: {
+    Register ImageReg = MI.getOperand(2).getReg();
+    SPIRVType *TypeDef = ST.getSPIRVGlobalRegistry()->getResultType(ImageReg);
+    if (imageTypeHasUnknownFormat(TypeDef))
+      Reqs.addCapability(SPIRV::Capability::StorageImageReadWithoutFormat);
+    break;
+  }
+
   default:
     break;
   }

llvm/test/CodeGen/SPIRV/hlsl-resources/BufferLoad.ll

@@ -0,0 +1,66 @@
+; RUN: llc -O0 -verify-machineinstrs -mtriple=spirv-vulkan-library %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv-vulkan-library %s -o - -filetype=obj | spirv-val %}
+
+; CHECK-NOT: OpCapability StorageImageReadWithoutFormat
+
+; CHECK-DAG: OpDecorate [[IntBufferVar:%[0-9]+]] DescriptorSet 16
+; CHECK-DAG: OpDecorate [[IntBufferVar]] Binding 7
+
+; CHECK-DAG: [[int:%[0-9]+]] = OpTypeInt 32 0
+; CHECK-DAG: [[zero:%[0-9]+]] = OpConstant [[int]] 0
+; CHECK-DAG: [[v4_int:%[0-9]+]] = OpTypeVector [[int]] 4
+; CHECK-DAG: [[v2_int:%[0-9]+]] = OpTypeVector [[int]] 2
+; CHECK-DAG: [[RWBufferTypeInt:%[0-9]+]] = OpTypeImage [[int]] Buffer 2 0 0 2 R32i {{$}}
+; CHECK-DAG: [[IntBufferPtrType:%[0-9]+]] = OpTypePointer UniformConstant [[RWBufferTypeInt]]
+; CHECK-DAG: [[IntBufferVar]] = OpVariable [[IntBufferPtrType]] UniformConstant
+
+; CHECK: {{%[0-9]+}} = OpFunction {{%[0-9]+}} DontInline {{%[0-9]+}}
+; CHECK-NEXT: OpLabel
+define void @RWBufferLoad_Vec4_I32() #0 {
+; CHECK: [[buffer:%[0-9]+]] = OpLoad [[RWBufferTypeInt]] [[IntBufferVar]]
+  %buffer0 = call target("spirv.Image", i32, 5, 2, 0, 0, 2, 24)
+      @llvm.spv.handle.fromBinding.tspirv.Image_i32_5_2_0_0_2_24(
+          i32 16, i32 7, i32 1, i32 0, i1 false)
+
+; CHECK: OpImageRead [[v4_int]] [[buffer]] [[zero]]
+  %data0 = call <4 x i32> @llvm.spv.typedBufferLoad(
+      target("spirv.Image", i32, 5, 2, 0, 0, 2, 24) %buffer0, i32 0)
+
+  ret void
+}
+
+; CHECK: {{%[0-9]+}} = OpFunction {{%[0-9]+}} DontInline {{%[0-9]+}}
+; CHECK-NEXT: OpLabel
+define void @RWBufferLoad_I32() #0 {
+; CHECK: [[buffer:%[0-9]+]] = OpLoad [[RWBufferTypeInt]] [[IntBufferVar]]
+  %buffer1 = call target("spirv.Image", i32, 5, 2, 0, 0, 2, 24)
+      @llvm.spv.handle.fromBinding.tspirv.Image_i32_5_2_0_0_2_24(
+          i32 16, i32 7, i32 1, i32 0, i1 false)
+
+; CHECK: [[V:%[0-9]+]] = OpImageRead [[v4_int]] [[buffer]] [[zero]]
+; CHECK: OpCompositeExtract [[int]] [[V]] 0
+  %data1 = call i32 @llvm.spv.typedBufferLoad(
+      target("spirv.Image", i32, 5, 2, 0, 0, 2, 24) %buffer1, i32 0)
+
+  ret void
+}
+
+; CHECK: {{%[0-9]+}} = OpFunction {{%[0-9]+}} DontInline {{%[0-9]+}}
+; CHECK-NEXT: OpLabel
+define void @RWBufferLoad_Vec2_I32() #0 {
+; CHECK: [[buffer:%[0-9]+]] = OpLoad [[RWBufferTypeInt]] [[IntBufferVar]]
+  %buffer0 = call target("spirv.Image", i32, 5, 2, 0, 0, 2, 24)
+      @llvm.spv.handle.fromBinding.tspirv.Image_i32_5_2_0_0_2_24(
+          i32 16, i32 7, i32 1, i32 0, i1 false)
+
+; CHECK: [[V:%[0-9]+]] = OpImageRead [[v4_int]] [[buffer]] [[zero]]
+; CHECK: [[e0:%[0-9]+]] = OpCompositeExtract [[int]] [[V]] 0
+; CHECK: [[e1:%[0-9]+]] = OpCompositeExtract [[int]] [[V]] 1
+; CHECK: OpCompositeConstruct [[v2_int]] [[e0]] [[e1]]
+  %data0 = call <2 x i32> @llvm.spv.typedBufferLoad(
+      target("spirv.Image", i32, 5, 2, 0, 0, 2, 24) %buffer0, i32 0)
+
+  ret void
+}
+
+attributes #0 = { convergent noinline norecurse "frame-pointer"="all" "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" "no-trapping-math"="true" "stack-protector-buffer-size"="8" }

llvm/test/CodeGen/SPIRV/hlsl-resources/UnknownBufferLoad.ll

@@ -0,0 +1,30 @@
+; RUN: llc -O0 -verify-machineinstrs -mtriple=spirv1.6-vulkan1.3-library %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv1.6-vulkan1.3-library %s -o - -filetype=obj | spirv-val %}
+
+; CHECK: OpCapability StorageImageReadWithoutFormat
+; CHECK-DAG: OpDecorate [[IntBufferVar:%[0-9]+]] DescriptorSet 16
+; CHECK-DAG: OpDecorate [[IntBufferVar]] Binding 7
+
+; CHECK-DAG: [[int:%[0-9]+]] = OpTypeInt 32 0
+; CHECK-DAG: [[zero:%[0-9]+]] = OpConstant [[int]] 0
+; CHECK-DAG: [[v4_int:%[0-9]+]] = OpTypeVector [[int]] 4
+; CHECK-DAG: [[RWBufferTypeInt:%[0-9]+]] = OpTypeImage [[int]] Buffer 2 0 0 2 Unknown {{$}}
+; CHECK-DAG: [[IntBufferPtrType:%[0-9]+]] = OpTypePointer UniformConstant [[RWBufferTypeInt]]
+; CHECK-DAG: [[IntBufferVar]] = OpVariable [[IntBufferPtrType]] UniformConstant
+
+; CHECK: {{%[0-9]+}} = OpFunction {{%[0-9]+}} DontInline {{%[0-9]+}}
+; CHECK-NEXT: OpLabel
+define void @RWBufferLoad_Vec4_I32() #0 {
+; CHECK: [[buffer:%[0-9]+]] = OpLoad [[RWBufferTypeInt]] [[IntBufferVar]]
+  %buffer0 = call target("spirv.Image", i32, 5, 2, 0, 0, 2, 0)
+      @llvm.spv.handle.fromBinding.tspirv.Image_f32_5_2_0_0_2_0(
+          i32 16, i32 7, i32 1, i32 0, i1 false)
+
+; CHECK: OpImageRead [[v4_int]] [[buffer]] [[zero]]
+  %data0 = call <4 x i32> @llvm.spv.typedBufferLoad(
+      target("spirv.Image", i32, 5, 2, 0, 0, 2, 0) %buffer0, i32 0)
+
+  ret void
+}
+
+attributes #0 = { convergent noinline norecurse "frame-pointer"="all" "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" "no-trapping-math"="true" "stack-protector-buffer-size"="8" }

llvm/test/CodeGen/SPIRV/read_image.ll

@@ -1,5 +1,7 @@
 ; RUN: llc -O0 -mtriple=spirv64-unknown-unknown %s -o - | FileCheck %s --check-prefix=CHECK-SPIRV
 
+; CHECK-SPIRV: OpCapability StorageImageReadWithoutFormat
+
 ; CHECK-SPIRV: %[[#IntTy:]] = OpTypeInt
 ; CHECK-SPIRV: %[[#IVecTy:]] = OpTypeVector %[[#IntTy]]
 ; CHECK-SPIRV: %[[#FloatTy:]] = OpTypeFloat

llvm/test/CodeGen/SPIRV/transcoding/OpImageReadMS.ll

@@ -1,5 +1,6 @@
 ; RUN: llc -O0 -mtriple=spirv32-unknown-unknown %s -o - | FileCheck %s --check-prefix=CHECK-SPIRV
 
+; CHECK-SPIRV: OpCapability StorageImageReadWithoutFormat
 ; CHECK-SPIRV: %[[#]] = OpImageRead %[[#]] %[[#]] %[[#]] Sample %[[#]]
 
 define spir_kernel void @sample_test(target("spirv.Image", void, 1, 0, 0, 1, 0, 0, 0) %source, i32 %sampler, <4 x float> addrspace(1)* nocapture %results) {