[llvm] [DirectX] introducing lowering for `bufferUpdateCounter` (PR #115041)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-directx

Author: None (joaosaffran)

<details>
<summary>Changes</summary>

- Adding custom lowering for `bufferUpdateCounter`
- introduces llvm intrinsic `int_dx_updateCounter`
- adds tests

Closes #<!-- -->92147

---
Full diff: https://github.com/llvm/llvm-project/pull/115041.diff


4 Files Affected:

- (modified) llvm/include/llvm/IR/IntrinsicsDirectX.td (+3) 
- (modified) llvm/lib/Target/DirectX/DXIL.td (+9) 
- (modified) llvm/lib/Target/DirectX/DXILOpLowering.cpp (+25) 
- (added) llvm/test/CodeGen/DirectX/updateCounter.ll (+31) 


``````````diff
diff --git a/llvm/include/llvm/IR/IntrinsicsDirectX.td b/llvm/include/llvm/IR/IntrinsicsDirectX.td
index e30d37f69f781e..d460fe44e004d4 100644
--- a/llvm/include/llvm/IR/IntrinsicsDirectX.td
+++ b/llvm/include/llvm/IR/IntrinsicsDirectX.td
@@ -35,6 +35,9 @@ def int_dx_typedBufferLoad_checkbit
 def int_dx_typedBufferStore
     : DefaultAttrsIntrinsic<[], [llvm_any_ty, llvm_i32_ty, llvm_anyvector_ty]>;
 
+def int_dx_updateCounter
+    : DefaultAttrsIntrinsic<[], [llvm_any_ty, llvm_anyint_ty]>;
+    
 // Cast between target extension handle types and dxil-style opaque handles
 def int_dx_cast_handle : Intrinsic<[llvm_any_ty], [llvm_any_ty]>;
 
diff --git a/llvm/lib/Target/DirectX/DXIL.td b/llvm/lib/Target/DirectX/DXIL.td
index af12b74351058e..edfadd41d84216 100644
--- a/llvm/lib/Target/DirectX/DXIL.td
+++ b/llvm/lib/Target/DirectX/DXIL.td
@@ -730,6 +730,15 @@ def BufferStore : DXILOp<69, bufferStore> {
   let stages = [Stages<DXIL1_0, [all_stages]>];
 }
 
+def UpdateCounter : DXILOp<70, bufferUpdateCounter> {
+  let Doc = "increments/decrements a buffer counter";
+  let arguments = [
+    HandleTy, Int8Ty
+  ];
+  let result = VoidTy;
+  let stages = [Stages<DXIL1_0, [all_stages]>];
+}
+
 def CheckAccessFullyMapped : DXILOp<71, checkAccessFullyMapped> {
   let Doc = "checks whether a Sample, Gather, or Load operation "
             "accessed mapped tiles in a tiled resource";
diff --git a/llvm/lib/Target/DirectX/DXILOpLowering.cpp b/llvm/lib/Target/DirectX/DXILOpLowering.cpp
index 8acc9c1efa08c0..0dd3a8dc1ad4ce 100644
--- a/llvm/lib/Target/DirectX/DXILOpLowering.cpp
+++ b/llvm/lib/Target/DirectX/DXILOpLowering.cpp
@@ -463,6 +463,28 @@ class OpLowerer {
     });
   }
 
+  [[nodiscard]] bool lowerUpdateCounter(Function &F) {
+    IRBuilder<> &IRB = OpBuilder.getIRB();
+
+    return replaceFunction(F, [&](CallInst *CI) -> Error {
+      IRB.SetInsertPoint(CI);
+      Value *Handle =
+          createTmpHandleCast(CI->getArgOperand(0), OpBuilder.getHandleType());
+      Value *Op1 = CI->getArgOperand(1);
+
+      std::array<Value *, 2> Args{Handle, Op1};
+
+      Expected<CallInst *> OpCall =
+          OpBuilder.tryCreateOp(OpCode::UpdateCounter, Args, CI->getName());
+
+      if (Error E = OpCall.takeError())
+        return E;
+
+      CI->eraseFromParent();
+      return Error::success();
+    });
+  }
+
   [[nodiscard]] bool lowerTypedBufferStore(Function &F) {
     IRBuilder<> &IRB = OpBuilder.getIRB();
     Type *Int8Ty = IRB.getInt8Ty();
@@ -600,6 +622,9 @@ class OpLowerer {
       case Intrinsic::dx_typedBufferStore:
         HasErrors |= lowerTypedBufferStore(F);
         break;
+      case Intrinsic::dx_updateCounter:
+        HasErrors |= lowerUpdateCounter(F);
+        break;
       // TODO: this can be removed when
       // https://github.com/llvm/llvm-project/issues/113192 is fixed
       case Intrinsic::dx_splitdouble:
diff --git a/llvm/test/CodeGen/DirectX/updateCounter.ll b/llvm/test/CodeGen/DirectX/updateCounter.ll
new file mode 100644
index 00000000000000..459668c0a3cb51
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/updateCounter.ll
@@ -0,0 +1,31 @@
+; RUN: opt -S -dxil-op-lower %s | FileCheck %s
+
+
+target triple = "dxil-pc-shadermodel6.6-compute"
+
+define void @update_counter_decrement_vector() {
+  %buffer = call target("dx.TypedBuffer", <4 x float>, 0, 0, 0)
+      @llvm.dx.handle.fromBinding.tdx.TypedBuffer_v4f32_0_0_0(
+          i32 0, i32 0, i32 1, i32 0, i1 false)
+
+  call void @llvm.dx.updateCounter(target("dx.TypedBuffer", <4 x float>, 0, 0, 0) %buffer, i8 -1)
+  ret void
+}
+
+define void @update_counter_increment_vector() {
+  %buffer = call target("dx.TypedBuffer", <4 x float>, 0, 0, 0)
+      @llvm.dx.handle.fromBinding.tdx.TypedBuffer_v4f32_0_0_0(
+          i32 0, i32 0, i32 1, i32 0, i1 false)
+
+  call void @llvm.dx.updateCounter(target("dx.TypedBuffer", <4 x float>, 0, 0, 0) %buffer, i8 1)
+  ret void
+}
+
+define void @update_counter_decrement_scalar() {
+  %buffer = call target("dx.RawBuffer", i8, 0, 0)
+      @llvm.dx.handle.fromBinding.tdx.RawBuffer_i8_0_0t(
+          i32 1, i32 8, i32 1, i32 0, i1 false)
+
+  call void @llvm.dx.updateCounter(target("dx.RawBuffer", i8, 0, 0) %buffer, i8 -1)
+  ret void
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/115041