[SPIRV] Implement handle_fromBinding intrinsic.

[s-perron]

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

[llvmbot]

@llvm/pr-subscribers-llvm-ir

@llvm/pr-subscribers-backend-spir-v

Author: Steven Perron (s-perron)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/111052.diff

6 Files Affected:

• (modified) llvm/docs/SPIRVUsage.rst (+7)
• (modified) llvm/include/llvm/IR/IntrinsicsSPIRV.td (+10)
• (modified) llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.cpp (+23)
• (modified) llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.h (+3)
• (modified) llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp (+53)
• (added) llvm/test/CodeGen/SPIRV/HlslBufferLoad.ll (+28)

diff --git a/llvm/docs/SPIRVUsage.rst b/llvm/docs/SPIRVUsage.rst
index 5485bb6195c3d4..f3031c809b40bd 100644
--- a/llvm/docs/SPIRVUsage.rst
+++ b/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 `ararySize > 1`, then the binding represents and 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:
 
diff --git a/llvm/include/llvm/IR/IntrinsicsSPIRV.td b/llvm/include/llvm/IR/IntrinsicsSPIRV.td
index 0567efd8a5d7af..767cd1453102b4 100644
--- a/llvm/include/llvm/IR/IntrinsicsSPIRV.td
+++ b/llvm/include/llvm/IR/IntrinsicsSPIRV.td
@@ -84,4 +84,14 @@ let TargetPrefix = "spv" in {
     [IntrNoMem, Commutative] >;
   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]>;
+
+  // 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]>;
 }
diff --git a/llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.cpp b/llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.cpp
index f35c2435e60a4d..c5cede7745ac52 100644
--- a/llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.cpp
@@ -713,6 +713,29 @@ Register SPIRVGlobalRegistry::buildGlobalVariable(
   return Reg;
 }
 
+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. We use the result register of the type in the name.
+  std::string name = ("__resource_" + Twine(VarType->getOperand(0).getReg()) +
+                      "_" + 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,
diff --git a/llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.h b/llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.h
index ace5cfe91ebe48..4e4b57c9502ca7 100644
--- a/llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.h
+++ b/llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.h
@@ -518,6 +518,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,
diff --git a/llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp b/llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp
index 3917ad180b87fc..5fda4fde35925c 100644
--- a/llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp
@@ -44,6 +44,18 @@ using ExtInstList =
 
 namespace {
 
+uint64_t getUnsignedConstantValueFromReg(llvm::Register Reg,
+                                         const llvm::MachineRegisterInfo &MRI) {
+  llvm::SPIRVType *ConstTy = MRI.getVRegDef(Reg);
+  assert(ConstTy && ConstTy->getOpcode() == llvm::SPIRV::ASSIGN_TYPE &&
+         ConstTy->getOperand(1).isReg());
+  llvm::Register ConstReg = ConstTy->getOperand(1).getReg();
+  const llvm::MachineInstr *Const = MRI.getVRegDef(ConstReg);
+  assert(Const && Const->getOpcode() == llvm::TargetOpcode::G_CONSTANT);
+  const llvm::APInt &Val = Const->getOperand(1).getCImm()->getValue();
+  return Val.getZExtValue();
+}
+
 #define GET_GLOBALISEL_PREDICATE_BITSET
 #include "SPIRVGenGlobalISel.inc"
 #undef GET_GLOBALISEL_PREDICATE_BITSET
@@ -232,6 +244,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 +267,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
@@ -2547,6 +2565,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);
@@ -2558,6 +2580,37 @@ bool SPIRVInstructionSelector::selectIntrinsic(Register ResVReg,
   return true;
 }
 
+void SPIRVInstructionSelector::selectHandleFromBinding(Register &ResVReg,
+                                                       const SPIRVType *ResType,
+                                                       MachineInstr &I) const {
+
+  uint32_t Set =
+      getUnsignedConstantValueFromReg(I.getOperand(2).getReg(), *MRI);
+  uint32_t Binding =
+      getUnsignedConstantValueFromReg(I.getOperand(3).getReg(), *MRI);
+  uint32_t ArraySize =
+      getUnsignedConstantValueFromReg(I.getOperand(4).getReg(), *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 {
diff --git a/llvm/test/CodeGen/SPIRV/HlslBufferLoad.ll b/llvm/test/CodeGen/SPIRV/HlslBufferLoad.ll
new file mode 100644
index 00000000000000..90ec6997c9cf73
--- /dev/null
+++ b/llvm/test/CodeGen/SPIRV/HlslBufferLoad.ll
@@ -0,0 +1,28 @@
+; 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 [[BufferVar:%[0-9]+]] DescriptorSet 16
+; CHECK-DAG: OpDecorate [[BufferVar]] Binding 7
+
+; CHECK: [[float:%[0-9]+]] = OpTypeFloat 32
+; CHECK: [[RWBufferType:%[0-9]+]] = OpTypeImage [[float]] Buffer 2 0 0 2 R32i {{$}}
+; CHECK: [[BufferPtrType:%[0-9]+]] = OpTypePointer UniformConstant [[RWBufferType]]
+; CHECK: [[BufferVar]] = OpVariable [[BufferPtrType]] UniformConstant
+
+; CHECK: {{%[0-9]+}} = OpFunction {{%[0-9]+}} DontInline {{%[0-9]+}}
+; CHECK-NEXT: OpLabel
+define void @RWBufferLoad() #0 {
+; CHECK-NEXT: [[buffer:%[0-9]+]] = OpLoad [[RWBufferType]] [[BufferVar]]
+  %buffer0 = call target("spirv.Image", float, 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 [[RWBufferType]] [[BufferVar]]
+  %buffer1 = call target("spirv.Image", float, 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" }
\ No newline at end of file