[all-commits] [llvm/llvm-project] 41bd35: [mlir python] Port Python core code to nanobind. (...

[Peter Hawkins via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 41bd35b58bb482fd466aa4b13aa44a810ad6470f
      https://github.com/llvm/llvm-project/commit/41bd35b58bb482fd466aa4b13aa44a810ad6470f
  Author: Peter Hawkins <phawkins at google.com>
  Date:   2024-12-18 (Wed, 18 Dec 2024)

  Changed paths:
    M mlir/cmake/modules/MLIRDetectPythonEnv.cmake
    M mlir/include/mlir/Bindings/Python/IRTypes.h
    M mlir/include/mlir/Bindings/Python/PybindAdaptors.h
    M mlir/lib/Bindings/Python/Globals.h
    M mlir/lib/Bindings/Python/IRAffine.cpp
    M mlir/lib/Bindings/Python/IRAttributes.cpp
    M mlir/lib/Bindings/Python/IRCore.cpp
    M mlir/lib/Bindings/Python/IRInterfaces.cpp
    M mlir/lib/Bindings/Python/IRModule.cpp
    M mlir/lib/Bindings/Python/IRModule.h
    M mlir/lib/Bindings/Python/IRTypes.cpp
    M mlir/lib/Bindings/Python/MainModule.cpp
    A mlir/lib/Bindings/Python/NanobindUtils.h
    M mlir/lib/Bindings/Python/Pass.cpp
    M mlir/lib/Bindings/Python/Pass.h
    R mlir/lib/Bindings/Python/PybindUtils.h
    M mlir/lib/Bindings/Python/Rewrite.cpp
    M mlir/lib/Bindings/Python/Rewrite.h
    M mlir/python/CMakeLists.txt
    M mlir/python/requirements.txt
    M mlir/test/python/ir/symbol_table.py
    M utils/bazel/WORKSPACE
    M utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

  Log Message:
  -----------
  [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.



To unsubscribe from these emails, change your notification settings at https://github.com/llvm/llvm-project/settings/notifications