[LFI] Add MCLFIRewriter infrastructure (#172906)

This is the second patch in the LFI series, adding the following
features:

* `MCLFIRewriter` class, which will be used to perform LFI rewrites on a
per-architecture basis. Each architecture where LFI is supported will
implement a subclass (for example,
`Target/AArch64/MCTargetDesc/AArch64MCLFIRewriter.cpp`) that will
implement the `rewriteInst` function to perform the actual rewriting
(the AArch64 version will be added in the next PR). The generic rewriter
class provides some instruction info utilities (`mayLoad`, `mayStore`)
and is used to call `rewriteInst` during instruction emission. It also
provides `onLabel` which allows the rewriter to know possible branch
targets, making certain optimizations like guard elimination possible.
* LFI streamer initialization that marks object files with a NOTE
section to indicate that the object file is using LFI.
* A basic LFI assembly parser that introduces the `.lfi_rewrite_disable`
and `.lfi_rewrite_enable` directives that can be used to control whether
rewriting is enabled or not in hand-written assembly.

llvm/docs/LFI.rst

@@ -398,6 +398,33 @@ In certain cases, guards may be hoisted outside of loops.
 |                       |                               |
 +-----------------------+-------------------------------+
 
+Assembler Directives
+====================
+
+The LFI assembler supports the following directives for controlling the
+rewriter.
+
+``.lfi_rewrite_disable``
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+Disables LFI assembly rewrites for all subsequent instructions, until
+``.lfi_rewrite_enable`` is used. This can be useful for hand-written assembly
+that is already safe and should not be modified by the rewriter.
+
+``.lfi_rewrite_enable``
+~~~~~~~~~~~~~~~~~~~~~~~
+
+Re-enables LFI assembly rewrites after a previous ``.lfi_rewrite_disable``.
+
+Example:
+
+.. code-block:: gas
+
+  .lfi_rewrite_disable
+  // No rewrites applied here.
+  ldr x0, [x27, w1, uxtw]
+  .lfi_rewrite_enable
+
 References
 ++++++++++
 

llvm/include/llvm/MC/MCLFI.h

@@ -0,0 +1,26 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+/// LFI-specific code for MC.
+///
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Compiler.h"
+
+namespace llvm {
+
+class MCContext;
+class MCStreamer;
+class Triple;
+
+LLVM_ABI void initializeLFIMCStreamer(MCStreamer &Streamer, MCContext &Ctx,
+                                      const Triple &TheTriple);
+
+} // namespace llvm

llvm/include/llvm/MC/MCLFIRewriter.h

@@ -0,0 +1,68 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+/// This file declares the MCLFIRewriter class, an abstract class that
+/// encapsulates the rewriting logic for MCInsts.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_MC_MCLFIREWRITER_H
+#define LLVM_MC_MCLFIREWRITER_H
+
+#include "llvm/MC/MCContext.h"
+#include "llvm/MC/MCInstrInfo.h"
+#include "llvm/MC/MCRegisterInfo.h"
+#include "llvm/Support/Compiler.h"
+
+namespace llvm {
+class MCInst;
+class MCSubtargetInfo;
+class MCStreamer;
+class MCSymbol;
+
+class MCLFIRewriter {
+private:
+  MCContext &Ctx;
+
+protected:
+  bool Enabled = true;
+  std::unique_ptr<MCInstrInfo> InstInfo;
+  std::unique_ptr<MCRegisterInfo> RegInfo;
+
+public:
+  MCLFIRewriter(MCContext &Ctx, std::unique_ptr<MCRegisterInfo> &&RI,
+                std::unique_ptr<MCInstrInfo> &&II)
+      : Ctx(Ctx), InstInfo(std::move(II)), RegInfo(std::move(RI)) {}
+
+  LLVM_ABI void error(const MCInst &Inst, const char Msg[]);
+
+  void disable() { Enabled = false; }
+  void enable() { Enabled = true; }
+
+  LLVM_ABI bool isCall(const MCInst &Inst) const;
+  LLVM_ABI bool isBranch(const MCInst &Inst) const;
+  LLVM_ABI bool isIndirectBranch(const MCInst &Inst) const;
+  LLVM_ABI bool isReturn(const MCInst &Inst) const;
+
+  LLVM_ABI bool mayLoad(const MCInst &Inst) const;
+  LLVM_ABI bool mayStore(const MCInst &Inst) const;
+
+  LLVM_ABI bool mayModifyRegister(const MCInst &Inst, MCRegister Reg) const;
+
+  virtual ~MCLFIRewriter() = default;
+  virtual bool rewriteInst(const MCInst &Inst, MCStreamer &Out,
+                           const MCSubtargetInfo &STI) = 0;
+
+  // Called when a label is emitted. Used for optimizations that require
+  // information about jump targets, such as guard elimination.
+  virtual void onLabel(const MCSymbol *Symbol) {}
+};
+
+} // namespace llvm
+#endif

llvm/include/llvm/MC/MCParser/MCAsmParserExtension.h

@@ -11,6 +11,7 @@
 
 #include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/MC/MCLFIRewriter.h"
 #include "llvm/MC/MCParser/MCAsmParser.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/SMLoc.h"
@@ -127,6 +128,7 @@ MCAsmParserExtension *createCOFFMasmParser();
 MCAsmParserExtension *createGOFFAsmParser();
 MCAsmParserExtension *createXCOFFAsmParser();
 MCAsmParserExtension *createWasmAsmParser();
+MCAsmParserExtension *createLFIAsmParser(MCLFIRewriter *Exp);
 
 } // end namespace llvm
 

