[HLSL] Implement DXILResourceBindingAnalysis

DXILResourceBindingAnalysis analyses all 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 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,108 @@ class DXILResourceWrapperPass : public ModulePass {
 
 ModulePass *createDXILResourceWrapperPassPass();
 
+//===----------------------------------------------------------------------===//
+
+// DXILResourceBindingsInfo 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 DXILResourceBindingsInfo {
+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 ResClass;
+    llvm::SmallVector<RegisterSpace> Spaces;
+    BindingSpaces(dxil::ResourceClass ResClass) : ResClass(ResClass) {
+      // initialize space0
+      Spaces.emplace_back(0);
+    }
+  };
+
+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:
+  DXILResourceBindingsInfo()
+      : SRVSpaces(dxil::ResourceClass::SRV),
+        UAVSpaces(dxil::ResourceClass::UAV),
+        CBufferSpaces(dxil::ResourceClass::CBuffer),
+        SamplerSpaces(dxil::ResourceClass::Sampler), ImplicitBinding(false),
+        OverlappingBinding(false) {}
+
+  bool containsImplicitBinding() const { return ImplicitBinding; }
+  bool containsOverlappingBinding() 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;
+};
+
+class DXILResourceBindingAnalysis
+    : public AnalysisInfoMixin<DXILResourceBindingAnalysis> {
+  friend AnalysisInfoMixin<DXILResourceBindingAnalysis>;
+
+  static AnalysisKey Key;
+
+public:
+  using Result = DXILResourceBindingsInfo;
+
+  DXILResourceBindingsInfo run(Module &M, ModuleAnalysisManager &AM);
+};
+
 } // 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],
+          [IntrNoMem]>;
+
 def int_dx_resource_getpointer
     : DefaultAttrsIntrinsic<[llvm_anyptr_ty], [llvm_any_ty, llvm_i32_ty],
                             [IntrNoMem]>;

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,121 @@ SmallVector<dxil::ResourceInfo *> DXILResourceMap::findByUse(const Value *Key) {
 
 //===----------------------------------------------------------------------===//
 
+void DXILResourceBindingsInfo::populate(Module &M, DXILResourceTypeMap &DRTM) {
+  struct Binding {
+    ResourceClass ResClass;
+    uint32_t Space;
+    uint32_t LowerBound;
+    uint32_t UpperBound;
+    Binding(ResourceClass RC, uint32_t Sp, uint32_t LB, uint32_t UB)
+        : ResClass(RC), Space(Sp), LowerBound(LB), UpperBound(UB) {}
+  };
+  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: {
+      if (!F.user_empty())
+        ImplicitBinding = true;
+      break;
+    }
+    }
+  }
+
+  // sort all the collected bindings
+  llvm::stable_sort(Bindings, [](auto &LHS, auto &RHS) {
+    return std::tie(LHS.ResClass, LHS.Space, LHS.LowerBound) <
+           std::tie(RHS.ResClass, RHS.Space, RHS.LowerBound);
+  });
+
+  // remove duplicates
+  llvm::unique(Bindings, [](auto &LHS, auto &RHS) {
+    return std::tie(LHS.ResClass, LHS.Space, LHS.LowerBound, LHS.UpperBound) ==
+           std::tie(RHS.ResClass, RHS.Space, RHS.LowerBound, RHS.UpperBound);
+  });
+
+  // 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->ResClass != B.ResClass)
+      // move to the next resource class spaces
+      BS = &getBindingSpaces(B.ResClass);
+
+    RegisterSpace *S = &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 {
+      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 +1007,14 @@ DXILResourceMap DXILResourceAnalysis::run(Module &M,
   return Data;
 }
 
+DXILResourceBindingsInfo
+DXILResourceBindingAnalysis::run(Module &M, ModuleAnalysisManager &AM) {
+  DXILResourceBindingsInfo 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);

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
   )