[mlir python] Dis-aggregating MLIRPythonExtension.Core

[mikeurbach]

Since the work last year to refactor how MLIR Python extensions are linked, we have had this comment about dis-aggregating MLIRPythonExtension.Core:

llvm-project/mlir/python/CMakeLists.txt

Line 279 in 6ccc273

MLIRCAPIRegistration # TODO: See about dis-aggregating

My understanding is this dependency on MLIRCAPIRegistration is pulling in all the MLIR dialects, translations, and conversions:

llvm-project/mlir/lib/CAPI/Registration/CMakeLists.txt

Lines 1 to 14 in 10affe7

	# Dialect registration.
	get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
	get_property(translation_libs GLOBAL PROPERTY MLIR_TRANSLATION_LIBS)
	get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
	add_mlir_upstream_c_api_library(MLIRCAPIRegistration
	Registration.cpp
	LINK_LIBS PUBLIC
	MLIRCAPIIR
	MLIRLLVMToLLVMIRTranslation
	${dialect_libs}
	${translation_libs}
	${conversion_libs}
	)

Some downstream projects may not want all of that, and may just want to depend on the "core" MLIR Python extension. This is definitely the case in the CIRCT project.

I wanted to open an issue before diving in to see if anyone has thoughts about how to go about this.

cc @stellaraccident

[llvmbot]

@llvm/issue-subscribers-mlir-python

[stellaraccident]

Yes! I started a patch a few weeks ago to break this apart and got it to the point that I could use it (but needed more work to land). Let me try to remember what computer i stashed that branch on :)

[stellaraccident]

(by the numbers it only helped iree a little bit, but iree is already using most of mlir. I expect folks who are more just using a couple of dialects would benefit immensely)

[mikeurbach]

Awesome, thanks @stellaraccident! I'm happy to help here if I can.

In CIRCT, I don't think our Python libraries are going to need anything more than the builtin dialect, so we should be able to show the benefit for projects on the other end of the spectrum.

[stellaraccident]

Here's the patch I started: stellaraccident@93daca1

This was enough to have IREE select only the parts it wanted so I could assess. I expect it would also let you see that in circt.

You'll see that a key thing that needs to be done is to fix the commented out IRCore.cpp bit for mlirRegisterAllDialects. I had some thoughts about this -- none of them written down.

Here is the idea:

• Break out a separate extension for _mlirRegisterAllDialects.so which exposed a register_all_dialects(ctx) function that did this.
• Define a global hook for adding a context customization callback (i.e. mlir._mlir_libs._mlir.globals.add_context_creation_hook(f)). The context creation hook would invoke all of these callbacks whenever a new context was created.
• Add something to _mlir_libs.__init__.py like:

# Register upstream dialects if available.
try:
  from . import _mlirRegisterAllDialects
  def _register_all_dialects_hook(context):
    _mlirRegisterAllDialects.register_all_dialects(context)
except ModuleNotFoundError:
  pass
else:
  _mlir.globals.add_context_creation_hook(_register_all_dialects_hook)

# Defer to a project-level initializer for downstreams to customize.
try:
  from . import _site_initializer
except ModuleNotFoundError:
  pass

The idea is that Python level glue would inject the context initialization logic. We would have logic to do this automatically if the installation was built with the _mlirRegisterAllDialects extension (like upstream is), or downstreams can provide their own _site_initializer.py that does whatever they want. The end result would be that downstreams could set themselves up so that any created contexts are configured as needed by the project, removing a fair amount of setup boiler-plate for folks who opt-in.

There's some subtlety here and I am either happy to review if someone wants to take a stab, or I can find some time to implement this myself.

[joker-eph]

_mlir.globals.add_context_creation_hook(_register_all_dialects_hook)

That would make any new context automatically register all dialects, why would we do this?
I'm concerned about modularity in general and this kind of global hooks goes fairly against it.

[mikeurbach]

Well, isn't that what happens here today?

llvm-project/mlir/lib/Bindings/Python/IRCore.cpp

Line 477 in 8c4a07c

mlirRegisterAllDialects(context);

I am also concerned about the modularity, but I was worried about backwards compatibility in the event projects have come to depend on this.

My understanding of Stella's sketch is projects that actually want all MLIR dialects loaded can use the new, separated _mlirRegisterAllDialects.so, but projects who don't could just not build/depend on that, and not use this global context hook at all if they don't want.

I think we've actually asked for a similar mechanism before for CIRCT, and met the same concerns, so I'm not super attached to that. I just want a safe path to not compiling most of MLIR for the Python extensions, and as a secondary concern, not registering most of MLIR into my Python contexts.

[joker-eph]

Well, isn't that what happens here today?

