llvm · citymarina · May 8, 2025 · Mar 17, 2025
diff --git a/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp b/llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp
@@ -836,21 +836,22 @@ bool AArch64ExpandPseudo::expandCALL_RVMARKER(
     MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI) {
   // Expand CALL_RVMARKER pseudo to:
   // - a branch to the call target, followed by
-  // - the special `mov x29, x29` marker, and
+  // - the special `mov x29, x29` marker, if necessary, and
   // - another branch, to the runtime function
   // Mark the sequence as bundle, to avoid passes moving other code in between.
   MachineInstr &MI = *MBBI;
   MachineOperand &RVTarget = MI.getOperand(0);
+  bool DoEmitMarker = MI.getOperand(1).getImm();
   assert(RVTarget.isGlobal() && "invalid operand for attached call");
 
   MachineInstr *OriginalCall = nullptr;
 
   if (MI.getOpcode() == AArch64::BLRA_RVMARKER) {
     // ptrauth call.
-    const MachineOperand &CallTarget = MI.getOperand(1);
-    const MachineOperand &Key = MI.getOperand(2);
-    const MachineOperand &IntDisc = MI.getOperand(3);
-    const MachineOperand &AddrDisc = MI.getOperand(4);
+    const MachineOperand &CallTarget = MI.getOperand(2);
+    const MachineOperand &Key = MI.getOperand(3);
+    const MachineOperand &IntDisc = MI.getOperand(4);
+    const MachineOperand &AddrDisc = MI.getOperand(5);
 
     assert((Key.getImm() == AArch64PACKey::IA ||
             Key.getImm() == AArch64PACKey::IB) &&
@@ -859,19 +860,20 @@ bool AArch64ExpandPseudo::expandCALL_RVMARKER(
     MachineOperand Ops[] = {CallTarget, Key, IntDisc, AddrDisc};
 
     OriginalCall = createCallWithOps(MBB, MBBI, TII, AArch64::BLRA, Ops,
-                                     /*RegMaskStartIdx=*/5);
+                                     /*RegMaskStartIdx=*/6);
   } else {
     assert(MI.getOpcode() == AArch64::BLR_RVMARKER && "unknown rvmarker MI");
-    OriginalCall = createCall(MBB, MBBI, TII, MI.getOperand(1),
+    OriginalCall = createCall(MBB, MBBI, TII, MI.getOperand(2),
                               // Regmask starts after the RV and call targets.
-                              /*RegMaskStartIdx=*/2);
+                              /*RegMaskStartIdx=*/3);
   }
 
-  BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(AArch64::ORRXrs))
-                     .addReg(AArch64::FP, RegState::Define)
-                     .addReg(AArch64::XZR)
-                     .addReg(AArch64::FP)
-                     .addImm(0);
+  if (DoEmitMarker)
+    BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(AArch64::ORRXrs))
+        .addReg(AArch64::FP, RegState::Define)
+        .addReg(AArch64::XZR)
+        .addReg(AArch64::FP)
+        .addImm(0);
 
   auto *RVCall = BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(AArch64::BL))
                      .add(RVTarget)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -9540,6 +9540,13 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
     Function *ARCFn = *objcarc::getAttachedARCFunction(CLI.CB);
     auto GA = DAG.getTargetGlobalAddress(ARCFn, DL, PtrVT);
     Ops.insert(Ops.begin() + 1, GA);
