[llvm] [mlir] [mlir][Python] create MLIRPythonSupport (PR #171775)
via llvm-commits
llvm-commits at lists.llvm.org
Fri Jan 2 11:59:12 PST 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Maksim Levental (makslevental)
<details>
<summary>Changes</summary>
# What
This PR adds a shared library `MLIRPythonSupport` which contains all of the CRTP classes ike `PyConcreteValue`, `PyConcreteType`, `PyConcreteAttribute`, as well as other useful code like `Defaulting*` and etc enabling their reuse in downstream projects. Downstream projects can now do
```c++
struct PyTestType : mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyConcreteType<PyTestType> {
...
};
class PyTestAttr : public mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyConcreteAttribute<PyTestAttr> {
...
}
NB_MODULE(_mlirPythonTestNanobind, m) {
PyTestType::bind(m);
PyTestAttr::bind(m);
}
```
instead of using the discordant alternative `mlir_type_subclass`/`mlir_attr_subclass` (same goes for `PyConcreteValue`/`mlir_value_subclass`).
# Why
This PR is mostly code motion (along with CMake) but before I describe the changes I want to state the goals/benefits:
1. Currently upstream "core" extensions and "dialect" extensions ([all of the `Dialect*` extensions here](https://github.com/llvm/llvm-project/tree/d7c734b5a14bd91e1c76e2ce0014c19f9deef487/mlir/lib/Bindings/Python)) are a two-tier system;
**a**. [core extensions](https://github.com/llvm/llvm-project/blob/main/mlir/lib/Bindings/Python/IRTypes.cpp#L361) enjoy first class support as far as type inference[^3], type stub generation, and ease of implementation, while dialect extensions [have poorer support](https://reviews.llvm.org/D150927), incorrect type stub generation much more tedious (boilerplate) implementation;
**b**. Crucially, this two-tiered system is reflected in the fact that **the two sets of types/attributes are not in the same Python object hierarchy**. To wit: `isinstance(..., Type)` and `isinstance(..., Attribute)` are not supported for the dialect extensions[^2];
**c**. Since these types are not exposed in public headers, downstream users (dialect extensions or not) cannot write functions that overload on e.g. `PyFloat8*Type` - that's quite a [useful feature](https://github.com/nod-ai/PI/blob/fdbee98df8376f47818e6b47e1cf089528c9d48d/cpp_ext/TorchOps.cpp#L29-L69)!
2. The dialect extensions incur a sizeable performance penalty relative to the core extensions in that every single trip across the wire (either `python->cpp` or `cpp->python`) requires work in addition to nanobind's own casting/construction pipeline;
**a**. When going from `python->cpp`, [we extract the capsule object from the Python object](https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h#L219C24-L219C46) and then extract from the capsule the `Mlir*` opaque struct/ptr. This side isn't so onerous;
**b**. When going from `cpp->python` we call long-hand call Python `import` APIs and construct the Python object using `_CAPICreate`. Note, there at least 2 `attr` calls incurred in addition to `_CAPICreate`; this is already much more [efficiently handled by nanobind itself](https://github.com/wjakob/nanobind/blob/4ba51fcf795971c5d603d875ae4184bc0c9bd8e6/src/nb_internals.h#L381-L382)!
3. This division blocks various features: in some configurations[^1] we trigger a circular import bug because "dialect" types and attributes perform an [import of the root `_mlir` module](https://github.com/llvm/llvm-project/blob/bd9651bf78f2b1713a8203e0bd5b97f7ff199924/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h#L585) when they are created (the types themselves, not even instances of those types). This blocks type stub generation for dialect extensions (i.e., the reason we currently only generate type stubs for `_mlir`).
# How
Prior this was not done/possible because of "ODR" issues but I have resolved those issues; the basic idea for how we solve this is "move things we want to share into shared libraries":
1. Move IRCore (stuff like `PyConcreteValue`, `PyConcreteType`, `PyConcreteAttribute`) into `MLIRPythonSupport`;
- Note, we move the rest of the things in `IRModule.h` (renamed to `IRCore.h`) because `PyConcreteValue`, `PyConcreteType`, `PyConcreteAttribute` depend on them. This makes for a bigger PR than one would hope for but ultimately I think we should give people access to these classes to use as they see fit (specifically inherit from, but also liberally use in bindings signatures instead of the opaque `Mlir*` struct wrappers).
2. Put all of this code into a nested namespace `MLIR_BINDINGS_PYTHON_DOMAIN` which is determined by a compile time define (and tied to `MLIR_BINDINGS_PYTHON_NB_DOMAIN`). This is necessary in order to prevent conflicts on both symbol name **and** typeid (necessary for nanobind to not double register binded types) between multiple bindings libraries (e.g., `torch-mlir`, and `jax`). Note [nanobind doesn't support `module_local` like pybind11](https://nanobind.readthedocs.io/en/latest/porting.html#removed-features). It does support `NB_DOMAIN` but that is not sufficient for disambiguating typeids across projects (to wit: we currently define `NB_DOMAIN` and it was still necessary to move everything to a nested namespace);
3. Build the [nanobind library itself as a shared object](https://github.com/wjakob/nanobind/blob/master/cmake/nanobind-config.cmake#L127) (and link it to both the extensions and `MLIRPythonSupport`).
4. CMake to make this work, in-tree, out-of-tree, downstream, upstream, etc.
# Testing
Three tests are added here
1. `PythonTestModuleNanobind` is ported to use `PyConcreteType<PyTestType>` instead of `mlir_type_subclass` and `PyConcreteAttribute<PyTestAttr>` instead of `mlir_atrr_subclass`, verifying this works for non-core extensions in-tree;
2. `StandaloneExtensionNanobind` is ported to use `struct PyCustomType : mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyConcreteType<PyCustomType>` instead of `mlir_type_subclass` verifying this works for non-core extensions out-of-tree;
3. `StandaloneExtensionNanobind`'s `smoketest` is extended to also load another bindings package (namely `mlir`) verifying `MLIR_BINDINGS_PYTHON_DOMAIN` successfully disambiguates symbols and typeids.
I have also tested this downstream: https://github.com/llvm/eudsl/pull/287 as well run the following builder bots:
mlir-nvidia-gcc7: https://lab.llvm.org/buildbot/#/buildrequests/6654424?redirect_to_build=true
I have also tested against IREE: https://github.com/iree-org/iree/pull/21916
# Integration
It is highly recommended to set the CMake var `MLIR_BINDINGS_PYTHON_NB_DOMAIN` (which will also determine `MLIR_BINDINGS_PYTHON_DOMAIN`) to something unique for each downstream. This can also be passed explicitly to `add_mlir_python_modules` if your project builds multiple bindings packages. I added a `WARNING` to this effect in `AddMLIRPython.cmake`.
[^3]: Python values being typed correctly when exiting from cpp;
[^1]: Specifically when the modules are imported using `importlib`, which occurs with nanobind's [stubgen](https://github.com/wjakob/nanobind/blob/master/src/stubgen.py#L965);
[^2]: The workaround we implemented was a class method for the dialect bindings called `Class.isinstance(...)`;
---
Patch is 219.86 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/171775.diff
28 Files Affected:
- (modified) mlir/cmake/modules/AddMLIRPython.cmake (+190-31)
- (modified) mlir/docs/Bindings/Python.md (+7)
- (modified) mlir/examples/standalone/include/Standalone-c/Dialects.h (+7)
- (modified) mlir/examples/standalone/lib/CAPI/Dialects.cpp (+13)
- (modified) mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp (+27)
- (modified) mlir/examples/standalone/test/python/smoketest.py (+4)
- (modified) mlir/include/mlir-c/Support.h (+2)
- (renamed) mlir/include/mlir/Bindings/Python/Globals.h (+7-9)
- (renamed) mlir/include/mlir/Bindings/Python/IRCore.h (+584-102)
- (modified) mlir/include/mlir/Bindings/Python/IRTypes.h (+6-3)
- (renamed) mlir/include/mlir/Bindings/Python/NanobindUtils.h ()
- (modified) mlir/lib/Bindings/Python/DialectSMT.cpp (+1-1)
- (renamed) mlir/lib/Bindings/Python/Globals.cpp (+16-4)
- (modified) mlir/lib/Bindings/Python/IRAffine.cpp (+30-10)
- (modified) mlir/lib/Bindings/Python/IRAttributes.cpp (+16-14)
- (modified) mlir/lib/Bindings/Python/IRCore.cpp (+916-1068)
- (modified) mlir/lib/Bindings/Python/IRInterfaces.cpp (+3-3)
- (modified) mlir/lib/Bindings/Python/IRTypes.cpp (+25-17)
- (modified) mlir/lib/Bindings/Python/MainModule.cpp (+34-126)
- (modified) mlir/lib/Bindings/Python/Pass.cpp (+28-15)
- (modified) mlir/lib/Bindings/Python/Pass.h (+3-2)
- (modified) mlir/lib/Bindings/Python/Rewrite.cpp (+54-58)
- (modified) mlir/lib/Bindings/Python/Rewrite.h (+3-3)
- (modified) mlir/python/CMakeLists.txt (+12-10)
- (modified) mlir/test/Examples/standalone/test.wheel.toy (+2)
- (modified) mlir/test/python/dialects/python_test.py (+24-12)
- (modified) mlir/test/python/lib/PythonTestModuleNanobind.cpp (+46-24)
- (modified) utils/bazel/llvm-project-overlay/mlir/BUILD.bazel (+1-1)
``````````diff
diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index 8c301faf0941a..f4d078dfe7118 100644
--- a/mlir/cmake/modules/AddMLIRPython.cmake
+++ b/mlir/cmake/modules/AddMLIRPython.cmake
@@ -228,7 +228,7 @@ endfunction()
# aggregate dylib that is linked against.
function(declare_mlir_python_extension name)
cmake_parse_arguments(ARG
- ""
+ "_PRIVATE_SUPPORT_LIB"
"ROOT_DIR;MODULE_NAME;ADD_TO_PARENT"
"SOURCES;PRIVATE_LINK_LIBS;EMBED_CAPI_LINK_LIBS"
${ARGN})
@@ -236,6 +236,11 @@ function(declare_mlir_python_extension name)
if(NOT ARG_ROOT_DIR)
set(ARG_ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
endif()
+ if(ARG__PRIVATE_SUPPORT_LIB)
+ set(SOURCES_TYPE "support")
+ else()
+ set(SOURCES_TYPE "extension")
+ endif()
set(_install_destination "src/python/${name}")
add_library(${name} INTERFACE)
@@ -243,7 +248,7 @@ function(declare_mlir_python_extension name)
# 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"
- mlir_python_SOURCES_TYPE extension
+ mlir_python_SOURCES_TYPE "${SOURCES_TYPE}"
mlir_python_EXTENSION_MODULE_NAME "${ARG_MODULE_NAME}"
mlir_python_EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
mlir_python_DEPENDS ""
@@ -297,6 +302,58 @@ function(_mlir_python_install_sources name source_root_dir destination)
)
endfunction()
+function(build_nanobind_lib)
+ cmake_parse_arguments(ARG
+ ""
+ "INSTALL_COMPONENT;INSTALL_DESTINATION;OUTPUT_DIRECTORY;MLIR_BINDINGS_PYTHON_NB_DOMAIN"
+ ""
+ ${ARGN})
+
+ # Only build in free-threaded mode if the Python ABI supports it.
+ # See https://github.com/wjakob/nanobind/blob/4ba51fcf795971c5d603d875ae4184bc0c9bd8e6/cmake/nanobind-config.cmake#L363-L371.
+ if (NB_ABI MATCHES "[0-9]t")
+ set(_ft "-ft")
+ endif()
+ # nanobind does a string match on the suffix to figure out whether to build
+ # the lib with free threading...
+ set(NB_LIBRARY_TARGET_NAME "nanobind${_ft}-${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}")
+ set(NB_LIBRARY_TARGET_NAME "${NB_LIBRARY_TARGET_NAME}" PARENT_SCOPE)
+ nanobind_build_library(${NB_LIBRARY_TARGET_NAME} AS_SYSINCLUDE)
+ target_compile_definitions(${NB_LIBRARY_TARGET_NAME}
+ PRIVATE
+ NB_DOMAIN=${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}
+ )
+ if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+ # nanobind handles this correctly for MacOS by explicitly setting -U for all the necessary Python symbols
+ # (see https://github.com/wjakob/nanobind/blob/master/cmake/darwin-ld-cpython.sym)
+ # but since we set -z,defs in llvm/cmake/modules/HandleLLVMOptions.cmake:340 for all Linux shlibs
+ # we need to negate it here (we could have our own linux-ld-cpython.sym but that would be too much
+ # maintenance).
+ target_link_options(${NB_LIBRARY_TARGET_NAME} PRIVATE "LINKER:-z,undefs")
+ endif()
+ # nanobind configures with LTO for shared build which doesn't work everywhere
+ # (see https://github.com/llvm/llvm-project/issues/139602).
+ if(NOT LLVM_ENABLE_LTO)
+ set_target_properties(${NB_LIBRARY_TARGET_NAME} PROPERTIES
+ INTERPROCEDURAL_OPTIMIZATION_RELEASE OFF
+ INTERPROCEDURAL_OPTIMIZATION_MINSIZEREL OFF
+ )
+ endif()
+ set_target_properties(${NB_LIBRARY_TARGET_NAME} PROPERTIES
+ LIBRARY_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
+ BINARY_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
+ # Needed for windows (and doesn't hurt others).
+ RUNTIME_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
+ ARCHIVE_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
+ )
+ mlir_python_setup_extension_rpath(${NB_LIBRARY_TARGET_NAME})
+ install(TARGETS ${NB_LIBRARY_TARGET_NAME}
+ COMPONENT ${ARG_INSTALL_COMPONENT}
+ LIBRARY DESTINATION "${ARG_INSTALL_DESTINATION}"
+ RUNTIME DESTINATION "${ARG_INSTALL_DESTINATION}"
+ )
+endfunction()
+
# Function: add_mlir_python_modules
# Adds python modules to a project, building them from a list of declared
# source groupings (see declare_mlir_python_sources and
@@ -308,6 +365,11 @@ endfunction()
# for non-relocatable modules or a deeper directory tree for relocatable.
# INSTALL_PREFIX: Prefix into the install tree for installing the package.
# Typically mirrors the path above but without an absolute path.
+# MLIR_BINDINGS_PYTHON_NB_DOMAIN: nanobind (and MLIR) domain within which
+# extensions will be compiled. This determines whether this package
+# will share nanobind types with other bindings packages. Expected to be unique
+# per project (and per specific set of bindings, for projects with multiple
+# bindings packages).
# DECLARED_SOURCES: List of declared source groups to include. The entire
# DAG of source modules is included.
# COMMON_CAPI_LINK_LIBS: List of dylibs (typically one) to make every
@@ -315,11 +377,32 @@ endfunction()
function(add_mlir_python_modules name)
cmake_parse_arguments(ARG
""
- "ROOT_PREFIX;INSTALL_PREFIX"
+ "ROOT_PREFIX;INSTALL_PREFIX;MLIR_BINDINGS_PYTHON_NB_DOMAIN"
"COMMON_CAPI_LINK_LIBS;DECLARED_SOURCES"
${ARGN})
+
+ # TODO(max): do the same for MLIR_PYTHON_PACKAGE_PREFIX?
+ if((NOT ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN) AND MLIR_BINDINGS_PYTHON_NB_DOMAIN)
+ set(ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN ${MLIR_BINDINGS_PYTHON_NB_DOMAIN})
+ endif()
+ if((NOT ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN) OR ("${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}" STREQUAL ""))
+ message(WARNING "MLIR_BINDINGS_PYTHON_NB_DOMAIN CMake var is not set - setting to a default `mlir`.\
+ It is highly recommend to set this to something unique so that your project's Python bindings do not collide with\
+ others'. You also pass explicitly to `add_mlir_python_modules`.\
+ See https://github.com/llvm/llvm-project/pull/171775 for more information.")
+ set(ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN "mlir")
+ endif()
+
+ # This call sets NB_LIBRARY_TARGET_NAME.
+ build_nanobind_lib(
+ INSTALL_COMPONENT ${name}
+ INSTALL_DESTINATION "${ARG_INSTALL_PREFIX}/_mlir_libs"
+ OUTPUT_DIRECTORY "${ARG_ROOT_PREFIX}/_mlir_libs"
+ MLIR_BINDINGS_PYTHON_NB_DOMAIN ${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}
+ )
+
# Helper to process an individual target.
- function(_process_target modules_target sources_target)
+ function(_process_target modules_target sources_target support_libs)
get_target_property(_source_type ${sources_target} mlir_python_SOURCES_TYPE)
if(_source_type STREQUAL "pure")
@@ -337,16 +420,20 @@ function(add_mlir_python_modules name)
get_target_property(_module_name ${sources_target} mlir_python_EXTENSION_MODULE_NAME)
# 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}"
+ add_mlir_python_extension(${_extension_target} "${_module_name}" ${NB_LIBRARY_TARGET_NAME}
INSTALL_COMPONENT ${modules_target}
INSTALL_DIR "${ARG_INSTALL_PREFIX}/_mlir_libs"
OUTPUT_DIRECTORY "${ARG_ROOT_PREFIX}/_mlir_libs"
+ MLIR_BINDINGS_PYTHON_NB_DOMAIN ${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}
LINK_LIBS PRIVATE
${sources_target}
${ARG_COMMON_CAPI_LINK_LIBS}
+ ${support_libs}
)
add_dependencies(${modules_target} ${_extension_target})
mlir_python_setup_extension_rpath(${_extension_target})
+ elseif(_source_type STREQUAL "support")
+ # do nothing because already built
else()
message(SEND_ERROR "Unrecognized source type '${_source_type}' for python source target ${sources_target}")
return()
@@ -356,8 +443,36 @@ function(add_mlir_python_modules name)
# Build the modules target.
add_custom_target(${name} ALL)
_flatten_mlir_python_targets(_flat_targets ${ARG_DECLARED_SOURCES})
+
+ # Build all support libs first.
+ set(_mlir_python_support_libs)
+ foreach(sources_target ${_flat_targets})
+ get_target_property(_source_type ${sources_target} mlir_python_SOURCES_TYPE)
+ if(_source_type STREQUAL "support")
+ get_target_property(_module_name ${sources_target} mlir_python_EXTENSION_MODULE_NAME)
+ # Use a similar mechanism as nanobind to help the runtime loader pick the correct lib.
+ set(_module_name "${_module_name}-${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}")
+ set(_extension_target "${name}.extension.${_module_name}.so")
+ add_mlir_python_extension(${_extension_target} "${_module_name}" ${NB_LIBRARY_TARGET_NAME}
+ INSTALL_COMPONENT ${name}
+ INSTALL_DIR "${ARG_INSTALL_PREFIX}/_mlir_libs"
+ OUTPUT_DIRECTORY "${ARG_ROOT_PREFIX}/_mlir_libs"
+ MLIR_BINDINGS_PYTHON_NB_DOMAIN ${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}
+ _PRIVATE_SUPPORT_LIB
+ LINK_LIBS PRIVATE
+ LLVMSupport
+ ${sources_target}
+ ${ARG_COMMON_CAPI_LINK_LIBS}
+ )
+ add_dependencies(${name} ${_extension_target})
+ mlir_python_setup_extension_rpath(${_extension_target})
+ list(APPEND _mlir_python_support_libs "${_extension_target}")
+ endif()
+ endforeach()
+
+ # Build extensions.
foreach(sources_target ${_flat_targets})
- _process_target(${name} ${sources_target})
+ _process_target(${name} ${sources_target} "${_mlir_python_support_libs}")
endforeach()
# Create an install target.
@@ -622,7 +737,7 @@ function(add_mlir_python_common_capi_library name)
set_target_properties(${name} PROPERTIES
LIBRARY_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
BINARY_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
- # Needed for windows (and don't hurt others).
+ # Needed for windows (and doesn't hurt others).
RUNTIME_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
ARCHIVE_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
)
@@ -742,10 +857,10 @@ endfunction()
################################################################################
# Build python extension
################################################################################
-function(add_mlir_python_extension libname extname)
+function(add_mlir_python_extension libname extname nb_library_target_name)
cmake_parse_arguments(ARG
- ""
- "INSTALL_COMPONENT;INSTALL_DIR;OUTPUT_DIRECTORY"
+ "_PRIVATE_SUPPORT_LIB"
+ "INSTALL_COMPONENT;INSTALL_DIR;OUTPUT_DIRECTORY;MLIR_BINDINGS_PYTHON_NB_DOMAIN"
"SOURCES;LINK_LIBS"
${ARGN})
if(ARG_UNPARSED_ARGUMENTS)
@@ -761,10 +876,41 @@ 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}
+ if(ARG__PRIVATE_SUPPORT_LIB)
+ add_library(${libname} SHARED ${ARG_SOURCES})
+ if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+ # nanobind handles this correctly for MacOS by explicitly setting -U for all the necessary Python symbols
+ # (see https://github.com/wjakob/nanobind/blob/master/cmake/darwin-ld-cpython.sym)
+ # but since we set -z,defs in llvm/cmake/modules/HandleLLVMOptions.cmake:340 for all Linux shlibs
+ # we need to negate it here (we could have our own linux-ld-cpython.sym but that would be too much
+ # maintenance).
+ target_link_options(${libname} PRIVATE "LINKER:-z,undefs")
+ endif()
+ nanobind_link_options(${libname})
+ target_compile_definitions(${libname}
+ PRIVATE
+ NB_DOMAIN=${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}
+ MLIR_CAPI_BUILDING_LIBRARY=1
+ )
+ if(MSVC)
+ set_property(TARGET ${libname} PROPERTY WINDOWS_EXPORT_ALL_SYMBOLS ON)
+ endif()
+ else()
+ nanobind_add_module(${libname}
+ NB_DOMAIN ${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}
+ FREE_THREADED
+ NB_SHARED
+ ${ARG_SOURCES}
+ )
+ target_compile_definitions(${libname}
+ PRIVATE
+ MLIR_BINDINGS_PYTHON_DOMAIN=${MLIR_BINDINGS_PYTHON_NB_DOMAIN}
+ )
+ endif()
+ target_link_libraries(${libname} PRIVATE ${nb_library_target_name})
+ target_compile_definitions(${libname}
+ PRIVATE
+ MLIR_BINDINGS_PYTHON_DOMAIN=${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}
)
if(APPLE)
# In llvm/cmake/modules/HandleLLVMOptions.cmake:268 we set -Wl,-flat_namespace which breaks
@@ -778,29 +924,28 @@ function(add_mlir_python_extension libname extname)
# 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}
+ target_compile_options(${nb_library_target_name}
PRIVATE
-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
- -Wno-missing-field-initializers
+ -Wno-cast-qual
+ -Wno-covered-switch-default
+ -Wno-deprecated-literal-operator
+ -Wno-nested-anon-types
+ -Wno-unused-parameter
+ -Wno-zero-length-array
+ -Wno-missing-field-initializers
${eh_rtti_enable})
target_compile_options(${libname}
PRIVATE
-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
- -Wno-missing-field-initializers
+ -Wno-cast-qual
+ -Wno-covered-switch-default
+ -Wno-deprecated-literal-operator
+ -Wno-nested-anon-types
+ -Wno-unused-parameter
+ -Wno-zero-length-array
+ -Wno-missing-field-initializers
${eh_rtti_enable})
endif()
@@ -818,12 +963,26 @@ function(add_mlir_python_extension libname extname)
target_compile_options(${libname} PRIVATE ${eh_rtti_enable})
+ # Quoting CMake:
+ #
+ # "If you use it on normal shared libraries which other targets link against, on some platforms a
+ # linker will insert a full path to the library (as specified at link time) into the dynamic section of the
+ # dependent binary. Therefore, once installed, dynamic loader may eventually fail to locate the library
+ # for the binary."
+ #
+ # So for support libs we do need an SO name but for extensions we do not (they're MODULEs anyway -
+ # i.e., can't be linked against, only loaded).
+ if (ARG__PRIVATE_SUPPORT_LIB)
+ set(_no_soname OFF)
+ else ()
+ set(_no_soname ON)
+ endif ()
# Configure the output to match python expectations.
set_target_properties(
${libname} PROPERTIES
LIBRARY_OUTPUT_DIRECTORY ${ARG_OUTPUT_DIRECTORY}
OUTPUT_NAME "${extname}"
- NO_SONAME ON
+ NO_SONAME ${_no_soname}
)
if(WIN32)
diff --git a/mlir/docs/Bindings/Python.md b/mlir/docs/Bindings/Python.md
index 4f4f531f7723c..4278774933a4a 100644
--- a/mlir/docs/Bindings/Python.md
+++ b/mlir/docs/Bindings/Python.md
@@ -37,6 +37,13 @@
LLVM ERROR: ... unregistered/uninitialized dialect/type/pass ...`
```
+* **`MLIR_BINDINGS_PYTHON_NB_DOMAIN`**: `STRING`
+
+ nanobind (and MLIR) domain within which extensions will be compiled.
+ This determines whether this package will share nanobind types with other bindings packages.
+ Expected to be unique per project (and per specific set of bindings, for projects with multiple bindings packages).
+ Can also be passed explicitly to `add_mlir_python_modules`.
+
### Recommended development practices
It is recommended to use a Python virtual environment. Many ways exist for this,
diff --git a/mlir/examples/standalone/include/Standalone-c/Dialects.h b/mlir/examples/standalone/include/Standalone-c/Dialects.h
index b3e47752ccc69..5aa9e004cb9fe 100644
--- a/mlir/examples/standalone/include/Standalone-c/Dialects.h
+++ b/mlir/examples/standalone/include/Standalone-c/Dialects.h
@@ -17,6 +17,13 @@ extern "C" {
MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Standalone, standalone);
+MLIR_CAPI_EXPORTED MlirType mlirStandaloneCustomTypeGet(MlirContext ctx,
+ MlirStringRef value);
+
+MLIR_CAPI_EXPORTED bool mlirStandaloneTypeIsACustomType(MlirType t);
+
+MLIR_CAPI_EXPORTED MlirTypeID mlirStandaloneCustomTypeGetTypeID(void);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/examples/standalone/lib/CAPI/Dialects.cpp b/mlir/examples/standalone/lib/CAPI/Dialects.cpp
index 98006e81a3d26..4de55ba485490 100644
--- a/mlir/examples/standalone/lib/CAPI/Dialects.cpp
+++ b/mlir/examples/standalone/lib/CAPI/Dialects.cpp
@@ -9,7 +9,20 @@
#include "Standalone-c/Dialects.h"
#include "Standalone/StandaloneDialect.h"
+#include "Standalone/StandaloneTypes.h"
#include "mlir/CAPI/Registration.h"
MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Standalone, standalone,
mlir::standalone::StandaloneDialect)
+
+MlirType mlirStandaloneCustomTypeGet(MlirContext ctx, MlirStringRef value) {
+ return wrap(mlir::standalone::CustomType::get(unwrap(ctx), unwrap(value)));
+}
+
+bool mlirStandaloneTypeIsACustomType(MlirType t) {
+ return llvm::isa<mlir::standalone::CustomType>(unwrap(t));
+}
+
+MlirTypeID mlirStandaloneCustomTypeGetTypeID() {
+ return wrap(mlir::standalone::CustomType::getTypeID());
+}
diff --git a/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp b/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp
index 0ec6cdfa7994b..c568369913595 100644
--- a/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp
+++ b/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp
@@ -11,17 +11,44 @@
#include "Standalone-c/Dialects.h"
#include "mlir-c/Dialect/Arith.h"
+#include "mlir/Bindings/Python/IRCore.h"
#include "mlir/Bindings/Python/Nanobind.h"
#include "mlir/Bindings/Python/NanobindAdaptors.h"
namespace nb = nanobind;
+struct PyCustomType
+ : mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyConcreteType<PyCustomType> {
+ static constexpr IsAFunctionTy isaFunction = mlirStandaloneTypeIsACustomType;
+ static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+ mlirStandaloneCustomTypeGetTypeID;
+ static constexpr const char *pyClassName = "CustomType";
+ using PyConcreteType::PyConcreteType;
+
+ static void bindDerived(ClassTy &c) {
+ c.def_static(
+ "get",
+ [](const std::string &value,
+ mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::DefaultingPyMlirContext
+ context) {
+ return PyCustomType(
+ context->getRef(),
+ mlirStandaloneCustomTypeGet(
+ context.get()->get(),
+ mlirStringRefCreateFromCString(value.c_str())));
+ },
+ nb::arg("value"), nb::arg("context").none() = nb::none());
+ }
+};
+
NB_MODULE(_standaloneDialectsNanobind, m) {
//===--------------------------------------------------------------------===//
// standalone dialect
//===--------------------------------------------------------------------===//
auto standaloneM = m.def_submodule("standalone");
+ PyCustomType::bind(standaloneM);
+
standaloneM.def(
"register_dialects",
[](MlirContext context, bool load) {
diff --git a/mlir/examples/standalone/test/python/smoketest.py b/mlir/examples/standalone/test/python/smoketest.py
index 09040eb2f45dc..fe4e40e6e8a99 100644
--- a/mlir/examples/standalone/test/python/smoketest.py
+++ b/mlir/examples/standalone/test/python/smoketest.py
@@ -19,6 +19,10 @@
# CHECK: standalone.foo %[[C2]] : i32
print(str(standalone_module), file=sys.stderr)
+ custom_type = standalone_d.CustomType.get("foo")
+ # CHECK: !standalone.custom<"foo">
+ print(custom_type, file=sys.stderr)
+
# CHECK: Testing mlir package
print("Testing mlir package", file=sys.stderr)
diff --git a/mlir/include/mlir-c/Support.h b/mlir/include/mlir-c/Support.h
index 78fc94f93439e..6abd8894227c3 100644
--- a/mlir/include/mlir-c/Support.h
+++ b/mlir/include/mlir-c/Support.h
@@ -46,6 +46,8 @@
#define MLIR_CAPI_EXPORTED __attribute__((visibility("default")))
#endif
+#d...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/171775
More information about the llvm-commits
mailing list