diff --git a/mlir/examples/standalone/python/CMakeLists.txt b/mlir/examples/standalone/python/CMakeLists.txt index cb10518e94e33..4dc4f7ae0cff6 100644 --- a/mlir/examples/standalone/python/CMakeLists.txt +++ b/mlir/examples/standalone/python/CMakeLists.txt @@ -28,6 +28,8 @@ declare_mlir_python_extension(StandalonePythonSources.Pybind11Extension StandaloneExtensionPybind11.cpp EMBED_CAPI_LINK_LIBS StandaloneCAPI + MLIRCAPITransforms + MLIRCAPIArith PYTHON_BINDINGS_LIBRARY pybind11 ) @@ -38,6 +40,8 @@ declare_mlir_python_extension(StandalonePythonSources.NanobindExtension StandaloneExtensionNanobind.cpp EMBED_CAPI_LINK_LIBS StandaloneCAPI + MLIRCAPITransforms + MLIRCAPIArith PYTHON_BINDINGS_LIBRARY nanobind GENERATE_TYPE_STUBS ) @@ -54,9 +58,6 @@ add_mlir_python_common_capi_library(StandalonePythonCAPI RELATIVE_INSTALL_ROOT "../../../.." DECLARED_SOURCES StandalonePythonSources - # TODO: Remove this in favor of showing fine grained registration once - # available. - MLIRPythonExtension.RegisterEverything MLIRPythonSources.Core MLIRPythonSources.Dialects.builtin ) @@ -70,9 +71,6 @@ add_mlir_python_modules(StandalonePythonModules INSTALL_PREFIX "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}" DECLARED_SOURCES StandalonePythonSources - # TODO: Remove this in favor of showing fine grained registration once - # available. - MLIRPythonExtension.RegisterEverything MLIRPythonSources.Core MLIRPythonSources.Dialects.builtin COMMON_CAPI_LINK_LIBS diff --git a/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp b/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp index 189ebac368bf5..63b49d03cb835 100644 --- a/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp +++ b/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp @@ -10,6 +10,7 @@ //===----------------------------------------------------------------------===// #include "Standalone-c/Dialects.h" +#include "mlir-c/Dialect/Arith.h" #include "mlir/Bindings/Python/Nanobind.h" #include "mlir/Bindings/Python/NanobindAdaptors.h" @@ -22,12 +23,16 @@ NB_MODULE(_standaloneDialectsNanobind, m) { auto standaloneM = m.def_submodule("standalone"); standaloneM.def( - "register_dialect", + "register_dialects", [](MlirContext context, bool load) { - MlirDialectHandle handle = mlirGetDialectHandle__standalone__(); - mlirDialectHandleRegisterDialect(handle, context); + MlirDialectHandle standaloneHandle = + mlirGetDialectHandle__standalone__(); + MlirDialectHandle arithHandle = mlirGetDialectHandle__arith__(); + mlirDialectHandleRegisterDialect(standaloneHandle, context); + mlirDialectHandleRegisterDialect(arithHandle, context); if (load) { - mlirDialectHandleLoadDialect(handle, context); + mlirDialectHandleLoadDialect(standaloneHandle, context); + mlirDialectHandleLoadDialect(arithHandle, 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 397db4c20e743..88493f4f18f00 100644 --- a/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp +++ b/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp @@ -10,6 +10,7 @@ //===----------------------------------------------------------------------===// #include "Standalone-c/Dialects.h" +#include "mlir-c/Dialect/Arith.h" #include "mlir/Bindings/Python/PybindAdaptors.h" using namespace mlir::python::adaptors; @@ -21,12 +22,16 @@ PYBIND11_MODULE(_standaloneDialectsPybind11, m) { auto standaloneM = m.def_submodule("standalone"); standaloneM.def( - "register_dialect", + "register_dialects", [](MlirContext context, bool load) { - MlirDialectHandle handle = mlirGetDialectHandle__standalone__(); - mlirDialectHandleRegisterDialect(handle, context); + MlirDialectHandle standaloneHandle = + mlirGetDialectHandle__standalone__(); + MlirDialectHandle arithHandle = mlirGetDialectHandle__arith__(); + mlirDialectHandleRegisterDialect(standaloneHandle, context); + mlirDialectHandleRegisterDialect(arithHandle, context); if (load) { - mlirDialectHandleLoadDialect(handle, context); + mlirDialectHandleLoadDialect(standaloneHandle, context); + mlirDialectHandleLoadDialect(arithHandle, context); } }, py::arg("context") = py::none(), py::arg("load") = true); diff --git a/mlir/examples/standalone/test/python/smoketest.py b/mlir/examples/standalone/test/python/smoketest.py index bd40c65d16164..6d41ecc11c540 100644 --- a/mlir/examples/standalone/test/python/smoketest.py +++ b/mlir/examples/standalone/test/python/smoketest.py @@ -14,7 +14,7 @@ with Context(): - standalone_d.register_dialect() + standalone_d.register_dialects() module = Module.parse( """ %0 = arith.constant 2 : i32