+
+    // We may or may not need to emit both the marker and the retain/claim call.
+    // Do what the frontend tells us: if the rvmarker module flag is present,
+    // emit the marker.  Always emit the call regardless.
+    // Tell the pseudo expansion using an additional boolean op.
+    SDValue DoEmitMarker = DAG.getTargetConstant(true, DL, MVT::i32);
+    Ops.insert(Ops.begin() + 2, DoEmitMarker);
   } else if (CallConv == CallingConv::ARM64EC_Thunk_X64) {
     Opc = AArch64ISD::CALL_ARM64EC_TO_X64;
   } else if (GuardWithBTI) {

diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.td b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
@@ -754,10 +754,11 @@ def AArch64authtcret: SDNode<"AArch64ISD::AUTH_TC_RETURN",
 
 def AArch64authcall_rvmarker : SDNode<"AArch64ISD::AUTH_CALL_RVMARKER",
                                  SDTypeProfile<0, -1, [SDTCisPtrTy<0>,
-                                                       SDTCisPtrTy<1>,
-                                                       SDTCisVT<2, i32>,
-                                                       SDTCisVT<3, i64>,
-                                                       SDTCisVT<4, i64>]>,
+                                                       SDTCisVT<1, i32>,
+                                                       SDTCisPtrTy<2>,
+                                                       SDTCisVT<3, i32>,
+                                                       SDTCisVT<4, i64>,
+                                                       SDTCisVT<5, i64>]>,
                                  [SDNPHasChain, SDNPOptInGlue, SDNPOutGlue,
                                   SDNPVariadic]>;
 
@@ -1896,9 +1897,9 @@ let Predicates = [HasPAuth] in {
   }
 
   def BLRA_RVMARKER : Pseudo<
