From 344d7ba781c553ffd69f590df9a562fcb57878de Mon Sep 17 00:00:00 2001 From: erwei-xilinx Date: Mon, 13 Jan 2025 03:39:54 -0800 Subject: [PATCH 1/6] Switch from Pybind11 to Nanobind --- python/AIEMLIRModule.cpp | 56 +++++++++++++++++++++------------------- python/AIERTModule.cpp | 13 +++++----- python/CMakeLists.txt | 10 +++++++ python/XRTModule.cpp | 45 +++++++++++++++++--------------- 4 files changed, 70 insertions(+), 54 deletions(-) diff --git a/python/AIEMLIRModule.cpp b/python/AIEMLIRModule.cpp index 8f27cd227a5..588a90f7fe2 100644 --- a/python/AIEMLIRModule.cpp +++ b/python/AIEMLIRModule.cpp @@ -15,11 +15,10 @@ #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 "llvm/ADT/Twine.h" -#include -#include -#include +#include #include #include @@ -27,11 +26,11 @@ #include #include -using namespace mlir::python::adaptors; -namespace py = pybind11; -using namespace py::literals; +using namespace mlir::python; +namespace nb = nanobind; +using namespace nb::literals; -PYBIND11_MODULE(_aie, m) { +NB_MODULE(_aie, m) { aieRegisterAllPasses(); @@ -48,33 +47,35 @@ PYBIND11_MODULE(_aie, m) { "registry"_a); // AIE types bindings - mlir_type_subclass(m, "ObjectFifoType", aieTypeIsObjectFifoType) + nanobind_adaptors::mlir_type_subclass(m, "ObjectFifoType", + aieTypeIsObjectFifoType) .def_classmethod( "get", - [](const py::object &cls, const MlirType type) { + [](const nb::object &cls, const MlirType type) { return cls(aieObjectFifoTypeGet(type)); }, "Get an instance of ObjectFifoType with given element type.", - "self"_a, "type"_a = py::none()); + "self"_a, "type"_a = nb::none()); - mlir_type_subclass(m, "ObjectFifoSubviewType", aieTypeIsObjectFifoSubviewType) + nanobind_adaptors::mlir_type_subclass(m, "ObjectFifoSubviewType", + aieTypeIsObjectFifoSubviewType) .def_classmethod( "get", - [](const py::object &cls, const MlirType type) { + [](const nb::object &cls, const MlirType type) { return cls(aieObjectFifoSubviewTypeGet(type)); }, "Get an instance of ObjectFifoSubviewType with given element type.", - "self"_a, "type"_a = py::none()); + "self"_a, "type"_a = nb::none()); auto stealCStr = [](MlirStringRef mlirString) { if (!mlirString.data || mlirString.length == 0) throw std::runtime_error("couldn't translate"); std::string cpp(mlirString.data, mlirString.length); free((void *)mlirString.data); - py::handle pyS = PyUnicode_DecodeLatin1(cpp.data(), cpp.length(), nullptr); + nb::handle pyS = PyUnicode_DecodeLatin1(cpp.data(), cpp.length(), nullptr); if (!pyS) - throw py::error_already_set(); - return py::reinterpret_steal(pyS); + throw nb::python_error(); + return nb::steal(pyS); }; m.def( @@ -101,8 +102,11 @@ PYBIND11_MODULE(_aie, m) { if (mlirLogicalResultIsFailure(aieTranslateToCDODirect( op, {workDirPath.data(), workDirPath.size()}, bigendian, emitUnified, cdoDebug, aieSim, xaieDebug, enableCores))) - throw py::value_error("Failed to generate cdo because: " + - scope.takeMessage()); + throw nb::value_error( + (llvm::Twine("Failed to generate cdo because: ") + + llvm::Twine(scope.takeMessage())) + .str() + .c_str()); }, "module"_a, "work_dir_path"_a, "bigendian"_a = false, "emit_unified"_a = false, "cdo_debug"_a = false, "aiesim"_a = false, @@ -110,8 +114,8 @@ PYBIND11_MODULE(_aie, m) { m.def( "transaction_binary_to_mlir", - [](MlirContext ctx, py::bytes bytes) { - std::string s = bytes; + [](MlirContext ctx, nb::bytes bytes) { + std::string s = static_cast(bytes.data()); MlirStringRef bin = {s.data(), s.size()}; return aieTranslateBinaryToTxn(ctx, bin); }, @@ -120,9 +124,9 @@ PYBIND11_MODULE(_aie, m) { m.def( "npu_instgen", [&stealCStr](MlirOperation op) { - py::str npuInstructions = stealCStr(aieTranslateToNPU(op)); + nb::str npuInstructions = stealCStr(aieTranslateToNPU(op)); auto individualInstructions = - npuInstructions.attr("split")().cast(); + nb::cast(npuInstructions.attr("split")()); for (size_t i = 0; i < individualInstructions.size(); ++i) individualInstructions[i] = individualInstructions[i].attr("strip")(); return individualInstructions; @@ -132,10 +136,10 @@ PYBIND11_MODULE(_aie, m) { m.def( "generate_control_packets", [&stealCStr](MlirOperation op) { - py::str ctrlPackets = + nb::str ctrlPackets = stealCStr(aieTranslateControlPacketsToUI32Vec(op)); auto individualInstructions = - ctrlPackets.attr("split")().cast(); + nb::cast(ctrlPackets.attr("split")()); for (size_t i = 0; i < individualInstructions.size(); ++i) individualInstructions[i] = individualInstructions[i].attr("strip")(); return individualInstructions; @@ -171,7 +175,7 @@ PYBIND11_MODULE(_aie, m) { m.def("get_target_model", [](uint32_t d) -> PyAieTargetModel { return aieGetTargetModel(d); }); - py::class_(m, "AIETargetModel", py::module_local()) + nb::class_(m, "AIETargetModel") .def( "columns", [](PyAieTargetModel &self) { diff --git a/python/AIERTModule.cpp b/python/AIERTModule.cpp index 0f627a6ad4c..2aecbb5d0af 100644 --- a/python/AIERTModule.cpp +++ b/python/AIERTModule.cpp @@ -13,13 +13,12 @@ #include "aie/Bindings/PyTypes.h" -#include -#include +#include #include -namespace py = pybind11; -using namespace py::literals; +namespace nb = nanobind; +using namespace nb::literals; class PyAIERTControl { public: @@ -31,10 +30,10 @@ class PyAIERTControl { AieRtControl ctl; }; -PYBIND11_MODULE(_aiert, m) { +NB_MODULE(_aiert, m) { - py::class_(m, "AIERTControl", py::module_local()) - .def(py::init(), "target_model"_a) + nb::class_(m, "AIERTControl") + .def(nb::init(), "target_model"_a) .def("start_transaction", [](PyAIERTControl &self) { aieRtStartTransaction(self.ctl); }) .def("export_serialized_transaction", diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 1df5584816a..9b737f10ef6 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -191,6 +191,9 @@ if (AIE_ENABLE_PYTHON_PASSES) PRIVATE_LINK_LIBS ${_py_libs} + + PYTHON_BINDINGS_LIBRARY + nanobind ) target_include_directories( AIEPythonExtensions.MLIR @@ -282,6 +285,8 @@ else () AIECAPI PRIVATE_LINK_LIBS LLVMSupport + PYTHON_BINDINGS_LIBRARY + nanobind ) if(AIE_ENABLE_XRT_PYTHON_BINDINGS) @@ -298,6 +303,8 @@ else () LLVMSupport xrt_coreutil uuid + PYTHON_BINDINGS_LIBRARY + nanobind ) target_include_directories(AIEPythonExtensions.XRT INTERFACE ${XRT_INCLUDE_DIR}) target_link_directories(AIEPythonExtensions.XRT INTERFACE ${XRT_LIB_DIR}) @@ -314,6 +321,9 @@ else () PRIVATE_LINK_LIBS LLVMSupport + + PYTHON_BINDINGS_LIBRARY + nanobind ) add_mlir_python_common_capi_library(AIEAggregateCAPI diff --git a/python/XRTModule.cpp b/python/XRTModule.cpp index fb200f66502..9f9f49759a3 100644 --- a/python/XRTModule.cpp +++ b/python/XRTModule.cpp @@ -13,17 +13,19 @@ #include "xrt/xrt_device.h" #include "xrt/xrt_kernel.h" -#include -#include -#include -#include +#include +#include +#include +#include #include +#include +#include #include #include -namespace py = pybind11; -using namespace py::literals; +namespace nb = nanobind; +using namespace nb::literals; // group_id 0 is for npu instructions // group_id 1 is for number of npu instructions @@ -55,16 +57,16 @@ class PyXCLBin { } template - std::vector - mmapBuffers(std::vector> shapes) { + std::vector> + mmapBuffers(std::vector> shapes) { this->buffers.reserve(shapes.size()); - std::vector views; + std::vector> views; views.reserve(shapes.size()); auto initAndViewBuffer = [this]( - std::vector shape, int groupId, + std::vector shape, int groupId, std::vector> &buffers, - std::vector &views) { + std::vector> &views) { int nElements = std::accumulate(shape.begin(), shape.end(), 1, std::multiplies<>()); int nBytes = nElements * sizeof(ElementT); @@ -79,12 +81,13 @@ class PyXCLBin { std::vector strides_{1}; for (int i = shape.size() - 1; i > 0; i--) strides_.push_back(strides_.back() * shape[i]); - std::vector strides; + std::vector strides; // stride in bytes std::transform(strides_.rbegin(), strides_.rend(), std::back_inserter(strides), [](int s) { return s * sizeof(ElementT); }); - views.push_back(py::memoryview::from_buffer(buf, shape, strides)); + views.push_back(nb::ndarray(buf, shape.size(), shape.data(), nb::handle(), + strides.data())); }; for (size_t i = 0; i < shapes.size(); ++i) @@ -140,22 +143,22 @@ class PyXCLBin { std::unique_ptr run_; }; -PYBIND11_MODULE(_xrt, m) { +NB_MODULE(_xrt, m) { - py::class_(m, "XCLBin", py::module_local()) - .def(py::init(), + nb::class_(m, "XCLBin") + .def(nb::init(), "xclbin_path"_a, "kernel_name"_a, "device_index"_a = 0) .def("load_npu_instructions", &PyXCLBin::loadNPUInstructions, "insts"_a) .def("sync_buffers_to_device", &PyXCLBin::syncBuffersToDevice) .def("sync_buffers_from_device", &PyXCLBin::syncBuffersFromDevice) .def("run", &PyXCLBin::run) .def("_run_only_npu_instructions", &PyXCLBin::_runOnlyNpuInstructions) - .def("wait", &PyXCLBin::wait, "timeout"_a = py::none()) + .def("wait", &PyXCLBin::wait, "timeout"_a = nb::none()) .def( "mmap_buffers", - [](PyXCLBin &self, const std::vector> &shapes, - const py::object &npFormat) { - auto npy = py::module_::import("numpy"); + [](PyXCLBin &self, const std::vector> &shapes, + const nb::object &npFormat) { + auto npy = nb::module_::import_("numpy"); if (npFormat.is(npy.attr("int16"))) return self.mmapBuffers(shapes); if (npFormat.is(npy.attr("int32"))) @@ -167,7 +170,7 @@ PYBIND11_MODULE(_xrt, m) { if (npFormat.is(npy.attr("float64"))) return self.mmapBuffers(shapes); throw std::runtime_error("unsupported np format: " + - py::repr(npFormat).cast()); + nb::cast(nb::repr(npFormat))); }, "shapes"_a, "np_format"_a) .def("_get_buffer_host_address", [](PyXCLBin &self, size_t idx) { From 9643a0fd81e69481d47a7ed7c2828478e8de3ba1 Mon Sep 17 00:00:00 2001 From: erwei-xilinx Date: Mon, 13 Jan 2025 15:44:48 -0800 Subject: [PATCH 2/6] Fixup txn2mlir error Co-authored-by: Jeff Fifield --- python/AIEMLIRModule.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/AIEMLIRModule.cpp b/python/AIEMLIRModule.cpp index 588a90f7fe2..51c8f0441b7 100644 --- a/python/AIEMLIRModule.cpp +++ b/python/AIEMLIRModule.cpp @@ -115,8 +115,8 @@ NB_MODULE(_aie, m) { m.def( "transaction_binary_to_mlir", [](MlirContext ctx, nb::bytes bytes) { - std::string s = static_cast(bytes.data()); - MlirStringRef bin = {s.data(), s.size()}; + MlirStringRef bin = {static_cast(bytes.data()), + bytes.size()}; return aieTranslateBinaryToTxn(ctx, bin); }, "ctx"_a, "binary"_a); From 693645c13449bb58180c26d422d2f5864dc9a4dd Mon Sep 17 00:00:00 2001 From: erwei-xilinx Date: Mon, 13 Jan 2025 17:33:52 -0800 Subject: [PATCH 3/6] Remove pybind from requirements.txt --- python/requirements.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/python/requirements.txt b/python/requirements.txt index 5d2f0799c80..d9befb1f82e 100644 --- a/python/requirements.txt +++ b/python/requirements.txt @@ -7,7 +7,6 @@ lit numpy>=1.19.5, <=1.26 # 2.1 would be nice for typing improvements, but it doesn't seem to work well with pyxrt pandas psutil -pybind11>=2.9.0, <=2.10.3 rich setuptools wheel From ef04bab63641d4ef40de650aed97c6f11de4873b Mon Sep 17 00:00:00 2001 From: erwei-xilinx Date: Mon, 13 Jan 2025 17:42:18 -0800 Subject: [PATCH 4/6] Remove pybind dependence --- python/CMakeLists.txt | 4 ---- 1 file changed, 4 deletions(-) diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 9b737f10ef6..1c59f002504 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -151,12 +151,9 @@ if (AIE_ENABLE_PYTHON_PASSES) # Missed one... get_target_property(RegisterEverythingSources MLIRPythonExtension.RegisterEverything INTERFACE_SOURCES) - set(PYBINDINGS_SRC "${MLIR_INSTALL_PREFIX}/src/python/MLIRPythonExtension.Core") - set(_py_srcs ${CMAKE_CURRENT_SOURCE_DIR}/AIEMLIRModule.cpp # Python passes - ${CMAKE_CURRENT_SOURCE_DIR}/PybindTypes.cpp ${CMAKE_CURRENT_SOURCE_DIR}/PythonPass.cpp ${CMAKE_CURRENT_SOURCE_DIR}/RouterPass.cpp @@ -197,7 +194,6 @@ if (AIE_ENABLE_PYTHON_PASSES) ) target_include_directories( AIEPythonExtensions.MLIR - INTERFACE $ ) if (AIE_ENABLE_XRT_PYTHON_BINDINGS) target_include_directories(AIEPythonExtensions.MLIR INTERFACE ${XRT_INCLUDE_DIR}) From df4aa896b09ddfa9375b12a589bacf7464b23a8d Mon Sep 17 00:00:00 2001 From: erwei-xilinx Date: Mon, 13 Jan 2025 17:52:53 -0800 Subject: [PATCH 5/6] Revert "Remove pybind dependence" This reverts commit ef04bab63641d4ef40de650aed97c6f11de4873b. --- python/CMakeLists.txt | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 1c59f002504..9b737f10ef6 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -151,9 +151,12 @@ if (AIE_ENABLE_PYTHON_PASSES) # Missed one... get_target_property(RegisterEverythingSources MLIRPythonExtension.RegisterEverything INTERFACE_SOURCES) + set(PYBINDINGS_SRC "${MLIR_INSTALL_PREFIX}/src/python/MLIRPythonExtension.Core") + set(_py_srcs ${CMAKE_CURRENT_SOURCE_DIR}/AIEMLIRModule.cpp # Python passes + ${CMAKE_CURRENT_SOURCE_DIR}/PybindTypes.cpp ${CMAKE_CURRENT_SOURCE_DIR}/PythonPass.cpp ${CMAKE_CURRENT_SOURCE_DIR}/RouterPass.cpp @@ -194,6 +197,7 @@ if (AIE_ENABLE_PYTHON_PASSES) ) target_include_directories( AIEPythonExtensions.MLIR + INTERFACE $ ) if (AIE_ENABLE_XRT_PYTHON_BINDINGS) target_include_directories(AIEPythonExtensions.MLIR INTERFACE ${XRT_INCLUDE_DIR}) From 5daa4668fa4570588bc24dfc777d29648dbfb2cc Mon Sep 17 00:00:00 2001 From: erwei-xilinx Date: Mon, 13 Jan 2025 17:53:58 -0800 Subject: [PATCH 6/6] Revert "Remove pybind from requirements.txt" This reverts commit 693645c13449bb58180c26d422d2f5864dc9a4dd. --- python/requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/python/requirements.txt b/python/requirements.txt index d9befb1f82e..5d2f0799c80 100644 --- a/python/requirements.txt +++ b/python/requirements.txt @@ -7,6 +7,7 @@ lit numpy>=1.19.5, <=1.26 # 2.1 would be nice for typing improvements, but it doesn't seem to work well with pyxrt pandas psutil +pybind11>=2.9.0, <=2.10.3 rich setuptools wheel