Revert "Revert "[mlir python] Add nanobind support (#119232)

Reverts revert #118517 after (hopefully) fixing builders
(llvm/llvm-zorg#328,
llvm/llvm-zorg#327)

This reverts commit 61bf308.

Author: makslevental

mlir/cmake/modules/AddMLIRPython.cmake

@@ -114,10 +114,11 @@ endfunction()
 #   EMBED_CAPI_LINK_LIBS: Dependent CAPI libraries that this extension depends
 #     on. These will be collected for all extensions and put into an
 #     aggregate dylib that is linked against.
+#   PYTHON_BINDINGS_LIBRARY: Either pybind11 or nanobind.
 function(declare_mlir_python_extension name)
   cmake_parse_arguments(ARG
     ""
-    "ROOT_DIR;MODULE_NAME;ADD_TO_PARENT"
+    "ROOT_DIR;MODULE_NAME;ADD_TO_PARENT;PYTHON_BINDINGS_LIBRARY"
     "SOURCES;PRIVATE_LINK_LIBS;EMBED_CAPI_LINK_LIBS"
     ${ARGN})
 
@@ -126,15 +127,20 @@ function(declare_mlir_python_extension name)
   endif()
   set(_install_destination "src/python/${name}")
 
+  if(NOT ARG_PYTHON_BINDINGS_LIBRARY)
+    set(ARG_PYTHON_BINDINGS_LIBRARY "pybind11")
+  endif()
+
   add_library(${name} INTERFACE)
   set_target_properties(${name} PROPERTIES
     # Yes: Leading-lowercase property names are load bearing and the recommended
     # way to do this: https://gitlab.kitware.com/cmake/cmake/-/issues/19261
-    EXPORT_PROPERTIES "mlir_python_SOURCES_TYPE;mlir_python_EXTENSION_MODULE_NAME;mlir_python_EMBED_CAPI_LINK_LIBS;mlir_python_DEPENDS"
+    EXPORT_PROPERTIES "mlir_python_SOURCES_TYPE;mlir_python_EXTENSION_MODULE_NAME;mlir_python_EMBED_CAPI_LINK_LIBS;mlir_python_DEPENDS;mlir_python_BINDINGS_LIBRARY"
     mlir_python_SOURCES_TYPE extension
     mlir_python_EXTENSION_MODULE_NAME "${ARG_MODULE_NAME}"
     mlir_python_EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
     mlir_python_DEPENDS ""
+    mlir_python_BINDINGS_LIBRARY "${ARG_PYTHON_BINDINGS_LIBRARY}"
   )
 
   # Set the interface source and link_libs properties of the target
@@ -223,12 +229,14 @@ function(add_mlir_python_modules name)
     elseif(_source_type STREQUAL "extension")
       # Native CPP extension.
       get_target_property(_module_name ${sources_target} mlir_python_EXTENSION_MODULE_NAME)
+      get_target_property(_bindings_library ${sources_target} mlir_python_BINDINGS_LIBRARY)
       # Transform relative source to based on root dir.
       set(_extension_target "${modules_target}.extension.${_module_name}.dso")
       add_mlir_python_extension(${_extension_target} "${_module_name}"
         INSTALL_COMPONENT ${modules_target}
         INSTALL_DIR "${ARG_INSTALL_PREFIX}/_mlir_libs"
         OUTPUT_DIRECTORY "${ARG_ROOT_PREFIX}/_mlir_libs"
+        PYTHON_BINDINGS_LIBRARY ${_bindings_library}
         LINK_LIBS PRIVATE
           ${sources_target}
           ${ARG_COMMON_CAPI_LINK_LIBS}
@@ -634,7 +642,7 @@ endfunction()
 function(add_mlir_python_extension libname extname)
   cmake_parse_arguments(ARG
   ""
-  "INSTALL_COMPONENT;INSTALL_DIR;OUTPUT_DIRECTORY"
+  "INSTALL_COMPONENT;INSTALL_DIR;OUTPUT_DIRECTORY;PYTHON_BINDINGS_LIBRARY"
   "SOURCES;LINK_LIBS"
   ${ARGN})
   if(ARG_UNPARSED_ARGUMENTS)
@@ -644,9 +652,16 @@ function(add_mlir_python_extension libname extname)
   # The actual extension library produces a shared-object or DLL and has
   # sources that must be compiled in accordance with pybind11 needs (RTTI and
   # exceptions).
-  pybind11_add_module(${libname}
-    ${ARG_SOURCES}
-  )
+  if(NOT DEFINED ARG_PYTHON_BINDINGS_LIBRARY OR ARG_PYTHON_BINDINGS_LIBRARY STREQUAL "pybind11")
+    pybind11_add_module(${libname}
+      ${ARG_SOURCES}
+    )
+  elseif(ARG_PYTHON_BINDINGS_LIBRARY STREQUAL "nanobind")
+    nanobind_add_module(${libname}
+      NB_DOMAIN mlir
+      ${ARG_SOURCES}
+    )
+  endif()
 
   # The extension itself must be compiled with RTTI and exceptions enabled.
   # Also, some warning classes triggered by pybind11 are disabled.

mlir/cmake/modules/MLIRDetectPythonEnv.cmake

@@ -21,6 +21,12 @@ macro(mlir_configure_python_dev_packages)
 
     find_package(Python3 ${LLVM_MINIMUM_PYTHON_VERSION}
       COMPONENTS Interpreter ${_python_development_component} REQUIRED)
+
+    # It's a little silly to detect Python a second time, but nanobind's cmake
+    # code looks for Python_ not Python3_.
+    find_package(Python ${LLVM_MINIMUM_PYTHON_VERSION}
+      COMPONENTS Interpreter ${_python_development_component} REQUIRED)
+
     unset(_python_development_component)
     message(STATUS "Found python include dirs: ${Python3_INCLUDE_DIRS}")
     message(STATUS "Found python libraries: ${Python3_LIBRARIES}")
@@ -31,6 +37,13 @@ macro(mlir_configure_python_dev_packages)
     message(STATUS "Python prefix = '${PYTHON_MODULE_PREFIX}', "
                   "suffix = '${PYTHON_MODULE_SUFFIX}', "
                   "extension = '${PYTHON_MODULE_EXTENSION}")
+
+    mlir_detect_nanobind_install()
+    find_package(nanobind 2.2 CONFIG REQUIRED)
+    message(STATUS "Found nanobind v${nanobind_VERSION}: ${nanobind_INCLUDE_DIR}")
+    message(STATUS "Python prefix = '${PYTHON_MODULE_PREFIX}', "
+                  "suffix = '${PYTHON_MODULE_SUFFIX}', "
+                  "extension = '${PYTHON_MODULE_EXTENSION}")
   endif()
 endmacro()
 
@@ -58,3 +71,29 @@ function(mlir_detect_pybind11_install)
     set(pybind11_DIR "${PACKAGE_DIR}" PARENT_SCOPE)
   endif()
 endfunction()
+
+
+# Detects a nanobind package installed in the current python environment
+# and sets variables to allow it to be found. This allows nanobind to be
+# installed via pip, which typically yields a much more recent version than
+# the OS install, which will be available otherwise.
+function(mlir_detect_nanobind_install)
+  if(nanobind_DIR)
+    message(STATUS "Using explicit nanobind cmake directory: ${nanobind_DIR} (-Dnanobind_DIR to change)")
+  else()
+    message(STATUS "Checking for nanobind in python path...")
+    execute_process(
+      COMMAND "${Python3_EXECUTABLE}"
+      -c "import nanobind;print(nanobind.cmake_dir(), end='')"
+      WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+      RESULT_VARIABLE STATUS
+      OUTPUT_VARIABLE PACKAGE_DIR
+      ERROR_QUIET)
+    if(NOT STATUS EQUAL "0")
+      message(STATUS "not found (install via 'pip install nanobind' or set nanobind_DIR)")
+      return()
+    endif()
+    message(STATUS "found (${PACKAGE_DIR})")
+    set(nanobind_DIR "${PACKAGE_DIR}" PARENT_SCOPE)
+  endif()
+endfunction()

mlir/docs/Bindings/Python.md

@@ -1138,12 +1138,14 @@ attributes and types must connect to the relevant C APIs for building and
 inspection, which must be provided first. Bindings for `Attribute` and `Type`
 subclasses can be defined using
 [`include/mlir/Bindings/Python/PybindAdaptors.h`](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/PybindAdaptors.h)
-utilities that mimic pybind11 API for defining functions and properties. These
-bindings are to be included in a separate pybind11 module. The utilities also
-provide automatic casting between C API handles `MlirAttribute` and `MlirType`
-and their Python counterparts so that the C API handles can be used directly in
-binding implementations. The methods and properties provided by the bindings
-should follow the principles discussed above.
+or
+[`include/mlir/Bindings/Python/NanobindAdaptors.h`](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h)
+utilities that mimic pybind11/nanobind API for defining functions and
+properties. These bindings are to be included in a separate module. The
+utilities also provide automatic casting between C API handles `MlirAttribute`
+and `MlirType` and their Python counterparts so that the C API handles can be
+used directly in binding implementations. The methods and properties provided by
+the bindings should follow the principles discussed above.
 
 The attribute and type bindings for a dialect can be located in
 `lib/Bindings/Python/Dialect<Name>.cpp` and should be compiled into a separate
@@ -1179,7 +1181,9 @@ make the passes available along with the dialect.
 Dialect functionality other than IR objects or passes, such as helper functions,
 can be exposed to Python similarly to attributes and types. C API is expected to
 exist for this functionality, which can then be wrapped using pybind11 and
-`[include/mlir/Bindings/Python/PybindAdaptors.h](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/PybindAdaptors.h)`
+`[include/mlir/Bindings/Python/PybindAdaptors.h](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/PybindAdaptors.h)`,
+or nanobind and
+`[include/mlir/Bindings/Python/NanobindAdaptors.h](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h)`
 utilities to connect to the rest of Python API. The bindings can be located in a
-separate pybind11 module or in the same module as attributes and types, and
+separate module or in the same module as attributes and types, and
 loaded along with the dialect.

mlir/examples/standalone/python/CMakeLists.txt

@@ -17,18 +17,32 @@ declare_mlir_dialect_python_bindings(
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir_standalone"
   TD_FILE dialects/StandaloneOps.td
   SOURCES
-    dialects/standalone.py
+    dialects/standalone_pybind11.py
+    dialects/standalone_nanobind.py
   DIALECT_NAME standalone)
 
-declare_mlir_python_extension(StandalonePythonSources.Extension
-  MODULE_NAME _standaloneDialects
+
+declare_mlir_python_extension(StandalonePythonSources.Pybind11Extension
+  MODULE_NAME _standaloneDialectsPybind11
+  ADD_TO_PARENT StandalonePythonSources
+  SOURCES
+    StandaloneExtensionPybind11.cpp
+  EMBED_CAPI_LINK_LIBS
+    StandaloneCAPI
+  PYTHON_BINDINGS_LIBRARY pybind11
+)
+
+declare_mlir_python_extension(StandalonePythonSources.NanobindExtension
+  MODULE_NAME _standaloneDialectsNanobind
   ADD_TO_PARENT StandalonePythonSources
   SOURCES
-    StandaloneExtension.cpp
+    StandaloneExtensionNanobind.cpp
   EMBED_CAPI_LINK_LIBS
     StandaloneCAPI
+  PYTHON_BINDINGS_LIBRARY nanobind
 )
 
+
 ################################################################################
 # Common CAPI
 ################################################################################

mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp

@@ -0,0 +1,35 @@
+//===- StandaloneExtension.cpp - Extension module -------------------------===//
+//
+// This is the nanobind version of the example module. There is also a pybind11
+// example in StandaloneExtensionPybind11.cpp.
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include <nanobind/nanobind.h>
+
+#include "Standalone-c/Dialects.h"
+#include "mlir/Bindings/Python/NanobindAdaptors.h"
+
+namespace nb = nanobind;
+
+NB_MODULE(_standaloneDialectsNanobind, m) {
+  //===--------------------------------------------------------------------===//
+  // standalone dialect
+  //===--------------------------------------------------------------------===//
+  auto standaloneM = m.def_submodule("standalone");
+
+  standaloneM.def(
+      "register_dialect",
+      [](MlirContext context, bool load) {
+        MlirDialectHandle handle = mlirGetDialectHandle__standalone__();
+        mlirDialectHandleRegisterDialect(handle, context);
+        if (load) {
+          mlirDialectHandleLoadDialect(handle, context);
+        }
+      },
+      nb::arg("context").none() = nb::none(), nb::arg("load") = true);
+}

mlir/examples/standalone/python/StandaloneExtension.cpp renamed to mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp

@@ -1,4 +1,7 @@
-//===- StandaloneExtension.cpp - Extension module -------------------------===//
+//===- StandaloneExtensionPybind11.cpp - Extension module -----------------===//
+//
+// This is the pybind11 version of the example module. There is also a nanobind
+// example in StandaloneExtensionNanobind.cpp.
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -11,7 +14,7 @@
 
 using namespace mlir::python::adaptors;
 
-PYBIND11_MODULE(_standaloneDialects, m) {
+PYBIND11_MODULE(_standaloneDialectsPybind11, m) {
   //===--------------------------------------------------------------------===//
   // standalone dialect
   //===--------------------------------------------------------------------===//

mlir/examples/standalone/python/mlir_standalone/dialects/standalone.py renamed to mlir/examples/standalone/python/mlir_standalone/dialects/standalone_nanobind.py

@@ -3,4 +3,4 @@
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from ._standalone_ops_gen import *
-from .._mlir_libs._standaloneDialects.standalone import *
+from .._mlir_libs._standaloneDialectsNanobind.standalone import *

mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py

@@ -0,0 +1,6 @@
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from ._standalone_ops_gen import *
+from .._mlir_libs._standaloneDialectsPybind11.standalone import *

mlir/examples/standalone/test/python/smoketest.py

@@ -1,7 +1,17 @@
-# RUN: %python %s | FileCheck %s
+# RUN: %python %s pybind11 | FileCheck %s
+# RUN: %python %s nanobind | FileCheck %s
 
+import sys
 from mlir_standalone.ir import *
-from mlir_standalone.dialects import builtin as builtin_d, standalone as standalone_d
+from mlir_standalone.dialects import builtin as builtin_d
+
+if sys.argv[1] == "pybind11":
+    from mlir_standalone.dialects import standalone_pybind11 as standalone_d
+elif sys.argv[1] == "nanobind":
+    from mlir_standalone.dialects import standalone_nanobind as standalone_d
+else:
+    raise ValueError("Expected either pybind11 or nanobind as arguments")
+
 
 with Context():
     standalone_d.register_dialect()

mlir/include/mlir/Bindings/Python/Diagnostics.h

@@ -0,0 +1,59 @@
+//===- Diagnostics.h - Helpers for diagnostics in Python bindings ---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_BINDINGS_PYTHON_DIAGNOSTICS_H
+#define MLIR_BINDINGS_PYTHON_DIAGNOSTICS_H
+
+#include <cassert>
+#include <string>
+
+#include "mlir-c/Diagnostics.h"
+#include "mlir-c/IR.h"
+#include "llvm/ADT/StringRef.h"
+
+namespace mlir {
+namespace python {
+
+/// 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
+
+#endif // MLIR_BINDINGS_PYTHON_DIAGNOSTICS_H