-
Notifications
You must be signed in to change notification settings - Fork 15.2k
[MLIR][Python] enable piecemeal dialect registration directly from Python #161118
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
[MLIR][Python] enable piecemeal dialect registration directly from Python #161118
Conversation
e19a7b4
to
16aca5e
Compare
99683dd
to
a3272cb
Compare
4cb401a
to
72f15f3
Compare
@llvm/pr-subscribers-mlir Author: Maksim Levental (makslevental) ChangesTL;DR: This PR enables registering dialects piecemeal (in the Python bindings) without using The PR is large but it's mostly boilerplate. The way the new registration mechanism works is by exposing a
The core of it looks like this: from .._mlir_libs import get_dialect_registry as _get_dialect_registry
from .._mlir_libs._capi import register_dialect as _register_dialect
_dialect_registry = _get_dialect_registry()
if "mlirGetDialectHandle__quant__" not in _dialect_registry.dialect_names:
_register_dialect("mlirGetDialectHandle__quant__", _dialect_registry) I.e., if someone chooses to forgo using The rest of the PR simply adds the boilerplate necessary for users to actually make these calls for all the dialects. Some notes:
Patch is 37.37 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/161118.diff 38 Files Affected:
diff --git a/mlir/examples/standalone/pyproject.toml b/mlir/examples/standalone/pyproject.toml
index 5a1e6e86513c3..90535030df817 100644
--- a/mlir/examples/standalone/pyproject.toml
+++ b/mlir/examples/standalone/pyproject.toml
@@ -37,6 +37,8 @@ cmake.source-dir = "."
# This is for installing/distributing the python bindings target and only the python bindings target.
build.targets = ["StandalonePythonModules"]
install.components = ["StandalonePythonModules"]
+# The default is true but make it explicit to highlight that this option exists (turn off for debug symbols).
+install.strip = true
[tool.scikit-build.cmake.define]
# Optional
@@ -51,6 +53,9 @@ LLVM_USE_LINKER = { env = "LLVM_USE_LINKER", default = "" }
CMAKE_VISIBILITY_INLINES_HIDDEN = "ON"
CMAKE_C_VISIBILITY_PRESET = "hidden"
CMAKE_CXX_VISIBILITY_PRESET = "hidden"
+# Disables generation of "version soname" (i.e. libFoo.so.<version>),
+# which causes pure duplication of various shlibs for Python wheels.
+CMAKE_PLATFORM_NO_VERSIONED_SONAME = "ON"
# Non-optional (alternatively you could use CMAKE_PREFIX_PATH here).
MLIR_DIR = { env = "MLIR_DIR", default = "" }
diff --git a/mlir/examples/standalone/python/CMakeLists.txt b/mlir/examples/standalone/python/CMakeLists.txt
index 905c944939756..6d86d0e4ce3a8 100644
--- a/mlir/examples/standalone/python/CMakeLists.txt
+++ b/mlir/examples/standalone/python/CMakeLists.txt
@@ -65,7 +65,8 @@ add_mlir_python_common_capi_library(StandalonePythonCAPI
DECLARED_SOURCES
StandalonePythonSources
MLIRPythonSources.Core
- MLIRPythonSources.Dialects.builtin
+ EMBED_LIBS
+ MLIRCAPIQuant
)
################################################################################
@@ -138,6 +139,8 @@ set(_declared_sources
StandalonePythonSources
MLIRPythonSources.Core
MLIRPythonSources.Dialects.builtin
+ MLIRPythonSources.Dialects.arith
+ MLIRPythonSources.Dialects.quant
)
# For an external projects build, the MLIRPythonExtension.Core.type_stub_gen
# target already exists and can just be added to DECLARED_SOURCES.
diff --git a/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp b/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp
index 0ec6cdfa7994b..e28744b32ccfe 100644
--- a/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp
+++ b/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp
@@ -32,7 +32,7 @@ NB_MODULE(_standaloneDialectsNanobind, m) {
mlirDialectHandleRegisterDialect(standaloneHandle, context);
if (load) {
mlirDialectHandleLoadDialect(arithHandle, context);
- mlirDialectHandleRegisterDialect(standaloneHandle, context);
+ mlirDialectHandleLoadDialect(standaloneHandle, context);
}
},
nb::arg("context").none() = nb::none(), nb::arg("load") = true,
diff --git a/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp b/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
index da8c2167dc36b..db7645bcf5331 100644
--- a/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
+++ b/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
@@ -31,7 +31,7 @@ PYBIND11_MODULE(_standaloneDialectsPybind11, m) {
mlirDialectHandleRegisterDialect(standaloneHandle, context);
if (load) {
mlirDialectHandleLoadDialect(arithHandle, context);
- mlirDialectHandleRegisterDialect(standaloneHandle, context);
+ mlirDialectHandleLoadDialect(standaloneHandle, context);
}
},
py::arg("context") = py::none(), py::arg("load") = true);
diff --git a/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_nanobind.py b/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_nanobind.py
index 6218720951c82..3eb011ff4e87e 100644
--- a/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_nanobind.py
+++ b/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_nanobind.py
@@ -4,3 +4,10 @@
from ._standalone_ops_gen import *
from .._mlir_libs._standaloneDialectsNanobind.standalone import *
+
+from .._mlir_libs import get_dialect_registry as _get_dialect_registry
+from .._mlir_libs._capi import register_dialect as _register_dialect
+
+_dialect_registry = _get_dialect_registry()
+if "mlirGetDialectHandle__quant__" not in _dialect_registry.dialect_names:
+ _register_dialect("mlirGetDialectHandle__quant__", _dialect_registry)
diff --git a/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py b/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py
index bfb98e404e13f..2a259b71b33be 100644
--- a/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py
+++ b/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py
@@ -4,3 +4,10 @@
from ._standalone_ops_gen import *
from .._mlir_libs._standaloneDialectsPybind11.standalone import *
+
+from .._mlir_libs import get_dialect_registry as _get_dialect_registry
+from .._mlir_libs._capi import register_dialect as _register_dialect
+
+_dialect_registry = _get_dialect_registry()
+if "mlirGetDialectHandle__quant__" not in _dialect_registry.dialect_names:
+ _register_dialect("mlirGetDialectHandle__quant__", _dialect_registry)
diff --git a/mlir/examples/standalone/test/python/smoketest.py b/mlir/examples/standalone/test/python/smoketest.py
index 26d84fd63e947..3628522360c74 100644
--- a/mlir/examples/standalone/test/python/smoketest.py
+++ b/mlir/examples/standalone/test/python/smoketest.py
@@ -11,15 +11,26 @@
else:
raise ValueError("Expected either pybind11 or nanobind as arguments")
+from mlir_standalone.dialects import arith, quant
-with Context():
+
+with Context(), Location.unknown():
standalone_d.register_dialects()
- module = Module.parse(
- """
- %0 = arith.constant 2 : i32
- %1 = standalone.foo %0 : i32
- """
+ f32 = F32Type.get()
+ i8 = IntegerType.get_signless(8)
+ i32 = IntegerType.get_signless(32)
+ uniform = quant.UniformQuantizedType.get(
+ quant.UniformQuantizedType.FLAG_SIGNED, i8, f32, 0.99872, 127, -8, 7
)
- # CHECK: %[[C:.*]] = arith.constant 2 : i32
- # CHECK: standalone.foo %[[C]] : i32
+
+ module = Module.create()
+ with InsertionPoint(module.body):
+ two_i32 = arith.constant(i32, 2)
+ standalone_d.foo(two_i32)
+ two_f32 = arith.constant(f32, 2.0)
+ quant.qcast(uniform, two_f32)
+ # CHECK: %[[TWOI32:.*]] = arith.constant 2 : i32
+ # CHECK: standalone.foo %[[TWOI32]] : i32
+ # CHECK: %[[TWOF32:.*]] = arith.constant 2.000000e+00 : f32
+ # CHECK: quant.qcast %[[TWOF32]] : f32 to !quant.uniform<i8<-8:7>:f32, 9.987200e-01:127>
print(str(module))
diff --git a/mlir/include/mlir-c/Bindings/Python/Interop.h b/mlir/include/mlir-c/Bindings/Python/Interop.h
index a33190c380d37..89559da689017 100644
--- a/mlir/include/mlir-c/Bindings/Python/Interop.h
+++ b/mlir/include/mlir-c/Bindings/Python/Interop.h
@@ -84,6 +84,8 @@
#define MLIR_PYTHON_CAPSULE_VALUE MAKE_MLIR_PYTHON_QUALNAME("ir.Value._CAPIPtr")
#define MLIR_PYTHON_CAPSULE_TYPEID \
MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID._CAPIPtr")
+#define MLIR_PYTHON_CAPSULE_DIALECT_HANDLE \
+ MAKE_MLIR_PYTHON_QUALNAME("ir.DialectHandle._CAPIPtr")
/** Attribute on MLIR Python objects that expose their C-API pointer.
* This will be a type-specific capsule created as per one of the helpers
@@ -457,6 +459,13 @@ static inline MlirValue mlirPythonCapsuleToValue(PyObject *capsule) {
return value;
}
+static inline MlirDialectHandle
+mlirPythonCapsuleToDialectHandle(PyObject *capsule) {
+ void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_DIALECT_HANDLE);
+ MlirDialectHandle handle = {ptr};
+ return handle;
+}
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir-c/Dialect/Affine.h b/mlir/include/mlir-c/Dialect/Affine.h
new file mode 100644
index 0000000000000..b2bf5aad44de9
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/Affine.h
@@ -0,0 +1,36 @@
+//===-- mlir-c/Dialect/Affine.h - C API for Affine dialect --------*- C -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM
+// Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header declares the C interface for registering and accessing the
+// Affine dialect. A dialect should be registered with a context to make it
+// available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_AFFINE_H
+#define MLIR_C_DIALECT_AFFINE_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Affine, affine);
+
+MLIR_CAPI_EXPORTED void
+mlirAffineRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_AFFINE_H
diff --git a/mlir/include/mlir-c/Dialect/Bufferization.h b/mlir/include/mlir-c/Dialect/Bufferization.h
new file mode 100644
index 0000000000000..41af7a294eb5c
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/Bufferization.h
@@ -0,0 +1,36 @@
+//===-- mlir-c/Dialect/Bufferization.h - C API for Bufferization dialect --===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM
+// Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header declares the C interface for registering and accessing the
+// Bufferization dialect. A dialect should be registered with a context to make
+// it available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_BUFFERIZATION_H
+#define MLIR_C_DIALECT_BUFFERIZATION_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Bufferization, bufferization);
+
+MLIR_CAPI_EXPORTED void mlirBufferizationRegisterTransformDialectExtension(
+ MlirDialectRegistry registry);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_BUFFERIZATION_H
diff --git a/mlir/include/mlir-c/Dialect/Complex.h b/mlir/include/mlir-c/Dialect/Complex.h
new file mode 100644
index 0000000000000..e51a67346a6ee
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/Complex.h
@@ -0,0 +1,33 @@
+//===-- mlir-c/Dialect/Complex.h - C API for Complex dialect ------*- C -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM
+// Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header declares the C interface for registering and accessing the
+// Complex dialect. A dialect should be registered with a context to make it
+// available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_COMPLEX_H
+#define MLIR_C_DIALECT_COMPLEX_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Complex, complex);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_COMPLEX_H
diff --git a/mlir/include/mlir-c/Dialect/GPU.h b/mlir/include/mlir-c/Dialect/GPU.h
index 321c1122c3370..44d6308c270b2 100644
--- a/mlir/include/mlir-c/Dialect/GPU.h
+++ b/mlir/include/mlir-c/Dialect/GPU.h
@@ -63,6 +63,9 @@ mlirGPUObjectAttrHasKernels(MlirAttribute mlirObjectAttr);
MLIR_CAPI_EXPORTED MlirAttribute
mlirGPUObjectAttrGetKernels(MlirAttribute mlirObjectAttr);
+MLIR_CAPI_EXPORTED void
+mlirGPURegisterTransformDialectExtension(MlirDialectRegistry registry);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir-c/Dialect/Linalg.h b/mlir/include/mlir-c/Dialect/Linalg.h
index 339e63d667c5e..a3fa3a93a70c7 100644
--- a/mlir/include/mlir-c/Dialect/Linalg.h
+++ b/mlir/include/mlir-c/Dialect/Linalg.h
@@ -55,6 +55,9 @@ mlirLinalgGetIndexingMapsAttribute(MlirOperation op);
MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Linalg, linalg);
+MLIR_CAPI_EXPORTED void
+mlirLinalgRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir-c/Dialect/MemRef.h b/mlir/include/mlir-c/Dialect/MemRef.h
index 087a4b3f85b1d..38184404d6853 100644
--- a/mlir/include/mlir-c/Dialect/MemRef.h
+++ b/mlir/include/mlir-c/Dialect/MemRef.h
@@ -26,6 +26,9 @@ extern "C" {
MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(MemRef, memref);
+MLIR_CAPI_EXPORTED void
+mlirMemRefRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir-c/Dialect/NVGPU.h b/mlir/include/mlir-c/Dialect/NVGPU.h
index e58015a4a3421..89783d3758f9a 100644
--- a/mlir/include/mlir-c/Dialect/NVGPU.h
+++ b/mlir/include/mlir-c/Dialect/NVGPU.h
@@ -29,6 +29,9 @@ MLIR_CAPI_EXPORTED MlirType mlirNVGPUTensorMapDescriptorTypeGet(
MlirContext ctx, MlirType tensorMemrefType, int swizzle, int l2promo,
int oobFill, int interleave);
+MLIR_CAPI_EXPORTED void
+mlirNVGPURegisterTransformDialectExtension(MlirDialectRegistry registry);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir-c/Dialect/PDL.h b/mlir/include/mlir-c/Dialect/PDL.h
index 6ad2e2da62d87..eb611964e278b 100644
--- a/mlir/include/mlir-c/Dialect/PDL.h
+++ b/mlir/include/mlir-c/Dialect/PDL.h
@@ -66,6 +66,9 @@ MLIR_CAPI_EXPORTED bool mlirTypeIsAPDLValueType(MlirType type);
MLIR_CAPI_EXPORTED MlirType mlirPDLValueTypeGet(MlirContext ctx);
+MLIR_CAPI_EXPORTED void
+mlirPDLRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir-c/Dialect/SMT.h b/mlir/include/mlir-c/Dialect/SMT.h
index 0ad64746f148b..e6d8e14efed6e 100644
--- a/mlir/include/mlir-c/Dialect/SMT.h
+++ b/mlir/include/mlir-c/Dialect/SMT.h
@@ -104,6 +104,9 @@ mlirSMTAttrGetBVCmpPredicate(MlirContext ctx, MlirStringRef str);
MLIR_CAPI_EXPORTED MlirAttribute mlirSMTAttrGetIntPredicate(MlirContext ctx,
MlirStringRef str);
+MLIR_CAPI_EXPORTED void
+mlirSMTRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir-c/Dialect/SparseTensor.h b/mlir/include/mlir-c/Dialect/SparseTensor.h
index c816c1b58690e..0701211375a48 100644
--- a/mlir/include/mlir-c/Dialect/SparseTensor.h
+++ b/mlir/include/mlir-c/Dialect/SparseTensor.h
@@ -107,6 +107,9 @@ mlirSparseTensorEncodingAttrBuildLvlType(
const enum MlirSparseTensorLevelPropertyNondefault *properties,
unsigned propSize, unsigned n, unsigned m);
+MLIR_CAPI_EXPORTED void
+mlirSparseTensorRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir-c/Dialect/Tensor.h b/mlir/include/mlir-c/Dialect/Tensor.h
index 74cbc5a6f0fe1..5891ade3bf3fd 100644
--- a/mlir/include/mlir-c/Dialect/Tensor.h
+++ b/mlir/include/mlir-c/Dialect/Tensor.h
@@ -18,6 +18,9 @@ extern "C" {
MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Tensor, tensor);
+MLIR_CAPI_EXPORTED void
+mlirTensorRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir-c/Dialect/Tosa.h b/mlir/include/mlir-c/Dialect/Tosa.h
new file mode 100644
index 0000000000000..ed55577996604
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/Tosa.h
@@ -0,0 +1,33 @@
+//===-- mlir-c/Dialect/Tosa.h - C API for Tosa dialect ----------*- C ---*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM
+// Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header declares the C interface for registering and accessing the
+// Tosa dialect. A dialect should be registered with a context to make it
+// available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_TOSA_H
+#define MLIR_C_DIALECT_TOSA_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Tosa, tosa);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_TOSA_H
diff --git a/mlir/include/mlir-c/Dialect/UB.h b/mlir/include/mlir-c/Dialect/UB.h
new file mode 100644
index 0000000000000..74159f0c705de
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/UB.h
@@ -0,0 +1,33 @@
+//===-- mlir-c/Dialect/UB.h - C API for UB dialect ----------------*- C -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM
+// Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header declares the C interface for registering and accessing the
+// UB dialect. A dialect should be registered with a context to make it
+// available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_UB_H
+#define MLIR_C_DIALECT_UB_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(UB, ub);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_UB_H
diff --git a/mlir/include/mlir-c/Dialect/Vector.h b/mlir/include/mlir-c/Dialect/Vector.h
index 6256c82d172ce..b7248b7985057 100644
--- a/mlir/include/mlir-c/Dialect/Vector.h
+++ b/mlir/include/mlir-c/Dialect/Vector.h
@@ -26,6 +26,9 @@ extern "C" {
MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Vector, vector);
+MLIR_CAPI_EXPORTED void
+mlirVectorRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir-c/IR.h b/mlir/include/mlir-c/IR.h
index 061d7620ba077..55cc86accb8a0 100644
--- a/mlir/include/mlir-c/IR.h
+++ b/mlir/include/mlir-c/IR.h
@@ -66,6 +66,7 @@ DEFINE_C_API_STRUCT(MlirLocation, const void);
DEFINE_C_API_STRUCT(MlirModule, const void);
DEFINE_C_API_STRUCT(MlirType, const void);
DEFINE_C_API_STRUCT(MlirValue, const void);
+DEFINE_C_API_STRUCT(MlirDialectHandle, const void);
#undef DEFINE_C_API_STRUCT
@@ -207,11 +208,6 @@ MLIR_CAPI_EXPORTED MlirStringRef mlirDialectGetNamespace(MlirDialect dialect);
// registration schemes.
//===----------------------------------------------------------------------===//
-struct MlirDialectHandle {
- const void *ptr;
-};
-typedef struct MlirDialectHandle MlirDialectHandle;
-
#define MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Name, Namespace) \
MLIR_CAPI_EXPORTED MlirDialectHandle mlirGetDialectHandle__##Namespace##__( \
void)
@@ -233,6 +229,11 @@ MLIR_CAPI_EXPORTED void mlirDialectHandleRegisterDialect(MlirDialectHandle,
MLIR_CAPI_EXPORTED MlirDialect mlirDialectHandleLoadDialect(MlirDialectHandle,
MlirContext);
+/// Checks if the dialec...
[truncated]
|
For context: here is a previous PR that attempted to do the same thing #74252. |
72f15f3
to
2dd5d63
Compare
969af55
to
4a54377
Compare
54f76f4
to
b703e53
Compare
b703e53
to
f8e5414
Compare
TL;DR: This PR enables registering dialects piecemeal (in the Python bindings) without using
_site_initialize_0
or anything like that.The PR is large but it's mostly boilerplate. The way the new registration mechanism works is by exposing a
ctypes
-based CAPI formlirGetDialectHandle__DIALECT__
and addingDialectRegistry.insert_dialect
. The three changes to focus on areThe core of it looks like this:
I.e., if someone chooses to forgo using
MLIRPythonExtension.RegisterEverything
, they can drop these registration calls somewhere (see the changes to the standalone example which now registers thequant
dialect using this path) and everything will work.The rest of the PR simply adds the boilerplate necessary for users to actually make these calls for all the dialects.
Some notes:
MLIRPythonExtension.RegisterEverything
you will not transitively link your chosen dialects. Thus you will need to addMLIRCAPI<DIALECT>
to either yourdeclare_mlir_python_extension
invocation or youradd_mlir_python_common_capi_library
invocation (see the changes to examples/standalone/python/CMakeLists.txt);MLIRPythonExtension.Core
(i.e.,MLIRPythonExtension.Core
still depends on buildings and linking libs for all the dialects). This is because in addition to registering dialects,MLIRPythonExtension.RegisterEverything
registers all passes, all dialect extensions, all interfaces, and all LLVMIR translations. This is tedious but doable (cf the other PR);ctypes
- the added test shows how to createPyCapsule
s fromctypes.c_void_p
.