[RFC] -load-pass-plugin for mlir-opt

This proposal is about implementing functionality similar to LLVM opt’s -load/-load-pass-plugin, i.e. runtime loading and linking of rewrite passes. Such functionality enables downstream users to load and link against precompiled mlir-opt binaries, thus decreasing maintenance, particularly with respect to compile times. In addition, this functionality reduces the barrier to entry for developers that aren’t familiar with LLVM/MLIR in general. Finally, albeit of marginal value, it removes that disparity between MLIR and LLVM (which might come as a surprise to novices).

Metaphorically, the goal is to enable construction of an mlir-tutor analog to llvm-tutor.

Proof of Concept

I have a very hacky, minimal working example here. With the (egregious) hacking there I’m able to write this pass

struct BbqPass
    : public PassWrapper<BbqPass, OperationPass<ModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(BbqPass)

  StringRef getArgument() const final { return "bbq"; }
  StringRef getDescription() const final {
    return "my bbq pass the description";
  }
  void runOnOperation() override {
    auto f = getOperation();
    std::cerr << "-- dump first op in module\n";
    f.getBodyRegion().getOps().begin()->dump();
  }
};

compile it out-of-tree on Ubuntu 22 (against mlir-opt compiled prior) using this CMake

cmake_minimum_required(VERSION 3.13.4)
project(HelloWorldMLIR LANGUAGES CXX C)
set(CMAKE_C_STANDARD 11)
set(CMAKE_CXX_STANDARD 17)

set(LT_LLVM_INSTALL_DIR "" CACHE PATH "LLVM installation directory")

set(LLVM_DIR "${LT_LLVM_INSTALL_DIR}/lib/cmake/llvm/")
set(MLIR_DIR "${LT_LLVM_INSTALL_DIR}/lib/cmake/mlir/")

find_package(LLVM REQUIRED CONFIG)
find_package(MLIR REQUIRED CONFIG)

include_directories(${LLVM_INCLUDE_DIRS})
include_directories(${MLIR_INCLUDE_DIRS})

list(APPEND CMAKE_MODULE_PATH "${MLIR_CMAKE_DIR}")
list(APPEND CMAKE_MODULE_PATH "${LLVM_CMAKE_DIR}")

include(AddLLVM)
include(AddMLIR)

add_llvm_pass_plugin(HelloWorldMLIR HelloWorldMLIR.cpp)

and run it like this

mlir-opt -load-pass-plugin $(pwd)/HelloWorldMLIR.so --pass-pipeline=bbq linear.mlir

Easy things

mlir/include/mlir/Pass/PassPlugin.h and mlir/lib/Pass/PassPlugin.cpp are basically a simple matter of s/LLVM/MLIR/g (module case). The only questions here (I think) would be

  1. Whether to adopt the full void (*RegisterPassBuilderCallbacks)(PassBuilder &) (substituting PassManager for PassBuilder). Currently I do not pass the PassManager to the plugin because I couldn’t figure out how to do that and add the pass argument to the pass registry before cl::ParseCommandLineOptions is called (and without that ordering --pass-pipeline=bbq fails).
  2. Whether to try to unify with llvm/include/llvm/Passes/PassPlugin.h in order to avoid the blatant code duplication. One way this can be accomplished is actually by just obliging plugin devs to ignore the PassBuilder& (i.e. no updates have to be made to llvm/Passes/PassPlugin.h at all), but this route is probably a bad idea.

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp is also just a matter of adding the cl::list<std::string> passPlugins("load-pass-plugin").

Hard things

The hardest part of getting this working was solving the runtime symbol resolution on the plugin side. Partially this might be because I am not a LLVM CMake conventions connoisseur so I wasn’t able to immediately discern the correct LLVM CMake macros to use; although I got things to work by using add_llvm_pass_plugin(HelloWorldMLIR HelloWorldMLIR.cpp) in the plugin’s CMake and

   DEPENDS
   ${LIBS}
