[AMDGPU] Add amdgpu_cs_chain[_preserve] CCs to IR & verifier

[rovka]

Summary

Add the amdgpu_cs_chain and amdgpu_cs_chain_preserve keywords to
LLVM IR and make sure we can parse and print them. Also make sure we
perform some basic checks in the IR verifier - similar to what we check
for many of the other AMDGPU calling conventions, plus the additional
restriction that we can't have direct calls to functions with these
calling conventions.

[foad]

Please upload the patch with full context.

llvm/lib/IR/Verifier.cpp

3292

Should check Call.getCallingConv(), instead of relying on whether or not we found a Callee function.

[arsenm]

llvm/test/Verifier/amdgpu-cc.ll

217

Is a regular pointer to the stack allowed?

[nhaehnle]

llvm/test/Verifier/amdgpu-cc.ll

217

You mean a ptr addrspace(5)? Yes, it's allowed, and we plan to use it.

It's technically not going to be a *stack* pointer as far as LLVM is concerned, because the stack is "dormant"/"inexistent" between chain functions. The plan is to split the wave's scratch allocation into two parts, the bottom part under the control of LLVM and the top part under the control of the frontend, and so passing pointers to that top part can still be meaningful even though the stack "doesn't exist" at chain function boundaries as far as LLVM is concerned.

In principle, we also allow pointers to other address spaces, although I don't see us using those at least for the time being.

[rovka]

Address review comments.

[arsenm]

llvm/include/llvm/IR/CallingConv.h

244–245

Should say how the _Preserve variant differs

[rovka]

Mention what amdgpu_cs_chain_preserve preserves.

[rovka]

Any other comments or is this good to go?

[arsenm]

llvm/test/Verifier/amdgpu-cc.ll

233

Check a call to a function pointer

llvm/include/llvm/AsmParser/LLToken.h

enum Kind {
   kw_cxx_fast_tlscc,
   kw_amdgpu_vs,
   kw_amdgpu_ls,
   kw_amdgpu_hs,
   kw_amdgpu_es,
   kw_amdgpu_gs,
   kw_amdgpu_ps,
   kw_amdgpu_cs,
+  kw_amdgpu_cs_chain,
+  kw_amdgpu_cs_chain_preserve,
   kw_amdgpu_kernel,
   kw_amdgpu_gfx,
   kw_tailcc,
 
   // Attributes:
   kw_attributes,
   kw_sync,
   kw_async,

llvm/include/llvm/IR/CallingConv.h

enum {
     M68k_INTR = 101,
 
     /// Preserve X0-X13, X19-X29, SP, Z0-Z31, P0-P15.
     AArch64_SME_ABI_Support_Routines_PreserveMost_From_X0 = 102,
 
     /// Preserve X2-X15, X19-X29, SP, Z0-Z31, P0-P15.
     AArch64_SME_ABI_Support_Routines_PreserveMost_From_X2 = 103,
 
+    /// Used on AMDGPUs to give the middle-end more control over argument
+    /// placement.
+    AMDGPU_CS_Chain = 104,
+
+    /// Used on AMDGPUs to give the middle-end more control over argument
+    /// placement. Preserves active lane values for input VGPRs.
+    AMDGPU_CS_ChainPreserve = 105,
+
     /// The highest possible ID. Must be some 2^k - 1.
     MaxID = 1023
   };
 
 } // end namespace CallingConv
 
 } // end namespace llvm
 
 #endif // LLVM_IR_CALLINGCONV_H

llvm/lib/AsmParser/LLLexer.cpp

#define KEYWORD(STR)                                                           \
   KEYWORD(cxx_fast_tlscc);
   KEYWORD(amdgpu_vs);
   KEYWORD(amdgpu_ls);
   KEYWORD(amdgpu_hs);
   KEYWORD(amdgpu_es);
   KEYWORD(amdgpu_gs);
   KEYWORD(amdgpu_ps);
   KEYWORD(amdgpu_cs);
+  KEYWORD(amdgpu_cs_chain);
+  KEYWORD(amdgpu_cs_chain_preserve);
   KEYWORD(amdgpu_kernel);
   KEYWORD(amdgpu_gfx);
   KEYWORD(tailcc);
 
   KEYWORD(cc);
   KEYWORD(c);
 
   KEYWORD(attributes);

