blockmap: Split apart return edges from other dynamic edges.

This is required for dealing with compressed returns in the ykpt
decoder.

Author: vext01

llvm/include/llvm/CodeGen/TargetInstrInfo.h

@@ -1444,6 +1444,16 @@ class TargetInstrInfo : public MCInstrInfo {
     return false;
   }
 
+  /// Returns true if MI is a far call.
+  ///
+  /// YKFIXME: x86 only.
+  virtual bool isFarCall(const MachineInstr &MI) const { return false; }
+
+  /// Returns true if MI is a far return.
+  ///
+  /// YKFIXME: x86 only.
+  virtual bool isFarRet(const MachineInstr &MI) const { return false; }
+
   /// Returns true if the tail call can be made conditional on BranchCond.
   virtual bool canMakeTailCallConditional(SmallVectorImpl<MachineOperand> &Cond,
                                           const MachineInstr &TailCall) const {

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -1508,18 +1508,20 @@ void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF) {
         Unconditional = 0,
         // Choice of two successors.
         Conditional = 1,
-        // A control flow edge known only at runtime, e.g. a return or an
-        // indirect
+        // A return edge.
+        Return = 2,
+        // Any other control flow edge known only at runtime, e.g. an indirect
         // branch.
-        Dynamic = 2,
+        Dynamic = 3,
       };
 
       // Then depending of the `SuccessorKind` there is a payload describing the
       // successors:
       //
       // `Unconditional`: [target-address: u64]
       // `Conditional`:   [taken-address: u64, not-taken-address: u64]
-      // `Dynamic`:       [] (i.e. nothing, because nothing is statically known)
+      // `Return`:        []
+      // `Dynamic`:       []
 
       MachineBasicBlock *TBB = nullptr, *FBB = nullptr;
       SmallVector<MachineOperand> Conds;
@@ -1573,7 +1575,12 @@ void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF) {
         // Wasn't a branch or a fall-thru. Must be a different kind of
         // terminator.
         const MachineInstr *LastI = &*MBB.instr_rbegin();
-        if (LastI->isReturn() || LastI->isIndirectBranch()) {
+        if (LastI->isReturn()) {
+          // Ensure we are only working with near returns. See comment
+          // elsewhere about far calls for reasoning.
+          assert(!MF.getSubtarget().getInstrInfo()->isFarRet(*LastI));
+          OutStreamer->emitInt8(Return);
+        } else if (LastI->isIndirectBranch()) {
           OutStreamer->emitInt8(Dynamic);
         } else {
           std::string Msg = "unhandled block terminator in block: ";
@@ -1799,6 +1806,13 @@ void AsmPrinter::emitFunctionBody() {
             // Statically known function address.
             CallTargetSym = getSymbol(CallOpnd.getGlobal());
           }
+
+          // Ensure we are only working with near calls. This matters because
+          // Intel PT optimises near calls, and it simplifies our implementation
+          // if we can disregard far calls. Far calls are a hangover from
+          // 16-bit X86 that we are unlikley to encounter anyway.
+          assert(!MF->getSubtarget().getInstrInfo()->isFarCall(MI));
+
           assert(YkCallMarkerSyms.find(&MBB) != YkCallMarkerSyms.end());
           YkCallMarkerSyms[&MBB].push_back({YkPreCallSym, CallTargetSym});
         }

llvm/lib/Target/X86/X86InstrInfo.cpp

@@ -2913,6 +2913,33 @@ bool X86InstrInfo::isUnconditionalTailCall(const MachineInstr &MI) const {
   }
 }
 
+bool X86InstrInfo::isFarCall(const MachineInstr &MI) const {
+  switch (MI.getOpcode()) {
+  case X86::FARCALL16i:
+  case X86::FARCALL16m:
+  case X86::FARCALL32i:
+  case X86::FARCALL32m:
+  case X86::FARCALL64m:
+    return true;
+  default:
+    return false;
+  }
+}
+
+bool X86InstrInfo::isFarRet(const MachineInstr &MI) const {
+  switch (MI.getOpcode()) {
+  case X86::LRET16:
+  case X86::LRET32:
+  case X86::LRET64:
+  case X86::LRETI16:
+  case X86::LRETI32:
+  case X86::LRETI64:
+    return true;
+  default:
+    return false;
+  }
+}
+
 bool X86InstrInfo::canMakeTailCallConditional(
     SmallVectorImpl<MachineOperand> &BranchCond,
     const MachineInstr &TailCall) const {

llvm/lib/Target/X86/X86InstrInfo.h

@@ -316,6 +316,8 @@ class X86InstrInfo final : public X86GenInstrInfo {
 
   // Branch analysis.
   bool isUnconditionalTailCall(const MachineInstr &MI) const override;
+  bool isFarCall(const MachineInstr &MI) const override;
+  bool isFarRet(const MachineInstr &MI) const override;
   bool canMakeTailCallConditional(SmallVectorImpl<MachineOperand> &Cond,
                                   const MachineInstr &TailCall) const override;
   void replaceBranchWithTailCall(MachineBasicBlock &MBB,