Python bindings for out-of-tree projects

I would like to define Python bindings for an out-of-tree MLIR client that uses custom dialects. It needs quite deep IR inspection capabilities – traversing the IR with a mix of standard+custom operations, inspecting and setting builtin+custom attributes.

Currently, it doesn’t look like the bindings are set up for such a use case. I started some refactoring to achieve the following state:

  • most C++ bindings code is available as a library so that out-of-tree projects can #include headers related to bindings and use PyMlir* classes in their bindings code;
  • the bindings entrypoint, i.e. the _mlir module definition, is a separate library that depends on the former and only defines the extension module;
  • same for the out-of-tree project, which has its own entrypoint library defining the _external_project module, and other extension libraries for its dialects.

My project has a custom build system so I would appreciate if somebody reflected this library layout in CMake. Currently it is still building a combined library.

An extra design point is making external dialects available in the context. Currently, my approach resorts to something like

with mlir.ir.Context():
  external_project.load_dialects()
  # do stuff

but this looks obnoxious. I am considering to provide a dialect registry inside the bindings that can be populated when the dialect package is imported and appended to all live contexts (Python keeps track of those). This could look like

import external_project.dialects.external_dialect
with mlir.ir.Context():
  # do stuff directly

or, if we want a more opt-in registration behavior, like

import external_project.dialects.unregistered_dialect
import external_project.dialects.external_dialect.registration
with mlir.ir.Context():
  # unregistered_dialect is not registered 
  # external_dialect IS registered when the package was imported

Thoughts or suggestions?

@mikeurbach @jdd @stellaraccident

Thanks for opening this discussion @ftynse. I recently experimented with something very similar in CIRCT, and I think we have the same end state in mind. A couple specific points:

I started prototyping this, and I can try to help out on that front. As a concrete use-case CIRCT is using CMake and will make use of the refactored library.

I would add builtin+custom types as well, if that wasn’t implied. I don’t think it will be too different, and PyConcreteType was the first refactoring I had in mind as a follow up to ⚙ D101063 [mlir] Move PyConcreteAttribute to header. NFC.

This is the same pattern we have been using: circt/rtl.py at 1ae59c60d52176d8ad65d386900fc94857104c08 · llvm/circt · GitHub

This came up recently in code review. Having a mechanism to register dialects when the package is imported sounds good to me. I don’t have strong opinions about exactly how it works, but your examples both seem like an improvement I would be happy to use.

One last point, which @jdd mentioned in D101063, but is probably worth discussing here in general. In the end state, would the refactored headers be moved out from lib/ and somewhere into include/? This seems more natural to me, but I think it’s worth discussing exactly where to put them.

I’m quite concerned about the “pollution” associated with implicitly registering/loading a dialect in every available context. This leads to surprising behavior when trying to compose libraries / components in the same process, which is why we eliminated it entirely from MLIR. It isn’t clear to me why this should be reintroduced here?

Note also that registration and loading are different: it isn’t clear in your last option that external_dialect is loaded in the context or nor?

It also does not really address the “global pollution” aspect as it seems that every possible context created in the process in the future (even in unrelated libraries) will register (/load?) this dialect.

Thanks! I suppose we can split this by having one patch that rearranges the files but keeps the overall cmake structure and another patch that just changes the cmake structure.

Certainly. Attributes and types use the same mechanism.

I support the idea of putting them in include/mlir/Bindings/Python, which we already use for .td files anyway. The headers were originally “private implementation” so they were put in lib/ for the same reason as PassDetail.h and the likes.

The rationale I heard before was that loading all dialects is too expensive, not that it creates composability problems. Do you mind giving some examples of the problems you’ve seen?

I’m concerned about boilerplate and uncommon patterns in a productivity-oriented language. import something is the common way for Python to get something available without having to manually deal with package initialization.

Whatever lets me do external_dialect.Attribute.get() without any extra action. Well, I can live with with mlir.ir.Context(loadAll=True) if necessary.

Only those contexts visible to Python. Again, this sounds like an extension of the common behavior in Python, importing something at the top level makes it globally available.

No: it is first a composability problem, the performance is a side-effect that was added there. But the performance aspect is a symptom of the lack of composability.

Performance remains: loading more dialects will affect performance. But also behavior can change, for example loading a dialect adds new canonicalizer rules. There is also an aspect of potential name collisions.

I’m concerned about boilerplate and uncommon patterns in a productivity-oriented language. import something is the common way for Python to get something available without having to manually deal with package initialization.

I don’t see this as “package initialization”. This would imply global state and I’m against global registration/initialization: this is about instance initialization instead: every single new context has to be initialized/setup correctly / explicitly.

Users can “hide” this in their own coherent framework (like npcomp, etc.) but I’d really like MLIR (including MLIR python) to avoid locking ourselves into this.

Joining a couple of parallel threads:

Cross posting a comment from the former:

Generic concerns from me:

  • This does not scale beyond a single out of tree project controlling the entire Python environment by providing the mlir native libraries and its project libraries (i.e. circt). If another MLIR project wanted to co-exist, it would need to be built against the same installed MLIR libraries.
  • The compatibility rules for pybind11 to interop at the C++ level between multiple shared-library extensions are arcane and hard to satisfy for things not intimately built together (i.e. basically the tuple of (compiler, standard library, compile options, etc) all have to match and are checked specifically. Penalty for violation is that the types at the Python level are actually forked (the binding is actually stored in a dict in the Python builtin module with a key hard-coded to the above parameters). This usually causes chaos to ensue and is not completely obvious.
  • We still have the lingering TypeID linkage issue (I’ve been hunting it down for 6 months and haven’t run it to ground completely yet). This will work one day and not the next, based on past experience. This bug unfortunately can happen when you obey all of the other rules. I have to assume we’ll track it down eventually…
  • Other projects who have gone down this path have eventually had to grow a platform specific loader stub which can adapt things at runtime and work around issues that come from distributing binaries.
  • Those other projects have typically also had to embrace a deployment strategy that lets them pip install headers/dev libraries and rework end projects to build against that (which provides a bit of a source of truth to make sure everything is at the same revision/config options/etc). LLVM’s source instability will hurt here since we have very limited and unpredictable compatibility windows between projects.

An answer that neatly side-steps this for projects like circt and npcomp, which merely want to use MLIR for their own ends, and have limited benefit from joining on a single shared installation, is to move the MLIR Python binding in a direction so that it can be privately embedded in such a project (i.e. as an circt.mlir private module vs a shared dependency on mlir). I’ve tried to keep the source tree in a rough shape to make such an evolution practical at some point, but it would be definite work.

Forks of the conversation:

This thread is dealing with a) decoupling dialect registration for out of tree projects, and b) externalizing some of the helpers used within the _mlir.so module for use outside (PyConcrete{Type|Attribute}). I don’t have recent state on the former. On the latter, I would like to do an experiment where I rewrite the useful helpers in a way that doesn’t leak internals – and then use that.

FYI - I finally got annoyed enough to write out a better system for custom attributes and types, and this one is compatible with out-of-tree from the get-go:

See ESIModule.cpp for how it is used. Strict improvement, I think.