[mlir python] Port Python core code to nanobind.

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 wjakob/nanobind#806 to land in
  nanobind first. Without that fix, importing the MLIR modules will
  fail.
* 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/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

@@ -64,8 +64,7 @@ static py::object mlirApiObjectToCapsule(py::handle apiObject) {
 // ownership is unclear.
 
 /// Casts object <-> MlirAffineMap.
-template <>
-struct type_caster<MlirAffineMap> {
+template <> struct type_caster<MlirAffineMap> {
   PYBIND11_TYPE_CASTER(MlirAffineMap, _("MlirAffineMap"));
   bool load(handle src, bool) {
     py::object capsule = mlirApiObjectToCapsule(src);
@@ -86,8 +85,7 @@ struct type_caster<MlirAffineMap> {
 };
 
 /// Casts object <-> MlirAttribute.
-template <>
-struct type_caster<MlirAttribute> {
+template <> struct type_caster<MlirAttribute> {
   PYBIND11_TYPE_CASTER(MlirAttribute, _("MlirAttribute"));
   bool load(handle src, bool) {
     py::object capsule = mlirApiObjectToCapsule(src);
@@ -106,8 +104,7 @@ struct type_caster<MlirAttribute> {
 };
 
 /// Casts object -> MlirBlock.
-template <>
-struct type_caster<MlirBlock> {
+template <> struct type_caster<MlirBlock> {
   PYBIND11_TYPE_CASTER(MlirBlock, _("MlirBlock"));
   bool load(handle src, bool) {
     py::object capsule = mlirApiObjectToCapsule(src);
@@ -117,8 +114,7 @@ struct type_caster<MlirBlock> {
 };
 
 /// Casts object -> MlirContext.
-template <>
-struct type_caster<MlirContext> {
+template <> struct type_caster<MlirContext> {
   PYBIND11_TYPE_CASTER(MlirContext, _("MlirContext"));
   bool load(handle src, bool) {
     if (src.is_none()) {
@@ -137,8 +133,7 @@ struct type_caster<MlirContext> {
 };
 
 /// Casts object <-> MlirDialectRegistry.
-template <>
-struct type_caster<MlirDialectRegistry> {
+template <> struct type_caster<MlirDialectRegistry> {
   PYBIND11_TYPE_CASTER(MlirDialectRegistry, _("MlirDialectRegistry"));
   bool load(handle src, bool) {
     py::object capsule = mlirApiObjectToCapsule(src);
@@ -156,8 +151,7 @@ struct type_caster<MlirDialectRegistry> {
 };
 
 /// Casts object <-> MlirLocation.
-template <>
-struct type_caster<MlirLocation> {
+template <> struct type_caster<MlirLocation> {
   PYBIND11_TYPE_CASTER(MlirLocation, _("MlirLocation"));
   bool load(handle src, bool) {
     if (src.is_none()) {
@@ -181,8 +175,7 @@ struct type_caster<MlirLocation> {
 };
 
 /// Casts object <-> MlirModule.
-template <>
-struct type_caster<MlirModule> {
+template <> struct type_caster<MlirModule> {
   PYBIND11_TYPE_CASTER(MlirModule, _("MlirModule"));
   bool load(handle src, bool) {
     py::object capsule = mlirApiObjectToCapsule(src);
@@ -200,8 +193,7 @@ struct type_caster<MlirModule> {
 };
 
 /// Casts object <-> MlirFrozenRewritePatternSet.
-template <>
-struct type_caster<MlirFrozenRewritePatternSet> {
+template <> struct type_caster<MlirFrozenRewritePatternSet> {
   PYBIND11_TYPE_CASTER(MlirFrozenRewritePatternSet,
                        _("MlirFrozenRewritePatternSet"));
   bool load(handle src, bool) {
@@ -221,8 +213,7 @@ struct type_caster<MlirFrozenRewritePatternSet> {
 };
 
 /// Casts object <-> MlirOperation.
-template <>
-struct type_caster<MlirOperation> {
+template <> struct type_caster<MlirOperation> {
   PYBIND11_TYPE_CASTER(MlirOperation, _("MlirOperation"));
   bool load(handle src, bool) {
     py::object capsule = mlirApiObjectToCapsule(src);
@@ -242,8 +233,7 @@ struct type_caster<MlirOperation> {
 };
 
 /// Casts object <-> MlirValue.
-template <>
-struct type_caster<MlirValue> {
+template <> struct type_caster<MlirValue> {
   PYBIND11_TYPE_CASTER(MlirValue, _("MlirValue"));
   bool load(handle src, bool) {
     py::object capsule = mlirApiObjectToCapsule(src);
@@ -264,8 +254,7 @@ struct type_caster<MlirValue> {
 };
 
 /// Casts object -> MlirPassManager.
-template <>
-struct type_caster<MlirPassManager> {
+template <> struct type_caster<MlirPassManager> {
   PYBIND11_TYPE_CASTER(MlirPassManager, _("MlirPassManager"));
   bool load(handle src, bool) {
     py::object capsule = mlirApiObjectToCapsule(src);
@@ -275,8 +264,7 @@ struct type_caster<MlirPassManager> {
 };
 
 /// Casts object <-> MlirTypeID.
-template <>
-struct type_caster<MlirTypeID> {
+template <> struct type_caster<MlirTypeID> {
   PYBIND11_TYPE_CASTER(MlirTypeID, _("MlirTypeID"));
   bool load(handle src, bool) {
     py::object capsule = mlirApiObjectToCapsule(src);
@@ -296,8 +284,7 @@ struct type_caster<MlirTypeID> {
 };
 
 /// Casts object <-> MlirType.
-template <>
-struct type_caster<MlirType> {
+template <> struct type_caster<MlirType> {
   PYBIND11_TYPE_CASTER(MlirType, _("MlirType"));
   bool load(handle src, bool) {
     py::object capsule = mlirApiObjectToCapsule(src);
@@ -374,9 +361,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 +373,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,7 +9,7 @@
 #ifndef MLIR_BINDINGS_PYTHON_GLOBALS_H
 #define MLIR_BINDINGS_PYTHON_GLOBALS_H
 
-#include "PybindUtils.h"
+#include "NanobindUtils.h"
 
 #include "mlir-c/IR.h"
 #include "mlir/CAPI/Support.h"
@@ -57,71 +57,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;