[MachineOutliner] Don't outline functions starting with PATCHABLE_FUNCTION_ENTER/FENTRL_CALL

[MaskRay]

Summary

MachineOutliner may outline a "patchable-function-entry" function whose body has
a TargetOpcode::PATCHABLE_FUNCTION_ENTER MachineInstr. This is incorrect because
the special code sequence must stay unchanged to be used at run-time.
Avoid outlining PATCHABLE_FUNCTION_ENTER. While here, avoid outlining FENTRY_CALL too
(which doesn't reproduce currently) to allow phase ordering flexibility.

Fixes #52635

[paquette]

LGTM

llvm/include/llvm/CodeGen/TargetInstrInfo.h

public:
   getOutliningType(MachineBasicBlock::iterator &MIT, unsigned Flags) const {
     llvm_unreachable(
         "Target didn't implement TargetInstrInfo::getOutliningType!");
   }
 
   /// Optional target hook that returns true if \p MBB is safe to outline from,
   /// and returns any target-specific information in \p Flags.
   virtual bool isMBBSafeToOutlineFrom(MachineBasicBlock &MBB,
-                                      unsigned &Flags) const {
-    return true;
-  }
+                                      unsigned &Flags) const;
 
   /// Insert a custom frame for outlined functions.
   virtual void buildOutlinedFrame(MachineBasicBlock &MBB, MachineFunction &MF,
                                   const outliner::OutlinedFunction &OF) const {
     llvm_unreachable(
         "Target didn't implement TargetInstrInfo::buildOutlinedFrame!");
   }
 

llvm/lib/CodeGen/TargetInstrInfo.cpp

if (ParentFn.hasFnAttribute("target-features"))
     F.addFnAttr(ParentFn.getFnAttribute("target-features"));
 
   // Set nounwind, so we don't generate eh_frame.
   if (llvm::all_of(Candidates, [](const outliner::Candidate &C) {
         return C.getMF()->getFunction().hasFnAttribute(Attribute::NoUnwind);
       }))
     F.addFnAttr(Attribute::NoUnwind);
 }
+
+bool TargetInstrInfo::isMBBSafeToOutlineFrom(MachineBasicBlock &MBB,
+                                             unsigned &Flags) const {
+  // Some instrumentations create special TargetOpcode at the start which
+  // expands to special code sequences which must be present.
+  auto First = MBB.getFirstNonDebugInstr();
+  if (First != MBB.end() &&
+      (First->getOpcode() == TargetOpcode::FENTRY_CALL ||
+       First->getOpcode() == TargetOpcode::PATCHABLE_FUNCTION_ENTER))
+    return false;
+
+  return true;
+}

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

