[Mlir-commits] [llvm] [mlir] Revert "[MLIR][Python] use `FetchContent_Declare` for nanobind and remove pybind (#161230)" (PR #162309)
Maksim Levental
llvmlistbot at llvm.org
Tue Oct 7 08:50:17 PDT 2025
https://github.com/makslevental created https://github.com/llvm/llvm-project/pull/162309
This reverts commit 84a214856ad989f37af19f5e8aaa9ec2346dde6f.
This gives us more time to work out the alternative and also people to migrate
>From 0d8252ff23d7ff94ec879094da104798bc58daf4 Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Tue, 7 Oct 2025 08:45:00 -0700
Subject: [PATCH] Revert "[MLIR][Python] use `FetchContent_Declare` for
nanobind and remove pybind (#161230)"
This reverts commit 84a214856ad989f37af19f5e8aaa9ec2346dde6f.
---
.ci/all_requirements.txt | 8 +
mlir/cmake/modules/AddMLIRPython.cmake | 127 ++--
mlir/cmake/modules/MLIRDetectPythonEnv.cmake | 91 ++-
mlir/docs/Dialects/Linalg/OpDSL.md | 4 +-
mlir/examples/standalone/pyproject.toml | 4 +-
.../examples/standalone/python/CMakeLists.txt | 18 +
.../python/StandaloneExtensionPybind11.cpp | 38 ++
.../dialects/standalone_pybind11.py | 6 +
.../standalone/test/python/smoketest.py | 11 +-
.../mlir/Bindings/Python/PybindAdaptors.h | 616 ++++++++++++++++++
mlir/python/CMakeLists.txt | 61 +-
mlir/python/mlir/dialects/python_test.py | 11 +-
mlir/python/requirements.txt | 2 +
mlir/test/python/dialects/python_test.py | 31 +-
mlir/test/python/lib/CMakeLists.txt | 1 +
.../python/lib/PythonTestModulePybind11.cpp | 118 ++++
.../update_core_linalg_named_ops.sh.in | 2 +-
.../llvm-project-overlay/mlir/BUILD.bazel | 54 +-
18 files changed, 1082 insertions(+), 121 deletions(-)
create mode 100644 mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
create mode 100644 mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py
create mode 100644 mlir/include/mlir/Bindings/Python/PybindAdaptors.h
create mode 100644 mlir/test/python/lib/PythonTestModulePybind11.cpp
diff --git a/.ci/all_requirements.txt b/.ci/all_requirements.txt
index 313ab1076bb26..ac9682a09bec1 100644
--- a/.ci/all_requirements.txt
+++ b/.ci/all_requirements.txt
@@ -194,6 +194,10 @@ ml-dtypes==0.5.1 ; python_version < "3.13" \
--hash=sha256:d13755f8e8445b3870114e5b6240facaa7cb0c3361e54beba3e07fa912a6e12b \
--hash=sha256:fd918d4e6a4e0c110e2e05be7a7814d10dc1b95872accbf6512b80a109b71ae1
# via -r mlir/python/requirements.txt
+nanobind==2.9.2 \
+ --hash=sha256:c37957ffd5eac7eda349cff3622ecd32e5ee1244ecc912c99b5bc8188bafd16e \
+ --hash=sha256:e7608472de99d375759814cab3e2c94aba3f9ec80e62cfef8ced495ca5c27d6e
+ # via -r mlir/python/requirements.txt
numpy==2.0.2 \
--hash=sha256:0123ffdaa88fa4ab64835dcbde75dcdf89c453c922f18dced6e27c90d1d0ec5a \
--hash=sha256:11a76c372d1d37437857280aa142086476136a8c0f373b2e648ab2c8f18fb195 \
@@ -295,6 +299,10 @@ pyasn1-modules==0.4.2 \
--hash=sha256:29253a9207ce32b64c3ac6600edc75368f98473906e8fd1043bd6b5b1de2c14a \
--hash=sha256:677091de870a80aae844b1ca6134f54652fa2c8c5a52aa396440ac3106e941e6
# via google-auth
+pybind11==2.13.6 \
+ --hash=sha256:237c41e29157b962835d356b370ededd57594a26d5894a795960f0047cb5caf5 \
+ --hash=sha256:ba6af10348c12b24e92fa086b39cfba0eff619b61ac77c406167d813b096d39a
+ # via -r mlir/python/requirements.txt
pyyaml==6.0.1 \
--hash=sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5 \
--hash=sha256:062582fca9fabdd2c8b54a3ef1c978d786e0f6b3a1510e0ac93ef59e0ddae2bc \
diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index ea34f948cd3dd..fa6aec8a603a9 100644
--- a/mlir/cmake/modules/AddMLIRPython.cmake
+++ b/mlir/cmake/modules/AddMLIRPython.cmake
@@ -123,12 +123,12 @@ function(mlir_generate_type_stubs)
"IMPORT_PATHS;DEPENDS_TARGETS;OUTPUTS;DEPENDS_TARGET_SRC_DEPS"
${ARGN})
- # for people installing a distro (e.g., pip install) of nanobind
+ # for people doing find_package(nanobind)
if(EXISTS ${nanobind_DIR}/../src/stubgen.py)
set(NB_STUBGEN "${nanobind_DIR}/../src/stubgen.py")
elseif(EXISTS ${nanobind_DIR}/../stubgen.py)
set(NB_STUBGEN "${nanobind_DIR}/../stubgen.py")
- # for people using nanobind git source tree (e.g., FetchContent_Declare and FetchContent_MakeAvailable)
+ # for people using FetchContent_Declare and FetchContent_MakeAvailable
elseif(EXISTS ${nanobind_SOURCE_DIR}/src/stubgen.py)
set(NB_STUBGEN "${nanobind_SOURCE_DIR}/src/stubgen.py")
elseif(EXISTS ${nanobind_SOURCE_DIR}/stubgen.py)
@@ -226,10 +226,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})
@@ -238,15 +239,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
@@ -335,12 +341,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}
@@ -745,7 +753,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)
@@ -753,7 +761,7 @@ function(add_mlir_python_extension libname extname)
endif()
# The extension itself must be compiled with RTTI and exceptions enabled.
- # Also, some warning classes triggered by nanobind are disabled.
+ # Also, some warning classes triggered by pybind11 are disabled.
set(eh_rtti_enable)
if (MSVC)
set(eh_rtti_enable /EHsc /GR)
@@ -761,53 +769,62 @@ function(add_mlir_python_extension libname extname)
set(eh_rtti_enable -frtti -fexceptions)
endif ()
- nanobind_add_module(${libname}
- NB_DOMAIN ${MLIR_BINDINGS_PYTHON_NB_DOMAIN}
- FREE_THREADED
- ${ARG_SOURCES}
- )
+ # 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).
+ 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_BINDINGS_PYTHON_NB_DOMAIN}
+ FREE_THREADED
+ ${ARG_SOURCES}
+ )
- if (NOT MLIR_DISABLE_CONFIGURE_PYTHON_DEV_PACKAGES
- AND (LLVM_COMPILER_IS_GCC_COMPATIBLE OR CLANG_CL))
- # Avoid some warnings from upstream nanobind.
- # If a superproject set MLIR_DISABLE_CONFIGURE_PYTHON_DEV_PACKAGES, let
- # the super project handle compile options as it wishes.
- get_property(NB_LIBRARY_TARGET_NAME TARGET ${libname} PROPERTY LINK_LIBRARIES)
- target_compile_options(${NB_LIBRARY_TARGET_NAME}
- PRIVATE
- -Wall -Wextra -Wpedantic
- -Wno-c++98-compat-extra-semi
- -Wno-cast-qual
- -Wno-covered-switch-default
- -Wno-deprecated-literal-operator
- -Wno-nested-anon-types
- -Wno-unused-parameter
- -Wno-zero-length-array
- ${eh_rtti_enable})
-
- target_compile_options(${libname}
- PRIVATE
- -Wall -Wextra -Wpedantic
- -Wno-c++98-compat-extra-semi
- -Wno-cast-qual
- -Wno-covered-switch-default
- -Wno-deprecated-literal-operator
- -Wno-nested-anon-types
- -Wno-unused-parameter
- -Wno-zero-length-array
- ${eh_rtti_enable})
- endif()
+ if (NOT MLIR_DISABLE_CONFIGURE_PYTHON_DEV_PACKAGES
+ AND (LLVM_COMPILER_IS_GCC_COMPATIBLE OR CLANG_CL))
+ # Avoid some warnings from upstream nanobind.
+ # If a superproject set MLIR_DISABLE_CONFIGURE_PYTHON_DEV_PACKAGES, let
+ # the super project handle compile options as it wishes.
+ get_property(NB_LIBRARY_TARGET_NAME TARGET ${libname} PROPERTY LINK_LIBRARIES)
+ target_compile_options(${NB_LIBRARY_TARGET_NAME}
+ PRIVATE
+ -Wall -Wextra -Wpedantic
+ -Wno-c++98-compat-extra-semi
+ -Wno-cast-qual
+ -Wno-covered-switch-default
+ -Wno-deprecated-literal-operator
+ -Wno-nested-anon-types
+ -Wno-unused-parameter
+ -Wno-zero-length-array
+ ${eh_rtti_enable})
+
+ target_compile_options(${libname}
+ PRIVATE
+ -Wall -Wextra -Wpedantic
+ -Wno-c++98-compat-extra-semi
+ -Wno-cast-qual
+ -Wno-covered-switch-default
+ -Wno-deprecated-literal-operator
+ -Wno-nested-anon-types
+ -Wno-unused-parameter
+ -Wno-zero-length-array
+ ${eh_rtti_enable})
+ endif()
- if(APPLE)
- # NanobindAdaptors.h uses PyClassMethod_New to build `pure_subclass`es but nanobind
- # doesn't declare this API as undefined in its linker flags. So we need to declare it as such
- # for downstream users that do not do something like `-undefined dynamic_lookup`.
- # Same for the rest.
- target_link_options(${libname} PUBLIC
- "LINKER:-U,_PyClassMethod_New"
- "LINKER:-U,_PyCode_Addr2Location"
- "LINKER:-U,_PyFrame_GetLasti"
- )
+ if(APPLE)
+ # NanobindAdaptors.h uses PyClassMethod_New to build `pure_subclass`es but nanobind
+ # doesn't declare this API as undefined in its linker flags. So we need to declare it as such
+ # for downstream users that do not do something like `-undefined dynamic_lookup`.
+ # Same for the rest.
+ target_link_options(${libname} PUBLIC
+ "LINKER:-U,_PyClassMethod_New"
+ "LINKER:-U,_PyCode_Addr2Location"
+ "LINKER:-U,_PyFrame_GetLasti"
+ )
+ endif()
endif()
target_compile_options(${libname} PRIVATE ${eh_rtti_enable})
@@ -845,11 +862,11 @@ function(add_mlir_python_extension libname extname)
if(WIN32)
# On Windows, pyconfig.h (and by extension python.h) hardcode the version of the
# python library which will be used for linkage depending on the flavor of the build.
- # nanobind has a workaround which depends on the definition of Py_DEBUG (if Py_DEBUG
- # is not passed in as a compile definition, nanobind undefs _DEBUG when including
+ # pybind11 has a workaround which depends on the definition of Py_DEBUG (if Py_DEBUG
+ # is not passed in as a compile definition, pybind11 undefs _DEBUG when including
# python.h, so that the release python library would be used).
- # Since mlir uses nanobind, we can leverage their workaround by never directly
- # pyconfig.h or python.h and instead relying on the nanobind headers to include the
+ # Since mlir uses pybind11, we can leverage their workaround by never directly
+ # pyconfig.h or python.h and instead relying on the pybind11 headers to include the
# necessary python headers. This results in mlir always linking against the
# release python library via the (undocumented) cmake property Python3_LIBRARY_RELEASE.
target_link_libraries(${libname} PRIVATE ${Python3_LIBRARY_RELEASE})
diff --git a/mlir/cmake/modules/MLIRDetectPythonEnv.cmake b/mlir/cmake/modules/MLIRDetectPythonEnv.cmake
index edbad2e6b5d62..d18f8c06158b2 100644
--- a/mlir/cmake/modules/MLIRDetectPythonEnv.cmake
+++ b/mlir/cmake/modules/MLIRDetectPythonEnv.cmake
@@ -46,20 +46,81 @@ macro(mlir_configure_python_dev_packages)
message(STATUS "Found python include dirs: ${Python3_INCLUDE_DIRS}")
message(STATUS "Found python libraries: ${Python3_LIBRARIES}")
message(STATUS "Found numpy v${Python3_NumPy_VERSION}: ${Python3_NumPy_INCLUDE_DIRS}")
- message(STATUS "Python extension suffix for modules: '${Python3_SOABI}'")
- if(nanobind_DIR)
- message(STATUS "Using explicit nanobind cmake directory: ${nanobind_DIR} (-Dnanobind_DIR to change)")
- find_package(nanobind 2.9 CONFIG REQUIRED)
- else()
- include(FetchContent)
- FetchContent_Declare(
- nanobind
- GIT_REPOSITORY https://github.com/wjakob/nanobind.git
- GIT_TAG v2.9.0
- GIT_SHALLOW TRUE
- )
- FetchContent_MakeAvailable(nanobind)
- endif()
- message(STATUS "Found nanobind: ${NB_DIR}")
+ mlir_detect_pybind11_install()
+ find_package(pybind11 2.10 CONFIG REQUIRED)
+ message(STATUS "Found pybind11 v${pybind11_VERSION}: ${pybind11_INCLUDE_DIR}")
+ message(STATUS "Python prefix = '${PYTHON_MODULE_PREFIX}', "
+ "suffix = '${PYTHON_MODULE_SUFFIX}', "
+ "extension = '${PYTHON_MODULE_EXTENSION}")
+
+ mlir_detect_nanobind_install()
+ find_package(nanobind 2.9 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()
+
+# Detects a pybind11 package installed in the current python environment
+# and sets variables to allow it to be found. This allows pybind11 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_pybind11_install)
+ if(pybind11_DIR)
+ message(STATUS "Using explicit pybind11 cmake directory: ${pybind11_DIR} (-Dpybind11_DIR to change)")
+ else()
+ message(STATUS "Checking for pybind11 in python path...")
+ execute_process(
+ COMMAND "${Python3_EXECUTABLE}"
+ -c "import pybind11;print(pybind11.get_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 pybind11' or set pybind11_DIR)")
+ return()
+ endif()
+ message(STATUS "found (${PACKAGE_DIR})")
+ 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)
+ execute_process(
+ COMMAND "${Python3_EXECUTABLE}"
+ -c "import nanobind;print(nanobind.include_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()
+ set(nanobind_INCLUDE_DIR "${PACKAGE_DIR}" PARENT_SCOPE)
+ endif()
+endfunction()
diff --git a/mlir/docs/Dialects/Linalg/OpDSL.md b/mlir/docs/Dialects/Linalg/OpDSL.md
index 5d7e2740ec90b..b892bbe427a18 100644
--- a/mlir/docs/Dialects/Linalg/OpDSL.md
+++ b/mlir/docs/Dialects/Linalg/OpDSL.md
@@ -16,7 +16,7 @@ corresponding `linalg.generic` IR for the composition.
## Basic usage
The tool is bundled with the MLIR Python bindings. To use from the CMake build
-tree, MLIR must be built with Python bindings enabled
+tree, MLIR must be build with Python bindings enabled
(`-DMLIR_ENABLE_BINDINGS_PYTHON=ON`). Then add the `python` directory in the
build tree to your `PYTHONPATH` environment variable (i.e. `export
PYTHONPATH=$PWD/build/tools/mlir/python_packages/mlir_core`). Optionally, use an
@@ -24,7 +24,7 @@ installed MLIR package, if available, to avoid building.
```shell
# Dump the `core_named_ops.py` module as YAML.
-python -m mlir.dialects.linalg.opdsl.dump_oplib.ops.core_named_ops
+python -m mlir.dialects.linalg.opdsl.dump_oplib .ops.core_named_ops
```
Alternatively, run the `$PWD/build/bin/update_core_linalg_named_ops.sh` script,
diff --git a/mlir/examples/standalone/pyproject.toml b/mlir/examples/standalone/pyproject.toml
index 75e215369ed55..5a1e6e86513c3 100644
--- a/mlir/examples/standalone/pyproject.toml
+++ b/mlir/examples/standalone/pyproject.toml
@@ -23,7 +23,9 @@ Discussions = "https://discourse.llvm.org/"
[build-system]
requires = [
"scikit-build-core>=0.10.7",
- "typing_extensions>=4.12.2"
+ "typing_extensions>=4.12.2",
+ "nanobind>=2.9, <3.0",
+ "pybind11>=2.10.0, <=2.13.6",
]
build-backend = "scikit_build_core.build"
diff --git a/mlir/examples/standalone/python/CMakeLists.txt b/mlir/examples/standalone/python/CMakeLists.txt
index 108c343714421..905c944939756 100644
--- a/mlir/examples/standalone/python/CMakeLists.txt
+++ b/mlir/examples/standalone/python/CMakeLists.txt
@@ -16,10 +16,27 @@ declare_mlir_dialect_python_bindings(
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir_standalone"
TD_FILE dialects/StandaloneOps.td
SOURCES
+ dialects/standalone_pybind11.py
dialects/standalone_nanobind.py
_mlir_libs/_standaloneDialectsNanobind/py.typed
DIALECT_NAME standalone)
+
+declare_mlir_python_extension(StandalonePythonSources.Pybind11Extension
+ MODULE_NAME _standaloneDialectsPybind11
+ ADD_TO_PARENT StandalonePythonSources
+ SOURCES
+ StandaloneExtensionPybind11.cpp
+ PRIVATE_LINK_LIBS
+ LLVMSupport
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIIR
+ MLIRCAPIArith
+ MLIRCAPITransforms
+ StandaloneCAPI
+ PYTHON_BINDINGS_LIBRARY pybind11
+)
+
declare_mlir_python_extension(StandalonePythonSources.NanobindExtension
MODULE_NAME _standaloneDialectsNanobind
ADD_TO_PARENT StandalonePythonSources
@@ -32,6 +49,7 @@ declare_mlir_python_extension(StandalonePythonSources.NanobindExtension
MLIRCAPIArith
MLIRCAPITransforms
StandaloneCAPI
+ PYTHON_BINDINGS_LIBRARY nanobind
)
diff --git a/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp b/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
new file mode 100644
index 0000000000000..da8c2167dc36b
--- /dev/null
+++ b/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
@@ -0,0 +1,38 @@
+//===- 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.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "Standalone-c/Dialects.h"
+#include "mlir-c/Dialect/Arith.h"
+#include "mlir/Bindings/Python/PybindAdaptors.h"
+
+using namespace mlir::python::adaptors;
+
+PYBIND11_MODULE(_standaloneDialectsPybind11, m) {
+ //===--------------------------------------------------------------------===//
+ // standalone dialect
+ //===--------------------------------------------------------------------===//
+ auto standaloneM = m.def_submodule("standalone");
+
+ standaloneM.def(
+ "register_dialects",
+ [](MlirContext context, bool load) {
+ MlirDialectHandle arithHandle = mlirGetDialectHandle__arith__();
+ MlirDialectHandle standaloneHandle =
+ mlirGetDialectHandle__standalone__();
+ mlirDialectHandleRegisterDialect(arithHandle, context);
+ mlirDialectHandleRegisterDialect(standaloneHandle, context);
+ if (load) {
+ mlirDialectHandleLoadDialect(arithHandle, context);
+ mlirDialectHandleRegisterDialect(standaloneHandle, context);
+ }
+ },
+ py::arg("context") = py::none(), py::arg("load") = true);
+}
diff --git a/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py b/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py
new file mode 100644
index 0000000000000..bfb98e404e13f
--- /dev/null
+++ b/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 *
diff --git a/mlir/examples/standalone/test/python/smoketest.py b/mlir/examples/standalone/test/python/smoketest.py
index f8819841fac45..26d84fd63e947 100644
--- a/mlir/examples/standalone/test/python/smoketest.py
+++ b/mlir/examples/standalone/test/python/smoketest.py
@@ -1,7 +1,16 @@
+# RUN: %python %s pybind11 | FileCheck %s
# RUN: %python %s nanobind | FileCheck %s
+import sys
from mlir_standalone.ir import *
-from mlir_standalone.dialects import standalone_nanobind as standalone_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_dialects()
diff --git a/mlir/include/mlir/Bindings/Python/PybindAdaptors.h b/mlir/include/mlir/Bindings/Python/PybindAdaptors.h
new file mode 100644
index 0000000000000..edc69774be922
--- /dev/null
+++ b/mlir/include/mlir/Bindings/Python/PybindAdaptors.h
@@ -0,0 +1,616 @@
+//===- PybindAdaptors.h - Interop with MLIR APIs via pybind11 -------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+// This file contains adaptors for clients of the core MLIR Python APIs to
+// interop via MLIR CAPI types, using pybind11. The facilities here do not
+// depend on implementation details of the MLIR Python API and do not introduce
+// C++-level dependencies with it (requiring only Python and CAPI-level
+// dependencies).
+//
+// It is encouraged to be used both in-tree and out-of-tree. For in-tree use
+// cases, it should be used for dialect implementations (versus relying on
+// Pybind-based internals of the core libraries).
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_BINDINGS_PYTHON_PYBINDADAPTORS_H
+#define MLIR_BINDINGS_PYTHON_PYBINDADAPTORS_H
+
+#include <pybind11/functional.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/pytypes.h>
+#include <pybind11/stl.h>
+
+#include "mlir-c/Bindings/Python/Interop.h"
+#include "mlir-c/Diagnostics.h"
+#include "mlir-c/IR.h"
+
+#include "llvm/ADT/Twine.h"
+
+namespace py = pybind11;
+using namespace py::literals;
+
+// Raw CAPI type casters need to be declared before use, so always include them
+// first.
+namespace pybind11 {
+namespace detail {
+
+/// Helper to convert a presumed MLIR API object to a capsule, accepting either
+/// an explicit Capsule (which can happen when two C APIs are communicating
+/// directly via Python) or indirectly by querying the MLIR_PYTHON_CAPI_PTR_ATTR
+/// attribute (through which supported MLIR Python API objects export their
+/// contained API pointer as a capsule). Throws a type error if the object is
+/// neither. This is intended to be used from type casters, which are invoked
+/// with a raw handle (unowned). The returned object's lifetime may not extend
+/// beyond the apiObject handle without explicitly having its refcount increased
+/// (i.e. on return).
+static py::object mlirApiObjectToCapsule(py::handle apiObject) {
+ if (PyCapsule_CheckExact(apiObject.ptr()))
+ return py::reinterpret_borrow<py::object>(apiObject);
+ if (!py::hasattr(apiObject, MLIR_PYTHON_CAPI_PTR_ATTR)) {
+ auto repr = py::repr(apiObject).cast<std::string>();
+ throw py::type_error(
+ (llvm::Twine("Expected an MLIR object (got ") + repr + ").").str());
+ }
+ return apiObject.attr(MLIR_PYTHON_CAPI_PTR_ATTR);
+}
+
+// Note: Currently all of the following support cast from py::object to the
+// Mlir* C-API type, but only a few light-weight, context-bound ones
+// implicitly cast the other way because the use case has not yet emerged and
+// ownership is unclear.
+
+/// Casts object <-> MlirAffineMap.
+template <>
+struct type_caster<MlirAffineMap> {
+ PYBIND11_TYPE_CASTER(MlirAffineMap, _("MlirAffineMap"));
+ bool load(handle src, bool) {
+ py::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToAffineMap(capsule.ptr());
+ if (mlirAffineMapIsNull(value)) {
+ return false;
+ }
+ return !mlirAffineMapIsNull(value);
+ }
+ static handle cast(MlirAffineMap v, return_value_policy, handle) {
+ py::object capsule =
+ py::reinterpret_steal<py::object>(mlirPythonAffineMapToCapsule(v));
+ return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("AffineMap")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .release();
+ }
+};
+
+/// Casts object <-> MlirAttribute.
+template <>
+struct type_caster<MlirAttribute> {
+ PYBIND11_TYPE_CASTER(MlirAttribute, _("MlirAttribute"));
+ bool load(handle src, bool) {
+ py::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToAttribute(capsule.ptr());
+ return !mlirAttributeIsNull(value);
+ }
+ static handle cast(MlirAttribute v, return_value_policy, handle) {
+ py::object capsule =
+ py::reinterpret_steal<py::object>(mlirPythonAttributeToCapsule(v));
+ return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Attribute")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .attr(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR)()
+ .release();
+ }
+};
+
+/// Casts object -> MlirBlock.
+template <>
+struct type_caster<MlirBlock> {
+ PYBIND11_TYPE_CASTER(MlirBlock, _("MlirBlock"));
+ bool load(handle src, bool) {
+ py::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToBlock(capsule.ptr());
+ return !mlirBlockIsNull(value);
+ }
+};
+
+/// Casts object -> MlirContext.
+template <>
+struct type_caster<MlirContext> {
+ PYBIND11_TYPE_CASTER(MlirContext, _("MlirContext"));
+ bool load(handle src, bool) {
+ if (src.is_none()) {
+ // Gets the current thread-bound context.
+ // TODO: This raises an error of "No current context" currently.
+ // Update the implementation to pretty-print the helpful error that the
+ // core implementations print in this case.
+ src = py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Context")
+ .attr("current");
+ }
+ py::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToContext(capsule.ptr());
+ return !mlirContextIsNull(value);
+ }
+};
+
+/// Casts object <-> MlirDialectRegistry.
+template <>
+struct type_caster<MlirDialectRegistry> {
+ PYBIND11_TYPE_CASTER(MlirDialectRegistry, _("MlirDialectRegistry"));
+ bool load(handle src, bool) {
+ py::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToDialectRegistry(capsule.ptr());
+ return !mlirDialectRegistryIsNull(value);
+ }
+ static handle cast(MlirDialectRegistry v, return_value_policy, handle) {
+ py::object capsule = py::reinterpret_steal<py::object>(
+ mlirPythonDialectRegistryToCapsule(v));
+ return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("DialectRegistry")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .release();
+ }
+};
+
+/// Casts object <-> MlirLocation.
+template <>
+struct type_caster<MlirLocation> {
+ PYBIND11_TYPE_CASTER(MlirLocation, _("MlirLocation"));
+ bool load(handle src, bool) {
+ if (src.is_none()) {
+ // Gets the current thread-bound context.
+ src = py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Location")
+ .attr("current");
+ }
+ py::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToLocation(capsule.ptr());
+ return !mlirLocationIsNull(value);
+ }
+ static handle cast(MlirLocation v, return_value_policy, handle) {
+ py::object capsule =
+ py::reinterpret_steal<py::object>(mlirPythonLocationToCapsule(v));
+ return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Location")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .release();
+ }
+};
+
+/// Casts object <-> MlirModule.
+template <>
+struct type_caster<MlirModule> {
+ PYBIND11_TYPE_CASTER(MlirModule, _("MlirModule"));
+ bool load(handle src, bool) {
+ py::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToModule(capsule.ptr());
+ return !mlirModuleIsNull(value);
+ }
+ static handle cast(MlirModule v, return_value_policy, handle) {
+ py::object capsule =
+ py::reinterpret_steal<py::object>(mlirPythonModuleToCapsule(v));
+ return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Module")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .release();
+ };
+};
+
+/// Casts object <-> MlirFrozenRewritePatternSet.
+template <>
+struct type_caster<MlirFrozenRewritePatternSet> {
+ PYBIND11_TYPE_CASTER(MlirFrozenRewritePatternSet,
+ _("MlirFrozenRewritePatternSet"));
+ bool load(handle src, bool) {
+ py::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToFrozenRewritePatternSet(capsule.ptr());
+ return value.ptr != nullptr;
+ }
+ static handle cast(MlirFrozenRewritePatternSet v, return_value_policy,
+ handle) {
+ py::object capsule = py::reinterpret_steal<py::object>(
+ mlirPythonFrozenRewritePatternSetToCapsule(v));
+ return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("rewrite"))
+ .attr("FrozenRewritePatternSet")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .release();
+ };
+};
+
+/// Casts object <-> MlirOperation.
+template <>
+struct type_caster<MlirOperation> {
+ PYBIND11_TYPE_CASTER(MlirOperation, _("MlirOperation"));
+ bool load(handle src, bool) {
+ py::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToOperation(capsule.ptr());
+ return !mlirOperationIsNull(value);
+ }
+ static handle cast(MlirOperation v, return_value_policy, handle) {
+ if (v.ptr == nullptr)
+ return py::none();
+ py::object capsule =
+ py::reinterpret_steal<py::object>(mlirPythonOperationToCapsule(v));
+ return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Operation")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .release();
+ };
+};
+
+/// Casts object <-> MlirValue.
+template <>
+struct type_caster<MlirValue> {
+ PYBIND11_TYPE_CASTER(MlirValue, _("MlirValue"));
+ bool load(handle src, bool) {
+ py::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToValue(capsule.ptr());
+ return !mlirValueIsNull(value);
+ }
+ static handle cast(MlirValue v, return_value_policy, handle) {
+ if (v.ptr == nullptr)
+ return py::none();
+ py::object capsule =
+ py::reinterpret_steal<py::object>(mlirPythonValueToCapsule(v));
+ return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Value")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .attr(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR)()
+ .release();
+ };
+};
+
+/// Casts object -> MlirPassManager.
+template <>
+struct type_caster<MlirPassManager> {
+ PYBIND11_TYPE_CASTER(MlirPassManager, _("MlirPassManager"));
+ bool load(handle src, bool) {
+ py::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToPassManager(capsule.ptr());
+ return !mlirPassManagerIsNull(value);
+ }
+};
+
+/// Casts object <-> MlirTypeID.
+template <>
+struct type_caster<MlirTypeID> {
+ PYBIND11_TYPE_CASTER(MlirTypeID, _("MlirTypeID"));
+ bool load(handle src, bool) {
+ py::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToTypeID(capsule.ptr());
+ return !mlirTypeIDIsNull(value);
+ }
+ static handle cast(MlirTypeID v, return_value_policy, handle) {
+ if (v.ptr == nullptr)
+ return py::none();
+ py::object capsule =
+ py::reinterpret_steal<py::object>(mlirPythonTypeIDToCapsule(v));
+ return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("TypeID")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .release();
+ };
+};
+
+/// Casts object <-> MlirType.
+template <>
+struct type_caster<MlirType> {
+ PYBIND11_TYPE_CASTER(MlirType, _("MlirType"));
+ bool load(handle src, bool) {
+ py::object capsule = mlirApiObjectToCapsule(src);
+ value = mlirPythonCapsuleToType(capsule.ptr());
+ return !mlirTypeIsNull(value);
+ }
+ static handle cast(MlirType t, return_value_policy, handle) {
+ py::object capsule =
+ py::reinterpret_steal<py::object>(mlirPythonTypeToCapsule(t));
+ return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Type")
+ .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .attr(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR)()
+ .release();
+ }
+};
+
+} // namespace detail
+} // namespace pybind11
+
+namespace mlir {
+namespace python {
+namespace adaptors {
+
+/// Provides a facility like py::class_ for defining a new class in a scope,
+/// but this allows extension of an arbitrary Python class, defining methods
+/// on it is a similar way. Classes defined in this way are very similar to
+/// if defined in Python in the usual way but use Pybind11 machinery to do
+/// it. These are not "real" Pybind11 classes but pure Python classes with no
+/// relation to a concrete C++ class.
+///
+/// Derived from a discussion upstream:
+/// https://github.com/pybind/pybind11/issues/1193
+/// (plus a fair amount of extra curricular poking)
+/// TODO: If this proves useful, see about including it in pybind11.
+class pure_subclass {
+public:
+ pure_subclass(py::handle scope, const char *derivedClassName,
+ const py::object &superClass) {
+ py::object pyType =
+ py::reinterpret_borrow<py::object>((PyObject *)&PyType_Type);
+ py::object metaclass = pyType(superClass);
+ py::dict attributes;
+
+ thisClass =
+ metaclass(derivedClassName, py::make_tuple(superClass), attributes);
+ scope.attr(derivedClassName) = thisClass;
+ }
+
+ template <typename Func, typename... Extra>
+ pure_subclass &def(const char *name, Func &&f, const Extra &...extra) {
+ py::cpp_function cf(
+ std::forward<Func>(f), py::name(name), py::is_method(thisClass),
+ py::sibling(py::getattr(thisClass, name, py::none())), extra...);
+ thisClass.attr(cf.name()) = cf;
+ return *this;
+ }
+
+ template <typename Func, typename... Extra>
+ pure_subclass &def_property_readonly(const char *name, Func &&f,
+ const Extra &...extra) {
+ py::cpp_function cf(
+ std::forward<Func>(f), py::name(name), py::is_method(thisClass),
+ py::sibling(py::getattr(thisClass, name, py::none())), extra...);
+ auto builtinProperty =
+ py::reinterpret_borrow<py::object>((PyObject *)&PyProperty_Type);
+ thisClass.attr(name) = builtinProperty(cf);
+ return *this;
+ }
+
+ template <typename Func, typename... Extra>
+ pure_subclass &def_staticmethod(const char *name, Func &&f,
+ const Extra &...extra) {
+ static_assert(!std::is_member_function_pointer<Func>::value,
+ "def_staticmethod(...) called with a non-static member "
+ "function pointer");
+ py::cpp_function cf(std::forward<Func>(f), py::name(name),
+ py::scope(thisClass), extra...);
+ thisClass.attr(cf.name()) = py::staticmethod(cf);
+ return *this;
+ }
+
+ template <typename Func, typename... Extra>
+ pure_subclass &def_classmethod(const char *name, Func &&f,
+ const Extra &...extra) {
+ static_assert(!std::is_member_function_pointer<Func>::value,
+ "def_classmethod(...) called with a non-static member "
+ "function pointer");
+ py::cpp_function cf(std::forward<Func>(f), py::name(name),
+ py::scope(thisClass), extra...);
+ thisClass.attr(cf.name()) =
+ py::reinterpret_borrow<py::object>(PyClassMethod_New(cf.ptr()));
+ return *this;
+ }
+
+ py::object get_class() const { return thisClass; }
+
+protected:
+ py::object superClass;
+ py::object thisClass;
+};
+
+/// Creates a custom subclass of mlir.ir.Attribute, implementing a casting
+/// constructor and type checking methods.
+class mlir_attribute_subclass : public pure_subclass {
+public:
+ using IsAFunctionTy = bool (*)(MlirAttribute);
+ using GetTypeIDFunctionTy = MlirTypeID (*)();
+
+ /// Subclasses by looking up the super-class dynamically.
+ mlir_attribute_subclass(py::handle scope, const char *attrClassName,
+ IsAFunctionTy isaFunction,
+ GetTypeIDFunctionTy getTypeIDFunction = nullptr)
+ : mlir_attribute_subclass(
+ scope, attrClassName, isaFunction,
+ py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("Attribute"),
+ getTypeIDFunction) {}
+
+ /// Subclasses with a provided mlir.ir.Attribute super-class. This must
+ /// be used if the subclass is being defined in the same extension module
+ /// as the mlir.ir class (otherwise, it will trigger a recursive
+ /// initialization).
+ mlir_attribute_subclass(py::handle scope, const char *typeClassName,
+ IsAFunctionTy isaFunction, const py::object &superCls,
+ GetTypeIDFunctionTy getTypeIDFunction = nullptr)
+ : pure_subclass(scope, typeClassName, superCls) {
+ // Casting constructor. Note that it hard, if not impossible, to properly
+ // call chain to parent `__init__` in pybind11 due to its special handling
+ // for init functions that don't have a fully constructed self-reference,
+ // which makes it impossible to forward it to `__init__` of a superclass.
+ // Instead, provide a custom `__new__` and call that of a superclass, which
+ // eventually calls `__init__` of the superclass. Since attribute subclasses
+ // have no additional members, we can just return the instance thus created
+ // without amending it.
+ std::string captureTypeName(
+ typeClassName); // As string in case if typeClassName is not static.
+ py::cpp_function newCf(
+ [superCls, isaFunction, captureTypeName](py::object cls,
+ py::object otherAttribute) {
+ MlirAttribute rawAttribute = py::cast<MlirAttribute>(otherAttribute);
+ if (!isaFunction(rawAttribute)) {
+ auto origRepr = py::repr(otherAttribute).cast<std::string>();
+ throw std::invalid_argument(
+ (llvm::Twine("Cannot cast attribute to ") + captureTypeName +
+ " (from " + origRepr + ")")
+ .str());
+ }
+ py::object self = superCls.attr("__new__")(cls, otherAttribute);
+ return self;
+ },
+ py::name("__new__"), py::arg("cls"), py::arg("cast_from_attr"));
+ thisClass.attr("__new__") = newCf;
+
+ // 'isinstance' method.
+ def_staticmethod(
+ "isinstance",
+ [isaFunction](MlirAttribute other) { return isaFunction(other); },
+ py::arg("other_attribute"));
+ def("__repr__", [superCls, captureTypeName](py::object self) {
+ return py::repr(superCls(self))
+ .attr("replace")(superCls.attr("__name__"), captureTypeName);
+ });
+ if (getTypeIDFunction) {
+ def_staticmethod("get_static_typeid",
+ [getTypeIDFunction]() { return getTypeIDFunction(); });
+ py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
+ getTypeIDFunction())(pybind11::cpp_function(
+ [thisClass = thisClass](const py::object &mlirAttribute) {
+ return thisClass(mlirAttribute);
+ }));
+ }
+ }
+};
+
+/// Creates a custom subclass of mlir.ir.Type, implementing a casting
+/// constructor and type checking methods.
+class mlir_type_subclass : public pure_subclass {
+public:
+ using IsAFunctionTy = bool (*)(MlirType);
+ using GetTypeIDFunctionTy = MlirTypeID (*)();
+
+ /// Subclasses by looking up the super-class dynamically.
+ mlir_type_subclass(py::handle scope, const char *typeClassName,
+ IsAFunctionTy isaFunction,
+ GetTypeIDFunctionTy getTypeIDFunction = nullptr)
+ : mlir_type_subclass(
+ scope, typeClassName, isaFunction,
+ py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir")).attr("Type"),
+ getTypeIDFunction) {}
+
+ /// Subclasses with a provided mlir.ir.Type super-class. This must
+ /// be used if the subclass is being defined in the same extension module
+ /// as the mlir.ir class (otherwise, it will trigger a recursive
+ /// initialization).
+ mlir_type_subclass(py::handle scope, const char *typeClassName,
+ IsAFunctionTy isaFunction, const py::object &superCls,
+ GetTypeIDFunctionTy getTypeIDFunction = nullptr)
+ : pure_subclass(scope, typeClassName, superCls) {
+ // Casting constructor. Note that it hard, if not impossible, to properly
+ // call chain to parent `__init__` in pybind11 due to its special handling
+ // for init functions that don't have a fully constructed self-reference,
+ // which makes it impossible to forward it to `__init__` of a superclass.
+ // Instead, provide a custom `__new__` and call that of a superclass, which
+ // eventually calls `__init__` of the superclass. Since attribute subclasses
+ // have no additional members, we can just return the instance thus created
+ // without amending it.
+ std::string captureTypeName(
+ typeClassName); // As string in case if typeClassName is not static.
+ py::cpp_function newCf(
+ [superCls, isaFunction, captureTypeName](py::object cls,
+ py::object otherType) {
+ MlirType rawType = py::cast<MlirType>(otherType);
+ if (!isaFunction(rawType)) {
+ auto origRepr = py::repr(otherType).cast<std::string>();
+ throw std::invalid_argument((llvm::Twine("Cannot cast type to ") +
+ captureTypeName + " (from " +
+ origRepr + ")")
+ .str());
+ }
+ py::object self = superCls.attr("__new__")(cls, otherType);
+ return self;
+ },
+ py::name("__new__"), py::arg("cls"), py::arg("cast_from_type"));
+ thisClass.attr("__new__") = newCf;
+
+ // 'isinstance' method.
+ def_staticmethod(
+ "isinstance",
+ [isaFunction](MlirType other) { return isaFunction(other); },
+ py::arg("other_type"));
+ def("__repr__", [superCls, captureTypeName](py::object self) {
+ return py::repr(superCls(self))
+ .attr("replace")(superCls.attr("__name__"), captureTypeName);
+ });
+ if (getTypeIDFunction) {
+ // 'get_static_typeid' method.
+ // This is modeled as a static method instead of a static property because
+ // `def_property_readonly_static` is not available in `pure_subclass` and
+ // we do not want to introduce the complexity that pybind uses to
+ // implement it.
+ def_staticmethod("get_static_typeid",
+ [getTypeIDFunction]() { return getTypeIDFunction(); });
+ py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
+ getTypeIDFunction())(pybind11::cpp_function(
+ [thisClass = thisClass](const py::object &mlirType) {
+ return thisClass(mlirType);
+ }));
+ }
+ }
+};
+
+/// Creates a custom subclass of mlir.ir.Value, implementing a casting
+/// constructor and type checking methods.
+class mlir_value_subclass : public pure_subclass {
+public:
+ using IsAFunctionTy = bool (*)(MlirValue);
+
+ /// Subclasses by looking up the super-class dynamically.
+ mlir_value_subclass(py::handle scope, const char *valueClassName,
+ IsAFunctionTy isaFunction)
+ : mlir_value_subclass(
+ scope, valueClassName, isaFunction,
+ py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir")).attr("Value")) {
+ }
+
+ /// Subclasses with a provided mlir.ir.Value super-class. This must
+ /// be used if the subclass is being defined in the same extension module
+ /// as the mlir.ir class (otherwise, it will trigger a recursive
+ /// initialization).
+ mlir_value_subclass(py::handle scope, const char *valueClassName,
+ IsAFunctionTy isaFunction, const py::object &superCls)
+ : pure_subclass(scope, valueClassName, superCls) {
+ // Casting constructor. Note that it hard, if not impossible, to properly
+ // call chain to parent `__init__` in pybind11 due to its special handling
+ // for init functions that don't have a fully constructed self-reference,
+ // which makes it impossible to forward it to `__init__` of a superclass.
+ // Instead, provide a custom `__new__` and call that of a superclass, which
+ // eventually calls `__init__` of the superclass. Since attribute subclasses
+ // have no additional members, we can just return the instance thus created
+ // without amending it.
+ std::string captureValueName(
+ valueClassName); // As string in case if valueClassName is not static.
+ py::cpp_function newCf(
+ [superCls, isaFunction, captureValueName](py::object cls,
+ py::object otherValue) {
+ MlirValue rawValue = py::cast<MlirValue>(otherValue);
+ if (!isaFunction(rawValue)) {
+ auto origRepr = py::repr(otherValue).cast<std::string>();
+ throw std::invalid_argument((llvm::Twine("Cannot cast value to ") +
+ captureValueName + " (from " +
+ origRepr + ")")
+ .str());
+ }
+ py::object self = superCls.attr("__new__")(cls, otherValue);
+ return self;
+ },
+ py::name("__new__"), py::arg("cls"), py::arg("cast_from_value"));
+ thisClass.attr("__new__") = newCf;
+
+ // 'isinstance' method.
+ def_staticmethod(
+ "isinstance",
+ [isaFunction](MlirValue other) { return isaFunction(other); },
+ py::arg("other_value"));
+ }
+};
+
+} // namespace adaptors
+
+} // namespace python
+} // namespace mlir
+
+#endif // MLIR_BINDINGS_PYTHON_PYBINDADAPTORS_H
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index cea5b252e0faf..9f5246de6bda0 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -440,11 +440,11 @@ declare_mlir_dialect_python_bindings(
DIALECT_NAME smt)
declare_mlir_dialect_python_bindings(
- ADD_TO_PARENT MLIRPythonSources.Dialects
- ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
- TD_FILE dialects/SPIRVOps.td
- SOURCES dialects/spirv.py
- DIALECT_NAME spirv)
+ ADD_TO_PARENT MLIRPythonSources.Dialects
+ ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+ TD_FILE dialects/SPIRVOps.td
+ SOURCES dialects/spirv.py
+ DIALECT_NAME spirv)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -501,6 +501,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
MODULE_NAME _mlir
ADD_TO_PARENT MLIRPythonSources.Core
ROOT_DIR "${PYTHON_SOURCE_DIR}"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
MainModule.cpp
IRAffine.cpp
@@ -539,6 +540,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
MODULE_NAME _mlirRegisterEverything
ROOT_DIR "${PYTHON_SOURCE_DIR}"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
RegisterEverything.cpp
PRIVATE_LINK_LIBS
@@ -549,10 +551,11 @@ declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
MLIRCAPIRegisterEverything
)
-declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
MODULE_NAME _mlirDialectsLinalg
ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
ROOT_DIR "${PYTHON_SOURCE_DIR}"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
DialectLinalg.cpp
PRIVATE_LINK_LIBS
@@ -562,10 +565,11 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Nanobind
MLIRCAPILinalg
)
-declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
MODULE_NAME _mlirDialectsGPU
ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
ROOT_DIR "${PYTHON_SOURCE_DIR}"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
DialectGPU.cpp
PRIVATE_LINK_LIBS
@@ -575,10 +579,11 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Nanobind
MLIRCAPIGPU
)
-declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
MODULE_NAME _mlirDialectsLLVM
ADD_TO_PARENT MLIRPythonSources.Dialects.llvm
ROOT_DIR "${PYTHON_SOURCE_DIR}"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
DialectLLVM.cpp
PRIVATE_LINK_LIBS
@@ -588,10 +593,11 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Nanobind
MLIRCAPILLVM
)
-declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
MODULE_NAME _mlirDialectsQuant
ADD_TO_PARENT MLIRPythonSources.Dialects.quant
ROOT_DIR "${PYTHON_SOURCE_DIR}"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
DialectQuant.cpp
PRIVATE_LINK_LIBS
@@ -601,10 +607,11 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Nanobind
MLIRCAPIQuant
)
-declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
MODULE_NAME _mlirDialectsNVGPU
ADD_TO_PARENT MLIRPythonSources.Dialects.nvgpu
ROOT_DIR "${PYTHON_SOURCE_DIR}"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
DialectNVGPU.cpp
PRIVATE_LINK_LIBS
@@ -614,10 +621,11 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Nanobind
MLIRCAPINVGPU
)
-declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
MODULE_NAME _mlirDialectsPDL
ADD_TO_PARENT MLIRPythonSources.Dialects.pdl
ROOT_DIR "${PYTHON_SOURCE_DIR}"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
DialectPDL.cpp
PRIVATE_LINK_LIBS
@@ -627,10 +635,11 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Nanobind
MLIRCAPIPDL
)
-declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
MODULE_NAME _mlirDialectsSparseTensor
ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
ROOT_DIR "${PYTHON_SOURCE_DIR}"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
DialectSparseTensor.cpp
PRIVATE_LINK_LIBS
@@ -640,10 +649,11 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Nanobind
MLIRCAPISparseTensor
)
-declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
MODULE_NAME _mlirDialectsTransform
ADD_TO_PARENT MLIRPythonSources.Dialects.transform
ROOT_DIR "${PYTHON_SOURCE_DIR}"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
DialectTransform.cpp
PRIVATE_LINK_LIBS
@@ -653,10 +663,11 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Nanobind
MLIRCAPITransformDialect
)
-declare_mlir_python_extension(MLIRPythonExtension.Dialects.IRDL.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.IRDL.Pybind
MODULE_NAME _mlirDialectsIRDL
ADD_TO_PARENT MLIRPythonSources.Dialects.irdl
ROOT_DIR "${PYTHON_SOURCE_DIR}"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
DialectIRDL.cpp
PRIVATE_LINK_LIBS
@@ -670,6 +681,7 @@ declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
MODULE_NAME _mlirAsyncPasses
ADD_TO_PARENT MLIRPythonSources.Dialects.async
ROOT_DIR "${PYTHON_SOURCE_DIR}"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
AsyncPasses.cpp
PRIVATE_LINK_LIBS
@@ -683,6 +695,7 @@ if(MLIR_ENABLE_EXECUTION_ENGINE)
MODULE_NAME _mlirExecutionEngine
ADD_TO_PARENT MLIRPythonSources.ExecutionEngine
ROOT_DIR "${PYTHON_SOURCE_DIR}"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
ExecutionEngineModule.cpp
PRIVATE_LINK_LIBS
@@ -696,6 +709,7 @@ declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
MODULE_NAME _mlirGPUPasses
ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
ROOT_DIR "${PYTHON_SOURCE_DIR}"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
GPUPasses.cpp
PRIVATE_LINK_LIBS
@@ -708,6 +722,7 @@ declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
MODULE_NAME _mlirLinalgPasses
ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
ROOT_DIR "${PYTHON_SOURCE_DIR}"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
LinalgPasses.cpp
PRIVATE_LINK_LIBS
@@ -716,10 +731,11 @@ declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
MLIRCAPILinalg
)
-declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Nanobind
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
MODULE_NAME _mlirDialectsSMT
ADD_TO_PARENT MLIRPythonSources.Dialects.smt
ROOT_DIR "${PYTHON_SOURCE_DIR}"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
DialectSMT.cpp
# Headers must be included explicitly so they are installed.
@@ -736,6 +752,7 @@ declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
MODULE_NAME _mlirSparseTensorPasses
ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
ROOT_DIR "${PYTHON_SOURCE_DIR}"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
SparseTensorPasses.cpp
PRIVATE_LINK_LIBS
@@ -748,6 +765,7 @@ declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
MODULE_NAME _mlirTransformInterpreter
ADD_TO_PARENT MLIRPythonSources.Dialects.transform
ROOT_DIR "${PYTHON_SOURCE_DIR}"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
TransformInterpreter.cpp
PRIVATE_LINK_LIBS
@@ -789,10 +807,23 @@ if(MLIR_INCLUDE_TESTS)
ADD_TO_PARENT MLIRPythonTestSources.Dialects.PythonTest
SOURCES "dialects/_python_test_ops_gen.py")
+ declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtensionPybind11
+ MODULE_NAME _mlirPythonTestPybind11
+ ADD_TO_PARENT MLIRPythonTestSources.Dialects
+ ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
+ PYTHON_BINDINGS_LIBRARY pybind11
+ SOURCES
+ PythonTestModulePybind11.cpp
+ PRIVATE_LINK_LIBS
+ LLVMSupport
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIPythonTestDialect
+ )
declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtensionNanobind
MODULE_NAME _mlirPythonTestNanobind
ADD_TO_PARENT MLIRPythonTestSources.Dialects
ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
+ PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
PythonTestModuleNanobind.cpp
PRIVATE_LINK_LIBS
diff --git a/mlir/python/mlir/dialects/python_test.py b/mlir/python/mlir/dialects/python_test.py
index 56d3c0f7a5465..9380896c8c06e 100644
--- a/mlir/python/mlir/dialects/python_test.py
+++ b/mlir/python/mlir/dialects/python_test.py
@@ -5,7 +5,12 @@
from ._python_test_ops_gen import *
-def register_python_test_dialect(registry):
- from .._mlir_libs import _mlirPythonTestNanobind
+def register_python_test_dialect(registry, use_nanobind):
+ if use_nanobind:
+ from .._mlir_libs import _mlirPythonTestNanobind
- _mlirPythonTestNanobind.register_dialect(registry)
+ _mlirPythonTestNanobind.register_dialect(registry)
+ else:
+ from .._mlir_libs import _mlirPythonTestPybind11
+
+ _mlirPythonTestPybind11.register_dialect(registry)
diff --git a/mlir/python/requirements.txt b/mlir/python/requirements.txt
index 5ff9500e50127..abe09259bb1e8 100644
--- a/mlir/python/requirements.txt
+++ b/mlir/python/requirements.txt
@@ -1,4 +1,6 @@
+nanobind>=2.9, <3.0
numpy>=1.19.5, <=2.1.2
+pybind11>=2.10.0, <=2.13.6
PyYAML>=5.4.0, <=6.0.1
ml_dtypes>=0.1.0, <=0.6.0; python_version<"3.13" # provides several NumPy dtype extensions, including the bf16
ml_dtypes>=0.5.0, <=0.6.0; python_version>="3.13"
diff --git a/mlir/test/python/dialects/python_test.py b/mlir/test/python/dialects/python_test.py
index 5a9acc7dcf6bb..1194e32c960c8 100644
--- a/mlir/test/python/dialects/python_test.py
+++ b/mlir/test/python/dialects/python_test.py
@@ -1,4 +1,5 @@
-# RUN: %PYTHON %s | FileCheck %s
+# RUN: %PYTHON %s pybind11 | FileCheck %s
+# RUN: %PYTHON %s nanobind | FileCheck %s
import sys
import typing
from typing import Union, Optional
@@ -9,14 +10,26 @@
import mlir.dialects.tensor as tensor
import mlir.dialects.arith as arith
-from mlir._mlir_libs._mlirPythonTestNanobind import (
- TestAttr,
- TestType,
- TestTensorValue,
- TestIntegerRankedTensorType,
-)
-
-test.register_python_test_dialect(get_dialect_registry())
+if sys.argv[1] == "pybind11":
+ from mlir._mlir_libs._mlirPythonTestPybind11 import (
+ TestAttr,
+ TestType,
+ TestTensorValue,
+ TestIntegerRankedTensorType,
+ )
+
+ test.register_python_test_dialect(get_dialect_registry(), use_nanobind=False)
+elif sys.argv[1] == "nanobind":
+ from mlir._mlir_libs._mlirPythonTestNanobind import (
+ TestAttr,
+ TestType,
+ TestTensorValue,
+ TestIntegerRankedTensorType,
+ )
+
+ test.register_python_test_dialect(get_dialect_registry(), use_nanobind=True)
+else:
+ raise ValueError("Expected pybind11 or nanobind as argument")
def run(f):
diff --git a/mlir/test/python/lib/CMakeLists.txt b/mlir/test/python/lib/CMakeLists.txt
index f51a7b44e64dd..9a813dace2f54 100644
--- a/mlir/test/python/lib/CMakeLists.txt
+++ b/mlir/test/python/lib/CMakeLists.txt
@@ -1,6 +1,7 @@
set(LLVM_OPTIONAL_SOURCES
PythonTestCAPI.cpp
PythonTestDialect.cpp
+ PythonTestModulePybind11.cpp
PythonTestModuleNanobind.cpp
)
diff --git a/mlir/test/python/lib/PythonTestModulePybind11.cpp b/mlir/test/python/lib/PythonTestModulePybind11.cpp
new file mode 100644
index 0000000000000..94a5f5178d16e
--- /dev/null
+++ b/mlir/test/python/lib/PythonTestModulePybind11.cpp
@@ -0,0 +1,118 @@
+//===- PythonTestModule.cpp - Python extension for the PythonTest dialect -===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+// This is the pybind11 edition of the PythonTest dialect module.
+//===----------------------------------------------------------------------===//
+
+#include "PythonTestCAPI.h"
+#include "mlir-c/BuiltinAttributes.h"
+#include "mlir-c/BuiltinTypes.h"
+#include "mlir-c/IR.h"
+#include "mlir/Bindings/Python/PybindAdaptors.h"
+
+namespace py = pybind11;
+using namespace mlir::python::adaptors;
+using namespace pybind11::literals;
+
+static bool mlirTypeIsARankedIntegerTensor(MlirType t) {
+ return mlirTypeIsARankedTensor(t) &&
+ mlirTypeIsAInteger(mlirShapedTypeGetElementType(t));
+}
+
+PYBIND11_MODULE(_mlirPythonTestPybind11, m) {
+ m.def(
+ "register_python_test_dialect",
+ [](MlirContext context, bool load) {
+ MlirDialectHandle pythonTestDialect =
+ mlirGetDialectHandle__python_test__();
+ mlirDialectHandleRegisterDialect(pythonTestDialect, context);
+ if (load) {
+ mlirDialectHandleLoadDialect(pythonTestDialect, context);
+ }
+ },
+ py::arg("context"), py::arg("load") = true);
+
+ m.def(
+ "register_dialect",
+ [](MlirDialectRegistry registry) {
+ MlirDialectHandle pythonTestDialect =
+ mlirGetDialectHandle__python_test__();
+ mlirDialectHandleInsertDialect(pythonTestDialect, registry);
+ },
+ py::arg("registry"));
+
+ mlir_attribute_subclass(m, "TestAttr",
+ mlirAttributeIsAPythonTestTestAttribute,
+ mlirPythonTestTestAttributeGetTypeID)
+ .def_classmethod(
+ "get",
+ [](const py::object &cls, MlirContext ctx) {
+ return cls(mlirPythonTestTestAttributeGet(ctx));
+ },
+ py::arg("cls"), py::arg("context") = py::none());
+
+ mlir_type_subclass(m, "TestType", mlirTypeIsAPythonTestTestType,
+ mlirPythonTestTestTypeGetTypeID)
+ .def_classmethod(
+ "get",
+ [](const py::object &cls, MlirContext ctx) {
+ return cls(mlirPythonTestTestTypeGet(ctx));
+ },
+ py::arg("cls"), py::arg("context") = py::none());
+
+ auto typeCls =
+ mlir_type_subclass(m, "TestIntegerRankedTensorType",
+ mlirTypeIsARankedIntegerTensor,
+ py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr("RankedTensorType"))
+ .def_classmethod(
+ "get",
+ [](const py::object &cls, std::vector<int64_t> shape,
+ unsigned width, MlirContext ctx) {
+ MlirAttribute encoding = mlirAttributeGetNull();
+ return cls(mlirRankedTensorTypeGet(
+ shape.size(), shape.data(), mlirIntegerTypeGet(ctx, width),
+ encoding));
+ },
+ "cls"_a, "shape"_a, "width"_a, "context"_a = py::none());
+
+ assert(py::hasattr(typeCls.get_class(), "static_typeid") &&
+ "TestIntegerRankedTensorType has no static_typeid");
+
+ MlirTypeID mlirRankedTensorTypeID = mlirRankedTensorTypeGetTypeID();
+
+ py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(mlirRankedTensorTypeID,
+ "replace"_a = true)(
+ pybind11::cpp_function([typeCls](const py::object &mlirType) {
+ return typeCls.get_class()(mlirType);
+ }));
+
+ auto valueCls = mlir_value_subclass(m, "TestTensorValue",
+ mlirTypeIsAPythonTestTestTensorValue)
+ .def("is_null", [](MlirValue &self) {
+ return mlirValueIsNull(self);
+ });
+
+ py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr(MLIR_PYTHON_CAPI_VALUE_CASTER_REGISTER_ATTR)(
+ mlirRankedTensorTypeID)(
+ pybind11::cpp_function([valueCls](const py::object &valueObj) {
+ py::object capsule = mlirApiObjectToCapsule(valueObj);
+ MlirValue v = mlirPythonCapsuleToValue(capsule.ptr());
+ MlirType t = mlirValueGetType(v);
+ // This is hyper-specific in order to exercise/test registering a
+ // value caster from cpp (but only for a single test case; see
+ // testTensorValue python_test.py).
+ if (mlirShapedTypeHasStaticShape(t) &&
+ mlirShapedTypeGetDimSize(t, 0) == 1 &&
+ mlirShapedTypeGetDimSize(t, 1) == 2 &&
+ mlirShapedTypeGetDimSize(t, 2) == 3)
+ return valueCls.get_class()(valueObj);
+ return valueObj;
+ }));
+}
diff --git a/mlir/tools/mlir-linalg-ods-gen/update_core_linalg_named_ops.sh.in b/mlir/tools/mlir-linalg-ods-gen/update_core_linalg_named_ops.sh.in
index 0bb6a209b6c18..da4db3971987d 100755
--- a/mlir/tools/mlir-linalg-ods-gen/update_core_linalg_named_ops.sh.in
+++ b/mlir/tools/mlir-linalg-ods-gen/update_core_linalg_named_ops.sh.in
@@ -26,7 +26,7 @@ export PYTHONPATH="$python_package_dir"
OUTPUT="$(
echo "### AUTOGENERATED from core_named_ops.py" && \
echo "### To regenerate, run: bin/update_core_linalg_named_ops.sh" && \
- "$python_exe" -m mlir.dialects.linalg.opdsl.dump_oplib.ops.core_named_ops \
+ "$python_exe" -m mlir.dialects.linalg.opdsl.dump_oplib .ops.core_named_ops \
)"
echo "$OUTPUT" > "$dest_file"
echo "Success."
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index 0c77a1e62723a..422c29fc9c4d5 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -1049,39 +1049,55 @@ filegroup(
)
cc_library(
- name = "MLIRBindingsPythonNanobindHeaders",
+ name = "MLIRBindingsPythonHeaders",
includes = [
"include",
],
textual_hdrs = [":MLIRBindingsPythonHeaderFiles"],
deps = [
":CAPIIRHeaders",
- "@nanobind",
+ "@pybind11",
"@rules_python//python/cc:current_py_cc_headers",
],
)
-alias(
- name = "MLIRBindingsPythonHeaders",
- actual = ":MLIRBindingsPythonNanobindHeaders",
+cc_library(
+ name = "MLIRBindingsPythonHeadersAndDeps",
+ includes = [
+ "include",
+ ],
+ textual_hdrs = [":MLIRBindingsPythonHeaderFiles"],
+ deps = [
+ ":CAPIIR",
+ "@pybind11",
+ "@rules_python//python/cc:current_py_cc_headers",
+ ],
)
cc_library(
- name = "MLIRBindingsPythonNanobindHeadersAndDeps",
+ name = "MLIRBindingsPythonNanobindHeaders",
includes = [
"include",
],
textual_hdrs = [":MLIRBindingsPythonHeaderFiles"],
deps = [
- ":CAPIIR",
+ ":CAPIIRHeaders",
"@nanobind",
"@rules_python//python/cc:current_py_cc_headers",
],
)
-alias(
- name = "MLIRBindingsPythonHeadersAndDeps",
- actual = ":MLIRBindingsPythonNanobindHeadersAndDeps",
+cc_library(
+ name = "MLIRBindingsPythonNanobindHeadersAndDeps",
+ includes = [
+ "include",
+ ],
+ textual_hdrs = [":MLIRBindingsPythonHeaderFiles"],
+ deps = [
+ ":CAPIIR",
+ "@nanobind",
+ "@rules_python//python/cc:current_py_cc_headers",
+ ],
)
# These flags are needed for pybind11 to work.
@@ -1131,7 +1147,7 @@ cc_library(
":CAPIIR",
":CAPIInterfaces",
":CAPITransforms",
- ":MLIRBindingsPythonHeadersAndDeps",
+ ":MLIRBindingsPythonNanobindHeadersAndDeps",
":Support",
":config",
"//llvm:Support",
@@ -1154,7 +1170,7 @@ cc_library(
":CAPIDebugHeaders",
":CAPIIRHeaders",
":CAPITransformsHeaders",
- ":MLIRBindingsPythonHeaders",
+ ":MLIRBindingsPythonNanobindHeaders",
":Support",
":config",
"//llvm:Support",
@@ -1204,7 +1220,7 @@ cc_binary(
linkstatic = 0,
deps = [
":CAPIIR",
- ":MLIRBindingsPythonHeadersAndDeps",
+ ":MLIRBindingsPythonNanobindHeadersAndDeps",
"@nanobind",
],
)
@@ -1222,7 +1238,7 @@ cc_binary(
deps = [
":CAPIIR",
":CAPILinalg",
- ":MLIRBindingsPythonHeadersAndDeps",
+ ":MLIRBindingsPythonNanobindHeadersAndDeps",
"@nanobind",
],
)
@@ -1237,7 +1253,7 @@ cc_binary(
deps = [
":CAPIIR",
":CAPILLVM",
- ":MLIRBindingsPythonHeadersAndDeps",
+ ":MLIRBindingsPythonNanobindHeadersAndDeps",
"@nanobind",
],
)
@@ -1252,7 +1268,7 @@ cc_binary(
deps = [
":CAPIIR",
":CAPIQuant",
- ":MLIRBindingsPythonHeadersAndDeps",
+ ":MLIRBindingsPythonNanobindHeadersAndDeps",
"@nanobind",
],
)
@@ -1267,7 +1283,7 @@ cc_binary(
deps = [
":CAPIIR",
":CAPISparseTensor",
- ":MLIRBindingsPythonHeadersAndDeps",
+ ":MLIRBindingsPythonNanobindHeadersAndDeps",
"@nanobind",
],
)
@@ -1282,7 +1298,7 @@ cc_binary(
linkstatic = 0,
deps = [
":CAPIExecutionEngine",
- ":MLIRBindingsPythonHeadersAndDeps",
+ ":MLIRBindingsPythonNanobindHeadersAndDeps",
"@nanobind",
"@rules_python//python/cc:current_py_cc_headers",
],
@@ -1298,7 +1314,7 @@ cc_binary(
linkstatic = 0,
deps = [
":CAPILinalg",
- ":MLIRBindingsPythonHeadersAndDeps",
+ ":MLIRBindingsPythonNanobindHeadersAndDeps",
"@nanobind",
"@rules_python//python/cc:current_py_cc_headers",
],
More information about the Mlir-commits
mailing list