[HLSL] Analyze update counter usage

Author: V-FEXrt

llvm/include/llvm/Analysis/DXILResource.h

@@ -18,6 +18,10 @@
 #include "llvm/Support/Alignment.h"
 #include "llvm/Support/DXILABI.h"
 
+#include <algorithm>
+#include <cstdint>
+#include <unordered_map>
+
 namespace llvm {
 class CallInst;
 class DataLayout;
@@ -407,6 +411,69 @@ class DXILResourceTypeMap {
   }
 };
 
+enum ResourceCounterDirection {
+  Increment,
+  Decrement,
+  Unknown,
+};
+
+class DXILResourceCounterDirectionMap {
+  std::vector<std::pair<dxil::ResourceBindingInfo, ResourceCounterDirection>> CounterDirections;
+
+public:
+  bool invalidate(Module &M, const PreservedAnalyses &PA,
+                  ModuleAnalysisManager::Invalidator &Inv);
+
+   void populate(Module &M, ModuleAnalysisManager &AM);
+
+   ResourceCounterDirection operator[](const dxil::ResourceBindingInfo &Info) const {
+    auto Lower = std::lower_bound(CounterDirections.begin(), CounterDirections.end(), std::pair{Info, ResourceCounterDirection::Unknown},  [](auto lhs, auto rhs){
+                                  return lhs.first  < rhs.first;
+    });
+
+    if (Lower == CounterDirections.end()) {
+      return ResourceCounterDirection::Unknown;
+    }
+
+    if (Lower->first != Info) {
+      return ResourceCounterDirection::Unknown;
+    }
+
+    return Lower->second;
+   }
+};
+
+class DXILResourceCounterDirectionAnalysis
+    : public AnalysisInfoMixin<DXILResourceCounterDirectionAnalysis> {
+  friend AnalysisInfoMixin<DXILResourceCounterDirectionAnalysis>;
+
+  static AnalysisKey Key;
+
+public:
+  using Result = DXILResourceCounterDirectionMap;
+
+  DXILResourceCounterDirectionMap run(Module &M, ModuleAnalysisManager &AM) {
+    DXILResourceCounterDirectionMap DRCDM{};
+    DRCDM.populate(M, AM);
+    return DRCDM;
+  }
+};
+
+class DXILResourceCounterDirectionWrapperPass : public ImmutablePass {
+  DXILResourceCounterDirectionMap DRCDM;
+
+  virtual void anchor();
+
+public:
+  static char ID;
+  DXILResourceCounterDirectionWrapperPass();
+
+  DXILResourceCounterDirectionMap &getResourceCounterDirectionMap() { return DRCDM; }
+  const DXILResourceCounterDirectionMap &getResourceCounterDirectionMap() const { return DRCDM; }
+};
+
+ModulePass *createDXILResourceCounterDirectionWrapperPassPass();
+
 class DXILResourceTypeAnalysis
     : public AnalysisInfoMixin<DXILResourceTypeAnalysis> {
   friend AnalysisInfoMixin<DXILResourceTypeAnalysis>;

llvm/include/llvm/InitializePasses.h

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

llvm/lib/Analysis/DXILResource.cpp

@@ -19,6 +19,7 @@
 #include "llvm/IR/Module.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/FormatVariadic.h"
+#include <algorithm>
 
 #define DEBUG_TYPE "dxil-resource"
 
@@ -818,8 +819,88 @@ DXILBindingMap::findByUse(const Value *Key) const {
 
 //===----------------------------------------------------------------------===//
 
+static bool isUpdateCounterIntrinsic(Function &F) {
+  return F.getIntrinsicID() == Intrinsic::dx_resource_updatecounter;
+}
+
+void DXILResourceCounterDirectionMap::populate(Module &M, ModuleAnalysisManager &AM) {
+    DXILBindingMap &DBM = AM.getResult<DXILResourceBindingAnalysis>(M);
+    CounterDirections.clear();
+
+    for (Function &F : M.functions()) {
+      if (!isUpdateCounterIntrinsic(F))
+        continue;
+
+      for (const User *U : F.users()) {
+        const CallInst *CI = dyn_cast<CallInst>(U);
+        assert(CI && "Users of dx_resource_updateCounter must be call instrs");
+
+        // Determine if the use is an increment or decrement
+        Value *CountArg = CI->getArgOperand(1);
+        ConstantInt *CountValue = dyn_cast<ConstantInt>(CountArg);
+        int64_t CountLiteral = CountValue->getSExtValue();
+
+        ResourceCounterDirection Direction = ResourceCounterDirection::Unknown;
+        if (CountLiteral > 0) {
+          Direction = ResourceCounterDirection::Increment;
+        }
+        if (CountLiteral < 0) {
+          Direction = ResourceCounterDirection::Decrement;
+        }
+
+
+        // Collect all potential creation points for the handle arg
+        Value *HandleArg = CI->getArgOperand(0);
+        SmallVector<dxil::ResourceBindingInfo> RBInfos = DBM.findByUse(HandleArg);
+        for(const dxil::ResourceBindingInfo RBInfo : RBInfos) {
+          CounterDirections.emplace_back(RBInfo, Direction);
+        }
+      }
+    }
+
+    // An entry that is not in the map is considered unknown so its wasted
+    // overhead and increased complexity to keep it so it should be removed.
+    const auto RemoveEnd = std::remove_if(CounterDirections.begin(), CounterDirections.end(), [](const auto& Item) { 
+         return Item.second == ResourceCounterDirection::Unknown;
+    });
+
+    // Order for fast lookup
+    std::sort(CounterDirections.begin(), RemoveEnd);
+
+    // Remove the duplicate entries. Since direction is considered for equality
+    // a unique resource with more than one direction will not be deduped.
+    const auto UniqueEnd = std::unique(CounterDirections.begin(), RemoveEnd);
+
+    // Actually erase the items invalidated by remove_if + unique
+    CounterDirections.erase(UniqueEnd, CounterDirections.end());
+
+    // If any duplicate entries still exist at this point then it must be a
+    // resource that was both incremented and decremented which is not allowed.
+    const auto DuplicateEntry = std::adjacent_find(CounterDirections.begin(), CounterDirections.end(), [](const auto& LHS, const auto& RHS){
+      return LHS.first == RHS.first;
+    });
+    if (DuplicateEntry == CounterDirections.end())
+        return;
+
+    // TODO: Emit an error message
+    assert(CounterDirections.size() == 1 && "dups found");
+}
+
+bool DXILResourceCounterDirectionMap::invalidate(Module &M, const PreservedAnalyses &PA,
+                                     ModuleAnalysisManager::Invalidator &Inv) {
+  // Passes that introduce resource types must explicitly invalidate this pass.
+  // auto PAC = PA.getChecker<DXILResourceTypeAnalysis>();
+  // return !PAC.preservedWhenStateless();
+  return false;
+}
+
+void DXILResourceCounterDirectionWrapperPass::anchor() {}
+
+//===----------------------------------------------------------------------===//
+
 AnalysisKey DXILResourceTypeAnalysis::Key;
 AnalysisKey DXILResourceBindingAnalysis::Key;
+AnalysisKey DXILResourceCounterDirectionAnalysis::Key;
 
 DXILBindingMap DXILResourceBindingAnalysis::run(Module &M,
                                                 ModuleAnalysisManager &AM) {
@@ -838,6 +919,13 @@ DXILResourceBindingPrinterPass::run(Module &M, ModuleAnalysisManager &AM) {
   return PreservedAnalyses::all();
 }
 
+INITIALIZE_PASS(DXILResourceCounterDirectionWrapperPass, "dxil-resource-counter",
+                "DXIL Resource Counter Analysis", false, true)
+
+DXILResourceCounterDirectionWrapperPass::DXILResourceCounterDirectionWrapperPass() : ImmutablePass(ID) {
+  initializeDXILResourceTypeWrapperPassPass(*PassRegistry::getPassRegistry());
+}
+
 void DXILResourceTypeWrapperPass::anchor() {}
 
 DXILResourceTypeWrapperPass::DXILResourceTypeWrapperPass() : ImmutablePass(ID) {
@@ -847,11 +935,16 @@ DXILResourceTypeWrapperPass::DXILResourceTypeWrapperPass() : ImmutablePass(ID) {
 INITIALIZE_PASS(DXILResourceTypeWrapperPass, "dxil-resource-type",
                 "DXIL Resource Type Analysis", false, true)
 char DXILResourceTypeWrapperPass::ID = 0;
+char DXILResourceCounterDirectionWrapperPass::ID = 0;
 
 ModulePass *llvm::createDXILResourceTypeWrapperPassPass() {
   return new DXILResourceTypeWrapperPass();
 }
 
+ModulePass *llvm::createDXILResourceCounterDirectionWrapperPassPass() {
+  return new DXILResourceCounterDirectionWrapperPass();
+}
+
 DXILResourceBindingWrapperPass::DXILResourceBindingWrapperPass()
     : ModulePass(ID) {
   initializeDXILResourceBindingWrapperPassPass(

llvm/unittests/Target/DirectX/UniqueResourceFromUseTests.cpp

@@ -35,6 +35,7 @@ class UniqueResourceFromUseTest : public testing::Test {
     PB->registerModuleAnalyses(*MAM);
     MAM->registerPass([&] { return DXILResourceTypeAnalysis(); });
     MAM->registerPass([&] { return DXILResourceBindingAnalysis(); });
+    MAM->registerPass([&] { return DXILResourceCounterDirectionAnalysis(); });
   }
 
   virtual void TearDown() {
@@ -280,4 +281,47 @@ declare target("dx.RawBuffer", float, 1, 0) @ind.func(target("dx.RawBuffer", flo
   }
 }
 
+TEST_F(UniqueResourceFromUseTest, TestResourceCounter) {
+  StringRef Assembly = R"(
+define void @main() {
+entry:
+  %handle = call target("dx.RawBuffer", float, 1, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_f32_1_0t(i32 1, i32 2, i32 3, i32 4, i1 false)
+  call i32 @llvm.dx.resource.updatecounter.tdx.RawBuffer_f32_1_0t(target("dx.RawBuffer", float, 1, 0) %handle, i8 -1)
+  call i32 @llvm.dx.resource.updatecounter.tdx.RawBuffer_f32_1_0t(target("dx.RawBuffer", float, 1, 0) %handle, i8 -1)
+  call i32 @llvm.dx.resource.updatecounter.tdx.RawBuffer_f32_1_0t(target("dx.RawBuffer", float, 1, 0) %handle, i8 -1)
+  call i32 @llvm.dx.resource.updatecounter.tdx.RawBuffer_f32_1_0t(target("dx.RawBuffer", float, 1, 0) %handle, i8 -1)
+  call i32 @llvm.dx.resource.updatecounter.tdx.RawBuffer_f32_1_0t(target("dx.RawBuffer", float, 1, 0) %handle, i8 -1)
+  call i32 @llvm.dx.resource.updatecounter.tdx.RawBuffer_f32_1_0t(target("dx.RawBuffer", float, 1, 0) %handle, i8 -1)
+  call void @a.func(target("dx.RawBuffer", float, 1, 0) %handle)
+  ret void
+}
+
+declare target("dx.RawBuffer", float, 1, 0) @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_f32_1_0t(i32, i32, i32, i32, i1)
+declare i32 @llvm.dx.resource.updatecounter.tdx.RawBuffer_f32_1_0t(target("dx.RawBuffer", float, 1, 0), i8)
+declare void @a.func(target("dx.RawBuffer", float, 1, 0) %handle)
+  )";
+
+  LLVMContext Context;
+  SMDiagnostic Error;
+  auto M = parseAssemblyString(Assembly, Error, Context);
+  ASSERT_TRUE(M) << "Bad assembly?";
+
+  const DXILBindingMap &DBM = MAM->getResult<DXILResourceBindingAnalysis>(*M);
+  const DXILResourceCounterDirectionMap &DCDM = MAM->getResult<DXILResourceCounterDirectionAnalysis>(*M);
+
+  for (const Function &F : M->functions()) {
+    if (F.getName() != "a.func") {
+      continue;
+    }
+
+    for (const User *U : F.users()) {
+      const CallInst *CI = cast<CallInst>(U);
+      const Value *Handle = CI->getArgOperand(0);
+      const auto Bindings = DBM.findByUse(Handle);
+      ASSERT_EQ(Bindings.size(), 1u);
+      ASSERT_EQ(DCDM[Bindings.front()], ResourceCounterDirection::Decrement);
+    }
+  }
+}
+
 } // namespace