[HLSL] Implementation of DXILResourceImplicitBinding pass

This pass takes advantage of the DXILResourceBinding analysis and assigns
register slots to resources that do not have explicit binding.

Part 2/2 of #136786

Closes #136786

Author: hekota

llvm/include/llvm/Analysis/DXILResource.h

@@ -632,12 +632,15 @@ class DXILResourceBindingInfo {
     RegisterSpace(uint32_t Space) : Space(Space) {
       FreeRanges.emplace_back(0, UINT32_MAX);
     }
+    // Size == -1 means unbounded array
+    bool findAvailableBinding(int32_t Size, uint32_t *RegSlot);
   };
 
   struct BindingSpaces {
     dxil::ResourceClass ResClass;
     llvm::SmallVector<RegisterSpace> Spaces;
     BindingSpaces(dxil::ResourceClass ResClass) : ResClass(ResClass) {}
+    RegisterSpace &getOrInsertSpace(uint32_t Space);
   };
 
 private:
@@ -658,6 +661,7 @@ class DXILResourceBindingInfo {
         OverlappingBinding(false) {}
 
   bool hasImplicitBinding() const { return ImplicitBinding; }
+  void setHasImplicitBinding(bool Value) { ImplicitBinding = Value; }
   bool hasOverlappingBinding() const { return OverlappingBinding; }
 
   BindingSpaces &getBindingSpaces(dxil::ResourceClass RC) {
@@ -673,6 +677,10 @@ class DXILResourceBindingInfo {
     }
   }
 
+  // Size == -1 means unbounded array
+  bool findAvailableBinding(dxil::ResourceClass RC, uint32_t Space,
+                            int32_t Size, uint32_t *RegSlot);
+
   friend class DXILResourceBindingAnalysis;
   friend class DXILResourceBindingWrapperPass;
 };

llvm/include/llvm/InitializePasses.h

@@ -85,6 +85,7 @@ void initializeDCELegacyPassPass(PassRegistry &);
 void initializeDXILMetadataAnalysisWrapperPassPass(PassRegistry &);
 void initializeDXILMetadataAnalysisWrapperPrinterPass(PassRegistry &);
 void initializeDXILResourceBindingWrapperPassPass(PassRegistry &);
+void initializeDXILResourceImplicitBindingLegacyPass(PassRegistry &);
 void initializeDXILResourceTypeWrapperPassPass(PassRegistry &);
 void initializeDXILResourceWrapperPassPass(PassRegistry &);
 void initializeDeadMachineInstructionElimPass(PassRegistry &);

llvm/lib/Analysis/Analysis.cpp

@@ -27,6 +27,7 @@ void llvm::initializeAnalysis(PassRegistry &Registry) {
   initializeCallGraphViewerPass(Registry);
   initializeCycleInfoWrapperPassPass(Registry);
   initializeDXILMetadataAnalysisWrapperPassPass(Registry);
+  initializeDXILResourceWrapperPassPass(Registry);
   initializeDXILResourceBindingWrapperPassPass(Registry);
   initializeDXILResourceTypeWrapperPassPass(Registry);
   initializeDXILResourceWrapperPassPass(Registry);

llvm/lib/Analysis/DXILResource.cpp

@@ -995,6 +995,59 @@ void DXILResourceBindingInfo::populate(Module &M, DXILResourceTypeMap &DRTM) {
   }
 }
 
