[WIP] Tracking PR for LFI integrated rewriter (LLVM 20)

clang/lib/Basic/Targets/AArch64.cpp

@@ -413,6 +413,10 @@ void AArch64TargetInfo::getTargetDefines(const LangOptions &Opts,
     Builder.defineMacro("__aarch64__");
   }
 
+  if (getTriple().isAArch64LFI()) {
+    Builder.defineMacro("__LFI__");
+  }
+
   // Inline assembly supports AArch64 flag outputs.
   Builder.defineMacro("__GCC_ASM_FLAG_OUTPUTS__");
 

clang/lib/Driver/CMakeLists.txt

@@ -64,6 +64,7 @@ add_clang_library(clangDriver
   ToolChains/Hexagon.cpp
   ToolChains/HLSL.cpp
   ToolChains/Hurd.cpp
+  ToolChains/LFILinux.cpp
   ToolChains/Linux.cpp
   ToolChains/MipsLinux.cpp
   ToolChains/MinGW.cpp

clang/lib/Driver/Driver.cpp

@@ -28,6 +28,7 @@
 #include "ToolChains/Haiku.h"
 #include "ToolChains/Hexagon.h"
 #include "ToolChains/Hurd.h"
+#include "ToolChains/LFILinux.h"
 #include "ToolChains/Lanai.h"
 #include "ToolChains/Linux.h"
 #include "ToolChains/MSP430.h"
@@ -6654,6 +6655,8 @@ const ToolChain &Driver::getToolChain(const ArgList &Args,
         TC = std::make_unique<toolchains::VEToolChain>(*this, Target, Args);
       else if (Target.isOHOSFamily())
         TC = std::make_unique<toolchains::OHOS>(*this, Target, Args);
+      else if (Target.isLFI())
+        TC = std::make_unique<toolchains::LFILinuxToolChain>(*this, Target, Args);
       else
         TC = std::make_unique<toolchains::Linux>(*this, Target, Args);
       break;

clang/lib/Driver/ToolChains/LFILinux.cpp

@@ -0,0 +1,29 @@
+//===-- LFILinux.cpp - LFI ToolChain Implementations ------------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "LFILinux.h"
+#include "clang/Driver/Driver.h"
+#include "clang/Driver/Options.h"
+
+using namespace clang::driver::toolchains;
+using namespace llvm::opt;
+
+void LFILinuxToolChain::AddCXXStdlibLibArgs(const ArgList &Args,
+                                            ArgStringList &CmdArgs) const {
+  switch (GetCXXStdlibType(Args)) {
+  case ToolChain::CST_Libcxx:
+    CmdArgs.push_back("-lc++");
+    if (Args.hasArg(options::OPT_fexperimental_library))
+      CmdArgs.push_back("-lc++experimental");
+    CmdArgs.push_back("-lc++abi");
+    break;
+  case ToolChain::CST_Libstdcxx:
+    CmdArgs.push_back("-lstdc++");
+    break;
+  }
+}

clang/lib/Driver/ToolChains/LFILinux.h

@@ -0,0 +1,35 @@
+//===--- LFILinux.h - LFI ToolChain Implementations -------------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_LIB_DRIVER_TOOLCHAINS_LFI_LINUX_H
+#define LLVM_CLANG_LIB_DRIVER_TOOLCHAINS_LFI_LINUX_H
+
+#include "Linux.h"
+
+namespace clang {
+namespace driver {
+namespace toolchains {
+
+class LLVM_LIBRARY_VISIBILITY LFILinuxToolChain : public Linux {
+public:
+  LFILinuxToolChain(const Driver &D, const llvm::Triple &Triple,
+                    const llvm::opt::ArgList &Args)
+      : Linux(D, Triple, Args) {
+    ExtraOpts.push_back("-z");
+    ExtraOpts.push_back("separate-code");
+  }
+
+  void AddCXXStdlibLibArgs(const llvm::opt::ArgList &Args,
+                           llvm::opt::ArgStringList &CmdArgs) const override;
+};
+
+} // end namespace toolchains
+} // end namespace driver
+} // end namespace clang
+
+#endif // LLVM_CLANG_LIB_DRIVER_TOOLCHAINS_LFI_LINUX_H

clang/tools/driver/cc1as_main.cpp

@@ -27,6 +27,7 @@
 #include "llvm/MC/MCCodeEmitter.h"
 #include "llvm/MC/MCContext.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/MCAsmParser.h"
@@ -577,6 +578,8 @@ static bool ExecuteAssemblerImpl(AssemblerInvocation &Opts,
     Str.reset(TheTarget->createMCObjectStreamer(
         T, Ctx, std::move(MAB), std::move(OW), std::move(CE), *STI));
     Str.get()->initSections(Opts.NoExecStack, *STI);
+    if (T.isLFI())
+      initializeLFIMCStreamer(*Str.get(), Ctx, T);
     if (T.isOSBinFormatMachO() && T.isOSDarwin()) {
       Triple *TVT = Opts.DarwinTargetVariantTriple
                         ? &*Opts.DarwinTargetVariantTriple

compiler-rt/cmake/Modules/AllSupportedArchDefs.cmake

@@ -1,4 +1,4 @@
-set(ARM64 aarch64)
+set(ARM64 aarch64 aarch64_lfi)
 set(ARM32 arm armhf)
 set(HEXAGON hexagon)
 set(X86 i386)

compiler-rt/cmake/builtin-config-ix.cmake

@@ -58,7 +58,7 @@ else()
 endif()
 
 set(AMDGPU amdgcn)
-set(ARM64 aarch64)
+set(ARM64 aarch64 aarch64_lfi)
 set(ARM32 arm armhf armv4t armv5te armv6 armv6m armv7m armv7em armv7 armv7s armv7k armv8m.base armv8m.main armv8.1m.main)
 set(AVR avr)
 set(HEXAGON hexagon)

compiler-rt/lib/builtins/CMakeLists.txt

@@ -649,6 +649,7 @@ set(armv7k_SOURCES ${arm_SOURCES})
 set(arm64_SOURCES ${aarch64_SOURCES})
 set(arm64e_SOURCES ${aarch64_SOURCES})
 set(arm64_32_SOURCES ${aarch64_SOURCES})
+set(aarch64_lfi_SOURCES ${aarch64_SOURCES})
 
 # macho_embedded archs
 set(armv6m_SOURCES ${thumb1_SOURCES})

compiler-rt/lib/hwasan/hwasan_interceptors.cpp

@@ -382,7 +382,9 @@ InternalLongjmp(__hw_register_buf env, int retval) {
       "ldp	x21, x22, [%0, #2<<3];"
       "ldp	x23, x24, [%0, #4<<3];"
       "ldp	x25, x26, [%0, #6<<3];"
+#ifndef __LFI__
       "ldp	x27, x28, [%0, #8<<3];"
+#endif
       "ldp	x29, x30, [%0, #10<<3];"
       "ldp	 d8,  d9, [%0, #14<<3];"
       "ldp	d10, d11, [%0, #16<<3];"

libunwind/src/UnwindRegistersRestore.S

@@ -655,8 +655,13 @@ DEFINE_LIBUNWIND_FUNCTION(__libunwind_Registers_arm64_jumpto)
   ldp    x20,x21, [x0, #0x0A0]
   ldp    x22,x23, [x0, #0x0B0]
   ldp    x24,x25, [x0, #0x0C0]
+#ifdef __LFI__
+  ldp    x26,xzr, [x0, #0x0D0]
+  ldp    xzr,x29, [x0, #0x0E0]
+#else
   ldp    x26,x27, [x0, #0x0D0]
   ldp    x28,x29, [x0, #0x0E0]
+#endif
   ldr    x30,     [x0, #0x100]  // restore pc into lr
 #if defined(__ARM_FP) && __ARM_FP != 0
   ldp    d0, d1,  [x0, #0x110]

llvm/include/llvm/MC/MCLFI.h

@@ -0,0 +1,22 @@
+//===- MCLFI.h - LFI-specific code for MC -----------------------*- 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
+//
+// This file was written by the Native Client authors, modified for LFI.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Support/CommandLine.h"
+
+namespace llvm {
+
+class MCContext;
+class MCStreamer;
+class Triple;
+
+void initializeLFIMCStreamer(MCStreamer &Streamer, MCContext &Ctx,
+                             const Triple &TheTriple);
+
+} // namespace llvm

llvm/include/llvm/MC/MCLFIExpander.h

@@ -0,0 +1,79 @@
+//===- llvm/MC/MCLFIExpander.h ----------------------------------*- 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
+//
+// This file was written by the Native Client authors, modified for LFI.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares the MCLFIExpander class. This is an abstract
+// class that encapsulates the expansion logic for MCInsts, and holds
+// state such as available scratch registers.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_MC_MCLFIEXPANDER_H
+#define LLVM_MC_MCLFIEXPANDER_H
+
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/MC/MCContext.h"
+#include "llvm/MC/MCInstrInfo.h"
+#include "llvm/MC/MCRegisterInfo.h"
+
+namespace llvm {
+class MCInst;
+class MCSubtargetInfo;
+class MCStreamer;
+
+class MCLFIExpander {
+private:
+  SmallVector<MCRegister, 2> ScratchRegs;
+  MCContext &Ctx;
+  bool Enabled = true;
+
+protected:
+  std::unique_ptr<MCInstrInfo> InstInfo;
+  std::unique_ptr<MCRegisterInfo> RegInfo;
+  void invalidateScratchRegs(const MCInst &Inst);
+  MCRegister getScratchReg(int index);
+  unsigned numScratchRegs() const;
+  virtual bool isValidScratchRegister(MCRegister Reg) const = 0;
+
+public:
+  MCLFIExpander(MCContext &Ctx, std::unique_ptr<MCRegisterInfo> &&RI,
+                std::unique_ptr<MCInstrInfo> &&II)
+      : Ctx(Ctx), InstInfo(std::move(II)), RegInfo(std::move(RI)) {}
+
+  void Error(const MCInst &Inst, const char msg[]);
+
+  bool addScratchReg(MCRegister Reg);
+  void clearScratchRegs();
+
+  void disable();
+  void enable();
+  bool isEnabled();
+
+  bool isPseudo(const MCInst &Inst) const;
+
+  bool mayAffectControlFlow(const MCInst &Inst) const;
+  bool isCall(const MCInst &Inst) const;
+  bool isBranch(const MCInst &Inst) const;
+  bool isIndirectBranch(const MCInst &Inst) const;
+  bool isReturn(const MCInst &Inst) const;
+  bool isVariadic(const MCInst &Inst) const;
+
+  bool mayLoad(const MCInst &Inst) const;
+  bool mayStore(const MCInst &Inst) const;
+
+  bool mayModifyRegister(const MCInst &Inst, MCRegister Reg) const;
+  bool explicitlyModifiesRegister(const MCInst &Inst, MCRegister Reg) const;
+
+  virtual ~MCLFIExpander() = default;
+  virtual bool expandInst(const MCInst &Inst, MCStreamer &Out,
+                          const MCSubtargetInfo &STI) = 0;
+};
+
+} // namespace llvm
+#endif

llvm/include/llvm/MC/MCStreamer.h

@@ -21,6 +21,7 @@
 #include "llvm/MC/MCDwarf.h"
 #include "llvm/MC/MCFragment.h"
 #include "llvm/MC/MCLinkerOptimizationHint.h"
+#include "llvm/MC/MCLFIExpander.h"
 #include "llvm/MC/MCPseudoProbe.h"
 #include "llvm/MC/MCWinEH.h"
 #include "llvm/Support/Error.h"
@@ -276,6 +277,8 @@ class MCStreamer {
   /// Returns true if the .cv_loc directive is in the right section.
   bool checkCVLocSection(unsigned FuncId, unsigned FileNo, SMLoc Loc);
 
+  std::unique_ptr<MCLFIExpander> LFIExpander;
+
 public:
   MCStreamer(const MCStreamer &) = delete;
   MCStreamer &operator=(const MCStreamer &) = delete;
@@ -293,6 +296,10 @@ class MCStreamer {
     return StartTokLocPtr ? *StartTokLocPtr : SMLoc();
   }
 
+  void setLFIExpander(MCLFIExpander *Exp) { LFIExpander.reset(Exp); }
+
+  MCLFIExpander *getLFIExpander() { return LFIExpander.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 MCLFIExpander;
 class MCObjectWriter;
 class MCRegisterInfo;
 class MCRelocationInfo;
@@ -230,6 +231,10 @@ class Target {
       mca::InstrumentManager *(*)(const MCSubtargetInfo &STI,
                                   const MCInstrInfo &MCII);
 
+  using MCLFIExpanderCtorTy = MCLFIExpander *(*)(
+      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.
@@ -340,6 +345,10 @@ class Target {
   /// InstrumentManager, if registered (default = nullptr).
   InstrumentManagerCtorTy InstrumentManagerCtorFn = nullptr;
 
+  // MCLFIExpanderCtorFn - Construction function for this target's
+  // MCLFIExpander, if registered.
+  MCLFIExpanderCtorTy MCLFIExpanderCtorFn;
+
 public:
   Target() = default;
 
@@ -538,6 +547,13 @@ class Target {
                                      const MCSubtargetInfo &STI, bool, bool,
                                      bool) const;
 
+  void createMCLFIExpander(MCStreamer &S,
+                           std::unique_ptr<MCRegisterInfo> &&RegInfo,
+                           std::unique_ptr<MCInstrInfo> &&InstInfo) const {
+    if (MCLFIExpanderCtorFn)
+      MCLFIExpanderCtorFn(S, std::move(RegInfo), std::move(InstInfo));
+  }
+
   MCStreamer *createAsmStreamer(MCContext &Ctx,
                                 std::unique_ptr<formatted_raw_ostream> OS,
                                 MCInstPrinter *IP,
@@ -1008,6 +1024,11 @@ struct TargetRegistry {
     T.InstrumentManagerCtorFn = Fn;
   }
 
+  static void RegisterMCLFIExpander(Target &T,
+                                    Target::MCLFIExpanderCtorTy Fn) {
+    T.MCLFIExpanderCtorFn = Fn;
+  }
+
   /// @}
 };
 

llvm/include/llvm/TargetParser/Triple.h

@@ -148,6 +148,8 @@ class Triple {
     AArch64SubArch_arm64e,
     AArch64SubArch_arm64ec,
 
+    AArch64SubArch_lfi,
+
     KalimbaSubArch_v3,
     KalimbaSubArch_v4,
     KalimbaSubArch_v5,
@@ -676,6 +678,16 @@ class Triple {
            getSubArch() == Triple::AArch64SubArch_arm64ec;
   }
 
+  bool isLFI() const {
+    return isAArch64LFI();
+  }
+
+  /// Checks if we're targeting the AArch64 LFI subarch.
+  bool isAArch64LFI() const {
+    return getArch() == Triple::aarch64 &&
+           getSubArch() == Triple::AArch64SubArch_lfi;
+  }
+
   bool isWindowsCoreCLREnvironment() const {
     return isOSWindows() && getEnvironment() == Triple::CoreCLR;
   }

llvm/lib/MC/CMakeLists.txt

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