[llvm] [mlir] [mlir python] Port in-tree dialects to nanobind. (PR #119924)
via llvm-commits
llvm-commits at lists.llvm.org
Fri Dec 13 12:52:44 PST 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Peter Hawkins (hawkinsp)
<details>
<summary>Changes</summary>
This is a companion to #<!-- -->118583, although it can be landed independently because since #<!-- -->117922 dialects do not have to use the same Python binding framework as the Python core code.
This PR ports all of the in-tree dialect and pass extensions to nanobind, with the exception of those that remain for testing pybind11 support. It would make sense to merge this PR after merging #<!-- -->118583, if we have agreed that we are migrating the core to nanobind.
This PR also:
* removes CollectDiagnosticsToStringScope from NanobindAdaptors.h. This was overlooked in a previous PR and it is duplicated in Diagnostics.h.
---
Patch is 63.07 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/119924.diff
19 Files Affected:
- (modified) mlir/include/mlir/Bindings/Python/NanobindAdaptors.h (-34)
- (modified) mlir/lib/Bindings/Python/AsyncPasses.cpp (+2-3)
- (modified) mlir/lib/Bindings/Python/DialectGPU.cpp (+24-22)
- (modified) mlir/lib/Bindings/Python/DialectLLVM.cpp (+33-25)
- (modified) mlir/lib/Bindings/Python/DialectLinalg.cpp (+7-5)
- (modified) mlir/lib/Bindings/Python/DialectNVGPU.cpp (+11-10)
- (modified) mlir/lib/Bindings/Python/DialectPDL.cpp (+21-23)
- (modified) mlir/lib/Bindings/Python/DialectQuant.cpp (+37-38)
- (modified) mlir/lib/Bindings/Python/DialectSparseTensor.cpp (+21-21)
- (modified) mlir/lib/Bindings/Python/DialectTransform.cpp (+24-25)
- (modified) mlir/lib/Bindings/Python/ExecutionEngineModule.cpp (+46-42)
- (modified) mlir/lib/Bindings/Python/GPUPasses.cpp (+2-3)
- (modified) mlir/lib/Bindings/Python/LinalgPasses.cpp (+2-2)
- (modified) mlir/lib/Bindings/Python/RegisterEverything.cpp (+3-2)
- (modified) mlir/lib/Bindings/Python/SparseTensorPasses.cpp (+2-2)
- (modified) mlir/lib/Bindings/Python/TransformInterpreter.cpp (+22-21)
- (modified) mlir/python/CMakeLists.txt (+15)
- (modified) mlir/test/python/dialects/sparse_tensor/dialect.py (+1-1)
- (modified) utils/bazel/llvm-project-overlay/mlir/BUILD.bazel (+12-11)
``````````diff
diff --git a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
index 5e01cebcb09c91..71b12b4dfc637f 100644
--- a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
+++ b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
@@ -631,40 +631,6 @@ class mlir_value_subclass : public pure_subclass {
} // namespace nanobind_adaptors
-/// RAII scope intercepting all diagnostics into a string. The message must be
-/// checked before this goes out of scope.
-class CollectDiagnosticsToStringScope {
-public:
- explicit CollectDiagnosticsToStringScope(MlirContext ctx) : context(ctx) {
- handlerID = mlirContextAttachDiagnosticHandler(ctx, &handler, &errorMessage,
- /*deleteUserData=*/nullptr);
- }
- ~CollectDiagnosticsToStringScope() {
- assert(errorMessage.empty() && "unchecked error message");
- mlirContextDetachDiagnosticHandler(context, handlerID);
- }
-
- [[nodiscard]] std::string takeMessage() { return std::move(errorMessage); }
-
-private:
- static MlirLogicalResult handler(MlirDiagnostic diag, void *data) {
- auto printer = +[](MlirStringRef message, void *data) {
- *static_cast<std::string *>(data) +=
- llvm::StringRef(message.data, message.length);
- };
- MlirLocation loc = mlirDiagnosticGetLocation(diag);
- *static_cast<std::string *>(data) += "at ";
- mlirLocationPrint(loc, printer, data);
- *static_cast<std::string *>(data) += ": ";
- mlirDiagnosticPrint(diag, printer, data);
- return mlirLogicalResultSuccess();
- }
-
- MlirContext context;
- MlirDiagnosticHandlerID handlerID;
- std::string errorMessage = "";
-};
-
} // namespace python
} // namespace mlir
diff --git a/mlir/lib/Bindings/Python/AsyncPasses.cpp b/mlir/lib/Bindings/Python/AsyncPasses.cpp
index b611a758dbbb37..540fbf5bbe7290 100644
--- a/mlir/lib/Bindings/Python/AsyncPasses.cpp
+++ b/mlir/lib/Bindings/Python/AsyncPasses.cpp
@@ -8,14 +8,13 @@
#include "mlir-c/Dialect/Async.h"
-#include <pybind11/detail/common.h>
-#include <pybind11/pybind11.h>
+#include <nanobind/nanobind.h>
// -----------------------------------------------------------------------------
// Module initialization.
// -----------------------------------------------------------------------------
-PYBIND11_MODULE(_mlirAsyncPasses, m) {
+NB_MODULE(_mlirAsyncPasses, m) {
m.doc() = "MLIR Async Dialect Passes";
// Register all Async passes on load.
diff --git a/mlir/lib/Bindings/Python/DialectGPU.cpp b/mlir/lib/Bindings/Python/DialectGPU.cpp
index 560a54bcd15919..f0154f1b110bdc 100644
--- a/mlir/lib/Bindings/Python/DialectGPU.cpp
+++ b/mlir/lib/Bindings/Python/DialectGPU.cpp
@@ -9,21 +9,23 @@
#include "mlir-c/Dialect/GPU.h"
#include "mlir-c/IR.h"
#include "mlir-c/Support.h"
-#include "mlir/Bindings/Python/PybindAdaptors.h"
+#include "mlir/Bindings/Python/NanobindAdaptors.h"
-#include <pybind11/detail/common.h>
-#include <pybind11/pybind11.h>
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/optional.h>
+
+namespace nb = nanobind;
+using namespace nanobind::literals;
-namespace py = pybind11;
using namespace mlir;
using namespace mlir::python;
-using namespace mlir::python::adaptors;
+using namespace mlir::python::nanobind_adaptors;
// -----------------------------------------------------------------------------
// Module initialization.
// -----------------------------------------------------------------------------
-PYBIND11_MODULE(_mlirDialectsGPU, m) {
+NB_MODULE(_mlirDialectsGPU, m) {
m.doc() = "MLIR GPU Dialect";
//===-------------------------------------------------------------------===//
// AsyncTokenType
@@ -34,11 +36,11 @@ PYBIND11_MODULE(_mlirDialectsGPU, m) {
mlirGPUAsyncTokenType.def_classmethod(
"get",
- [](py::object cls, MlirContext ctx) {
+ [](nb::object cls, MlirContext ctx) {
return cls(mlirGPUAsyncTokenTypeGet(ctx));
},
- "Gets an instance of AsyncTokenType in the same context", py::arg("cls"),
- py::arg("ctx") = py::none());
+ "Gets an instance of AsyncTokenType in the same context", nb::arg("cls"),
+ nb::arg("ctx").none() = nb::none());
//===-------------------------------------------------------------------===//
// ObjectAttr
@@ -47,12 +49,12 @@ PYBIND11_MODULE(_mlirDialectsGPU, m) {
mlir_attribute_subclass(m, "ObjectAttr", mlirAttributeIsAGPUObjectAttr)
.def_classmethod(
"get",
- [](py::object cls, MlirAttribute target, uint32_t format,
- py::bytes object, std::optional<MlirAttribute> mlirObjectProps,
+ [](nb::object cls, MlirAttribute target, uint32_t format,
+ nb::bytes object, std::optional<MlirAttribute> mlirObjectProps,
std::optional<MlirAttribute> mlirKernelsAttr) {
- py::buffer_info info(py::buffer(object).request());
- MlirStringRef objectStrRef =
- mlirStringRefCreate(static_cast<char *>(info.ptr), info.size);
+ MlirStringRef objectStrRef = mlirStringRefCreate(
+ static_cast<char *>(const_cast<void *>(object.data())),
+ object.size());
return cls(mlirGPUObjectAttrGetWithKernels(
mlirAttributeGetContext(target), target, format, objectStrRef,
mlirObjectProps.has_value() ? *mlirObjectProps
@@ -61,7 +63,7 @@ PYBIND11_MODULE(_mlirDialectsGPU, m) {
: MlirAttribute{nullptr}));
},
"cls"_a, "target"_a, "format"_a, "object"_a,
- "properties"_a = py::none(), "kernels"_a = py::none(),
+ "properties"_a.none() = nb::none(), "kernels"_a.none() = nb::none(),
"Gets a gpu.object from parameters.")
.def_property_readonly(
"target",
@@ -73,18 +75,18 @@ PYBIND11_MODULE(_mlirDialectsGPU, m) {
"object",
[](MlirAttribute self) {
MlirStringRef stringRef = mlirGPUObjectAttrGetObject(self);
- return py::bytes(stringRef.data, stringRef.length);
+ return nb::bytes(stringRef.data, stringRef.length);
})
.def_property_readonly("properties",
- [](MlirAttribute self) {
+ [](MlirAttribute self) -> nb::object {
if (mlirGPUObjectAttrHasProperties(self))
- return py::cast(
+ return nb::cast(
mlirGPUObjectAttrGetProperties(self));
- return py::none().cast<py::object>();
+ return nb::none();
})
- .def_property_readonly("kernels", [](MlirAttribute self) {
+ .def_property_readonly("kernels", [](MlirAttribute self) -> nb::object {
if (mlirGPUObjectAttrHasKernels(self))
- return py::cast(mlirGPUObjectAttrGetKernels(self));
- return py::none().cast<py::object>();
+ return nb::cast(mlirGPUObjectAttrGetKernels(self));
+ return nb::none();
});
}
diff --git a/mlir/lib/Bindings/Python/DialectLLVM.cpp b/mlir/lib/Bindings/Python/DialectLLVM.cpp
index cccf1370b8cc87..93714c43e153cf 100644
--- a/mlir/lib/Bindings/Python/DialectLLVM.cpp
+++ b/mlir/lib/Bindings/Python/DialectLLVM.cpp
@@ -12,15 +12,23 @@
#include "mlir-c/IR.h"
#include "mlir-c/Support.h"
#include "mlir/Bindings/Python/Diagnostics.h"
-#include "mlir/Bindings/Python/PybindAdaptors.h"
+#include "mlir/Bindings/Python/NanobindAdaptors.h"
+
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/optional.h>
+#include <nanobind/stl/string.h>
+#include <nanobind/stl/vector.h>
+
+namespace nb = nanobind;
+
+using namespace nanobind::literals;
-namespace py = pybind11;
using namespace llvm;
using namespace mlir;
using namespace mlir::python;
-using namespace mlir::python::adaptors;
+using namespace mlir::python::nanobind_adaptors;
-void populateDialectLLVMSubmodule(const pybind11::module &m) {
+void populateDialectLLVMSubmodule(const nanobind::module_ &m) {
//===--------------------------------------------------------------------===//
// StructType
@@ -31,35 +39,35 @@ void populateDialectLLVMSubmodule(const pybind11::module &m) {
llvmStructType.def_classmethod(
"get_literal",
- [](py::object cls, const std::vector<MlirType> &elements, bool packed,
+ [](nb::object cls, const std::vector<MlirType> &elements, bool packed,
MlirLocation loc) {
CollectDiagnosticsToStringScope scope(mlirLocationGetContext(loc));
MlirType type = mlirLLVMStructTypeLiteralGetChecked(
loc, elements.size(), elements.data(), packed);
if (mlirTypeIsNull(type)) {
- throw py::value_error(scope.takeMessage());
+ throw nb::value_error(scope.takeMessage().c_str());
}
return cls(type);
},
- "cls"_a, "elements"_a, py::kw_only(), "packed"_a = false,
- "loc"_a = py::none());
+ "cls"_a, "elements"_a, nb::kw_only(), "packed"_a = false,
+ "loc"_a.none() = nb::none());
llvmStructType.def_classmethod(
"get_identified",
- [](py::object cls, const std::string &name, MlirContext context) {
+ [](nb::object cls, const std::string &name, MlirContext context) {
return cls(mlirLLVMStructTypeIdentifiedGet(
context, mlirStringRefCreate(name.data(), name.size())));
},
- "cls"_a, "name"_a, py::kw_only(), "context"_a = py::none());
+ "cls"_a, "name"_a, nb::kw_only(), "context"_a.none() = nb::none());
llvmStructType.def_classmethod(
"get_opaque",
- [](py::object cls, const std::string &name, MlirContext context) {
+ [](nb::object cls, const std::string &name, MlirContext context) {
return cls(mlirLLVMStructTypeOpaqueGet(
context, mlirStringRefCreate(name.data(), name.size())));
},
- "cls"_a, "name"_a, "context"_a = py::none());
+ "cls"_a, "name"_a, "context"_a.none() = nb::none());
llvmStructType.def(
"set_body",
@@ -67,22 +75,22 @@ void populateDialectLLVMSubmodule(const pybind11::module &m) {
MlirLogicalResult result = mlirLLVMStructTypeSetBody(
self, elements.size(), elements.data(), packed);
if (!mlirLogicalResultIsSuccess(result)) {
- throw py::value_error(
+ throw nb::value_error(
"Struct body already set to different content.");
}
},
- "elements"_a, py::kw_only(), "packed"_a = false);
+ "elements"_a, nb::kw_only(), "packed"_a = false);
llvmStructType.def_classmethod(
"new_identified",
- [](py::object cls, const std::string &name,
+ [](nb::object cls, const std::string &name,
const std::vector<MlirType> &elements, bool packed, MlirContext ctx) {
return cls(mlirLLVMStructTypeIdentifiedNewGet(
ctx, mlirStringRefCreate(name.data(), name.length()),
elements.size(), elements.data(), packed));
},
- "cls"_a, "name"_a, "elements"_a, py::kw_only(), "packed"_a = false,
- "context"_a = py::none());
+ "cls"_a, "name"_a, "elements"_a, nb::kw_only(), "packed"_a = false,
+ "context"_a.none() = nb::none());
llvmStructType.def_property_readonly(
"name", [](MlirType type) -> std::optional<std::string> {
@@ -93,12 +101,12 @@ void populateDialectLLVMSubmodule(const pybind11::module &m) {
return StringRef(stringRef.data, stringRef.length).str();
});
- llvmStructType.def_property_readonly("body", [](MlirType type) -> py::object {
+ llvmStructType.def_property_readonly("body", [](MlirType type) -> nb::object {
// Don't crash in absence of a body.
if (mlirLLVMStructTypeIsOpaque(type))
- return py::none();
+ return nb::none();
- py::list body;
+ nb::list body;
for (intptr_t i = 0, e = mlirLLVMStructTypeGetNumElementTypes(type); i < e;
++i) {
body.append(mlirLLVMStructTypeGetElementType(type, i));
@@ -119,24 +127,24 @@ void populateDialectLLVMSubmodule(const pybind11::module &m) {
mlir_type_subclass(m, "PointerType", mlirTypeIsALLVMPointerType)
.def_classmethod(
"get",
- [](py::object cls, std::optional<unsigned> addressSpace,
+ [](nb::object cls, std::optional<unsigned> addressSpace,
MlirContext context) {
CollectDiagnosticsToStringScope scope(context);
MlirType type = mlirLLVMPointerTypeGet(
context, addressSpace.has_value() ? *addressSpace : 0);
if (mlirTypeIsNull(type)) {
- throw py::value_error(scope.takeMessage());
+ throw nb::value_error(scope.takeMessage().c_str());
}
return cls(type);
},
- "cls"_a, "address_space"_a = py::none(), py::kw_only(),
- "context"_a = py::none())
+ "cls"_a, "address_space"_a.none() = nb::none(), nb::kw_only(),
+ "context"_a.none() = nb::none())
.def_property_readonly("address_space", [](MlirType type) {
return mlirLLVMPointerTypeGetAddressSpace(type);
});
}
-PYBIND11_MODULE(_mlirDialectsLLVM, m) {
+NB_MODULE(_mlirDialectsLLVM, m) {
m.doc() = "MLIR LLVM Dialect";
populateDialectLLVMSubmodule(m);
diff --git a/mlir/lib/Bindings/Python/DialectLinalg.cpp b/mlir/lib/Bindings/Python/DialectLinalg.cpp
index 2e54ebeb61fb10..4ac031ba2b145b 100644
--- a/mlir/lib/Bindings/Python/DialectLinalg.cpp
+++ b/mlir/lib/Bindings/Python/DialectLinalg.cpp
@@ -6,22 +6,24 @@
//
//===----------------------------------------------------------------------===//
+#include <nanobind/nanobind.h>
+
#include "mlir-c/Dialect/Linalg.h"
#include "mlir-c/IR.h"
-#include "mlir/Bindings/Python/PybindAdaptors.h"
+#include "mlir/Bindings/Python/NanobindAdaptors.h"
-namespace py = pybind11;
+namespace nb = nanobind;
-static void populateDialectLinalgSubmodule(py::module m) {
+static void populateDialectLinalgSubmodule(nb::module_ m) {
m.def(
"fill_builtin_region",
[](MlirOperation op) { mlirLinalgFillBuiltinNamedOpRegion(op); },
- py::arg("op"),
+ nb::arg("op"),
"Fill the region for `op`, which is assumed to be a builtin named Linalg "
"op.");
}
-PYBIND11_MODULE(_mlirDialectsLinalg, m) {
+NB_MODULE(_mlirDialectsLinalg, m) {
m.doc() = "MLIR Linalg dialect.";
populateDialectLinalgSubmodule(m);
diff --git a/mlir/lib/Bindings/Python/DialectNVGPU.cpp b/mlir/lib/Bindings/Python/DialectNVGPU.cpp
index 754e0a75b0abc7..6937c686277cad 100644
--- a/mlir/lib/Bindings/Python/DialectNVGPU.cpp
+++ b/mlir/lib/Bindings/Python/DialectNVGPU.cpp
@@ -6,35 +6,36 @@
//
//===----------------------------------------------------------------------===//
+#include <nanobind/nanobind.h>
+
#include "mlir-c/Dialect/NVGPU.h"
#include "mlir-c/IR.h"
-#include "mlir/Bindings/Python/PybindAdaptors.h"
-#include <pybind11/pybind11.h>
+#include "mlir/Bindings/Python/NanobindAdaptors.h"
-namespace py = pybind11;
+namespace nb = nanobind;
using namespace llvm;
using namespace mlir;
using namespace mlir::python;
-using namespace mlir::python::adaptors;
+using namespace mlir::python::nanobind_adaptors;
-static void populateDialectNVGPUSubmodule(const pybind11::module &m) {
+static void populateDialectNVGPUSubmodule(const nb::module_ &m) {
auto nvgpuTensorMapDescriptorType = mlir_type_subclass(
m, "TensorMapDescriptorType", mlirTypeIsANVGPUTensorMapDescriptorType);
nvgpuTensorMapDescriptorType.def_classmethod(
"get",
- [](py::object cls, MlirType tensorMemrefType, int swizzle, int l2promo,
+ [](nb::object cls, MlirType tensorMemrefType, int swizzle, int l2promo,
int oobFill, int interleave, MlirContext ctx) {
return cls(mlirNVGPUTensorMapDescriptorTypeGet(
ctx, tensorMemrefType, swizzle, l2promo, oobFill, interleave));
},
"Gets an instance of TensorMapDescriptorType in the same context",
- py::arg("cls"), py::arg("tensor_type"), py::arg("swizzle"),
- py::arg("l2promo"), py::arg("oob_fill"), py::arg("interleave"),
- py::arg("ctx") = py::none());
+ nb::arg("cls"), nb::arg("tensor_type"), nb::arg("swizzle"),
+ nb::arg("l2promo"), nb::arg("oob_fill"), nb::arg("interleave"),
+ nb::arg("ctx").none() = nb::none());
}
-PYBIND11_MODULE(_mlirDialectsNVGPU, m) {
+NB_MODULE(_mlirDialectsNVGPU, m) {
m.doc() = "MLIR NVGPU dialect.";
populateDialectNVGPUSubmodule(m);
diff --git a/mlir/lib/Bindings/Python/DialectPDL.cpp b/mlir/lib/Bindings/Python/DialectPDL.cpp
index 8d3f9a7ab1d6ac..56a7968c06384a 100644
--- a/mlir/lib/Bindings/Python/DialectPDL.cpp
+++ b/mlir/lib/Bindings/Python/DialectPDL.cpp
@@ -6,21 +6,19 @@
//
//===----------------------------------------------------------------------===//
+#include <nanobind/nanobind.h>
+
#include "mlir-c/Dialect/PDL.h"
#include "mlir-c/IR.h"
-#include "mlir/Bindings/Python/PybindAdaptors.h"
-#include <pybind11/cast.h>
-#include <pybind11/detail/common.h>
-#include <pybind11/pybind11.h>
-#include <pybind11/pytypes.h>
+#include "mlir/Bindings/Python/NanobindAdaptors.h"
-namespace py = pybind11;
+namespace nb = nanobind;
using namespace llvm;
using namespace mlir;
using namespace mlir::python;
-using namespace mlir::python::adaptors;
+using namespace mlir::python::nanobind_adaptors;
-void populateDialectPDLSubmodule(const pybind11::module &m) {
+void populateDialectPDLSubmodule(const nanobind::module_ &m) {
//===-------------------------------------------------------------------===//
// PDLType
//===-------------------------------------------------------------------===//
@@ -35,11 +33,11 @@ void populateDialectPDLSubmodule(const pybind11::module &m) {
mlir_type_subclass(m, "AttributeType", mlirTypeIsAPDLAttributeType);
attributeType.def_classmethod(
"get",
- [](py::object cls, MlirContext ctx) {
+ [](nb::object cls, MlirContext ctx) {
return cls(mlirPDLAttributeTypeGet(ctx));
},
- "Get an instance of AttributeType in given context.", py::arg("cls"),
- py::arg("context") = py::none());
+ "Get an instance of AttributeType in given context.", nb::arg("cls"),
+ nb::arg("context").none() = nb::none());
//===-------------------------------------------------------------------===//
// OperationType
@@ -49,11 +47,11 @@ void populateDialectPDLSubmodule(const pybind11::module &m) {
mlir_type_subclass(m, "OperationType", mlirTypeIsAPDLOperationType);
operationType.def_classmethod(
"get",
- [](py::object cls, MlirContext ctx) {
+ [](nb::object cls, MlirContext ctx) {
return cls(mlirPDLOperationTypeGet(ctx));
},
- "Get an instance of OperationType in given context.", py::arg("cls"),
- py::arg("context") = py::none());
+ "Get an instance of OperationType in given context.", nb::arg("cls"),
+ nb::arg("context").none() = nb::none());
//===-------------------------------------------------------------------===//
// RangeType
@@ -62,12 +60,12 @@ void populateDialectPDLSubmodule(const pybind11::module &m) {
auto rangeType = mlir_type_subclass(m, "RangeType", mlirTypeIsAPDLRangeType);
rangeType.def_classmethod(
"get",
- [](py::object cls, MlirType elementType) {
+ [](nb::object cls, MlirType elementType) {
return cls(mlirPDLRangeTypeGet(elementType));
},
"Gets an instance of RangeType in the same context as the provided "
"element type.",
- py::arg("cls"), py::arg("element_type"));
+ nb::arg("cls"), nb::arg("element_type"));
rangeType.def_property_readonly(
"element_type",
[](MlirType type) { return mlirPDLRangeTypeGetElementType(type); },
@@ -80,11 +78,11 @@ void populateDialectPDLSubmodule(const pybind11::module &m) {
auto typeType = mlir_type_subclass(m, "TypeType", mlirTypeIsAPDLTypeType);
typeType.def_classmethod(
"get",
- [](py::object cls, MlirContext ctx) {
+ [](nb::object cls, MlirContext ctx) {
return cls(mlirPDLTypeTypeGet(ctx));
},
- "Get an instance of TypeType in given context.", py::arg("cls"),
- py::arg("context") = py::none());
+ "Get an instance of TypeType in given context.",...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/119924
More information about the llvm-commits
mailing list