mirror of
https://github.com/intel/llvm.git
synced 2026-01-27 06:06:34 +08:00
[mlir python] Port Python core code to nanobind. (#118583)
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 (https://github.com/wjakob/nanobind/pull/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.
This commit is contained in:
@@ -448,6 +448,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
|
||||
MODULE_NAME _mlir
|
||||
ADD_TO_PARENT MLIRPythonSources.Core
|
||||
ROOT_DIR "${PYTHON_SOURCE_DIR}"
|
||||
PYTHON_BINDINGS_LIBRARY nanobind
|
||||
SOURCES
|
||||
MainModule.cpp
|
||||
IRAffine.cpp
|
||||
@@ -463,7 +464,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
|
||||
Globals.h
|
||||
IRModule.h
|
||||
Pass.h
|
||||
PybindUtils.h
|
||||
NanobindUtils.h
|
||||
Rewrite.h
|
||||
PRIVATE_LINK_LIBS
|
||||
LLVMSupport
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
nanobind>=2.0, <3.0
|
||||
nanobind>=2.4, <3.0
|
||||
numpy>=1.19.5, <=2.1.2
|
||||
pybind11>=2.10.0, <=2.13.6
|
||||
PyYAML>=5.4.0, <=6.0.1
|
||||
|
||||
Reference in New Issue
Block a user