if (MF.getTarget().getMCAsmInfo()->usesWindowsCFI())
     return false;
 
   // It's safe to outline from MF.
   return true;
 }
 
 bool AArch64InstrInfo::isMBBSafeToOutlineFrom(MachineBasicBlock &MBB,
                                               unsigned &Flags) const {
+  if (!TargetInstrInfo::isMBBSafeToOutlineFrom(MBB, Flags))
+    return false;
   // Check if LR is available through all of the MBB. If it's not, then set
   // a flag.
   assert(MBB.getParent()->getRegInfo().tracksLiveness() &&
          "Suitable Machine Function for outlining must track liveness");
   LiveRegUnits LRU(getRegisterInfo());
 
   std::for_each(MBB.rbegin(), MBB.rend(),
                 [&LRU](MachineInstr &MI) { LRU.accumulate(MI); });

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

bool RISCVInstrInfo::isFunctionSafeToOutlineFrom(
 
   // It's safe to outline from MF.
   return true;
 }
 
 bool RISCVInstrInfo::isMBBSafeToOutlineFrom(MachineBasicBlock &MBB,
                                             unsigned &Flags) const {
   // More accurate safety checking is done in getOutliningCandidateInfo.
-  return true;
+  return TargetInstrInfo::isMBBSafeToOutlineFrom(MBB, Flags);
 }
 
 // Enum values indicating how an outlined call should be constructed.
 enum MachineOutlinerConstructionID {
   MachineOutlinerDefault
 };
 
 outliner::OutlinedFunction RISCVInstrInfo::getOutliningCandidateInfo(

llvm/test/CodeGen/AArch64/machine-outliner-patchable.ll

+; RUN: llc < %s -verify-machineinstrs -enable-machine-outliner | FileCheck %s
+
+target triple = "aarch64-unknown-linux-gnu"
+
+declare void @foo(i32, i32, i32, i32) minsize
+
+;; TargetOpcode::FENTRY_CALL at the start of the function expands to a __fentry__
+;; call which must be present. Don't outline it.
+define void @fentry0(i1 %a) nounwind "fentry-call"="true" {
+; CHECK-LABEL: fentry0:
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:    # FEntry call
+; CHECK:       // %bb.1:
+; CHECK-NEXT:    bl OUTLINED_FUNCTION_1
+entry:
+  br i1 %a, label %if.then, label %if.end
+if.then:
+  call void @foo(i32 1, i32 2, i32 3, i32 4)
+  br label %if.end
+if.end:
+  call void @foo(i32 5, i32 6, i32 7, i32 8)
+  ret void
+}
+
+define void @fentry1(i1 %a) nounwind "fentry-call"="true" {
+; CHECK-LABEL: fentry1:
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:    # FEntry call
+; CHECK:       // %bb.1:
+; CHECK-NEXT:    bl OUTLINED_FUNCTION_1
+entry:
+  br i1 %a, label %if.then, label %if.end
+if.then:
+  call void @foo(i32 1, i32 2, i32 3, i32 4)
+  br label %if.end
+if.end:
+  call void @foo(i32 5, i32 6, i32 7, i32 8)
+  ret void
+}
+
+;; TargetOpcode::PATCHABLE_FUNCTION_ENTER at the start of the function expands to
+;; NOPs which must be present. Don't outline them.
+define void @patchable0(i1 %a) nounwind "patchable-function-entry"="2" {
+; CHECK-LABEL: patchable0:
+; CHECK-NEXT:  .Lfunc_begin0:
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:    nop
+; CHECK-NEXT:    nop
+; CHECK:       // %bb.1:
+; CHECK-NEXT:    bl OUTLINED_FUNCTION_1
+entry:
+  br i1 %a, label %if.then, label %if.end
+if.then:
+  call void @foo(i32 1, i32 2, i32 3, i32 4)
+  br label %if.end
+if.end:
+  call void @foo(i32 5, i32 6, i32 7, i32 8)
+  ret void
+}
+
+define void @patchable1(i1 %a) nounwind "patchable-function-entry"="2" {
+; CHECK-LABEL: patchable1:
+; CHECK-NEXT:  .Lfunc_begin1:
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:    nop
+; CHECK-NEXT:    nop
+; CHECK:       // %bb.1:
+; CHECK-NEXT:    bl OUTLINED_FUNCTION_1
+entry:
+  br i1 %a, label %if.then, label %if.end
+if.then:
+  call void @foo(i32 1, i32 2, i32 3, i32 4)
+  br label %if.end
+if.end:
+  call void @foo(i32 5, i32 6, i32 7, i32 8)
+  ret void
+}
+
+;; Similar to "patchable-function-entry".
+define void @xray0(i1 %a) nounwind "function-instrument"="xray-always" {
+; CHECK-LABEL: xray0:
+; CHECK-NEXT:  .Lfunc_begin2:
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:  .p2align 2
+; CHECK-NEXT:  .Lxray_sled_0:
+; CHECK:       // %bb.1:
+; CHECK-NEXT:    bl OUTLINED_FUNCTION_1
+entry:
+  br i1 %a, label %if.then, label %if.end
+if.then:
+  call void @foo(i32 1, i32 2, i32 3, i32 4)
+  br label %if.end
+if.end:
+  call void @foo(i32 5, i32 6, i32 7, i32 8)
+  ret void
+}
+
+define void @xray1(i1 %a) nounwind "function-instrument"="xray-always" {
+; CHECK-LABEL: xray1:
+; CHECK-NEXT:  .Lfunc_begin3:
+; CHECK-NEXT:  // %bb.0:
+; CHECK-NEXT:  .p2align 2
+; CHECK-NEXT:  .Lxray_sled_2:
+; CHECK:       // %bb.1:
+; CHECK-NEXT:    bl OUTLINED_FUNCTION_1
+entry:
+  br i1 %a, label %if.then, label %if.end
+if.then:
+  call void @foo(i32 1, i32 2, i32 3, i32 4)
+  br label %if.end
+if.end:
+  call void @foo(i32 5, i32 6, i32 7, i32 8)
+  ret void
+}

llvm/test/CodeGen/RISCV/machine-outliner-patchable.ll

+; RUN: llc < %s -verify-machineinstrs -enable-machine-outliner | FileCheck %s
+
+target triple = "riscv64-unknown-linux-gnu"
+
+declare void @foo(i32, i32, i32, i32) minsize
+
+;; TargetOpcode::FENTRY_CALL at the start of the function expands to a __fentry__
+;; call which must be present. Don't outline it.
+define void @fentry0(i1 %a) nounwind "fentry-call"="true" {
+; CHECK-LABEL: fentry0:
+; CHECK-NEXT:  # %bb.0:
+; CHECK-NEXT:    # FEntry call
+; CHECK:       # %bb.1:
+; CHECK-NEXT:    call t0, OUTLINED_FUNCTION_1
+entry:
+  br i1 %a, label %if.then, label %if.end
+if.then:
+  call void @foo(i32 1, i32 2, i32 3, i32 4)
+  br label %if.end
+if.end:
+  call void @foo(i32 5, i32 6, i32 7, i32 8)
+  ret void
+}
+
+define void @fentry1(i1 %a) nounwind "fentry-call"="true" {
+; CHECK-LABEL: fentry1:
+; CHECK-NEXT:  # %bb.0:
+; CHECK-NEXT:    # FEntry call
+; CHECK:       # %bb.1:
+; CHECK-NEXT:    call t0, OUTLINED_FUNCTION_1
+entry:
+  br i1 %a, label %if.then, label %if.end
+if.then:
+  call void @foo(i32 1, i32 2, i32 3, i32 4)
+  br label %if.end
+if.end:
+  call void @foo(i32 5, i32 6, i32 7, i32 8)
+  ret void
+}
+
+;; TargetOpcode::PATCHABLE_FUNCTION_ENTER at the start of the function expands to
+;; NOPs which must be present. Don't outline them.
+define void @patchable0(i1 %a) nounwind "patchable-function-entry"="2" {
+; CHECK-LABEL: patchable0:
+; CHECK-NEXT:  .Lfunc_begin0:
+; CHECK-NEXT:  # %bb.0:
+; CHECK-NEXT:    nop
+; CHECK-NEXT:    nop
+; CHECK:       # %bb.1:
+; CHECK-NEXT:    call t0, OUTLINED_FUNCTION_1
+entry:
+  br i1 %a, label %if.then, label %if.end
+if.then:
+  call void @foo(i32 1, i32 2, i32 3, i32 4)
+  br label %if.end
+if.end:
+  call void @foo(i32 5, i32 6, i32 7, i32 8)
+  ret void
+}
+
+define void @patchable1(i1 %a) nounwind "patchable-function-entry"="2" {
+; CHECK-LABEL: patchable1:
+; CHECK-NEXT:  .Lfunc_begin1:
+; CHECK-NEXT:  # %bb.0:
+; CHECK-NEXT:    nop
+; CHECK-NEXT:    nop
+; CHECK:       # %bb.1:
+; CHECK-NEXT:    call t0, OUTLINED_FUNCTION_1
+entry:
+  br i1 %a, label %if.then, label %if.end
+if.then:
+  call void @foo(i32 1, i32 2, i32 3, i32 4)
+  br label %if.end
+if.end:
+  call void @foo(i32 5, i32 6, i32 7, i32 8)
+  ret void
+}