[SPIRV] Implement handle_fromBinding intrinsic. (#111052)

Implement the intrinsic `llvm.spv.handle.fromBinding`, which returns the
handle for a global resource. This involves creating a global variable
that matches the return-type, set, and binding in the call, and
returning the handle to that resource.

This commit implements the scalar version. It does not handle arrays of
resources yet. It also does not handle storage buffers yet. We do not
have the type for the storage buffers designed yet.

Part of #81036

Author: s-perron

llvm/docs/SPIRVUsage.rst

@@ -381,6 +381,13 @@ SPIR-V backend, along with their descriptions and argument details.
      - Pointer
      - `[8-bit Integer]`
      - Creates a resource handle for graphics or compute resources. Facilitates the management and use of resources in shaders.
+   * - `int_spv_handle_fromBinding`
+     - spirv.Image
+     - `[32-bit Integer set, 32-bit Integer binding, 32-bit Integer arraySize, 32-bit Integer index, bool isUniformIndex]`
+     - Returns the handle for the resource at the given set and binding.\
+       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.
 
 .. _spirv-builtin-functions:
 

llvm/include/llvm/IR/IntrinsicsSPIRV.td

@@ -85,4 +85,14 @@ let TargetPrefix = "spv" in {
   def int_spv_wave_is_first_lane : DefaultAttrsIntrinsic<[llvm_i1_ty], [], [IntrConvergent]>;
   def int_spv_sign : DefaultAttrsIntrinsic<[LLVMScalarOrSameVectorWidth<0, llvm_i32_ty>], [llvm_any_ty], [IntrNoMem]>;
   def int_spv_radians : DefaultAttrsIntrinsic<[LLVMMatchType<0>], [llvm_anyfloat_ty], [IntrNoMem]>;
+
+  // Create resource handle given the binding information. Returns a 
+  // type appropriate for the kind of resource given the set id, binding id,
+  // array size of the binding, as well as an index and an indicator
+  // whether that index may be non-uniform.
+  def int_spv_handle_fromBinding
+      : DefaultAttrsIntrinsic<
+            [llvm_any_ty],
+            [llvm_i32_ty, llvm_i32_ty, llvm_i32_ty, llvm_i32_ty, llvm_i1_ty],
+            [IntrNoMem]>;
 }

llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.cpp

@@ -713,6 +713,53 @@ Register SPIRVGlobalRegistry::buildGlobalVariable(
   return Reg;
 }
 
+static std::string buildSpirvTypeName(const SPIRVType *Type,
+                                      MachineIRBuilder &MIRBuilder) {
+  switch (Type->getOpcode()) {
+  case SPIRV::OpTypeImage: {
+    Register SampledTypeReg = Type->getOperand(1).getReg();
+    auto *SampledType = MIRBuilder.getMRI()->getUniqueVRegDef(SampledTypeReg);
+    std::string TypeName =
+        "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::OpTypeFloat:
+    return ("f" + Twine(Type->getOperand(1).getImm())).str();
+  case SPIRV::OpTypeInt:
+    if (Type->getOperand(2).getImm())
+      return ("i" + Twine(Type->getOperand(1).getImm())).str();
+    return ("u" + Twine(Type->getOperand(1).getImm())).str();
+  default:
+    llvm_unreachable("Trying to the the name of an unknown type.");
+  }
+}
+
+Register SPIRVGlobalRegistry::getOrCreateGlobalVariableWithBinding(
+    const SPIRVType *VarType, uint32_t Set, uint32_t Binding,
+    MachineIRBuilder &MIRBuilder) {
+  SPIRVType *VarPointerTypeReg = getOrCreateSPIRVPointerType(
+      VarType, MIRBuilder, SPIRV::StorageClass::UniformConstant);
+  Register VarReg =
+      MIRBuilder.getMRI()->createVirtualRegister(&SPIRV::iIDRegClass);
+
+  // TODO: The name should come from the llvm-ir, but how that name will be
+  // passed from the HLSL to the backend has not been decided. Using this place
+  // holder for now.
+  std::string Name = ("__resource_" + buildSpirvTypeName(VarType, MIRBuilder) +
+                      "_" + Twine(Set) + "_" + Twine(Binding))
+                         .str();
+  buildGlobalVariable(VarReg, VarPointerTypeReg, Name, nullptr,
+                      SPIRV::StorageClass::UniformConstant, nullptr, false,
+                      false, SPIRV::LinkageType::Import, MIRBuilder, false);
+
+  buildOpDecorate(VarReg, MIRBuilder, SPIRV::Decoration::DescriptorSet, {Set});
+  buildOpDecorate(VarReg, MIRBuilder, SPIRV::Decoration::Binding, {Binding});
+  return VarReg;
+}
+
 SPIRVType *SPIRVGlobalRegistry::getOpTypeArray(uint32_t NumElems,
                                                SPIRVType *ElemType,
                                                MachineIRBuilder &MIRBuilder,

llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.h

@@ -519,6 +519,9 @@ class SPIRVGlobalRegistry {
                                SPIRV::LinkageType::LinkageType LinkageType,
                                MachineIRBuilder &MIRBuilder,
                                bool IsInstSelector);
+  Register getOrCreateGlobalVariableWithBinding(const SPIRVType *VarType,
+                                                uint32_t Set, uint32_t Binding,
+                                                MachineIRBuilder &MIRBuilder);
 
   // Convenient helpers for getting types with check for duplicates.
   SPIRVType *getOrCreateSPIRVIntegerType(unsigned BitWidth,

llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp

@@ -232,6 +232,9 @@ class SPIRVInstructionSelector : public InstructionSelector {
 
   bool selectUnmergeValues(MachineInstr &I) const;
 
+  void selectHandleFromBinding(Register &ResVReg, const SPIRVType *ResType,
+                               MachineInstr &I) const;
+
   // Utilities
   Register buildI32Constant(uint32_t Val, MachineInstr &I,
                             const SPIRVType *ResType = nullptr) const;
@@ -252,6 +255,9 @@ class SPIRVInstructionSelector : public InstructionSelector {
                                           uint32_t Opcode) const;
   MachineInstrBuilder buildConstGenericPtr(MachineInstr &I, Register SrcPtr,
                                            SPIRVType *SrcPtrTy) const;
+  Register buildPointerToResource(const SPIRVType *ResType, uint32_t Set,
+                                  uint32_t Binding, uint32_t ArraySize,
+                                  MachineIRBuilder MIRBuilder) const;
 };
 
 } // end anonymous namespace
@@ -2549,6 +2555,10 @@ bool SPIRVInstructionSelector::selectIntrinsic(Register ResVReg,
   // Discard internal intrinsics.
   case Intrinsic::spv_value_md:
     break;
+  case Intrinsic::spv_handle_fromBinding: {
+    selectHandleFromBinding(ResVReg, ResType, I);
+    return true;
+  }
   default: {
     std::string DiagMsg;
     raw_string_ostream OS(DiagMsg);
@@ -2560,6 +2570,34 @@ bool SPIRVInstructionSelector::selectIntrinsic(Register ResVReg,
   return true;
 }
 
+void SPIRVInstructionSelector::selectHandleFromBinding(Register &ResVReg,
+                                                       const SPIRVType *ResType,
+                                                       MachineInstr &I) const {
+
+  uint32_t Set = foldImm(I.getOperand(2), MRI);
+  uint32_t Binding = foldImm(I.getOperand(3), MRI);
+  uint32_t ArraySize = foldImm(I.getOperand(4), MRI);
+
+  MachineIRBuilder MIRBuilder(I);
+  Register VarReg =
+      buildPointerToResource(ResType, Set, Binding, ArraySize, MIRBuilder);
+
+  // 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.
+  BuildMI(*I.getParent(), I, I.getDebugLoc(), TII.get(SPIRV::OpLoad))
+      .addDef(ResVReg)
+      .addUse(GR.getSPIRVTypeID(ResType))
+      .addUse(VarReg);
+}
+
+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);
+}
+
 bool SPIRVInstructionSelector::selectAllocaArray(Register ResVReg,
                                                  const SPIRVType *ResType,
                                                  MachineInstr &I) const {
@@ -2863,15 +2901,7 @@ bool SPIRVInstructionSelector::selectSpvThreadId(Register ResVReg,
   // Get Thread ID index. Expecting operand is a constant immediate value,
   // wrapped in a type assignment.
   assert(I.getOperand(2).isReg());
-  Register ThreadIdReg = I.getOperand(2).getReg();
-  SPIRVType *ConstTy = this->MRI->getVRegDef(ThreadIdReg);
-  assert(ConstTy && ConstTy->getOpcode() == SPIRV::ASSIGN_TYPE &&
-         ConstTy->getOperand(1).isReg());
-  Register ConstReg = ConstTy->getOperand(1).getReg();
-  const MachineInstr *Const = this->MRI->getVRegDef(ConstReg);
-  assert(Const && Const->getOpcode() == TargetOpcode::G_CONSTANT);
-  const llvm::APInt &Val = Const->getOperand(1).getCImm()->getValue();
-  const uint32_t ThreadId = Val.getZExtValue();
+  const uint32_t ThreadId = foldImm(I.getOperand(2), MRI);
 
   // Extract the thread ID from the loaded vector value.
   MachineBasicBlock &BB = *I.getParent();

llvm/test/CodeGen/SPIRV/HlslBufferLoad.ll

@@ -0,0 +1,58 @@
+; RUN: llc -verify-machineinstrs -O0 -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-DAG: OpDecorate [[IntBufferVar:%[0-9]+]] DescriptorSet 16
+; CHECK-DAG: OpDecorate [[IntBufferVar]] Binding 7
+; CHECK-DAG: OpDecorate [[FloatBufferVar:%[0-9]+]] DescriptorSet 16
+; CHECK-DAG: OpDecorate [[FloatBufferVar]] Binding 7
+
+; CHECK-DAG: [[float:%[0-9]+]] = OpTypeFloat 32
+; CHECK-DAG: [[int:%[0-9]+]] = OpTypeInt 32 0
+; CHECK-DAG: [[RWBufferTypeInt:%[0-9]+]] = OpTypeImage [[int]] Buffer 2 0 0 2 R32i {{$}}
+; CHECK-DAG: [[RWBufferTypeFloat:%[0-9]+]] = OpTypeImage [[float]] Buffer 2 0 0 2 R32f {{$}}
+; CHECK-DAG: [[IntBufferPtrType:%[0-9]+]] = OpTypePointer UniformConstant [[RWBufferTypeInt]]
+; CHECK-DAG: [[FloatBufferPtrType:%[0-9]+]] = OpTypePointer UniformConstant [[RWBufferTypeFloat]]
+; CHECK-DAG: [[IntBufferVar]] = OpVariable [[IntBufferPtrType]] UniformConstant
+; CHECK-DAG: [[FloatBufferVar]] = OpVariable [[FloatBufferPtrType]] UniformConstant
+
+; CHECK: {{%[0-9]+}} = OpFunction {{%[0-9]+}} DontInline {{%[0-9]+}}
+; CHECK-NEXT: OpLabel
+define void @RWBufferLoad() #0 {
+; CHECK-NEXT: [[buffer:%[0-9]+]] = OpLoad [[RWBufferTypeInt]] [[IntBufferVar]]
+  %buffer0 = call target("spirv.Image", i32, 5, 2, 0, 0, 2, 24)
+      @llvm.spv.handle.fromBinding.tspirv.Image_f32_5_2_0_0_2_24(
+          i32 16, i32 7, i32 1, i32 0, i1 false)
+
+; Make sure we use the same variable with multiple loads.
+; CHECK-NEXT: [[buffer:%[0-9]+]] = OpLoad [[RWBufferTypeInt]] [[IntBufferVar]]
+  %buffer1 = call target("spirv.Image", i32, 5, 2, 0, 0, 2, 24)
+      @llvm.spv.handle.fromBinding.tspirv.Image_f32_5_2_0_0_2_24(
+          i32 16, i32 7, i32 1, i32 0, i1 false)
+  ret void
+}
+
+; CHECK: {{%[0-9]+}} = OpFunction {{%[0-9]+}} DontInline {{%[0-9]+}}
+; CHECK-NEXT: OpLabel
+define void @UseDifferentGlobalVar() #0 {
+; Make sure we use a different variable from the first function. They have
+; different types.
+; CHECK-NEXT: [[buffer:%[0-9]+]] = OpLoad [[RWBufferTypeFloat]] [[FloatBufferVar]]
+  %buffer0 = call target("spirv.Image", float, 5, 2, 0, 0, 2, 3)
+      @llvm.spv.handle.fromBinding.tspirv.Image_f32_5_2_0_0_2_3(
+          i32 16, i32 7, i32 1, i32 0, i1 false)
+  ret void
+}
+
+; CHECK: {{%[0-9]+}} = OpFunction {{%[0-9]+}} DontInline {{%[0-9]+}}
+; CHECK-NEXT: OpLabel
+define void @ReuseGlobalVarFromFirstFunction() #0 {
+; Make sure we use the same variable as the first function. They should be the
+; same in case one function calls the other.
+; CHECK-NEXT: [[buffer:%[0-9]+]] = OpLoad [[RWBufferTypeInt]] [[IntBufferVar]]
+  %buffer1 = call target("spirv.Image", i32, 5, 2, 0, 0, 2, 24)
+      @llvm.spv.handle.fromBinding.tspirv.Image_f32_5_2_0_0_2_24(
+          i32 16, i32 7, i32 1, i32 0, 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" }