From 5d8eabc4e2fd55a8ea3731c5ba31d39c16a7d7f9 Mon Sep 17 00:00:00 2001 From: Peter Hawkins 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} | 4 +- utils/bazel/WORKSPACE | 18 + .../llvm-project-overlay/mlir/BUILD.bazel | 50 +- utils/bazel/third_party_build/nanobind.BUILD | 25 + utils/bazel/third_party_build/robin_map.BUILD | 12 + 25 files changed, 1184 insertions(+), 105 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} (96%) 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.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 + +#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 +#include + +#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(data) += + llvm::StringRef(message.data, message.length); + }; + MlirLocation loc = mlirDiagnosticGetLocation(diag); + *static_cast(data) += "at "; + mlirLocationPrint(loc, printer, data); + *static_cast(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 +#include + +#include + +#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(apiObject); + if (!nanobind::hasattr(apiObject, MLIR_PYTHON_CAPI_PTR_ATTR)) { + std::string repr = nanobind::cast(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 { + 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(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 { + 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(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 { + 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 { + 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 { + 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( + 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 { + 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(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 { + 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(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 { + 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( + 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 { + 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(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 { + 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(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 { + 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 { + 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(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 { + 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(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((PyObject *)&PyType_Type); + nanobind::object metaclass = pyType(superClass); + nanobind::dict attributes; + + thisClass = metaclass(derivedClassName, nanobind::make_tuple(superClass), + attributes); + scope.attr(derivedClassName) = thisClass; + } + + template + pure_subclass &def(const char *name, Func &&f, const Extra &...extra) { + nanobind::object cf = nanobind::cpp_function( + std::forward(f), nanobind::name(name), nanobind::is_method(), + nanobind::scope(thisClass), extra...); + thisClass.attr(name) = cf; + return *this; + } + + template + pure_subclass &def_property_readonly(const char *name, Func &&f, + const Extra &...extra) { + nanobind::object cf = nanobind::cpp_function( + std::forward(f), nanobind::name(name), nanobind::is_method(), + nanobind::scope(thisClass), extra...); + auto builtinProperty = + nanobind::borrow((PyObject *)&PyProperty_Type); + thisClass.attr(name) = builtinProperty(cf); + return *this; + } + + template + pure_subclass &def_staticmethod(const char *name, Func &&f, + const Extra &...extra) { + static_assert(!std::is_member_function_pointer::value, + "def_staticmethod(...) called with a non-static member " + "function pointer"); + nanobind::object cf = nanobind::cpp_function( + std::forward(f), + nanobind::name(name), // nanobind::scope(thisClass), + extra...); + thisClass.attr(name) = cf; + return *this; + } + + template + pure_subclass &def_classmethod(const char *name, Func &&f, + const Extra &...extra) { + static_assert(!std::is_member_function_pointer::value, + "def_classmethod(...) called with a non-static member " + "function pointer"); + nanobind::object cf = nanobind::cpp_function( + std::forward(f), + nanobind::name(name), // nanobind::scope(thisClass), + extra...); + thisClass.attr(name) = + nanobind::borrow(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(otherAttribute); + if (!isaFunction(rawAttribute)) { + auto origRepr = + nanobind::cast(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(otherType); + if (!isaFunction(rawType)) { + auto origRepr = + nanobind::cast(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(otherValue); + if (!isaFunction(rawValue)) { + auto origRepr = + nanobind::cast(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(data) += + llvm::StringRef(message.data, message.length); + }; + MlirLocation loc = mlirDiagnosticGetLocation(diag); + *static_cast(data) += "at "; + mlirLocationPrint(loc, printer, data); + *static_cast(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(data) += - llvm::StringRef(message.data, message.length); - }; - MlirLocation loc = mlirDiagnosticGetLocation(diag); - *static_cast(data) += "at "; - mlirLocationPrint(loc, printer, data); - *static_cast(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 + #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 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 +#include + #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 -#include - 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..7c504d04be0d13 --- /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 +#include + +#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 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.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 96% rename from mlir/test/python/lib/PythonTestModule.cpp rename to mlir/test/python/lib/PythonTestModulePybind11.cpp index a4f538dcb55944..94a5f5178d16e8 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) { 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/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"], +)