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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Robert Imschweiler (ro-i)

<details>
<summary>Changes</summary>

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.

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


2 Files Affected:

- (modified) llvm/lib/Target/AMDGPU/SIISelLowering.cpp (+20-3) 
- (added) llvm/test/CodeGen/AMDGPU/llvm.amdgcn.workitem.id-undefined-attr.ll (+36) 


``````````diff
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 }

``````````

</details>


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