AMDGPU: Teach isArgPassedInSGPR() about cs_chain* calling convention

[ruiling]

This cs_chain and cs_chain_preserve use InReg attribute to indicate argument passed through SGPR.

[llvmbot]

@llvm/pr-subscribers-backend-amdgpu

@llvm/pr-subscribers-llvm-analysis

Changes

This cs_chain and cs_chain_preserve use InReg attribute to indicate argument passed through SGPR.

---

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

2 Files Affected:

• (modified) llvm/lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp (+4)
• (modified) llvm/test/Analysis/UniformityAnalysis/AMDGPU/kernel-args.ll (+21)

diff --git a/llvm/lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp b/llvm/lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp
index 13ef48a284beb61..f624cc53a952082 100644
--- a/llvm/lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp
@@ -2531,6 +2531,8 @@ bool isArgPassedInSGPR(const Argument *A) {
   case CallingConv::AMDGPU_PS:
   case CallingConv::AMDGPU_CS:
   case CallingConv::AMDGPU_Gfx:
+  case CallingConv::AMDGPU_CS_Chain:
+  case CallingConv::AMDGPU_CS_ChainPreserve:
     // For non-compute shaders, SGPR inputs are marked with either inreg or
     // byval. Everything else is in VGPRs.
     return A->hasAttribute(Attribute::InReg) ||
@@ -2556,6 +2558,8 @@ bool isArgPassedInSGPR(const CallBase *CB, unsigned ArgNo) {
   case CallingConv::AMDGPU_PS:
   case CallingConv::AMDGPU_CS:
   case CallingConv::AMDGPU_Gfx:
+  case CallingConv::AMDGPU_CS_Chain:
+  case CallingConv::AMDGPU_CS_ChainPreserve:
     // For non-compute shaders, SGPR inputs are marked with either inreg or
     // byval. Everything else is in VGPRs.
     return CB->paramHasAttr(ArgNo, Attribute::InReg) ||
diff --git a/llvm/test/Analysis/UniformityAnalysis/AMDGPU/kernel-args.ll b/llvm/test/Analysis/UniformityAnalysis/AMDGPU/kernel-args.ll
index 7473353d054d1db..4b014f969257bdb 100644
--- a/llvm/test/Analysis/UniformityAnalysis/AMDGPU/kernel-args.ll
+++ b/llvm/test/Analysis/UniformityAnalysis/AMDGPU/kernel-args.ll
@@ -36,4 +36,25 @@ define void @test_c(ptr addrspace(5) byval([4 x <16 x i8>]) %arg0, float inreg %
   ret void
 }
 
+; CHECK-LABEL: for function 'test_amdgpu_cs_chain':
+; CHECK-DAG: DIVERGENT:  ptr addrspace(4) %arg0
+; CHECK-DAG: DIVERGENT:  <2 x i32> %arg3
+; CHECK-DAG: DIVERGENT:  <3 x i32> %arg4
+; CHECK-DAG: DIVERGENT:  float %arg5
+; CHECK-DAG: DIVERGENT:  i32 %arg6
+; CHECK-NOT: DIVERGENT
+define amdgpu_cs_chain void @test_amdgpu_cs_chain(ptr addrspace(4) byref([4 x <16 x i8>]) %arg0, float inreg %arg1, i32 inreg %arg2, <2 x i32> %arg3, <3 x i32> %arg4, float %arg5, i32 %arg6) #0 {
+  ret void
+}
+
+; CHECK-LABEL: for function 'test_amdgpu_cs_chain_preserve':
+; CHECK-DAG: DIVERGENT:  ptr addrspace(4) %arg0
+; CHECK-DAG: DIVERGENT:  <2 x i32> %arg3
+; CHECK-DAG: DIVERGENT:  <3 x i32> %arg4
+; CHECK-DAG: DIVERGENT:  float %arg5
+; CHECK-DAG: DIVERGENT:  i32 %arg6
+; CHECK-NOT: DIVERGENT
+define amdgpu_cs_chain_preserve void @test_amdgpu_cs_chain_preserve(ptr addrspace(4) byref([4 x <16 x i8>]) %arg0, float inreg %arg1, i32 inreg %arg2, <2 x i32> %arg3, <3 x i32> %arg4, float %arg5, i32 %arg6) #0 {
+  ret void
+}
 attributes #0 = { nounwind }