llvm/include/llvm/MC/MCStreamer.h

@@ -19,6 +19,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/MC/MCDirectives.h"
 #include "llvm/MC/MCDwarf.h"
+#include "llvm/MC/MCLFIRewriter.h"
 #include "llvm/MC/MCLinkerOptimizationHint.h"
 #include "llvm/MC/MCPseudoProbe.h"
 #include "llvm/MC/MCSection.h"
@@ -291,6 +292,8 @@ protected:
   /// Returns true if the .cv_loc directive is in the right section.
   bool checkCVLocSection(unsigned FuncId, unsigned FileNo, SMLoc Loc);
 
+  std::unique_ptr<MCLFIRewriter> LFIRewriter;
+
 public:
   MCStreamer(const MCStreamer &) = delete;
   MCStreamer &operator=(const MCStreamer &) = delete;
@@ -308,6 +311,12 @@ public:
     return StartTokLocPtr ? *StartTokLocPtr : SMLoc();
   }
 
+  void setLFIRewriter(std::unique_ptr<MCLFIRewriter> Rewriter) {
+    LFIRewriter = std::move(Rewriter);
+  }
+
+  MCLFIRewriter *getLFIRewriter() { return LFIRewriter.get(); }
+
   /// State management
   ///
   virtual void reset();

llvm/include/llvm/MC/TargetRegistry.h

@@ -46,6 +46,7 @@ class MCDisassembler;
 class MCInstPrinter;
 class MCInstrAnalysis;
 class MCInstrInfo;
+class MCLFIRewriter;
 class MCObjectWriter;
 class MCRegisterInfo;
 class MCRelocationInfo;
@@ -237,6 +238,11 @@ public:
       mca::InstrumentManager *(*)(const MCSubtargetInfo &STI,
                                   const MCInstrInfo &MCII);
 
+  using MCLFIRewriterCtorTy =
+      MCLFIRewriter *(*)(MCStreamer & S,
+                         std::unique_ptr<MCRegisterInfo> &&RegInfo,
+                         std::unique_ptr<MCInstrInfo> &&InstInfo);
+
 private:
   /// Next - The next registered target in the linked list, maintained by the
   /// TargetRegistry.
@@ -351,6 +357,10 @@ private:
   /// InstrumentManager, if registered (default = nullptr).
   InstrumentManagerCtorTy InstrumentManagerCtorFn = nullptr;
 
+  // MCLFIRewriterCtorFn - Construction function for this target's
+  // MCLFIRewriter, if registered (default = nullptr).
+  MCLFIRewriterCtorTy MCLFIRewriterCtorFn = nullptr;
+
 public:
   Target() = default;
 
@@ -566,6 +576,13 @@ public:
     return nullptr;
   }
 
+  void createMCLFIRewriter(MCStreamer &S,
+                           std::unique_ptr<MCRegisterInfo> &&RegInfo,
+                           std::unique_ptr<MCInstrInfo> &&InstInfo) const {
+    if (MCLFIRewriterCtorFn)
+      MCLFIRewriterCtorFn(S, std::move(RegInfo), std::move(InstInfo));
+  }
+
   /// createMCRelocationInfo - Create a target specific MCRelocationInfo.
   ///
   /// \param TT The target triple.
@@ -1021,6 +1038,10 @@ struct TargetRegistry {
     T.InstrumentManagerCtorFn = Fn;
   }
 
+  static void RegisterMCLFIRewriter(Target &T, Target::MCLFIRewriterCtorTy Fn) {
+    T.MCLFIRewriterCtorFn = Fn;
+  }
+
   /// @}
 };
 

llvm/lib/MC/CMakeLists.txt

@@ -31,6 +31,8 @@ add_llvm_component_library(LLVMMC
   MCInstrAnalysis.cpp
   MCInstrDesc.cpp
   MCInstrInfo.cpp
+  MCLFI.cpp
+  MCLFIRewriter.cpp
   MCLabel.cpp
   MCLinkerOptimizationHint.cpp
   MCMachOStreamer.cpp

llvm/lib/MC/MCAsmStreamer.cpp

@@ -19,6 +19,7 @@
 #include "llvm/MC/MCExpr.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/MC/MCInstPrinter.h"
+#include "llvm/MC/MCLFIRewriter.h"
 #include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCObjectWriter.h"
 #include "llvm/MC/MCPseudoProbe.h"
@@ -2494,6 +2495,9 @@ void MCAsmStreamer::AddEncodingComment(const MCInst &Inst,
 
 void MCAsmStreamer::emitInstruction(const MCInst &Inst,
                                     const MCSubtargetInfo &STI) {
+  if (LFIRewriter && LFIRewriter->rewriteInst(Inst, *this, STI))
+    return;
+
   if (CurFrag) {
     MCSection *Sec = getCurrentSectionOnly();
     Sec->setHasInstructions(true);

llvm/lib/MC/MCLFI.cpp

@@ -0,0 +1,78 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+/// LFI-specific MC implementation.
+///
+//===----------------------------------------------------------------------===//
+
+#include "llvm/MC/MCLFI.h"
+#include "llvm/BinaryFormat/ELF.h"
+#include "llvm/MC/MCContext.h"
+#include "llvm/MC/MCLFIRewriter.h"
+#include "llvm/MC/MCSectionELF.h"
+#include "llvm/MC/MCStreamer.h"
+#include "llvm/MC/TargetRegistry.h"
+#include "llvm/Support/Alignment.h"
+#include "llvm/TargetParser/Triple.h"
+
+static const char NoteNamespace[] = "LFI";
+
+namespace llvm {
+
+cl::opt<bool> FlagEnableRewriting("lfi-enable-rewriter",
+                                  cl::desc("Enable rewriting for LFI."),
+                                  cl::init(true));
+
+void initializeLFIMCStreamer(MCStreamer &Streamer, MCContext &Ctx,
+                             const Triple &TheTriple) {
+  assert(TheTriple.isLFI());
+  const char *NoteName;
+  const char *NoteArch;
+  switch (TheTriple.getArch()) {
+  case Triple::aarch64:
+    NoteName = ".note.LFI.ABI.aarch64";
+    NoteArch = "aarch64";
+    break;
+  default:
+    reportFatalUsageError("Unsupported architecture for LFI");
+  }
+
+  std::string Error; // empty
+  const Target *TheTarget = TargetRegistry::lookupTarget(TheTriple, Error);
+
+  // Create the Target specific MCLFIRewriter.
+  assert(TheTarget != nullptr);
+  if (FlagEnableRewriting) {
+    TheTarget->createMCLFIRewriter(
+        Streamer,
+        std::unique_ptr<MCRegisterInfo>(TheTarget->createMCRegInfo(TheTriple)),
+        std::unique_ptr<MCInstrInfo>(TheTarget->createMCInstrInfo()));
+  }
+
+  // Emit an ELF Note section in its own COMDAT group which identifies LFI
+  // object files.
+  MCSectionELF *Note = Ctx.getELFSection(NoteName, ELF::SHT_NOTE,
+                                         ELF::SHF_ALLOC | ELF::SHF_GROUP, 0,
+                                         NoteName, /*IsComdat=*/true);
+
+  Streamer.pushSection();
+  Streamer.switchSection(Note);
+  Streamer.emitIntValue(strlen(NoteNamespace) + 1, 4);
+  Streamer.emitIntValue(strlen(NoteArch) + 1, 4);
+  Streamer.emitIntValue(ELF::NT_VERSION, 4);
+  Streamer.emitBytes(NoteNamespace);
+  Streamer.emitIntValue(0, 1); // NUL terminator
+  Streamer.emitValueToAlignment(Align(4));
+  Streamer.emitBytes(NoteArch);
+  Streamer.emitIntValue(0, 1); // NUL terminator
+  Streamer.emitValueToAlignment(Align(4));
+  Streamer.popSection();
+}
+
+} // namespace llvm

llvm/lib/MC/MCLFIRewriter.cpp

@@ -0,0 +1,54 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+/// This file implements the MCLFIRewriter class, a base class that
+/// encapsulates the rewriting logic for MCInsts.
+///
+//===----------------------------------------------------------------------===//
+
+#include "llvm/MC/MCLFIRewriter.h"
+#include "llvm/MC/MCInst.h"
+#include "llvm/MC/MCInstrInfo.h"
+#include "llvm/MC/MCRegisterInfo.h"
+
+namespace llvm {
+
+void MCLFIRewriter::error(const MCInst &Inst, const char Msg[]) {
+  Ctx.reportError(Inst.getLoc(), Msg);
+}
+
+bool MCLFIRewriter::isCall(const MCInst &Inst) const {
+  return InstInfo->get(Inst.getOpcode()).isCall();
+}
+
+bool MCLFIRewriter::isBranch(const MCInst &Inst) const {
+  return InstInfo->get(Inst.getOpcode()).isBranch();
+}
+
+bool MCLFIRewriter::isIndirectBranch(const MCInst &Inst) const {
+  return InstInfo->get(Inst.getOpcode()).isIndirectBranch();
+}
+
+bool MCLFIRewriter::isReturn(const MCInst &Inst) const {
+  return InstInfo->get(Inst.getOpcode()).isReturn();
+}
+
+bool MCLFIRewriter::mayLoad(const MCInst &Inst) const {
+  return InstInfo->get(Inst.getOpcode()).mayLoad();
+}
+
+bool MCLFIRewriter::mayStore(const MCInst &Inst) const {
+  return InstInfo->get(Inst.getOpcode()).mayStore();
+}
+
+bool MCLFIRewriter::mayModifyRegister(const MCInst &Inst,
+                                      MCRegister Reg) const {
+  return InstInfo->get(Inst.getOpcode()).hasDefOfPhysReg(Inst, Reg, *RegInfo);
+}
+} // namespace llvm

llvm/lib/MC/MCObjectStreamer.cpp

@@ -15,6 +15,7 @@
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCDwarf.h"
 #include "llvm/MC/MCExpr.h"
+#include "llvm/MC/MCLFIRewriter.h"
 #include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCObjectWriter.h"
 #include "llvm/MC/MCSFrame.h"
@@ -398,6 +399,9 @@ bool MCObjectStreamer::mayHaveInstructions(MCSection &Sec) const {
 
 void MCObjectStreamer::emitInstruction(const MCInst &Inst,
                                        const MCSubtargetInfo &STI) {
+  if (LFIRewriter && LFIRewriter->rewriteInst(Inst, *this, STI))
+    return;
+
   MCStreamer::emitInstruction(Inst, STI);
 
   MCSection *Sec = getCurrentSectionOnly();

llvm/lib/MC/MCParser/AsmParser.cpp

@@ -120,6 +120,7 @@ private:
   SourceMgr::DiagHandlerTy SavedDiagHandler;
   void *SavedDiagContext;
   std::unique_ptr<MCAsmParserExtension> PlatformParser;
+  std::unique_ptr<MCAsmParserExtension> LFIParser;
   SMLoc StartTokLoc;
   std::optional<SMLoc> CFIStartProcLoc;
 
@@ -788,6 +789,10 @@ AsmParser::AsmParser(SourceMgr &SM, MCContext &Ctx, MCStreamer &Out,
   }
 
   PlatformParser->Initialize(*this);
+  if (Out.getLFIRewriter()) {
+    LFIParser.reset(createLFIAsmParser(Out.getLFIRewriter()));
+    LFIParser->Initialize(*this);
+  }
   initializeDirectiveKindMap();
   initializeCVDefRangeTypeMap();
 }

llvm/lib/MC/MCParser/CMakeLists.txt

@@ -6,6 +6,7 @@ add_llvm_component_library(LLVMMCParser
   GOFFAsmParser.cpp
   DarwinAsmParser.cpp
   ELFAsmParser.cpp
+  LFIAsmParser.cpp
   MCAsmParser.cpp
   MCAsmParserExtension.cpp
   MCTargetAsmParser.cpp

llvm/lib/MC/MCParser/LFIAsmParser.cpp

@@ -0,0 +1,70 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+/// LFI assembly parser.
+///
+//===----------------------------------------------------------------------===//
+
+#include "llvm/MC/MCLFIRewriter.h"
+#include "llvm/MC/MCParser/MCAsmParserExtension.h"
+#include "llvm/MC/MCStreamer.h"
+
+using namespace llvm;
+
+class LFIAsmParser : public MCAsmParserExtension {
+  MCLFIRewriter *Rewriter;
+  template <bool (LFIAsmParser::*HandlerMethod)(StringRef, SMLoc)>
+  void addDirectiveHandler(StringRef Directive) {
+    MCAsmParser::ExtensionDirectiveHandler Handler =
+        std::make_pair(this, HandleDirective<LFIAsmParser, HandlerMethod>);
+
+    getParser().addDirectiveHandler(Directive, Handler);
+  }
+
+public:
+  LFIAsmParser(MCLFIRewriter *Exp) : Rewriter(Exp) {}
+  void Initialize(MCAsmParser &Parser) override {
+    // Call the base implementation.
+    MCAsmParserExtension::Initialize(Parser);
+    addDirectiveHandler<&LFIAsmParser::parseRewriteDisable>(
+        ".lfi_rewrite_disable");
+    addDirectiveHandler<&LFIAsmParser::parseRewriteEnable>(
+        ".lfi_rewrite_enable");
+  }
+
+  /// ::= {.lfi_rewrite_disable}
+  bool parseRewriteDisable(StringRef Directive, SMLoc Loc) {
+    getParser().checkForValidSection();
+    if (getLexer().isNot(AsmToken::EndOfStatement))
+      return TokError("unexpected token");
+    Lex();
+
+    Rewriter->disable();
+
+    return false;
+  }
+
+  /// ::= {.lfi_rewrite_enable}
+  bool parseRewriteEnable(StringRef Directive, SMLoc Loc) {
+    getParser().checkForValidSection();
+    if (getLexer().isNot(AsmToken::EndOfStatement))
+      return TokError("unexpected token");
+    Lex();
+
+    Rewriter->enable();
+
+    return false;
+  }
+};
+
+namespace llvm {
+MCAsmParserExtension *createLFIAsmParser(MCLFIRewriter *Exp) {
+  return new LFIAsmParser(Exp);
+}
+} // namespace llvm

llvm/lib/MC/MCStreamer.cpp

@@ -399,6 +399,9 @@ void MCStreamer::emitLabel(MCSymbol *Symbol, SMLoc Loc) {
 
   Symbol->setFragment(&getCurrentSectionOnly()->getDummyFragment());
 
+  if (LFIRewriter)
+    LFIRewriter->onLabel(Symbol);
+
   MCTargetStreamer *TS = getTargetStreamer();
   if (TS)
     TS->emitLabel(Symbol);

llvm/lib/MC/TargetRegistry.cpp

@@ -13,6 +13,7 @@
 #include "llvm/MC/MCCodeEmitter.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCInstPrinter.h"
+#include "llvm/MC/MCLFI.h"
 #include "llvm/MC/MCObjectStreamer.h"
 #include "llvm/MC/MCObjectWriter.h"
 #include "llvm/Support/raw_ostream.h"
@@ -76,6 +77,8 @@ MCStreamer *Target::createMCObjectStreamer(
   }
   if (ObjectTargetStreamerCtorFn)
     ObjectTargetStreamerCtorFn(*S, STI);
+  if (T.isLFI())
+    initializeLFIMCStreamer(*S, Ctx, T);
   return S;
 }
 

llvm/tools/llvm-mc/llvm-mc.cpp

@@ -21,6 +21,7 @@
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCInstPrinter.h"
 #include "llvm/MC/MCInstrInfo.h"
+#include "llvm/MC/MCLFI.h"
 #include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCObjectWriter.h"
 #include "llvm/MC/MCParser/AsmLexer.h"
@@ -626,6 +627,11 @@ int main(int argc, char **argv) {
     Str.reset(TheTarget->createAsmStreamer(Ctx, std::move(FOut), std::move(IP),
                                            std::move(CE), std::move(MAB)));
 
+    Triple T(TripleName);
+    if (T.isLFI()) {
+      Str->initSections(NoExecStack, *STI);
+      initializeLFIMCStreamer(*Str.get(), Ctx, T);
+    }
   } else if (FileType == OFT_Null) {
     Str.reset(TheTarget->createNullStreamer(Ctx));
   } else {