[HLSL] Anaylze updateCounter usage

Author: V-FEXrt

llvm/include/llvm/Analysis/DXILResource.h

@@ -463,7 +463,7 @@ class DXILResourceMap {
   /// ambiguous so multiple creation instructions may be returned. The resulting
   /// ResourceInfo can be used to depuplicate unique handles that
   /// reference the same resource
-  SmallVector<dxil::ResourceInfo> findByUse(const Value *Key) const;
+  SmallVector<dxil::ResourceInfo *> findByUse(const Value *Key);
 
   const_iterator find(const CallInst *Key) const {
     auto Pos = CallMap.find(Key);

llvm/lib/Analysis/DXILResource.cpp

@@ -697,6 +697,9 @@ bool DXILResourceTypeMap::invalidate(Module &M, const PreservedAnalyses &PA,
 }
 
 //===----------------------------------------------------------------------===//
+static bool isUpdateCounterIntrinsic(Function &F) {
+  return F.getIntrinsicID() == Intrinsic::dx_resource_updatecounter;
+}
 
 void DXILResourceMap::populate(Module &M, DXILResourceTypeMap &DRTM) {
   SmallVector<std::tuple<CallInst *, ResourceInfo, ResourceTypeInfo>> CIToInfos;
@@ -775,6 +778,42 @@ void DXILResourceMap::populate(Module &M, DXILResourceTypeMap &DRTM) {
     // Adjust the resource binding to use the next ID.
     RI.setBindingID(NextID++);
   }
+
+  for (Function &F : M.functions()) {
+    if (!isUpdateCounterIntrinsic(F))
+      continue;
+
+    LLVM_DEBUG(dbgs() << "Update Counter Function: " << F.getName() << "\n");
+
+    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 = cast<ConstantInt>(CountArg);
+      int64_t CountLiteral = CountValue->getSExtValue();
+
+      // 0 is an unknown direction and shouldn't result in an insert
+      if (CountLiteral == 0)
+        continue;
+
+      ResourceCounterDirection Direction = ResourceCounterDirection::Decrement;
+      if (CountLiteral > 0)
+        Direction = ResourceCounterDirection::Increment;
+
+      // Collect all potential creation points for the handle arg
+      Value *HandleArg = CI->getArgOperand(0);
+      SmallVector<ResourceInfo *> RBInfos = findByUse(HandleArg);
+      for (ResourceInfo *RBInfo : RBInfos) {
+        if (RBInfo->CounterDirection == ResourceCounterDirection::Unknown ||
+            RBInfo->CounterDirection == Direction)
+          RBInfo->CounterDirection = Direction;
+        else
+          RBInfo->CounterDirection = ResourceCounterDirection::Invalid;
+      }
+    }
+  }
 }
 
 void DXILResourceMap::print(raw_ostream &OS, DXILResourceTypeMap &DRTM,
@@ -793,10 +832,9 @@ void DXILResourceMap::print(raw_ostream &OS, DXILResourceTypeMap &DRTM,
   }
 }
 
-SmallVector<dxil::ResourceInfo>
-DXILResourceMap::findByUse(const Value *Key) const {
+SmallVector<dxil::ResourceInfo *> DXILResourceMap::findByUse(const Value *Key) {
   if (const PHINode *Phi = dyn_cast<PHINode>(Key)) {
-    SmallVector<dxil::ResourceInfo> Children;
+    SmallVector<dxil::ResourceInfo *> Children;
     for (const Value *V : Phi->operands()) {
       Children.append(findByUse(V));
     }
@@ -810,9 +848,10 @@ DXILResourceMap::findByUse(const Value *Key) const {
   switch (CI->getIntrinsicID()) {
   // Found the create, return the binding
   case Intrinsic::dx_resource_handlefrombinding: {
-    const auto *It = find(CI);
+    auto Pos = CallMap.find(CI);
+    ResourceInfo *It = &Infos[Pos->second];
     assert(It != Infos.end() && "HandleFromBinding must be in resource map");
-    return {*It};
+    return {It};
   }
   default:
     break;
@@ -821,7 +860,7 @@ DXILResourceMap::findByUse(const Value *Key) const {
   // Check if any of the parameters are the resource we are following. If so
   // keep searching. If none of them are return an empty list
   const Type *UseType = CI->getType();
-  SmallVector<dxil::ResourceInfo> Children;
+  SmallVector<dxil::ResourceInfo *> Children;
   for (const Value *V : CI->args()) {
     if (V->getType() != UseType)
       continue;

llvm/test/Analysis/DXILResource/buffer-frombinding.ll

@@ -69,14 +69,15 @@ define void @test_typedbuffer() {
   %uav1 = call target("dx.TypedBuffer", <4 x float>, 1, 0, 0)
               @llvm.dx.resource.handlefrombinding.tdx.TypedBuffer_f32_1_0(
                   i32 3, i32 5, i32 1, i32 0, i1 false)
+  call i32 @llvm.dx.resource.updatecounter(target("dx.TypedBuffer", <4 x float>, 1, 0, 0) %uav1, i8 -1)
   ; CHECK: Resource [[UAV1:[0-9]+]]:
   ; CHECK:   Binding:
   ; CHECK:     Record ID: 1
   ; CHECK:     Space: 3
   ; CHECK:     Lower Bound: 5
   ; CHECK:     Size: 1
   ; CHECK:   Globally Coherent: 0
-  ; CHECK:   Counter Direction: Unknown
+  ; CHECK:   Counter Direction: Decrement
   ; CHECK:   Class: UAV
   ; CHECK:   Kind: TypedBuffer
   ; CHECK:   IsROV: 0
@@ -92,14 +93,15 @@ define void @test_typedbuffer() {
   %uav2_2 = call target("dx.TypedBuffer", <4 x float>, 1, 0, 0)
               @llvm.dx.resource.handlefrombinding.tdx.TypedBuffer_f32_1_0(
                   i32 4, i32 0, i32 10, i32 5, i1 false)
+  call i32 @llvm.dx.resource.updatecounter(target("dx.TypedBuffer", <4 x float>, 1, 0, 0) %uav2_2, i8 1)
   ; CHECK: Resource [[UAV2:[0-9]+]]:
   ; CHECK:   Binding:
   ; CHECK:     Record ID: 2
   ; CHECK:     Space: 4
   ; CHECK:     Lower Bound: 0
   ; CHECK:     Size: 10
   ; CHECK:   Globally Coherent: 0
-  ; CHECK:   Counter Direction: Unknown
+  ; CHECK:   Counter Direction: Increment
   ; CHECK:   Class: UAV
   ; CHECK:   Kind: TypedBuffer
   ; CHECK:   IsROV: 0