[mlir][python] allow upstream dialect registration #74252
Conversation
|
@llvm/pr-subscribers-backend-amdgpu @llvm/pr-subscribers-mlir Author: Maksim Levental (makslevental) ChangesThis (currently just a sketch) PR allows users that do not build the Right now the API is lame - pass a string that's associated with the dialect (the dialect namespace) and the m.def(
"add_dialect_to_dialect_registry",
[](MlirDialectRegistry registry, const std::string &dialectNamespace) {
#define MLIR_DEFINE_CAPI_DIALECT_REGISTRATION_(NAMESPACE) \
if (dialectNamespace == #NAMESPACE) { \
mlirDialectHandleInsertDialect(mlirGetDialectHandle__##NAMESPACE##__(), \
registry); \
return; \
}Before I do more manicuring on the API I just want to make sure this is what everyone wants. Also I'm not sure how to test this upstream because it's not like I can turn off Full diff: https://github.com/llvm/llvm-project/pull/74252.diff 3 Files Affected:
diff --git a/mlir/include/mlir-c/RegisterEverything.h b/mlir/include/mlir-c/RegisterEverything.h
index ea2ea86449727..f894419ecb1e4 100644
--- a/mlir/include/mlir-c/RegisterEverything.h
+++ b/mlir/include/mlir-c/RegisterEverything.h
@@ -31,6 +31,32 @@ MLIR_CAPI_EXPORTED void mlirRegisterAllLLVMTranslations(MlirContext context);
/// Register all compiler passes of MLIR.
MLIR_CAPI_EXPORTED void mlirRegisterAllPasses(void);
+#define MLIR_DEFINE_CAPI_DIALECT_REGISTRATION_(NAMESPACE) \
+ MLIR_CAPI_EXPORTED MlirDialectHandle mlirGetDialectHandle__##NAMESPACE##__();
+
+#define FORALL_DIALECTS(_) \
+ _(acc) \
+ _(affine) \
+ _(amx) \
+ _(arm_neon) \
+ _(arm_sme) \
+ _(arm_sve) \
+ _(bufferization) \
+ _(complex) \
+ _(emitc) \
+ _(index) \
+ _(irdl) \
+ _(mesh) \
+ _(spirv) \
+ _(tosa) \
+ _(ub) \
+ _(x86vector)
+
+FORALL_DIALECTS(MLIR_DEFINE_CAPI_DIALECT_REGISTRATION_)
+
+#undef MLIR_DEFINE_CAPI_DIALECT_REGISTRATION_
+#undef FORALL_DIALECTS
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/lib/Bindings/Python/MainModule.cpp b/mlir/lib/Bindings/Python/MainModule.cpp
index 17272472ccca4..4605ccc0f2935 100644
--- a/mlir/lib/Bindings/Python/MainModule.cpp
+++ b/mlir/lib/Bindings/Python/MainModule.cpp
@@ -12,6 +12,31 @@
#include "IRModule.h"
#include "Pass.h"
+#include "mlir-c/Dialect/AMDGPU.h"
+#include "mlir-c/Dialect/Arith.h"
+#include "mlir-c/Dialect/Async.h"
+#include "mlir-c/Dialect/ControlFlow.h"
+#include "mlir-c/Dialect/Func.h"
+#include "mlir-c/Dialect/GPU.h"
+#include "mlir-c/Dialect/LLVM.h"
+#include "mlir-c/Dialect/Linalg.h"
+#include "mlir-c/Dialect/MLProgram.h"
+#include "mlir-c/Dialect/Math.h"
+#include "mlir-c/Dialect/MemRef.h"
+#include "mlir-c/Dialect/NVGPU.h"
+#include "mlir-c/Dialect/NVVM.h"
+#include "mlir-c/Dialect/OpenMP.h"
+#include "mlir-c/Dialect/PDL.h"
+#include "mlir-c/Dialect/Quant.h"
+#include "mlir-c/Dialect/ROCDL.h"
+#include "mlir-c/Dialect/SCF.h"
+#include "mlir-c/Dialect/Shape.h"
+#include "mlir-c/Dialect/SparseTensor.h"
+#include "mlir-c/Dialect/Tensor.h"
+#include "mlir-c/Dialect/Transform.h"
+#include "mlir-c/Dialect/Vector.h"
+#include "mlir-c/RegisterEverything.h"
+
namespace py = pybind11;
using namespace mlir;
using namespace py::literals;
@@ -65,6 +90,63 @@ PYBIND11_MODULE(_mlir, m) {
},
"dialect_class"_a,
"Class decorator for registering a custom Dialect wrapper");
+ m.def(
+ "add_dialect_to_dialect_registry",
+ [](MlirDialectRegistry registry, const std::string &dialectNamespace) {
+
+#define MLIR_DEFINE_CAPI_DIALECT_REGISTRATION_(NAMESPACE) \
+ if (dialectNamespace == #NAMESPACE) { \
+ mlirDialectHandleInsertDialect(mlirGetDialectHandle__##NAMESPACE##__(), \
+ registry); \
+ return; \
+ }
+
+#define FORALL_DIALECTS(_) \
+ _(acc) \
+ _(affine) \
+ _(amdgpu) \
+ _(amx) \
+ _(arith) \
+ _(arm_neon) \
+ _(arm_sme) \
+ _(arm_sve) \
+ _(async) \
+ _(bufferization) \
+ _(cf) \
+ _(complex) \
+ _(emitc) \
+ _(func) \
+ _(gpu) \
+ _(index) \
+ _(irdl) \
+ _(linalg) \
+ _(llvm) \
+ _(math) \
+ _(memref) \
+ _(mesh) \
+ _(ml_program) \
+ _(nvgpu) \
+ _(nvvm) \
+ _(omp) \
+ _(pdl) \
+ _(quant) \
+ _(rocdl) \
+ _(scf) \
+ _(shape) \
+ _(spirv) \
+ _(tensor) \
+ _(tosa) \
+ _(ub) \
+ _(vector) \
+ _(x86vector)
+ FORALL_DIALECTS(MLIR_DEFINE_CAPI_DIALECT_REGISTRATION_)
+
+#undef MLIR_DEFINE_CAPI_DIALECT_REGISTRATION_
+#undef FORALL_DIALECTS
+ throw std::runtime_error("unknown dialect namespace: " +
+ dialectNamespace);
+ },
+ "dialect_registry"_a, "dialect_namespace"_a);
m.def(
"register_operation",
[](const py::object &dialectClass, bool replace) -> py::cpp_function {
diff --git a/mlir/lib/CAPI/RegisterEverything/RegisterEverything.cpp b/mlir/lib/CAPI/RegisterEverything/RegisterEverything.cpp
index c1c4a418b2552..debebe58ad64f 100644
--- a/mlir/lib/CAPI/RegisterEverything/RegisterEverything.cpp
+++ b/mlir/lib/CAPI/RegisterEverything/RegisterEverything.cpp
@@ -9,13 +9,43 @@
#include "mlir-c/RegisterEverything.h"
#include "mlir/CAPI/IR.h"
+#include "mlir/CAPI/Registration.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/InitAllDialects.h"
#include "mlir/InitAllExtensions.h"
#include "mlir/InitAllPasses.h"
#include "mlir/Target/LLVMIR/Dialect/All.h"
-#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
-#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
+
+using namespace mlir;
+
+#define MLIR_DEFINE_CAPI_DIALECT_REGISTRATION_(NAMESPACE, NAME) \
+ MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(NAME, NAMESPACE, \
+ NAMESPACE::NAME##Dialect)
+
+#define FORALL_DIALECTS(_) \
+ _(acc, OpenACC) \
+ _(affine, Affine) \
+ _(amx, AMX) \
+ _(arith, Arith) \
+ _(arm_neon, ArmNeon) \
+ _(arm_sme, ArmSME) \
+ _(arm_sve, ArmSVE) \
+ _(bufferization, Bufferization) \
+ _(complex, Complex) \
+ _(emitc, EmitC) \
+ _(mlir, DLTI) \
+ _(index, Index) \
+ _(irdl, IRDL) \
+ _(mesh, Mesh) \
+ _(spirv, SPIRV) \
+ _(tosa, Tosa) \
+ _(ub, UB) \
+ _(x86vector, X86Vector)
+
+FORALL_DIALECTS(MLIR_DEFINE_CAPI_DIALECT_REGISTRATION_)
+
+#undef MLIR_DEFINE_CAPI_DIALECT_REGISTRATION_
+#undef FORALL_DIALECTS
void mlirRegisterAllDialects(MlirDialectRegistry registry) {
mlir::registerAllDialects(*unwrap(registry));
|
makslevental
force-pushed
the
upstream_dialect_registration
branch
from
b0ad6cc to
69d1958
Compare
makslevental
force-pushed
the
upstream_dialect_registration
branch
from
69d1958 to
f732927
Compare
| #include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h" | ||
| #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h" |
There was a problem hiding this comment.
Just dead code (or, alternatively, an IWYU failure somewhere else...)
makslevental
force-pushed
the
upstream_dialect_registration
branch
from
f732927 to
bc9e216
Compare
makslevental
force-pushed
the
upstream_dialect_registration
branch
2 times, most recently
from
f53d4ff to
bb78956
Compare
makslevental
force-pushed
the
upstream_dialect_registration
branch
from
bb78956 to
719695d
Compare
There was a problem hiding this comment.
I'm not sure about the intent here, seems to be the usual confusion with dialect registration and loading. Does this intend to register the dialects, i.e. make them available to the parser? Does this rather intend to load the dialects into the context, i.e. allow one to construct objects from these dialects? For the latter, I'd look at what downstreams already do with various load_<downstream>_dialects(context) functions and reuse that, modulo some homogenization, for upstream dialects.
As a vague design suggestion, I would find it convenient to have from mlir.dialects import foo to have foo loaded in whatever context I create afterwards and from mlir.dialects.foo import parser to have foo registered (but not loaded) in whatever context I create afterwards + have a mechanism to override this behavior should we need a clean context for some reason. But maybe it's too magic for other folks.
| # Dialects | ||
| MLIRCAPIAMDGPU | ||
| MLIRCAPIArith | ||
| MLIRCAPIAsync | ||
| MLIRCAPIControlFlow | ||
| MLIRCAPIFunc | ||
| MLIRCAPIGPU | ||
| MLIRCAPILLVM | ||
| MLIRCAPILinalg | ||
| MLIRCAPIMLProgram | ||
| MLIRCAPIMath | ||
| MLIRCAPIMemRef | ||
| MLIRCAPINVGPU | ||
| MLIRCAPINVVM | ||
| MLIRCAPIOpenMP | ||
| MLIRCAPIPDL | ||
| MLIRCAPIQuant | ||
| MLIRCAPIROCDL | ||
| MLIRCAPISCF | ||
| MLIRCAPIShape | ||
| MLIRCAPISparseTensor | ||
| MLIRCAPITensor | ||
| MLIRCAPITransformDialect | ||
| MLIRCAPIVector | ||
| MLIRCAPIRemainingDialects |
There was a problem hiding this comment.
I don't think we want "Core" bindings to depend on all dialects...
There was a problem hiding this comment.
The issue is that right now most of the dialects don't have their own pybind modules. I was trying to avoid introducing 20 new modules that just have register_dialect in them. I also couldn't put these in RegisterEverything because the mere presence of that module triggers the omnibus loading.
There was a problem hiding this comment.
Yeah. I need to activate some old caches for this, but I think it was this issue that kept me from doing this the first time around.
If we can get the dependency dag such that RegisterEverything carries the weight of "everything" but is strictly optional, I think that is what we are going for.
For upstream dialects, registration is sufficient because the auto-loading mechanisms are infallible. That's why I proposed this version rather than the one alluded to in the issue (in terms of loading) because this version can indeed be a side-effecting import, which as you say is very fine UX.
That's doable since now the default registry is an exposed global. |
|
@hawkinsp just pinging again to get your feedback. |
I think the issue is that we use this in various jit and multi-client flows where "after" has no defined meaning -- it is all concurrent. I'm certainly not against doing something more ergonomic for this class of things but would like to avoid non-deterministic behavior in concurrent situations. |
|
I've been studying this a bit and digging deeper I find the current situation unfortunate and inconsistent. Atm it seems we have a :
I also see that only a tiny subset of dialects have a dedicated module (in mlir/lib/Bindings/Python/DialectXXX.cpp) making things inconsistent. I imagine this could be tablegen'd but a low-tech solution can already improve things significantly. |
This is related to "loading" dialects when it comes to the python bindings i.e., how
Linking is a rabbit hole that I've explored pretty extensively. The issue (as I see it) 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 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 If the question is really about linking i.e. people want to build bindings that don't have any object code for
I guess your issue is you'd like for there to be a something like 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.
You can actually go further - if you turn on
You could but you're not going to save yourself any linked baggage as I mentioned above.
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. After writing all this out I'm basically in the same place I was when I opened this PR: what exactly do people want???. I think if someone (@nicolasvasilache or @hawkinsp) could articulate what their goal is it would go a long way to deciding how to proceed. Just "fixing asymmetry" is a nice goal but not great motivation to rearchitect :) |
|
This part of the system was, frankly, never finished. Every stakeholder wanted something else and the desire to not end up in a linking-the-world-required state in the upstream world made this hard. I think the solution may be relatively simple: Treat "RegisterEverything" as something that everyone will want one of -- but make it so that everyone can have their own. Right now "RegisterEverything" is hard coded to be literally all upstream dialects, but what if there was just an easy way to, given a list of dialects, generate the registration handles and statically populate the name->handle map. Then upstream uses that rule to make its "RegisterEverything". Downstreams generate their own C code for that using the same generator, etc. I'd much prefer one C implementation file that needed to be generated for a list of dialects vs the current state that it trends to where there are tons of little pieces that have to fit together. If we had that, it could replace the ad-hoc _site_initialize stuff for most users, and the way upstream does it would be the same as downstreams. Could ultimately be one CMake target to generate the Edit: Or we go even simpler. If you want your very own Basically, instead of trying to factor |
This (currently just a sketch) PR allows users that do not build the
MLIRPythonExtension.RegisterEverythingtarget to still register upstream dialects from Python (discussed in #74245). It uses theget_dialect_registryAPI introduced in #72488 and adds/exposesMlirDialectHandles for all upstream dialects (i.e., fills in those that weren't already generating).Right now the API is lame - pass a string that's associated with the dialect (the dialect namespace) and the
MlirDialectHandleis fetched/gotten/materialized internally (by calling the appropriate C API). Ignoring whether we should keep it like this, the basic idea is captured in them.def:m.def( "add_dialect_to_dialect_registry", [](MlirDialectRegistry registry, const std::string &dialectNamespace) { #define MLIR_DEFINE_CAPI_DIALECT_REGISTRATION_(NAMESPACE) \ if (dialectNamespace == #NAMESPACE) { \ mlirDialectHandleInsertDialect(mlirGetDialectHandle__##NAMESPACE##__(), \ registry); \ return; \ }where (currently!) the user would do
ir.add_dialect_to_dialect_registry(get_dialect_registry(), "arith")and achieve the desired effect i.e., the nextir.Context()will include the arith dialect in its dialect registry. At least I believe that's the desired effect (gleaned from #74245). So before I do more manicuring on the API I just want to make sure this is what everyone wants.Also I'm not sure how to test this upstream because it's not like I can turn offMLIRPythonExtension.RegisterEverythingfor the core bindings. There is theStandaloneexample that conveniently a TODO about moving to finer grained registration so maybe now is the time to check off that TODO?This is now tested by removing
MLIRPythonExtension.RegisterEverythingfrom the standalone example and callingadd_dialect_to_dialect_registry(get_dialect_registry(), "arith")explicitly (the smoketest has a module with anarithop in it).cc @hawkinsp