[AArch64] Fix verifier error when outlining indirect calls

[olista01]

Summary

The MachineOutliner for AArch64 transforms indirect calls into indirect
tail calls, replacing the call with the TCRETURNri pseudo-instruction.
This pseudo lowers to a BR, but has the isCall and isReturn flags set.

The problem is that TCRETURNri takes a tcGPR64 as the register argument,
to prevent indiret tail-calls from using caller-saved registers. The
indirect calls transformed by the outliner could use caller-saved
registers. This is fine, because the outliner ensures that the register
is available at all call sites. However, this causes a verifier failure
when the register is not in tcGPR64. The fix is to add a new
pseudo-instruction like TCRETURNri, but which accepts any GPR.

[rengolin]

Are there no other ways to identify an outliner function, say a function flag or a needed pattern, so we can change the verifier instead, and not fail when called-saved regs are used?

[olista01]

I don't think adding a special case to the verifier makes sense, when it's easy enough to generate code which is valid by the normal rules.

I did consider the alternative of just generating a BR instruction directly, since the outliner runs very late so the isCall/isReturn flags are unlikely to be used in a later pass, but I'm not sure if the verifier would like that either.

[paquette]

In D52829#1254842, @olista01 wrote:

I don't think adding a special case to the verifier makes sense, when it's easy enough to generate code which is valid by the normal rules.

I think this is probably really the only thing the outliner does that upsets the verifier. So, from that perspective, I think that this is fine.

OTOH, having a function flag or something like @rengolin suggested is probably a bit cleaner in the long term. I slightly prefer this position since it seems like it should be the verifier's responsibility to understand what to do with an outlined function. Also, if it does turn out that there are other weird cases, it'd be kind of overkill to add special outlining instructions everywhere just to appease the verifier.

So, if it's not too much overhead, I think that adding an outlined function flag or something would probably be better, but for the short-term I'm fine with this.

In D52829#1254842, @olista01 wrote:

I did consider the alternative of just generating a BR instruction directly, since the outliner runs very late so the isCall/isReturn flags are unlikely to be used in a later pass, but I'm not sure if the verifier would like that either.

That's how it used to be, and the verifier didn't like that because the BR doesn't count as a terminator. :(

[rengolin]

In D52829#1255435, @paquette wrote:

OTOH, having a function flag or something like @rengolin suggested is probably a bit cleaner in the long term. I slightly prefer this position since it seems like it should be the verifier's responsibility to understand what to do with an outlined function. Also, if it does turn out that there are other weird cases, it'd be kind of overkill to add special outlining instructions everywhere just to appease the verifier.

Yes, that's what I was worried about.

So, if it's not too much overhead, I think that adding an outlined function flag or something would probably be better, but for the short-term I'm fine with this.

SGTM, thanks!

[paquette]

As long as you add a FIXME saying that we should move this responsibility over to the verifier, this LGTM.

llvm/trunk/lib/Target/AArch64/AArch64AsmPrinter.cpp

if (isVerbose() && OutStreamer->hasRawTextSupport()) {
       OutStreamer->EmitRawText(StringRef(OS.str()));
     }
     return;
   }
 
   // Tail calls use pseudo instructions so they have the proper code-gen
   // attributes (isCall, isReturn, etc.). We lower them to the real
   // instruction here.
