[llvm] [HLSL] Add UsedByAtomic64 shader flag (PR #211691)

Joshua Batista via llvm-commits llvm-commits at lists.llvm.org
Fri Jul 24 14:48:51 PDT 2026


https://github.com/bob80905 updated https://github.com/llvm/llvm-project/pull/211691

>From 1f1fd9ab6db2c5a0d072de9c5994d35d1b8eb5f8 Mon Sep 17 00:00:00 2001
From: Joshua Batista <jbatista at microsoft.com>
Date: Thu, 23 Jul 2026 15:50:51 -0700
Subject: [PATCH 1/4] first attempt

---
 llvm/include/llvm/Analysis/DXILResource.h     |  3 +
 llvm/lib/Analysis/DXILResource.cpp            | 51 +++++++++++++++++
 .../lib/Target/DirectX/DXContainerGlobals.cpp |  4 +-
 .../Analysis/DXILResource/has-atomic64-use.ll | 56 +++++++++++++++++++
 .../PSVResources-UsedByAtomic64.ll            | 47 ++++++++++++++++
 5 files changed, 158 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/Analysis/DXILResource/has-atomic64-use.ll
 create mode 100644 llvm/test/CodeGen/DirectX/ContainerData/PSVResources-UsedByAtomic64.ll

diff --git a/llvm/include/llvm/Analysis/DXILResource.h b/llvm/include/llvm/Analysis/DXILResource.h
index 0a637d1a4bcc3..d0c6ea44b6add 100644
--- a/llvm/include/llvm/Analysis/DXILResource.h
+++ b/llvm/include/llvm/Analysis/DXILResource.h
@@ -410,6 +410,7 @@ class ResourceInfo {
 public:
   bool GloballyCoherent = false;
   ResourceCounterDirection CounterDirection = ResourceCounterDirection::Unknown;
+  bool HasAtomic64Use = false;
 
   ResourceInfo(uint32_t RecordID, uint32_t Space, uint32_t LowerBound,
                uint32_t Size, TargetExtType *HandleTy, StringRef Name = "",
@@ -516,6 +517,8 @@ class DXILResourceMap {
   void populateResourceInfos(Module &M, DXILResourceTypeMap &DRTM);
   /// Analyze and populate the directions of the resource counters.
   void populateCounterDirections(Module &M);
+  /// Detect 64-bit atomic uses of resources and set `HasAtomic64Use`.
+  void populateAtomicUses(Module &M);
 
   /// Resolves a resource handle into a vector of ResourceInfos that
   /// represent the possible unique creations of the handle. Certain cases are
diff --git a/llvm/lib/Analysis/DXILResource.cpp b/llvm/lib/Analysis/DXILResource.cpp
index a427a261d82d3..4baeb306ee0f1 100644
--- a/llvm/lib/Analysis/DXILResource.cpp
+++ b/llvm/lib/Analysis/DXILResource.cpp
@@ -14,6 +14,7 @@
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/DiagnosticInfo.h"
+#include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/IntrinsicsDirectX.h"
@@ -804,6 +805,7 @@ void ResourceInfo::print(raw_ostream &OS, dxil::ResourceTypeInfo &RTI,
      << "    Size: " << Binding.Size << "\n";
 
   OS << "  Globally Coherent: " << GloballyCoherent << "\n";
+  OS << "  Has Atomic64 Use: " << HasAtomic64Use << "\n";
   OS << "  Counter Direction: ";
 
   switch (CounterDirection) {
@@ -984,9 +986,58 @@ void DXILResourceMap::populateCounterDirections(Module &M) {
   }
 }
 
+void DXILResourceMap::populateAtomicUses(Module &M) {
+  auto FindResourceHandle = [](Value *Ptr) -> Value * {
+    Ptr = Ptr->stripPointerCasts();
+    while (auto *GEP = dyn_cast<GetElementPtrInst>(Ptr))
+      Ptr = GEP->getPointerOperand()->stripPointerCasts();
+    auto *II = dyn_cast<IntrinsicInst>(Ptr);
+    if (II && II->getIntrinsicID() == Intrinsic::dx_resource_getpointer)
+      return II->getArgOperand(0);
+    return nullptr;
+  };
+
+  auto MarkHasAtomic64UseFromHandle = [this](Value *Handle) {
+    if (!Handle)
+      return;
+    for (ResourceInfo *RI : findByUse(Handle))
+      RI->HasAtomic64Use = true;
+  };
+
+  // Detect on `atomicrmw`/`cmpxchg` before `DXILResourceAccess` lowers them.
+  for (Function &F : M.functions()) {
+    for (Instruction &I : instructions(F)) {
+      if (auto *AI = dyn_cast<AtomicRMWInst>(&I)) {
+        if (!AI->getValOperand()->getType()->isIntegerTy(64))
+          continue;
+        MarkHasAtomic64UseFromHandle(
+            FindResourceHandle(AI->getPointerOperand()));
+      } else if (auto *CX = dyn_cast<AtomicCmpXchgInst>(&I)) {
+        if (!CX->getNewValOperand()->getType()->isIntegerTy(64))
+          continue;
+        MarkHasAtomic64UseFromHandle(
+            FindResourceHandle(CX->getPointerOperand()));
+      }
+    }
+  }
+
+  // Also handle the post-lowering `llvm.dx.resource.atomic.binop` form.
+  for (Function &F : M.functions()) {
+    if (F.getIntrinsicID() != Intrinsic::dx_resource_atomic_binop)
+      continue;
+    for (User *U : F.users()) {
+      auto *CI = dyn_cast<CallInst>(U);
+      if (!CI || !CI->getType()->isIntegerTy(64))
+        continue;
+      MarkHasAtomic64UseFromHandle(CI->getArgOperand(0));
+    }
+  }
+}
+
 void DXILResourceMap::populate(Module &M, DXILResourceTypeMap &DRTM) {
   populateResourceInfos(M, DRTM);
   populateCounterDirections(M);
+  populateAtomicUses(M);
 }
 
 void DXILResourceMap::print(raw_ostream &OS, DXILResourceTypeMap &DRTM,
diff --git a/llvm/lib/Target/DirectX/DXContainerGlobals.cpp b/llvm/lib/Target/DirectX/DXContainerGlobals.cpp
index 1753b3e3f3a1a..c4cc76d457367 100644
--- a/llvm/lib/Target/DirectX/DXContainerGlobals.cpp
+++ b/llvm/lib/Target/DirectX/DXContainerGlobals.cpp
@@ -325,9 +325,7 @@ void DXContainerGlobals::addResourcesForPSV(Module &M, PSVRuntimeInfo &PSV) {
       ResType = dxbc::PSV::ResourceType::UAVRaw;
 
     dxbc::PSV::ResourceFlags Flags;
-    // TODO: Add support for dxbc::PSV::ResourceFlag::UsedByAtomic64, tracking
-    // with https://github.com/llvm/llvm-project/issues/104392
-    Flags.Flags = 0u;
+    Flags.Bits.UsedByAtomic64 = RI.HasAtomic64Use;
 
     PSV.Resources.push_back(
         MakeBinding(Binding, ResType, TypeInfo.getResourceKind(), Flags));
diff --git a/llvm/test/Analysis/DXILResource/has-atomic64-use.ll b/llvm/test/Analysis/DXILResource/has-atomic64-use.ll
new file mode 100644
index 0000000000000..1d32da7fb8d57
--- /dev/null
+++ b/llvm/test/Analysis/DXILResource/has-atomic64-use.ll
@@ -0,0 +1,56 @@
+; RUN: opt -S -disable-output -passes="print<dxil-resources>" < %s 2>&1 | FileCheck %s
+
+; Verifies that the DXILResourceMap analysis detects 64-bit atomic uses on
+; UAV resources and sets `HasAtomic64Use` on the corresponding ResourceInfo.
+;
+; * `Atomic64UAV`  is used by a 64-bit `atomicrmw`  -> Has Atomic64 Use: 1
+; * `Atomic32UAV`  is used by a 32-bit `atomicrmw`  -> Has Atomic64 Use: 0
+; * `PlainUAV`     is only stored to (no atomics)   -> Has Atomic64 Use: 0
+
+target triple = "dxil-pc-shadermodel6.6-compute"
+
+define void @main() {
+  ; RWByteAddressBuffer Atomic64UAV : register(u0);
+  %h64 = call target("dx.RawBuffer", i8, 1, 0)
+      @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_i8_1_0t(
+          i32 0, i32 0, i32 1, i32 0, ptr null)
+  ; CHECK:      Binding:
+  ; CHECK:        Record ID: 0
+  ; CHECK:        Space: 0
+  ; CHECK:        Lower Bound: 0
+  ; CHECK:        Size: 1
+  ; CHECK:      Has Atomic64 Use: 1
+  %p64 = call ptr @llvm.dx.resource.getpointer(
+      target("dx.RawBuffer", i8, 1, 0) %h64, i32 0)
+  %old64 = atomicrmw add ptr %p64, i64 1 monotonic
+
+  ; RWByteAddressBuffer Atomic32UAV : register(u1);
+  %h32 = call target("dx.RawBuffer", i8, 1, 0)
+      @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_i8_1_0t(
+          i32 0, i32 1, i32 1, i32 0, ptr null)
+  ; CHECK:      Binding:
+  ; CHECK:        Record ID: 1
+  ; CHECK:        Space: 0
+  ; CHECK:        Lower Bound: 1
+  ; CHECK:        Size: 1
+  ; CHECK:      Has Atomic64 Use: 0
+  %p32 = call ptr @llvm.dx.resource.getpointer(
+      target("dx.RawBuffer", i8, 1, 0) %h32, i32 0)
+  %old32 = atomicrmw add ptr %p32, i32 1 monotonic
+
+  ; RWBuffer<int> PlainUAV : register(u2);
+  %hp = call target("dx.TypedBuffer", i32, 1, 0, 1)
+      @llvm.dx.resource.handlefrombinding.tdx.TypedBuffer_i32_1_0_1t(
+          i32 0, i32 2, i32 1, i32 0, ptr null)
+  ; CHECK:      Binding:
+  ; CHECK:        Record ID: 2
+  ; CHECK:        Space: 0
+  ; CHECK:        Lower Bound: 2
+  ; CHECK:        Size: 1
+  ; CHECK:      Has Atomic64 Use: 0
+  %pp = call ptr @llvm.dx.resource.getpointer(
+      target("dx.TypedBuffer", i32, 1, 0, 1) %hp, i32 0)
+  store i32 42, ptr %pp
+
+  ret void
+}
diff --git a/llvm/test/CodeGen/DirectX/ContainerData/PSVResources-UsedByAtomic64.ll b/llvm/test/CodeGen/DirectX/ContainerData/PSVResources-UsedByAtomic64.ll
new file mode 100644
index 0000000000000..6213e1ae9ca30
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/ContainerData/PSVResources-UsedByAtomic64.ll
@@ -0,0 +1,47 @@
+; RUN: llc %s -disable-dxil-remove-unused-resources --filetype=obj -o - | obj2yaml | FileCheck %s
+
+; Two UAVs: `AtomicUAV` is used by a 64-bit atomicrmw, `PlainUAV` is only
+; loaded/stored (no atomics). Only `AtomicUAV` should have the
+; `UsedByAtomic64` PSV resource flag set.
+
+target triple = "dxil-pc-shadermodel6.6-compute"
+
+; CHECK: Resources:
+
+; RWByteAddressBuffer AtomicUAV : register(u0);
+; CHECK:        - Type:            UAVRaw
+; CHECK:          Space:           0
+; CHECK:          LowerBound:      0
+; CHECK:          UpperBound:      0
+; CHECK:          Kind:            RawBuffer
+; CHECK:          Flags:
+; CHECK:            UsedByAtomic64:  true
+
+; RWBuffer<int> PlainUAV : register(u1);
+; CHECK:        - Type:            UAVTyped
+; CHECK:          Space:           0
+; CHECK:          LowerBound:      1
+; CHECK:          UpperBound:      1
+; CHECK:          Kind:            TypedBuffer
+; CHECK:          Flags:
+; CHECK:            UsedByAtomic64:  false
+
+define void @main() #0 {
+  %atomic = call target("dx.RawBuffer", i8, 1, 0)
+      @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_i8_1_0t(
+          i32 0, i32 0, i32 1, i32 0, ptr null)
+  %atomicPtr = call ptr @llvm.dx.resource.getpointer(
+      target("dx.RawBuffer", i8, 1, 0) %atomic, i32 0)
+  %old = atomicrmw add ptr %atomicPtr, i64 1 monotonic
+
+  %plain = call target("dx.TypedBuffer", i32, 1, 0, 1)
+      @llvm.dx.resource.handlefrombinding.tdx.TypedBuffer_i32_1_0_1t(
+          i32 0, i32 1, i32 1, i32 0, ptr null)
+  %plainPtr = call ptr @llvm.dx.resource.getpointer(
+      target("dx.TypedBuffer", i32, 1, 0, 1) %plain, i32 0)
+  store i32 42, ptr %plainPtr
+
+  ret void
+}
+
+attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }

>From 8f1ecf6955dab1754795d5e1c8142e5b87f15348 Mon Sep 17 00:00:00 2001
From: Joshua Batista <jbatista at microsoft.com>
Date: Thu, 23 Jul 2026 16:17:05 -0700
Subject: [PATCH 2/4] add missing header

---
 llvm/lib/Analysis/DXILResource.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/llvm/lib/Analysis/DXILResource.cpp b/llvm/lib/Analysis/DXILResource.cpp
index 4baeb306ee0f1..24601dd1582ee 100644
--- a/llvm/lib/Analysis/DXILResource.cpp
+++ b/llvm/lib/Analysis/DXILResource.cpp
@@ -17,6 +17,7 @@
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/IntrinsicsDirectX.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"

>From 4a1a6764c435f0fe0eff1e8ab9be4d68b3d33486 Mon Sep 17 00:00:00 2001
From: Joshua Batista <jbatista at microsoft.com>
Date: Thu, 23 Jul 2026 16:25:34 -0700
Subject: [PATCH 3/4] clang format

---
 llvm/lib/Analysis/DXILResource.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Analysis/DXILResource.cpp b/llvm/lib/Analysis/DXILResource.cpp
index 24601dd1582ee..c4b7b6b0c948e 100644
--- a/llvm/lib/Analysis/DXILResource.cpp
+++ b/llvm/lib/Analysis/DXILResource.cpp
@@ -16,8 +16,8 @@
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
-#include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/IntrinsicsDirectX.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"

>From d93f36f7b7caee86a595cf78b78ec2463264f1d8 Mon Sep 17 00:00:00 2001
From: Joshua Batista <jbatista at microsoft.com>
Date: Fri, 24 Jul 2026 14:48:33 -0700
Subject: [PATCH 4/4] address Finn

---
 llvm/include/llvm/Analysis/DXILResource.h |   7 +-
 llvm/lib/Analysis/DXILResource.cpp        | 111 +++++++++-------------
 2 files changed, 48 insertions(+), 70 deletions(-)

diff --git a/llvm/include/llvm/Analysis/DXILResource.h b/llvm/include/llvm/Analysis/DXILResource.h
index d0c6ea44b6add..a41c6f970c241 100644
--- a/llvm/include/llvm/Analysis/DXILResource.h
+++ b/llvm/include/llvm/Analysis/DXILResource.h
@@ -515,10 +515,9 @@ class DXILResourceMap {
   void populate(Module &M, DXILResourceTypeMap &DRTM);
   /// Populate the map given the resource binding calls in the given module.
   void populateResourceInfos(Module &M, DXILResourceTypeMap &DRTM);
-  /// Analyze and populate the directions of the resource counters.
-  void populateCounterDirections(Module &M);
-  /// Detect 64-bit atomic uses of resources and set `HasAtomic64Use`.
-  void populateAtomicUses(Module &M);
+  /// Analyze uses to fill in per-resource dynamic state — counter directions
+  /// and 64-bit atomic use — in a single walk of the module's instructions.
+  void populateFromInstructions(Module &M);
 
   /// Resolves a resource handle into a vector of ResourceInfos that
   /// represent the possible unique creations of the handle. Certain cases are
diff --git a/llvm/lib/Analysis/DXILResource.cpp b/llvm/lib/Analysis/DXILResource.cpp
index c4b7b6b0c948e..687ca3fa64ba0 100644
--- a/llvm/lib/Analysis/DXILResource.cpp
+++ b/llvm/lib/Analysis/DXILResource.cpp
@@ -837,10 +837,6 @@ bool DXILResourceTypeMap::invalidate(Module &M, const PreservedAnalyses &PA,
 }
 
 //===----------------------------------------------------------------------===//
-static bool isUpdateCounterIntrinsic(Function &F) {
-  return F.getIntrinsicID() == Intrinsic::dx_resource_updatecounter;
-}
-
 StringRef dxil::getResourceNameFromBindingCall(CallInst *CI) {
   Value *Op = nullptr;
   switch (CI->getCalledFunction()->getIntrinsicID()) {
@@ -948,46 +944,7 @@ void DXILResourceMap::populateResourceInfos(Module &M,
   }
 }
 
-void DXILResourceMap::populateCounterDirections(Module &M) {
-  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;
-        else if (RBInfo->CounterDirection != Direction) {
-          RBInfo->CounterDirection = ResourceCounterDirection::Invalid;
-          HasInvalidDirection = true;
-        }
-      }
-    }
-  }
-}
-
-void DXILResourceMap::populateAtomicUses(Module &M) {
+void DXILResourceMap::populateFromInstructions(Module &M) {
   auto FindResourceHandle = [](Value *Ptr) -> Value * {
     Ptr = Ptr->stripPointerCasts();
     while (auto *GEP = dyn_cast<GetElementPtrInst>(Ptr))
@@ -1005,40 +962,62 @@ void DXILResourceMap::populateAtomicUses(Module &M) {
       RI->HasAtomic64Use = true;
   };
 
-  // Detect on `atomicrmw`/`cmpxchg` before `DXILResourceAccess` lowers them.
+  auto RecordCounterDirection = [this](const CallInst *CI) {
+    ConstantInt *CountValue = cast<ConstantInt>(CI->getArgOperand(1));
+    int64_t CountLiteral = CountValue->getSExtValue();
+    if (CountLiteral == 0)
+      return;
+    ResourceCounterDirection Direction =
+        CountLiteral > 0 ? ResourceCounterDirection::Increment
+                         : ResourceCounterDirection::Decrement;
+    for (ResourceInfo *RBInfo : findByUse(CI->getArgOperand(0))) {
+      if (RBInfo->CounterDirection == ResourceCounterDirection::Unknown)
+        RBInfo->CounterDirection = Direction;
+      else if (RBInfo->CounterDirection != Direction) {
+        RBInfo->CounterDirection = ResourceCounterDirection::Invalid;
+        HasInvalidDirection = true;
+      }
+    }
+  };
+
+  // Single walk of every instruction in the module. Handles both `atomicrmw`
+  // / `cmpxchg` (when this analysis runs before `DXILResourceAccess`) and the
+  // lowered `llvm.dx.resource.atomic.binop` form (when it runs after), since
+  // `DXILResourceAccess` does not preserve `DXILResourceWrapperPass` and any
+  // downstream request re-populates the map.
   for (Function &F : M.functions()) {
     for (Instruction &I : instructions(F)) {
       if (auto *AI = dyn_cast<AtomicRMWInst>(&I)) {
-        if (!AI->getValOperand()->getType()->isIntegerTy(64))
-          continue;
-        MarkHasAtomic64UseFromHandle(
-            FindResourceHandle(AI->getPointerOperand()));
+        if (AI->getValOperand()->getType()->isIntegerTy(64))
+          MarkHasAtomic64UseFromHandle(
+              FindResourceHandle(AI->getPointerOperand()));
       } else if (auto *CX = dyn_cast<AtomicCmpXchgInst>(&I)) {
-        if (!CX->getNewValOperand()->getType()->isIntegerTy(64))
+        if (CX->getNewValOperand()->getType()->isIntegerTy(64))
+          MarkHasAtomic64UseFromHandle(
+              FindResourceHandle(CX->getPointerOperand()));
+      } else if (auto *CI = dyn_cast<CallInst>(&I)) {
+        Function *Called = CI->getCalledFunction();
+        if (!Called)
           continue;
-        MarkHasAtomic64UseFromHandle(
-            FindResourceHandle(CX->getPointerOperand()));
+        switch (Called->getIntrinsicID()) {
+        case Intrinsic::dx_resource_updatecounter:
+          RecordCounterDirection(CI);
+          break;
+        case Intrinsic::dx_resource_atomic_binop:
+          if (CI->getType()->isIntegerTy(64))
+            MarkHasAtomic64UseFromHandle(CI->getArgOperand(0));
+          break;
+        default:
+          break;
+        }
       }
     }
   }
-
-  // Also handle the post-lowering `llvm.dx.resource.atomic.binop` form.
-  for (Function &F : M.functions()) {
-    if (F.getIntrinsicID() != Intrinsic::dx_resource_atomic_binop)
-      continue;
-    for (User *U : F.users()) {
-      auto *CI = dyn_cast<CallInst>(U);
-      if (!CI || !CI->getType()->isIntegerTy(64))
-        continue;
-      MarkHasAtomic64UseFromHandle(CI->getArgOperand(0));
-    }
-  }
 }
 
 void DXILResourceMap::populate(Module &M, DXILResourceTypeMap &DRTM) {
   populateResourceInfos(M, DRTM);
-  populateCounterDirections(M);
-  populateAtomicUses(M);
+  populateFromInstructions(M);
 }
 
 void DXILResourceMap::print(raw_ostream &OS, DXILResourceTypeMap &DRTM,



More information about the llvm-commits mailing list