+// returns false if binding could not be found in given space
+bool DXILResourceBindingInfo::findAvailableBinding(dxil::ResourceClass RC,
+                                                   uint32_t Space, int32_t Size,
+                                                   uint32_t *RegSlot) {
+  BindingSpaces &BS = getBindingSpaces(RC);
+  RegisterSpace &RS = BS.getOrInsertSpace(Space);
+  return RS.findAvailableBinding(Size, RegSlot);
+}
+
+DXILResourceBindingInfo::RegisterSpace &
+DXILResourceBindingInfo::BindingSpaces::getOrInsertSpace(uint32_t Space) {
+  for (auto *I = Spaces.begin(); I != Spaces.end(); ++I) {
+    if (I->Space == Space)
+      return *I;
+    if (I->Space < Space)
+      continue;
+    return *Spaces.insert(I, Space);
+  }
+  return Spaces.emplace_back(Space);
+}
+
+bool DXILResourceBindingInfo::RegisterSpace::findAvailableBinding(
+    int32_t Size, uint32_t *RegSlot) {
+  assert((Size == -1 || Size > 0) && "invalid size");
+
+  if (FreeRanges.empty())
+    return false;
+
+  // unbounded array
+  if (Size == -1) {
+    BindingRange &Last = FreeRanges.back();
+    if (Last.UpperBound != UINT32_MAX)
+      // this space is already occupied by an unbounded array
+      return false;
+    *RegSlot = Last.LowerBound;
+    FreeRanges.pop_back();
+    return true;
+  }
+
+  // single resource or fixed-size array
+  for (BindingRange &R : FreeRanges) {
+    if (R.UpperBound - R.LowerBound + 1 < (uint32_t)Size)
+      continue;
+    *RegSlot = R.LowerBound;
+    // This might create a range where (LowerBound == UpperBound + 1), but
+    // that's ok.
+    R.LowerBound += Size;
+    return true;
+  }
+
+  return false;
+}
+
 //===----------------------------------------------------------------------===//
 
 AnalysisKey DXILResourceTypeAnalysis::Key;

llvm/lib/Target/DirectX/CMakeLists.txt

