[Mlir-commits] [llvm] [mlir] [mlir python] Add nanobind support for standalone dialects. (PR #117922)
Peter Hawkins
llvmlistbot at llvm.org
Mon Dec 2 11:47:41 PST 2024
https://github.com/hawkinsp updated https://github.com/llvm/llvm-project/pull/117922
>From 2f0a8a169e081839858ba8fe75c13d3008ab7ac0 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins at google.com>
Date: Wed, 27 Nov 2024 20:17:40 +0000
Subject: [PATCH] [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.
---
mlir/cmake/modules/AddMLIRPython.cmake | 27 +-
mlir/cmake/modules/MLIRDetectPythonEnv.cmake | 39 +
mlir/docs/Bindings/Python.md | 20 +-
.../examples/standalone/python/CMakeLists.txt | 22 +-
.../python/StandaloneExtensionNanobind.cpp | 35 +
...on.cpp => StandaloneExtensionPybind11.cpp} | 7 +-
.../{standalone.py => standalone_nanobind.py} | 2 +-
.../dialects/standalone_pybind11.py | 6 +
.../standalone/test/python/smoketest.py | 14 +-
.../mlir/Bindings/Python/Diagnostics.h | 59 ++
.../mlir/Bindings/Python/NanobindAdaptors.h | 671 ++++++++++++++++++
.../mlir/Bindings/Python/PybindAdaptors.h | 43 +-
mlir/lib/Bindings/Python/DialectLLVM.cpp | 4 +-
.../Bindings/Python/TransformInterpreter.cpp | 7 +-
mlir/python/CMakeLists.txt | 23 +-
mlir/python/mlir/dialects/python_test.py | 17 +-
mlir/python/requirements.txt | 1 +
mlir/test/python/dialects/python_test.py | 59 +-
mlir/test/python/lib/CMakeLists.txt | 3 +-
.../python/lib/PythonTestModuleNanobind.cpp | 121 ++++
...odule.cpp => PythonTestModulePybind11.cpp} | 6 +-
utils/bazel/WORKSPACE | 18 +
.../llvm-project-overlay/mlir/BUILD.bazel | 50 +-
.../mlir/python/BUILD.bazel | 3 +-
utils/bazel/third_party_build/nanobind.BUILD | 25 +
utils/bazel/third_party_build/robin_map.BUILD | 12 +
26 files changed, 1187 insertions(+), 107 deletions(-)
create mode 100644 mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp
rename mlir/examples/standalone/python/{StandaloneExtension.cpp => StandaloneExtensionPybind11.cpp} (81%)
rename mlir/examples/standalone/python/mlir_standalone/dialects/{standalone.py => standalone_nanobind.py} (78%)
create mode 100644 mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py
create mode 100644 mlir/include/mlir/Bindings/Python/Diagnostics.h
create mode 100644 mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
create mode 100644 mlir/test/python/lib/PythonTestModuleNanobind.cpp
rename mlir/test/python/lib/{PythonTestModule.cpp => PythonTestModulePybind11.cpp} (94%)
create mode 100644 utils/bazel/third_party_build/nanobind.BUILD
create mode 100644 utils/bazel/third_party_build/robin_map.BUILD
diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index 7b91f43e2d57fd..67619a90c90be9 100644
--- a/mlir/cmake/modules/AddMLIRPython.cmake
+++ b/mlir/cmake/modules/AddMLIRPython.cmake
@@ -114,10 +114,11 @@ endfunction()
# EMBED_CAPI_LINK_LIBS: Dependent CAPI libraries that this extension depends
# on. These will be collected for all extensions and put into an
# aggregate dylib that is linked against.
+# PYTHON_BINDINGS_LIBRARY: Either pybind11 or nanobind.
function(declare_mlir_python_extension name)
cmake_parse_arguments(ARG
""
- "ROOT_DIR;MODULE_NAME;ADD_TO_PARENT"
+ "ROOT_DIR;MODULE_NAME;ADD_TO_PARENT;PYTHON_BINDINGS_LIBRARY"
"SOURCES;PRIVATE_LINK_LIBS;EMBED_CAPI_LINK_LIBS"
${ARGN})
@@ -126,15 +127,20 @@ function(declare_mlir_python_extension name)
endif()
set(_install_destination "src/python/${name}")
+ if(NOT ARG_PYTHON_BINDINGS_LIBRARY)
+ set(ARG_PYTHON_BINDINGS_LIBRARY "pybind11")
+ endif()
+
add_library(${name} INTERFACE)
set_target_properties(${name} PROPERTIES
# Yes: Leading-lowercase property names are load bearing and the recommended
# way to do this: https://gitlab.kitware.com/cmake/cmake/-/issues/19261
- EXPORT_PROPERTIES "mlir_python_SOURCES_TYPE;mlir_python_EXTENSION_MODULE_NAME;mlir_python_EMBED_CAPI_LINK_LIBS;mlir_python_DEPENDS"
+ EXPORT_PROPERTIES "mlir_python_SOURCES_TYPE;mlir_python_EXTENSION_MODULE_NAME;mlir_python_EMBED_CAPI_LINK_LIBS;mlir_python_DEPENDS;mlir_python_BINDINGS_LIBRARY"
mlir_python_SOURCES_TYPE extension
mlir_python_EXTENSION_MODULE_NAME "${ARG_MODULE_NAME}"
mlir_python_EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
mlir_python_DEPENDS ""
+ mlir_python_BINDINGS_LIBRARY "${ARG_PYTHON_BINDINGS_LIBRARY}"
)
# Set the interface source and link_libs properties of the target
@@ -223,12 +229,14 @@ function(add_mlir_python_modules name)
elseif(_source_type STREQUAL "extension")
# Native CPP extension.
get_target_property(_module_name ${sources_target} mlir_python_EXTENSION_MODULE_NAME)
+ get_target_property(_bindings_library ${sources_target} mlir_python_BINDINGS_LIBRARY)
# Transform relative source to based on root dir.
set(_extension_target "${modules_target}.extension.${_module_name}.dso")
add_mlir_python_extension(${_extension_target} "${_module_name}"
INSTALL_COMPONENT ${modules_target}
INSTALL_DIR "${ARG_INSTALL_PREFIX}/_mlir_libs"
OUTPUT_DIRECTORY "${ARG_ROOT_PREFIX}/_mlir_libs"
+ PYTHON_BINDINGS_LIBRARY ${_bindings_library}
LINK_LIBS PRIVATE
${sources_target}
${ARG_COMMON_CAPI_LINK_LIBS}
@@ -634,7 +642,7 @@ endfunction()
function(add_mlir_python_extension libname extname)
cmake_parse_arguments(ARG
""
- "INSTALL_COMPONENT;INSTALL_DIR;OUTPUT_DIRECTORY"
+ "INSTALL_COMPONENT;INSTALL_DIR;OUTPUT_DIRECTORY;PYTHON_BINDINGS_LIBRARY"
"SOURCES;LINK_LIBS"
${ARGN})
if(ARG_UNPARSED_ARGUMENTS)
@@ -644,9 +652,16 @@ function(add_mlir_python_extension libname extname)
# The actual extension library produces a shared-object or DLL and has
# sources that must be compiled in accordance with pybind11 needs (RTTI and
# exceptions).
- pybind11_add_module(${libname}
- ${ARG_SOURCES}
- )
+ if(NOT DEFINED ARG_PYTHON_BINDINGS_LIBRARY OR ARG_PYTHON_BINDINGS_LIBRARY STREQUAL "pybind11")
+ pybind11_add_module(${libname}
+ ${ARG_SOURCES}
+ )
+ elseif(ARG_PYTHON_BINDINGS_LIBRARY STREQUAL "nanobind")
+ nanobind_add_module(${libname}
+ NB_DOMAIN mlir
+ ${ARG_SOURCES}
+ )
+ endif()
# The extension itself must be compiled with RTTI and exceptions enabled.
# Also, some warning classes triggered by pybind11 are disabled.
diff --git a/mlir/cmake/modules/MLIRDetectPythonEnv.cmake b/mlir/cmake/modules/MLIRDetectPythonEnv.cmake
index 05397b7a1e1c75..c62ac7fa615ea6 100644
--- a/mlir/cmake/modules/MLIRDetectPythonEnv.cmake
+++ b/mlir/cmake/modules/MLIRDetectPythonEnv.cmake
@@ -21,6 +21,12 @@ macro(mlir_configure_python_dev_packages)
find_package(Python3 ${LLVM_MINIMUM_PYTHON_VERSION}
COMPONENTS Interpreter ${_python_development_component} REQUIRED)
+
+ # It's a little silly to detect Python a second time, but nanobind's cmake
+ # code looks for Python_ not Python3_.
+ find_package(Python ${LLVM_MINIMUM_PYTHON_VERSION}
+ COMPONENTS Interpreter ${_python_development_component} REQUIRED)
+
unset(_python_development_component)
message(STATUS "Found python include dirs: ${Python3_INCLUDE_DIRS}")
message(STATUS "Found python libraries: ${Python3_LIBRARIES}")
@@ -31,6 +37,13 @@ macro(mlir_configure_python_dev_packages)
message(STATUS "Python prefix = '${PYTHON_MODULE_PREFIX}', "
"suffix = '${PYTHON_MODULE_SUFFIX}', "
"extension = '${PYTHON_MODULE_EXTENSION}")
+
+ mlir_detect_nanobind_install()
+ find_package(nanobind 2.2 CONFIG REQUIRED)
+ message(STATUS "Found nanobind v${nanobind_VERSION}: ${nanobind_INCLUDE_DIR}")
+ message(STATUS "Python prefix = '${PYTHON_MODULE_PREFIX}', "
+ "suffix = '${PYTHON_MODULE_SUFFIX}', "
+ "extension = '${PYTHON_MODULE_EXTENSION}")
endif()
endmacro()
@@ -58,3 +71,29 @@ function(mlir_detect_pybind11_install)
set(pybind11_DIR "${PACKAGE_DIR}" PARENT_SCOPE)
endif()
endfunction()
+
+
+# Detects a nanobind package installed in the current python environment
+# and sets variables to allow it to be found. This allows nanobind to be
+# installed via pip, which typically yields a much more recent version than
+# the OS install, which will be available otherwise.
+function(mlir_detect_nanobind_install)
+ if(nanobind_DIR)
+ message(STATUS "Using explicit nanobind cmake directory: ${nanobind_DIR} (-Dnanobind_DIR to change)")
+ else()
+ message(STATUS "Checking for nanobind in python path...")
+ execute_process(
+ COMMAND "${Python3_EXECUTABLE}"
+ -c "import nanobind;print(nanobind.cmake_dir(), end='')"
+ WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+ RESULT_VARIABLE STATUS
+ OUTPUT_VARIABLE PACKAGE_DIR
+ ERROR_QUIET)
+ if(NOT STATUS EQUAL "0")
+ message(STATUS "not found (install via 'pip install nanobind' or set nanobind_DIR)")
+ return()
+ endif()
+ message(STATUS "found (${PACKAGE_DIR})")
+ set(nanobind_DIR "${PACKAGE_DIR}" PARENT_SCOPE)
+ endif()
+endfunction()
diff --git a/mlir/docs/Bindings/Python.md b/mlir/docs/Bindings/Python.md
index 6e52c4deaad9aa..a0bd1cac118bad 100644
--- a/mlir/docs/Bindings/Python.md
+++ b/mlir/docs/Bindings/Python.md
@@ -1138,12 +1138,14 @@ attributes and types must connect to the relevant C APIs for building and
inspection, which must be provided first. Bindings for `Attribute` and `Type`
subclasses can be defined using
[`include/mlir/Bindings/Python/PybindAdaptors.h`](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/PybindAdaptors.h)
-utilities that mimic pybind11 API for defining functions and properties. These
-bindings are to be included in a separate pybind11 module. The utilities also
-provide automatic casting between C API handles `MlirAttribute` and `MlirType`
-and their Python counterparts so that the C API handles can be used directly in
-binding implementations. The methods and properties provided by the bindings
-should follow the principles discussed above.
+or
+[`include/mlir/Bindings/Python/NanobindAdaptors.h`](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h)
+utilities that mimic pybind11/nanobind API for defining functions and
+properties. These bindings are to be included in a separate module. The
+utilities also provide automatic casting between C API handles `MlirAttribute`
+and `MlirType` and their Python counterparts so that the C API handles can be
+used directly in binding implementations. The methods and properties provided by
+the bindings should follow the principles discussed above.
The attribute and type bindings for a dialect can be located in
`lib/Bindings/Python/Dialect<Name>.cpp` and should be compiled into a separate
@@ -1179,7 +1181,9 @@ make the passes available along with the dialect.
Dialect functionality other than IR objects or passes, such as helper functions,
can be exposed to Python similarly to attributes and types. C API is expected to
exist for this functionality, which can then be wrapped using pybind11 and
-`[include/mlir/Bindings/Python/PybindAdaptors.h](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/PybindAdaptors.h)`
+`[include/mlir/Bindings/Python/PybindAdaptors.h](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/PybindAdaptors.h)`,
+or nanobind and
+`[include/mlir/Bindings/Python/NanobindAdaptors.h](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h)`
utilities to connect to the rest of Python API. The bindings can be located in a
-separate pybind11 module or in the same module as attributes and types, and
+separate module or in the same module as attributes and types, and
loaded along with the dialect.
diff --git a/mlir/examples/standalone/python/CMakeLists.txt b/mlir/examples/standalone/python/CMakeLists.txt
index a8c43827a5a375..69c82fd9135798 100644
--- a/mlir/examples/standalone/python/CMakeLists.txt
+++ b/mlir/examples/standalone/python/CMakeLists.txt
@@ -17,18 +17,32 @@ declare_mlir_dialect_python_bindings(
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir_standalone"
TD_FILE dialects/StandaloneOps.td
SOURCES
- dialects/standalone.py
+ dialects/standalone_pybind11.py
+ dialects/standalone_nanobind.py
DIALECT_NAME standalone)
-declare_mlir_python_extension(StandalonePythonSources.Extension
- MODULE_NAME _standaloneDialects
+
+declare_mlir_python_extension(StandalonePythonSources.Pybind11Extension
+ MODULE_NAME _standaloneDialectsPybind11
+ ADD_TO_PARENT StandalonePythonSources
+ SOURCES
+ StandaloneExtensionPybind11.cpp
+ EMBED_CAPI_LINK_LIBS
+ StandaloneCAPI
+ PYTHON_BINDINGS_LIBRARY pybind11
+)
+
+declare_mlir_python_extension(StandalonePythonSources.NanobindExtension
+ MODULE_NAME _standaloneDialectsNanobind
ADD_TO_PARENT StandalonePythonSources
SOURCES
- StandaloneExtension.cpp
+ StandaloneExtensionNanobind.cpp
EMBED_CAPI_LINK_LIBS
StandaloneCAPI
+ PYTHON_BINDINGS_LIBRARY nanobind
)
+
################################################################################
# Common CAPI
################################################################################
diff --git a/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp b/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp
new file mode 100644
index 00000000000000..6d83dc585dcd1d
--- /dev/null
+++ b/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp
@@ -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);
+}
diff --git a/mlir/examples/standalone/python/StandaloneExtension.cpp b/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
similarity index 81%
rename from mlir/examples/standalone/python/StandaloneExtension.cpp
rename to mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
index 5e83060cd48d82..397db4c20e7432 100644
--- a/mlir/examples/standalone/python/StandaloneExtension.cpp
+++ b/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
@@ -1,4 +1,7 @@
-//===- StandaloneExtension.cpp - Extension module -------------------------===//
+//===- StandaloneExtensionPybind11.cpp - Extension module -----------------===//
+//
+// This is the pybind11 version of the example module. There is also a nanobind
+// example in StandaloneExtensionNanobind.cpp.
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
@@ -11,7 +14,7 @@
using namespace mlir::python::adaptors;
-PYBIND11_MODULE(_standaloneDialects, m) {
+PYBIND11_MODULE(_standaloneDialectsPybind11, m) {
//===--------------------------------------------------------------------===//
// standalone dialect
//===--------------------------------------------------------------------===//
diff --git a/mlir/examples/standalone/python/mlir_standalone/dialects/standalone.py b/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_nanobind.py
similarity index 78%
rename from mlir/examples/standalone/python/mlir_standalone/dialects/standalone.py
rename to mlir/examples/standalone/python/mlir_standalone/dialects/standalone_nanobind.py
index c958b2ac193682..6218720951c82a 100644
--- a/mlir/examples/standalone/python/mlir_standalone/dialects/standalone.py
+++ b/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_nanobind.py
@@ -3,4 +3,4 @@
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from ._standalone_ops_gen import *
-from .._mlir_libs._standaloneDialects.standalone import *
+from .._mlir_libs._standaloneDialectsNanobind.standalone import *
diff --git a/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py b/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py
new file mode 100644
index 00000000000000..bfb98e404e13f2
--- /dev/null
+++ b/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py
@@ -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 *
diff --git a/mlir/examples/standalone/test/python/smoketest.py b/mlir/examples/standalone/test/python/smoketest.py
index 08e08cbd2fe24c..bd40c65d161645 100644
--- a/mlir/examples/standalone/test/python/smoketest.py
+++ b/mlir/examples/standalone/test/python/smoketest.py
@@ -1,7 +1,17 @@
-# RUN: %python %s | FileCheck %s
+# RUN: %python %s pybind11 | FileCheck %s
+# RUN: %python %s nanobind | FileCheck %s
+import sys
from mlir_standalone.ir import *
-from mlir_standalone.dialects import builtin as builtin_d, standalone as standalone_d
+from mlir_standalone.dialects import builtin as builtin_d
+
+if sys.argv[1] == "pybind11":
+ from mlir_standalone.dialects import standalone_pybind11 as standalone_d
+elif sys.argv[1] == "nanobind":
+ from mlir_standalone.dialects import standalone_nanobind as standalone_d
+else:
+ raise ValueError("Expected either pybind11 or nanobind as arguments")
+
with Context():
standalone_d.register_dialect()
diff --git a/mlir/include/mlir/Bindings/Python/Diagnostics.h b/mlir/include/mlir/Bindings/Python/Diagnostics.h
new file mode 100644
index 00000000000000..ea80e14dde0f3a
--- /dev/null
+++ b/mlir/include/mlir/Bindings/Python/Diagnostics.h
@@ -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
diff --git a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
new file mode 100644
index 00000000000000..5e01cebcb09c91
--- /dev/null
+++ b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
@@ -0,0 +1,671 @@
+//===- NanobindAdaptors.h - Interop with MLIR APIs via nanobind -----------===//
+//
+// 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 file contains adaptors for clients of the core MLIR Python APIs to
+// interop via MLIR CAPI types, using nanobind. The facilities here do not
+// depend on implementation details of the MLIR Python API and do not introduce
+// C++-level dependencies with it (requiring only Python and CAPI-level
+// dependencies).
+//
+// It is encouraged to be used both in-tree and out-of-tree. For in-tree use
+// cases, it should be used for dialect implementations (versus relying on
+// Pybind-based internals of the core libraries).
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_BINDINGS_PYTHON_NANOBINDADAPTORS_H
+#define MLIR_BINDINGS_PYTHON_NANOBINDADAPTORS_H
+
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/string.h>
+
+#include <cstdint>
+
+#include "mlir-c/Bindings/Python/Interop.h"
+#include "mlir-c/Diagnostics.h"
+#include "mlir-c/IR.h"
+#include "llvm/ADT/Twine.h"
+
+// Raw CAPI type casters need to be declared before use, so always include them
+// first.
+namespace nanobind {
+namespace detail {
+
+/// Helper to convert a presumed MLIR API object to a capsule, accepting either
+/// an explicit Capsule (which can happen when two C APIs are communicating
+/// directly via Python) or indirectly by querying the MLIR_PYTHON_CAPI_PTR_ATTR
+/// attribute (through which supported MLIR Python API objects export their
+/// contained API pointer as a capsule). Throws a type error if the object is
+/// neither. This is intended to be used from type casters, which are invoked
+/// with a raw handle (unowned). The returned object's lifetime may not extend
+/// beyond the apiObject handle without explicitly having its refcount increased
+/// (i.e. on return).
+static nanobind::object mlirApiObjectToCapsule(nanobind::handle apiObject) {
+ if (PyCapsule_CheckExact(apiObject.ptr()))
+ return nanobind::borrow<nanobind::object>(apiObject);
+ if (!nanobind::hasattr(apiObject, MLIR_PYTHON_CAPI_PTR_ATTR)) {
+ std::string repr = nanobind::cast<std::string>(nanobind::repr(apiObject));
+ throw nanobind::type_error(
+ (llvm::Twine("Expected an MLIR object (got ") + repr + ").")
+ .str()
+ .c_str());
+ }
+ return apiObject.attr(MLIR_PYTHON_CAPI_PTR_ATTR);
+}
+
+// Note: Currently all of the following support cast from nanobind::object to
+// the Mlir* C-API type, but only a few light-weight, context-bound ones
+// implicitly cast the other way because the use case has not yet emerged and
+// ownership is unclear.
+
+/// Casts object <-> MlirAffineMap.
+template <>
+struct type_caster<MlirAffineMap> {
+ NB_TYPE_CASTER(MlirAffineMap, const_name("MlirAffineMap"));
+ bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
+ nanobind::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToAffineMap(capsule.ptr());
+ if (mlirAffineMapIsNull(value)) {
+ return false;
+ }
+ return !mlirAffineMapIsNull(value);
+ }
+ static handle from_cpp(MlirAffineMap v, rv_policy,
+ cleanup_list *cleanup) noexcept {
+ nanobind::object capsule =
+ nanobind::steal<nanobind::object>(mlirPythonAffineMapToCapsule(v));
+ return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("AffineMap")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .release();
+ }
+};
+
+/// Casts object <-> MlirAttribute.
+template <>
+struct type_caster<MlirAttribute> {
+ NB_TYPE_CASTER(MlirAttribute, const_name("MlirAttribute"));
+ bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
+ nanobind::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToAttribute(capsule.ptr());
+ return !mlirAttributeIsNull(value);
+ }
+ static handle from_cpp(MlirAttribute v, rv_policy,
+ cleanup_list *cleanup) noexcept {
+ nanobind::object capsule =
+ nanobind::steal<nanobind::object>(mlirPythonAttributeToCapsule(v));
+ return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Attribute")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .attr(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR)()
+ .release();
+ }
+};
+
+/// Casts object -> MlirBlock.
+template <>
+struct type_caster<MlirBlock> {
+ NB_TYPE_CASTER(MlirBlock, const_name("MlirBlock"));
+ bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
+ nanobind::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToBlock(capsule.ptr());
+ return !mlirBlockIsNull(value);
+ }
+};
+
+/// Casts object -> MlirContext.
+template <>
+struct type_caster<MlirContext> {
+ NB_TYPE_CASTER(MlirContext, const_name("MlirContext"));
+ bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
+ if (src.is_none()) {
+ // Gets the current thread-bound context.
+ // TODO: This raises an error of "No current context" currently.
+ // Update the implementation to pretty-print the helpful error that the
+ // core implementations print in this case.
+ src = nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Context")
+ .attr("current");
+ }
+ nanobind::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToContext(capsule.ptr());
+ return !mlirContextIsNull(value);
+ }
+};
+
+/// Casts object <-> MlirDialectRegistry.
+template <>
+struct type_caster<MlirDialectRegistry> {
+ NB_TYPE_CASTER(MlirDialectRegistry, const_name("MlirDialectRegistry"));
+ bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
+ nanobind::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToDialectRegistry(capsule.ptr());
+ return !mlirDialectRegistryIsNull(value);
+ }
+ static handle from_cpp(MlirDialectRegistry v, rv_policy,
+ cleanup_list *cleanup) noexcept {
+ nanobind::object capsule = nanobind::steal<nanobind::object>(
+ mlirPythonDialectRegistryToCapsule(v));
+ return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("DialectRegistry")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .release();
+ }
+};
+
+/// Casts object <-> MlirLocation.
+template <>
+struct type_caster<MlirLocation> {
+ NB_TYPE_CASTER(MlirLocation, const_name("MlirLocation"));
+ bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
+ if (src.is_none()) {
+ // Gets the current thread-bound context.
+ src = nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Location")
+ .attr("current");
+ }
+ nanobind::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToLocation(capsule.ptr());
+ return !mlirLocationIsNull(value);
+ }
+ static handle from_cpp(MlirLocation v, rv_policy,
+ cleanup_list *cleanup) noexcept {
+ nanobind::object capsule =
+ nanobind::steal<nanobind::object>(mlirPythonLocationToCapsule(v));
+ return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Location")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .release();
+ }
+};
+
+/// Casts object <-> MlirModule.
+template <>
+struct type_caster<MlirModule> {
+ NB_TYPE_CASTER(MlirModule, const_name("MlirModule"));
+ bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
+ nanobind::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToModule(capsule.ptr());
+ return !mlirModuleIsNull(value);
+ }
+ static handle from_cpp(MlirModule v, rv_policy,
+ cleanup_list *cleanup) noexcept {
+ nanobind::object capsule =
+ nanobind::steal<nanobind::object>(mlirPythonModuleToCapsule(v));
+ return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Module")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .release();
+ };
+};
+
+/// Casts object <-> MlirFrozenRewritePatternSet.
+template <>
+struct type_caster<MlirFrozenRewritePatternSet> {
+ NB_TYPE_CASTER(MlirFrozenRewritePatternSet,
+ const_name("MlirFrozenRewritePatternSet"));
+ bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
+ nanobind::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToFrozenRewritePatternSet(capsule.ptr());
+ return value.ptr != nullptr;
+ }
+ static handle from_cpp(MlirFrozenRewritePatternSet v, rv_policy, handle) {
+ nanobind::object capsule = nanobind::steal<nanobind::object>(
+ mlirPythonFrozenRewritePatternSetToCapsule(v));
+ return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("rewrite"))
+ .attr("FrozenRewritePatternSet")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .release();
+ };
+};
+
+/// Casts object <-> MlirOperation.
+template <>
+struct type_caster<MlirOperation> {
+ NB_TYPE_CASTER(MlirOperation, const_name("MlirOperation"));
+ bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
+ nanobind::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToOperation(capsule.ptr());
+ return !mlirOperationIsNull(value);
+ }
+ static handle from_cpp(MlirOperation v, rv_policy,
+ cleanup_list *cleanup) noexcept {
+ if (v.ptr == nullptr)
+ return nanobind::none();
+ nanobind::object capsule =
+ nanobind::steal<nanobind::object>(mlirPythonOperationToCapsule(v));
+ return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Operation")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .release();
+ };
+};
+
+/// Casts object <-> MlirValue.
+template <>
+struct type_caster<MlirValue> {
+ NB_TYPE_CASTER(MlirValue, const_name("MlirValue"));
+ bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
+ nanobind::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToValue(capsule.ptr());
+ return !mlirValueIsNull(value);
+ }
+ static handle from_cpp(MlirValue v, rv_policy,
+ cleanup_list *cleanup) noexcept {
+ if (v.ptr == nullptr)
+ return nanobind::none();
+ nanobind::object capsule =
+ nanobind::steal<nanobind::object>(mlirPythonValueToCapsule(v));
+ return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Value")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .attr(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR)()
+ .release();
+ };
+};
+
+/// Casts object -> MlirPassManager.
+template <>
+struct type_caster<MlirPassManager> {
+ NB_TYPE_CASTER(MlirPassManager, const_name("MlirPassManager"));
+ bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
+ nanobind::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToPassManager(capsule.ptr());
+ return !mlirPassManagerIsNull(value);
+ }
+};
+
+/// Casts object <-> MlirTypeID.
+template <>
+struct type_caster<MlirTypeID> {
+ NB_TYPE_CASTER(MlirTypeID, const_name("MlirTypeID"));
+ bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
+ nanobind::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToTypeID(capsule.ptr());
+ return !mlirTypeIDIsNull(value);
+ }
+ static handle from_cpp(MlirTypeID v, rv_policy,
+ cleanup_list *cleanup) noexcept {
+ if (v.ptr == nullptr)
+ return nanobind::none();
+ nanobind::object capsule =
+ nanobind::steal<nanobind::object>(mlirPythonTypeIDToCapsule(v));
+ return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("TypeID")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .release();
+ };
+};
+
+/// Casts object <-> MlirType.
+template <>
+struct type_caster<MlirType> {
+ NB_TYPE_CASTER(MlirType, const_name("MlirType"));
+ bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
+ nanobind::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToType(capsule.ptr());
+ return !mlirTypeIsNull(value);
+ }
+ static handle from_cpp(MlirType t, rv_policy,
+ cleanup_list *cleanup) noexcept {
+ nanobind::object capsule =
+ nanobind::steal<nanobind::object>(mlirPythonTypeToCapsule(t));
+ return nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Type")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .attr(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR)()
+ .release();
+ }
+};
+
+} // namespace detail
+} // namespace nanobind
+
+namespace mlir {
+namespace python {
+namespace nanobind_adaptors {
+
+/// Provides a facility like nanobind::class_ for defining a new class in a
+/// scope, but this allows extension of an arbitrary Python class, defining
+/// methods on it is a similar way. Classes defined in this way are very similar
+/// to if defined in Python in the usual way but use nanobind machinery to
+/// do it. These are not "real" nanobind classes but pure Python classes
+/// with no relation to a concrete C++ class.
+///
+/// Derived from a discussion upstream:
+/// https://github.com/pybind/pybind11/issues/1193
+/// (plus a fair amount of extra curricular poking)
+/// TODO: If this proves useful, see about including it in nanobind.
+class pure_subclass {
+public:
+ pure_subclass(nanobind::handle scope, const char *derivedClassName,
+ const nanobind::object &superClass) {
+ nanobind::object pyType =
+ nanobind::borrow<nanobind::object>((PyObject *)&PyType_Type);
+ nanobind::object metaclass = pyType(superClass);
+ nanobind::dict attributes;
+
+ thisClass = metaclass(derivedClassName, nanobind::make_tuple(superClass),
+ attributes);
+ scope.attr(derivedClassName) = thisClass;
+ }
+
+ template <typename Func, typename... Extra>
+ pure_subclass &def(const char *name, Func &&f, const Extra &...extra) {
+ nanobind::object cf = nanobind::cpp_function(
+ std::forward<Func>(f), nanobind::name(name), nanobind::is_method(),
+ nanobind::scope(thisClass), extra...);
+ thisClass.attr(name) = cf;
+ return *this;
+ }
+
+ template <typename Func, typename... Extra>
+ pure_subclass &def_property_readonly(const char *name, Func &&f,
+ const Extra &...extra) {
+ nanobind::object cf = nanobind::cpp_function(
+ std::forward<Func>(f), nanobind::name(name), nanobind::is_method(),
+ nanobind::scope(thisClass), extra...);
+ auto builtinProperty =
+ nanobind::borrow<nanobind::object>((PyObject *)&PyProperty_Type);
+ thisClass.attr(name) = builtinProperty(cf);
+ return *this;
+ }
+
+ template <typename Func, typename... Extra>
+ pure_subclass &def_staticmethod(const char *name, Func &&f,
+ const Extra &...extra) {
+ static_assert(!std::is_member_function_pointer<Func>::value,
+ "def_staticmethod(...) called with a non-static member "
+ "function pointer");
+ nanobind::object cf = nanobind::cpp_function(
+ std::forward<Func>(f),
+ nanobind::name(name), // nanobind::scope(thisClass),
+ extra...);
+ thisClass.attr(name) = cf;
+ return *this;
+ }
+
+ template <typename Func, typename... Extra>
+ pure_subclass &def_classmethod(const char *name, Func &&f,
+ const Extra &...extra) {
+ static_assert(!std::is_member_function_pointer<Func>::value,
+ "def_classmethod(...) called with a non-static member "
+ "function pointer");
+ nanobind::object cf = nanobind::cpp_function(
+ std::forward<Func>(f),
+ nanobind::name(name), // nanobind::scope(thisClass),
+ extra...);
+ thisClass.attr(name) =
+ nanobind::borrow<nanobind::object>(PyClassMethod_New(cf.ptr()));
+ return *this;
+ }
+
+ nanobind::object get_class() const { return thisClass; }
+
+protected:
+ nanobind::object superClass;
+ nanobind::object thisClass;
+};
+
+/// Creates a custom subclass of mlir.ir.Attribute, implementing a casting
+/// constructor and type checking methods.
+class mlir_attribute_subclass : public pure_subclass {
+public:
+ using IsAFunctionTy = bool (*)(MlirAttribute);
+ using GetTypeIDFunctionTy = MlirTypeID (*)();
+
+ /// Subclasses by looking up the super-class dynamically.
+ mlir_attribute_subclass(nanobind::handle scope, const char *attrClassName,
+ IsAFunctionTy isaFunction,
+ GetTypeIDFunctionTy getTypeIDFunction = nullptr)
+ : mlir_attribute_subclass(
+ scope, attrClassName, isaFunction,
+ nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Attribute"),
+ getTypeIDFunction) {}
+
+ /// Subclasses with a provided mlir.ir.Attribute super-class. This must
+ /// be used if the subclass is being defined in the same extension module
+ /// as the mlir.ir class (otherwise, it will trigger a recursive
+ /// initialization).
+ mlir_attribute_subclass(nanobind::handle scope, const char *typeClassName,
+ IsAFunctionTy isaFunction,
+ const nanobind::object &superCls,
+ GetTypeIDFunctionTy getTypeIDFunction = nullptr)
+ : pure_subclass(scope, typeClassName, superCls) {
+ // Casting constructor. Note that it hard, if not impossible, to properly
+ // call chain to parent `__init__` in nanobind due to its special handling
+ // for init functions that don't have a fully constructed self-reference,
+ // which makes it impossible to forward it to `__init__` of a superclass.
+ // Instead, provide a custom `__new__` and call that of a superclass, which
+ // eventually calls `__init__` of the superclass. Since attribute subclasses
+ // have no additional members, we can just return the instance thus created
+ // without amending it.
+ std::string captureTypeName(
+ typeClassName); // As string in case if typeClassName is not static.
+ nanobind::object newCf = nanobind::cpp_function(
+ [superCls, isaFunction, captureTypeName](
+ nanobind::object cls, nanobind::object otherAttribute) {
+ MlirAttribute rawAttribute =
+ nanobind::cast<MlirAttribute>(otherAttribute);
+ if (!isaFunction(rawAttribute)) {
+ auto origRepr =
+ nanobind::cast<std::string>(nanobind::repr(otherAttribute));
+ throw std::invalid_argument(
+ (llvm::Twine("Cannot cast attribute to ") + captureTypeName +
+ " (from " + origRepr + ")")
+ .str());
+ }
+ nanobind::object self = superCls.attr("__new__")(cls, otherAttribute);
+ return self;
+ },
+ nanobind::name("__new__"), nanobind::arg("cls"),
+ nanobind::arg("cast_from_attr"));
+ thisClass.attr("__new__") = newCf;
+
+ // 'isinstance' method.
+ def_staticmethod(
+ "isinstance",
+ [isaFunction](MlirAttribute other) { return isaFunction(other); },
+ nanobind::arg("other_attribute"));
+ def("__repr__", [superCls, captureTypeName](nanobind::object self) {
+ return nanobind::repr(superCls(self))
+ .attr("replace")(superCls.attr("__name__"), captureTypeName);
+ });
+ if (getTypeIDFunction) {
+ def_staticmethod("get_static_typeid",
+ [getTypeIDFunction]() { return getTypeIDFunction(); });
+ nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
+ getTypeIDFunction())(nanobind::cpp_function(
+ [thisClass = thisClass](const nanobind::object &mlirAttribute) {
+ return thisClass(mlirAttribute);
+ }));
+ }
+ }
+};
+
+/// Creates a custom subclass of mlir.ir.Type, implementing a casting
+/// constructor and type checking methods.
+class mlir_type_subclass : public pure_subclass {
+public:
+ using IsAFunctionTy = bool (*)(MlirType);
+ using GetTypeIDFunctionTy = MlirTypeID (*)();
+
+ /// Subclasses by looking up the super-class dynamically.
+ mlir_type_subclass(nanobind::handle scope, const char *typeClassName,
+ IsAFunctionTy isaFunction,
+ GetTypeIDFunctionTy getTypeIDFunction = nullptr)
+ : mlir_type_subclass(
+ scope, typeClassName, isaFunction,
+ nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Type"),
+ getTypeIDFunction) {}
+
+ /// Subclasses with a provided mlir.ir.Type super-class. This must
+ /// be used if the subclass is being defined in the same extension module
+ /// as the mlir.ir class (otherwise, it will trigger a recursive
+ /// initialization).
+ mlir_type_subclass(nanobind::handle scope, const char *typeClassName,
+ IsAFunctionTy isaFunction,
+ const nanobind::object &superCls,
+ GetTypeIDFunctionTy getTypeIDFunction = nullptr)
+ : pure_subclass(scope, typeClassName, superCls) {
+ // Casting constructor. Note that it hard, if not impossible, to properly
+ // call chain to parent `__init__` in nanobind due to its special handling
+ // for init functions that don't have a fully constructed self-reference,
+ // which makes it impossible to forward it to `__init__` of a superclass.
+ // Instead, provide a custom `__new__` and call that of a superclass, which
+ // eventually calls `__init__` of the superclass. Since attribute subclasses
+ // have no additional members, we can just return the instance thus created
+ // without amending it.
+ std::string captureTypeName(
+ typeClassName); // As string in case if typeClassName is not static.
+ nanobind::object newCf = nanobind::cpp_function(
+ [superCls, isaFunction, captureTypeName](nanobind::object cls,
+ nanobind::object otherType) {
+ MlirType rawType = nanobind::cast<MlirType>(otherType);
+ if (!isaFunction(rawType)) {
+ auto origRepr =
+ nanobind::cast<std::string>(nanobind::repr(otherType));
+ throw std::invalid_argument((llvm::Twine("Cannot cast type to ") +
+ captureTypeName + " (from " +
+ origRepr + ")")
+ .str());
+ }
+ nanobind::object self = superCls.attr("__new__")(cls, otherType);
+ return self;
+ },
+ nanobind::name("__new__"), nanobind::arg("cls"),
+ nanobind::arg("cast_from_type"));
+ thisClass.attr("__new__") = newCf;
+
+ // 'isinstance' method.
+ def_staticmethod(
+ "isinstance",
+ [isaFunction](MlirType other) { return isaFunction(other); },
+ nanobind::arg("other_type"));
+ def("__repr__", [superCls, captureTypeName](nanobind::object self) {
+ return nanobind::repr(superCls(self))
+ .attr("replace")(superCls.attr("__name__"), captureTypeName);
+ });
+ if (getTypeIDFunction) {
+ // 'get_static_typeid' method.
+ // This is modeled as a static method instead of a static property because
+ // `def_property_readonly_static` is not available in `pure_subclass` and
+ // we do not want to introduce the complexity that pybind uses to
+ // implement it.
+ def_staticmethod("get_static_typeid",
+ [getTypeIDFunction]() { return getTypeIDFunction(); });
+ nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
+ getTypeIDFunction())(nanobind::cpp_function(
+ [thisClass = thisClass](const nanobind::object &mlirType) {
+ return thisClass(mlirType);
+ }));
+ }
+ }
+};
+
+/// Creates a custom subclass of mlir.ir.Value, implementing a casting
+/// constructor and type checking methods.
+class mlir_value_subclass : public pure_subclass {
+public:
+ using IsAFunctionTy = bool (*)(MlirValue);
+
+ /// Subclasses by looking up the super-class dynamically.
+ mlir_value_subclass(nanobind::handle scope, const char *valueClassName,
+ IsAFunctionTy isaFunction)
+ : mlir_value_subclass(
+ scope, valueClassName, isaFunction,
+ nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Value")) {}
+
+ /// Subclasses with a provided mlir.ir.Value super-class. This must
+ /// be used if the subclass is being defined in the same extension module
+ /// as the mlir.ir class (otherwise, it will trigger a recursive
+ /// initialization).
+ mlir_value_subclass(nanobind::handle scope, const char *valueClassName,
+ IsAFunctionTy isaFunction,
+ const nanobind::object &superCls)
+ : pure_subclass(scope, valueClassName, superCls) {
+ // Casting constructor. Note that it hard, if not impossible, to properly
+ // call chain to parent `__init__` in nanobind due to its special handling
+ // for init functions that don't have a fully constructed self-reference,
+ // which makes it impossible to forward it to `__init__` of a superclass.
+ // Instead, provide a custom `__new__` and call that of a superclass, which
+ // eventually calls `__init__` of the superclass. Since attribute subclasses
+ // have no additional members, we can just return the instance thus created
+ // without amending it.
+ std::string captureValueName(
+ valueClassName); // As string in case if valueClassName is not static.
+ nanobind::object newCf = nanobind::cpp_function(
+ [superCls, isaFunction, captureValueName](nanobind::object cls,
+ nanobind::object otherValue) {
+ MlirValue rawValue = nanobind::cast<MlirValue>(otherValue);
+ if (!isaFunction(rawValue)) {
+ auto origRepr =
+ nanobind::cast<std::string>(nanobind::repr(otherValue));
+ throw std::invalid_argument((llvm::Twine("Cannot cast value to ") +
+ captureValueName + " (from " +
+ origRepr + ")")
+ .str());
+ }
+ nanobind::object self = superCls.attr("__new__")(cls, otherValue);
+ return self;
+ },
+ nanobind::name("__new__"), nanobind::arg("cls"),
+ nanobind::arg("cast_from_value"));
+ thisClass.attr("__new__") = newCf;
+
+ // 'isinstance' method.
+ def_staticmethod(
+ "isinstance",
+ [isaFunction](MlirValue other) { return isaFunction(other); },
+ nanobind::arg("other_value"));
+ }
+};
+
+} // namespace nanobind_adaptors
+
+/// 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_NANOBINDADAPTORS_H
diff --git a/mlir/include/mlir/Bindings/Python/PybindAdaptors.h b/mlir/include/mlir/Bindings/Python/PybindAdaptors.h
index df4b9bf713592d..c8233355d1d67b 100644
--- a/mlir/include/mlir/Bindings/Python/PybindAdaptors.h
+++ b/mlir/include/mlir/Bindings/Python/PybindAdaptors.h
@@ -1,4 +1,4 @@
-//===- PybindAdaptors.h - Adaptors for interop with MLIR APIs -------------===//
+//===- PybindAdaptors.h - Interop with MLIR APIs via pybind11 -------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
@@ -6,9 +6,10 @@
//
//===----------------------------------------------------------------------===//
// This file contains adaptors for clients of the core MLIR Python APIs to
-// interop via MLIR CAPI types. The facilities here do not depend on
-// implementation details of the MLIR Python API and do not introduce C++-level
-// dependencies with it (requiring only Python and CAPI-level dependencies).
+// interop via MLIR CAPI types, using pybind11. The facilities here do not
+// depend on implementation details of the MLIR Python API and do not introduce
+// C++-level dependencies with it (requiring only Python and CAPI-level
+// dependencies).
//
// It is encouraged to be used both in-tree and out-of-tree. For in-tree use
// cases, it should be used for dialect implementations (versus relying on
@@ -611,40 +612,6 @@ class mlir_value_subclass : public pure_subclass {
} // namespace adaptors
-/// 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
diff --git a/mlir/lib/Bindings/Python/DialectLLVM.cpp b/mlir/lib/Bindings/Python/DialectLLVM.cpp
index 42a4c8c0793ba8..cccf1370b8cc87 100644
--- a/mlir/lib/Bindings/Python/DialectLLVM.cpp
+++ b/mlir/lib/Bindings/Python/DialectLLVM.cpp
@@ -6,11 +6,13 @@
//
//===----------------------------------------------------------------------===//
+#include <string>
+
#include "mlir-c/Dialect/LLVM.h"
#include "mlir-c/IR.h"
#include "mlir-c/Support.h"
+#include "mlir/Bindings/Python/Diagnostics.h"
#include "mlir/Bindings/Python/PybindAdaptors.h"
-#include <string>
namespace py = pybind11;
using namespace llvm;
diff --git a/mlir/lib/Bindings/Python/TransformInterpreter.cpp b/mlir/lib/Bindings/Python/TransformInterpreter.cpp
index f6b4532b1b6be4..0c8c0e0a965aa7 100644
--- a/mlir/lib/Bindings/Python/TransformInterpreter.cpp
+++ b/mlir/lib/Bindings/Python/TransformInterpreter.cpp
@@ -10,14 +10,15 @@
//
//===----------------------------------------------------------------------===//
+#include <pybind11/detail/common.h>
+#include <pybind11/pybind11.h>
+
#include "mlir-c/Dialect/Transform/Interpreter.h"
#include "mlir-c/IR.h"
#include "mlir-c/Support.h"
+#include "mlir/Bindings/Python/Diagnostics.h"
#include "mlir/Bindings/Python/PybindAdaptors.h"
-#include <pybind11/detail/common.h>
-#include <pybind11/pybind11.h>
-
namespace py = pybind11;
namespace {
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 23187f256455bb..e1b870b53ad25c 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -683,7 +683,9 @@ if(MLIR_INCLUDE_TESTS)
MLIRPythonTestSources.Dialects.PythonTest
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
ADD_TO_PARENT MLIRPythonTestSources.Dialects
- SOURCES dialects/python_test.py)
+ SOURCES
+ dialects/python_test.py
+ )
set(LLVM_TARGET_DEFINITIONS
"${MLIR_MAIN_SRC_DIR}/test/python/python_test_ops.td")
mlir_tablegen(
@@ -697,12 +699,25 @@ if(MLIR_INCLUDE_TESTS)
ADD_TO_PARENT MLIRPythonTestSources.Dialects.PythonTest
SOURCES "dialects/_python_test_ops_gen.py")
- declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtension
- MODULE_NAME _mlirPythonTest
+ declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtensionPybind11
+ MODULE_NAME _mlirPythonTestPybind11
+ ADD_TO_PARENT MLIRPythonTestSources.Dialects
+ ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
+ PYTHON_BINDINGS_LIBRARY pybind11
+ SOURCES
+ PythonTestModulePybind11.cpp
+ PRIVATE_LINK_LIBS
+ LLVMSupport
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIPythonTestDialect
+ )
+ declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtensionNanobind
+ MODULE_NAME _mlirPythonTestNanobind
ADD_TO_PARENT MLIRPythonTestSources.Dialects
ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
- PythonTestModule.cpp
+ PythonTestModuleNanobind.cpp
PRIVATE_LINK_LIBS
LLVMSupport
EMBED_CAPI_LINK_LIBS
diff --git a/mlir/python/mlir/dialects/python_test.py b/mlir/python/mlir/dialects/python_test.py
index b5baa80bc767fb..9380896c8c06e8 100644
--- a/mlir/python/mlir/dialects/python_test.py
+++ b/mlir/python/mlir/dialects/python_test.py
@@ -3,15 +3,14 @@
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from ._python_test_ops_gen import *
-from .._mlir_libs._mlirPythonTest import (
- TestAttr,
- TestType,
- TestTensorValue,
- TestIntegerRankedTensorType,
-)
-def register_python_test_dialect(registry):
- from .._mlir_libs import _mlirPythonTest
+def register_python_test_dialect(registry, use_nanobind):
+ if use_nanobind:
+ from .._mlir_libs import _mlirPythonTestNanobind
- _mlirPythonTest.register_dialect(registry)
+ _mlirPythonTestNanobind.register_dialect(registry)
+ else:
+ from .._mlir_libs import _mlirPythonTestPybind11
+
+ _mlirPythonTestPybind11.register_dialect(registry)
diff --git a/mlir/python/requirements.txt b/mlir/python/requirements.txt
index 272d066831f927..ab8a9122919e19 100644
--- a/mlir/python/requirements.txt
+++ b/mlir/python/requirements.txt
@@ -1,3 +1,4 @@
+nanobind>=2.0, <3.0
numpy>=1.19.5, <=2.1.2
pybind11>=2.10.0, <=2.13.6
PyYAML>=5.4.0, <=6.0.1
diff --git a/mlir/test/python/dialects/python_test.py b/mlir/test/python/dialects/python_test.py
index 948d1225ea489c..fd678f8321fd93 100644
--- a/mlir/test/python/dialects/python_test.py
+++ b/mlir/test/python/dialects/python_test.py
@@ -1,12 +1,33 @@
-# RUN: %PYTHON %s | FileCheck %s
+# RUN: %PYTHON %s pybind11 | FileCheck %s
+# RUN: %PYTHON %s nanobind | FileCheck %s
+import sys
from mlir.ir import *
import mlir.dialects.func as func
import mlir.dialects.python_test as test
import mlir.dialects.tensor as tensor
import mlir.dialects.arith as arith
-test.register_python_test_dialect(get_dialect_registry())
+if sys.argv[1] == "pybind11":
+ from mlir._mlir_libs._mlirPythonTestPybind11 import (
+ TestAttr,
+ TestType,
+ TestTensorValue,
+ TestIntegerRankedTensorType,
+ )
+
+ test.register_python_test_dialect(get_dialect_registry(), use_nanobind=False)
+elif sys.argv[1] == "nanobind":
+ from mlir._mlir_libs._mlirPythonTestNanobind import (
+ TestAttr,
+ TestType,
+ TestTensorValue,
+ TestIntegerRankedTensorType,
+ )
+
+ test.register_python_test_dialect(get_dialect_registry(), use_nanobind=True)
+else:
+ raise ValueError("Expected pybind11 or nanobind as argument")
def run(f):
@@ -308,7 +329,7 @@ def testOptionalOperandOp():
@run
def testCustomAttribute():
with Context() as ctx, Location.unknown():
- a = test.TestAttr.get()
+ a = TestAttr.get()
# CHECK: #python_test.test_attr
print(a)
@@ -325,11 +346,11 @@ def testCustomAttribute():
print(repr(op2.test_attr))
# The following cast must not assert.
- b = test.TestAttr(a)
+ b = TestAttr(a)
unit = UnitAttr.get()
try:
- test.TestAttr(unit)
+ TestAttr(unit)
except ValueError as e:
assert "Cannot cast attribute to TestAttr" in str(e)
else:
@@ -338,7 +359,7 @@ def testCustomAttribute():
# The following must trigger a TypeError from our adaptors and must not
# crash.
try:
- test.TestAttr(42)
+ TestAttr(42)
except TypeError as e:
assert "Expected an MLIR object" in str(e)
else:
@@ -347,7 +368,7 @@ def testCustomAttribute():
# The following must trigger a TypeError from pybind (therefore, not
# checking its message) and must not crash.
try:
- test.TestAttr(42, 56)
+ TestAttr(42, 56)
except TypeError:
pass
else:
@@ -357,12 +378,12 @@ def testCustomAttribute():
@run
def testCustomType():
with Context() as ctx:
- a = test.TestType.get()
+ a = TestType.get()
# CHECK: !python_test.test_type
print(a)
# The following cast must not assert.
- b = test.TestType(a)
+ b = TestType(a)
# Instance custom types should have typeids
assert isinstance(b.typeid, TypeID)
# Subclasses of ir.Type should not have a static_typeid
@@ -374,7 +395,7 @@ def testCustomType():
i8 = IntegerType.get_signless(8)
try:
- test.TestType(i8)
+ TestType(i8)
except ValueError as e:
assert "Cannot cast type to TestType" in str(e)
else:
@@ -383,7 +404,7 @@ def testCustomType():
# The following must trigger a TypeError from our adaptors and must not
# crash.
try:
- test.TestType(42)
+ TestType(42)
except TypeError as e:
assert "Expected an MLIR object" in str(e)
else:
@@ -392,7 +413,7 @@ def testCustomType():
# The following must trigger a TypeError from pybind (therefore, not
# checking its message) and must not crash.
try:
- test.TestType(42, 56)
+ TestType(42, 56)
except TypeError:
pass
else:
@@ -405,7 +426,7 @@ def testTensorValue():
with Context() as ctx, Location.unknown():
i8 = IntegerType.get_signless(8)
- class Tensor(test.TestTensorValue):
+ class Tensor(TestTensorValue):
def __str__(self):
return super().__str__().replace("Value", "Tensor")
@@ -425,9 +446,9 @@ def __str__(self):
# Classes of custom types that inherit from concrete types should have
# static_typeid
- assert isinstance(test.TestIntegerRankedTensorType.static_typeid, TypeID)
+ assert isinstance(TestIntegerRankedTensorType.static_typeid, TypeID)
# And it should be equal to the in-tree concrete type
- assert test.TestIntegerRankedTensorType.static_typeid == t.type.typeid
+ assert TestIntegerRankedTensorType.static_typeid == t.type.typeid
d = tensor.EmptyOp([1, 2, 3], IntegerType.get_signless(5)).result
# CHECK: Value(%{{.*}} = tensor.empty() : tensor<1x2x3xi5>)
@@ -491,7 +512,7 @@ def inferReturnTypeComponents():
@run
def testCustomTypeTypeCaster():
with Context() as ctx, Location.unknown():
- a = test.TestType.get()
+ a = TestType.get()
assert a.typeid is not None
b = Type.parse("!python_test.test_type")
@@ -500,7 +521,7 @@ def testCustomTypeTypeCaster():
# CHECK: TestType(!python_test.test_type)
print(repr(b))
- c = test.TestIntegerRankedTensorType.get([10, 10], 5)
+ c = TestIntegerRankedTensorType.get([10, 10], 5)
# CHECK: tensor<10x10xi5>
print(c)
# CHECK: TestIntegerRankedTensorType(tensor<10x10xi5>)
@@ -511,7 +532,7 @@ def testCustomTypeTypeCaster():
@register_type_caster(c.typeid)
def type_caster(pytype):
- return test.TestIntegerRankedTensorType(pytype)
+ return TestIntegerRankedTensorType(pytype)
except RuntimeError as e:
print(e)
@@ -530,7 +551,7 @@ def type_caster(pytype):
@register_type_caster(c.typeid, replace=True)
def type_caster(pytype):
- return test.TestIntegerRankedTensorType(pytype)
+ return TestIntegerRankedTensorType(pytype)
d = tensor.EmptyOp([10, 10], IntegerType.get_signless(5)).result
# CHECK: tensor<10x10xi5>
diff --git a/mlir/test/python/lib/CMakeLists.txt b/mlir/test/python/lib/CMakeLists.txt
index d7cbbfbc214772..198ed8211e773f 100644
--- a/mlir/test/python/lib/CMakeLists.txt
+++ b/mlir/test/python/lib/CMakeLists.txt
@@ -1,7 +1,8 @@
set(LLVM_OPTIONAL_SOURCES
PythonTestCAPI.cpp
PythonTestDialect.cpp
- PythonTestModule.cpp
+ PythonTestModulePybind11.cpp
+ PythonTestModuleNanobind.cpp
)
add_mlir_library(MLIRPythonTestDialect
diff --git a/mlir/test/python/lib/PythonTestModuleNanobind.cpp b/mlir/test/python/lib/PythonTestModuleNanobind.cpp
new file mode 100644
index 00000000000000..6c4a04a38c2e94
--- /dev/null
+++ b/mlir/test/python/lib/PythonTestModuleNanobind.cpp
@@ -0,0 +1,121 @@
+//===- PythonTestModuleNanobind.cpp - PythonTest dialect extension --------===//
+//
+// 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 is the nanobind edition of the PythonTest dialect module.
+//===----------------------------------------------------------------------===//
+
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/vector.h>
+
+#include "PythonTestCAPI.h"
+#include "mlir-c/BuiltinAttributes.h"
+#include "mlir-c/BuiltinTypes.h"
+#include "mlir-c/IR.h"
+#include "mlir/Bindings/Python/NanobindAdaptors.h"
+
+namespace nb = nanobind;
+using namespace mlir::python::nanobind_adaptors;
+
+static bool mlirTypeIsARankedIntegerTensor(MlirType t) {
+ return mlirTypeIsARankedTensor(t) &&
+ mlirTypeIsAInteger(mlirShapedTypeGetElementType(t));
+}
+
+NB_MODULE(_mlirPythonTestNanobind, m) {
+ m.def(
+ "register_python_test_dialect",
+ [](MlirContext context, bool load) {
+ MlirDialectHandle pythonTestDialect =
+ mlirGetDialectHandle__python_test__();
+ mlirDialectHandleRegisterDialect(pythonTestDialect, context);
+ if (load) {
+ mlirDialectHandleLoadDialect(pythonTestDialect, context);
+ }
+ },
+ nb::arg("context"), nb::arg("load") = true);
+
+ m.def(
+ "register_dialect",
+ [](MlirDialectRegistry registry) {
+ MlirDialectHandle pythonTestDialect =
+ mlirGetDialectHandle__python_test__();
+ mlirDialectHandleInsertDialect(pythonTestDialect, registry);
+ },
+ nb::arg("registry"));
+
+ mlir_attribute_subclass(m, "TestAttr",
+ mlirAttributeIsAPythonTestTestAttribute,
+ mlirPythonTestTestAttributeGetTypeID)
+ .def_classmethod(
+ "get",
+ [](const nb::object &cls, MlirContext ctx) {
+ return cls(mlirPythonTestTestAttributeGet(ctx));
+ },
+ nb::arg("cls"), nb::arg("context").none() = nb::none());
+
+ mlir_type_subclass(m, "TestType", mlirTypeIsAPythonTestTestType,
+ mlirPythonTestTestTypeGetTypeID)
+ .def_classmethod(
+ "get",
+ [](const nb::object &cls, MlirContext ctx) {
+ return cls(mlirPythonTestTestTypeGet(ctx));
+ },
+ nb::arg("cls"), nb::arg("context").none() = nb::none());
+
+ auto typeCls =
+ mlir_type_subclass(m, "TestIntegerRankedTensorType",
+ mlirTypeIsARankedIntegerTensor,
+ nb::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("RankedTensorType"))
+ .def_classmethod(
+ "get",
+ [](const nb::object &cls, std::vector<int64_t> shape,
+ unsigned width, MlirContext ctx) {
+ MlirAttribute encoding = mlirAttributeGetNull();
+ return cls(mlirRankedTensorTypeGet(
+ shape.size(), shape.data(), mlirIntegerTypeGet(ctx, width),
+ encoding));
+ },
+ nb::arg("cls"), nb::arg("shape"), nb::arg("width"),
+ nb::arg("context").none() = nb::none());
+
+ assert(nb::hasattr(typeCls.get_class(), "static_typeid") &&
+ "TestIntegerRankedTensorType has no static_typeid");
+
+ MlirTypeID mlirRankedTensorTypeID = mlirRankedTensorTypeGetTypeID();
+
+ nb::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
+ mlirRankedTensorTypeID, nb::arg("replace") = true)(
+ nanobind::cpp_function([typeCls](const nb::object &mlirType) {
+ return typeCls.get_class()(mlirType);
+ }));
+
+ auto valueCls = mlir_value_subclass(m, "TestTensorValue",
+ mlirTypeIsAPythonTestTestTensorValue)
+ .def("is_null", [](MlirValue &self) {
+ return mlirValueIsNull(self);
+ });
+
+ nb::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr(MLIR_PYTHON_CAPI_VALUE_CASTER_REGISTER_ATTR)(
+ mlirRankedTensorTypeID)(
+ nanobind::cpp_function([valueCls](const nb::object &valueObj) {
+ nb::object capsule = mlirApiObjectToCapsule(valueObj);
+ MlirValue v = mlirPythonCapsuleToValue(capsule.ptr());
+ MlirType t = mlirValueGetType(v);
+ // This is hyper-specific in order to exercise/test registering a
+ // value caster from cpp (but only for a single test case; see
+ // testTensorValue python_test_nanobind.py).
+ if (mlirShapedTypeHasStaticShape(t) &&
+ mlirShapedTypeGetDimSize(t, 0) == 1 &&
+ mlirShapedTypeGetDimSize(t, 1) == 2 &&
+ mlirShapedTypeGetDimSize(t, 2) == 3)
+ return valueCls.get_class()(valueObj);
+ return valueObj;
+ }));
+}
diff --git a/mlir/test/python/lib/PythonTestModule.cpp b/mlir/test/python/lib/PythonTestModulePybind11.cpp
similarity index 94%
rename from mlir/test/python/lib/PythonTestModule.cpp
rename to mlir/test/python/lib/PythonTestModulePybind11.cpp
index a4f538dcb55944..b4c8bd3d1dd6ee 100644
--- a/mlir/test/python/lib/PythonTestModule.cpp
+++ b/mlir/test/python/lib/PythonTestModulePybind11.cpp
@@ -5,6 +5,8 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
+// This is the pybind11 edition of the PythonTest dialect module.
+//===----------------------------------------------------------------------===//
#include "PythonTestCAPI.h"
#include "mlir-c/BuiltinAttributes.h"
@@ -21,7 +23,7 @@ static bool mlirTypeIsARankedIntegerTensor(MlirType t) {
mlirTypeIsAInteger(mlirShapedTypeGetElementType(t));
}
-PYBIND11_MODULE(_mlirPythonTest, m) {
+PYBIND11_MODULE(_mlirPythonTestPybind11, m) {
m.def(
"register_python_test_dialect",
[](MlirContext context, bool load) {
@@ -105,7 +107,7 @@ PYBIND11_MODULE(_mlirPythonTest, m) {
MlirType t = mlirValueGetType(v);
// This is hyper-specific in order to exercise/test registering a
// value caster from cpp (but only for a single test case; see
- // testTensorValue python_test.py).
+ // testTensorValue python_test_pybind11.py).
if (mlirShapedTypeHasStaticShape(t) &&
mlirShapedTypeGetDimSize(t, 0) == 1 &&
mlirShapedTypeGetDimSize(t, 1) == 2 &&
diff --git a/utils/bazel/WORKSPACE b/utils/bazel/WORKSPACE
index 7baca11eed3d39..66ba1ac1b17e1e 100644
--- a/utils/bazel/WORKSPACE
+++ b/utils/bazel/WORKSPACE
@@ -148,6 +148,24 @@ maybe(
url = "https://github.com/pybind/pybind11/archive/v2.10.3.zip",
)
+maybe(
+ http_archive,
+ name = "robin_map",
+ strip_prefix = "robin-map-1.3.0",
+ sha256 = "a8424ad3b0affd4c57ed26f0f3d8a29604f0e1f2ef2089f497f614b1c94c7236",
+ build_file = "@llvm-raw//utils/bazel/third_party_build:robin_map.BUILD",
+ url = "https://github.com/Tessil/robin-map/archive/refs/tags/v1.3.0.tar.gz",
+)
+
+maybe(
+ http_archive,
+ name = "nanobind",
+ build_file = "@llvm-raw//utils/bazel/third_party_build:nanobind.BUILD",
+ sha256 = "bfbfc7e5759f1669e4ddb48752b1ddc5647d1430e94614d6f8626df1d508e65a",
+ strip_prefix = "nanobind-2.2.0",
+ url = "https://github.com/wjakob/nanobind/archive/refs/tags/v2.2.0.tar.gz",
+)
+
load("@rules_python//python:repositories.bzl", "py_repositories", "python_register_toolchains")
py_repositories()
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index 179fed2f5e9a00..544becfa30b40f 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -932,7 +932,6 @@ exports_files(
filegroup(
name = "MLIRBindingsPythonHeaderFiles",
srcs = glob([
- "lib/Bindings/Python/*.h",
"include/mlir-c/Bindings/Python/*.h",
"include/mlir/Bindings/Python/*.h",
]),
@@ -942,12 +941,10 @@ cc_library(
name = "MLIRBindingsPythonHeaders",
includes = [
"include",
- "lib/Bindings/Python",
],
textual_hdrs = [":MLIRBindingsPythonHeaderFiles"],
deps = [
":CAPIIRHeaders",
- ":CAPITransformsHeaders",
"@pybind11",
"@rules_python//python/cc:current_py_cc_headers",
],
@@ -957,17 +954,41 @@ cc_library(
name = "MLIRBindingsPythonHeadersAndDeps",
includes = [
"include",
- "lib/Bindings/Python",
],
textual_hdrs = [":MLIRBindingsPythonHeaderFiles"],
deps = [
":CAPIIR",
- ":CAPITransforms",
"@pybind11",
"@rules_python//python/cc:current_py_cc_headers",
],
)
+cc_library(
+ name = "MLIRBindingsPythonNanobindHeaders",
+ includes = [
+ "include",
+ ],
+ textual_hdrs = [":MLIRBindingsPythonHeaderFiles"],
+ deps = [
+ ":CAPIIRHeaders",
+ "@nanobind",
+ "@rules_python//python/cc:current_py_cc_headers",
+ ],
+)
+
+cc_library(
+ name = "MLIRBindingsPythonNanobindHeadersAndDeps",
+ includes = [
+ "include",
+ ],
+ textual_hdrs = [":MLIRBindingsPythonHeaderFiles"],
+ deps = [
+ ":CAPIIR",
+ "@nanobind",
+ "@rules_python//python/cc:current_py_cc_headers",
+ ],
+)
+
# These flags are needed for pybind11 to work.
PYBIND11_COPTS = [
"-fexceptions",
@@ -993,16 +1014,25 @@ filegroup(
],
)
+filegroup(
+ name = "MLIRBindingsPythonCoreHeaders",
+ srcs = glob([
+ "lib/Bindings/Python/*.h",
+ ]),
+)
+
cc_library(
name = "MLIRBindingsPythonCore",
srcs = [":MLIRBindingsPythonSourceFiles"],
copts = PYBIND11_COPTS,
features = PYBIND11_FEATURES,
+ textual_hdrs = [":MLIRBindingsPythonCoreHeaders"],
deps = [
":CAPIAsync",
":CAPIDebug",
":CAPIIR",
":CAPIInterfaces",
+ ":CAPITransforms",
":MLIRBindingsPythonHeadersAndDeps",
":Support",
":config",
@@ -1017,10 +1047,12 @@ cc_library(
srcs = [":MLIRBindingsPythonSourceFiles"],
copts = PYBIND11_COPTS,
features = PYBIND11_FEATURES,
+ textual_hdrs = [":MLIRBindingsPythonCoreHeaders"],
deps = [
":CAPIAsyncHeaders",
":CAPIDebugHeaders",
":CAPIIRHeaders",
+ ":CAPITransformsHeaders",
":MLIRBindingsPythonHeaders",
":Support",
":config",
@@ -1050,6 +1082,9 @@ cc_binary(
# These flags are needed for pybind11 to work.
copts = PYBIND11_COPTS,
features = PYBIND11_FEATURES,
+ includes = [
+ "lib/Bindings/Python",
+ ],
linkshared = 1,
linkstatic = 0,
deps = [
@@ -1063,6 +1098,9 @@ cc_binary(
srcs = ["lib/Bindings/Python/DialectLinalg.cpp"],
copts = PYBIND11_COPTS,
features = PYBIND11_FEATURES,
+ includes = [
+ "lib/Bindings/Python",
+ ],
linkshared = 1,
linkstatic = 0,
deps = [
@@ -8448,9 +8486,9 @@ cc_library(
hdrs = ["include/mlir/Conversion/ConvertToLLVM/ToLLVMPass.h"],
includes = ["include"],
deps = [
+ ":Analysis",
":ConversionPassIncGen",
":ConvertToLLVMInterface",
- ":Analysis",
":IR",
":LLVMCommonConversion",
":LLVMDialect",
diff --git a/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
index 254cab0db4a5d6..9dcb6f86ba7f49 100644
--- a/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
@@ -765,7 +765,8 @@ gentbl_filegroup(
filegroup(
name = "PythonTestPyFiles",
srcs = [
- "mlir/dialects/python_test.py",
+ "mlir/dialects/python_test_nanobind.py",
+ "mlir/dialects/python_test_pybind11.py",
":PythonTestPyGen",
],
)
diff --git a/utils/bazel/third_party_build/nanobind.BUILD b/utils/bazel/third_party_build/nanobind.BUILD
new file mode 100644
index 00000000000000..262d14a040b87e
--- /dev/null
+++ b/utils/bazel/third_party_build/nanobind.BUILD
@@ -0,0 +1,25 @@
+cc_library(
+ name = "nanobind",
+ srcs = glob(
+ [
+ "src/*.cpp",
+ ],
+ exclude = ["src/nb_combined.cpp"],
+ ),
+ defines = [
+ "NB_BUILD=1",
+ "NB_SHARED=1",
+ ],
+ includes = ["include"],
+ textual_hdrs = glob(
+ [
+ "include/**/*.h",
+ "src/*.h",
+ ],
+ ),
+ visibility = ["//visibility:public"],
+ deps = [
+ "@robin_map",
+ "@rules_python//python/cc:current_py_cc_headers",
+ ],
+)
diff --git a/utils/bazel/third_party_build/robin_map.BUILD b/utils/bazel/third_party_build/robin_map.BUILD
new file mode 100644
index 00000000000000..b8d04beaed81f9
--- /dev/null
+++ b/utils/bazel/third_party_build/robin_map.BUILD
@@ -0,0 +1,12 @@
+cc_library(
+ name = "robin_map",
+ hdrs = [
+ "include/tsl/robin_growth_policy.h",
+ "include/tsl/robin_hash.h",
+ "include/tsl/robin_map.h",
+ "include/tsl/robin_set.h",
+ ],
+ includes = ["."],
+ strip_include_prefix = "include",
+ visibility = ["//visibility:public"],
+)
More information about the Mlir-commits
mailing list