Skip to content

Commit

Permalink
[Python] Switch CIRCT Python extension and dialects to nanobind.
Browse files Browse the repository at this point in the history
This follows the general nanobind porting guide and the upstream MLIR
changes in the same vein. For the most part, this is a very direct
migration accomplished with sed.

There were a couple minor differences.

Nanobind is more strict about implicit casting, which popped up in a
couple places:

* In MSFTModule, getNearestFreeInColumn was expecting a
CIRCTMSFTPrimitiveType (int), but was actually being passed
PrimitiveType enums at the call sites. Nanobind refused to
automatically cast enum to int, so it was done manually.
* In SVModule, SVAttributeAttr.get was expecting a bool for
emitAsComment, even though it was declared as having a default of
None. Nanobind refused to automatically cast None to bool, so it was
done manually, with a default of false.

Nanobind also prints enums slightly differently, so one FileCheck test
had to be updated for ESI.

Otherwise, the changes were purely mechanical.
  • Loading branch information
mikeurbach committed Jan 22, 2025
1 parent 19f9af0 commit e30ed45
Show file tree
Hide file tree
Showing 16 changed files with 379 additions and 392 deletions.
4 changes: 2 additions & 2 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -514,11 +514,11 @@ option(CIRCT_BINDINGS_PYTHON_ENABLED "Enables CIRCT Python bindings." OFF)
if(CIRCT_BINDINGS_PYTHON_ENABLED)
message(STATUS "CIRCT Python bindings are enabled.")
set(MLIR_DISABLE_CONFIGURE_PYTHON_DEV_PACKAGES 0)
mlir_detect_pybind11_install()
mlir_detect_nanobind_install()
# Prime the search like mlir_configure_python_dev_modules
find_package(Python3 3.8 COMPONENTS Interpreter Development)
find_package(Python3 3.8 COMPONENTS Interpreter Development.Module)
find_package(pybind11 2.10 CONFIG REQUIRED)
find_package(nanobind 2.4.0 CONFIG REQUIRED)
else()
message(STATUS "CIRCT Python bindings are disabled.")
# Lookup python either way as some integration tests use python without the
Expand Down
2 changes: 1 addition & 1 deletion integration_test/Bindings/Python/dialects/esi.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
assert (not bundle_type.resettable)
for bchan in bundle_type.channels:
print(bchan)
# CHECK: ('i16chan', <ChannelDirection.TO: 1>, Type(!esi.channel<i16>))
# CHECK: ('i16chan', ChannelDirection.TO, Type(!esi.channel<i16>))
print()