-  case AArch64::TCRETURNri: {
+  case AArch64::TCRETURNri:
+  case AArch64::TCRETURNriALL: {
     MCInst TmpInst;
     TmpInst.setOpcode(AArch64::BR);
     TmpInst.addOperand(MCOperand::createReg(MI->getOperand(0).getReg()));
     EmitToStreamer(*OutStreamer, TmpInst);
     return;
   }
   case AArch64::TCRETURNdi: {
     MCOperand Dest;

llvm/trunk/lib/Target/AArch64/AArch64InstrInfo.cpp

void AArch64InstrInfo::buildOutlinedFrame(
   // tail-call.
   if (OF.FrameConstructionID == MachineOutlinerThunk) {
     MachineInstr *Call = &*--MBB.instr_end();
     unsigned TailOpcode;
     if (Call->getOpcode() == AArch64::BL) {
       TailOpcode = AArch64::TCRETURNdi;
     } else {
       assert(Call->getOpcode() == AArch64::BLR);
-      TailOpcode = AArch64::TCRETURNri;
+      TailOpcode = AArch64::TCRETURNriALL;
     }
     MachineInstr *TC = BuildMI(MF, DebugLoc(), get(TailOpcode))
                             .add(Call->getOperand(0))
                             .addImm(0);
     MBB.insert(MBB.end(), TC);
     Call->eraseFromParent();
   }
 

llvm/trunk/lib/Target/AArch64/AArch64InstrInfo.td

 
 // Tail call return handling. These are all compiler pseudo-instructions,
 // so no encoding information or anything like that.
 let isCall = 1, isTerminator = 1, isReturn = 1, isBarrier = 1, Uses = [SP] in {
   def TCRETURNdi : Pseudo<(outs), (ins i64imm:$dst, i32imm:$FPDiff), []>,
                    Sched<[WriteBrReg]>;
   def TCRETURNri : Pseudo<(outs), (ins tcGPR64:$dst, i32imm:$FPDiff), []>,
                    Sched<[WriteBrReg]>;
+  // Indirect tail-call with any register allowed, used by MachineOutliner when
+  // this is proven safe.
+  // FIXME: If we have to add any more hacks like this, we should instead relax
+  // some verifier checks for outlined functions.
+  def TCRETURNriALL : Pseudo<(outs), (ins GPR64:$dst, i32imm:$FPDiff), []>,
+                      Sched<[WriteBrReg]>;
 }
 
 def : Pat<(AArch64tcret tcGPR64:$dst, (i32 timm:$FPDiff)),
           (TCRETURNri tcGPR64:$dst, imm:$FPDiff)>;
 def : Pat<(AArch64tcret tglobaladdr:$dst, (i32 timm:$FPDiff)),
           (TCRETURNdi texternalsym:$dst, imm:$FPDiff)>;
 def : Pat<(AArch64tcret texternalsym:$dst, (i32 timm:$FPDiff)),
           (TCRETURNdi texternalsym:$dst, imm:$FPDiff)>;
 
 include "AArch64InstrAtomics.td"
 include "AArch64SVEInstrInfo.td"

llvm/trunk/test/CodeGen/AArch64/machine-outliner-thunk.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -enable-machine-outliner -verify-machineinstrs | FileCheck %s
 
 target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64-pc-linux-gnu"
 
 declare i32 @thunk_called_fn(i32, i32, i32, i32)
 
 define i32 @a() {
 ; CHECK-LABEL: a:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
 ; CHECK-NEXT:    .cfi_offset w30, -16
-; CHECK-NEXT:    bl OUTLINED_FUNCTION_0
+; CHECK-NEXT:    bl [[OUTLINED_DIRECT:OUTLINED_FUNCTION_[0-9]+]]
 ; CHECK-NEXT:    add w0, w0, #8 // =8
 ; CHECK-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
 ; CHECK-NEXT:    ret
 entry:
   %call = tail call i32 @thunk_called_fn(i32 1, i32 2, i32 3, i32 4)
   %cx = add i32 %call, 8
   ret i32 %cx
 }
 
 define i32 @b() {
 ; CHECK-LABEL: b:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
 ; CHECK-NEXT:    .cfi_offset w30, -16
-; CHECK-NEXT:    bl OUTLINED_FUNCTION_0
+; CHECK-NEXT:    bl [[OUTLINED_DIRECT]]
 ; CHECK-NEXT:    add w0, w0, #88 // =88
 ; CHECK-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
 ; CHECK-NEXT:    ret
 entry:
   %call = tail call i32 @thunk_called_fn(i32 1, i32 2, i32 3, i32 4)
   %cx = add i32 %call, 88
   ret i32 %cx
 }
 
-; CHECK-LABEL: OUTLINED_FUNCTION_0:
+define hidden i32 @c(i32 (i32, i32, i32, i32)* %fptr) {
+; CHECK-LABEL: c:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    .cfi_offset w30, -16
+; CHECK-NEXT:    bl [[OUTLINED_INDIRECT:OUTLINED_FUNCTION_[0-9]+]]
+; CHECK-NEXT:    add w0, w0, #8 // =8
+; CHECK-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; CHECK-NEXT:    ret
+entry:
+  %call = tail call i32 %fptr(i32 1, i32 2, i32 3, i32 4)
+  %add = add nsw i32 %call, 8
+  ret i32 %add
+}
+
+define hidden i32 @d(i32 (i32, i32, i32, i32)* %fptr) {
+; CHECK-LABEL: d:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    .cfi_offset w30, -16
+; CHECK-NEXT:    bl [[OUTLINED_INDIRECT]]
+; CHECK-NEXT:    add w0, w0, #88 // =88
+; CHECK-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; CHECK-NEXT:    ret
+entry:
+  %call = tail call i32 %fptr(i32 1, i32 2, i32 3, i32 4)
+  %add = add nsw i32 %call, 88
+  ret i32 %add
+}
+
+; CHECK: [[OUTLINED_INDIRECT]]:
+; CHECK:        // %bb.0:
+; CHECK-NEXT:   mov     x8, x0
+; CHECK-NEXT:   orr     w0, wzr, #0x1
+; CHECK-NEXT:   orr     w1, wzr, #0x2
+; CHECK-NEXT:   orr     w2, wzr, #0x3
+; CHECK-NEXT:   orr     w3, wzr, #0x4
+; CHECK-NEXT:   br      x8
+
+; CHECK: [[OUTLINED_DIRECT]]:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    orr     w0, wzr, #0x1
 ; CHECK-NEXT:    orr     w1, wzr, #0x2
 ; CHECK-NEXT:    orr     w2, wzr, #0x3
 ; CHECK-NEXT:    orr     w3, wzr, #0x4
 ; CHECK-NEXT:    b       thunk_called_fn