llvm/lib/AsmParser/LLParser.cpp

 ///   ::= 'cxx_fast_tlscc'
 ///   ::= 'amdgpu_vs'
 ///   ::= 'amdgpu_ls'
 ///   ::= 'amdgpu_hs'
 ///   ::= 'amdgpu_es'
 ///   ::= 'amdgpu_gs'
 ///   ::= 'amdgpu_ps'
 ///   ::= 'amdgpu_cs'
+///   ::= 'amdgpu_cs_chain'
+///   ::= 'amdgpu_cs_chain_preserve'
 ///   ::= 'amdgpu_kernel'
 ///   ::= 'tailcc'
 ///   ::= 'cc' UINT
 ///
 bool LLParser::parseOptionalCallingConv(unsigned &CC) {
   switch (Lex.getKind()) {
   default:                       CC = CallingConv::C; return false;
   case lltok::kw_ccc:            CC = CallingConv::C; break;
bool LLParser::parseOptionalCallingConv(unsigned &CC) {
   case lltok::kw_amdgpu_vs:      CC = CallingConv::AMDGPU_VS; break;
   case lltok::kw_amdgpu_gfx:     CC = CallingConv::AMDGPU_Gfx; break;
   case lltok::kw_amdgpu_ls:      CC = CallingConv::AMDGPU_LS; break;
   case lltok::kw_amdgpu_hs:      CC = CallingConv::AMDGPU_HS; break;
   case lltok::kw_amdgpu_es:      CC = CallingConv::AMDGPU_ES; break;
   case lltok::kw_amdgpu_gs:      CC = CallingConv::AMDGPU_GS; break;
   case lltok::kw_amdgpu_ps:      CC = CallingConv::AMDGPU_PS; break;
   case lltok::kw_amdgpu_cs:      CC = CallingConv::AMDGPU_CS; break;
+  case lltok::kw_amdgpu_cs_chain:
+    CC = CallingConv::AMDGPU_CS_Chain;
+    break;
+  case lltok::kw_amdgpu_cs_chain_preserve:
+    CC = CallingConv::AMDGPU_CS_ChainPreserve;
+    break;
   case lltok::kw_amdgpu_kernel:  CC = CallingConv::AMDGPU_KERNEL; break;
   case lltok::kw_tailcc:         CC = CallingConv::Tail; break;
   case lltok::kw_cc: {
       Lex.Lex();
       return parseUInt32(CC);
     }
   }
 

llvm/lib/IR/AsmWriter.cpp

case CallingConv::DUMMY_HHVM_C:
     break;
   case CallingConv::AMDGPU_VS:     Out << "amdgpu_vs"; break;
   case CallingConv::AMDGPU_LS:     Out << "amdgpu_ls"; break;
   case CallingConv::AMDGPU_HS:     Out << "amdgpu_hs"; break;
   case CallingConv::AMDGPU_ES:     Out << "amdgpu_es"; break;
   case CallingConv::AMDGPU_GS:     Out << "amdgpu_gs"; break;
   case CallingConv::AMDGPU_PS:     Out << "amdgpu_ps"; break;
   case CallingConv::AMDGPU_CS:     Out << "amdgpu_cs"; break;
+  case CallingConv::AMDGPU_CS_Chain:
+    Out << "amdgpu_cs_chain";
+    break;
+  case CallingConv::AMDGPU_CS_ChainPreserve:
+    Out << "amdgpu_cs_chain_preserve";
+    break;
   case CallingConv::AMDGPU_KERNEL: Out << "amdgpu_kernel"; break;
   case CallingConv::AMDGPU_Gfx:    Out << "amdgpu_gfx"; break;
   }
 }
 
 enum PrefixType {
   GlobalPrefix,
   ComdatPrefix,

llvm/lib/IR/Verifier.cpp

case CallingConv::C:
     break;
   case CallingConv::X86_INTR: {
     Check(F.arg_empty() || Attrs.hasParamAttr(0, Attribute::ByVal),
           "Calling convention parameter requires byval", &F);
     break;
   }
   case CallingConv::AMDGPU_KERNEL:
   case CallingConv::SPIR_KERNEL:
+  case CallingConv::AMDGPU_CS_Chain:
+  case CallingConv::AMDGPU_CS_ChainPreserve:
     Check(F.getReturnType()->isVoidTy(),
           "Calling convention requires void return type", &F);
     [[fallthrough]];
   case CallingConv::AMDGPU_VS:
   case CallingConv::AMDGPU_HS:
   case CallingConv::AMDGPU_GS:
   case CallingConv::AMDGPU_PS:
   case CallingConv::AMDGPU_CS:
void Verifier::visitCallBase(CallBase &Call) {
 
   Function *Callee =
       dyn_cast<Function>(Call.getCalledOperand()->stripPointerCasts());
   bool IsIntrinsic = Callee && Callee->isIntrinsic();
   if (IsIntrinsic)
     Check(Callee->getValueType() == FTy,
           "Intrinsic called with incompatible signature", Call);
 
+  // Disallow calls to functions with the amdgpu_cs_chain[_preserve] calling
+  // convention.
+  auto CC = Call.getCallingConv();
+  Check(CC != CallingConv::AMDGPU_CS_Chain &&
+            CC != CallingConv::AMDGPU_CS_ChainPreserve,
+        "Direct calls to amdgpu_cs_chain/amdgpu_cs_chain_preserve functions "
+        "not allowed. Please use the @llvm.amdgpu.cs.chain intrinsic instead.",
+        Call);
+
   auto VerifyTypeAlign = [&](Type *Ty, const Twine &Message) {
     if (!Ty->isSized())
       return;
     Align ABIAlign = DL.getABITypeAlign(Ty);
     Align MaxAlign(ParamMaxAlignment);
     Check(ABIAlign <= MaxAlign,
           "Incorrect alignment of " + Message + " to called function!", Call);
   };

llvm/test/Assembler/amdgpu-cs-chain-cc.ll

+; RUN: llvm-as < %s | llvm-dis | llvm-as | llvm-dis | FileCheck %s
+
+; CHECK: amdgpu_cs_chain void @amdgpu_cs_chain_cc
+define amdgpu_cs_chain void @amdgpu_cs_chain_cc() {
+entry:
+  ret void
+}
+
+; CHECK: amdgpu_cs_chain_preserve void @amdgpu_cs_chain_preserve_cc
+define amdgpu_cs_chain_preserve void @amdgpu_cs_chain_preserve_cc() {
+entry:
+  ret void
+}

llvm/test/Verifier/amdgpu-cc.ll

define amdgpu_kernel void @inalloca_as0_cc_amdgpu_kernel(ptr inalloca(i32) %ptr) {
   ret void
 }
 
 ; CHECK: Calling convention disallows stack byref
 ; CHECK-NEXT: ptr @byref_as5_cc_amdgpu_kernel
 define amdgpu_kernel void @byref_as5_cc_amdgpu_kernel(ptr addrspace(5) byref(i32) %ptr) {
   ret void
 }
+
+; CHECK: Calling convention requires void return type
+; CHECK-NEXT: ptr @nonvoid_cc_amdgpu_cs_chain
+define amdgpu_cs_chain i32 @nonvoid_cc_amdgpu_cs_chain() {
+  ret i32 0
+}
+
+; CHECK: Calling convention does not support varargs or perfect forwarding!
+; CHECK-NEXT: ptr @varargs_amdgpu_cs_chain
+define amdgpu_cs_chain void @varargs_amdgpu_cs_chain(...) {
+  ret void
+}
+
+; CHECK: Calling convention does not allow sret
+; CHECK-NEXT: ptr @sret_cc_amdgpu_cs_chain_as0
+define amdgpu_cs_chain void @sret_cc_amdgpu_cs_chain_as0(ptr sret(i32) %ptr) {
+  ret void
+}
+
+; CHECK: Calling convention disallows byval
+; CHECK-NEXT: ptr @byval_cc_amdgpu_cs_chain
+define amdgpu_cs_chain void @byval_cc_amdgpu_cs_chain(ptr addrspace(1) byval(i32) %ptr) {
+  ret void
+}
+
+; CHECK: Calling convention disallows stack byref
+; CHECK-NEXT: ptr @byref_cc_amdgpu_cs_chain
+define amdgpu_cs_chain void @byref_cc_amdgpu_cs_chain(ptr addrspace(5) byref(i32) %ptr) {
+  ret void
+}
+
+; CHECK: Calling convention disallows preallocated
+; CHECK-NEXT: ptr @preallocated_cc_amdgpu_cs_chain
+define amdgpu_cs_chain void @preallocated_cc_amdgpu_cs_chain(ptr preallocated(i32) %ptr) {
+  ret void
+}
+
+; CHECK: Calling convention disallows inalloca
+; CHECK-NEXT: ptr @inalloca_cc_amdgpu_cs_chain
+define amdgpu_cs_chain void @inalloca_cc_amdgpu_cs_chain(ptr inalloca(i32) %ptr) {
+  ret void
+}
+
+; CHECK: Calling convention requires void return type
+; CHECK-NEXT: ptr @nonvoid_cc_amdgpu_cs_chain_preserve
+define amdgpu_cs_chain_preserve i32 @nonvoid_cc_amdgpu_cs_chain_preserve() {
+  ret i32 0
+}
+
+; CHECK: Calling convention does not support varargs or perfect forwarding!
+; CHECK-NEXT: ptr @varargs_amdgpu_cs_chain_preserve
+define amdgpu_cs_chain_preserve void @varargs_amdgpu_cs_chain_preserve(...) {
+  ret void
+}
+
+; CHECK: Calling convention does not allow sret
+; CHECK-NEXT: ptr @sret_cc_amdgpu_cs_chain_preserve_as0
+define amdgpu_cs_chain_preserve void @sret_cc_amdgpu_cs_chain_preserve_as0(ptr sret(i32) %ptr) {
+  ret void
+}
+
+; CHECK: Calling convention does not allow sret
+; CHECK-NEXT: ptr @sret_cc_amdgpu_cs_chain_preserve
+define amdgpu_cs_chain_preserve void @sret_cc_amdgpu_cs_chain_preserve(ptr addrspace(5) sret(i32) %ptr) {
+  ret void
+}
+
+; CHECK: Calling convention disallows byval
+; CHECK-NEXT: ptr @byval_cc_amdgpu_cs_chain_preserve
+define amdgpu_cs_chain_preserve void @byval_cc_amdgpu_cs_chain_preserve(ptr addrspace(1) byval(i32) %ptr) {
+  ret void
+}
+
+; CHECK: Calling convention disallows stack byref
+; CHECK-NEXT: ptr @byref_cc_amdgpu_cs_chain_preserve
+define amdgpu_cs_chain_preserve void @byref_cc_amdgpu_cs_chain_preserve(ptr addrspace(5) byref(i32) %ptr) {
+  ret void
+}
+
+; CHECK: Calling convention disallows preallocated
+; CHECK-NEXT: ptr @preallocated_cc_amdgpu_cs_chain_preserve
+define amdgpu_cs_chain_preserve void @preallocated_cc_amdgpu_cs_chain_preserve(ptr preallocated(i32) %ptr) {
+  ret void
+}
+
+; CHECK: Calling convention disallows inalloca
+; CHECK-NEXT: ptr @inalloca_cc_amdgpu_cs_chain_preserve
+define amdgpu_cs_chain_preserve void @inalloca_cc_amdgpu_cs_chain_preserve(ptr inalloca(i32) %ptr) {
+  ret void
+}
+
+declare amdgpu_cs_chain void @amdgpu_cs_chain_call_target()
+declare amdgpu_cs_chain_preserve void @amdgpu_cs_chain_preserve_call_target()
+
+define amdgpu_cs_chain void @cant_call_amdgpu_cs_chain_functions(ptr %f) {
+  ; CHECK: Direct calls to amdgpu_cs_chain/amdgpu_cs_chain_preserve functions not allowed. Please use the @llvm.amdgpu.cs.chain intrinsic instead.
+  ; CHECK-NEXT: call amdgpu_cs_chain
+  call amdgpu_cs_chain void @amdgpu_cs_chain_call_target()
+
+  ; CHECK: Direct calls to amdgpu_cs_chain/amdgpu_cs_chain_preserve functions not allowed. Please use the @llvm.amdgpu.cs.chain intrinsic instead.
+  ; CHECK-NEXT: call amdgpu_cs_chain_preserve
+  call amdgpu_cs_chain_preserve void @amdgpu_cs_chain_preserve_call_target()
+
+  ; CHECK: Direct calls to amdgpu_cs_chain/amdgpu_cs_chain_preserve functions not allowed. Please use the @llvm.amdgpu.cs.chain intrinsic instead.
+  ; CHECK-NEXT: call amdgpu_cs_chain
+  call amdgpu_cs_chain void %f()
+
+  ; CHECK: Direct calls to amdgpu_cs_chain/amdgpu_cs_chain_preserve functions not allowed. Please use the @llvm.amdgpu.cs.chain intrinsic instead.
+  ; CHECK-NEXT: call amdgpu_cs_chain
+  call amdgpu_cs_chain_preserve void %f()
+
+  ret void
+}