forked from llvm/llvm-project
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[mlir python] Add nanobind support for standalone dialects.
This PR allows out-of-tree dialects to write Python dialect modules using nanobind instead of pybind11. It may make sense to migrate in-tree dialects and some of the ODS Python infrastructure to nanobind, but that is a topic for a future change. This PR makes the following changes: * adds nanobind to the CMake and Bazel build systems. We also add robin_map to the Bazel build, which is a dependency of nanobind. * adds a PYTHON_BINDING_LIBRARY option to various CMake functions, such as declare_mlir_python_extension, allowing users to select a Python binding library. * creates a fork of mlir/include/mlir/Bindings/Python/PybindAdaptors.h named NanobindAdaptors.h. This plays the same role, using nanobind instead of pybind11. * splits CollectDiagnosticsToStringScope out of PybindAdaptors.h and into a new header mlir/include/mlir/Bindings/Python/Diagnostics.h, since it is code that is no way related to pybind11 or for that matter, Python. * changed the standalone Python extension example to have both pybind11 and nanobind variants. * changed mlir/python/mlir/dialects/python_test.py to have both pybind11 and nanobind variants. Notes: * A slightly unfortunate thing that I needed to do in the CMake integration was to use FindPython in addition to FindPython3, since nanobind's CMake integration expects the Python_ names for variables. Perhaps there's a better way to do this.
- Loading branch information
Showing
25 changed files
with
1,184 additions
and
105 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
35 changes: 35 additions & 0 deletions
35
mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,35 @@ | ||
| //===- StandaloneExtension.cpp - Extension module -------------------------===// | ||
| // | ||
| // This is the nanobind version of the example module. There is also a pybind11 | ||
| // example in StandaloneExtensionPybind11.cpp. | ||
| // | ||
| // 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 | ||
| // | ||
| //===----------------------------------------------------------------------===// | ||
|
|
||
| #include <nanobind/nanobind.h> | ||
|
|
||
| #include "Standalone-c/Dialects.h" | ||
| #include "mlir/Bindings/Python/NanobindAdaptors.h" | ||
|
|
||
| namespace nb = nanobind; | ||
|
|
||
| NB_MODULE(_standaloneDialectsNanobind, m) { | ||
| //===--------------------------------------------------------------------===// | ||
| // standalone dialect | ||
| //===--------------------------------------------------------------------===// | ||
| auto standaloneM = m.def_submodule("standalone"); | ||
|
|
||
| standaloneM.def( | ||
| "register_dialect", | ||
| [](MlirContext context, bool load) { | ||
| MlirDialectHandle handle = mlirGetDialectHandle__standalone__(); | ||
| mlirDialectHandleRegisterDialect(handle, context); | ||
| if (load) { | ||
| mlirDialectHandleLoadDialect(handle, context); | ||
| } | ||
| }, | ||
| nb::arg("context").none() = nb::none(), nb::arg("load") = true); | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
6 changes: 6 additions & 0 deletions
6
mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,6 @@ | ||
| # 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 | ||
|
|
||
| from ._standalone_ops_gen import * | ||
| from .._mlir_libs._standaloneDialectsPybind11.standalone import * |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,59 @@ | ||
| //===- Diagnostics.h - Helpers for diagnostics in Python bindings ---------===// | ||
| // | ||
| // 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 | ||
| // | ||
| //===----------------------------------------------------------------------===// | ||
|
|
||
| #ifndef MLIR_BINDINGS_PYTHON_DIAGNOSTICS_H | ||
| #define MLIR_BINDINGS_PYTHON_DIAGNOSTICS_H | ||
|
|
||
| #include <cassert> | ||
| #include <string> | ||
|
|
||
| #include "mlir-c/Diagnostics.h" | ||
| #include "mlir-c/IR.h" | ||
| #include "llvm/ADT/StringRef.h" | ||
|
|
||
| namespace mlir { | ||
| namespace python { | ||
|
|
||
| /// RAII scope intercepting all diagnostics into a string. The message must be | ||
| /// checked before this goes out of scope. | ||
| class CollectDiagnosticsToStringScope { | ||
| public: | ||
| explicit CollectDiagnosticsToStringScope(MlirContext ctx) : context(ctx) { | ||
| handlerID = mlirContextAttachDiagnosticHandler(ctx, &handler, &errorMessage, | ||
| /*deleteUserData=*/nullptr); | ||
| } | ||
| ~CollectDiagnosticsToStringScope() { | ||
| assert(errorMessage.empty() && "unchecked error message"); | ||
| mlirContextDetachDiagnosticHandler(context, handlerID); | ||
| } | ||
|
|
||
| [[nodiscard]] std::string takeMessage() { return std::move(errorMessage); } | ||
|
|
||
| private: | ||
| static MlirLogicalResult handler(MlirDiagnostic diag, void *data) { | ||
| auto printer = +[](MlirStringRef message, void *data) { | ||
| *static_cast<std::string *>(data) += | ||
| llvm::StringRef(message.data, message.length); | ||
| }; | ||
| MlirLocation loc = mlirDiagnosticGetLocation(diag); | ||
| *static_cast<std::string *>(data) += "at "; | ||
| mlirLocationPrint(loc, printer, data); | ||
| *static_cast<std::string *>(data) += ": "; | ||
| mlirDiagnosticPrint(diag, printer, data); | ||
| return mlirLogicalResultSuccess(); | ||
| } | ||
|
|
||
| MlirContext context; | ||
| MlirDiagnosticHandlerID handlerID; | ||
| std::string errorMessage = ""; | ||
| }; | ||
|
|
||
| } // namespace python | ||
| } // namespace mlir | ||
|
|
||
| #endif // MLIR_BINDINGS_PYTHON_DIAGNOSTICS_H |
Oops, something went wrong.