[llvm] [AMDGPU] Reject illegal buffer atomic data widths in SelectionDAG (PR #210925)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Arseniy Obolenskiy (aobolensk)

<details>
<summary>Changes</summary>

Continuation of the fix introduced in #<!-- -->210366, but this time for raw/struct buffer atomics instead of image atomics

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


2 Files Affected:

- (modified) llvm/lib/Target/AMDGPU/SIISelLowering.cpp (+22) 
- (added) llvm/test/CodeGen/AMDGPU/llvm.amdgcn.buffer.atomic.illegal-data-type.err.ll (+53) 


``````````diff
diff --git a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
index 8156de7a24b01..f92b1731c3f7d 100644
--- a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
+++ b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
@@ -11471,6 +11471,11 @@ SDValue SITargetLowering::lowerRawBufferAtomicIntrin(SDValue Op,
   SDLoc DL(Op);
 
   SDValue VData = Op.getOperand(2);
+  if (VData.getValueSizeInBits() != 32 && VData.getValueSizeInBits() != 64) {
+    SmallVector<EVT, 2> ResultTypes(Op->values());
+    return diagnoseUnsupportedImage(DAG, Op, ResultTypes, DL,
+                                    "unsupported buffer atomic data type");
+  }
   SDValue Rsrc = bufferRsrcPtrToVector(Op.getOperand(3), DAG);
   auto [VOffset, Offset] = splitBufferOffsets(Op.getOperand(4), DAG);
   auto SOffset = selectSOffset(Op.getOperand(5), DAG, Subtarget);
@@ -11499,6 +11504,11 @@ SITargetLowering::lowerStructBufferAtomicIntrin(SDValue Op, SelectionDAG &DAG,
   SDLoc DL(Op);
 
   SDValue VData = Op.getOperand(2);
+  if (VData.getValueSizeInBits() != 32 && VData.getValueSizeInBits() != 64) {
+    SmallVector<EVT, 2> ResultTypes(Op->values());
+    return diagnoseUnsupportedImage(DAG, Op, ResultTypes, DL,
+                                    "unsupported buffer atomic data type");
+  }
   SDValue Rsrc = bufferRsrcPtrToVector(Op.getOperand(3), DAG);
   auto [VOffset, Offset] = splitBufferOffsets(Op.getOperand(5), DAG);
   auto SOffset = selectSOffset(Op.getOperand(6), DAG, Subtarget);
@@ -11812,6 +11822,12 @@ SDValue SITargetLowering::LowerINTRINSIC_W_CHAIN(SDValue Op,
                                          AMDGPUISD::BUFFER_ATOMIC_COND_SUB_U32);
   case Intrinsic::amdgcn_raw_buffer_atomic_cmpswap:
   case Intrinsic::amdgcn_raw_ptr_buffer_atomic_cmpswap: {
+    SDValue Src = Op.getOperand(2);
+    if (Src.getValueSizeInBits() != 32 && Src.getValueSizeInBits() != 64) {
+      SmallVector<EVT, 2> ResultTypes(Op->values());
+      return diagnoseUnsupportedImage(DAG, Op, ResultTypes, DL,
+                                      "unsupported buffer atomic data type");
+    }
     SDValue Rsrc = bufferRsrcPtrToVector(Op.getOperand(4), DAG);
     auto [VOffset, Offset] = splitBufferOffsets(Op.getOperand(5), DAG);
     auto SOffset = selectSOffset(Op.getOperand(6), DAG, Subtarget);
@@ -11836,6 +11852,12 @@ SDValue SITargetLowering::LowerINTRINSIC_W_CHAIN(SDValue Op,
   }
   case Intrinsic::amdgcn_struct_buffer_atomic_cmpswap:
   case Intrinsic::amdgcn_struct_ptr_buffer_atomic_cmpswap: {
+    SDValue Src = Op.getOperand(2);
+    if (Src.getValueSizeInBits() != 32 && Src.getValueSizeInBits() != 64) {
+      SmallVector<EVT, 2> ResultTypes(Op->values());
+      return diagnoseUnsupportedImage(DAG, Op, ResultTypes, DL,
+                                      "unsupported buffer atomic data type");
+    }
     SDValue Rsrc = bufferRsrcPtrToVector(Op->getOperand(4), DAG);
     auto [VOffset, Offset] = splitBufferOffsets(Op.getOperand(6), DAG);
     auto SOffset = selectSOffset(Op.getOperand(7), DAG, Subtarget);
diff --git a/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.buffer.atomic.illegal-data-type.err.ll b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.buffer.atomic.illegal-data-type.err.ll
new file mode 100644
index 0000000000000..ea529f2f41e47
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.buffer.atomic.illegal-data-type.err.ll
@@ -0,0 +1,53 @@
+; RUN: not llc -mtriple=amdgpu9.0a -O0 -filetype=null %s 2>&1 | FileCheck %s
+
+; llvm.amdgcn.{raw,struct}.ptr.buffer.atomic.<op> only supports 32-bit and
+; 64-bit data.
+
+; CHECK: error: {{.*}}unsupported buffer atomic data type
+define amdgpu_kernel void @raw_add_v3i16(ptr addrspace(8) %rsrc, <3 x i16> %data, ptr addrspace(1) %out) {
+  %v = call <3 x i16> @llvm.amdgcn.raw.ptr.buffer.atomic.add.v3i16(<3 x i16> %data, ptr addrspace(8) %rsrc, i32 0, i32 0, i32 0)
+  store <3 x i16> %v, ptr addrspace(1) %out
+  ret void
+}
+
+; CHECK: error: {{.*}}unsupported buffer atomic data type
+define amdgpu_kernel void @raw_add_i128(ptr addrspace(8) %rsrc, i128 %data, ptr addrspace(1) %out) {
+  %v = call i128 @llvm.amdgcn.raw.ptr.buffer.atomic.add.i128(i128 %data, ptr addrspace(8) %rsrc, i32 0, i32 0, i32 0)
+  store i128 %v, ptr addrspace(1) %out
+  ret void
+}
+
+; CHECK: error: {{.*}}unsupported buffer atomic data type
+define amdgpu_kernel void @raw_fadd_bf16(ptr addrspace(8) %rsrc, bfloat %data, ptr addrspace(1) %out) {
+  %v = call bfloat @llvm.amdgcn.raw.ptr.buffer.atomic.fadd.bf16(bfloat %data, ptr addrspace(8) %rsrc, i32 0, i32 0, i32 0)
+  store bfloat %v, ptr addrspace(1) %out
+  ret void
+}
+
+; CHECK: error: {{.*}}unsupported buffer atomic data type
+define amdgpu_kernel void @struct_add_v3i16(ptr addrspace(8) %rsrc, <3 x i16> %data, i32 %idx, ptr addrspace(1) %out) {
+  %v = call <3 x i16> @llvm.amdgcn.struct.ptr.buffer.atomic.add.v3i16(<3 x i16> %data, ptr addrspace(8) %rsrc, i32 %idx, i32 0, i32 0, i32 0)
+  store <3 x i16> %v, ptr addrspace(1) %out
+  ret void
+}
+
+; CHECK: error: {{.*}}unsupported buffer atomic data type
+define amdgpu_kernel void @raw_cmpswap_i128(ptr addrspace(8) %rsrc, i128 %cmp, i128 %data, ptr addrspace(1) %out) {
+  %v = call i128 @llvm.amdgcn.raw.ptr.buffer.atomic.cmpswap.i128(i128 %cmp, i128 %data, ptr addrspace(8) %rsrc, i32 0, i32 0, i32 0)
+  store i128 %v, ptr addrspace(1) %out
+  ret void
+}
+
+; CHECK: error: {{.*}}unsupported buffer atomic data type
+define amdgpu_kernel void @struct_cmpswap_i128(ptr addrspace(8) %rsrc, i128 %cmp, i128 %data, i32 %idx, ptr addrspace(1) %out) {
+  %v = call i128 @llvm.amdgcn.struct.ptr.buffer.atomic.cmpswap.i128(i128 %cmp, i128 %data, ptr addrspace(8) %rsrc, i32 %idx, i32 0, i32 0, i32 0)
+  store i128 %v, ptr addrspace(1) %out
+  ret void
+}
+
+declare <3 x i16> @llvm.amdgcn.raw.ptr.buffer.atomic.add.v3i16(<3 x i16>, ptr addrspace(8), i32, i32, i32 immarg)
+declare i128 @llvm.amdgcn.raw.ptr.buffer.atomic.add.i128(i128, ptr addrspace(8), i32, i32, i32 immarg)
+declare bfloat @llvm.amdgcn.raw.ptr.buffer.atomic.fadd.bf16(bfloat, ptr addrspace(8), i32, i32, i32 immarg)
+declare <3 x i16> @llvm.amdgcn.struct.ptr.buffer.atomic.add.v3i16(<3 x i16>, ptr addrspace(8), i32, i32, i32, i32 immarg)
+declare i128 @llvm.amdgcn.raw.ptr.buffer.atomic.cmpswap.i128(i128, i128, ptr addrspace(8), i32, i32, i32 immarg)
+declare i128 @llvm.amdgcn.struct.ptr.buffer.atomic.cmpswap.i128(i128, i128, ptr addrspace(8), i32, i32, i32, i32 immarg)

``````````

</details>


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