[HLSL] Analyze update counter usage

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,77 @@ 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

@@ -9,6 +9,7 @@
 #include "llvm/Analysis/DXILResource.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/DiagnosticInfo.h"
@@ -19,6 +20,7 @@
 #include "llvm/IR/Module.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/FormatVariadic.h"
+#include <algorithm>
 
 #define DEBUG_TYPE "dxil-resource"
 
@@ -818,8 +820,92 @@ 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 an entry explicitly marked
+  // unknown
+  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;
+
+  StringRef Message = "RWStructuredBuffers may increment or decrement their "
+                      "counters, but not both.";
+  M.getContext().diagnose(DiagnosticInfoGeneric(Message));
+}
+
+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 +924,16 @@ 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 +943,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/lib/Passes/PassRegistry.def

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

llvm/lib/Target/DirectX/DXILOpLowering.cpp

@@ -816,6 +816,8 @@ class OpLowerer {
 PreservedAnalyses DXILOpLowering::run(Module &M, ModuleAnalysisManager &MAM) {
   DXILBindingMap &DBM = MAM.getResult<DXILResourceBindingAnalysis>(M);
   DXILResourceTypeMap &DRTM = MAM.getResult<DXILResourceTypeAnalysis>(M);
+  DXILResourceCounterDirectionMap &DRCDM =
+      MAM.getResult<DXILResourceCounterDirectionAnalysis>(M);
 
   bool MadeChanges = OpLowerer(M, DBM, DRTM).lowerIntrinsics();
   if (!MadeChanges)

llvm/test/CodeGen/DirectX/resource_counter_error.ll

@@ -0,0 +1,13 @@
+; RUN: not opt -S -passes='dxil-op-lower' -mtriple=dxil-pc-shadermodel6.3-library %s 2>&1 | FileCheck %s
+; CHECK:  RWStructuredBuffers may increment or decrement their counters, but not both.
+
+define void @inc_and_dec() {
+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)
+  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)