-        (outs), (ins i64imm:$rvfunc, GPR64noip:$Rn, i32imm:$Key, i64imm:$Disc,
-                     GPR64:$AddrDisc),
-        [(AArch64authcall_rvmarker tglobaladdr:$rvfunc,
+        (outs), (ins i64imm:$rvfunc, i32imm:$withmarker, GPR64noip:$Rn,
+                     i32imm:$Key, i64imm:$Disc, GPR64:$AddrDisc),
+        [(AArch64authcall_rvmarker tglobaladdr:$rvfunc, timm:$withmarker,
                                    GPR64noip:$Rn, timm:$Key, timm:$Disc,
                                    GPR64:$AddrDisc)]>, Sched<[]> {
     let isCodeGenOnly = 1;
@@ -3293,8 +3294,9 @@ def : Pat<(AArch64call GPR64noip:$Rn),
           (BLRNoIP GPR64noip:$Rn)>,
       Requires<[SLSBLRMitigation]>;
 
-def : Pat<(AArch64call_rvmarker (i64 tglobaladdr:$rvfunc), GPR64:$Rn),
-          (BLR_RVMARKER tglobaladdr:$rvfunc, GPR64:$Rn)>,
+def : Pat<(AArch64call_rvmarker (i64 tglobaladdr:$rvfunc),
+                                (i32 timm:$withmarker), GPR64:$Rn),
+          (BLR_RVMARKER tglobaladdr:$rvfunc, timm:$withmarker, GPR64:$Rn)>,
       Requires<[NoSLSBLRMitigation]>;
 
 def : Pat<(AArch64call_bti GPR64:$Rn),

diff --git a/llvm/lib/Target/AArch64/GISel/AArch64CallLowering.cpp b/llvm/lib/Target/AArch64/GISel/AArch64CallLowering.cpp
@@ -1364,6 +1364,13 @@ bool AArch64CallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
     Function *ARCFn = *objcarc::getAttachedARCFunction(Info.CB);
     MIB.addGlobalAddress(ARCFn);
     ++CalleeOpNo;
+
+    // We may or may not need to emit both the marker and the retain/claim call.
+    // Do what the frontend tells us: if the rvmarker module flag is present,
+    // emit the marker.  Always emit the call regardless.
+    // Tell the pseudo expansion using an additional boolean op.
+    MIB.addImm(true);
+    ++CalleeOpNo;
   } else if (Info.CFIType) {
     MIB->setCFIType(MF, Info.CFIType->getZExtValue());
   }

diff --git a/llvm/test/CodeGen/AArch64/expand-blr-rvmarker-pseudo.mir b/llvm/test/CodeGen/AArch64/expand-blr-rvmarker-pseudo.mir
@@ -29,6 +29,10 @@
     ret void
   }
 
+  define void @test_no_nop() {
+    ret void
+  }
+
   declare ptr @attachedcall()
 
   declare ptr @objc_retainAutoreleasedReturnValue()
@@ -54,7 +58,7 @@ body:             |
     bb.0:
         liveins: $lr, $x0
 
-        BLR_RVMARKER @attachedcall, $x0, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $x0
+        BLR_RVMARKER @attachedcall, 1, $x0, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $x0
         RET_ReallyLR implicit killed $w0
 ...
 
@@ -74,7 +78,7 @@ body:             |
     bb.0:
         liveins: $lr, $x0
 
-        BLR_RVMARKER @attachedcall, @foo, $x0, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $x0
+        BLR_RVMARKER @attachedcall, 1, @foo, $x0, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $x0
         RET_ReallyLR implicit killed $w0
 ...
 
@@ -94,7 +98,7 @@ body:             |
     bb.0:
         liveins: $lr, $x0, $x1, $x2
 
-        BLR_RVMARKER @attachedcall, @foo, $x0, $x1, $x2, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def $x0
+        BLR_RVMARKER @attachedcall, 1, @foo, $x0, $x1, $x2, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def $x0
         RET_ReallyLR
 ...
 
@@ -114,7 +118,7 @@ body:             |
     bb.0:
         liveins: $lr, $w0, $w1
 
-        BLR_RVMARKER @attachedcall, @foo, $w0, $w1, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $x0
+        BLR_RVMARKER @attachedcall, 1, @foo, $w0, $w1, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $x0
         RET_ReallyLR implicit killed $w0
 ...
 
@@ -135,7 +139,7 @@ body:             |
     bb.0:
         liveins: $lr, $x8, $w0, $w1
 
-        BLR_RVMARKER @attachedcall, $x8, $w0, $w1, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $x0
+        BLR_RVMARKER @attachedcall, 1, $x8, $w0, $w1, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $x0
         RET_ReallyLR implicit killed $w0
 ...
 
@@ -158,6 +162,28 @@ body:             |
   bb.0:
     liveins: $lr
 
-    BLR_RVMARKER @objc_retainAutoreleasedReturnValue, @foo, undef $x0, csr_darwin_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $x0
+    BLR_RVMARKER @objc_retainAutoreleasedReturnValue, 1, @foo, undef $x0, csr_darwin_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $x0
     RET_ReallyLR
 ...
+
+# CHECK-LABEL: : test_no_nop
+# CHECK:       bb.0:
+# CHECK-NEXT: liveins:
+# CHECK-NEXT: {{  $}}
+# CHECK-NEXT:    BUNDLE implicit-def $lr, implicit-def $w30, implicit-def $w30_hi, implicit-def $sp, implicit-def $wsp, implicit-def $wsp_hi, implicit-def dead $x0, implicit $x0, implicit $sp {
+# CHECK-NEXT:      BLR $x0, csr_aarch64_aapcs, implicit-def $lr, implicit $sp, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $x0
+# CHECK-NEXT:      BL @attachedcall, implicit-def $lr, implicit internal $sp
+# CHECK-NEXT:   }
+# CHECK-NEXT:    RET undef $lr, implicit killed $w0
+#
+name: test_no_nop
+callSites:
+  - {bb: 0, offset: 0, fwdArgRegs:
+    - { arg: 0, reg: '$x0' } }
+body:             |
+    bb.0:
+        liveins: $lr, $x0
+
+        BLR_RVMARKER @attachedcall, 0, $x0, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $x0
+        RET_ReallyLR implicit killed $w0
+...
diff --git a/llvm/test/CodeGen/AArch64/rvmarker-pseudo-expansion-and-outlining.mir b/llvm/test/CodeGen/AArch64/rvmarker-pseudo-expansion-and-outlining.mir
@@ -52,7 +52,7 @@ body:             |
   bb.0:
     liveins: $lr
 
-    BLR_RVMARKER @attachedcall, @cb1, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def $x0
+    BLR_RVMARKER @attachedcall, 1, @cb1, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def $x0
     $w12 = ORRWri $wzr, 1
     $w12 = ORRWri $wzr, 1
     $w12 = ORRWri $wzr, 1
@@ -71,7 +71,7 @@ body:             |
   bb.0:
     liveins: $lr, $x19, $x20, $lr
 
-    BLR_RVMARKER @attachedcall, @cb2, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def $x0
+    BLR_RVMARKER @attachedcall, 1, @cb2, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def $x0
     $w12 = ORRWri $wzr, 1
     $w12 = ORRWri $wzr, 1
     $w12 = ORRWri $wzr, 1