Skip to content
This repository was archived by the owner on Oct 11, 2025. It is now read-only.

Commit b394e53

Browse files
authored
[mlir python] Add nanobind support for standalone dialects. (#117922)
This PR allows out-of-tree dialects to write Python dialect modules using nanobind instead of pybind11. It may make sense to migrate in-tree dialects and some of the ODS Python infrastructure to nanobind, but that is a topic for a future change. This PR makes the following changes: * adds nanobind to the CMake and Bazel build systems. We also add robin_map to the Bazel build, which is a dependency of nanobind. * adds a PYTHON_BINDING_LIBRARY option to various CMake functions, such as declare_mlir_python_extension, allowing users to select a Python binding library. * creates a fork of mlir/include/mlir/Bindings/Python/PybindAdaptors.h named NanobindAdaptors.h. This plays the same role, using nanobind instead of pybind11. * splits CollectDiagnosticsToStringScope out of PybindAdaptors.h and into a new header mlir/include/mlir/Bindings/Python/Diagnostics.h, since it is code that is no way related to pybind11 or for that matter, Python. * changed the standalone Python extension example to have both pybind11 and nanobind variants. * changed mlir/python/mlir/dialects/python_test.py to have both pybind11 and nanobind variants. Notes: * A slightly unfortunate thing that I needed to do in the CMake integration was to use FindPython in addition to FindPython3, since nanobind's CMake integration expects the Python_ names for variables. Perhaps there's a better way to do this.
1 parent 26536e4 commit b394e53

File tree

8 files changed

+770
-55
lines changed

8 files changed

+770
-55
lines changed
Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
//===- Diagnostics.h - Helpers for diagnostics in Python bindings ---------===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#ifndef MLIR_BINDINGS_PYTHON_DIAGNOSTICS_H
10+
#define MLIR_BINDINGS_PYTHON_DIAGNOSTICS_H
11+
12+
#include <cassert>
13+
#include <string>
14+
15+
#include "mlir-c/Diagnostics.h"
16+
#include "mlir-c/IR.h"
17+
#include "llvm/ADT/StringRef.h"
18+
19+
namespace mlir {
20+
namespace python {
21+
22+
/// RAII scope intercepting all diagnostics into a string. The message must be
23+
/// checked before this goes out of scope.
24+
class CollectDiagnosticsToStringScope {
25+
public:
26+
explicit CollectDiagnosticsToStringScope(MlirContext ctx) : context(ctx) {
27+
handlerID = mlirContextAttachDiagnosticHandler(ctx, &handler, &errorMessage,
28+
/*deleteUserData=*/nullptr);
29+
}
30+
~CollectDiagnosticsToStringScope() {
31+
assert(errorMessage.empty() && "unchecked error message");
32+
mlirContextDetachDiagnosticHandler(context, handlerID);
33+
}
34+
35+
[[nodiscard]] std::string takeMessage() { return std::move(errorMessage); }
36+
37+
private:
38+
static MlirLogicalResult handler(MlirDiagnostic diag, void *data) {
39+
auto printer = +[](MlirStringRef message, void *data) {
40+
*static_cast<std::string *>(data) +=
41+
llvm::StringRef(message.data, message.length);
42+
};
43+
MlirLocation loc = mlirDiagnosticGetLocation(diag);
44+
*static_cast<std::string *>(data) += "at ";
45+
mlirLocationPrint(loc, printer, data);
46+
*static_cast<std::string *>(data) += ": ";
47+
mlirDiagnosticPrint(diag, printer, data);
48+
return mlirLogicalResultSuccess();
49+
}
50+
51+
MlirContext context;
52+
MlirDiagnosticHandlerID handlerID;
53+
std::string errorMessage = "";
54+
};
55+
56+
} // namespace python
57+
} // namespace mlir
58+
59+
#endif // MLIR_BINDINGS_PYTHON_DIAGNOSTICS_H

0 commit comments

Comments
 (0)