Port compiler and iree-dialects to nanobind (#19790)

Upstream MLIR ported its Python bindings from pybind11 to nanobind and
so do we. With this we have a single library for the binding across the
compiler, iree-dialects and the runtime.

The macOS build is expected to still fail, see #19591, but as upstream
switched to build the bindings for the internal dialects with nanobind
it does anyway without reverts carried forward in our LLVM fork.

CMakeLists.txt

@@ -646,7 +646,7 @@ if(IREE_BUILD_PYTHON_BINDINGS)
   # Some parts of the build use FindPython instead of FindPython3. Why? No
   # one knows, but they are different. So make sure to bootstrap this one too.
   # Not doing this here risks them diverging, which on multi-Python systems,
-  # can be troublesome. Note that pybind11 and nanobind require FindPython.
+  # can be troublesome. Note that nanobind requires FindPython.
   set(Python_EXECUTABLE "${Python3_EXECUTABLE}")
   find_package(Python 3.9 COMPONENTS Interpreter Development.Module NumPy REQUIRED)
 elseif(IREE_BUILD_COMPILER OR IREE_BUILD_TESTS)
@@ -785,30 +785,6 @@ if(IREE_BUILD_PYTHON_BINDINGS OR IREE_BUILD_COMPILER)
   FetchContent_MakeAvailable(nanobind)
 endif()
 
-# Both the IREE and MLIR Python bindings require pybind11. We initialize it here
-# at the top level so that everything uses ours consistently.
-if(IREE_BUILD_PYTHON_BINDINGS AND IREE_BUILD_COMPILER)
-  set(pybind11_VERSION 2.13.6)
-  include(FetchContent)
-  FetchContent_Declare(
-      pybind11
-      GIT_REPOSITORY https://github.com/pybind/pybind11
-      GIT_TAG        v${pybind11_VERSION}
-  )
-  set(PYBIND11_FINDPYTHON ON)
-  FetchContent_MakeAvailable(pybind11)
-  # pybind11 source fetches do not include find_package integration, which is
-  # a shame since sub-projects can require that to work. If we were using
-  # CMake 3.24, we could just add OVERRIDE_FIND_PACKAGE to the
-  # FetchContent_Declare call above and it would take care of doing the
-  # following to let subsequent sub-project find_package calls to resolve
-  # successfully.
-  set(pybind11_DIR "${pybind11_BINARY_DIR}")
-  file(WRITE "${pybind11_BINARY_DIR}/pybind11Config.cmake" "")
-  file(WRITE "${pybind11_BINARY_DIR}/pybind11ConfigVersion.cmake"
-       "set(PACKAGE_VERSION ${pybind11_VERSION})\nset(PACKAGE_VERSION_COMPATIBLE TRUE)")
-endif()
-
 if(NOT IREE_BUILD_COMPILER)
   message(STATUS "Not adding LLVM/MLIR because the configuration does not require it")
 else()

build_tools/python_deploy/README.md

@@ -93,7 +93,7 @@ Python 3.10.4
 # Option A: Build like a normal dev setup (i.e. if allergic to Python
 # packaging and to triage issues that do not implicate that).
 [root@c8f6d0041d79 ]# cd /work/iree
-[root@c8f6d0041d79 iree]# pip install wheel cmake ninja pybind11 numpy
+[root@c8f6d0041d79 iree]# pip install wheel cmake ninja numpy
 [root@c8f6d0041d79 iree]# cmake -GNinja -B ../iree-build/ -S . -DCMAKE_BUILD_TYPE=Release -DIREE_BUILD_PYTHON_BINDINGS=ON
 [root@c8f6d0041d79 iree]# cd ../iree-build/
 [root@c8f6d0041d79 iree-build]# ninja

compiler/bindings/python/CMakeLists.txt

@@ -176,6 +176,7 @@ declare_mlir_python_sources(IREECompilerPythonExtensions)
 declare_mlir_python_extension(IREECompilerPythonExtensions.Registration
   MODULE_NAME _site_initialize_0
   ADD_TO_PARENT IREECompilerPythonExtensions
+  PYTHON_BINDINGS_LIBRARY nanobind
   SOURCES
     IREECompilerRegistration.cpp
   EMBED_CAPI_LINK_LIBS
@@ -187,6 +188,7 @@ declare_mlir_python_extension(IREECompilerPythonExtensions.Registration
 declare_mlir_python_extension(IREECompilerPythonExtensions.CompilerDialects
   MODULE_NAME _ireeCompilerDialects
   ADD_TO_PARENT IREECompilerPythonExtensions
+  PYTHON_BINDINGS_LIBRARY nanobind
   SOURCES
     IREECompilerDialectsModule.cpp
   EMBED_CAPI_LINK_LIBS

compiler/bindings/python/IREECompilerDialectsModule.cpp

@@ -11,15 +11,17 @@
 #include "iree/compiler/dialects/iree_gpu.h"
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/IR.h"
-#include "mlir/Bindings/Python/PybindAdaptors.h"
+#include "mlir/Bindings/Python/Nanobind.h"
+#include "mlir/Bindings/Python/NanobindAdaptors.h"
 
 static const char *kCodegenModuleImportPath =
     MAKE_MLIR_PYTHON_QUALNAME("dialects.iree_codegen");
 static const char *kGpuModuleImportPath =
     MAKE_MLIR_PYTHON_QUALNAME("dialects.iree_gpu");
 
-namespace py = pybind11;
-using namespace mlir::python::adaptors;
+namespace py = nanobind;
+using namespace nanobind::literals;
+using namespace mlir::python::nanobind_adaptors;
 
 static std::vector<MlirOperation>
 ireeCodegenGetExecutableVariantOpsBinding(MlirModule module) {
@@ -39,7 +41,7 @@ ireeCodegenQueryMMAIntrinsicsBinding(MlirOperation op) {
   ireeCodegenQueryMMAIntrinsics(op, &numMMAs, mmaIntrinsics.data());
 
   py::object mmaIntrinsicEnum =
-      py::module_::import(kGpuModuleImportPath).attr("MMAIntrinsic");
+      py::module_::import_(kGpuModuleImportPath).attr("MMAIntrinsic");
   std::vector<py::object> mmaList(numMMAs);
   for (size_t i = 0; i < numMMAs; ++i) {
     mmaList[i] = mmaIntrinsicEnum(mmaIntrinsics[i]);
@@ -48,7 +50,7 @@ ireeCodegenQueryMMAIntrinsicsBinding(MlirOperation op) {
   return mmaList;
 }
 
-PYBIND11_MODULE(_ireeCompilerDialects, m) {
+NB_MODULE(_ireeCompilerDialects, m) {
   m.doc() = "iree-compiler dialects python extension";
 
   auto iree_codegen_module =
@@ -75,7 +77,7 @@ PYBIND11_MODULE(_ireeCompilerDialects, m) {
       .def_property_readonly("value", [](MlirAttribute self) -> py::object {
         uint32_t rawValue =
             ireeCodegenDispatchLoweringPassPipelineAttrGetValue(self);
-        return py::module_::import(kCodegenModuleImportPath)
+        return py::module_::import_(kCodegenModuleImportPath)
             .attr("DispatchLoweringPassPipeline")(rawValue);
       });
 
@@ -204,7 +206,7 @@ PYBIND11_MODULE(_ireeCompilerDialects, m) {
                              ireeGPUReorderWorkgroupsStrategyAttrGetValue)
       .def_property_readonly("value", [](MlirAttribute self) -> py::object {
         uint32_t rawValue = ireeGPUReorderWorkgroupsStrategyAttrGetValue(self);
-        return py::module_::import(kGpuModuleImportPath)
+        return py::module_::import_(kGpuModuleImportPath)
             .attr("ReorderWorkgroupsStrategy")(rawValue);
       });
 
@@ -296,7 +298,7 @@ PYBIND11_MODULE(_ireeCompilerDialects, m) {
                              [](MlirAttribute self) -> py::object {
                                uint32_t rawValue =
                                    ireeGPUMMAIntrinsicAttrGetValue(self);
-                               return py::module_::import(kGpuModuleImportPath)
+                               return py::module_::import_(kGpuModuleImportPath)
                                    .attr("MMAIntrinsic")(rawValue);
                              })
       .def_property_readonly("mma", [](MlirAttribute self) -> MlirAttribute {

compiler/bindings/python/IREECompilerRegistration.cpp

@@ -6,10 +6,10 @@
 
 #include "iree/compiler/mlir_interop.h"
 #include "mlir-c/IR.h"
-#include "mlir/Bindings/Python/PybindAdaptors.h"
+#include "mlir/Bindings/Python/NanobindAdaptors.h"
 
-namespace py = pybind11;
-using namespace mlir::python::adaptors;
+namespace py = nanobind;
+using namespace mlir::python::nanobind_adaptors;
 
 namespace {
 
@@ -21,14 +21,14 @@ class GlobalInitializer {
 
 } // namespace
 
-PYBIND11_MODULE(_site_initialize_0, m, py::mod_gil_not_used()) {
+NB_MODULE(_site_initialize_0, m) {
   m.doc() = "iree-compile registration";
 
   // Make sure that GlobalInitialize and GlobalShutdown are called with module
   // lifetime.
   py::class_<GlobalInitializer>(m, "_GlobalInitializer");
   m.attr("_global_init_hook") =
-      py::cast(new GlobalInitializer, py::return_value_policy::take_ownership);
+      py::cast(new GlobalInitializer, py::rv_policy::take_ownership);
 
   m.def("register_dialects", [](MlirDialectRegistry registry) {
     ireeCompilerRegisterDialects(registry);

llvm-external-projects/iree-dialects/python/CMakeLists.txt

@@ -40,6 +40,7 @@ declare_mlir_dialect_extension_python_bindings(
 declare_mlir_python_extension(IREEDialectsPythonExtensions.Main
   MODULE_NAME _ireeDialects
   ADD_TO_PARENT IREEDialectsPythonExtensions
+  PYTHON_BINDINGS_LIBRARY nanobind
   SOURCES
     IREEDialectsModule.cpp
   EMBED_CAPI_LINK_LIBS

llvm-external-projects/iree-dialects/python/IREEDialectsModule.cpp

@@ -9,15 +9,15 @@
 #include "mlir-c/BuiltinTypes.h"
 #include "mlir-c/Diagnostics.h"
 #include "mlir-c/RegisterEverything.h"
-#include "mlir/Bindings/Python/PybindAdaptors.h"
+#include "mlir/Bindings/Python/NanobindAdaptors.h"
 
-namespace py = pybind11;
-using namespace mlir::python::adaptors;
+namespace py = nanobind;
+using namespace mlir::python::nanobind_adaptors;
 
-PYBIND11_MODULE(_ireeDialects, m, py::mod_gil_not_used()) {
+NB_MODULE(_ireeDialects, m) {
   m.doc() = "iree-dialects main python extension";
 
-  auto irModule = py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"));
+  auto irModule = py::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"));
   auto typeClass = irModule.attr("Type");
 
   //===--------------------------------------------------------------------===//