bundle_type = esi.BundleType.get(
Expand Down
38 changes: 19 additions & 19 deletions lib/Bindings/Python/CIRCTModule.cpp
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
//===- CIRCTModule.cpp - Main pybind module -------------------------------===//
//===- CIRCTModule.cpp - Main nanobind module -----------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
Expand Down Expand Up @@ -35,14 +35,14 @@
#include "mlir-c/Bindings/Python/Interop.h"
#include "mlir-c/IR.h"
#include "mlir-c/Transforms.h"
#include "mlir/Bindings/Python/PybindAdaptors.h"
#include "mlir/Bindings/Python/NanobindAdaptors.h"

#include "llvm-c/ErrorHandling.h"
#include "llvm/Support/Signals.h"

#include "PybindUtils.h"
#include <pybind11/pybind11.h>
namespace py = pybind11;
#include "NanobindUtils.h"
#include <nanobind/nanobind.h>
namespace nb = nanobind;

static void registerPasses() {
registerArcPasses();
Expand All @@ -59,15 +59,15 @@ static void registerPasses() {
mlirRegisterTransformsCSE();
}

PYBIND11_MODULE(_circt, m) {
NB_MODULE(_circt, m) {
m.doc() = "CIRCT Python Native Extension";
registerPasses();
llvm::sys::PrintStackTraceOnErrorSignal(/*argv=*/"");
LLVMEnablePrettyStackTrace();

m.def(
"register_dialects",
[](py::object capsule) {
[](nb::object capsule) {
// Get the MlirContext capsule from PyMlirContext capsule.
auto wrappedCapsule = capsule.attr(MLIR_PYTHON_CAPI_PTR_ATTR);
MlirContext context = mlirPythonCapsuleToContext(wrappedCapsule.ptr());
Expand Down Expand Up @@ -145,9 +145,9 @@ PYBIND11_MODULE(_circt, m) {
},
"Register CIRCT dialects on a PyMlirContext.");

m.def("export_verilog", [](MlirModule mod, py::object fileObject) {
m.def("export_verilog", [](MlirModule mod, nb::object fileObject) {
circt::python::PyFileAccumulator accum(fileObject, false);
py::gil_scoped_release();
nb::gil_scoped_release();
mlirExportVerilog(mod, accum.getCallback(), accum.getUserData());
});

Expand All @@ -156,26 +156,26 @@ PYBIND11_MODULE(_circt, m) {
mlirExportSplitVerilog(mod, cDirectory);
});

py::module esi = m.def_submodule("_esi", "ESI API");
nb::module_ esi = m.def_submodule("_esi", "ESI API");
circt::python::populateDialectESISubmodule(esi);
py::module msft = m.def_submodule("_msft", "MSFT API");
nb::module_ msft = m.def_submodule("_msft", "MSFT API");
circt::python::populateDialectMSFTSubmodule(msft);
py::module hw = m.def_submodule("_hw", "HW API");
nb::module_ hw = m.def_submodule("_hw", "HW API");
circt::python::populateDialectHWSubmodule(hw);
py::module seq = m.def_submodule("_seq", "Seq API");
nb::module_ seq = m.def_submodule("_seq", "Seq API");
circt::python::populateDialectSeqSubmodule(seq);
py::module om = m.def_submodule("_om", "OM API");
nb::module_ om = m.def_submodule("_om", "OM API");
circt::python::populateDialectOMSubmodule(om);
py::module rtg = m.def_submodule("_rtg", "RTG API");
nb::module_ rtg = m.def_submodule("_rtg", "RTG API");
circt::python::populateDialectRTGSubmodule(rtg);
py::module rtgtool = m.def_submodule("_rtgtool", "RTGTool API");
nb::module_ rtgtool = m.def_submodule("_rtgtool", "RTGTool API");
circt::python::populateDialectRTGToolSubmodule(rtgtool);
#ifdef CIRCT_INCLUDE_TESTS
py::module rtgtest = m.def_submodule("_rtgtest", "RTGTest API");
nb::module_ rtgtest = m.def_submodule("_rtgtest", "RTGTest API");
circt::python::populateDialectRTGTestSubmodule(rtgtest);
#endif
py::module sv = m.def_submodule("_sv", "SV API");
nb::module_ sv = m.def_submodule("_sv", "SV API");
circt::python::populateDialectSVSubmodule(sv);
py::module support = m.def_submodule("_support", "CIRCT support");
nb::module_ support = m.def_submodule("_support", "CIRCT support");
circt::python::populateSupportSubmodule(support);
}
22 changes: 11 additions & 11 deletions lib/Bindings/Python/CIRCTModules.h
Original file line number Diff line number Diff line change
Expand Up @@ -13,23 +13,23 @@
#ifndef CIRCT_BINDINGS_PYTHON_CIRCTMODULES_H
#define CIRCT_BINDINGS_PYTHON_CIRCTMODULES_H

#include <pybind11/pybind11.h>
#include <nanobind/nanobind.h>

namespace circt {
namespace python {

void populateDialectESISubmodule(pybind11::module &m);
void populateDialectHWSubmodule(pybind11::module &m);
void populateDialectMSFTSubmodule(pybind11::module &m);
void populateDialectOMSubmodule(pybind11::module &m);
void populateDialectRTGSubmodule(pybind11::module &m);
void populateDialectRTGToolSubmodule(pybind11::module &m);
void populateDialectESISubmodule(nanobind::module_ &m);
void populateDialectHWSubmodule(nanobind::module_ &m);
void populateDialectMSFTSubmodule(nanobind::module_ &m);
void populateDialectOMSubmodule(nanobind::module_ &m);
void populateDialectRTGSubmodule(nanobind::module_ &m);
void populateDialectRTGToolSubmodule(nanobind::module_ &m);
#ifdef CIRCT_INCLUDE_TESTS
void populateDialectRTGTestSubmodule(pybind11::module &m);
void populateDialectRTGTestSubmodule(nanobind::module_ &m);
#endif
void populateDialectSeqSubmodule(pybind11::module &m);
void populateDialectSVSubmodule(pybind11::module &m);
void populateSupportSubmodule(pybind11::module &m);
void populateDialectSeqSubmodule(nanobind::module_ &m);
void populateDialectSVSubmodule(nanobind::module_ &m);
void populateSupportSubmodule(nanobind::module_ &m);

} // namespace python
} // namespace circt
Expand Down
4 changes: 3 additions & 1 deletion lib/Bindings/Python/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ set(PYTHON_BINDINGS_SOURCES
SVModule.cpp
# Headers must be included explicitly so they are installed.
CIRCTModules.h
PybindUtils.h
NanobindUtils.h
)

set(PYTHON_BINDINGS_LINK_LIBS
Expand Down Expand Up @@ -70,6 +70,8 @@ declare_mlir_python_extension(CIRCTBindingsPythonExtension
${PYTHON_BINDINGS_LINK_LIBS}
PRIVATE_LINK_LIBS
LLVMSupport
PYTHON_BINDINGS_LIBRARY
nanobind
)

add_dependencies(CIRCTBindingsPythonExtension circt-headers)
Expand Down
Loading

0 comments on commit e30ed45

Please sign in to comment.