+  SUPPORT_PLUGINS
+  )
+ target_link_libraries(mlir-opt PUBLIC ${LIBS})
+ export_executable_symbols_for_plugins(mlir-opt)
llvm_update_compile_flags(mlir-opt)

in mlir/tools/mlir-opt/CMakeLists.txt, I am not sure if this is the best way. Consequently, I am not sure if the travesty of manually splicing in mangled symbol names into llvm/utils/extract_symbols.py could have been avoided altogether. I suspect that that’s not the case and that some changes do need to be made to extract_symbols.py. But as I’m also not a DWARF/ELF connoisseur, as of right now, I do not know where/why/how those changes should be made. For example, one particularly confusing thing for me was the erroring out due to the plugin being unable to locate _ZTVN4mlir4PassE (i.e. the vtable for mlir::Pass). Investigating by

$ readelf -a --wide mlir-opt | grep _ZTVN4mlir4PassE
599262: 0000000005bc56b8   112 OBJECT  LOCAL DEFAULT   24 _ZTVN4mlir4PassE

where the LOCAL binding seems to the issue (after patching extract_symbols.py the binding becomes GLOBAL). Possibly there is some connection to this comment in Pass.h.

I’m happy to go read up on the necessary things (DWARF/ELF) in order to do the right things if people think there’s overall value.

Unknown things (to me)

Everything around all the other things that one might want to be “pluggable”, such as dialects and conversions.

Whether this somehow breaks the one definition rule (or any other C++/LLVM rules for that matter).

Probaby lots of other things.

Happy to hear all comments/criticisms/concerns.

1 Like

Bumping this since I think some people missed it last time due to labor day weekend.

I think this is very interesting, I’m mostly concerned about the practical aspect of how to make it work reliably consider our platform support matrix, and the windows limit on the number of exported symbols. Maybe this could work reliably when build everything as individual shared libraries?
I just would be cautious about exposing a feature that has lot of “gotchas” and subtle failure mode that are hard to debug.

I have the same concerns as @mehdi_amini in general here. I’m fairly supportive if it can work in a very consistent manor across the platform matrix, otherwise I’m a bit wary of pushing it as a “builtin” supported thing (at least in the upstream tools/infra).

@mehdi_amini @River707

I just would be cautious about exposing a feature that has lot of “gotchas” and subtle failure mode that are hard to debug.

I’m sensitive to this concern as well and wouldn’t propose to sacrifice “opacity/transparency” for convenience.

Maybe this could work reliably when build everything as individual shared libraries?

I’m not sure what this is (which build path) but if this can’t be implemented uniformly across the various platforms then I think the feature could still have value as an optional build path - a “compile mlir-opt once and afterwards only load your pass plugins” type thing.

I’m willing to put in quite a bit of effort on this, because it would measurably improve own my workflow so I’ll bring it up during the next ODM.

You get it by adding -DBUILD_SHARED_LIBS=ON on your cmake invocation (doc for LLVM CMake: Building LLVM with CMake — LLVM 18.0.0git documentation ).

(The doc seems to indicate this does not works reliably on Windows though)

We discussed it a bit at the end of the last Open Meeting (starting around 45m47s).

I would consider that a first step would be to get it to reliably work on Mac/Linux, get this behind an opt-in (-DMLIR_ENABLE_OPT_PLUGINS=ON) and have CMake error in configurations that aren’t supported (Windows, possibly some shared library configs if needed, etc.).

1 Like

Okay after quite a long time I have some progress to report on this:

In the above repo I take a different approach than extending mlir-opt while preserving the spirit of the idea: supporting easily extending the upstream set of passes with a standalone pass with minimal boilerplate (for readability/n00bs like me) and (more importantly) absolutely minimal compilation. To that end, two approaches are exhibited;

  1. A python extension that doesn’t use mlir-tblgen (opting instead for subclassing PassWrapper directly).
  2. A dylib/shlib/library that only depends on headers at compile time.

