[MachineOutliner] Annotation for outlined functions in AArch64

- Adding changes to support comments on outlined functions with outlining for the conditions through which it was outlined (e.g. Thunks, Tail calls)
- Adapts the emitFunctionHeader to print out a comment next to the header if the target specifies it based on information in MachineFunctionInfo
- Adds mir test for function annotiation

Differential Revision: https://reviews.llvm.org/D78062

include/llvm/CodeGen/AsmPrinter.h

@@ -681,6 +681,9 @@ private:
   /// This method emits the header for the current function.
   virtual void emitFunctionHeader();
 
+  /// This method emits a comment next to header for the current function.
+  virtual void emitFunctionHeaderComment();
+
   /// Emit a blob of inline asm to the output streamer.
   void
   emitInlineAsm(StringRef Str, const MCSubtargetInfo &STI,

lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -668,6 +668,8 @@ void AsmPrinter::emitDebugValue(const MCExpr *Value, unsigned Size) const {
   OutStreamer->emitValue(Value, Size);
 }
 
+void AsmPrinter::emitFunctionHeaderComment() {}
+
 /// EmitFunctionHeader - This method emits the header for the current
 /// function.
 void AsmPrinter::emitFunctionHeader() {
@@ -704,6 +706,7 @@ void AsmPrinter::emitFunctionHeader() {
   if (isVerbose()) {
     F.printAsOperand(OutStreamer->GetCommentOS(),
                    /*PrintType=*/false, F.getParent());
+    emitFunctionHeaderComment();
     OutStreamer->GetCommentOS() << '\n';
   }
 

lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -114,6 +114,8 @@ public:
 
   void emitInstruction(const MachineInstr *MI) override;
 
+  void emitFunctionHeaderComment() override;
+
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AsmPrinter::getAnalysisUsage(AU);
     AU.setPreservesAll();
@@ -241,6 +243,13 @@ void AArch64AsmPrinter::emitStartOfAsmFile(Module &M) {
   OutStreamer->SwitchSection(Cur);
 }
 
+void AArch64AsmPrinter::emitFunctionHeaderComment() {
+  const AArch64FunctionInfo *FI = MF->getInfo<AArch64FunctionInfo>();
+  Optional<std::string> OutlinerString = FI->getOutliningStyle();
+  if (OutlinerString != None)
+    OutStreamer->GetCommentOS() << ' ' << OutlinerString;
+}
+
 void AArch64AsmPrinter::LowerPATCHABLE_FUNCTION_ENTER(const MachineInstr &MI)
 {
   const Function &F = MF->getFunction();

lib/Target/AArch64/AArch64InstrInfo.cpp

@@ -6408,9 +6408,14 @@ static void signOutlinedFunction(MachineFunction &MF, MachineBasicBlock &MBB,
 void AArch64InstrInfo::buildOutlinedFrame(
     MachineBasicBlock &MBB, MachineFunction &MF,
     const outliner::OutlinedFunction &OF) const {
-  // For thunk outlining, rewrite the last instruction from a call to a
-  // tail-call.
-  if (OF.FrameConstructionID == MachineOutlinerThunk) {
+
+  AArch64FunctionInfo *FI = MF.getInfo<AArch64FunctionInfo>();
+
+  if (OF.FrameConstructionID == MachineOutlinerTailCall)
+    FI->setOutliningStyle("Tail Call");
+  else if (OF.FrameConstructionID == MachineOutlinerThunk) {
+    // For thunk outlining, rewrite the last instruction from a call to a
+    // tail-call.
     MachineInstr *Call = &*--MBB.instr_end();
     unsigned TailOpcode;
     if (Call->getOpcode() == AArch64::BL) {
@@ -6424,6 +6429,8 @@ void AArch64InstrInfo::buildOutlinedFrame(
                            .addImm(0);
     MBB.insert(MBB.end(), TC);
     Call->eraseFromParent();
+
+    FI->setOutliningStyle("Thunk");
   }
 
   bool IsLeafFunction = true;
@@ -6529,6 +6536,8 @@ void AArch64InstrInfo::buildOutlinedFrame(
   signOutlinedFunction(MF, MBB, ShouldSignReturnAddr,
                        ShouldSignReturnAddrWithAKey);
 
+  FI->setOutliningStyle("Function");
+
   // Did we have to modify the stack by saving the link register?
   if (OF.FrameConstructionID != MachineOutlinerDefault)
     return;

lib/Target/AArch64/AArch64MachineFunctionInfo.h

@@ -131,6 +131,10 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
   // stack slot.
   unsigned TaggedBasePointerOffset = 0;
 
+  /// OutliningStyle denotes, if a function was outined, how it was outlined,
+  /// e.g. Tail Call, Thunk, or Function if none apply.
+  Optional<std::string> OutliningStyle;
+
 public:
   AArch64FunctionInfo() = default;
 
@@ -179,6 +183,9 @@ public:
   void setLocalStackSize(uint64_t Size) { LocalStackSize = Size; }
   uint64_t getLocalStackSize() const { return LocalStackSize; }
 
+  void setOutliningStyle(std::string Style) { OutliningStyle = Style; }
+  Optional<std::string> getOutliningStyle() const { return OutliningStyle; }
+
   void setCalleeSavedStackSize(unsigned Size) {
     CalleeSavedStackSize = Size;
     HasCalleeSavedStackSize = true;

test/CodeGen/AArch64/machine-outliner-function-annotate.mir

@@ -0,0 +1,51 @@
+# RUN: llc -mtriple=aarch64--- -start-before=machine-outliner -enable-machine-outliner %s -o - | FileCheck %s
+
+# Check that a non tail called or thunk function is annotated properly with
+# only "Function"
+
+# CHECK-LABEL: OUTLINED_FUNCTION_0:
+# CHECK-SAME: // @OUTLINED_FUNCTION_0 Function
+# CHECK:      mov     w0, #1
+# CHECK-NEXT: mov     w1, #2
+# CHECK-NEXT: mov     w2, #3
+# CHECK-NEXT: mov     w3, #4
+# CHECK-NEXT: ret
+
+---
+name:            a
+alignment:       4
+tracksRegLiveness: true
+machineFunctionInfo:
+  hasRedZone:      false
+body:             |
+  bb.0:
+    frame-setup CFI_INSTRUCTION def_cfa_offset 32
+    frame-setup CFI_INSTRUCTION offset $w19, -8
+    frame-setup CFI_INSTRUCTION offset $w30, -16
+    $w0 = MOVZWi 1, 0
+    $w1 = MOVZWi 2, 0
+    $w2 = MOVZWi 3, 0
+    $w3 = MOVZWi 4, 0
+    renamable $w19 = MOVZWi 2, 0
+    RET undef $lr
+
+...
+---
+name:            b
+alignment:       4
+tracksRegLiveness: true
+machineFunctionInfo:
+  hasRedZone:      false
+body:             |
+  bb.0:
+    frame-setup CFI_INSTRUCTION def_cfa_offset 32
+    frame-setup CFI_INSTRUCTION offset $w19, -8
+    frame-setup CFI_INSTRUCTION offset $w30, -16
+    $w0 = MOVZWi 1, 0
+    $w1 = MOVZWi 2, 0
+    $w2 = MOVZWi 3, 0
+    $w3 = MOVZWi 4, 0
+    renamable $w19 = MOVZWi 1, 0
+    RET undef $lr
+
+...

test/CodeGen/AArch64/machine-outliner-tail.ll

@@ -1,6 +1,7 @@
 ; RUN: llc -verify-machineinstrs -enable-machine-outliner -mtriple=aarch64-linux-gnu < %s | FileCheck %s
 
-; CHECK: OUTLINED_FUNCTION_0:
+; CHECK-LABEL: OUTLINED_FUNCTION_0:
+; CHECK-SAME: // @OUTLINED_FUNCTION_0 Tail Call
 ; CHECK:      mov     w0, #1
 ; CHECK-NEXT: mov     w1, #2
 ; CHECK-NEXT: mov     w2, #3

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

@@ -71,6 +71,7 @@ entry:
 }
 
 ; CHECK: [[OUTLINED_INDIRECT]]:
+; CHECK-SAME: // @[[OUTLINED_INDIRECT]] Thunk
 ; CHECK:        // %bb.0:
 ; CHECK-NEXT:   mov     x8, x0
 ; CHECK-NEXT:   mov     w0, #1
@@ -80,6 +81,7 @@ entry:
 ; CHECK-NEXT:   br      x8
 
 ; CHECK: [[OUTLINED_DIRECT]]:
+; CHECK-SAME: // @[[OUTLINED_DIRECT]] Thunk
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov     w0, #1
 ; CHECK-NEXT:    mov     w1, #2