[SPIRV] Add pass to replace gethandlefromimplicitbinding (#146756)

The HLSL frontend generates call to the intrinsic
@llvm.spv.resource.handlefromimplicitbinding to be able to access a
resource where the set and binding were not explicitly given in the
source code. Determining the correct set and binding cannot be done
during Clang's codegen or earlier because in DXIL, they must first
remove resource that are not accessed before assigning binding locations
to the resource without an explicit binding.

We will follow their lead.

This is a change from DXC, where implicit binding for SPIR-V are
assigned before optimizations.

See https://github.com/llvm/wg-hlsl/pull/309

4 files changed, 163 insertions, 0 deletions

diff --git a/llvm/lib/Target/SPIRV/CMakeLists.txt b/llvm/lib/Target/SPIRV/CMakeLists.txt
index ba09451..6660de9 100644
--- a/llvm/lib/Target/SPIRV/CMakeLists.txt
+++ b/llvm/lib/Target/SPIRV/CMakeLists.txt
@@ -26,6 +26,7 @@ add_llvm_target(SPIRVCodeGen
   SPIRVGlobalRegistry.cpp
   SPIRVInstrInfo.cpp
   SPIRVInstructionSelector.cpp
+  SPIRVLegalizeImplicitBinding.cpp
   SPIRVStripConvergentIntrinsics.cpp
   SPIRVLegalizePointerCast.cpp
   SPIRVMergeRegionExitTargets.cpp
diff --git a/llvm/lib/Target/SPIRV/SPIRV.h b/llvm/lib/Target/SPIRV/SPIRV.h
index 1688fa3..1934e98 100644
--- a/llvm/lib/Target/SPIRV/SPIRV.h
+++ b/llvm/lib/Target/SPIRV/SPIRV.h
@@ -23,6 +23,7 @@ ModulePass *createSPIRVPrepareFunctionsPass(const SPIRVTargetMachine &TM);
 FunctionPass *createSPIRVStructurizerPass();
 FunctionPass *createSPIRVMergeRegionExitTargetsPass();
 FunctionPass *createSPIRVStripConvergenceIntrinsicsPass();
+ModulePass *createSPIRVLegalizeImplicitBindingPass();
 FunctionPass *createSPIRVLegalizePointerCastPass(SPIRVTargetMachine *TM);
 FunctionPass *createSPIRVRegularizerPass();
 FunctionPass *createSPIRVPreLegalizerCombiner();
@@ -49,6 +50,7 @@ void initializeSPIRVRegularizerPass(PassRegistry &);
 void initializeSPIRVMergeRegionExitTargetsPass(PassRegistry &);
 void initializeSPIRVPrepareFunctionsPass(PassRegistry &);
 void initializeSPIRVStripConvergentIntrinsicsPass(PassRegistry &);
+void initializeSPIRVLegalizeImplicitBindingPass(PassRegistry &);
 } // namespace llvm
 
 #endif // LLVM_LIB_TARGET_SPIRV_SPIRV_H
diff --git a/llvm/lib/Target/SPIRV/SPIRVLegalizeImplicitBinding.cpp b/llvm/lib/Target/SPIRV/SPIRVLegalizeImplicitBinding.cpp
new file mode 100644
index 0000000..0398e52
--- /dev/null
+++ b/llvm/lib/Target/SPIRV/SPIRVLegalizeImplicitBinding.cpp
@@ -0,0 +1,159 @@
+//===- SPIRVLegalizeImplicitBinding.cpp - Legalize implicit bindings ----*- C++
+//-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass legalizes the @llvm.spv.resource.handlefromimplicitbinding
+// intrinsic by replacing it with a call to
+// @llvm.spv.resource.handlefrombinding.
+//
+//===----------------------------------------------------------------------===//
+
+#include "SPIRV.h"
+#include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/InstVisitor.h"
+#include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/IntrinsicsSPIRV.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Pass.h"
+#include <algorithm>
+#include <vector>
+
+using namespace llvm;
+
+namespace {
+class SPIRVLegalizeImplicitBinding : public ModulePass {
+public:
+  static char ID;
+  SPIRVLegalizeImplicitBinding() : ModulePass(ID) {}
+
+  bool runOnModule(Module &M) override;
+
+private:
+  void collectBindingInfo(Module &M);
+  uint32_t getAndReserveFirstUnusedBinding(uint32_t DescSet);
+  void replaceImplicitBindingCalls(Module &M);
+
+  // A map from descriptor set to a bit vector of used binding numbers.
+  std::vector<BitVector> UsedBindings;
+  // A list of all implicit binding calls, to be sorted by order ID.
+  SmallVector<CallInst *, 16> ImplicitBindingCalls;
+};
+
+struct BindingInfoCollector : public InstVisitor<BindingInfoCollector> {
+  std::vector<BitVector> &UsedBindings;
+  SmallVector<CallInst *, 16> &ImplicitBindingCalls;
+
+  BindingInfoCollector(std::vector<BitVector> &UsedBindings,
+                       SmallVector<CallInst *, 16> &ImplicitBindingCalls)
+      : UsedBindings(UsedBindings), ImplicitBindingCalls(ImplicitBindingCalls) {
+  }
+
+  void visitCallInst(CallInst &CI) {
+    if (CI.getIntrinsicID() == Intrinsic::spv_resource_handlefrombinding) {
+      const uint32_t DescSet =
+          cast<ConstantInt>(CI.getArgOperand(0))->getZExtValue();
+      const uint32_t Binding =
+          cast<ConstantInt>(CI.getArgOperand(1))->getZExtValue();
+
+      if (UsedBindings.size() <= DescSet) {
+        UsedBindings.resize(DescSet + 1);
+        UsedBindings[DescSet].resize(64);
+      }
+      if (UsedBindings[DescSet].size() <= Binding) {
+        UsedBindings[DescSet].resize(2 * Binding + 1);
+      }
+      UsedBindings[DescSet].set(Binding);
+    } else if (CI.getIntrinsicID() ==
+               Intrinsic::spv_resource_handlefromimplicitbinding) {
+      ImplicitBindingCalls.push_back(&CI);
+    }
+  }
+};
+
+void SPIRVLegalizeImplicitBinding::collectBindingInfo(Module &M) {
+  BindingInfoCollector InfoCollector(UsedBindings, ImplicitBindingCalls);
+  InfoCollector.visit(M);
+
+  // Sort the collected calls by their order ID.
+  std::sort(
+      ImplicitBindingCalls.begin(), ImplicitBindingCalls.end(),
+      [](const CallInst *A, const CallInst *B) {
+        const uint32_t OrderIdArgIdx = 0;
+        const uint32_t OrderA =
+            cast<ConstantInt>(A->getArgOperand(OrderIdArgIdx))->getZExtValue();
+        const uint32_t OrderB =
+            cast<ConstantInt>(B->getArgOperand(OrderIdArgIdx))->getZExtValue();
+        return OrderA < OrderB;
+      });
+}
+
+uint32_t SPIRVLegalizeImplicitBinding::getAndReserveFirstUnusedBinding(
+    uint32_t DescSet) {
+  if (UsedBindings.size() <= DescSet) {
+    UsedBindings.resize(DescSet + 1);
+    UsedBindings[DescSet].resize(64);
+  }
+
+  int NewBinding = UsedBindings[DescSet].find_first_unset();
+  if (NewBinding == -1) {
+    NewBinding = UsedBindings[DescSet].size();
+    UsedBindings[DescSet].resize(2 * NewBinding + 1);
+  }
+
+  UsedBindings[DescSet].set(NewBinding);
+  return NewBinding;
+}
+
+void SPIRVLegalizeImplicitBinding::replaceImplicitBindingCalls(Module &M) {
+  for (CallInst *OldCI : ImplicitBindingCalls) {
+    IRBuilder<> Builder(OldCI);
+    const uint32_t DescSet =
+        cast<ConstantInt>(OldCI->getArgOperand(1))->getZExtValue();
+    const uint32_t NewBinding = getAndReserveFirstUnusedBinding(DescSet);
+
+    SmallVector<Value *, 8> Args;
+    Args.push_back(Builder.getInt32(DescSet));
+    Args.push_back(Builder.getInt32(NewBinding));
+
+    // Copy the remaining arguments from the old call.
+    for (uint32_t i = 2; i < OldCI->arg_size(); ++i) {
+      Args.push_back(OldCI->getArgOperand(i));
+    }
+
+    Function *NewFunc = Intrinsic::getOrInsertDeclaration(
+        &M, Intrinsic::spv_resource_handlefrombinding, OldCI->getType());
+    CallInst *NewCI = Builder.CreateCall(NewFunc, Args);
+    NewCI->setCallingConv(OldCI->getCallingConv());
+
+    OldCI->replaceAllUsesWith(NewCI);
+    OldCI->eraseFromParent();
+  }
+}
+
+bool SPIRVLegalizeImplicitBinding::runOnModule(Module &M) {
+  collectBindingInfo(M);
+  if (ImplicitBindingCalls.empty()) {
+    return false;
+  }
+
+  replaceImplicitBindingCalls(M);
+  return true;
+}
+} // namespace
+
+char SPIRVLegalizeImplicitBinding::ID = 0;
+
+INITIALIZE_PASS(SPIRVLegalizeImplicitBinding, "legalize-spirv-implicit-binding",
+                "Legalize SPIR-V implicit bindings", false, false)
+
+ModulePass *llvm::createSPIRVLegalizeImplicitBindingPass() {
+  return new SPIRVLegalizeImplicitBinding();
+}
\ No newline at end of file
diff --git a/llvm/lib/Target/SPIRV/SPIRVTargetMachine.cpp b/llvm/lib/Target/SPIRV/SPIRVTargetMachine.cpp
index d7cf211..e0bfb77 100644
--- a/llvm/lib/Target/SPIRV/SPIRVTargetMachine.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVTargetMachine.cpp
@@ -226,6 +226,7 @@ void SPIRVPassConfig::addIRPasses() {
 }
 
 void SPIRVPassConfig::addISelPrepare() {
+  addPass(createSPIRVLegalizeImplicitBindingPass());
   addPass(createSPIRVEmitIntrinsicsPass(&getTM<SPIRVTargetMachine>()));
   if (TM.getSubtargetImpl()->isLogicalSPIRV())
     addPass(createSPIRVLegalizePointerCastPass(&getTM<SPIRVTargetMachine>()));