Both approaches run the same small pass pipeline (func.func(convert-linalg-to-loops),dummy-pass) on a small linalg example module (dummy-pass decorates the root module with attributes {dummy.dummy}). Also, importantly, both approaches work on both Linux and MacOS.

The first isn’t very novel/interesting, since it’s basically a repackaging of how the various Test* libraries already work, but hopefully it’s useful for someone.

On the otherhand, the second took quite a bit of experimentation and does resemble the original approach. To emphasize: the second approach compiles a shared library libmlir_plugin.so that doesn’t statically link any upstream object code (thus only depends, at compile time, on headers). The way this is accomplished is by compiling libMLIR-C with fvisibility=default[1] so that all symbols inherited from various MLIR libraries[2] are available at runtime (and then effectively LD_PRELOAD=libMLIR-C.so/DYLD_INSERT_LIBRARIES=libMLIR-C.dylib) and compiling libmlir_plugin with -undefined dynamic_lookup/--unresolved-symbols=ignore-all. In fact, though exactly this works fine on MacOS (i.e., if you compile the libmlir_plugin as an executable and preload libMLIR-C), I couldn’t get ld to load symbols in the right order/with the right “searchability” on Linux, so I resort to runtime loading using a Python script with dlopen(libMLIR-C.so, RTLD_GLOBAL) [3].

Note the repo (as a MWE), vis-a-vis these visibility acrobatics, depends on my distros of LLVM/MLIR.

Summary: if you’re fine with driving OpPassManager (and have an appropriately compiled libMLIR-C at runtime), you can already today build functional (i.e., working/usable) “pass plugins” that don’t, at compile time, link against anything.

Next steps: it seems possible to actually do this more easily for mlir-opt but I think this approach is possibly more reasonable since it works today with only the addition of a few CMake flags (i.e., no code needs to be added to mlir-opt or anywhere else). Now maybe libMLIR-C isn’t actually the right vehicle for this functionality; certainly the global visibility/relocatability of name mangled C++ symbols in a library that advertises a C API is quite strange[4]. In which case a potential way forward is a new target that aggregates all of the useful/interesting MLIR libraries but doesn’t claim a C API.

Regarding compatibility with BUILD_SHARED_LIBS=ON: I could not make this work wth shared libs - something to do with visibility=default and fPIC and a mysterious unexpected PLT reloc type 0x00 error. But it’s possible I was doing something wrong (it was quite late into the hacking session…) so if someone knows better I’m happy to try again.

Question/comments/concerns/suggestions?

Thanks @mehdi_amini for the ping/reminder.


  1. Again, this all works in a straightforward manner on MacOS (by setting -DCMAKE_C_VISIBILITY_PRESET=default -DCMAKE_CXX_VISIBILITY_PRESET=default) but on Linux I had to resort to surgery inside AddMLIR.cmake. ↩︎

  2. All those under LINK_LIBS PUBLIC in the various CAPI libraries. ↩︎

  3. RTLD_GLOBAL… symbols are made available for the relocation processing of any other object”. ↩︎

  4. I only latched on to libMLIR-C as a target for this functionality because in other experiments I noticed that lib*MLIRAggregateCAPI had all of the necessary symbols. ↩︎

2 Likes

With the approach you are taking, you might be able to just add the absolute path to libMLIR-C.so to your plugin link and avoid the rest of the preload and undef symbol bits. If you do that and run ldd on your plugin, this should reveal that it is linking against that.

There is a fair bit of happy accident involved in this working: that shared library does not transitively include everything, so you may run into things that are undefined in it.

I am a bit confused by BUILD_SHARED_LIBS not working… If you build your plugin in that mode, it will always depend on the right stuff and it will only build your stuff. I routinely use this in development out of tree builds via find_package with just this effect…

In any case, congrats on getting something working. That’s often half the battle :slight_smile: