[IR][AArch64] Add 'swiftcoro' parameter attribute.

ahmedbougacha · ahmedbougacha · commit d257da7cddb9 · 2025-03-28T13:31:48.000-07:00
It doesn't have any really interesting treatment, other than
being passed in a fixed register.

In most of our AArch64 calling conventions, that's X23.

In effect, this is mostly similar to swiftself.

rdar://135984630
diff --git a/llvm/include/llvm/Bitcode/LLVMBitCodes.h b/llvm/include/llvm/Bitcode/LLVMBitCodes.h
@@ -758,6 +758,7 @@ enum AttributeKindCodes {
   ATTR_KIND_SANITIZE_NUMERICAL_STABILITY = 93,
   ATTR_KIND_INITIALIZES = 94,
   ATTR_KIND_HYBRID_PATCHABLE = 95,
+  ATTR_KIND_SWIFT_CORO = 103,
 };
 
 enum ComdatSelectionKindCodes {
diff --git a/llvm/include/llvm/CodeGen/TargetCallingConv.h b/llvm/include/llvm/CodeGen/TargetCallingConv.h
@@ -40,6 +40,7 @@ namespace ISD {
     unsigned IsSplitEnd : 1;   ///< Last part of a split
     unsigned IsSwiftSelf : 1;  ///< Swift self parameter
     unsigned IsSwiftAsync : 1;  ///< Swift async context parameter
+    unsigned IsSwiftCoro : 1;  ///< Swift coro parameter
     unsigned IsSwiftError : 1; ///< Swift error parameter
     unsigned IsCFGuardTarget : 1; ///< Control Flow Guard target
     unsigned IsHva : 1;        ///< HVA field for
@@ -63,6 +64,7 @@ namespace ISD {
         : IsZExt(0), IsSExt(0), IsInReg(0), IsSRet(0), IsByVal(0), IsByRef(0),
           IsNest(0), IsReturned(0), IsSplit(0), IsInAlloca(0),
           IsPreallocated(0), IsSplitEnd(0), IsSwiftSelf(0), IsSwiftAsync(0),
+          IsSwiftCoro(0),
           IsSwiftError(0), IsCFGuardTarget(0), IsHva(0), IsHvaStart(0),
           IsSecArgPass(0), MemAlign(0), OrigAlign(0),
           IsInConsecutiveRegsLast(0), IsInConsecutiveRegs(0),
@@ -100,6 +102,9 @@ namespace ISD {
     bool isSwiftAsync() const { return IsSwiftAsync; }
     void setSwiftAsync() { IsSwiftAsync = 1; }
 
+    bool isSwiftCoro() const { return IsSwiftCoro; }
+    void setSwiftCoro() { IsSwiftCoro = 1; }
+
     bool isSwiftError() const { return IsSwiftError; }
     void setSwiftError() { IsSwiftError = 1; }
 
diff --git a/llvm/include/llvm/CodeGen/TargetLowering.h b/llvm/include/llvm/CodeGen/TargetLowering.h
@@ -313,6 +313,7 @@ class TargetLoweringBase {
     bool IsSwiftSelf : 1;
     bool IsSwiftAsync : 1;
     bool IsSwiftError : 1;
+    bool IsSwiftCoro : 1;
     bool IsCFGuardTarget : 1;
     MaybeAlign Alignment = std::nullopt;
     Type *IndirectType = nullptr;
@@ -321,7 +322,8 @@ class TargetLoweringBase {
         : IsSExt(false), IsZExt(false), IsInReg(false), IsSRet(false),
           IsNest(false), IsByVal(false), IsByRef(false), IsInAlloca(false),
           IsPreallocated(false), IsReturned(false), IsSwiftSelf(false),
-          IsSwiftAsync(false), IsSwiftError(false), IsCFGuardTarget(false) {}
+          IsSwiftAsync(false), IsSwiftError(false), IsSwiftCoro(false),
+          IsCFGuardTarget(false) {}
 
     void setAttributes(const CallBase *Call, unsigned ArgIdx);
   };
diff --git a/llvm/include/llvm/IR/Attributes.td b/llvm/include/llvm/IR/Attributes.td
@@ -315,6 +315,9 @@ def SwiftSelf : EnumAttr<"swiftself", [ParamAttr]>;
 /// Argument is swift async context.
 def SwiftAsync : EnumAttr<"swiftasync", [ParamAttr]>;
 
+/// Argument is swift coro allocator.
+def SwiftCoro : EnumAttr<"swiftcoro", [ParamAttr]>;
+
 /// Function must be in a unwind table.
 def UWTable : IntAttr<"uwtable", [FnAttr]>;
 
diff --git a/llvm/include/llvm/Target/TargetCallingConv.td b/llvm/include/llvm/Target/TargetCallingConv.td
@@ -61,6 +61,11 @@ class CCIfSwiftAsync<CCAction A> : CCIf<"ArgFlags.isSwiftAsync()", A> {
 class CCIfSwiftError<CCAction A> : CCIf<"ArgFlags.isSwiftError()", A> {
 }
 
+/// CCIfSwiftCoro - If the current argument has swiftcoro parameter attribute,
+/// apply Action A.
+class CCIfSwiftCoro<CCAction A> : CCIf<"ArgFlags.isSwiftCoro()", A> {
+}
+
 /// CCIfCFGuardTarget - If the current argument has cfguardtarget parameter
 /// attribute, apply Action A.
 class CCIfCFGuardTarget<CCAction A> : CCIf<"ArgFlags.isCFGuardTarget()", A> {
diff --git a/llvm/lib/Bitcode/Reader/BitcodeReader.cpp b/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
@@ -2149,6 +2149,8 @@ static Attribute::AttrKind getAttrFromCode(uint64_t Code) {
     return Attribute::SwiftSelf;
   case bitc::ATTR_KIND_SWIFT_ASYNC:
     return Attribute::SwiftAsync;
+  case bitc::ATTR_KIND_SWIFT_CORO:
+    return Attribute::SwiftCoro;
   case bitc::ATTR_KIND_UW_TABLE:
     return Attribute::UWTable;
   case bitc::ATTR_KIND_VSCALE_RANGE:
diff --git a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
@@ -852,6 +852,8 @@ static uint64_t getAttrKindEncoding(Attribute::AttrKind Kind) {
     return bitc::ATTR_KIND_SWIFT_SELF;
   case Attribute::SwiftAsync:
     return bitc::ATTR_KIND_SWIFT_ASYNC;
+  case Attribute::SwiftCoro:
+    return bitc::ATTR_KIND_SWIFT_CORO;
   case Attribute::UWTable:
     return bitc::ATTR_KIND_UW_TABLE;
   case Attribute::VScaleRange:
diff --git a/llvm/lib/CodeGen/GlobalISel/CallLowering.cpp b/llvm/lib/CodeGen/GlobalISel/CallLowering.cpp
@@ -63,6 +63,8 @@ addFlagsUsingAttrFn(ISD::ArgFlagsTy &Flags,
     Flags.setSwiftAsync();
   if (AttrFn(Attribute::SwiftError))
     Flags.setSwiftError();
+  if (AttrFn(Attribute::SwiftCoro))
+    Flags.setSwiftCoro();
 }
 
 ISD::ArgFlagsTy CallLowering::getAttributesForArgIdx(const CallBase &Call,
diff --git a/llvm/lib/CodeGen/SelectionDAG/FastISel.cpp b/llvm/lib/CodeGen/SelectionDAG/FastISel.cpp
@@ -1053,6 +1053,8 @@ bool FastISel::lowerCallTo(CallLoweringInfo &CLI) {
       Flags.setSwiftAsync();
     if (Arg.IsSwiftError)
       Flags.setSwiftError();
+    if (Arg.IsSwiftCoro)
+      Flags.setSwiftCoro();
     if (Arg.IsCFGuardTarget)
       Flags.setCFGuardTarget();
     if (Arg.IsByVal)
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -10860,6 +10860,7 @@ TargetLowering::LowerCallTo(TargetLowering::CallLoweringInfo &CLI) const {
     Entry.IsSwiftSelf = false;
     Entry.IsSwiftAsync = false;
     Entry.IsSwiftError = false;
+    Entry.IsSwiftCoro = false;
     Entry.IsCFGuardTarget = false;
     Entry.Alignment = Alignment;
     CLI.getArgs().insert(CLI.getArgs().begin(), Entry);
@@ -10977,6 +10978,8 @@ TargetLowering::LowerCallTo(TargetLowering::CallLoweringInfo &CLI) const {
         Flags.setSwiftAsync();
       if (Args[i].IsSwiftError)
         Flags.setSwiftError();
+      if (Args[i].IsSwiftCoro)
+        Flags.setSwiftCoro();
       if (Args[i].IsCFGuardTarget)
         Flags.setCFGuardTarget();
       if (Args[i].IsByVal)
@@ -11522,6 +11525,8 @@ void SelectionDAGISel::LowerArguments(const Function &F) {
         Flags.setSwiftAsync();
       if (Arg.hasAttribute(Attribute::SwiftError))
         Flags.setSwiftError();
+      if (Arg.hasAttribute(Attribute::SwiftCoro))
+        Flags.setSwiftCoro();
       if (Arg.hasAttribute(Attribute::ByVal))
         Flags.setByVal();
       if (Arg.hasAttribute(Attribute::ByRef))
diff --git a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
@@ -123,6 +123,7 @@ void TargetLoweringBase::ArgListEntry::setAttributes(const CallBase *Call,
   IsSwiftSelf = Call->paramHasAttr(ArgIdx, Attribute::SwiftSelf);
   IsSwiftAsync = Call->paramHasAttr(ArgIdx, Attribute::SwiftAsync);
   IsSwiftError = Call->paramHasAttr(ArgIdx, Attribute::SwiftError);
+  IsSwiftCoro = Call->paramHasAttr(ArgIdx, Attribute::SwiftCoro);
   Alignment = Call->getParamStackAlign(ArgIdx);
   IndirectType = nullptr;
   assert(IsByVal + IsPreallocated + IsInAlloca + IsSRet <= 1 &&
diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
@@ -2185,6 +2185,7 @@ void Verifier::verifyFunctionAttrs(FunctionType *FT, AttributeList Attrs,
   bool SawSwiftSelf = false;
   bool SawSwiftAsync = false;
   bool SawSwiftError = false;
+  bool SawSwiftCoro = false;
 
   // Verify return value attributes.
   AttributeSet RetAttrs = Attrs.getRetAttrs();
@@ -2261,6 +2262,11 @@ void Verifier::verifyFunctionAttrs(FunctionType *FT, AttributeList Attrs,
       SawSwiftError = true;
     }
 
+    if (ArgAttrs.hasAttribute(Attribute::SwiftCoro)) {
+      Check(!SawSwiftCoro, "Cannot have multiple 'swiftcoro' parameters!", V);
+      SawSwiftCoro = true;
+    }
+
     if (ArgAttrs.hasAttribute(Attribute::InAlloca)) {
       Check(i == FT->getNumParams() - 1,
             "inalloca isn't on the last parameter!", V);
@@ -3878,6 +3884,7 @@ static AttrBuilder getParameterABIAttributes(LLVMContext& C, unsigned I, Attribu
   static const Attribute::AttrKind ABIAttrs[] = {
       Attribute::StructRet,  Attribute::ByVal,          Attribute::InAlloca,
       Attribute::InReg,      Attribute::StackAlignment, Attribute::SwiftSelf,
+      Attribute::SwiftCoro,
       Attribute::SwiftAsync, Attribute::SwiftError,     Attribute::Preallocated,
       Attribute::ByRef};
   AttrBuilder Copy(C);
diff --git a/llvm/lib/Target/AArch64/AArch64CallingConvention.td b/llvm/lib/Target/AArch64/AArch64CallingConvention.td
@@ -72,6 +72,9 @@ defvar AArch64_Common = [
   // preserved for normal function calls.
   CCIfSwiftAsync<CCIfType<[i64], CCAssignToReg<[X22]>>>,
 
+  // Pass SwiftCoro in X23.
+  CCIfSwiftCoro<CCIfType<[i64], CCAssignToReg<[X23]>>>,
+
   CCIfConsecutiveRegs<CCCustom<"CC_AArch64_Custom_Block">>,
 
   CCIfType<[nxv16i8, nxv8i16, nxv4i32, nxv2i64, nxv2f16, nxv4f16, nxv8f16,
@@ -241,6 +244,9 @@ def CC_AArch64_Arm64EC_Thunk : CallingConv<[
   // normal functions don't need to save it somewhere.
   CCIfSwiftAsync<CCIfType<[i64], CCAssignToReg<[X21]>>>,
 
+  // Pass SwiftCoro in X22.
+  CCIfSwiftCoro<CCIfType<[i64], CCAssignToReg<[X22]>>>,
+
   // The 'CFGuardTarget' parameter, if any, is passed in RAX (R8).
   CCIfCFGuardTarget<CCAssignToReg<[X8]>>,
 
@@ -374,6 +380,9 @@ def CC_AArch64_DarwinPCS : CallingConv<[
   // preserved for normal function calls.
   CCIfSwiftAsync<CCIfType<[i64], CCAssignToReg<[X22]>>>,
 
+  // Pass SwiftCoro in X23.
+  CCIfSwiftCoro<CCIfType<[i64], CCAssignToReg<[X23]>>>,
+
   CCIfConsecutiveRegs<CCCustom<"CC_AArch64_Custom_Block">>,
 
   // Handle i1, i8, i16, i32, i64, f32, f64 and v2f64 by passing in registers,
diff --git a/llvm/lib/Target/AArch64/AArch64FastISel.cpp b/llvm/lib/Target/AArch64/AArch64FastISel.cpp
@@ -2930,6 +2930,7 @@ bool AArch64FastISel::fastLowerArguments() {
         Arg.hasAttribute(Attribute::SwiftSelf) ||
         Arg.hasAttribute(Attribute::SwiftAsync) ||
         Arg.hasAttribute(Attribute::SwiftError) ||
+        Arg.hasAttribute(Attribute::SwiftCoro) ||
         Arg.hasAttribute(Attribute::Nest))
       return false;
 
@@ -3197,6 +3198,7 @@ bool AArch64FastISel::fastLowerCall(CallLoweringInfo &CLI) {
 
   for (auto Flag : CLI.OutFlags)
     if (Flag.isInReg() || Flag.isSRet() || Flag.isNest() || Flag.isByVal() ||
+        Flag.isSwiftCoro() ||
         Flag.isSwiftSelf() || Flag.isSwiftAsync() || Flag.isSwiftError())
       return false;
 
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -7701,7 +7701,7 @@ SDValue AArch64TargetLowering::LowerFormalArguments(
   if (CallConv == CallingConv::PreserveNone) {
     for (const ISD::InputArg &I : Ins) {
       if (I.Flags.isSwiftSelf() || I.Flags.isSwiftError() ||
-          I.Flags.isSwiftAsync()) {
+          I.Flags.isSwiftAsync() || I.Flags.isSwiftCoro()) {
         MachineFunction &MF = DAG.getMachineFunction();
         DAG.getContext()->diagnose(DiagnosticInfoUnsupported(
             MF.getFunction(),
@@ -8996,7 +8996,7 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
   if (CallConv == CallingConv::PreserveNone) {
     for (const ISD::OutputArg &O : Outs) {
       if (O.Flags.isSwiftSelf() || O.Flags.isSwiftError() ||
-          O.Flags.isSwiftAsync()) {
+          O.Flags.isSwiftAsync() || O.Flags.isSwiftCoro()) {
         MachineFunction &MF = DAG.getMachineFunction();
         DAG.getContext()->diagnose(DiagnosticInfoUnsupported(
             MF.getFunction(),
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp
@@ -1107,6 +1107,8 @@ WebAssemblyTargetLowering::LowerCall(CallLoweringInfo &CLI,
     SDValue &OutVal = OutVals[I];
     HasSwiftSelfArg |= Out.Flags.isSwiftSelf();
     HasSwiftErrorArg |= Out.Flags.isSwiftError();
+    if (Out.Flags.isSwiftCoro())
+      fail(DL, DAG, "WebAssembly hasn't implemented swiftcoro arguments");
     if (Out.Flags.isNest())
       fail(DL, DAG, "WebAssembly hasn't implemented nest arguments");
     if (Out.Flags.isInAlloca())
diff --git a/llvm/lib/Transforms/Utils/CodeExtractor.cpp b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
@@ -995,6 +995,7 @@ Function *CodeExtractor::constructFunction(const ValueSet &inputs,
       case Attribute::SwiftError:
       case Attribute::SwiftSelf:
       case Attribute::SwiftAsync:
+      case Attribute::SwiftCoro:
       case Attribute::ZExt:
       case Attribute::ImmArg:
       case Attribute::ByRef:
diff --git a/llvm/test/Bitcode/attributes.ll b/llvm/test/Bitcode/attributes.ll
@@ -292,6 +292,12 @@ define void @f50(ptr swiftself %0)
   ret void;
 }
 
+; CHECK: define void @swiftcoro(ptr swiftcoro %0)
+define void @swiftcoro(ptr swiftcoro %0)
+{
+  ret void;
+}
+
 ; CHECK: define i32 @f51(ptr swifterror %0)
 define i32 @f51(ptr swifterror %0)
 {
diff --git a/llvm/test/Bitcode/compatibility.ll b/llvm/test/Bitcode/compatibility.ll
@@ -516,10 +516,10 @@ declare cc96 void @f.cc96()
 ; CHECK: declare amdgpu_es void @f.cc96()
 declare amdgpu_es void @f.amdgpu_es()
 ; CHECK: declare amdgpu_es void @f.amdgpu_es()
-+declare cc124 void @f.cc124()
-+; CHECK: declare swiftcorocc void @f.cc124()
-+declare swiftcorocc void @f.swiftcorocc()
-+; CHECK: declare swiftcorocc void @f.swiftcorocc()
+declare cc124 void @f.cc124()
+; CHECK: declare swiftcorocc void @f.cc124()
+declare swiftcorocc void @f.swiftcorocc()
+; CHECK: declare swiftcorocc void @f.swiftcorocc()
 declare cc1023 void @f.cc1023()
 ; CHECK: declare cc1023 void @f.cc1023()
 
@@ -582,6 +582,8 @@ declare void @f.param.swiftasync(ptr swiftasync)
 ; CHECK: declare void @f.param.swiftasync(ptr swiftasync)
 declare void @f.param.swifterror(ptr swifterror)
 ; CHECK: declare void @f.param.swifterror(ptr swifterror)
+declare void @f.param.swiftcoro(ptr swiftcoro)
+; CHECK: declare void @f.param.swiftcoro(ptr swiftcoro)
 declare void @f.param.allocalign(i32 allocalign)
 ; CHECK: declare void @f.param.allocalign(i32 allocalign)
 declare void @f.param.allocptr(ptr allocptr)
diff --git a/llvm/test/CodeGen/AArch64/swiftcoro.ll b/llvm/test/CodeGen/AArch64/swiftcoro.ll
@@ -0,0 +1,70 @@
+; RUN: llc -verify-machineinstrs -mtriple=aarch64-apple-ios -o - %s | FileCheck --check-prefix=CHECK --check-prefix=OPT --check-prefix=OPTAARCH64 %s
+; RUN: llc -O0 -fast-isel -verify-machineinstrs -mtriple=aarch64-apple-ios -o - %s | FileCheck %s
+; RUN: llc -verify-machineinstrs -mtriple=aarch64-unknown-linux-gnu -o - %s | FileCheck --check-prefix=CHECK --check-prefix=OPT --check-prefix=OPTAARCH64 %s
+; RUN: llc -verify-machineinstrs -mtriple=arm64_32-apple-ios -o - %s | FileCheck --check-prefix=CHECK --check-prefix=OPT --check-prefix=OPTARM64_32 %s
+
+; Parameter with swiftcoro should be allocated to x23.
+; CHECK-LABEL: swiftcoro_param:
+; CHECK: mov x0, x23
+; CHECK-NEXT: ret
+define ptr @swiftcoro_param(ptr swiftcoro %addr0) {
+  ret ptr %addr0
+}
+
+; Check that x23 is used to pass a swiftcoro argument.
+; CHECK-LABEL: call_swiftcoro:
+; CHECK: mov x23, x0
+; CHECK: bl {{_?}}swiftcoro_param
+; CHECK: ret
+define ptr @call_swiftcoro(ptr %arg) {
+  %res = call ptr @swiftcoro_param(ptr swiftcoro %arg)
+  ret ptr %res
+}
+
+; x23 should be saved by the callee even if used for swiftcoro
+; CHECK-LABEL: swiftcoro_clobber:
+; CHECK: {{stp|str}} {{.*}}x23{{.*}}sp
+; ...
+; CHECK: {{ldp|ldr}} {{.*}}x23{{.*}}sp
+; CHECK: ret
+define ptr @swiftcoro_clobber(ptr swiftcoro %addr0) {
+  call void asm sideeffect "", "~{x23}"()
+  ret ptr %addr0
+}
+
+; Demonstrate that we do not need any movs when calling multiple functions
+; with swiftcoro argument.
+; CHECK-LABEL: swiftcoro_passthrough:
+; OPT-NOT: mov{{.*}}x23
+; OPT: bl {{_?}}swiftcoro_param
+; OPT-NOT: mov{{.*}}x23
+; OPT-NEXT: bl {{_?}}swiftcoro_param
+; OPT: ret
+define void @swiftcoro_passthrough(ptr swiftcoro %addr0) {
+  call ptr @swiftcoro_param(ptr swiftcoro %addr0)
+  call ptr @swiftcoro_param(ptr swiftcoro %addr0)
+  ret void
+}
+
+; We can use a tail call if the callee swiftcoro is the same as the caller one.
+; This should also work with fast-isel.
+; CHECK-LABEL: swiftcoro_tail:
+; OPTAARCH64: b {{_?}}swiftcoro_param
+; OPTAARCH64-NOT: ret
+; OPTARM64_32: b {{_?}}swiftcoro_param
+define ptr @swiftcoro_tail(ptr swiftcoro %addr0) {
+  call void asm sideeffect "", "~{x23}"()
+  %res = musttail call ptr @swiftcoro_param(ptr swiftcoro %addr0)
+  ret ptr %res
+}
+
+; We can not use a tail call if the callee swiftcoro is not the same as the
+; caller one.
+; CHECK-LABEL: swiftcoro_notail:
+; CHECK: mov x23, x0
+; CHECK: bl {{_?}}swiftcoro_param
+; CHECK: ret
+define ptr @swiftcoro_notail(ptr swiftcoro %addr0, ptr %addr1) nounwind {
+  %res = tail call ptr @swiftcoro_param(ptr swiftcoro %addr1)
+  ret ptr %res
+}
diff --git a/llvm/test/Verifier/swiftcoro.ll b/llvm/test/Verifier/swiftcoro.ll
@@ -0,0 +1,4 @@
+; RUN: not llvm-as %s -o /dev/null 2>&1 | FileCheck %s
+
+declare void @a(ptr swiftcoro %a, ptr swiftcoro %b)
+; CHECK: Cannot have multiple 'swiftcoro' parameters!

Original file line number	Diff line number	Diff line change
`@@ -63,6 +63,8 @@ addFlagsUsingAttrFn(ISD::ArgFlagsTy &Flags,`
`63`	`63`	`Flags.setSwiftAsync();`
`64`	`64`	`if (AttrFn(Attribute::SwiftError))`
`65`	`65`	`Flags.setSwiftError();`
	`66`	`+ if (AttrFn(Attribute::SwiftCoro))`
	`67`	`+ Flags.setSwiftCoro();`
`66`	`68`	`}`
`67`	`69`
`68`	`70`	`ISD::ArgFlagsTy CallLowering::getAttributesForArgIdx(const CallBase &Call,`