[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 16.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