[CodeGen] Add extension points for TargetPassConfig::addMachinePasses

[drti]

Summary

Allow out-of-tree code to customize a target using a
mechanism similar to addGlobalExtension in PassManagerBuilder but for
machine passes.

[drti]

Hi this is my first submission to the LLVM project. I selected reviewers by looking at the recent commit history of TargetPassConfig.h so I hope I haven't gone too far wrong with that. Evan Cheng is listed in the CODE_OWNERS.txt for "parts of code generator not covered by someone else" but I couldn't find his user name on Phabricator (perhaps he doesn't have one).

[ychen]

Could you add an example like llvm/examples/Bye and a test using llc -load?

llvm/include/llvm/CodeGen/TargetPassConfig.h

367

Why is Target needed? A plugin could query TargetPassConfig about which target it wants to register itself.

llvm/unittests/CodeGen/TargetExtensionTest.cpp

54

I think testing the default target should be enough. Why we test all targets here?

[drti]

Thanks very much @ychen for looking at my review, and thanks for pointing me at the examples/Bye/Bye.cpp example code. I didn't know we could do that kind of thing in the build and I'll post an update with a new example later this week.

I see your other two comments as being closely related to each other. If we handle target filtering in the TargetPassConfig implementation then it makes sense for the test to confirm that the filtering works, i.e. an extension registered for one target is not mistakenly invoked for a different one. If we put the responsibility on the client code to filter out unwanted targets then the test wouldn't have to care and any target machine (so long as it can be cast to LLVMTargetMachine, of course) would do.

I can see one benefit to doing the filtering in the implementation rather than the client - in makes it obvious that the filtering is necessary, i.e. the API forces the client code to select a specific target. Otherwise the client code could omit the filtering and then it would work during developer testing say, but crash mysteriously if invoked on an unsuitable target. Maybe that's unlikely, so please let me know if you'd prefer me to simpilfy the implementation by removing the filtering.

[ychen]

In D98591#2630098, @drti wrote:

so please let me know if you'd prefer me to simpilfy the implementation by removing the filtering.

Yeah, I still think the filtering should be removed. The plugin should not reply on a Target instance to work and ideally, the plugin should be able to cope with more than one target (most likely a machine-independent pass).

[drti]

In D98591#2627643, @ychen wrote:

Could you add an example like llvm/examples/Bye and a test using llc -load?

I've made the changes to remove the filtering and simplify the test but I'm having some trouble with the Bye example. When I enable this via cmake ... -DLLVM_BYE_LINK_INTO_TOOLS=ON I get some test failures from llvm-check. Is that expected?

Failed Tests (6):

LLVM :: Other/new-pm-O0-defaults.ll
LLVM :: Other/opt-O0-pipeline.ll
LLVM :: Other/opt-O2-pipeline.ll
LLVM :: Other/opt-O3-pipeline-enable-matrix.ll
LLVM :: Other/opt-O3-pipeline.ll
LLVM :: Other/opt-Os-pipeline.ll

Example error output:

• TEST 'LLVM :: Other/opt-O3-pipeline-enable-matrix.ll' FAILED ****

Script:

: 'RUN: at line 1'; /home/raoul/build/llvm/head-static/bin/opt -enable-new-pm=0 -O3 -enable-matrix -debug-pass=Structure < /home/raoul/src/llvm/llvm-project/llvm/test/Other/opt-O3-pipeline-enable-matrix.ll -o /dev/null 2>&1 | /home/raoul/build/llvm/head-static/bin/FileCheck /home/raoul/src/llvm/llvm-project/llvm/test/Other/opt-O3-pipeline-enable-matrix.ll

Exit Code: 1

Command Output (stderr):

/home/raoul/src/llvm/llvm-project/llvm/test/Other/opt-O3-pipeline-enable-matrix.ll:216:15: error: CHECK-NEXT: is not on the line after the previous match
; CHECK-NEXT: Canonicalize natural loops

^

<stdin>:212:2: note: 'next' match was here
Canonicalize natural loops
^
<stdin>:208:11: note: previous match ended here
Early CSE

^

<stdin>:209:1: note: non-matching line after previous match is here
Good Bye World Pass
^

Input file: <stdin>
Check file: /home/raoul/src/llvm/llvm-project/llvm/test/Other/opt-O3-pipeline-enable-matrix.ll

-dump-input=help explains the following input dump.

Input was:
<<<<<<

  .
  .
  .
207:  Lower the matrix intrinsics 
208:  Early CSE 
209:  Good Bye World Pass 
210:  Dominator Tree Construction 
211:  Natural Loop Information 
212:  Canonicalize natural loops

next:216 !~~~~~~~~~~~~~~~~~~~~~~~~~ error: match on wrong line

213:  LCSSA Verifier 
214:  Loop-Closed SSA Form Pass 
215:  Basic Alias Analysis (stateless AA impl) 
216:  Function Alias Analysis Results 
217:  Scalar Evolution Analysis 
  .
  .
  .

[drti]

[CodeGen] Simplify unit test for TargetPassConfig::addExtension and
add example plugin module

The example MachineExtension plugin module has a lit test which runs
if examples are enabled.

[drti]

@ychen I've updated this review with your suggested changes. I've removed the extension filtering in TargetPassConfig, simplified the test and added an example. The example doesn't try to use the new pass manager because the last time I checked it still wasn't useful for machine passes. What do you think?

[ychen]

@drti hi, could you please rebase the patch? The base commit seems not in the git history.

[drti]

[CodeGen] Add extension points for TargetPassConfig::addMachinePasses

Allow out-of-tree code to customize a target using a
mechanism similar to addGlobalExtension in PassManagerBuilder but for
machine passes.

[drti]

@ychen apologies for that, the base revision should now be efc013ec4d950a68e6f80dd98cda35e1a96a6fc8 from the main repo.

[drti]

[CodeGen] Add extension points for TargetPassConfig::addMachinePasses

Allow out-of-tree code to customize a target using a
mechanism similar to addGlobalExtension in PassManagerBuilder but for
machine passes.

[drti]

[CodeGen] Add extension points for TargetPassConfig::addMachinePasses

Allow out-of-tree code to customize a target using a
mechanism similar to addGlobalExtension in PassManagerBuilder but for
machine passes.

[drti]

[CodeGen] Add extension points for TargetPassConfig::addMachinePasses

Allow out-of-tree code to customize a target using a
mechanism similar to addGlobalExtension in PassManagerBuilder but for
machine passes.

[drti]

@ychen I've rebased and fixed some pre-merge build and formatting issues. Could you please have another look at this review? If it looks OK could I also ask you to commit it for me, please? It's my first submission and I don't have commit rights myself.
Regards,
Raoul

[drti]

Ping

[ychen]

llvm/examples/Bye/CMakeLists.txt

10 ↗

(On Diff #345076)

Why this?

llvm/examples/CMakeLists.txt

11

I think it's better to call this MachineBye to reflect that the usage/implementation are similar between the two. If you think it makes sense, make sure the class name is also renamed accordingly.

llvm/examples/MachineExtension/MachineExtension.cpp

33 ↗

(On Diff #345076)

No need for anything concrete there. I think printing some string using errs() should be enough.

62 ↗

(On Diff #345076)

Better to use lambda for this one-time registration.

llvm/include/llvm/CodeGen/TargetPassConfig.h

91

How about reducing these extension points to a list that very likely to be useful? Unlike IR pipeline, the order among machine passes is pretty firm, if there are no concrete use cases for an extension point, it may become dead code, or we would need to move/rename it in the future.

505

I'd prefer to call it addMachinePassExtentions.

518

format

529

format

llvm/include/llvm/Target/TargetMachine.h

170

asLLVMTargetMachine is not needed. If you could just static_cast to LLVMTargetMachine directly since TargetMachine object is almost never created.

[drti]

llvm/examples/Bye/CMakeLists.txt

10 ↗

(On Diff #345076)

On my Linux system at least (Fedora 33) the Bye example fails to build when configured with -DBUILD_SHARED_LIBS=true. Strictly speaking this belongs in a separate review but seeing as I'm fixing it for the new example I thought I'd sneak this one in.

The failures is because -zdefs is enabled via CMAKE_SHARED_LINKER_FLAGS in llvm/cmake/modules/HandleLLVMOptions.cmake - this linker flag effectively introduces the same restriction against undefined symbols in shared objects as applies to DLLs on Windows (where the Bye example is disabled)

Example error:
examples/Bye/CMakeFiles/Bye.dir/Bye.cpp.o:Bye.cpp:vtable for (anonymous namespace)::LegacyBye: error: undefined reference to 'llvm::Pass::getPassName() const'

llvm/include/llvm/CodeGen/TargetPassConfig.h

91

I agree that we should limit the number of extension points, but I think we also should try to cater for a range of possible usages. The list at the moment more-or-less mirrors the virtual function extension points for the TargetPassConfig derived classes, which I thought was a fair compromise. There's a tradeoff here of course, but I would tend to err on the side of having more rather than fewer, simply because it's a high barrier for someone working on out-of-tree code to get an extension point added if necessary, compared to an in-tree change that needs, say, a new virtual function extension point.

518

This was auto-formatted this way, either during arcanist processing or to satisfy some lint/format warnings. I'd be hesitant to change this, even though I agree it doesn't look "right".

529

Also auto-formatted this way.

llvm/include/llvm/Target/TargetMachine.h

170

I've seen quite a few places in the code doing the exact unchecked static_cast you suggest, but I hadn't wanted to follow that pattern myself. I didn't think this base class warrants the full LLVM RTTI support, and the virtual function seemed like an easy solution. I had intended to submit another review to introduce it in the existing use-cases that are simple enough (see, e.g. the related TODO comment in llvm/tools/llvm-exegesis/lib/Assembler.cpp and FIXME in llvm/tools/opt/opt.cpp)

[drti]

[CodeGen] Add extension points for TargetPassConfig::addMachinePasses

Allow out-of-tree code to customize a target using a
mechanism similar to addGlobalExtension in PassManagerBuilder but for
machine passes.

[drti]

llvm/include/llvm/CodeGen/TargetPassConfig.h

505

Well my reasoning was that addExtension "adds" an extension function and "apply" runs the functions (if any). It may also be that a callback doesn't add any passes for a particular target so strictly speaking all we're doing is invoking the callbacks. What do you think?

[drti]

ping

[drti]

ping

[drti]

ping

[drti]

ping

[drti]

ping

[drti]

@ychen could you please take another look at this review? I'll understand if you don't have time for this - perhaps you could let me know if I should look for another reviewer. It's been quite a while now...

[arsenm]

I'm sure this needs a rebase. Should also be reviewed by out of tree target maintainers

[drti]

Hi @arsenm thanks for taking a look at this. I'm in the process of rebasing and re-testing but it will take me a little while to get back up to speed on this. I think it makes sense to put it in the "Plan Changes" stage for now.

Regards,
@drti

[drti]

[CodeGen] Add extension points for TargetPassConfig::addMachinePasses
    
Allow out-of-tree code to customize a target using a
mechanism similar to addGlobalExtension in PassManagerBuilder but for
machine passes.

[drti]

[CodeGen] Add extension points for TargetPassConfig::addMachinePasses

Rebase to see if the lljit-with-thinlto-summaries.test has been fixed

[drti]

@dsanders I think this is in a good state for review again, could you please have a look? There were some outstanding comments from @ychen earlier which you might want to consider (e.g. he thought we needed fewer extension points). Thanks @drti

[arsenm]

llvm/include/llvm/Target/TargetMachine.h

440

I don't understand the point of this, there's no cast here

llvm/lib/CodeGen/TargetPassConfig.cpp

730

Is removal ever really needed?

[drti]

llvm/include/llvm/Target/TargetMachine.h

440

This overrides the baseclass version of the function which returns nullptr. I think ideally we'd have proper dyn_cast<LLVMTargetMachine> support but the TargetMachine base class isn't set up for this. The virtual function here is just an easy way to get safer downcasting rather than client code using an unchecked static_cast. @ychen was also questioning the utility of this, by the way, suggesting to use the static_cast like some existing code already does. What do you think?

llvm/lib/CodeGen/TargetPassConfig.cpp

730

I guess the only situation where it's useful is if shared objects/DLLs are dynamically unloaded - any registered extension would crash if not removed before unloading its code.

[drti]

Hi @arsenm I'm not really sure what the process is here. I'm happy to change the code just wanted some clarification on your opinion - have commented on the issues you raised.

[arsenm]

Seems plausible. Hopefully this won't get too in the way of new PM codegen support

llvm/include/llvm/CodeGen/TargetPassConfig.h

100

Can this be function_ref?

518

Might as well try hitting this with clang-format again

[drti]

Rebase to latest LLVM, fix formatting issues

[drti]

Stop test running on Windows using the new UNSUPPORTED syntax

[drti]

Hi @arsenm I've rebased this so it builds against the current HEAD and addressed a couple of inline comments. I don't have commit rights myself, but if you're happy with the current state could you please merge it for me or let me know what the next step is otherwise?

Regards,
@drti

llvm/include/llvm/CodeGen/TargetPassConfig.h

100

I don't think it can in general because we need to store a list of these and the function_ref docs say "This class does not own the callable, so it is not in general safe to store a function_ref"

llvm/examples/CMakeLists.txt

 add_subdirectory(BrainF)
 add_subdirectory(Fibonacci)
 add_subdirectory(HowToUseJIT)
 add_subdirectory(HowToUseLLJIT)
 add_subdirectory(IRTransforms)
 add_subdirectory(Kaleidoscope)
 add_subdirectory(ModuleMaker)
 add_subdirectory(OrcV2Examples)
 add_subdirectory(SpeculativeJIT)
 add_subdirectory(Bye)
+add_subdirectory(MachineBye)
 
 if(LLVM_ENABLE_EH AND (NOT WIN32) AND (NOT "${LLVM_NATIVE_ARCH}" STREQUAL "ARM"))
     add_subdirectory(ExceptionDemo)
 endif()
 
 if( HAVE_PTHREAD_H )
   add_subdirectory(ParallelJIT)
 endif( HAVE_PTHREAD_H )

llvm/examples/MachineBye/CMakeLists.txt

+# The plugin expects not to link against the Support and Core libraries,
+# but expects them to exist in the process loading the plugin. This doesn't
+# work with DLLs on Windows (where a shared library can't have undefined
+# references), so just skip this example on Windows.
+if (NOT WIN32)
+  # Remove any -zdefs flag (which disallows undefined symbols in
+  # shared objects) added by the main cmake files on Linux when
+  # BUILD_SHARED_LIBS is enabled
+  STRING(REPLACE "-Wl,-z,defs" "" CMAKE_SHARED_LINKER_FLAGS
+    "${CMAKE_SHARED_LINKER_FLAGS}")
+
+  # We have to build this shared so we can llc -load it
+  add_llvm_example_library(MachineBye
+    MachineBye.cpp
+    MODULE
+    DEPENDS
+    intrinsics_gen
+   )
+endif()

llvm/examples/MachineBye/MachineBye.cpp

+// Example showing how to manipulate a MachineFunction using your own pass via
+// the TargetPassConfig::addExtension mechanism
+
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineModuleInfo.h"
+#include "llvm/CodeGen/TargetPassConfig.h"
+#include "llvm/IR/LegacyPassManager.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/Host.h"
+#include "llvm/Target/TargetMachine.h"
+
+using namespace llvm;
+
+namespace {
+class MachineByePass : public MachineFunctionPass {
+public:
+  static char ID;
+  bool runOnMachineFunction(MachineFunction &MF) override;
+  MachineByePass();
+};
+
+char MachineByePass::ID = 0;
+
+MachineByePass::MachineByePass() : MachineFunctionPass(ID) {}
+
+bool MachineByePass::runOnMachineFunction(MachineFunction &MF) {
+  errs() << "MachineBye: ";
+  errs().write_escaped(MF.getName()) << '\n';
+  return false;
+}
+
+} // end of anonymous namespace
+
+static RegisterTargetExtension Registrar(TargetPassConfig::MPEP_EarlyAsPossible,
+                                         [](TargetPassConfig &TPC) {
+                                           TPC.addPass(new MachineByePass);
+                                         });

llvm/include/llvm/CodeGen/TargetPassConfig.h

 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CODEGEN_TARGETPASSCONFIG_H
 #define LLVM_CODEGEN_TARGETPASSCONFIG_H
 
 #include "llvm/Pass.h"
 #include "llvm/Support/CodeGen.h"
 #include <cassert>
+#include <functional>
 #include <string>
 
 namespace llvm {
 
+class Target;
 class LLVMTargetMachine;
 struct MachineSchedContext;
 class PassConfigImpl;
 class ScheduleDAGInstrs;
 class CSEConfigBase;
 class PassInstrumentationCallbacks;
 
 // The old pass manager infrastructure is hidden in a legacy namespace now.

 };
 
 
 /// Target-Independent Code Generator Pass Configuration Options.
 ///
 /// This is an ImmutablePass solely for the purpose of exposing CodeGen options
 /// to the internals of other CodeGen passes.
 class TargetPassConfig : public ImmutablePass {
+public:
+  /// Enum identifying when in the machine pass setup to add an
+  /// extension via a callback.
+  enum MachinePassExtensionPointTy {
+    MPEP_EarlyAsPossible,
+    MPEP_PreRegAlloc,
+    MPEP_PostRegAlloc,
+    MPEP_PreSched2,
+    MPEP_PreEmitPass,
+    MPEP_LateAsPossible
+  };
+
+  /// Callback function type for added extensions to a target.
+  using ExtensionFn = std::function<void(TargetPassConfig &)>;
+
+  /// ID used for removing a previously added target extension
+  using ExtensionID = int;
+
 private:
   PassManagerBase *PM = nullptr;
   AnalysisID StartBefore = nullptr;
   AnalysisID StartAfter = nullptr;
   AnalysisID StopBefore = nullptr;
   AnalysisID StopAfter = nullptr;
 
   unsigned StartBeforeInstanceNum = 0;
public:
 
   /// Check whether continuous CSE should be enabled in GISel passes.
   /// By default, it's enabled for non O0 levels.
   virtual bool isGISelCSEEnabled() const;
 
   /// Returns the CSEConfig object to use for the current optimization level.
   virtual std::unique_ptr<CSEConfigBase> getCSEConfig() const;
 
+  /// Add an extension to be applied to all Targets, returning a
+  /// nonzero ID for use with removeExtension. The preferred way to
+  /// register is via the RegisterTargetExtension class defined
+  /// below. If the extension only applies to particular target(s) the
+  /// ExtensionFn callback can check the properties of its
+  /// TargetPassConfig argument, e.g. via TPC.getTM<TargetMachine>().
+  static ExtensionID addExtension(MachinePassExtensionPointTy MPEP,
+                                  ExtensionFn Fn);
+
+  /// Remove a previously added extension. The preferred way to use
+  /// this is via the RegisterTargetExtension class defined below.
+  static void removeExtension(ExtensionID ID);
+
+  /// Add a pass to the PassManager if that pass is supposed to be run, as
+  /// determined by the StartAfter and StopAfter options. Takes ownership of the
+  /// pass.
+  void addPass(Pass *P);
+
 protected:
   // Helper to verify the analysis is really immutable.
   void setOpt(bool &Opt, bool Val);
 
   /// Return true if register allocator is specified by -regalloc=override.
   bool isCustomizedRegAlloc();
 
   /// Methods with trivial inline returns are convenient points in the common
protected:
 
   /// Utilities for targets to add passes to the pass manager.
   ///
 
   /// Add a CodeGen pass at this point in the pipeline after checking overrides.
   /// Return the pass that was added, or zero if no pass was added.
   AnalysisID addPass(AnalysisID PassID);
 
-  /// Add a pass to the PassManager if that pass is supposed to be run, as
-  /// determined by the StartAfter and StopAfter options. Takes ownership of the
-  /// pass.
-  void addPass(Pass *P);
-
   /// addMachinePasses helper to create the target-selected or overriden
   /// regalloc pass.
   virtual FunctionPass *createRegAllocPass(bool Optimized);
 
   /// Add core register allocator passes which do the actual register assignment
   /// and rewriting. \returns true if any passes were added.
   virtual bool addRegAssignAndRewriteFast();
   virtual bool addRegAssignAndRewriteOptimized();
+
+private:
+  /// Scan the set of extensions and call any registered for our
+  /// Target at the given extension point.
+  void applyAnyExtensions(MachinePassExtensionPointTy MPEP);
 };
 
 void registerCodeGenCallback(PassInstrumentationCallbacks &PIC,
                              LLVMTargetMachine &);
 
+/// Registers and de-registers a target extension function
+class RegisterTargetExtension {
+  TargetPassConfig::ExtensionID ExtensionID;
+
+public:
+  RegisterTargetExtension(TargetPassConfig::MachinePassExtensionPointTy MPEP,
+                          TargetPassConfig::ExtensionFn Fn)
+      : ExtensionID(TargetPassConfig::addExtension(MPEP, std::move(Fn))) {}
+
+  ~RegisterTargetExtension() {
+    if (ExtensionID)
+      TargetPassConfig::removeExtension(ExtensionID);
+  }
+
+  // Movable but not copyable
+  RegisterTargetExtension(RegisterTargetExtension &&Other)
+      : ExtensionID(Other.ExtensionID) {
+    Other.ExtensionID = 0;
+  }
+
+  RegisterTargetExtension(const RegisterTargetExtension &) = delete;
+  RegisterTargetExtension &operator=(const RegisterTargetExtension &) = delete;
+  RegisterTargetExtension &operator=(RegisterTargetExtension &&) = delete;
+};
+
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_TARGETPASSCONFIG_H

llvm/include/llvm/Target/TargetMachine.h

 
 namespace llvm {
 
 class AAManager;
 using ModulePassManager = PassManager<Module>;
 
 class Function;
 class GlobalValue;
+class LLVMTargetMachine;
 class MachineFunctionPassManager;
 class MachineFunctionAnalysisManager;
 class MachineModuleInfoWrapperPass;
 class Mangler;
 class MCAsmInfo;
 class MCContext;
 class MCInstrInfo;
 class MCRegisterInfo;
public:
   /// Parse out the target's MachineFunctionInfo from the YAML reprsentation.
   virtual bool parseMachineFunctionInfo(const yaml::MachineFunctionInfo &,
                                         PerFunctionMIParsingState &PFS,
                                         SMDiagnostic &Error,
                                         SMRange &SourceRange) const {
     return false;
   }
 
+  /// If this is an LLVMTargetMachine then return the downcast
+  /// pointer, otherwise return nullptr.
+  virtual LLVMTargetMachine *asLLVMTargetMachine() { return nullptr; }
+
+  /// Const version of asLLVMTargetMachine
+  const LLVMTargetMachine *asLLVMTargetMachine() const {
+    return const_cast<TargetMachine *>(this)->asLLVMTargetMachine();
+  }
+
   /// This method returns a pointer to the specified type of
   /// TargetSubtargetInfo.  In debug builds, it verifies that the object being
   /// returned is of the correct type.
   template <typename STC> const STC &getSubtarget(const Function &F) const {
     return *static_cast<const STC*>(getSubtargetImpl(F));
   }
 
   /// Create a DataLayout.
protected: // Can only create subclasses.
   LLVMTargetMachine(const Target &T, StringRef DataLayoutString,
                     const Triple &TT, StringRef CPU, StringRef FS,
                     const TargetOptions &Options, Reloc::Model RM,
                     CodeModel::Model CM, CodeGenOpt::Level OL);
 
   void initAsmInfo();
 
 public:
+  /// Provide safe downcast
+  LLVMTargetMachine *asLLVMTargetMachine() override { return this; }
+
   /// Get a TargetTransformInfo implementation for the target.
   ///
   /// The TTI returned uses the common code generator to answer queries about
   /// the IR.
   TargetTransformInfo getTargetTransformInfo(const Function &F) const override;
 
   /// Create a pass configuration object to be used by addPassToEmitX methods
   /// for generating a pipeline of CodeGen passes.

llvm/lib/CodeGen/TargetPassConfig.cpp

 #include "llvm/MC/MCTargetOptions.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CodeGen.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Discriminator.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/SaveAndRestore.h"
 #include "llvm/Support/Threading.h"
 #include "llvm/Support/VirtualFileSystem.h"
 #include "llvm/Support/WithColor.h"
 #include "llvm/Target/CGPassBuilderOption.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils.h"
 #include <cassert>
 #include <optional>
 #include <string>
+#include <tuple>
 
 using namespace llvm;
 
 static cl::opt<bool>
     EnableIPRA("enable-ipra", cl::init(false), cl::Hidden,
                cl::desc("Enable interprocedural register allocation "
                         "to reduce load/store at procedure calls."));
 static cl::opt<bool> DisablePostRASched("disable-post-ra", cl::Hidden,
static cl::opt<bool> DisableExpandReductions(
     "disable-expand-reductions", cl::init(false), cl::Hidden,
     cl::desc("Disable the expand reduction intrinsics pass from running"));
 
 /// Disable the select optimization pass.
 static cl::opt<bool> DisableSelectOptimize(
     "disable-select-optimize", cl::init(true), cl::Hidden,
     cl::desc("Disable the select-optimization pass from running"));
 
+using GlobalExtensionTy =
+    std::tuple<TargetPassConfig::MachinePassExtensionPointTy,
+               TargetPassConfig::ExtensionFn, TargetPassConfig::ExtensionID>;
+
+static ManagedStatic<SmallVector<GlobalExtensionTy, 2>> GlobalExtensions;
+
 /// Allow standard passes to be disabled by command line options. This supports
 /// simple binary flags that either suppress the pass or do nothing.
 /// i.e. -disable-mypass=false has no effect.
 /// These should be converted to boolOrDefault in order to use applyOverride.
 static IdentifyingPassPtr applyDisable(IdentifyingPassPtr PassID,
                                        bool Override) {
   if (Override)
     return IdentifyingPassPtr();

 
 bool TargetPassConfig::isPassSubstitutedOrOverridden(AnalysisID ID) const {
   IdentifyingPassPtr TargetID = getPassSubstitution(ID);
   IdentifyingPassPtr FinalPtr = overridePass(ID, TargetID);
   return !FinalPtr.isValid() || FinalPtr.isInstance() ||
       FinalPtr.getID() != ID;
 }
 
+auto TargetPassConfig::addExtension(MachinePassExtensionPointTy MPEP,
+                                    ExtensionFn Fn) -> ExtensionID {
+  static ExtensionID Counter = 0;
+  ExtensionID ID = ++Counter;
+  GlobalExtensions->push_back(std::make_tuple(MPEP, Fn, ID));
+  return ID;
+}
+
+void TargetPassConfig::removeExtension(ExtensionID ID) {
+  if (!GlobalExtensions.isConstructed())
+    return;
+
+  auto *I = llvm::find_if(*GlobalExtensions, [ID](const auto &Elem) {
+    return std::get<2>(Elem) == ID;
+  });
+
+  if (I != GlobalExtensions->end())
+    GlobalExtensions->erase(I);
+}
+
+void TargetPassConfig::applyAnyExtensions(MachinePassExtensionPointTy MPEP) {
+  if (GlobalExtensions.isConstructed() && !GlobalExtensions->empty()) {
+    for (const auto &Ext : *GlobalExtensions) {
+      if (std::get<0>(Ext) == MPEP)
+        std::get<1>(Ext)(*this);
+    }
+  }
+}
+
 /// Add a pass to the PassManager if that pass is supposed to be run.  If the
 /// Started/Stopped flags indicate either that the compilation should start at
 /// a later pass or that it should stop after an earlier pass, then do not add
 /// the pass.  Finally, compare the current pass against the StartAfter
 /// and StopAfter options and change the Started/Stopped flags accordingly.
 void TargetPassConfig::addPass(Pass *P) {
   assert(!Initialized && "PassConfig is immutable");
 

 /// target-specific fixups just before or after major stages. Additionally,
 /// targets have the flexibility to change pass order within a stage by
 /// overriding default implementation of add%Stage routines below. Each
 /// technique has maintainability tradeoffs because alternate pass orders are
 /// not well supported. addPre/Post works better if the target pass is easily
 /// tied to a common pass. But if it has subtle dependencies on multiple passes,
 /// the target should override the stage instead.
 ///
+/// External code can customize the passes added here for any target
+/// via the addGlobalExtension function.
+///
 /// TODO: We could use a single addPre/Post(ID) hook to allow pass injection
 /// before/after any target-independent pass. But it's currently overkill.
 void TargetPassConfig::addMachinePasses() {
   AddingMachinePasses = true;
 
+  applyAnyExtensions(MPEP_EarlyAsPossible);
+
   // Add passes that optimize machine instructions in SSA form.
   if (getOptLevel() != CodeGenOpt::None) {
     addMachineSSAOptimization();
   } else {
     // If the target requests it, assign local variables to stack slots relative
     // to one another and simplify frame index references where possible.
     addPass(&LocalStackSlotAllocationID);
   }
 
   if (TM->Options.EnableIPRA)
     addPass(createRegUsageInfoPropPass());
 
   // Run pre-ra passes.
   addPreRegAlloc();
+  applyAnyExtensions(MPEP_PreRegAlloc);
 
   // Debugifying the register allocator passes seems to provoke some
   // non-determinism that affects CodeGen and there doesn't seem to be a point
   // where it becomes safe again so stop debugifying here.
   DebugifyIsSafe = false;
 
   // Add a FSDiscriminator pass right before RA, so that we could get
   // more precise SampleFDO profile for RA.
void TargetPassConfig::addMachinePasses() {
   // Run register allocation and passes that are tightly coupled with it,
   // including phi elimination and scheduling.
   if (getOptimizeRegAlloc())
     addOptimizedRegAlloc();
   else
     addFastRegAlloc();
 
   // Run post-ra passes.
+  applyAnyExtensions(MPEP_PostRegAlloc);
   addPostRegAlloc();
 
   addPass(&RemoveRedundantDebugValuesID);
 
   addPass(&FixupStatepointCallerSavedID);
 
   // Insert prolog/epilog code.  Eliminate abstract frame index references...
   if (getOptLevel() != CodeGenOpt::None) {
void TargetPassConfig::addMachinePasses() {
   if (getOptLevel() != CodeGenOpt::None)
     addMachineLateOptimization();
 
   // Expand pseudo instructions before second scheduling pass.
   addPass(&ExpandPostRAPseudosID);
 
   // Run pre-sched2 passes.
   addPreSched2();
+  applyAnyExtensions(MPEP_PreSched2);
 
   if (EnableImplicitNullChecks)
     addPass(&ImplicitNullChecksID);
 
   // Second pass scheduler.
   // Let Target optionally insert this pass by itself at some other
   // point.
   if (getOptLevel() != CodeGenOpt::None &&
void TargetPassConfig::addMachinePasses() {
 
   // Insert before XRay Instrumentation.
   addPass(&FEntryInserterID);
 
   addPass(&XRayInstrumentationID);
   addPass(&PatchableFunctionID);
 
   addPreEmitPass();
+  applyAnyExtensions(MPEP_PreEmitPass);
 
   if (TM->Options.EnableIPRA)
     // Collect register usage information and produce a register mask of
     // clobbered registers, to be used to optimize call sites.
     addPass(createRegUsageInfoCollector());
 
   // FIXME: Some backends are incompatible with running the verifier after
   // addPreEmitPass.  Maybe only pass "false" here for those targets?
void TargetPassConfig::addMachinePasses() {
   if (!DisableCFIFixup && TM->Options.EnableCFIFixup)
     addPass(createCFIFixup());
 
   PM->add(createStackFrameLayoutAnalysisPass());
 
   // Add passes that directly emit MI after all other MI passes.
   addPreEmitPass2();
 
+  applyAnyExtensions(MPEP_LateAsPossible);
+
   AddingMachinePasses = false;
 }
 
 /// Add passes that optimize machine instructions in SSA form.
 void TargetPassConfig::addMachineSSAOptimization() {
   // Pre-ra tail duplication.
   addPass(&EarlyTailDuplicateID);
 

llvm/test/CMakeLists.txt

if(CMAKE_HOST_UNIX)
     list(APPEND LLVM_TEST_DEPENDS
       LLJITWithRemoteDebugging
       )
   endif()
   if (NOT WIN32)
     list(APPEND LLVM_TEST_DEPENDS
       Bye
       ExampleIRTransforms
+      MachineBye
       )
   endif()
 endif()
 
 if(TARGET ocaml_llvm)
   # Clear all non-OCaml cross-target dependencies when building out-of-tree.
   if(LLVM_OCAML_OUT_OF_TREE)
     set(LLVM_TEST_DEPENDS)

llvm/test/Feature/machine_bye.ll

+; RUN: llc < %s 2>&1 | FileCheck %s --check-prefixes BOTH
+; RUN: llc %loadmachinebye < %s 2>&1 | FileCheck %s --check-prefixes BOTH,LOAD
+; REQUIRES: examples
+; UNSUPPORTED: system-windows
+
+; LOAD: MachineBye:
+; BOTH-LABEL: foo
+define void @foo() {
+  ret void
+}

llvm/test/lit.cfg.py

else:
     config.available_features.add("can-execute")
 
 # Loadable module
 if config.has_plugins:
     config.available_features.add("plugins")
 
 if config.build_examples:
     config.available_features.add("examples")
+    config.substitutions.append(("%loadmachinebye",
+                                 "-load={}/MachineBye{}".format(config.llvm_shlib_dir,
+                                                                config.llvm_shlib_ext)))
 
 if config.linked_bye_extension:
     config.substitutions.append(("%llvmcheckext", "CHECK-EXT"))
     config.substitutions.append(("%loadbye", ""))
     config.substitutions.append(("%loadnewpmbye", ""))
 else:
     config.substitutions.append(("%llvmcheckext", "CHECK-NOEXT"))
     config.substitutions.append(

llvm/unittests/CodeGen/CMakeLists.txt

add_llvm_unittest(CodeGenTests
   MachineInstrTest.cpp
   MachineOperandTest.cpp
   RegAllocScoreTest.cpp
   PassManagerTest.cpp
   ScalableVectorMVTsTest.cpp
   SchedBoundary.cpp
   SelectionDAGAddressAnalysisTest.cpp
   TypeTraitsTest.cpp
+  TargetExtensionTest.cpp
   TargetOptionsTest.cpp
   TestAsmPrinter.cpp
   MLRegallocDevelopmentFeatures.cpp
   )
 
 add_subdirectory(GlobalISel)
 
 target_link_libraries(CodeGenTests PRIVATE LLVMTestingSupport)

llvm/unittests/CodeGen/TargetExtensionTest.cpp

+#include "llvm/CodeGen/MachineModuleInfo.h"
+#include "llvm/CodeGen/TargetPassConfig.h"
+#include "llvm/IR/LegacyPassManager.h"
+#include "llvm/IR/Module.h"
+#include "llvm/MC/TargetRegistry.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Target/TargetMachine.h"
+#include "llvm/Target/TargetOptions.h"
+#include "llvm/TargetParser/Host.h"
+#include "gtest/gtest.h"
+#include <memory>
+#include <vector>
+
+using namespace llvm;
+
+static void addMachinePasses(LLVMTargetMachine *LLVMTM) {
+  legacy::PassManager PM;
+  TargetPassConfig *TPC(LLVMTM->createPassConfig(PM));
+
+  auto MMIWMP = std::make_unique<MachineModuleInfoWrapperPass>(LLVMTM);
+  PM.add(TPC);
+  PM.add(MMIWMP.release());
+  TPC->addMachinePasses();
+}
+
+TEST(TargetExtensionTest, ExtensionOrdering) {
+  InitializeNativeTarget();
+
+  LLVMContext Context;
+
+  Module Mod("target-extensions-test", Context);
+  Mod.setTargetTriple(llvm::sys::getDefaultTargetTriple());
+  std::string Error;
+  const Target *Target(
+      TargetRegistry::lookupTarget(Mod.getTargetTriple(), Error));
+
+  ASSERT_NE(Target, nullptr);
+  std::unique_ptr<TargetMachine> TM(Target->createTargetMachine(
+      Mod.getTargetTriple(), "", "", TargetOptions{}, std::nullopt));
+  ASSERT_NE(TM.get(), nullptr);
+  LLVMTargetMachine *LTM(TM->asLLVMTargetMachine());
+
+  if (LTM) {
+    // Register extensions at all possible points
+    unsigned CallCount =
+        static_cast<unsigned>(TargetPassConfig::MPEP_EarlyAsPossible);
+    std::vector<RegisterTargetExtension> Registrars;
+
+    for (unsigned MPEP = TargetPassConfig::MPEP_EarlyAsPossible;
+         MPEP <= TargetPassConfig::MPEP_LateAsPossible; ++MPEP) {
+
+      // Register an extension at each extension point for this machine
+      // to allow checks for any enums we missed in the implementation
+      // and that they are called in the correct order
+      Registrars.emplace_back(
+          static_cast<TargetPassConfig::MachinePassExtensionPointTy>(MPEP),
+          [MPEP, &CallCount, &Mod](TargetPassConfig &TPC) {
+            if (TPC.getTM<TargetMachine>().getTargetTriple().str() ==
+                Mod.getTargetTriple()) {
+              // Check that we are called in expected order for the machine
+              EXPECT_EQ(CallCount, MPEP) << "Calls arrived out of order";
+              ++CallCount;
+            }
+          });
+    }
+
+    EXPECT_EQ(CallCount,
+              static_cast<unsigned>(TargetPassConfig::MPEP_EarlyAsPossible))
+        << "Initial callcount mismatch";
+
+    // Check that all registered extension points get called
+    addMachinePasses(LTM);
+
+    EXPECT_EQ(CallCount,
+              static_cast<unsigned>(TargetPassConfig::MPEP_LateAsPossible) + 1)
+        << "Check1 callcount mismatch";
+
+    // Now remove all extensions and check they are not called again
+    Registrars.clear();
+    addMachinePasses(LTM);
+
+    EXPECT_EQ(CallCount,
+              static_cast<unsigned>(TargetPassConfig::MPEP_LateAsPossible) + 1)
+        << "Final callcount mismatch";
+  }
+}