[SPIRV][HLSL] Handle arrays of resources

This commit adds the ability to get a particular resource from an array
of resources using the handle_fromBinding intrinsic.

The main changes are:

1. Create an array when generating the type.
2. Add capabilities from
[SPV_EXT_descriptor_indexing](https://htmlpreview.github.io/?https://github.com/KhronosGroup/SPIRV-Registry/blob/main/extensions/EXT/SPV_EXT_descriptor_indexing.html).

We are still missing the ability to declare a runtime array. That will
be done in a follow up PR.

llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.cpp

@@ -716,6 +716,16 @@ Register SPIRVGlobalRegistry::buildGlobalVariable(
 static std::string buildSpirvTypeName(const SPIRVType *Type,
                                       MachineIRBuilder &MIRBuilder) {
   switch (Type->getOpcode()) {
+  case SPIRV::OpTypeSampledImage: {
+    Register SampledTypeReg = Type->getOperand(1).getReg();
+    auto *SampledType = MIRBuilder.getMRI()->getUniqueVRegDef(SampledTypeReg);
+    std::string TypeName =
+        "sampled_image_" + buildSpirvTypeName(SampledType, MIRBuilder);
+    for (uint32_t I = 2; I < Type->getNumOperands(); ++I) {
+      TypeName = (TypeName + '_' + Twine(Type->getOperand(I).getImm())).str();
+    }
+    return TypeName;
+  }
   case SPIRV::OpTypeImage: {
     Register SampledTypeReg = Type->getOperand(1).getReg();
     auto *SampledType = MIRBuilder.getMRI()->getUniqueVRegDef(SampledTypeReg);
@@ -726,8 +736,18 @@ static std::string buildSpirvTypeName(const SPIRVType *Type,
     }
     return TypeName;
   }
+  case SPIRV::OpTypeArray: {
+    Register ElementTypeReg = Type->getOperand(1).getReg();
+    auto *ElementType = MIRBuilder.getMRI()->getUniqueVRegDef(ElementTypeReg);
+    uint32_t ArraySize = 32; // Type->getOperand(2).getCImm()->getZExtValue();
+    return (buildSpirvTypeName(ElementType, MIRBuilder) + Twine("[") +
+            Twine(ArraySize) + Twine("]"))
+        .str();
+  }
   case SPIRV::OpTypeFloat:
     return ("f" + Twine(Type->getOperand(1).getImm())).str();
+  case SPIRV::OpTypeSampler:
+    return ("sampler");
   case SPIRV::OpTypeInt:
     if (Type->getOperand(2).getImm())
       return ("i" + Twine(Type->getOperand(1).getImm())).str();

llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp

@@ -257,6 +257,7 @@ class SPIRVInstructionSelector : public InstructionSelector {
                                            SPIRVType *SrcPtrTy) const;
   Register buildPointerToResource(const SPIRVType *ResType, uint32_t Set,
                                   uint32_t Binding, uint32_t ArraySize,
+                                  Register IndexReg, bool IsNonUniform,
                                   MachineIRBuilder MIRBuilder) const;
 };
 
@@ -2577,10 +2578,15 @@ void SPIRVInstructionSelector::selectHandleFromBinding(Register &ResVReg,
   uint32_t Set = foldImm(I.getOperand(2), MRI);
   uint32_t Binding = foldImm(I.getOperand(3), MRI);
   uint32_t ArraySize = foldImm(I.getOperand(4), MRI);
+  Register IndexReg = I.getOperand(5).getReg();
+  bool IsNonUniform = ArraySize > 1 && foldImm(I.getOperand(6), MRI);
 
   MachineIRBuilder MIRBuilder(I);
-  Register VarReg =
-      buildPointerToResource(ResType, Set, Binding, ArraySize, MIRBuilder);
+  Register VarReg = buildPointerToResource(ResType, Set, Binding, ArraySize,
+                                           IndexReg, IsNonUniform, MIRBuilder);
+
+  if (IsNonUniform)
+    buildOpDecorate(ResVReg, I, TII, SPIRV::Decoration::NonUniformEXT, {});
 
   // TODO: For now we assume the resource is an image, which needs to be
   // loaded to get the handle. That will not be true for storage buffers.
@@ -2592,10 +2598,35 @@ void SPIRVInstructionSelector::selectHandleFromBinding(Register &ResVReg,
 
 Register SPIRVInstructionSelector::buildPointerToResource(
     const SPIRVType *ResType, uint32_t Set, uint32_t Binding,
-    uint32_t ArraySize, MachineIRBuilder MIRBuilder) const {
-  assert(ArraySize == 1 && "Resource arrays are not implemented yet.");
-  return GR.getOrCreateGlobalVariableWithBinding(ResType, Set, Binding,
-                                                 MIRBuilder);
+    uint32_t ArraySize, Register IndexReg, bool IsNonUniform,
+    MachineIRBuilder MIRBuilder) const {
+  if (ArraySize == 1)
+    return GR.getOrCreateGlobalVariableWithBinding(ResType, Set, Binding,
+                                                   MIRBuilder);
+
+  const SPIRVType *VarType = GR.getOrCreateSPIRVArrayType(
+      ResType, ArraySize, *MIRBuilder.getInsertPt(), TII);
+  Register VarReg = GR.getOrCreateGlobalVariableWithBinding(
+      VarType, Set, Binding, MIRBuilder);
+
+  SPIRVType *ResPointerType = GR.getOrCreateSPIRVPointerType(
+      ResType, MIRBuilder, SPIRV::StorageClass::UniformConstant);
+
+  Register AcReg = MRI->createVirtualRegister(&SPIRV::iIDRegClass);
+  if (IsNonUniform) {
+    // It is unclear which value needs to be marked an non-uniform, so both
+    // the index and the access changed are decorated as non-uniform.
+    buildOpDecorate(IndexReg, MIRBuilder, SPIRV::Decoration::NonUniformEXT, {});
+    buildOpDecorate(AcReg, MIRBuilder, SPIRV::Decoration::NonUniformEXT, {});
+  }
+
+  MIRBuilder.buildInstr(SPIRV::OpAccessChain)
+      .addDef(AcReg)
+      .addUse(GR.getSPIRVTypeID(ResPointerType))
+      .addUse(VarReg)
+      .addUse(IndexReg);
+
+  return AcReg;
 }
 
 bool SPIRVInstructionSelector::selectAllocaArray(Register ResVReg,

llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp

@@ -689,11 +689,31 @@ void RequirementHandler::initAvailableCapabilitiesForVulkan(
     const SPIRVSubtarget &ST) {
   addAvailableCaps({Capability::Shader, Capability::Linkage});
 
-  // Provided by all supported Vulkan versions.
+  // Core in Vulkan 1.1 and earlier.
   addAvailableCaps({Capability::Int16, Capability::Int64, Capability::Float16,
                     Capability::Float64, Capability::GroupNonUniform,
                     Capability::Image1D, Capability::SampledBuffer,
-                    Capability::ImageBuffer});
+                    Capability::ImageBuffer,
+                    Capability::UniformBufferArrayDynamicIndexing,
+                    Capability::SampledImageArrayDynamicIndexing,
+                    Capability::StorageBufferArrayDynamicIndexing,
+                    Capability::StorageImageArrayDynamicIndexing});
+
+  // Became core in Vulkan 1.2
+  if (ST.isAtLeastSPIRVVer(VersionTuple(1, 5))) {
+    addAvailableCaps(
+        {Capability::ShaderNonUniformEXT, Capability::RuntimeDescriptorArrayEXT,
+         Capability::InputAttachmentArrayDynamicIndexingEXT,
+         Capability::UniformTexelBufferArrayDynamicIndexingEXT,
+         Capability::StorageTexelBufferArrayDynamicIndexingEXT,
+         Capability::UniformBufferArrayNonUniformIndexingEXT,
+         Capability::SampledImageArrayNonUniformIndexingEXT,
+         Capability::StorageBufferArrayNonUniformIndexingEXT,
+         Capability::StorageImageArrayNonUniformIndexingEXT,
+         Capability::InputAttachmentArrayNonUniformIndexingEXT,
+         Capability::UniformTexelBufferArrayNonUniformIndexingEXT,
+         Capability::StorageTexelBufferArrayNonUniformIndexingEXT});
+  }
 }
 
 } // namespace SPIRV
@@ -729,6 +749,8 @@ static void addOpDecorateReqs(const MachineInstr &MI, unsigned DecIndex,
              Dec == SPIRV::Decoration::ImplementInRegisterMapINTEL) {
     Reqs.addExtension(
         SPIRV::Extension::SPV_INTEL_global_variable_fpga_decorations);
+  } else if (Dec == SPIRV::Decoration::NonUniformEXT) {
+    Reqs.addRequirements(SPIRV::Capability::ShaderNonUniformEXT);
   }
 }
 
@@ -848,6 +870,134 @@ static void AddAtomicFloatRequirements(const MachineInstr &MI,
   }
 }
 
+bool isUniformTexelBuffer(MachineInstr *ImageInst) {
+  if (ImageInst->getOpcode() != SPIRV::OpTypeImage)
+    return false;
+  uint32_t Dim = ImageInst->getOperand(2).getImm();
+  uint32_t Sampled = ImageInst->getOperand(6).getImm();
+  return Dim == SPIRV::Dim::DIM_Buffer && Sampled == 1;
+}
+
+bool isStorageTexelBuffer(MachineInstr *ImageInst) {
+  if (ImageInst->getOpcode() != SPIRV::OpTypeImage)
+    return false;
+  uint32_t Dim = ImageInst->getOperand(2).getImm();
+  uint32_t Sampled = ImageInst->getOperand(6).getImm();
+  return Dim == SPIRV::Dim::DIM_Buffer && Sampled == 2;
+}
+
+bool isSampledImage(MachineInstr *ImageInst) {
+  if (ImageInst->getOpcode() != SPIRV::OpTypeImage)
+    return false;
+  uint32_t Dim = ImageInst->getOperand(2).getImm();
+  uint32_t Sampled = ImageInst->getOperand(6).getImm();
+  return Dim != SPIRV::Dim::DIM_Buffer && Sampled == 1;
+}
+
+bool isInputAttachment(MachineInstr *ImageInst) {
+  if (ImageInst->getOpcode() != SPIRV::OpTypeImage)
+    return false;
+  uint32_t Dim = ImageInst->getOperand(2).getImm();
+  uint32_t Sampled = ImageInst->getOperand(6).getImm();
+  return Dim == SPIRV::Dim::DIM_SubpassData && Sampled == 2;
+}
+
+bool isStorageImage(MachineInstr *ImageInst) {
+  if (ImageInst->getOpcode() != SPIRV::OpTypeImage)
+    return false;
+  uint32_t Dim = ImageInst->getOperand(2).getImm();
+  uint32_t Sampled = ImageInst->getOperand(6).getImm();
+  return Dim != SPIRV::Dim::DIM_Buffer && Sampled == 2;
+}
+
+bool isCombinedImageSampler(MachineInstr *SampledImageInst) {
+  if (SampledImageInst->getOpcode() != SPIRV::OpTypeSampledImage)
+    return false;
+
+  const MachineRegisterInfo &MRI = SampledImageInst->getMF()->getRegInfo();
+  Register ImageReg = SampledImageInst->getOperand(1).getReg();
+  auto *ImageInst = MRI.getUniqueVRegDef(ImageReg);
+  return isSampledImage(ImageInst);
+}
+
+bool hasNonUniformDecoration(Register Reg, const MachineRegisterInfo &MRI) {
+  for (const auto &MI : MRI.reg_instructions(Reg)) {
+    if (MI.getOpcode() != SPIRV::OpDecorate)
+      continue;
+
+    uint32_t Dec = MI.getOperand(1).getImm();
+    if (Dec == SPIRV::Decoration::NonUniformEXT)
+      return true;
+  }
+  return false;
+}
+void addOpAccessChainReqs(const MachineInstr &instr,
+                          SPIRV::RequirementHandler &handler,
+                          const SPIRVSubtarget &subtarget) {
+  const MachineRegisterInfo &MRI = instr.getMF()->getRegInfo();
+  // Get the result type. If it is an image type, then the shader uses
+  // descriptor indexing. The appropriate capabilities will be added based
+  // on the specifics of the image.
+  Register ResTypeReg = instr.getOperand(1).getReg();
+  MachineInstr *ResTypeInst = MRI.getUniqueVRegDef(ResTypeReg);
+
+  assert(ResTypeInst->getOpcode() == SPIRV::OpTypePointer);
+  uint32_t StorageClass = ResTypeInst->getOperand(1).getImm();
+  if (StorageClass != SPIRV::StorageClass::StorageClass::UniformConstant &&
+      StorageClass != SPIRV::StorageClass::StorageClass::Uniform &&
+      StorageClass != SPIRV::StorageClass::StorageClass::StorageBuffer) {
+    return;
+  }
+
+  Register PointeeTypeReg = ResTypeInst->getOperand(2).getReg();
+  MachineInstr *PointeeType = MRI.getUniqueVRegDef(PointeeTypeReg);
+  if (PointeeType->getOpcode() != SPIRV::OpTypeImage &&
+      PointeeType->getOpcode() != SPIRV::OpTypeSampledImage &&
+      PointeeType->getOpcode() != SPIRV::OpTypeSampler)
+    return;
+
+  bool IsNonUniform =
+      hasNonUniformDecoration(instr.getOperand(0).getReg(), MRI);
+  if (isUniformTexelBuffer(PointeeType)) {
+    if (IsNonUniform)
+      handler.addRequirements(
+          SPIRV::Capability::UniformTexelBufferArrayNonUniformIndexingEXT);
+    else
+      handler.addRequirements(
+          SPIRV::Capability::UniformTexelBufferArrayDynamicIndexingEXT);
+  } else if (isInputAttachment(PointeeType)) {
+    if (IsNonUniform)
+      handler.addRequirements(
+          SPIRV::Capability::InputAttachmentArrayNonUniformIndexingEXT);
+    else
+      handler.addRequirements(
+          SPIRV::Capability::InputAttachmentArrayDynamicIndexingEXT);
+  } else if (isStorageTexelBuffer(PointeeType)) {
+    if (IsNonUniform)
+      handler.addRequirements(
+          SPIRV::Capability::StorageTexelBufferArrayNonUniformIndexingEXT);
+    else
+      handler.addRequirements(
+          SPIRV::Capability::StorageTexelBufferArrayDynamicIndexingEXT);
+  } else if (isSampledImage(PointeeType) ||
+             isCombinedImageSampler(PointeeType) ||
+             PointeeType->getOpcode() == SPIRV::OpTypeSampler) {
+    if (IsNonUniform)
+      handler.addRequirements(
+          SPIRV::Capability::SampledImageArrayNonUniformIndexingEXT);
+    else
+      handler.addRequirements(
+          SPIRV::Capability::SampledImageArrayDynamicIndexing);
+  } else if (isStorageImage(PointeeType)) {
+    if (IsNonUniform)
+      handler.addRequirements(
+          SPIRV::Capability::StorageImageArrayNonUniformIndexingEXT);
+    else
+      handler.addRequirements(
+          SPIRV::Capability::StorageImageArrayDynamicIndexing);
+  }
+}
+
 void addInstrRequirements(const MachineInstr &MI,
                           SPIRV::RequirementHandler &Reqs,
                           const SPIRVSubtarget &ST) {
@@ -967,11 +1117,16 @@ void addInstrRequirements(const MachineInstr &MI,
   case SPIRV::OpConstantSampler:
     Reqs.addCapability(SPIRV::Capability::LiteralSampler);
     break;
+  case SPIRV::OpAccessChain:
+    addOpAccessChainReqs(MI, Reqs, ST);
+    break;
   case SPIRV::OpTypeImage:
     addOpTypeImageReqs(MI, Reqs, ST);
     break;
   case SPIRV::OpTypeSampler:
-    Reqs.addCapability(SPIRV::Capability::ImageBasic);
+    if (ST.isOpenCLEnv()) {
+      Reqs.addCapability(SPIRV::Capability::ImageBasic);
+    }
     break;
   case SPIRV::OpTypeForwardPointer:
     // TODO: check if it's OpenCL's kernel.

llvm/lib/Target/SPIRV/SPIRVSymbolicOperands.td

@@ -355,7 +355,9 @@ defm GeometryPointSize : CapabilityOperand<24, 0, 0, [], [Geometry]>;
 defm ImageGatherExtended : CapabilityOperand<25, 0, 0, [], [Shader]>;
 defm StorageImageMultisample : CapabilityOperand<27, 0, 0, [], [Shader]>;
 defm UniformBufferArrayDynamicIndexing : CapabilityOperand<28, 0, 0, [], [Shader]>;
-defm SampledImageArrayDymnamicIndexing : CapabilityOperand<29, 0, 0, [], [Shader]>;
+defm SampledImageArrayDynamicIndexing : CapabilityOperand<29, 0, 0, [], [Shader]>;
+defm StorageBufferArrayDynamicIndexing : CapabilityOperand<30, 0, 0, [], [Shader]>;
+defm StorageImageArrayDynamicIndexing : CapabilityOperand<31, 0, 0, [], [Shader]>;
 defm ClipDistance : CapabilityOperand<32, 0, 0, [], [Shader]>;
 defm CullDistance : CapabilityOperand<33, 0, 0, [], [Shader]>;
 defm SampleRateShading : CapabilityOperand<35, 0, 0, [], [Shader]>;

llvm/test/CodeGen/SPIRV/CombinedSamplerImageDynIdx.ll

@@ -0,0 +1,40 @@
+; RUN: llc -verify-machineinstrs -O0 -mtriple=spirv1.5-vulkan-library %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv1.5-vulkan-library %s -o - -filetype=obj | spirv-val %}
+
+; CHECK: OpCapability Shader
+; CHECK-NEXT: OpCapability SampledImageArrayDynamicIndexing
+; CHECK-NEXT: OpCapability Sampled1D
+; CHECK-NOT: OpCapability
+
+; CHECK-DAG: OpDecorate [[Var:%[0-9]+]] DescriptorSet 3
+; CHECK-DAG: OpDecorate [[Var]] Binding 4
+
+; CHECK-DAG: [[int:%[0-9]+]] = OpTypeInt 32 0
+; CHECK-DAG: [[BufferType:%[0-9]+]] = OpTypeImage [[int]] 1D 2 0 0 1 R32i {{$}}
+; CHECK-DAG: [[CombindedType:%[0-9]+]] = OpTypeSampledImage [[BufferType]]
+; CHECK-DAG: [[BufferPtrType:%[0-9]+]] = OpTypePointer UniformConstant [[CombindedType]]
+; CHECK-DAG: [[ArraySize:%[0-9]+]] = OpConstant [[int]] 3
+; CHECK-DAG: [[One:%[0-9]+]] = OpConstant [[int]] 1
+; CHECK-DAG: [[Zero:%[0-9]+]] = OpConstant [[int]] 0
+; CHECK-DAG: [[BufferArrayType:%[0-9]+]] = OpTypeArray [[CombindedType]] [[ArraySize]]
+; CHECK-DAG: [[ArrayPtrType:%[0-9]+]] = OpTypePointer UniformConstant [[BufferArrayType]]
+; CHECK-DAG: [[Var]] = OpVariable [[ArrayPtrType]] UniformConstant
+
+; CHECK: {{%[0-9]+}} = OpFunction {{%[0-9]+}} DontInline {{%[0-9]+}}
+; CHECK-NEXT: OpLabel
+define void @main() #0 {
+; CHECK: [[ac:%[0-9]+]] = OpAccessChain [[BufferPtrType]] [[Var]] [[Zero]]
+; CHECK: [[buffer:%[0-9]+]] = OpLoad [[CombindedType]] [[ac]]
+  %buffer0 = call target("spirv.SampledImage", i32, 0, 2, 0, 0, 1, 24)
+      @llvm.spv.handle.fromBinding.tspirv.Image_f32_0_2_0_0_1_24(
+          i32 3, i32 4, i32 3, i32 0, i1 false)
+
+; CHECK: [[ac:%[0-9]+]] = OpAccessChain [[BufferPtrType]] [[Var]] [[One]]
+; CHECK: [[buffer:%[0-9]+]] = OpLoad [[CombindedType]] [[ac]]
+  %buffer1 = call target("spirv.SampledImage", i32, 0, 2, 0, 0, 1, 24)
+      @llvm.spv.handle.fromBinding.tspirv.Image_f32_0_2_0_0_1_24(
+          i32 3, i32 4, i32 3, i32 1, i1 false)
+  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/CombinedSamplerImageNonUniformIdx.ll

@@ -0,0 +1,47 @@
+; RUN: llc -verify-machineinstrs -O0 -mtriple=spirv1.5-vulkan-library %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv1.5-vulkan-library %s -o - -filetype=obj | spirv-val %}
+
+; CHECK: OpCapability Shader
+; CHECK: OpCapability ShaderNonUniform
+; CHECK-NEXT: OpCapability SampledImageArrayNonUniformIndexing
+; CHECK-NEXT: OpCapability Sampled1D
+; CHECK-NOT: OpCapability
+
+; CHECK-DAG: OpDecorate [[Var:%[0-9]+]] DescriptorSet 3
+; CHECK-DAG: OpDecorate [[Var]] Binding 4
+; CHECK: OpDecorate [[Zero:%[0-9]+]] NonUniform
+; CHECK: OpDecorate [[ac0:%[0-9]+]] NonUniform
+; CHECK: OpDecorate [[ld0:%[0-9]+]] NonUniform
+; CHECK: OpDecorate [[One:%[0-9]+]] NonUniform
+; CHECK: OpDecorate [[ac1:%[0-9]+]] NonUniform
+; CHECK: OpDecorate [[ld1:%[0-9]+]] NonUniform
+
+; CHECK-DAG: [[int:%[0-9]+]] = OpTypeInt 32 0
+; CHECK-DAG: [[BufferType:%[0-9]+]] = OpTypeImage [[int]] 1D 2 0 0 1 R32i {{$}}
+; CHECK-DAG: [[CombindedType:%[0-9]+]] = OpTypeSampledImage [[BufferType]]
+; CHECK-DAG: [[BufferPtrType:%[0-9]+]] = OpTypePointer UniformConstant [[CombindedType]]
+; CHECK-DAG: [[ArraySize:%[0-9]+]] = OpConstant [[int]] 3
+; CHECK-DAG: [[One]] = OpConstant [[int]] 1
+; CHECK-DAG: [[Zero]] = OpConstant [[int]] 0
+; CHECK-DAG: [[BufferArrayType:%[0-9]+]] = OpTypeArray [[CombindedType]] [[ArraySize]]
+; CHECK-DAG: [[ArrayPtrType:%[0-9]+]] = OpTypePointer UniformConstant [[BufferArrayType]]
+; CHECK-DAG: [[Var]] = OpVariable [[ArrayPtrType]] UniformConstant
+
+; CHECK: {{%[0-9]+}} = OpFunction {{%[0-9]+}} DontInline {{%[0-9]+}}
+; CHECK-NEXT: OpLabel
+define void @main() #0 {
+; CHECK: [[ac0]] = OpAccessChain [[BufferPtrType]] [[Var]] [[Zero]]
+; CHECK: [[ld0:%[0-9]+]] = OpLoad [[CombindedType]] [[ac0]]
+  %buffer0 = call target("spirv.SampledImage", i32, 0, 2, 0, 0, 1, 24)
+      @llvm.spv.handle.fromBinding.tspirv.Image_f32_0_2_0_0_1_24(
+          i32 3, i32 4, i32 3, i32 0, i1 true)
+
+; CHECK: [[ac1]] = OpAccessChain [[BufferPtrType]] [[Var]] [[One]]
+; CHECK: [[ld1]] = OpLoad [[CombindedType]] [[ac1]]
+  %buffer1 = call target("spirv.SampledImage", i32, 0, 2, 0, 0, 1, 24)
+      @llvm.spv.handle.fromBinding.tspirv.Image_f32_0_2_0_0_1_24(
+          i32 3, i32 4, i32 3, i32 1, i1 true)
+  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" }