[AArch64] Lower calls with rv_marker attribute .

[fhahn]

Summary

This patch adds support for lowering function calls with the
rv_marker attribute. The goal is to expand such calls to the
following sequence of instructions:

BL @fn
mov x29, x29

This sequence of instructions triggers Objective-C runtime optimizations,
hence we want to ensure no instructions get moved in between them.
This patch achieves that by adding a new CALL_RVMARKER ISD node,
which gets turned into the BLR_RVMARKER pseudo, which eventually gets
expanded into the sequence mentioned above. The sequence is then marked
as instruction bundle, to avoid anything being moved in between.

@ahatanak is working on using this attribute in the front- & middle-end.

Together with the front- & middle-end changes, this should address
PR31925 for AArch64.

[dmgreen]

Together with the front- & middle-end changes, this should address PR31925 for AArch64.

This one: https://bugs.llvm.org/show_bug.cgi?id=31925 ? I was looking for more context but didn't find it.

[fhahn]

In D92569#2433647, @dmgreen wrote:

Together with the front- & middle-end changes, this should address PR31925 for AArch64.

This one: https://bugs.llvm.org/show_bug.cgi?id=31925 ? I was looking for more context but didn't find it.

Yes it's the right bug, but unfortunately does not have too much helpful information. Some more details about this optimization can be found here: https://opensource.apple.com/source/objc4/objc4-493.9/runtime/objc-arr.mm (search for Fast handling of returned autoreleased values).

I'll try to summarize the optimization. The goal is to reduce the overhead of objc_autoreleaseReturnValue in Callee followed by objc_retainAutoreleasedReturnValue in Caller in the snippet below. If they are back-to-back, then the ret object does not have to be added to the autorelease pool and which is much quicker.

To detect that pattern, objc_autoreleaseReturnValue looks at the instruction after ret = callee() in Caller. If it is the special marker instruction (mov x29, x29 on AArch64), then it assumes that it can do the optimization. The compiler is responsible for inserting this special marker. Currently it injects the marker as inline assembly (https://github.com/llvm/llvm-project/blob/master/llvm/lib/Transforms/ObjCARC/ObjCARCContract.cpp#L483), but optimizations can move things in between he call and the marker. This patch ensures that nothing can move instructions between the call and the marker in the backend. @ahatanak will soon post a patch for the front- and middle-end.

Caller:
  ret = callee();
  objc_retainAutoreleasedReturnValue(ret);
  // use ret here

Callee:
  // compute ret
  [ret retain];
  return objc_autoreleaseReturnValue(ret);

[aemerson]

As discussed offline, this also needs an implementation for GlobalISel. You can just add check lines to the test with -global-isel -global-isel-abort=1 for the tests.

[fhahn]

Also move callsite info from BLR_RVMARKER to the expanded call. Add a MIR test case for BLR_RVMARKER expansion.

[fhahn]

Handle expanding calls with multiple arguments.

[fhahn]

In D92569#2435493, @aemerson wrote:

As discussed offline, this also needs an implementation for GlobalISel. You can just add check lines to the test with -global-isel -global-isel-abort=1 for the tests.

Yep, but I think that can be done as a separate patch. The rv_attribute is not generated for -O0 for now, but I will add GlobalISel support as soon as SelDAG support is wrapped up. I added a -global-isel RUN line.

[fhahn]

The patch that adds the attribute is available now as well: D92808

[fhahn]

ping :)

[t.p.northover]

I think this looks reasonable now.

llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp

bool expandCMP_SWAP_128(MachineBasicBlock &MBB,
                           MachineBasicBlock::iterator MBBI,
                           MachineBasicBlock::iterator &NextMBBI);
   bool expandSetTagLoop(MachineBasicBlock &MBB,
                         MachineBasicBlock::iterator MBBI,
                         MachineBasicBlock::iterator &NextMBBI);
   bool expandSVESpillFill(MachineBasicBlock &MBB,
                           MachineBasicBlock::iterator MBBI, unsigned Opc,
                           unsigned N);
+  bool expandCALL_RVMARKER(MachineBasicBlock &MBB,
+                           MachineBasicBlock::iterator MBBI);
 };
 
 } // end anonymous namespace
 
 char AArch64ExpandPseudo::ID = 0;
 
 INITIALIZE_PASS(AArch64ExpandPseudo, "aarch64-expand-pseudo",
                 AARCH64_EXPAND_PSEUDO_NAME, false, false)
BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(Opc))
             Opc == AArch64::LDR_ZXI ? RegState::Define : 0)
         .addReg(MI.getOperand(1).getReg(), getKillRegState(Kill))
         .addImm(ImmOffset);
   }
   MI.eraseFromParent();
   return true;
 }
 
+bool AArch64ExpandPseudo::expandCALL_RVMARKER(
+    MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI) {
+  // Expand CALL_RVMARKER pseudo to a branch, followed by the special `mov x29,
+  // x29` marker. Mark the sequence as bundle, to avoid passes moving other code
+  // in between.
+  MachineInstr &MI = *MBBI;
+
+  MachineInstr *OriginalCall;
+  MachineOperand &CallTarget = MI.getOperand(0);
+  assert((CallTarget.isGlobal() || CallTarget.isReg()) &&
+         "invalid operand for regular call");
+  unsigned Opc = CallTarget.isGlobal() ? AArch64::BL : AArch64::BLR;
+  OriginalCall = BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(Opc)).getInstr();
+  OriginalCall->addOperand(CallTarget);
+
+  unsigned RegMaskStartIdx = 1;
+  // Skip register arguments. Those are added during ISel, but are not
+  // needed for the concrete branch.
+  while (!MI.getOperand(RegMaskStartIdx).isRegMask()) {
+    assert(MI.getOperand(RegMaskStartIdx).isReg() &&
+           "should only skip register operands");
+    RegMaskStartIdx++;
+  }
+  for (; RegMaskStartIdx < MI.getNumOperands(); ++RegMaskStartIdx)
+    OriginalCall->addOperand(MI.getOperand(RegMaskStartIdx));
+
+  auto *Marker = BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(AArch64::ORRXrs))
+                     .addReg(AArch64::FP)
+                     .addReg(AArch64::XZR)
+                     .addReg(AArch64::FP)
+                     .addImm(0)
+                     .getInstr();
+  if (MI.shouldUpdateCallSiteInfo())
+    MBB.getParent()->moveCallSiteInfo(&MI, Marker);
+  MI.eraseFromParent();
+  finalizeBundle(MBB, OriginalCall->getIterator(),
+                 std::next(Marker->getIterator()));
+  return true;
+}
+
 /// If MBBI references a pseudo instruction that should be expanded here,
 /// do the expansion and return true.  Otherwise return false.
 bool AArch64ExpandPseudo::expandMI(MachineBasicBlock &MBB,
                                    MachineBasicBlock::iterator MBBI,
                                    MachineBasicBlock::iterator &NextMBBI) {
   MachineInstr &MI = *MBBI;
   unsigned Opcode = MI.getOpcode();
 
bool AArch64ExpandPseudo::expandMI(MachineBasicBlock &MBB,
    case AArch64::STR_ZZXI:
      return expandSVESpillFill(MBB, MBBI, AArch64::STR_ZXI, 2);
    case AArch64::LDR_ZZZZXI:
      return expandSVESpillFill(MBB, MBBI, AArch64::LDR_ZXI, 4);
    case AArch64::LDR_ZZZXI:
      return expandSVESpillFill(MBB, MBBI, AArch64::LDR_ZXI, 3);
    case AArch64::LDR_ZZXI:
      return expandSVESpillFill(MBB, MBBI, AArch64::LDR_ZXI, 2);
+   case AArch64::BLR_RVMARKER:
+     return expandCALL_RVMARKER(MBB, MBBI);
   }
   return false;
 }
 
 /// Iterate over the instructions in basic block MBB and expand any
 /// pseudo instructions.  Return true if anything was modified.
 bool AArch64ExpandPseudo::expandMBB(MachineBasicBlock &MBB) {
   bool Modified = false;

llvm/lib/Target/AArch64/AArch64ISelLowering.h

enum NodeType : unsigned {
 
   STG,
   STZG,
   ST2G,
   STZ2G,
 
   LDP,
   STP,
-  STNP
+  STNP,
+
+  // Pseudo for a OBJC call that gets emitted together with a special `mov
+  // x29, x29` marker instruction.
+  CALL_RVMARKER
 };
 
 } // end namespace AArch64ISD
 
 namespace {
 
 // Any instruction that defines a 32-bit result zeros out the high half of the
 // register. Truncate can be lowered to EXTRACT_SUBREG. CopyFromReg may

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

case AArch64ISD::FIRST_NUMBER:
     MAKE_CASE(AArch64ISD::SSTNT1_INDEX_PRED)
     MAKE_CASE(AArch64ISD::LDP)
     MAKE_CASE(AArch64ISD::STP)
     MAKE_CASE(AArch64ISD::STNP)
     MAKE_CASE(AArch64ISD::DUP_MERGE_PASSTHRU)
     MAKE_CASE(AArch64ISD::INDEX_VECTOR)
     MAKE_CASE(AArch64ISD::UABD)
     MAKE_CASE(AArch64ISD::SABD)
+    MAKE_CASE(AArch64ISD::CALL_RVMARKER)
   }
 #undef MAKE_CASE
   return nullptr;
 }
 
 MachineBasicBlock *
 AArch64TargetLowering::EmitF128CSEL(MachineInstr &MI,
                                     MachineBasicBlock *MBB) const {
AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
   // actual call instruction.
   if (IsTailCall) {
     MF.getFrameInfo().setHasTailCall();
     SDValue Ret = DAG.getNode(AArch64ISD::TC_RETURN, DL, NodeTys, Ops);
     DAG.addCallSiteInfo(Ret.getNode(), std::move(CSInfo));
     return Ret;
   }
 
+  unsigned CallOpc = AArch64ISD::CALL;
+  // Calls marked with "rv_marker" are special. They should be expanded to the
+  // call, directly followed by a special marker sequence. Use the CALL_RVMARKER
+  // to do that.
+  if (CLI.CB && CLI.CB->hasRetAttr("rv_marker")) {
+    assert(!IsTailCall && "tail calls cannot be marked with rv_marker");
+    CallOpc = AArch64ISD::CALL_RVMARKER;
+  }
+
   // Returns a chain and a flag for retval copy to use.
-  Chain = DAG.getNode(AArch64ISD::CALL, DL, NodeTys, Ops);
+  Chain = DAG.getNode(CallOpc, DL, NodeTys, Ops);
   DAG.addNoMergeSiteInfo(Chain.getNode(), CLI.NoMerge);
   InFlag = Chain.getValue(1);
   DAG.addCallSiteInfo(Chain.getNode(), std::move(CSInfo));
 
   uint64_t CalleePopBytes =
       DoesCalleeRestoreStack(CallConv, TailCallOpt) ? alignTo(NumBytes, 16) : 0;
 
   Chain = DAG.getCALLSEQ_END(Chain, DAG.getIntPtrConstant(NumBytes, DL, true),

llvm/lib/Target/AArch64/AArch64InstrInfo.td

 def AArch64callseq_end   : SDNode<"ISD::CALLSEQ_END",
                                 SDCallSeqEnd<[ SDTCisVT<0, i32>,
                                                SDTCisVT<1, i32> ]>,
                                 [SDNPHasChain, SDNPOptInGlue, SDNPOutGlue]>;
 def AArch64call          : SDNode<"AArch64ISD::CALL",
                                 SDTypeProfile<0, -1, [SDTCisPtrTy<0>]>,
                                 [SDNPHasChain, SDNPOptInGlue, SDNPOutGlue,
                                  SDNPVariadic]>;
+
+def AArch64call_rvmarker: SDNode<"AArch64ISD::CALL_RVMARKER",
+                             SDTypeProfile<0, -1, [SDTCisPtrTy<0>]>,
+                             [SDNPHasChain, SDNPOptInGlue, SDNPOutGlue,
+                              SDNPVariadic]>;
+
 def AArch64brcond        : SDNode<"AArch64ISD::BRCOND", SDT_AArch64Brcond,
                                 [SDNPHasChain]>;
 def AArch64cbz           : SDNode<"AArch64ISD::CBZ", SDT_AArch64cbz,
                                 [SDNPHasChain]>;
 def AArch64cbnz           : SDNode<"AArch64ISD::CBNZ", SDT_AArch64cbz,
                                 [SDNPHasChain]>;
 def AArch64tbz           : SDNode<"AArch64ISD::TBZ", SDT_AArch64tbz,
                                 [SDNPHasChain]>;

 // Default to the LR register.
 def : InstAlias<"ret", (RET LR)>;
 
 let isCall = 1, Defs = [LR], Uses = [SP] in {
   def BLR : BranchReg<0b0001, "blr", []>;
   def BLRNoIP : Pseudo<(outs), (ins GPR64noip:$Rn), []>,
                 Sched<[WriteBrReg]>,
                 PseudoInstExpansion<(BLR GPR64:$Rn)>;
+  def BLR_RVMARKER : Pseudo<(outs), (ins variable_ops), []>,
+                     Sched<[WriteBrReg]>;
 } // isCall
 
 def : Pat<(AArch64call GPR64:$Rn),
           (BLR GPR64:$Rn)>,
       Requires<[NoSLSBLRMitigation]>;
 def : Pat<(AArch64call GPR64noip:$Rn),
           (BLRNoIP GPR64noip:$Rn)>,
       Requires<[SLSBLRMitigation]>;
 
+def : Pat<(AArch64call_rvmarker GPR64:$Rn),
+          (BLR_RVMARKER GPR64:$Rn)>,
+      Requires<[NoSLSBLRMitigation]>;
+
 let isBranch = 1, isTerminator = 1, isBarrier = 1, isIndirectBranch = 1 in {
 def BR  : BranchReg<0b0000, "br", [(brind GPR64:$Rn)]>;
 } // isBranch, isTerminator, isBarrier, isIndirectBranch
 
 // Create a separate pseudo-instruction for codegen to use so that we don't
 // flag lr as used in every function. It'll be restored before the RET by the
 // epilogue if it's legitimately used.
 def RET_ReallyLR : Pseudo<(outs), (ins), [(AArch64retflag)]>,

llvm/test/CodeGen/AArch64/call-rv-marker.ll

+; RUN: llc -o - %s | FileCheck --check-prefix=SELDAG --check-prefix=CHECK %s
+; RUN: llc -global-isel -o - %s | FileCheck --check-prefix=GISEL --check-prefix=CHECK %s
+
+; TODO: support marker generation with GlobalISel
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "arm64-apple-iphoneos"
+
+declare i8* @foo0(i32)
+declare i8* @foo1()
+
+declare void @llvm.objc.release(i8*)
+declare void @objc_object(i8*)
+
+declare void @foo2(i8*)
+
+declare void @llvm.lifetime.start.p0i8(i64 immarg, i8* nocapture)
+
+declare %struct.S* @_ZN1SD1Ev(%struct.S* nonnull dereferenceable(1))
+
+declare void @llvm.lifetime.end.p0i8(i64 immarg, i8* nocapture)
+
+
+%struct.S = type { i8 }
+
+@g = global i8* null, align 8
+@fptr = global i8* ()* null, align 8
+
+define i8* @rv_marker_1() {
+; CHECK-LABEL:    rv_marker_1:
+; CHECK:           .cfi_offset w30, -16
+; CHECK-NEXT:      bl foo1
+; SELDAG-NEXT:     mov x29, x29
+; GISEL-NOT:       mov x29, x29
+;
+entry:
+  %call = call "rv_marker" i8* @foo1()
+  ret i8* %call
+}
+
+define void @rv_marker_2_select(i32 %c) {
+; CHECK-LABEL: rv_marker_2_select:
+; SELDAG:        cinc  w0, w8, eq
+; GISEL:         csinc w0, w8, wzr, eq
+; CHECK-NEXT:    bl  foo0
+; SELDAG-NEXT:   mov x29, x29
+; CHECK-NEXT:    ldr x30, [sp], #16
+; CHECK-NEXT:    b  foo2
+;
+entry:
+  %tobool.not = icmp eq i32 %c, 0
+  %.sink = select i1 %tobool.not, i32 2, i32 1
+  %call1 = call "rv_marker" i8* @foo0(i32 %.sink)
+  tail call void @foo2(i8* %call1)
+  ret void
+}
+
+define void @rv_marker_3() personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
+; CHECK-LABEL: rv_marker_3
+; CHECK:         .cfi_offset w30, -32
+; CHECK-NEXT:    bl  foo1
+; SELDAG-NEXT:   mov x29, x29
+;
+entry:
+  %call = call "rv_marker" i8* @foo1()
+  invoke void @objc_object(i8* %call) #5
+          to label %invoke.cont unwind label %lpad
+
+invoke.cont:                                      ; preds = %entry
+  tail call void @llvm.objc.release(i8* %call)
+  ret void
+
+lpad:                                             ; preds = %entry
+  %0 = landingpad { i8*, i32 }
+          cleanup
+  tail call void @llvm.objc.release(i8* %call)
+  resume { i8*, i32 } %0
+}
+
+define void @rv_marker_4() personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
+; CHECK-LABEL: rv_marker_4
+; CHECK:       .Ltmp3:
+; CHECK-NEXT:     bl  foo1
+; SELDAG-NEXT:    mov x29, x29
+; CHECK-NEXT: .Ltmp4:
+;
+entry:
+  %s = alloca %struct.S, align 1
+  %0 = getelementptr inbounds %struct.S, %struct.S* %s, i64 0, i32 0
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* nonnull %0) #2
+  %call = invoke "rv_marker" i8* @foo1()
+          to label %invoke.cont unwind label %lpad
+
+invoke.cont:                                      ; preds = %entry
+  invoke void @objc_object(i8* %call) #5
+          to label %invoke.cont2 unwind label %lpad1
+
+invoke.cont2:                                     ; preds = %invoke.cont
+  tail call void @llvm.objc.release(i8* %call)
+  %call3 = call %struct.S* @_ZN1SD1Ev(%struct.S* nonnull dereferenceable(1) %s)
+  call void @llvm.lifetime.end.p0i8(i64 1, i8* nonnull %0)
+  ret void
+
+lpad:                                             ; preds = %entry
+  %1 = landingpad { i8*, i32 }
+          cleanup
+  br label %ehcleanup
+
+lpad1:                                            ; preds = %invoke.cont
+  %2 = landingpad { i8*, i32 }
+          cleanup
+  tail call void @llvm.objc.release(i8* %call)
+  br label %ehcleanup
+
+ehcleanup:                                        ; preds = %lpad1, %lpad
+  %.pn = phi { i8*, i32 } [ %2, %lpad1 ], [ %1, %lpad ]
+  %call4 = call %struct.S* @_ZN1SD1Ev(%struct.S* nonnull dereferenceable(1) %s)
+  call void @llvm.lifetime.end.p0i8(i64 1, i8* nonnull %0)
+  resume { i8*, i32 } %.pn
+}
+
+define i8* @rv_marker_5_indirect_call() {
+; CHECK-LABEL: rv_marker_5_indirect_call
+; CHECK:         ldr [[ADDR:x[0-9]+]], [
+; CHECK-NEXT:    blr [[ADDR]]
+; SLEDAG-NEXT:   mov x29, x29
+; GISEL-NOT:     mov x29, x29
+;
+entry:
+  %0 = load i8* ()*, i8* ()** @fptr, align 8
+  %call = call "rv_marker" i8* %0()
+  tail call void @foo2(i8* %call)
+  ret i8* %call
+}
+
+declare void @foo(i64, i64, i64)
+
+define void @rv_marker_multiarg(i64 %a, i64 %b, i64 %c) {
+; CHECK-LABEL: rv_marker_multiarg
+; CHECK:        mov [[TMP:x[0-9]+]], x0
+; CHECK-NEXT:   mov x0, x2
+; CHECK-NEXT:   mov x2, [[TMP]]
+; CHECK-NEXT:   bl  foo
+; SELDAG-NEXT:  mov x29, x29
+; GISEL-NOT:    mov x29, x29
+  call "rv_marker" void @foo(i64 %c, i64 %b, i64 %a)
+  ret void
+}
+
+declare i32 @__gxx_personality_v0(...)

llvm/test/CodeGen/AArch64/expand-blr-rvmarker-pseudo.mir

+# RUN: llc -run-pass=aarch64-expand-pseudo -mtriple=arm64-apple-ios -o - -emit-call-site-info %s | FileCheck %s
+
+# CHECK-LABEL: test_1_callsite_info
+# CHECK:       bb.0.entry:
+# CHECK-NEXT:    BUNDLE implicit-def $lr, implicit-def $w30, implicit-def $sp, implicit-def $wsp, implicit-def dead $x0, implicit $x0, implicit $sp, implicit $fp, implicit $xzr {
+# 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:      ORRXrs $fp, $xzr, $fp, 0
+# CHECK-NEXT:   }
+# CHECK-NEXT:    RET undef $lr, implicit killed $w0
+---
+name: test_1_callsite_info
+callSites:
+  - {bb: 0, offset: 0, fwdArgRegs:
+    - { arg: 0, reg: '$x0' } }
+body:             |
+    bb.0.entry:
+        liveins: $lr, $x0
+
+        BLR_RVMARKER $x0, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $x0
+        RET_ReallyLR implicit killed $w0
+...