Indeed, and that a mistake IMO. This was done temporarily apparently (https://reviews.llvm.org/D88155 ) but as usual this is the kind of temporary that last (which is why I'm against this kind of shortcut in general).

and not use this global context hook at all if they don't want.

That means that this global context hook can't be added to the context creation method, I'm not sure how it is invoked then?

[stellaraccident]

I was also thinking of adding an option to the context constructor to use the global registration hook.

Leaving that aside, as a downstream: I want the ability to implicitly configure a context to my requirements and not leave that to user code. As an example: iree.compiler.ir.Context() should give me a context populated with dialects, extension hooks and whatever we dream up that is relevant to my project. In this case, we would not include the all dialects registration, but we would add a hook as part of init that set it up by default to our specifications.

But I think this should be overridable. Perhaps iree.compiler.ir.Context(default_initialize=False) gives you a plain context with no site specific initialization done. Channeling @ftynse, the python bindings should do the common thing by default, and specific projects should be able to define what that default thing is.

[stellaraccident]

(I am actually completely fine making the "site specific default" for upstream mlir require explicit registration, but I would still like the mechanism, and i would like to use it to let this disaggregation path keep behavior the same -- so that we can switch the user impacting behavior separately if desired)

[joker-eph]

Leaving that aside, as a downstream: I want the ability to implicitly configure a context to my requirements and not leave that to user code. As an example: iree.compiler.ir.Context() should give me a context populated with dialects, extension hooks and whatever we dream up that is relevant to my project

Sure, that makes sense, but I don't quite see how this justifies a MLIR-supplied global registration hook that would applies to all projects?
Why wouldn't iree.compiler.ir.Context() just register what it needs? Why is more registration needed?

[stellaraccident]

Because that is a native class that downstreams get verbatim from the upstream API: there is no mechanism for customizing what it does, unless if we build a python-ey hook in -- which is what I propose here.

To be clear, the only reason such a hook is coupled to this patch is because I feel strongly that I would like to fix the build layering separate from making a user impacting breaking change. I am also perfectly open to making this change but there is a fair bit more work to define dialect level registration hooks for upstream dialects to do it completely, and we need flexibility on the chicken/egg. I would have trouble doing it in one step, even if we wanted to. So in my mind, the customization hook is a two birds/one stone thing: downstreams want/benefit from it, and upstream can use it as an intermediate step to a more explicit API and build layering.

[stellaraccident]

Thinking about this some more, I think I can extend the mechanism get you where you want to go, Mehdi, but there is a lot of downstream debt in this area, and I'm not going to change the user behavior in one step. I'd be aiming for a more clear way to provide dialect registration entry points and a customization mechanism to enable good default behavior. I've got an idea on my mind but it probably needs a patch, not just a sketch to make it clear.

[ftynse]

In Python, I'd like to strike a good balance between avoiding the negative effects of global registration and having something usable out-of-the-box. Most Python packages are usable by pip install foo && python -e import foo without much configuration...

This may not be very pythonic, but can't we just inherit contexts? Have a mlir.ir.Context that doesn't register anything, mlir.ir.AllDialectsContext(mlir.ir.Context) that registers all dialects upstream, iree.compiler.ir.IREEContext that registers what it wants, etc.

[stellaraccident]

Yeah, we can do that if we want (will be a little gross since it will need to hack something into a module agree it is defined, but meh -- it would just be in an init file.

I think I do need to spend some time and see about exposing dialect registration python side Ina better way: right now, we have no build-layered compatible way to register individual upstream dialects. I'm trying to avoid having an extension per dialect, but maybe that becomes necessary :(

[stellaraccident]

I took a stab at a more complete implementation. PTAL: https://reviews.llvm.org/D128593

[stellaraccident]

I think I got the above patch to a final state, and IREE (which is what I tested with) no longer links in everything. PTAL.

[stellaraccident]

Note that this does implement the global hooks for registration and initialization that Mehdi didn't like, but it also separates the _BaseContext from the Context (which subclasses _BaseContext). The hooks only get applied to the Context. I work on a lot of downstreams, and I strongly believe that getting the default Context creation right for a distribution (which can now enforce whatever defaults it wants) is the right thing (note that because of the way namespacing works, things like IREE and CIRCT already have their own: iree.compiler.ir.Context and circt.ir.Context).

Separating the classes like this should allow other policies/usage to be done in the future if there is ever a case and preserves API layering.

I'm willing to debate names (i.e. _BaseContext vs Context, AllDialectsContext, etc), but I believe that what I have chosen satisfies the principle of least surprise for users, and it centralizes some of the hard-to-get-right site specific initialization.

[mikeurbach]

Thanks @stellaraccident! I will take a look and try out the patch tomorrow.

[ftynse]

Should we consider this fixed by 5e83a5b ?

[stellaraccident]

I think so. There is always more cleanup that can be done, but I think the issue is fixed.