[CodeGen] Async unwind - add a pass to fix CFI information

This pass inserts the necessary CFI instructions to compensate for the
inconsistency of the call-frame information caused by linear (non-CFG
aware) nature of the unwind tables.

Unlike the `CFIInstrInserer` pass, this one almost always emits only
`.cfi_remember_state`/`.cfi_restore_state`, which results in smaller
unwind tables and also transparently handles custom unwind info
extensions like CFA offset adjustement and save locations of SVE
registers.

This pass takes advantage of the constraints that LLVM imposes on the
placement of save/restore points (cf. `ShrinkWrap.cpp`):

  * there is a single basic block, containing the function prologue

  * possibly multiple epilogue blocks, where each epilogue block is
    complete and self-contained, i.e. CSR restore instructions (and the
    corresponding CFI instructions are not split across two or more
    blocks.

  * prologue and epilogue blocks are outside of any loops

Thus, during execution, at the beginning and at the end of each basic
block the function can be in one of two states:

  - "has a call frame", if the function has executed the prologue, or
     has not executed any epilogue

  - "does not have a call frame", if the function has not executed the
    prologue, or has executed an epilogue

These properties can be computed for each basic block by a single RPO
traversal.

In order to accommodate backends which do not generate unwind info in
epilogues we compute an additional property "strong no call frame on
entry" which is set for the entry point of the function and for every
block reachable from the entry along a path that does not execute the
prologue. If this property holds, it takes precedence over the "has a
call frame" property.

From the point of view of the unwind tables, the "has/does not have
call frame" state at beginning of each block is determined by the
state at the end of the previous block, in layout order.

Where these states differ, we insert compensating CFI instructions,
which come in two flavours:

- CFI instructions, which reset the unwind table state to the
    initial one.  This is done by a target specific hook and is
    expected to be trivial to implement, for example it could be:
```
     .cfi_def_cfa <sp>, 0
     .cfi_same_value <rN>
     .cfi_same_value <rN-1>
     ...
```
where `<rN>` are the callee-saved registers.

- CFI instructions, which reset the unwind table state to the one
    created by the function prologue. These are the sequence:
```
       .cfi_restore_state
       .cfi_remember_state
```
In this case we also insert a `.cfi_remember_state` after the
last CFI instruction in the function prologue.

Reviewed By: MaskRay, danielkiss, chill

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

llvm/include/llvm/CodeGen/CFIFixup.h

@@ -0,0 +1,38 @@
+//===-- CFIFixup.h - Insert CFI remember/restore instructions ---*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// Contains definition of the base CFIFixup pass.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGEN_CFIFIXUP_H
+#define LLVM_CODEGEN_CFIFIXUP_H
+
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/InitializePasses.h"
+
+namespace llvm {
+class CFIFixup : public MachineFunctionPass {
+public:
+  static char ID;
+
+  CFIFixup() : MachineFunctionPass(ID) {
+    initializeCFIFixupPass(*PassRegistry::getPassRegistry());
+  }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesAll();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+};
+} // namespace llvm
+
+#endif // LLVM_CODEGEN_CFIFIXUP_H

llvm/include/llvm/CodeGen/Passes.h

@@ -494,6 +494,9 @@ namespace llvm {
   // This pass expands indirectbr instructions.
   FunctionPass *createIndirectBrExpandPass();
 
+  /// Creates CFI Fixup pass. \see CFIFixup.cpp
+  FunctionPass *createCFIFixup();
+
   /// Creates CFI Instruction Inserter pass. \see CFIInstrInserter.cpp
   FunctionPass *createCFIInstrInserter();
 

llvm/include/llvm/CodeGen/TargetFrameLowering.h

@@ -223,6 +223,14 @@ class TargetFrameLowering {
   emitCalleeSavedFrameMovesFullCFA(MachineBasicBlock &MBB,
                                    MachineBasicBlock::iterator MBBI) const {}
 
+  /// Returns true if we may need to fix the unwind infportmation for the
+  /// function.
+  virtual bool enableCFIFixup(MachineFunction &MF) const;
+
+  /// Emit CFI instructions that recreate the state of the unwind information
+  /// upon fucntion entry.
+  virtual void resetCFIToInitialState(MachineBasicBlock &MBB) const {}
+
   /// Replace a StackProbe stub (if any) with the actual probe code inline
   virtual void inlineStackProbe(MachineFunction &MF,
                                 MachineBasicBlock &PrologueMBB) const {}

llvm/include/llvm/InitializePasses.h

@@ -103,6 +103,7 @@ void initializeCFGSimplifyPassPass(PassRegistry&);
 void initializeCFGuardPass(PassRegistry&);
 void initializeCFGuardLongjmpPass(PassRegistry&);
 void initializeCFGViewerLegacyPassPass(PassRegistry&);
+void initializeCFIFixupPass(PassRegistry&);
 void initializeCFIInstrInserterPass(PassRegistry&);
 void initializeCFLAndersAAWrapperPassPass(PassRegistry&);
 void initializeCFLSteensAAWrapperPassPass(PassRegistry&);

llvm/include/llvm/Target/TargetMachine.h

@@ -257,6 +257,8 @@ class TargetMachine {
     Options.SupportsDebugEntryValues = Enable;
   }
 
+  void setCFIFixup(bool Enable) { Options.EnableCFIFixup = Enable; }
+
   bool getAIXExtendedAltivecABI() const {
     return Options.EnableAIXExtendedAltivecABI;
   }

llvm/include/llvm/Target/TargetOptions.h

@@ -144,6 +144,7 @@ namespace llvm {
           ValueTrackingVariableLocations(false), ForceDwarfFrameSection(false),
           XRayOmitFunctionIndex(false), DebugStrictDwarf(false),
           Hotpatch(false), PPCGenScalarMASSEntries(false), JMCInstrument(false),
+          EnableCFIFixup(false),
           FPDenormalMode(DenormalMode::IEEE, DenormalMode::IEEE) {}
 
     /// DisableFramePointerElim - This returns true if frame pointer elimination
@@ -356,6 +357,9 @@ namespace llvm {
     /// Enable JustMyCode instrumentation.
     unsigned JMCInstrument : 1;
 
+    /// Enable the CFIFixup pass.
+    unsigned EnableCFIFixup : 1;
+
     /// Name of the stack usage file (i.e., .su file) if user passes
     /// -fstack-usage. If empty, it can be implied that -fstack-usage is not
     /// passed on the command line.

llvm/lib/CodeGen/AsmPrinter/AsmPrinterDwarf.cpp

@@ -273,6 +273,12 @@ void AsmPrinter::emitCFIInstruction(const MCCFIInstruction &Inst) const {
   case MCCFIInstruction::OpUndefined:
     OutStreamer->emitCFIUndefined(Inst.getRegister());
     break;
+  case MCCFIInstruction::OpRememberState:
+    OutStreamer->emitCFIRememberState();
+    break;
+  case MCCFIInstruction::OpRestoreState:
+    OutStreamer->emitCFIRestoreState();
+    break;
   }
 }
 

llvm/lib/CodeGen/CFIFixup.cpp

@@ -0,0 +1,215 @@
+//===------ CFIFixup.cpp - Insert CFI remember/restore instructions -------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+
+// This pass inserts the necessary  instructions to adjust for the inconsistency
+// of the call-frame information caused by final machine basic block layout.
+// The pass relies in constraints LLVM imposes on the placement of
+// save/restore points (cf. ShrinkWrap):
+// * there is a single basic block, containing the function prologue
+// * possibly multiple epilogue blocks, where each epilogue block is
+//   complete and self-contained, i.e. CSR restore instructions (and the
+//   corresponding CFI instructions are not split across two or more blocks.
+// * prologue and epilogue blocks are outside of any loops
+// Thus, during execution, at the beginning and at the end of each basic block
+// the function can be in one of two states:
+//  - "has a call frame", if the function has executed the prologue, and
+//    has not executed any epilogue
+//  - "does not have a call frame", if the function has not executed the
+//    prologue, or has executed an epilogue
+// which can be computed by a single RPO traversal.
+
+// In order to accommodate backends which do not generate unwind info in
+// epilogues we compute an additional property "strong no call frame on entry",
+// which is set for the entry point of the function and for every block
+// reachable from the entry along a path that does not execute the prologue. If
+// this property holds, it takes precedence over the "has a call frame"
+// property.
+
+// From the point of view of the unwind tables, the "has/does not have call
+// frame" state at beginning of each block is determined by the state at the end
+// of the previous block, in layout order. Where these states differ, we insert
+// compensating CFI instructions, which come in two flavours:
+
+//   - CFI instructions, which reset the unwind table state to the initial one.
+//     This is done by a target specific hook and is expected to be trivial
+//     to implement, for example it could be:
+//       .cfi_def_cfa <sp>, 0
+//       .cfi_same_value <rN>
+//       .cfi_same_value <rN-1>
+//       ...
+//     where <rN> are the callee-saved registers.
+//   - CFI instructions, which reset the unwind table state to the one
+//     created by the function prologue. These are
+//       .cfi_restore_state
+//       .cfi_remember_state
+//     In this case we also insert a `.cfi_remember_state` after the last CFI
+//     instruction in the function prologue.
+//
+// Known limitations:
+//  * the pass cannot handle an epilogue preceding the prologue in the basic
+//    block layout
+//  * the pass does not handle functions where SP is used as a frame pointer and
+//    SP adjustments up and down are done in different basic blocks (TODO)
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/CFIFixup.h"
+
+#include "llvm/ADT/PostOrderIterator.h"
+#include "llvm/ADT/SmallBitVector.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/CodeGen/TargetFrameLowering.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/CodeGen/TargetSubtargetInfo.h"
+#include "llvm/MC/MCAsmInfo.h"
+#include "llvm/MC/MCDwarf.h"
+#include "llvm/Target/TargetMachine.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "cfi-fixup"
+
+char CFIFixup::ID = 0;
+
+INITIALIZE_PASS(CFIFixup, "cfi-fixup",
+                "Insert CFI remember/restore state instructions", false, false)
+FunctionPass *llvm::createCFIFixup() { return new CFIFixup(); }
+
+static bool isPrologueCFIInstruction(const MachineInstr &MI) {
+  return MI.getOpcode() == TargetOpcode::CFI_INSTRUCTION &&
+         MI.getFlag(MachineInstr::FrameSetup);
+}
+
+static bool containsPrologue(const MachineBasicBlock &MBB) {
+  return llvm::any_of(MBB.instrs(), isPrologueCFIInstruction);
+}
+
+static bool containsEpilogue(const MachineBasicBlock &MBB) {
+  return llvm::any_of(llvm::reverse(MBB), [](const auto &MI) {
+    return MI.getOpcode() == TargetOpcode::CFI_INSTRUCTION &&
+           MI.getFlag(MachineInstr::FrameDestroy);
+  });
+}
+
+bool CFIFixup::runOnMachineFunction(MachineFunction &MF) {
+  const TargetFrameLowering &TFL = *MF.getSubtarget().getFrameLowering();
+  if (!TFL.enableCFIFixup(MF))
+    return false;
+
+  const unsigned NumBlocks = MF.getNumBlockIDs();
+  if (NumBlocks < 2)
+    return false;
+
+  struct BlockFlags {
+    bool StrongNoFrameOnEntry : 1;
+    bool HasFrameOnEntry : 1;
+    bool HasFrameOnExit : 1;
+  };
+  SmallVector<BlockFlags, 32> BlockInfo(NumBlocks, {false, false, false});
+  BlockInfo[0].StrongNoFrameOnEntry = true;
+
+  // Compute the presence/absence of frame at each basic block.
+  MachineBasicBlock *PrologueBlock = nullptr;
+  ReversePostOrderTraversal<MachineBasicBlock *> RPOT(&*MF.begin());
+  for (MachineBasicBlock *MBB : RPOT) {
+    BlockFlags &Info = BlockInfo[MBB->getNumber()];
+
+    // Set to true if the current block contains the prologue or the epilogue,
+    // respectively.
+    bool HasPrologue = false;
+    bool HasEpilogue = false;
+
+    if (!PrologueBlock && !Info.HasFrameOnEntry && containsPrologue(*MBB)) {
+      PrologueBlock = MBB;
+      HasPrologue = true;
+    }
+
+    if (Info.HasFrameOnEntry || HasPrologue)
+      HasEpilogue = containsEpilogue(*MBB);
+
+    // If the function has a call frame at the entry of the current block or the
+    // current block contains the prologue, then the function has a call frame
+    // at the exit of the block, unless the block contains the epilogue.
+    Info.HasFrameOnExit = (Info.HasFrameOnEntry || HasPrologue) && !HasEpilogue;
+
+    // Set the successors' state on entry.
+    for (MachineBasicBlock *Succ : MBB->successors()) {
+      BlockFlags &SuccInfo = BlockInfo[Succ->getNumber()];
+      SuccInfo.StrongNoFrameOnEntry |=
+          Info.StrongNoFrameOnEntry && !HasPrologue;
+      SuccInfo.HasFrameOnEntry = Info.HasFrameOnExit;
+    }
+  }
+
+  if (!PrologueBlock)
+    return false;
+
+  // Walk the blocks of the function in "physical" order.
+  // Every block inherits the frame state (as recorded in the unwind tables)
+  // of the previous block. If the intended frame state is different, insert
+  // compensating CFI instructions.
+  const TargetInstrInfo &TII = *MF.getSubtarget().getInstrInfo();
+  bool Change = false;
+  // `InsertPt` always points to the point in a preceding block where we have to
+  // insert a `.cfi_remember_state`, in the case that the current block needs a
+  // `.cfi_restore_state`.
+  MachineBasicBlock *InsertMBB = PrologueBlock;
+  MachineBasicBlock::iterator InsertPt = PrologueBlock->begin();
+  for (MachineInstr &MI : *PrologueBlock)
+    if (isPrologueCFIInstruction(MI))
+      InsertPt = std::next(MI.getIterator());
+
+  assert(InsertPt != PrologueBlock->begin() &&
+         "Inconsistent notion of \"prologue block\"");
+
+  // No point starting before the prologue block.
+  // TODO: the unwind tables will still be incorrect if an epilogue physically
+  // preceeds the prologue.
+  MachineFunction::iterator CurrBB = std::next(PrologueBlock->getIterator());
+  bool HasFrame = BlockInfo[PrologueBlock->getNumber()].HasFrameOnExit;
+  while (CurrBB != MF.end()) {
+    auto &Info = BlockInfo[CurrBB->getNumber()];
+#ifndef NDEBUG
+    if (!Info.StrongNoFrameOnEntry) {
+      for (auto *Pred : CurrBB->predecessors())
+        assert(Info.HasFrameOnEntry ==
+                   BlockInfo[Pred->getNumber()].HasFrameOnExit &&
+               "Inconsistent call frame state");
+    }
+#endif
+    if (!Info.StrongNoFrameOnEntry && Info.HasFrameOnEntry && !HasFrame) {
+      // Reset to the "after prologue" state.
+
+      // Insert a `.cfi_remember_state` into the last block known to have a
+      // stack frame.
+      unsigned CFIIndex =
+          MF.addFrameInst(MCCFIInstruction::createRememberState(nullptr));
+      BuildMI(*InsertMBB, InsertPt, DebugLoc(),
+              TII.get(TargetOpcode::CFI_INSTRUCTION))
+          .addCFIIndex(CFIIndex);
+      // Insert a `.cfi_restore_state` at the beginning of the current block.
+      CFIIndex = MF.addFrameInst(MCCFIInstruction::createRestoreState(nullptr));
+      InsertPt = BuildMI(*CurrBB, CurrBB->begin(), DebugLoc(),
+                         TII.get(TargetOpcode::CFI_INSTRUCTION))
+                     .addCFIIndex(CFIIndex);
+      ++InsertPt;
+      InsertMBB = &*CurrBB;
+      Change = true;
+    } else if ((Info.StrongNoFrameOnEntry || !Info.HasFrameOnEntry) &&
+               HasFrame) {
+      // Reset to the state upon function entry.
+      TFL.resetCFIToInitialState(*CurrBB);
+      Change = true;
+    }
+
+    HasFrame = Info.HasFrameOnExit;
+    ++CurrBB;
+  }
+
+  return Change;
+}

llvm/lib/CodeGen/CMakeLists.txt

@@ -38,6 +38,7 @@ add_llvm_component_library(LLVMCodeGen
   CalcSpillWeights.cpp
   CallingConvLower.cpp
   CFGuardLongjmp.cpp
+  CFIFixup.cpp
   CFIInstrInserter.cpp
   CodeGen.cpp
   CodeGenCommonISel.cpp

llvm/lib/CodeGen/CodeGen.cpp

@@ -24,6 +24,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeBranchFolderPassPass(Registry);
   initializeBranchRelaxationPass(Registry);
   initializeCFGuardLongjmpPass(Registry);
+  initializeCFIFixupPass(Registry);
   initializeCFIInstrInserterPass(Registry);
   initializeCheckDebugMachineModulePass(Registry);
   initializeCodeGenPreparePass(Registry);

llvm/lib/CodeGen/TailDuplicator.cpp

@@ -383,8 +383,9 @@ void TailDuplicator::duplicateInstruction(
   // Allow duplication of CFI instructions.
   if (MI->isCFIInstruction()) {
     BuildMI(*PredBB, PredBB->end(), PredBB->findDebugLoc(PredBB->begin()),
-      TII->get(TargetOpcode::CFI_INSTRUCTION)).addCFIIndex(
-      MI->getOperand(0).getCFIIndex());
+            TII->get(TargetOpcode::CFI_INSTRUCTION))
+        .addCFIIndex(MI->getOperand(0).getCFIIndex())
+        .setMIFlags(MI->getFlags());
     return;
   }
   MachineInstr &NewMI = TII->duplicate(*PredBB, PredBB->end(), *MI);

llvm/lib/CodeGen/TargetFrameLoweringImpl.cpp

@@ -21,6 +21,8 @@
 #include "llvm/IR/CallingConv.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstrTypes.h"
+#include "llvm/MC/MCAsmInfo.h"
+#include "llvm/MC/MCRegisterInfo.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
@@ -36,6 +38,11 @@ bool TargetFrameLowering::enableCalleeSaveSkip(const MachineFunction &MF) const
   return false;
 }
 
+bool TargetFrameLowering::enableCFIFixup(MachineFunction &MF) const {
+  return MF.needsFrameMoves() &&
+         !MF.getTarget().getMCAsmInfo()->usesWindowsCFI();
+}
+
 /// Returns the displacement from the frame register to the stack
 /// frame of the specified index, along with the frame register used
 /// (in output arg FrameReg). This is the default implementation which

llvm/lib/CodeGen/TargetPassConfig.cpp

@@ -138,6 +138,11 @@ static cl::opt<RunOutliner> EnableMachineOutliner(
                           "Disable all outlining"),
                // Sentinel value for unspecified option.
                clEnumValN(RunOutliner::AlwaysOutline, "", "")));
+// Disable the pass to fix unwind information. Whether the pass is included in
+// the pipeline is controlled via the target options, this option serves as
+// manual override.
+static cl::opt<bool> DisableCFIFixup("disable-cfi-fixup", cl::Hidden,
+                                     cl::desc("Disable the CFI fixup pass"));
 // Enable or disable FastISel. Both options are needed, because
 // FastISel is enabled by default with -fast, and we wish to be
 // able to enable or disable fast-isel independently from -O0.
@@ -1275,6 +1280,9 @@ void TargetPassConfig::addMachinePasses() {
     addPass(createMachineFunctionSplitterPass());
   }
 
+  if (!DisableCFIFixup && TM->Options.EnableCFIFixup)
+    addPass(createCFIFixup());
+
   // Add passes that directly emit MI after all other MI passes.
   addPreEmitPass2();