Skip to content

Commit

Permalink
[mlir python] Port Python core code to nanobind.
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@hawkinsp
hawkinsp committed Dec 4, 2024
1 parent 5d8eabc commit 522a30d
Show file tree
Hide file tree
Showing 18 changed files with 1,838 additions and 1,576 deletions.
2 changes: 1 addition & 1 deletion mlir/include/mlir/Bindings/Python/IRTypes.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {

Expand Down
49 changes: 17 additions & 32 deletions mlir/include/mlir/Bindings/Python/PybindAdaptors.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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()) {
Expand All @@ -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);
Expand All @@ -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()) {
Expand All @@ -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);
Expand All @@ -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) {
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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);
Expand Down Expand Up @@ -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;
}
Expand All @@ -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;
Expand Down
32 changes: 16 additions & 16 deletions mlir/lib/Bindings/Python/Globals.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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;
Expand Down
Loading

0 comments on commit 522a30d

Please sign in to comment.