[llvm] [AMDGPU] SelDAG: fix lowering of undefined workitem intrinsics (PR #126058)

[Robert Imschweiler via llvm-commits]

https://github.com/ro-i created https://github.com/llvm/llvm-project/pull/126058

GlobalISel already handles undefined workitem.id.{x,y,z} intrinsics, SelDAG failed in AMDGPUISelLowering.cpp due to a failed assertion in `AMDGPUTargetLowering::loadInputValue`: `Arg && "Attempting to load missing argument"`. This commit changes the behavior of SelDAG to instead use a zero constant.

This LLVM defect was identified via the AMD Fuzzing project.

>From 9c8af8a0a76019069d04b93608ac0ebd4e2df234 Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Wed, 5 Feb 2025 08:02:43 -0600
Subject: [PATCH] [AMDGPU] SelDAG: fix lowering of undefined workitem
 intrinsics

GlobalISel already handles undefined workitem.id.{x,y,z} intrinsics, SelDAG failed in AMDGPUISelLowering.cpp due to a failed assertion in `AMDGPUTargetLowering::loadInputValue`: `Arg && "Attempting to load missing argument"`.
This commit changes the behavior of SelDAG to instead use a zero constant.

This LLVM defect was identified via the AMD Fuzzing project.
---
 llvm/lib/Target/AMDGPU/SIISelLowering.cpp     | 23 ++++++++++--
 .../llvm.amdgcn.workitem.id-undefined-attr.ll | 36 +++++++++++++++++++
 2 files changed, 56 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/CodeGen/AMDGPU/llvm.amdgcn.workitem.id-undefined-attr.ll

diff --git a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
index b632c50dae0e359..47654890ff50e74 100644
--- a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
+++ b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
@@ -8790,11 +8790,28 @@ SDValue SITargetLowering::LowerINTRINSIC_WO_CHAIN(SDValue Op,
                              AMDGPUFunctionArgInfo::LDS_KERNEL_ID);
   }
   case Intrinsic::amdgcn_workitem_id_x:
-    return lowerWorkitemID(DAG, Op, 0, MFI->getArgInfo().WorkItemIDX);
+    if (!MFI->getArgInfo().WorkItemIDX) {
+      // It's undefined behavior if a function marked with the amdgpu-no-*
+      // attributes uses the corresponding intrinsic.
+      return DAG.getConstant(0, SDLoc(Op),
+                             EVT::getIntegerVT(*DAG.getContext(), 32));
+    } else {
+      return lowerWorkitemID(DAG, Op, 0, MFI->getArgInfo().WorkItemIDX);
+    }
   case Intrinsic::amdgcn_workitem_id_y:
-    return lowerWorkitemID(DAG, Op, 1, MFI->getArgInfo().WorkItemIDY);
+    if (!MFI->getArgInfo().WorkItemIDY) {
+      return DAG.getConstant(0, SDLoc(Op),
+                             EVT::getIntegerVT(*DAG.getContext(), 32));
+    } else {
+      return lowerWorkitemID(DAG, Op, 1, MFI->getArgInfo().WorkItemIDY);
+    }
   case Intrinsic::amdgcn_workitem_id_z:
-    return lowerWorkitemID(DAG, Op, 2, MFI->getArgInfo().WorkItemIDZ);
+    if (!MFI->getArgInfo().WorkItemIDZ) {
+      return DAG.getConstant(0, SDLoc(Op),
+                             EVT::getIntegerVT(*DAG.getContext(), 32));
+    } else {
+      return lowerWorkitemID(DAG, Op, 2, MFI->getArgInfo().WorkItemIDZ);
+    }
   case Intrinsic::amdgcn_wavefrontsize:
     return DAG.getConstant(MF.getSubtarget<GCNSubtarget>().getWavefrontSize(),
                            SDLoc(Op), MVT::i32);
diff --git a/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.workitem.id-undefined-attr.ll b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.workitem.id-undefined-attr.ll
new file mode 100644
index 000000000000000..0e6a448fcae5a96
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.workitem.id-undefined-attr.ll
@@ -0,0 +1,36 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -mtriple=amdgcn -mcpu=gfx942 -O0 --global-isel=false -o - %s | FileCheck %s
+
+declare i32 @llvm.amdgcn.workitem.id.x() #0
+declare i32 @llvm.amdgcn.workitem.id.y() #0
+declare i32 @llvm.amdgcn.workitem.id.z() #0
+
+define amdgpu_ps i32 @test_workitem_id_x() #1 {
+; CHECK-LABEL: test_workitem_id_x:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    s_mov_b32 s0, 0
+; CHECK-NEXT:    ; return to shader part epilog
+  %id = call i32 @llvm.amdgcn.workitem.id.x()
+  ret i32 %id
+}
+
+define amdgpu_ps i32 @test_workitem_id_y() #1 {
+; CHECK-LABEL: test_workitem_id_y:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    s_mov_b32 s0, 0
+; CHECK-NEXT:    ; return to shader part epilog
+  %id = call i32 @llvm.amdgcn.workitem.id.y()
+  ret i32 %id
+}
+
+define amdgpu_ps i32 @test_workitem_id_z() #1 {
+; CHECK-LABEL: test_workitem_id_z:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    s_mov_b32 s0, 0
+; CHECK-NEXT:    ; return to shader part epilog
+  %id = call i32 @llvm.amdgcn.workitem.id.z()
+  ret i32 %id
+}
+
+attributes #0 = { nounwind readnone }
+attributes #1 = { nounwind }