@@ -31,6 +31,7 @@ add_llvm_target(DirectXCodeGen
   DXILPrepare.cpp
   DXILPrettyPrinter.cpp
   DXILResourceAccess.cpp
+  DXILResourceImplicitBinding.cpp
   DXILShaderFlags.cpp
   DXILTranslateMetadata.cpp
   DXILRootSignature.cpp

llvm/lib/Target/DirectX/DXILResourceImplicitBinding.cpp

@@ -0,0 +1,182 @@
+//===- DXILResourceImplicitBinding.cpp -----------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "DXILResourceImplicitBinding.h"
+#include "DirectX.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Analysis/DXILResource.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/DiagnosticInfo.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicsDirectX.h"
+#include "llvm/IR/Module.h"
+#include "llvm/InitializePasses.h"
+#include <cstdint>
+
+#define DEBUG_TYPE "dxil-resource-implicit-binding"
+
+using namespace llvm;
+using namespace llvm::dxil;
+
+namespace {
+
+static void diagnoseImplicitBindingNotFound(CallInst *ImplBindingCall) {
+  Function *F = ImplBindingCall->getFunction();
+  LLVMContext &Context = F->getParent()->getContext();
+  // FIXME: include the name of the resource in the error message
+  // (llvm/llvm-project#137868)
+  Context.diagnose(
+      DiagnosticInfoGenericWithLoc("resource cannot be allocated", *F,
+                                   ImplBindingCall->getDebugLoc(), DS_Error));
+}
+
+static bool assignBindings(Module &M, DXILResourceBindingInfo &DRBI,
+                           DXILResourceTypeMap &DRTM) {
+  struct ImplBindingCall {
+    int OrderID;
+    CallInst *Call;
+    ImplBindingCall(int OrderID, CallInst *Call)
+        : OrderID(OrderID), Call(Call) {}
+  };
+  SmallVector<ImplBindingCall> Calls;
+  SmallVector<Function *> FunctionsToMaybeRemove;
+
+  // collect all of the llvm.dx.resource.handlefromImplicitbinding calls
+  for (Function &F : M.functions()) {
+    if (!F.isDeclaration())
+      continue;
+
+    if (F.getIntrinsicID() != Intrinsic::dx_resource_handlefromimplicitbinding)
+      continue;
+
+    for (User *U : F.users()) {
+      if (CallInst *CI = dyn_cast<CallInst>(U)) {
+        int OrderID = cast<ConstantInt>(CI->getArgOperand(0))->getZExtValue();
+        Calls.emplace_back(OrderID, CI);
+      }
+    }
+    FunctionsToMaybeRemove.emplace_back(&F);
+  }
+
+  // sort all the collected implicit bindings by OrderID
+  llvm::stable_sort(
+      Calls, [](auto &LHS, auto &RHS) { return LHS.OrderID < RHS.OrderID; });
+
+  // iterate over sorted calls, find binding for each new OrderID and replace
+  // each call with dx_resource_handlefrombinding using the new binding
+  int LastOrderID = -1;
+  llvm::TargetExtType *HandleTy = nullptr;
+  ConstantInt *RegSlotOp = nullptr;
+  bool AllBindingsAssigned = true;
+  bool Changed = false;
+
+  for (auto &IB : Calls) {
+    IRBuilder<> Builder(IB.Call);
+
+    if (IB.OrderID != LastOrderID) {
+      LastOrderID = IB.OrderID;
+      HandleTy = cast<TargetExtType>(IB.Call->getType());
+      ResourceTypeInfo &RTI = DRTM[HandleTy];
+
+      uint32_t Space =
+          cast<ConstantInt>(IB.Call->getArgOperand(1))->getZExtValue();
+      int32_t Size =
+          cast<ConstantInt>(IB.Call->getArgOperand(2))->getZExtValue();
+
+      uint32_t RegSlot;
+      RegSlotOp = nullptr;
+      if (!DRBI.findAvailableBinding(RTI.getResourceClass(), Space, Size,
+                                     &RegSlot)) {
+        diagnoseImplicitBindingNotFound(IB.Call);
+        AllBindingsAssigned = false;
+        continue;
+      }
+      RegSlotOp = ConstantInt::get(Builder.getInt32Ty(), RegSlot);
+    }
+
+    if (!RegSlotOp)
+      continue;
+
+    auto *NewCall = Builder.CreateIntrinsic(
+        HandleTy, Intrinsic::dx_resource_handlefrombinding,
+        {IB.Call->getOperand(1),   /* space */
+         RegSlotOp,                /* register slot */
+         IB.Call->getOperand(2),   /* size */
+         IB.Call->getOperand(3),   /* index */
+         IB.Call->getOperand(4)}); /* non-uniform flag */
+    IB.Call->replaceAllUsesWith(NewCall);
+    IB.Call->eraseFromParent();
+    Changed = true;
+  }
+
+  for (Function *F : FunctionsToMaybeRemove) {
+    if (F->user_empty()) {
+      F->eraseFromParent();
+      Changed = true;
+    }
+  }
+
+  DRBI.setHasImplicitBinding(!AllBindingsAssigned);
+  return Changed;
+}
+
+} // end anonymous namespace
+
+PreservedAnalyses DXILResourceImplicitBinding::run(Module &M,
+                                                   ModuleAnalysisManager &AM) {
+
+  PreservedAnalyses PA;
+
+  DXILResourceBindingInfo &DRBI = AM.getResult<DXILResourceBindingAnalysis>(M);
+  DXILResourceTypeMap &DRTM = AM.getResult<DXILResourceTypeAnalysis>(M);
+  if (DRBI.hasImplicitBinding())
+    if (assignBindings(M, DRBI, DRTM))
+      return PA;
+  return PreservedAnalyses::all();
+}
+
+namespace {
+
+class DXILResourceImplicitBindingLegacy : public ModulePass {
+public:
+  DXILResourceImplicitBindingLegacy() : ModulePass(ID) {}
+
+  bool runOnModule(Module &M) override {
+    DXILResourceTypeMap &DRTM =
+        getAnalysis<DXILResourceTypeWrapperPass>().getResourceTypeMap();
+    DXILResourceBindingInfo &DRBI =
+        getAnalysis<DXILResourceBindingWrapperPass>().getBindingInfo();
+
+    if (DRBI.hasImplicitBinding())
+      return assignBindings(M, DRBI, DRTM);
+    return false;
+  }
+
+  static char ID; // Pass identification.
+  void getAnalysisUsage(llvm::AnalysisUsage &AU) const override {
+    AU.addRequired<DXILResourceTypeWrapperPass>();
+    AU.addRequired<DXILResourceBindingWrapperPass>();
+  }
+};
+
+char DXILResourceImplicitBindingLegacy::ID = 0;
+} // end anonymous namespace
+
+INITIALIZE_PASS_BEGIN(DXILResourceImplicitBindingLegacy, DEBUG_TYPE,
+                      "DXIL Resource Implicit Binding", false, false)
+INITIALIZE_PASS_DEPENDENCY(DXILResourceTypeWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(DXILResourceBindingWrapperPass)
+INITIALIZE_PASS_END(DXILResourceImplicitBindingLegacy, DEBUG_TYPE,
+                    "DXIL Resource Implicit Binding", false, false)
+
+ModulePass *llvm::createDXILResourceImplicitBindingLegacyPass() {
+  return new DXILResourceImplicitBindingLegacy();
+}

llvm/lib/Target/DirectX/DXILResourceImplicitBinding.h

@@ -0,0 +1,29 @@
+//===- DXILResourceImplicitBindings.h --_____________-----------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// \file Assign register slots to resources without explicit binding.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_TARGET_DIRECTX_DXILRESOURCEIMPLICITBINDING_H
+#define LLVM_LIB_TARGET_DIRECTX_DXILRESOURCEIMPLICITBINDING_H
+
+#include "llvm/IR/PassManager.h"
+#include "llvm/Pass.h"
+
+namespace llvm {
+
+class DXILResourceImplicitBinding
+    : public PassInfoMixin<DXILResourceImplicitBinding> {
+public:
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &);
+};
+
+} // namespace llvm
+
+#endif // LLVM_LIB_TARGET_DIRECTX_DXILRESOURCEIMPLICITBINDING_H

llvm/lib/Target/DirectX/DirectX.h

@@ -78,6 +78,12 @@ void initializeDXILResourceAccessLegacyPass(PassRegistry &);
 /// Pass to update resource accesses to use load/store directly.
 FunctionPass *createDXILResourceAccessLegacyPass();
 
+/// Initializer for DXILResourceImplicitBindingLegacyPass
+void initializeDXILResourceImplicitBindingLegacyPass(PassRegistry &);
+
+/// Pass to assign register slots to resources without binding.
+ModulePass *createDXILResourceImplicitBindingLegacyPass();
+
 /// Initializer for DXILTranslateMetadata.
 void initializeDXILTranslateMetadataLegacyPass(PassRegistry &);
 

llvm/lib/Target/DirectX/DirectXPassRegistry.def

@@ -30,6 +30,7 @@ MODULE_PASS("dxil-intrinsic-expansion", DXILIntrinsicExpansion())
 MODULE_PASS("dxil-op-lower", DXILOpLowering())
 MODULE_PASS("dxil-pretty-printer", DXILPrettyPrinterPass(dbgs()))
 MODULE_PASS("dxil-translate-metadata", DXILTranslateMetadata())
+MODULE_PASS("dxil-resource-implicit-binding", DXILResourceImplicitBinding())
 // TODO: rename to print<foo> after NPM switch
 MODULE_PASS("print-dx-shader-flags", dxil::ShaderFlagsAnalysisPrinter(dbgs()))
 MODULE_PASS("print<dxil-root-signature>", dxil::RootSignatureAnalysisPrinter(dbgs()))

llvm/lib/Target/DirectX/DirectXTargetMachine.cpp

@@ -21,6 +21,7 @@
 #include "DXILOpLowering.h"
 #include "DXILPrettyPrinter.h"
 #include "DXILResourceAccess.h"
+#include "DXILResourceImplicitBinding.h"
 #include "DXILRootSignature.h"
 #include "DXILShaderFlags.h"
 #include "DXILTranslateMetadata.h"
@@ -62,6 +63,7 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeDirectXTarget() {
   initializeDXContainerGlobalsPass(*PR);
   initializeDXILOpLoweringLegacyPass(*PR);
   initializeDXILResourceAccessLegacyPass(*PR);
+  initializeDXILResourceImplicitBindingLegacyPass(*PR);
   initializeDXILTranslateMetadataLegacyPass(*PR);
   initializeShaderFlagsAnalysisWrapperPass(*PR);
   initializeRootSignatureAnalysisWrapperPass(*PR);
@@ -99,6 +101,7 @@ class DirectXPassConfig : public TargetPassConfig {
   FunctionPass *createTargetRegisterAllocator(bool) override { return nullptr; }
   void addCodeGenPrepare() override {
     addPass(createDXILFinalizeLinkageLegacyPass());
+    addPass(createDXILResourceImplicitBindingLegacyPass());
     addPass(createDXILIntrinsicExpansionLegacyPass());
     addPass(createDXILCBufferAccessLegacyPass());
     addPass(createDXILDataScalarizationLegacyPass());