[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
(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.

Author: hawkinsp

mlir/cmake/modules/MLIRDetectPythonEnv.cmake

@@ -39,7 +39,7 @@ macro(mlir_configure_python_dev_packages)
                   "extension = '${PYTHON_MODULE_EXTENSION}")
 
     mlir_detect_nanobind_install()
-    find_package(nanobind 2.2 CONFIG REQUIRED)
+    find_package(nanobind 2.4 CONFIG REQUIRED)
     message(STATUS "Found nanobind v${nanobind_VERSION}: ${nanobind_INCLUDE_DIR}")
     message(STATUS "Python prefix = '${PYTHON_MODULE_PREFIX}', "
                   "suffix = '${PYTHON_MODULE_SUFFIX}', "

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 {
 

mlir/include/mlir/Bindings/Python/PybindAdaptors.h

@@ -374,9 +374,8 @@ class pure_subclass {
     static_assert(!std::is_member_function_pointer<Func>::value,
                   "def_staticmethod(...) called with a non-static member "
                   "function pointer");
-    py::cpp_function cf(
-        std::forward<Func>(f), py::name(name), py::scope(thisClass),
-        py::sibling(py::getattr(thisClass, name, py::none())), extra...);
+    py::cpp_function cf(std::forward<Func>(f), py::name(name),
+                        py::scope(thisClass), extra...);
     thisClass.attr(cf.name()) = py::staticmethod(cf);
     return *this;
   }
@@ -387,9 +386,8 @@ class pure_subclass {
     static_assert(!std::is_member_function_pointer<Func>::value,
                   "def_classmethod(...) called with a non-static member "
                   "function pointer");
-    py::cpp_function cf(
-        std::forward<Func>(f), py::name(name), py::scope(thisClass),
-        py::sibling(py::getattr(thisClass, name, py::none())), extra...);
+    py::cpp_function cf(std::forward<Func>(f), py::name(name),
+                        py::scope(thisClass), extra...);
     thisClass.attr(cf.name()) =
         py::reinterpret_borrow<py::object>(PyClassMethod_New(cf.ptr()));
     return *this;

mlir/lib/Bindings/Python/Globals.h

@@ -9,18 +9,17 @@
 #ifndef MLIR_BINDINGS_PYTHON_GLOBALS_H
 #define MLIR_BINDINGS_PYTHON_GLOBALS_H
 
-#include "PybindUtils.h"
+#include <optional>
+#include <string>
+#include <vector>
 
+#include "NanobindUtils.h"
 #include "mlir-c/IR.h"
 #include "mlir/CAPI/Support.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSet.h"
 
-#include <optional>
-#include <string>
-#include <vector>
-
 namespace mlir {
 namespace python {
 
@@ -57,71 +56,71 @@ class PyGlobals {
   /// Raises an exception if the mapping already exists and replace == false.
   /// This is intended to be called by implementation code.
   void registerAttributeBuilder(const std::string &attributeKind,
-                                pybind11::function pyFunc,
+                                nanobind::callable pyFunc,
                                 bool replace = false);
 
   /// Adds a user-friendly type caster. Raises an exception if the mapping
   /// already exists and replace == false. This is intended to be called by
   /// implementation code.
-  void registerTypeCaster(MlirTypeID mlirTypeID, pybind11::function typeCaster,
+  void registerTypeCaster(MlirTypeID mlirTypeID, nanobind::callable typeCaster,
                           bool replace = false);
 
   /// Adds a user-friendly value caster. Raises an exception if the mapping
   /// already exists and replace == false. This is intended to be called by
   /// implementation code.
   void registerValueCaster(MlirTypeID mlirTypeID,
-                           pybind11::function valueCaster,
+                           nanobind::callable valueCaster,
                            bool replace = false);
 
   /// Adds a concrete implementation dialect class.
   /// Raises an exception if the mapping already exists.
   /// This is intended to be called by implementation code.
   void registerDialectImpl(const std::string &dialectNamespace,
-                           pybind11::object pyClass);
+                           nanobind::object pyClass);
 
   /// Adds a concrete implementation operation class.
   /// Raises an exception if the mapping already exists and replace == false.
   /// This is intended to be called by implementation code.
   void registerOperationImpl(const std::string &operationName,
-                             pybind11::object pyClass, bool replace = false);
+                             nanobind::object pyClass, bool replace = false);
 
   /// Returns the custom Attribute builder for Attribute kind.
-  std::optional<pybind11::function>
+  std::optional<nanobind::callable>
   lookupAttributeBuilder(const std::string &attributeKind);
 
   /// Returns the custom type caster for MlirTypeID mlirTypeID.
-  std::optional<pybind11::function> lookupTypeCaster(MlirTypeID mlirTypeID,
+  std::optional<nanobind::callable> lookupTypeCaster(MlirTypeID mlirTypeID,
                                                      MlirDialect dialect);
 
   /// Returns the custom value caster for MlirTypeID mlirTypeID.
-  std::optional<pybind11::function> lookupValueCaster(MlirTypeID mlirTypeID,
+  std::optional<nanobind::callable> lookupValueCaster(MlirTypeID mlirTypeID,
                                                       MlirDialect dialect);
 
   /// Looks up a registered dialect class by namespace. Note that this may
   /// trigger loading of the defining module and can arbitrarily re-enter.
-  std::optional<pybind11::object>
+  std::optional<nanobind::object>
   lookupDialectClass(const std::string &dialectNamespace);
 
   /// Looks up a registered operation class (deriving from OpView) by operation
   /// name. Note that this may trigger a load of the dialect, which can
   /// arbitrarily re-enter.
-  std::optional<pybind11::object>
+  std::optional<nanobind::object>
   lookupOperationClass(llvm::StringRef operationName);
 
 private:
   static PyGlobals *instance;
   /// Module name prefixes to search under for dialect implementation modules.
   std::vector<std::string> dialectSearchPrefixes;
   /// Map of dialect namespace to external dialect class object.
-  llvm::StringMap<pybind11::object> dialectClassMap;
+  llvm::StringMap<nanobind::object> dialectClassMap;
   /// Map of full operation name to external operation class object.
-  llvm::StringMap<pybind11::object> operationClassMap;
+  llvm::StringMap<nanobind::object> operationClassMap;
   /// Map of attribute ODS name to custom builder.
-  llvm::StringMap<pybind11::object> attributeBuilderMap;
+  llvm::StringMap<nanobind::callable> attributeBuilderMap;
   /// Map of MlirTypeID to custom type caster.
-  llvm::DenseMap<MlirTypeID, pybind11::object> typeCasterMap;
+  llvm::DenseMap<MlirTypeID, nanobind::callable> typeCasterMap;
   /// Map of MlirTypeID to custom value caster.
-  llvm::DenseMap<MlirTypeID, pybind11::object> valueCasterMap;
+  llvm::DenseMap<MlirTypeID, nanobind::callable> valueCasterMap;
   /// Set of dialect namespaces that we have attempted to import implementation
   /// modules for.
   llvm::StringSet<> loadedDialectModules;