[Mlir-commits] [mlir] [mlir][python] allow upstream dialect registration (PR #74252)
Maksim Levental
llvmlistbot at llvm.org
Wed Dec 20 12:43:23 PST 2023
makslevental wrote:
> there is a `@_ods_cext.register_dialect` that seems to be a registration decorator, it seems related to https://github.com/llvm/llvm-project/blob/main/mlir/lib/Bindings/Python/MainModule.cpp#L59 (?) but I am not sure what to do with it
This is related to "loading" dialects when it comes to the python bindings i.e., how `op.operation.opview` materializes one of classes in the generated bindings. That decorator itself doesn't do much except decouple the dialect name from generated module name (`arith` vs `_arith_ops_gen.py`).
> the current register the world and/or link the world does not seem appealing.
Linking is a rabbit hole that I've explored pretty extensively. The issue (as I see them) is you want to enable dialects (upstream and downstream) to implement bindings posthoc i.e., in addition to whatever builtin dialects are currently handled by the `_mlir` extension. The current theory/approach on how to accomplish this necessitates many/several separate C extensions to be built, which leads to this asymmetry where upstream dialects are registered automatically/wholesale and downstream require explicit `register_dialect(ctx, /*load*/true)` call. Note that the separate modules in upstream for `transform` and `pdl` and `quant` don't have anything to do with registration but are about exposing various types.
If the question is about piece-meal registration then this PR could be the vehicle but I don't think sprawling out into ~20 separate C extension modules (dumped into `_mlir_libs`) is the answer. It wouldn't be hard to turn off `register_all` and just have ~20 separate `register_arith_dialect`, `register_scf_dialect`, etc.
If the question is really about linking i.e. people want to build bindings that don't have any object code for `affine` but do for `scf` that's a much much harder task because right now all that code doesn't live in the C extension but actually in the AggregateCAPI omnibus shlib which is opt-in _from each dialect's perspective_ (that's what `ENABLE_AGGREGATION` on `add_mlir_public_c_api_library` does). So by the time you get to building the bindings that shlib is already fixed.
> context.get_dialect_descriptor is wrapping `getOrLoadDialect` (why the name change btw?) and requires pre-registration.
I guess your issue is you'd like for there to be a something like `context.register_and_load_and_get_dialect_descriptor(dialect_name)`? Or just `context.register_and_load(dialect_name)`? The problem is you can't do that unless you just somewhere have a map from `dialect_name` -> `MlirDialectHandle` for all the dialects you're interested in. That's certainly possible for upstream dialects but isn't for downstream (where you have to actually generate/sign up for `mlirGetDialectHandle__YOURDIALECT__`). And then doing it just for upstream would again further enshrine the asymmetry. If these were generated by default then we could actually put a global map somewhere in the AggregateCAPI and add to it.
And to be honest, we should just start generating these. We've talked about generating C API for a while now and so we should just do it and this is a reasonable enough testcase.
> context.allow_unregistered_dialects, I am sure many people are using this blanket solution
You can actually go further - if you turn on `allow_unregistered_dialects` you can just copy-paste generated bindings from somewhere and they'll generate legal "unregistered IR" that'll then (assuming you copy-pasted recently) verify in a context that does have the dialects registered - I've done this for triton when I want to generate triton IR using my bindings but I don't want to build/link triton (so I just pass textual IR to their package/module/C extension).
> I also see that only a tiny subset of dialects have a dedicated module (in mlir/lib/Bindings/Python/DialectXXX.cpp) making things inconsistent. Could we just add a module per dialect with a single `register_dialect(bool load = false)` function ? This would also provide a location for grounding all future dialect module things and all per-dialect pybind things.
You could but you're not going to save yourself any linked baggage as I mentioned above.
> I imagine this could be tablegen'd but a low-tech solution can already improve things significantly.
Then you'll end up in a world where some people want to write their own C extensions (in order to pick and choose which types/attributes etc to expose) and some people don't. Not a serious barrier except I cringe at the CMake shenaningans to line up all the object files.
https://github.com/llvm/llvm-project/pull/74252
More information about the Mlir-commits
mailing list