[mlir:python] Add a mechanism to register upstream dialects from Python

[hawkinsp]

Currently there seem to be two ways to register the upstream dialects in a build of MLIR Python:
a) use mlirRegisterEverything, which works, but links in many dialects and transforms that might not be necessary.
b) define a _site_initialize module, which works, but requires that you write a custom pybind11 module.

Wouldn't it be nice if one could do neither (a) nor (b), and instead just explicitly register dialects like arith with a ir.Context from Python, without having to write a C++ Python extension to do so?

I'm imagining adding an API along the lines of:

arith.register_dialect(context)

to the generated code for a dialect, but I don't much mind exactly what the API is, so long as it exists. That appears to be the pattern often used by out-of-tree dialects.

@stellaraccident

[stellaraccident]

Yeah, I don't remember why we didn't finish that. I think we were missing some per dialect generated code.

[llvmbot]

@llvm/issue-subscribers-mlir-python

Author: Peter Hawkins (hawkinsp)

Currently there seem to be two ways to register the upstream dialects in a build of MLIR Python: a) use `mlirRegisterEverything`, which works, but links in many dialects and transforms that might not be necessary. b) define a `_site_initialize` module, which works, but requires that you write a custom pybind11 module.

Wouldn't it be nice if one could do neither (a) nor (b), and instead just explicitly register dialects like arith with a ir.Context from Python, without having write a C++ Python extension to do so?

I'm imagining adding an API along the lines of:

arith.register_dialect(context)

to the generated code for a dialect, but I don't much mind exactly what the API is, so long as it exists. That appears to be the pattern often used by out-of-tree dialects.

@stellaraccident

[makslevental]

Well you can't do automatically

arith.register_dialect(context)

because that's contingent on automatically having a context which I don't think is feasible/tenable? But with #72488 you don't actually need/want that; you can just have

arith.register_dialect(get_dialect_registry())

at module top level in arith.py or _arith_ops_gen.py (modulo a little refactoring of the C API).

I think it's a good idea! Only issue is some people might feel (not me personally) that side-effecting imports are a bad thing; see debate here #72338.

[stellaraccident]

I think Peter is just looking for any API that lets him register a dialect from "user code" (vs as part of project setup).

I haven't been tracking this part closely for years, since I've been mostly in the "project setup" camp. But the use case seems valid and I expect that some of the work that Maks has done maybe gets this close?

[jpienaar]

Yes, I think we just didn't need it yet given "project camp" (to use that phrasing 🙂) but should be easy to generate. I'm in between automatically doing it on import or not - mostly as folks keep being surprised about the inclusion based activation I did for type inference (but python folks may be different).

[stellaraccident]

(re imports: in many codebases, ordering of imports is incidental, but use of the context is not. I really don't want to be debugging why some contexts have the dialects registered and some don't based on global import order or scoping, which are a completely disconnected concept. I would be open to some form of "auto setup" but it needs to be optional. Context setup in a full compiler or standalone system is a deliberate business that goes beyond just registering all dialects. It was my original hope that the register everything library would enable things up for people who want it to be more adhoc)

In any case, sounds like we should have the API. Then we can debate about sugar and defaults.

[hawkinsp]

Indeed, I don't much mind about "automatic". I'm going to call ir.Context() myself, and once I have that context, I'm happy to register every dialect I need manually.

[makslevental]

I'm sketching it out right now - just a matter of where to neatly put the parts/pieces.

[stellaraccident]

Lol, I think that "just a matter" was why I withdrew into the hedge on this a couple of years ago :) Thanks for looking at it fresh.

[makslevental]

See #74252.