[DirectX] Implement DXILResourceBindingAnalysis (#137258)

`DXILResourceBindingAnalysis` analyses explicit resource bindings in the
module and puts together lists of used virtual register spaces and
available virtual register slot ranges for each binding type. It also
stores additional information found during the analysis such as whether
the module uses implicit bindings or if any of the bindings overlap.

This information will be used in `DXILResourceImplicitBindings` pass
(coming soon) to assign register slots to resources with implicit
bindings, and in a post-optimization validation pass that will raise
diagnostic about overlapping bindings.

Part 1/2 of #136786

Author: hekota

llvm/include/llvm/Analysis/DXILResource.h

@@ -10,13 +10,16 @@
 #define LLVM_ANALYSIS_DXILRESOURCE_H
 
 #include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Alignment.h"
 #include "llvm/Support/DXILABI.h"
+#include <climits>
+#include <cstdint>
 
 namespace llvm {
 class CallInst;
@@ -586,6 +589,125 @@ class DXILResourceWrapperPass : public ModulePass {
 
 ModulePass *createDXILResourceWrapperPassPass();
 
+//===----------------------------------------------------------------------===//
+
+// DXILResourceBindingInfo stores the results of DXILResourceBindingAnalysis
+// which analyses all llvm.dx.resource.handlefrombinding calls in the module
+// and puts together lists of used virtual register spaces and available
+// virtual register slot ranges for each binding type.
+// It also stores additional information found during the analysis such as
+// whether the module uses implicit bindings or if any of the bindings overlap.
+//
+// This information will be used in DXILResourceImplicitBindings pass to assign
+// register slots to resources with implicit bindings, and in a
+// post-optimization validation pass that will raise diagnostic about
+// overlapping bindings.
+//
+// For example for these resource bindings:
+//
+// RWBuffer<float> A[10] : register(u3);
+// RWBuffer<float> B[] : register(u5, space2)
+//
+// The analysis result for UAV binding type will look like this:
+//
+// UAVSpaces {
+//   ResClass = ResourceClass::UAV,
+//   Spaces = {
+//     { Space = 0, FreeRanges = {{ 0, 2 }, { 13, UINT32_MAX }} },
+//     { Space = 2, FreeRanges = {{ 0, 4 }} }
+//   }
+// }
+//
+class DXILResourceBindingInfo {
+public:
+  struct BindingRange {
+    uint32_t LowerBound;
+    uint32_t UpperBound;
+    BindingRange(uint32_t LB, uint32_t UB) : LowerBound(LB), UpperBound(UB) {}
+  };
+
+  struct RegisterSpace {
+    uint32_t Space;
+    SmallVector<BindingRange> FreeRanges;
+    RegisterSpace(uint32_t Space) : Space(Space) {
+      FreeRanges.emplace_back(0, UINT32_MAX);
+    }
+  };
+
+  struct BindingSpaces {
+    dxil::ResourceClass RC;
+    llvm::SmallVector<RegisterSpace> Spaces;
+    BindingSpaces(dxil::ResourceClass RC) : RC(RC) {}
+  };
+
+private:
+  BindingSpaces SRVSpaces, UAVSpaces, CBufferSpaces, SamplerSpaces;
+  bool ImplicitBinding;
+  bool OverlappingBinding;
+
+  // Populate the resource binding info given explicit resource binding calls
+  // in the module.
+  void populate(Module &M, DXILResourceTypeMap &DRTM);
+
+public:
+  DXILResourceBindingInfo()
+      : SRVSpaces(dxil::ResourceClass::SRV),
+        UAVSpaces(dxil::ResourceClass::UAV),
+        CBufferSpaces(dxil::ResourceClass::CBuffer),
+        SamplerSpaces(dxil::ResourceClass::Sampler), ImplicitBinding(false),
+        OverlappingBinding(false) {}
+
+  bool hasImplicitBinding() const { return ImplicitBinding; }
+  bool hasOverlappingBinding() const { return OverlappingBinding; }
+
+  BindingSpaces &getBindingSpaces(dxil::ResourceClass RC) {
+    switch (RC) {
+    case dxil::ResourceClass::SRV:
+      return SRVSpaces;
+    case dxil::ResourceClass::UAV:
+      return UAVSpaces;
+    case dxil::ResourceClass::CBuffer:
+      return CBufferSpaces;
+    case dxil::ResourceClass::Sampler:
+      return SamplerSpaces;
+    }
+  }
+
+  friend class DXILResourceBindingAnalysis;
+  friend class DXILResourceBindingWrapperPass;
+};
+
+class DXILResourceBindingAnalysis
+    : public AnalysisInfoMixin<DXILResourceBindingAnalysis> {
+  friend AnalysisInfoMixin<DXILResourceBindingAnalysis>;
+
+  static AnalysisKey Key;
+
+public:
+  using Result = DXILResourceBindingInfo;
+
+  DXILResourceBindingInfo run(Module &M, ModuleAnalysisManager &AM);
+};
+
+class DXILResourceBindingWrapperPass : public ModulePass {
+  std::unique_ptr<DXILResourceBindingInfo> BindingInfo;
+
+public:
+  static char ID;
+
+  DXILResourceBindingWrapperPass();
+  ~DXILResourceBindingWrapperPass() override;
+
+  DXILResourceBindingInfo &getBindingInfo() { return *BindingInfo; }
+  const DXILResourceBindingInfo &getBindingInfo() const { return *BindingInfo; }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override;
+  bool runOnModule(Module &M) override;
+  void releaseMemory() override;
+};
+
+ModulePass *createDXILResourceBindingWrapperPassPass();
+
 } // namespace llvm
 
 #endif // LLVM_ANALYSIS_DXILRESOURCE_H

llvm/include/llvm/IR/IntrinsicsDirectX.td

@@ -27,6 +27,15 @@ def int_dx_resource_handlefrombinding
           [llvm_i32_ty, llvm_i32_ty, llvm_i32_ty, llvm_i32_ty, llvm_i1_ty],
           [IntrNoMem]>;
 
+// Create resource handle with implicit binding in given register space.
+// Returns a `target("dx.")` type appropriate for the kind of resource and
+// the range size and index of the binding.
+def int_dx_resource_handlefromimplicitbinding
+    : DefaultAttrsIntrinsic<
+          [llvm_any_ty],
+          [llvm_i32_ty, llvm_i32_ty, llvm_i32_ty, llvm_i32_ty, llvm_i1_ty],
+          [IntrNoMem]>;
+
 def int_dx_resource_getpointer
     : DefaultAttrsIntrinsic<[llvm_anyptr_ty], [llvm_any_ty, llvm_i32_ty],
                             [IntrNoMem]>;

llvm/include/llvm/InitializePasses.h

@@ -84,8 +84,9 @@ void initializeDAHPass(PassRegistry &);
 void initializeDCELegacyPassPass(PassRegistry &);
 void initializeDXILMetadataAnalysisWrapperPassPass(PassRegistry &);
 void initializeDXILMetadataAnalysisWrapperPrinterPass(PassRegistry &);
-void initializeDXILResourceWrapperPassPass(PassRegistry &);
+void initializeDXILResourceBindingWrapperPassPass(PassRegistry &);
 void initializeDXILResourceTypeWrapperPassPass(PassRegistry &);
+void initializeDXILResourceWrapperPassPass(PassRegistry &);
 void initializeDeadMachineInstructionElimPass(PassRegistry &);
 void initializeDebugifyMachineModulePass(PassRegistry &);
 void initializeDependenceAnalysisWrapperPassPass(PassRegistry &);

llvm/lib/Analysis/Analysis.cpp

@@ -27,8 +27,9 @@ void llvm::initializeAnalysis(PassRegistry &Registry) {
   initializeCallGraphViewerPass(Registry);
   initializeCycleInfoWrapperPassPass(Registry);
   initializeDXILMetadataAnalysisWrapperPassPass(Registry);
-  initializeDXILResourceWrapperPassPass(Registry);
+  initializeDXILResourceBindingWrapperPassPass(Registry);
   initializeDXILResourceTypeWrapperPassPass(Registry);
+  initializeDXILResourceWrapperPassPass(Registry);
   initializeDependenceAnalysisWrapperPassPass(Registry);
   initializeDominanceFrontierWrapperPassPass(Registry);
   initializeDomViewerWrapperPassPass(Registry);

llvm/lib/Analysis/DXILResource.cpp

@@ -8,7 +8,9 @@
 
 #include "llvm/Analysis/DXILResource.h"
 #include "llvm/ADT/APInt.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/DiagnosticInfo.h"
@@ -19,6 +21,8 @@
 #include "llvm/IR/Module.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/FormatVariadic.h"
+#include <climits>
+#include <cstdint>
 
 #define DEBUG_TYPE "dxil-resource"
 
@@ -879,8 +883,126 @@ SmallVector<dxil::ResourceInfo *> DXILResourceMap::findByUse(const Value *Key) {
 
 //===----------------------------------------------------------------------===//
 
+void DXILResourceBindingInfo::populate(Module &M, DXILResourceTypeMap &DRTM) {
+  struct Binding {
+    ResourceClass RC;
+    uint32_t Space;
+    uint32_t LowerBound;
+    uint32_t UpperBound;
+    Binding(ResourceClass RC, uint32_t Space, uint32_t LowerBound,
+            uint32_t UpperBound)
+        : RC(RC), Space(Space), LowerBound(LowerBound), UpperBound(UpperBound) {
+    }
+  };
+  SmallVector<Binding> Bindings;
+
+  // collect all of the llvm.dx.resource.handlefrombinding calls;
+  // make a note if there is llvm.dx.resource.handlefromimplicitbinding
+  for (Function &F : M.functions()) {
+    if (!F.isDeclaration())
+      continue;
+
+    switch (F.getIntrinsicID()) {
+    default:
+      continue;
+    case Intrinsic::dx_resource_handlefrombinding: {
+      auto *HandleTy = cast<TargetExtType>(F.getReturnType());
+      ResourceTypeInfo &RTI = DRTM[HandleTy];
+
+      for (User *U : F.users())
+        if (CallInst *CI = dyn_cast<CallInst>(U)) {
+          uint32_t Space =
+              cast<ConstantInt>(CI->getArgOperand(0))->getZExtValue();
+          uint32_t LowerBound =
+              cast<ConstantInt>(CI->getArgOperand(1))->getZExtValue();
+          int32_t Size =
+              cast<ConstantInt>(CI->getArgOperand(2))->getZExtValue();
+
+          // negative size means unbounded resource array;
+          // upper bound register overflow should be detected in Sema
+          assert((Size < 0 || (unsigned)LowerBound + Size - 1 <= UINT32_MAX) &&
+                 "upper bound register overflow");
+          uint32_t UpperBound = Size < 0 ? UINT32_MAX : LowerBound + Size - 1;
+          Bindings.emplace_back(RTI.getResourceClass(), Space, LowerBound,
+                                UpperBound);
+        }
+      break;
+    }
+    case Intrinsic::dx_resource_handlefromimplicitbinding: {
+      ImplicitBinding = true;
+      break;
+    }
+    }
+  }
+
+  // sort all the collected bindings
+  llvm::stable_sort(Bindings, [](auto &LHS, auto &RHS) {
+    return std::tie(LHS.RC, LHS.Space, LHS.LowerBound) <
+           std::tie(RHS.RC, RHS.Space, RHS.LowerBound);
+  });
+
+  // remove duplicates
+  Binding *NewEnd = llvm::unique(Bindings, [](auto &LHS, auto &RHS) {
+    return std::tie(LHS.RC, LHS.Space, LHS.LowerBound, LHS.UpperBound) ==
+           std::tie(RHS.RC, RHS.Space, RHS.LowerBound, RHS.UpperBound);
+  });
+  if (NewEnd != Bindings.end())
+    Bindings.erase(NewEnd);
+
+  // Go over the sorted bindings and build up lists of free register ranges
+  // for each binding type and used spaces. Bindings are sorted by resource
+  // class, space, and lower bound register slot.
+  BindingSpaces *BS = &SRVSpaces;
+  for (unsigned I = 0, E = Bindings.size(); I != E; ++I) {
+    Binding &B = Bindings[I];
+
+    if (BS->RC != B.RC)
+      // move to the next resource class spaces
+      BS = &getBindingSpaces(B.RC);
+
+    RegisterSpace *S = BS->Spaces.empty() ? &BS->Spaces.emplace_back(B.Space)
+                                          : &BS->Spaces.back();
+    assert(S->Space <= B.Space && "bindings not sorted correctly?");
+    if (B.Space != S->Space)
+      // add new space
+      S = &BS->Spaces.emplace_back(B.Space);
+
+    // the space is full - set flag to report overlapping binding later
+    if (S->FreeRanges.empty()) {
+      OverlappingBinding = true;
+      continue;
+    }
+
+    // adjust the last free range lower bound, split it in two, or remove it
+    BindingRange &LastFreeRange = S->FreeRanges.back();
+    assert(LastFreeRange.UpperBound == UINT32_MAX);
+    if (LastFreeRange.LowerBound == B.LowerBound) {
+      if (B.UpperBound < UINT32_MAX)
+        LastFreeRange.LowerBound = B.UpperBound + 1;
+      else
+        S->FreeRanges.pop_back();
+    } else if (LastFreeRange.LowerBound < B.LowerBound) {
+      LastFreeRange.UpperBound = B.LowerBound - 1;
+      if (B.UpperBound < UINT32_MAX)
+        S->FreeRanges.emplace_back(B.UpperBound + 1, UINT32_MAX);
+    } else {
+      // FIXME: This only detects overlapping bindings that are not an exact
+      // match (llvm/llvm-project#110723)
+      OverlappingBinding = true;
+      if (B.UpperBound < UINT32_MAX)
+        LastFreeRange.LowerBound =
+            std::max(LastFreeRange.LowerBound, B.UpperBound + 1);
+      else
+        S->FreeRanges.pop_back();
+    }
+  }
+}
+
+//===----------------------------------------------------------------------===//
+
 AnalysisKey DXILResourceTypeAnalysis::Key;
 AnalysisKey DXILResourceAnalysis::Key;
+AnalysisKey DXILResourceBindingAnalysis::Key;
 
 DXILResourceMap DXILResourceAnalysis::run(Module &M,
                                           ModuleAnalysisManager &AM) {
@@ -890,6 +1012,14 @@ DXILResourceMap DXILResourceAnalysis::run(Module &M,
   return Data;
 }
 
+DXILResourceBindingInfo
+DXILResourceBindingAnalysis::run(Module &M, ModuleAnalysisManager &AM) {
+  DXILResourceBindingInfo Data;
+  DXILResourceTypeMap &DRTM = AM.getResult<DXILResourceTypeAnalysis>(M);
+  Data.populate(M, DRTM);
+  return Data;
+}
+
 PreservedAnalyses DXILResourcePrinterPass::run(Module &M,
                                                ModuleAnalysisManager &AM) {
   DXILResourceMap &DRM = AM.getResult<DXILResourceAnalysis>(M);
@@ -952,3 +1082,33 @@ char DXILResourceWrapperPass::ID = 0;
 ModulePass *llvm::createDXILResourceWrapperPassPass() {
   return new DXILResourceWrapperPass();
 }
+
+DXILResourceBindingWrapperPass::DXILResourceBindingWrapperPass()
+    : ModulePass(ID) {}
+
+DXILResourceBindingWrapperPass::~DXILResourceBindingWrapperPass() = default;
+
+void DXILResourceBindingWrapperPass::getAnalysisUsage(AnalysisUsage &AU) const {
+  AU.addRequiredTransitive<DXILResourceTypeWrapperPass>();
+  AU.setPreservesAll();
+}
+
+bool DXILResourceBindingWrapperPass::runOnModule(Module &M) {
+  BindingInfo.reset(new DXILResourceBindingInfo());
+
+  DXILResourceTypeMap &DRTM =
+      getAnalysis<DXILResourceTypeWrapperPass>().getResourceTypeMap();
+  BindingInfo->populate(M, DRTM);
+
+  return false;
+}
+
+void DXILResourceBindingWrapperPass::releaseMemory() { BindingInfo.reset(); }
+
+INITIALIZE_PASS(DXILResourceBindingWrapperPass, "dxil-resource-binding",
+                "DXIL Resource Binding Analysis", false, true)
+char DXILResourceBindingWrapperPass::ID = 0;
+
+ModulePass *llvm::createDXILResourceBindingWrapperPassPass() {
+  return new DXILResourceWrapperPass();
+}

llvm/lib/Passes/PassRegistry.def

@@ -24,6 +24,7 @@ MODULE_ANALYSIS("ctx-prof-analysis", CtxProfAnalysis())
 MODULE_ANALYSIS("dxil-metadata", DXILMetadataAnalysis())
 MODULE_ANALYSIS("dxil-resources", DXILResourceAnalysis())
 MODULE_ANALYSIS("dxil-resource-type", DXILResourceTypeAnalysis())
+MODULE_ANALYSIS("dxil-resource-bindings", DXILResourceBindingAnalysis())
 MODULE_ANALYSIS("inline-advisor", InlineAdvisorAnalysis())
 MODULE_ANALYSIS("ir-similarity", IRSimilarityAnalysis())
 MODULE_ANALYSIS("last-run-tracking", LastRunTrackingAnalysis())

llvm/unittests/Target/DirectX/CMakeLists.txt

@@ -22,4 +22,5 @@ add_llvm_target_unittest(DirectXTests
   PointerTypeAnalysisTests.cpp
   UniqueResourceFromUseTests.cpp
   RegisterCostTests.cpp
+  ResourceBindingAnalysisTests.cpp
   )