[mlir python] Port Python core code to nanobind. #118583
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@llvm/pr-subscribers-mlir Author: Peter Hawkins (hawkinsp) ChangesThis is a draft PR, which will require an RFC to land. Why? https://nanobind.readthedocs.io/en/latest/why.html says it better For a complicated Google-internal LLM model in JAX, this change improves the MLIR To a large extent, this is a mechanical change, for instance changing Notes:
Patch is 423.50 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/118583.diff 41 Files Affected:
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/IRTypes.h b/mlir/include/mlir/Bindings/Python/IRTypes.h
index 9afad4c23b3f35..ba9642cf2c6a2d 100644
--- a/mlir/include/mlir/Bindings/Python/IRTypes.h
+++ b/mlir/include/mlir/Bindings/Python/IRTypes.h
@@ -9,7 +9,7 @@
#ifndef MLIR_BINDINGS_PYTHON_IRTYPES_H
#define MLIR_BINDINGS_PYTHON_IRTYPES_H
-#include "mlir/Bindings/Python/PybindAdaptors.h"
+#include "mlir/Bindings/Python/NanobindAdaptors.h"
namespace mlir {
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/Twi...
[truncated]
|
|
(Note this builds on top of the reverted #117922, so that would need to be reapplied first) |
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
|
Nice! Maks or Peter, I think the PR description could basically just be copy/pasted as an RFC. I would probably add a section for expected steps needed by downstreams (which I think is "none but you may want to port your dialects and eliminate your pybind dep once things land") and rollout plan (which could just be "we will let the final PR sit for at least five days, notify here, and await/incorporate feedback from downstream testing during that period"). I don't anticipate major hurdles but I do know that upstream tests are not comprehensive for all corners of the API or how it is integrated in the wild. Adapting to that is up to the downstreams, but the quality will be better by having a testing period. |
|
The required nanobind change was released as part of nanobind 2.4.0, so this PR now pins nanobind 2.4. |
aa9b9fb to
2aa2dd2
Compare
marbre
added a commit
to marbre/iree
that referenced
this pull request
Dec 6, 2024
Initalizes nanobind at the top-level as MLIR is switching over to nanobind with llvm/llvm-project#118583.
marbre
added a commit
to marbre/iree
that referenced
this pull request
Dec 9, 2024
Initalizes nanobind at the top-level as MLIR is switching over to nanobind with llvm/llvm-project#118583.
marbre
added a commit
to iree-org/iree
that referenced
this pull request
Dec 9, 2024
Initalizes nanobind at the top-level as MLIR is switching over to nanobind with llvm/llvm-project#118583.
|
Rebased on top of main, now #117922 has been relanded. |
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Dec 12, 2024
I'm working towards moving the MLIR Python core code to use nanobind instead of pybind11: * llvm/llvm-project#117922, which was merged recently, allows downstream Python dialect extensions to be defined using either pybind11 or nanobind. * llvm/llvm-project#118583 is a PR in review that ports the Python core code to use nanobind instead of pybind11. This PR migrates StableHLO and related dialects to use nanobind rather than pybind11, with the goal of migrating JAX away from pybind11. PiperOrigin-RevId: 705197067
copybara-service bot
pushed a commit
to openxla/xla
that referenced
this pull request
Dec 17, 2024
I'm working towards moving the MLIR Python core code to use nanobind instead of pybind11: * llvm/llvm-project#117922, which was merged recently, allows downstream Python dialect extensions to be defined using either pybind11 or nanobind. * llvm/llvm-project#118583 is a PR in review that ports the Python core code to use nanobind instead of pybind11. This PR migrates StableHLO and related dialects to use nanobind rather than pybind11, with the goal of migrating JAX away from pybind11. PiperOrigin-RevId: 705197067
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 17, 2024
I'm working towards moving the MLIR Python core code to use nanobind instead of pybind11: * llvm/llvm-project#117922, which was merged recently, allows downstream Python dialect extensions to be defined using either pybind11 or nanobind. * llvm/llvm-project#118583 is a PR in review that ports the Python core code to use nanobind instead of pybind11. This PR migrates StableHLO and related dialects to use nanobind rather than pybind11, with the goal of migrating JAX away from pybind11. PiperOrigin-RevId: 705197067
copybara-service bot
pushed a commit
to openxla/xla
that referenced
this pull request
Dec 18, 2024
I'm working towards moving the MLIR Python core code to use nanobind instead of pybind11: * llvm/llvm-project#117922, which was merged recently, allows downstream Python dialect extensions to be defined using either pybind11 or nanobind. * llvm/llvm-project#118583 is a PR in review that ports the Python core code to use nanobind instead of pybind11. This PR migrates StableHLO and related dialects to use nanobind rather than pybind11, with the goal of migrating JAX away from pybind11. PiperOrigin-RevId: 705197067
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 18, 2024
I'm working towards moving the MLIR Python core code to use nanobind instead of pybind11: * llvm/llvm-project#117922, which was merged recently, allows downstream Python dialect extensions to be defined using either pybind11 or nanobind. * llvm/llvm-project#118583 is a PR in review that ports the Python core code to use nanobind instead of pybind11. This PR migrates StableHLO and related dialects to use nanobind rather than pybind11, with the goal of migrating JAX away from pybind11. PiperOrigin-RevId: 705197067
copybara-service bot
pushed a commit
to openxla/xla
that referenced
this pull request
Dec 18, 2024
I'm working towards moving the MLIR Python core code to use nanobind instead of pybind11: * llvm/llvm-project#117922, which was merged recently, allows downstream Python dialect extensions to be defined using either pybind11 or nanobind. * llvm/llvm-project#118583 is a PR in review that ports the Python core code to use nanobind instead of pybind11. This PR migrates StableHLO and related dialects to use nanobind rather than pybind11, with the goal of migrating JAX away from pybind11. PiperOrigin-RevId: 705197067
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 18, 2024
I'm working towards moving the MLIR Python core code to use nanobind instead of pybind11: * llvm/llvm-project#117922, which was merged recently, allows downstream Python dialect extensions to be defined using either pybind11 or nanobind. * llvm/llvm-project#118583 is a PR in review that ports the Python core code to use nanobind instead of pybind11. This PR migrates StableHLO and related dialects to use nanobind rather than pybind11, with the goal of migrating JAX away from pybind11. PiperOrigin-RevId: 705197067
copybara-service bot
pushed a commit
to openxla/xla
that referenced
this pull request
Dec 18, 2024
I'm working towards moving the MLIR Python core code to use nanobind instead of pybind11: * llvm/llvm-project#117922, which was merged recently, allows downstream Python dialect extensions to be defined using either pybind11 or nanobind. * llvm/llvm-project#118583 is a PR in review that ports the Python core code to use nanobind instead of pybind11. This PR migrates StableHLO and related dialects to use nanobind rather than pybind11, with the goal of migrating JAX away from pybind11. PiperOrigin-RevId: 705197067
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 18, 2024
I'm working towards moving the MLIR Python core code to use nanobind instead of pybind11: * llvm/llvm-project#117922, which was merged recently, allows downstream Python dialect extensions to be defined using either pybind11 or nanobind. * llvm/llvm-project#118583 is a PR in review that ports the Python core code to use nanobind instead of pybind11. This PR migrates StableHLO and related dialects to use nanobind rather than pybind11, with the goal of migrating JAX away from pybind11. PiperOrigin-RevId: 705197067
copybara-service bot
pushed a commit
to tensorflow/mlir-hlo
that referenced
this pull request
Dec 18, 2024
I'm working towards moving the MLIR Python core code to use nanobind instead of pybind11: * llvm/llvm-project#117922, which was merged recently, allows downstream Python dialect extensions to be defined using either pybind11 or nanobind. * llvm/llvm-project#118583 is a PR in review that ports the Python core code to use nanobind instead of pybind11. This PR migrates StableHLO and related dialects to use nanobind rather than pybind11, with the goal of migrating JAX away from pybind11. PiperOrigin-RevId: 707537037
copybara-service bot
pushed a commit
to openxla/shardy
that referenced
this pull request
Dec 18, 2024
I'm working towards moving the MLIR Python core code to use nanobind instead of pybind11: * llvm/llvm-project#117922, which was merged recently, allows downstream Python dialect extensions to be defined using either pybind11 or nanobind. * llvm/llvm-project#118583 is a PR in review that ports the Python core code to use nanobind instead of pybind11. This PR migrates StableHLO and related dialects to use nanobind rather than pybind11, with the goal of migrating JAX away from pybind11. PiperOrigin-RevId: 707537037
copybara-service bot
pushed a commit
to openxla/xla
that referenced
this pull request
Dec 18, 2024
I'm working towards moving the MLIR Python core code to use nanobind instead of pybind11: * llvm/llvm-project#117922, which was merged recently, allows downstream Python dialect extensions to be defined using either pybind11 or nanobind. * llvm/llvm-project#118583 is a PR in review that ports the Python core code to use nanobind instead of pybind11. This PR migrates StableHLO and related dialects to use nanobind rather than pybind11, with the goal of migrating JAX away from pybind11. PiperOrigin-RevId: 707537037
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Dec 18, 2024
I'm working towards moving the MLIR Python core code to use nanobind instead of pybind11: * llvm/llvm-project#117922, which was merged recently, allows downstream Python dialect extensions to be defined using either pybind11 or nanobind. * llvm/llvm-project#118583 is a PR in review that ports the Python core code to use nanobind instead of pybind11. This PR migrates StableHLO and related dialects to use nanobind rather than pybind11, with the goal of migrating JAX away from pybind11. PiperOrigin-RevId: 707537037
jpienaar
approved these changes
Dec 18, 2024
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/177/builds/10232 Here is the relevant piece of the build log for the reference |
jpienaar
added a commit
that referenced
this pull request
Dec 18, 2024
This reverts commit 41bd35b. Breakage detected, rolling back.
hawkinsp
added a commit
to hawkinsp/llvm-project
that referenced
this pull request
Dec 18, 2024
Relands llvm#118583, with a fix for Python 3.8 compatibility. It was not possible to set the buffer protocol accessers via slots in Python 3.8. Why? https://nanobind.readthedocs.io/en/latest/why.html says it better than I can, but my primary motivation for this change is to improve MLIR IR construction time from JAX. For a complicated Google-internal LLM model in JAX, this change improves the MLIR lowering time by around 5s (out of around 30s), which is a significant speedup for simply switching binding frameworks. To a large extent, this is a mechanical change, for instance changing `pybind11::` to `nanobind::`. Notes: * this PR needs Nanobind 2.4.0, because it needs a bug fix (wjakob/nanobind#806) that landed in that release. * this PR does not port the in-tree dialect extension modules. They can be ported in a future PR. * I removed the py::sibling() annotations from def_static and def_class in `PybindAdapters.h`. These ask pybind11 to try to form an overload with an existing method, but it's not possible to form mixed pybind11/nanobind overloads this ways and the parent class is now defined in nanobind. Better solutions may be possible here. * nanobind does not contain an exact equivalent of pybind11's buffer protocol support. It was not hard to add a nanobind implementation of a similar API. * nanobind is pickier about casting to std::vector<bool>, expecting that the input is a sequence of bool types, not truthy values. In a couple of places I added code to support truthy values during casting. * nanobind distinguishes bytes (`nb::bytes`) from strings (e.g., `std::string`). This required nb::bytes overloads in a few places.
hawkinsp
added a commit
to hawkinsp/llvm-project
that referenced
this pull request
Dec 18, 2024
Relands llvm#118583, with a fix for Python 3.8 compatibility. It was not possible to set the buffer protocol accessers via slots in Python 3.8. Why? https://nanobind.readthedocs.io/en/latest/why.html says it better than I can, but my primary motivation for this change is to improve MLIR IR construction time from JAX. For a complicated Google-internal LLM model in JAX, this change improves the MLIR lowering time by around 5s (out of around 30s), which is a significant speedup for simply switching binding frameworks. To a large extent, this is a mechanical change, for instance changing `pybind11::` to `nanobind::`. Notes: * this PR needs Nanobind 2.4.0, because it needs a bug fix (wjakob/nanobind#806) that landed in that release. * this PR does not port the in-tree dialect extension modules. They can be ported in a future PR. * I removed the py::sibling() annotations from def_static and def_class in `PybindAdapters.h`. These ask pybind11 to try to form an overload with an existing method, but it's not possible to form mixed pybind11/nanobind overloads this ways and the parent class is now defined in nanobind. Better solutions may be possible here. * nanobind does not contain an exact equivalent of pybind11's buffer protocol support. It was not hard to add a nanobind implementation of a similar API. * nanobind is pickier about casting to std::vector<bool>, expecting that the input is a sequence of bool types, not truthy values. In a couple of places I added code to support truthy values during casting. * nanobind distinguishes bytes (`nb::bytes`) from strings (e.g., `std::string`). This required nb::bytes overloads in a few places.
|
Yeah I had replied to @jpienaar on email, and meant to chime in here yesterday. I am starting to look at the CIRCT integrate, but don't expect and problems and am in favor of moving forward. |
|
Thanks @hawkinsp for the heavy lift and others for the pre-work to get such a big change sequence landed smoothly. Much appreciated. Happy new year. |
hawkinsp
added a commit
to hawkinsp/llvm-project
that referenced
this pull request
Dec 19, 2024
Relands llvm#118583, with a fix for Python 3.8 compatibility. It was not possible to set the buffer protocol accessers via slots in Python 3.8. Why? https://nanobind.readthedocs.io/en/latest/why.html says it better than I can, but my primary motivation for this change is to improve MLIR IR construction time from JAX. For a complicated Google-internal LLM model in JAX, this change improves the MLIR lowering time by around 5s (out of around 30s), which is a significant speedup for simply switching binding frameworks. To a large extent, this is a mechanical change, for instance changing `pybind11::` to `nanobind::`. Notes: * this PR needs Nanobind 2.4.0, because it needs a bug fix (wjakob/nanobind#806) that landed in that release. * this PR does not port the in-tree dialect extension modules. They can be ported in a future PR. * I removed the py::sibling() annotations from def_static and def_class in `PybindAdapters.h`. These ask pybind11 to try to form an overload with an existing method, but it's not possible to form mixed pybind11/nanobind overloads this ways and the parent class is now defined in nanobind. Better solutions may be possible here. * nanobind does not contain an exact equivalent of pybind11's buffer protocol support. It was not hard to add a nanobind implementation of a similar API. * nanobind is pickier about casting to std::vector<bool>, expecting that the input is a sequence of bool types, not truthy values. In a couple of places I added code to support truthy values during casting. * nanobind distinguishes bytes (`nb::bytes`) from strings (e.g., `std::string`). This required nb::bytes overloads in a few places.
|
Thanks Stella! For those following along: we reverted this PR because we found it failed to compile on 3.8 from one of the postsubmit build bots. New attempt is in #120473, which fixes Python 3.8 support. |
hawkinsp
added a commit
to hawkinsp/llvm-project
that referenced
this pull request
Dec 19, 2024
Relands llvm#118583, with a fix for Python 3.8 compatibility. It was not possible to set the buffer protocol accessers via slots in Python 3.8. Why? https://nanobind.readthedocs.io/en/latest/why.html says it better than I can, but my primary motivation for this change is to improve MLIR IR construction time from JAX. For a complicated Google-internal LLM model in JAX, this change improves the MLIR lowering time by around 5s (out of around 30s), which is a significant speedup for simply switching binding frameworks. To a large extent, this is a mechanical change, for instance changing `pybind11::` to `nanobind::`. Notes: * this PR needs Nanobind 2.4.0, because it needs a bug fix (wjakob/nanobind#806) that landed in that release. * this PR does not port the in-tree dialect extension modules. They can be ported in a future PR. * I removed the py::sibling() annotations from def_static and def_class in `PybindAdapters.h`. These ask pybind11 to try to form an overload with an existing method, but it's not possible to form mixed pybind11/nanobind overloads this ways and the parent class is now defined in nanobind. Better solutions may be possible here. * nanobind does not contain an exact equivalent of pybind11's buffer protocol support. It was not hard to add a nanobind implementation of a similar API. * nanobind is pickier about casting to std::vector<bool>, expecting that the input is a sequence of bool types, not truthy values. In a couple of places I added code to support truthy values during casting. * nanobind distinguishes bytes (`nb::bytes`) from strings (e.g., `std::string`). This required nb::bytes overloads in a few places.
jpienaar
pushed a commit
that referenced
this pull request
Dec 19, 2024
Relands #118583, with a fix for Python 3.8 compatibility. It was not possible to set the buffer protocol accessers via slots in Python 3.8. Why? https://nanobind.readthedocs.io/en/latest/why.html says it better than I can, but my primary motivation for this change is to improve MLIR IR construction time from JAX. For a complicated Google-internal LLM model in JAX, this change improves the MLIR lowering time by around 5s (out of around 30s), which is a significant speedup for simply switching binding frameworks. To a large extent, this is a mechanical change, for instance changing `pybind11::` to `nanobind::`. Notes: * this PR needs Nanobind 2.4.0, because it needs a bug fix (wjakob/nanobind#806) that landed in that release. * this PR does not port the in-tree dialect extension modules. They can be ported in a future PR. * I removed the py::sibling() annotations from def_static and def_class in `PybindAdapters.h`. These ask pybind11 to try to form an overload with an existing method, but it's not possible to form mixed pybind11/nanobind overloads this ways and the parent class is now defined in nanobind. Better solutions may be possible here. * nanobind does not contain an exact equivalent of pybind11's buffer protocol support. It was not hard to add a nanobind implementation of a similar API. * nanobind is pickier about casting to std::vector<bool>, expecting that the input is a sequence of bool types, not truthy values. In a couple of places I added code to support truthy values during casting. * nanobind distinguishes bytes (`nb::bytes`) from strings (e.g., `std::string`). This required nb::bytes overloads in a few places.
|
Thanks for the extra effort. In the new year, we should see about bumping the min python version now that 3.8 is eol. But always easier to do that explicitly vs inline with a patch like this. |
hawkinsp
added a commit
to hawkinsp/llvm-project
that referenced
this pull request
Dec 19, 2024
This is a companion to llvm#118583, although it can be landed independently because since llvm#117922 dialects do not have to use the same Python binding framework as the Python core code. This PR ports all of the in-tree dialect and pass extensions to nanobind, with the exception of those that remain for testing pybind11 support. It would make sense to merge this PR after merging llvm#118583, if we have agreed that we are migrating the core to nanobind. This PR also: * removes CollectDiagnosticsToStringScope from NanobindAdaptors.h. This was overlooked in a previous PR and it is duplicated in Diagnostics.h.
hawkinsp
added a commit
to hawkinsp/llvm-project
that referenced
this pull request
Dec 20, 2024
This is a companion to llvm#118583, although it can be landed independently because since llvm#117922 dialects do not have to use the same Python binding framework as the Python core code. This PR ports all of the in-tree dialect and pass extensions to nanobind, with the exception of those that remain for testing pybind11 support. It would make sense to merge this PR after merging llvm#118583, if we have agreed that we are migrating the core to nanobind. This PR also: * removes CollectDiagnosticsToStringScope from NanobindAdaptors.h. This was overlooked in a previous PR and it is duplicated in Diagnostics.h.
jpienaar
added a commit
to hawkinsp/llvm-project
that referenced
this pull request
Dec 21, 2024
This is a companion to llvm#118583, although it can be landed independently because since llvm#117922 dialects do not have to use the same Python binding framework as the Python core code. This PR ports all of the in-tree dialect and pass extensions to nanobind, with the exception of those that remain for testing pybind11 support. It would make sense to merge this PR after merging llvm#118583, if we have agreed that we are migrating the core to nanobind. This PR also: * removes CollectDiagnosticsToStringScope from NanobindAdaptors.h. This was overlooked in a previous PR and it is duplicated in Diagnostics.h. --------- Co-authored-by: Jacques Pienaar <jpienaar@google.com>
jpienaar
added a commit
that referenced
this pull request
Dec 21, 2024
This is a companion to #118583, although it can be landed independently because since #117922 dialects do not have to use the same Python binding framework as the Python core code. This PR ports all of the in-tree dialect and pass extensions to nanobind, with the exception of those that remain for testing pybind11 support. This PR also: * removes CollectDiagnosticsToStringScope from NanobindAdaptors.h. This was overlooked in a previous PR and it is duplicated in Diagnostics.h. --------- Co-authored-by: Jacques Pienaar <jpienaar@google.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Why? https://nanobind.readthedocs.io/en/latest/why.html says it better
than I can, but my primary motivation for this change is to improve MLIR
IR construction time from JAX.
For a complicated Google-internal LLM model in JAX, this change improves the MLIR
lowering time by around 5s (out of around 30s), which is a significant
speedup for simply switching binding frameworks.
To a large extent, this is a mechanical change, for instance changing
pybind11::to
nanobind::.Notes:
PybindAdapters.h. These ask pybind11 to try to form an overload with an existing method, but it's not possible to form mixed pybind11/nanobind overloads this ways and the parent class is now defined in nanobind. Better solutions may be possible here.nb::bytes) from strings (e.g.,std::string). This required nb::bytes overloads in a few places.