[llvm] AMDGPU: Teach isArgPassedInSGPR() about cs_chain* calling convention (PR #67086)

[via llvm-commits]

https://github.com/ruiling created https://github.com/llvm/llvm-project/pull/67086

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

>From f86dbf30b0339f3fd073b962ebead984a6a36ea2 Mon Sep 17 00:00:00 2001
From: Ruiling Song <ruiling.song at amd.com>
Date: Fri, 22 Sep 2023 11:34:46 +0800
Subject: [PATCH] AMDGPU: Teach isArgPassedInSGPR() about cs_chain* calling
 convention

This cs_chain and cs_chain_preserve use InReg attribute to indicate
argument passed through SGPR.
---
 .../Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp    |  4 ++++
 .../UniformityAnalysis/AMDGPU/kernel-args.ll  | 21 +++++++++++++++++++
 2 files changed, 25 insertions(+)

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 }