[DirectX] Implement DXILResourceBindingAnalysis

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
   )