[llvm] [DirectX] Fix DXIL expansion of 64-bit ByteAddressBuffer stores (PR #224136)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-directx

Author: Farzon Lotfi (farzonl)

<details>
<summary>Changes</summary>

Fixes #<!-- -->223558

In llvm/lib/Target/DirectX/DXILIntrinsicExpansion.cpp we were not doing Create add with the  byte index offsets because were were looking at the wrong argument.

- Fixes DXIL expansion of 64-bit ByteAddressBuffer stores by advancing the byte index instead of the poison offset.
- Adds regression coverage for wide double stores with offsets confirmed as 0, 16, 32, 48.

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


3 Files Affected:

- (modified) llvm/lib/Target/DirectX/DXILIntrinsicExpansion.cpp (+5-1) 
- (modified) llvm/test/CodeGen/DirectX/RawBufferStoreDouble.ll (+26) 
- (added) llvm/test/CodeGen/DirectX/ResourceAccess/store-rawbuffer-wide-double.ll (+22) 


``````````diff
diff --git a/llvm/lib/Target/DirectX/DXILIntrinsicExpansion.cpp b/llvm/lib/Target/DirectX/DXILIntrinsicExpansion.cpp
index 95781680ee0de..2a1416deccc3a 100644
--- a/llvm/lib/Target/DirectX/DXILIntrinsicExpansion.cpp
+++ b/llvm/lib/Target/DirectX/DXILIntrinsicExpansion.cpp
@@ -969,7 +969,11 @@ static bool expandBufferStoreIntrinsic(CallInst *Orig, bool IsRaw) {
     if (IsRaw) {
       StoreIntrinsic = Intrinsic::dx_resource_store_rawbuffer;
       Value *Tmp = Builder.getInt32(4 * Base);
-      Args.push_back(Builder.CreateAdd(Orig->getOperand(2), Tmp));
+      Value *Offset = Orig->getOperand(2);
+      Args.push_back(Offset);
+      unsigned AddressArg = isa<PoisonValue>(Offset) ? 1 : 2;
+      if (Base != 0)
+        Args[AddressArg] = Builder.CreateAdd(Args[AddressArg], Tmp);
     }
 
     SmallVector<int, 4> Mask;
diff --git a/llvm/test/CodeGen/DirectX/RawBufferStoreDouble.ll b/llvm/test/CodeGen/DirectX/RawBufferStoreDouble.ll
index 30a8039820e28..d076d7d167f9e 100644
--- a/llvm/test/CodeGen/DirectX/RawBufferStoreDouble.ll
+++ b/llvm/test/CodeGen/DirectX/RawBufferStoreDouble.ll
@@ -103,3 +103,29 @@ define void @storev4f64(<4 x double> %0, i32 %index) {
       <4 x double> %0)
   ret void
 }
+
+define void @storev4f64_byte(<4 x double> %0, i32 %index) {
+  ; CHECK: [[Buf:%.*]] = tail call target("dx.RawBuffer", i8, 1, 0)
+  ; CHECK-SAME: @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_i8_1_0t(
+  ; CHECK-SAME: i32 0, i32 0, i32 1, i32 0, ptr null)
+  %buffer = tail 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)
+
+  ; CHECK63: call void @llvm.dx.resource.store.rawbuffer
+  ; CHECK63-SAME: target("dx.RawBuffer", i8, 1, 0) [[Buf]], i32 %index, i32 poison, <4 x double> %0)
+
+  ; CHECK62: [[SD:%.*]] = call { <4 x i32>, <4 x i32> } @llvm.dx.splitdouble.v4i32(<4 x double> %0)
+  ; CHECK62: [[Lo:%.*]] = extractvalue { <4 x i32>, <4 x i32> } [[SD]], 0
+  ; CHECK62: [[Hi:%.*]] = extractvalue { <4 x i32>, <4 x i32> } [[SD]], 1
+  ; CHECK62: [[Vec:%.*]] = shufflevector <4 x i32> [[Lo]], <4 x i32> [[Hi]], <8 x i32> <i32 0, i32 4, i32 1, i32 5, i32 2, i32 6, i32 3, i32 7>
+  ; CHECK62: [[Low:%.*]] = shufflevector <8 x i32> [[Vec]], <8 x i32> poison, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+  ; CHECK62: call void @llvm.dx.resource.store.rawbuffer.tdx.RawBuffer_i8_1_0t.v4i32(target("dx.RawBuffer", i8, 1, 0) [[Buf]], i32 %index, i32 poison, <4 x i32> [[Low]])
+  ; CHECK62: [[NextIndex:%.*]] = add i32 %index, 16
+  ; CHECK62: [[High:%.*]] = shufflevector <8 x i32> [[Vec]], <8 x i32> poison, <4 x i32> <i32 4, i32 5, i32 6, i32 7>
+  ; CHECK62: call void @llvm.dx.resource.store.rawbuffer.tdx.RawBuffer_i8_1_0t.v4i32(target("dx.RawBuffer", i8, 1, 0) [[Buf]], i32 [[NextIndex]], i32 poison, <4 x i32> [[High]])
+  call void @llvm.dx.resource.store.rawbuffer(
+      target("dx.RawBuffer", i8, 1, 0) %buffer, i32 %index, i32 poison,
+      <4 x double> %0)
+  ret void
+}
diff --git a/llvm/test/CodeGen/DirectX/ResourceAccess/store-rawbuffer-wide-double.ll b/llvm/test/CodeGen/DirectX/ResourceAccess/store-rawbuffer-wide-double.ll
new file mode 100644
index 0000000000000..307b1b832eaf4
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/ResourceAccess/store-rawbuffer-wide-double.ll
@@ -0,0 +1,22 @@
+; RUN: opt -S -dxil-resource-access -dxil-intrinsic-expansion \
+; RUN:   -mtriple=dxil-pc-shadermodel6.2-compute %s | FileCheck %s
+
+define void @storev8f64_byte(i32 %index, <8 x double> %data) {
+  %buffer = call target("dx.RawBuffer", i8, 1, 0)
+      @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, ptr null)
+
+  ; CHECK-NOT: @llvm.dx.resource.getpointer
+  %ptr = call ptr @llvm.dx.resource.getpointer(
+      target("dx.RawBuffer", i8, 1, 0) %buffer, i32 %index)
+
+  ; CHECK: call void [[STORE:@llvm\.dx\.resource\.store\.rawbuffer[^(]*]]([[BUFFER_TY:target\("dx\.RawBuffer", i8, 1, 0\)]] %buffer, i32 %index, [[STORE_TAIL:i32 poison, <4 x i32>]]
+  ; CHECK: %[[INDEX16:.*]] = add i32 %index, 16
+  ; CHECK: call void [[STORE]]([[BUFFER_TY]] %buffer, i32 %[[INDEX16]], [[STORE_TAIL]]
+  ; CHECK: %[[INDEX32:.*]] = add i32 %index, 32
+  ; CHECK: call void [[STORE]]([[BUFFER_TY]] %buffer, i32 %[[INDEX32]], [[STORE_TAIL]]
+  ; CHECK: %[[INDEX48:.*]] = add i32 %[[INDEX32]], 16
+  ; CHECK: call void [[STORE]]([[BUFFER_TY]] %buffer, i32 %[[INDEX48]], [[STORE_TAIL]]
+  store <8 x double> %data, ptr %ptr
+
+  ret void
+}

``````````

</details>


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