[AMDGPU] Change handling of unsupported non-compute shaders with HSA (#126798)

Previous handling in `SITargetLowering::LowerFormalArguments` only
reported a diagnostic message and continued execution by returning a
non-usable `SDValue`. This results in llvm crashing later with an
unrelated error. This commit changes the detection of an unsupported
non-compute shader to be a fatal error right away.

As an example situation, take the usage of an `amdgpu_ps` function and
the `amdgcn-unknown-amdhsa` target triple.
```
define amdgpu_ps void @foo(ptr %p, i32 %i) {
        store i32 %i, ptr %p
        ret void
}
```
Compiling this code (with `llc -mtriple=amdgcn-unknown-amdhsa
-mcpu=gfx942`, for example) fails with:
```
error: <unknown>:0:0: in function foo void (ptr, i32): unsupported non-compute shaders with HSA

llc:
[...]/git/trunk21.0/llvm-project/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp:11790:
void llvm::SelectionDAGISel::LowerArguments(const llvm::Function&):
Assertion `InVals.size() == Ins.size() && "LowerFormalArguments didn't emit the correct number of values!"' failed.
[...]
```

Author: ro-i

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

@@ -2822,12 +2822,13 @@ SDValue SITargetLowering::LowerFormalArguments(
   const Function &Fn = MF.getFunction();
   FunctionType *FType = MF.getFunction().getFunctionType();
   SIMachineFunctionInfo *Info = MF.getInfo<SIMachineFunctionInfo>();
+  bool IsError = false;
 
   if (Subtarget->isAmdHsaOS() && AMDGPU::isGraphics(CallConv)) {
     DiagnosticInfoUnsupported NoGraphicsHSA(
         Fn, "unsupported non-compute shaders with HSA", DL.getDebugLoc());
     DAG.getContext()->diagnose(NoGraphicsHSA);
-    return DAG.getEntryNode();
+    IsError = true;
   }
 
   SmallVector<ISD::InputArg, 16> Splits;
@@ -2936,7 +2937,7 @@ SDValue SITargetLowering::LowerFormalArguments(
 
   for (unsigned i = 0, e = Ins.size(), ArgIdx = 0; i != e; ++i) {
     const ISD::InputArg &Arg = Ins[i];
-    if (Arg.isOrigArg() && Skipped[Arg.getOrigArgIndex()]) {
+    if ((Arg.isOrigArg() && Skipped[Arg.getOrigArgIndex()]) || IsError) {
       InVals.push_back(DAG.getUNDEF(Arg.VT));
       continue;
     }

llvm/test/CodeGen/AMDGPU/no-hsa-graphics-shaders.ll

@@ -1,19 +1,65 @@
-; RUN: not llc -mtriple=amdgcn-unknown-amdhsa < %s 2>&1 | FileCheck %s
+; RUN: not llc -mtriple=amdgcn-unknown-amdhsa -O0 -filetype=null < %s 2>&1 | FileCheck %s
 
-; CHECK: in function pixel_s{{.*}}: unsupported non-compute shaders with HSA
-define amdgpu_ps void @pixel_shader() #0 {
+@I = global i32 42
+@P = global ptr @I
+
+; CHECK: error: <unknown>:0:0: in function pixel_shader_zero_args void (): unsupported non-compute shaders with HSA
+; CHECK: error: <unknown>:0:0: in function pixel_shader_one_arg void (ptr): unsupported non-compute shaders with HSA
+; CHECK: error: <unknown>:0:0: in function pixel_shader_two_args void (ptr, i32): unsupported non-compute shaders with HSA
+; CHECK: error: <unknown>:0:0: in function vertex_shader_zero_args void (): unsupported non-compute shaders with HSA
+; CHECK: error: <unknown>:0:0: in function vertex_shader_one_arg void (ptr): unsupported non-compute shaders with HSA
+; CHECK: error: <unknown>:0:0: in function vertex_shader_two_args void (ptr, i32): unsupported non-compute shaders with HSA
+; CHECK: error: <unknown>:0:0: in function geometry_shader_zero_args void (): unsupported non-compute shaders with HSA
+; CHECK: error: <unknown>:0:0: in function geometry_shader_one_arg void (ptr): unsupported non-compute shaders with HSA
+; CHECK: error: <unknown>:0:0: in function geometry_shader_two_args void (ptr, i32): unsupported non-compute shaders with HSA
+
+define amdgpu_ps void @pixel_shader_zero_args() {
+  %i = load i32, ptr @I
+  store i32 %i, ptr @P
+  ret void
+}
+
+define amdgpu_ps void @pixel_shader_one_arg(ptr %p) {
+  %i = load i32, ptr @I
+  store i32 %i, ptr %p
   ret void
 }
 
-; CHECK: in function vertex_s{{.*}}: unsupported non-compute shaders with HSA
-define amdgpu_vs void @vertex_shader() #0 {
+define amdgpu_ps void @pixel_shader_two_args(ptr %p, i32 %i) {
+  store i32 %i, ptr %p
   ret void
 }
 
-; CHECK: in function geometry_s{{.*}}: unsupported non-compute shaders with HSA
-define amdgpu_gs void @geometry_shader() #0 {
+define amdgpu_vs void @vertex_shader_zero_args() {
+  %i = load i32, ptr @I
+  store i32 %i, ptr @P
   ret void
 }
 
-!llvm.module.flags = !{!0}
-!0 = !{i32 1, !"amdhsa_code_object_version", i32 400}
+define amdgpu_vs void @vertex_shader_one_arg(ptr %p) {
+  %i = load i32, ptr @I
+  store i32 %i, ptr %p
+  ret void
+}
+
+define amdgpu_vs void @vertex_shader_two_args(ptr %p, i32 %i) {
+  store i32 %i, ptr %p
+  ret void
+}
+
+define amdgpu_gs void @geometry_shader_zero_args() {
+  %i = load i32, ptr @I
+  store i32 %i, ptr @P
+  ret void
+}
+
+define amdgpu_gs void @geometry_shader_one_arg(ptr %p) {
+  %i = load i32, ptr @I
+  store i32 %i, ptr %p
+  ret void
+}
+
+define amdgpu_gs void @geometry_shader_two_args(ptr %p, i32 %i) {
+  store i32 %i, ptr %p
+  ret void
+}