shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
Reverts revert #118517 after (hopefully) fixing builders
(https://github.com/llvm/llvm-zorg/pull/328,
https://github.com/llvm/llvm-zorg/pull/327)

This reverts commit 61bf308cf2fc32452f14861c102ace89f5f36fec.
This commit is contained in:
Maksim Levental 2024-12-09 16:37:43 -05:00 committed by GitHub
parent 1d4b5c161f
commit 392622d084
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
25 changed files with 1184 additions and 105 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
mlir/cmake/modules/AddMLIRPython.cmake
View File

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

39
mlir/cmake/modules/MLIRDetectPythonEnv.cmake
View File

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

20
mlir/docs/Bindings/Python.md
View File

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

22
mlir/examples/standalone/python/CMakeLists.txt
View File

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

35
mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp Normal file
View File

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

7
mlir/examples/standalone/python/StandaloneExtension.cpp → mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
View File

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

2
mlir/examples/standalone/python/mlir_standalone/dialects/standalone.py → mlir/examples/standalone/python/mlir_standalone/dialects/standalone_nanobind.py
View File

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

6
mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py Normal file
View File

@ -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 *

14
mlir/examples/standalone/test/python/smoketest.py
View File

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

59
mlir/include/mlir/Bindings/Python/Diagnostics.h Normal file
View File

@ -0,0 +1,59 @@
//===- Diagnostics.h - Helpers for diagnostics in Python bindings ---------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
#ifndef MLIR_BINDINGS_PYTHON_DIAGNOSTICS_H
#define MLIR_BINDINGS_PYTHON_DIAGNOSTICS_H
#include <cassert>
#include <string>
#include "mlir-c/Diagnostics.h"
#include "mlir-c/IR.h"
#include "llvm/ADT/StringRef.h"
namespace mlir {
namespace python {
/// RAII scope intercepting all diagnostics into a string. The message must be
/// checked before this goes out of scope.
class CollectDiagnosticsToStringScope {
public:
explicit CollectDiagnosticsToStringScope(MlirContext ctx) : context(ctx) {
handlerID = mlirContextAttachDiagnosticHandler(ctx, &handler, &errorMessage,
/*deleteUserData=*/nullptr);
}
~CollectDiagnosticsToStringScope() {
assert(errorMessage.empty() && "unchecked error message");
mlirContextDetachDiagnosticHandler(context, handlerID);
}
[[nodiscard]] std::string takeMessage() { return std::move(errorMessage); }
private:
static MlirLogicalResult handler(MlirDiagnostic diag, void *data) {
auto printer = +[](MlirStringRef message, void *data) {
*static_cast<std::string *>(data) +=
llvm::StringRef(message.data, message.length);
};
MlirLocation loc = mlirDiagnosticGetLocation(diag);
*static_cast<std::string *>(data) += "at ";
mlirLocationPrint(loc, printer, data);
*static_cast<std::string *>(data) += ": ";
mlirDiagnosticPrint(diag, printer, data);
return mlirLogicalResultSuccess();
}
MlirContext context;
MlirDiagnosticHandlerID handlerID;
std::string errorMessage = "";
};
} // namespace python
} // namespace mlir
#endif // MLIR_BINDINGS_PYTHON_DIAGNOSTICS_H

671
mlir/include/mlir/Bindings/Python/NanobindAdaptors.h Normal file
View File

@ -0,0 +1,671 @@
//===- NanobindAdaptors.h - Interop with MLIR APIs via nanobind -----------===//
//
// 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 nanobind. 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_NANOBINDADAPTORS_H
#define MLIR_BINDINGS_PYTHON_NANOBINDADAPTORS_H
#include <nanobind/nanobind.h>
#include <nanobind/stl/string.h>
#include <cstdint>
#include "mlir-c/Bindings/Python/Interop.h"
#include "mlir-c/Diagnostics.h"
#include "mlir-c/IR.h"
#include "llvm/ADT/Twine.h"
// Raw CAPI type casters need to be declared before use, so always include them
// first.
namespace nanobind {
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 nanobind::object mlirApiObjectToCapsule(nanobind::handle apiObject) {
if (PyCapsule_CheckExact(apiObject.ptr()))
return nanobind::borrow<nanobind::object>(apiObject);
if (!nanobind::hasattr(apiObject, MLIR_PYTHON_CAPI_PTR_ATTR)) {
std::string repr = nanobind::cast<std::string>(nanobind::repr(apiObject));
throw nanobind::type_error(
(llvm::Twine("Expected an MLIR object (got ") + repr + ").")
.str()
.c_str());
}
return apiObject.attr(MLIR_PYTHON_CAPI_PTR_ATTR);
}
// Note: Currently all of the following support cast from nanobind::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> {
NB_TYPE_CASTER(MlirAffineMap, const_name("MlirAffineMap"));
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
nanobind::object capsule = mlirApiObjectToCapsule(src);
value = mlirPythonCapsuleToAffineMap(capsule.ptr());
if (mlirAffineMapIsNull(value)) {
return false;
}
return !mlirAffineMapIsNull(value);
}
static handle from_cpp(MlirAffineMap v, rv_policy,
cleanup_list *cleanup) noexcept {
nanobind::object capsule =
nanobind::steal<nanobind::object>(mlirPythonAffineMapToCapsule(v));
return nanobind::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> {
NB_TYPE_CASTER(MlirAttribute, const_name("MlirAttribute"));
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
nanobind::object capsule = mlirApiObjectToCapsule(src);
value = mlirPythonCapsuleToAttribute(capsule.ptr());
return !mlirAttributeIsNull(value);
}
static handle from_cpp(MlirAttribute v, rv_policy,
cleanup_list *cleanup) noexcept {
nanobind::object capsule =
nanobind::steal<nanobind::object>(mlirPythonAttributeToCapsule(v));
return nanobind::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> {
NB_TYPE_CASTER(MlirBlock, const_name("MlirBlock"));
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
nanobind::object capsule = mlirApiObjectToCapsule(src);
value = mlirPythonCapsuleToBlock(capsule.ptr());
return !mlirBlockIsNull(value);
}
};
/// Casts object -> MlirContext.
template <>
struct type_caster<MlirContext> {
NB_TYPE_CASTER(MlirContext, const_name("MlirContext"));
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
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 = nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
.attr("Context")
.attr("current");
}
nanobind::object capsule = mlirApiObjectToCapsule(src);
value = mlirPythonCapsuleToContext(capsule.ptr());
return !mlirContextIsNull(value);
}
};
/// Casts object <-> MlirDialectRegistry.
template <>
struct type_caster<MlirDialectRegistry> {
NB_TYPE_CASTER(MlirDialectRegistry, const_name("MlirDialectRegistry"));
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
nanobind::object capsule = mlirApiObjectToCapsule(src);
value = mlirPythonCapsuleToDialectRegistry(capsule.ptr());
return !mlirDialectRegistryIsNull(value);
}
static handle from_cpp(MlirDialectRegistry v, rv_policy,
cleanup_list *cleanup) noexcept {
nanobind::object capsule = nanobind::steal<nanobind::object>(
mlirPythonDialectRegistryToCapsule(v));
return nanobind::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> {
NB_TYPE_CASTER(MlirLocation, const_name("MlirLocation"));
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
if (src.is_none()) {
// Gets the current thread-bound context.
src = nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
.attr("Location")
.attr("current");
}
nanobind::object capsule = mlirApiObjectToCapsule(src);
value = mlirPythonCapsuleToLocation(capsule.ptr());
return !mlirLocationIsNull(value);
}
static handle from_cpp(MlirLocation v, rv_policy,
cleanup_list *cleanup) noexcept {
nanobind::object capsule =
nanobind::steal<nanobind::object>(mlirPythonLocationToCapsule(v));
return nanobind::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> {
NB_TYPE_CASTER(MlirModule, const_name("MlirModule"));
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
nanobind::object capsule = mlirApiObjectToCapsule(src);
value = mlirPythonCapsuleToModule(capsule.ptr());
return !mlirModuleIsNull(value);
}
static handle from_cpp(MlirModule v, rv_policy,
cleanup_list *cleanup) noexcept {
nanobind::object capsule =
nanobind::steal<nanobind::object>(mlirPythonModuleToCapsule(v));
return nanobind::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> {
NB_TYPE_CASTER(MlirFrozenRewritePatternSet,
const_name("MlirFrozenRewritePatternSet"));
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
nanobind::object capsule = mlirApiObjectToCapsule(src);
value = mlirPythonCapsuleToFrozenRewritePatternSet(capsule.ptr());
return value.ptr != nullptr;
}
static handle from_cpp(MlirFrozenRewritePatternSet v, rv_policy, handle) {
nanobind::object capsule = nanobind::steal<nanobind::object>(
mlirPythonFrozenRewritePatternSetToCapsule(v));
return nanobind::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> {
NB_TYPE_CASTER(MlirOperation, const_name("MlirOperation"));
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
nanobind::object capsule = mlirApiObjectToCapsule(src);
value = mlirPythonCapsuleToOperation(capsule.ptr());
return !mlirOperationIsNull(value);
}
static handle from_cpp(MlirOperation v, rv_policy,
cleanup_list *cleanup) noexcept {
if (v.ptr == nullptr)
return nanobind::none();
nanobind::object capsule =
nanobind::steal<nanobind::object>(mlirPythonOperationToCapsule(v));
return nanobind::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> {
NB_TYPE_CASTER(MlirValue, const_name("MlirValue"));
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
nanobind::object capsule = mlirApiObjectToCapsule(src);
value = mlirPythonCapsuleToValue(capsule.ptr());
return !mlirValueIsNull(value);
}
static handle from_cpp(MlirValue v, rv_policy,
cleanup_list *cleanup) noexcept {
if (v.ptr == nullptr)
return nanobind::none();
nanobind::object capsule =
nanobind::steal<nanobind::object>(mlirPythonValueToCapsule(v));
return nanobind::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> {
NB_TYPE_CASTER(MlirPassManager, const_name("MlirPassManager"));
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
nanobind::object capsule = mlirApiObjectToCapsule(src);
value = mlirPythonCapsuleToPassManager(capsule.ptr());
return !mlirPassManagerIsNull(value);
}
};
/// Casts object <-> MlirTypeID.
template <>
struct type_caster<MlirTypeID> {
NB_TYPE_CASTER(MlirTypeID, const_name("MlirTypeID"));
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
nanobind::object capsule = mlirApiObjectToCapsule(src);
value = mlirPythonCapsuleToTypeID(capsule.ptr());
return !mlirTypeIDIsNull(value);
}
static handle from_cpp(MlirTypeID v, rv_policy,
cleanup_list *cleanup) noexcept {
if (v.ptr == nullptr)
return nanobind::none();
nanobind::object capsule =
nanobind::steal<nanobind::object>(mlirPythonTypeIDToCapsule(v));
return nanobind::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> {
NB_TYPE_CASTER(MlirType, const_name("MlirType"));
bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) {
nanobind::object capsule = mlirApiObjectToCapsule(src);
value = mlirPythonCapsuleToType(capsule.ptr());
return !mlirTypeIsNull(value);
}
static handle from_cpp(MlirType t, rv_policy,
cleanup_list *cleanup) noexcept {
nanobind::object capsule =
nanobind::steal<nanobind::object>(mlirPythonTypeToCapsule(t));
return nanobind::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 nanobind
namespace mlir {
namespace python {
namespace nanobind_adaptors {
/// Provides a facility like nanobind::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 nanobind machinery to
/// do it. These are not "real" nanobind 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 nanobind.
class pure_subclass {
public:
pure_subclass(nanobind::handle scope, const char *derivedClassName,
const nanobind::object &superClass) {
nanobind::object pyType =
nanobind::borrow<nanobind::object>((PyObject *)&PyType_Type);
nanobind::object metaclass = pyType(superClass);
nanobind::dict attributes;
thisClass = metaclass(derivedClassName, nanobind::make_tuple(superClass),
attributes);
scope.attr(derivedClassName) = thisClass;
}
template <typename Func, typename... Extra>
pure_subclass &def(const char *name, Func &&f, const Extra &...extra) {
nanobind::object cf = nanobind::cpp_function(
std::forward<Func>(f), nanobind::name(name), nanobind::is_method(),
nanobind::scope(thisClass), extra...);
thisClass.attr(name) = cf;
return *this;
}
template <typename Func, typename... Extra>
pure_subclass &def_property_readonly(const char *name, Func &&f,
const Extra &...extra) {
nanobind::object cf = nanobind::cpp_function(
std::forward<Func>(f), nanobind::name(name), nanobind::is_method(),
nanobind::scope(thisClass), extra...);
auto builtinProperty =
nanobind::borrow<nanobind::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");
nanobind::object cf = nanobind::cpp_function(
std::forward<Func>(f),
nanobind::name(name), // nanobind::scope(thisClass),
extra...);
thisClass.attr(name) = 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");
nanobind::object cf = nanobind::cpp_function(
std::forward<Func>(f),
nanobind::name(name), // nanobind::scope(thisClass),
extra...);
thisClass.attr(name) =
nanobind::borrow<nanobind::object>(PyClassMethod_New(cf.ptr()));
return *this;
}
nanobind::object get_class() const { return thisClass; }
protected:
nanobind::object superClass;
nanobind::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(nanobind::handle scope, const char *attrClassName,
IsAFunctionTy isaFunction,
GetTypeIDFunctionTy getTypeIDFunction = nullptr)
: mlir_attribute_subclass(
scope, attrClassName, isaFunction,
nanobind::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(nanobind::handle scope, const char *typeClassName,
IsAFunctionTy isaFunction,
const nanobind::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 nanobind 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.
nanobind::object newCf = nanobind::cpp_function(
[superCls, isaFunction, captureTypeName](
nanobind::object cls, nanobind::object otherAttribute) {
MlirAttribute rawAttribute =
nanobind::cast<MlirAttribute>(otherAttribute);
if (!isaFunction(rawAttribute)) {
auto origRepr =
nanobind::cast<std::string>(nanobind::repr(otherAttribute));
throw std::invalid_argument(
(llvm::Twine("Cannot cast attribute to ") + captureTypeName +
" (from " + origRepr + ")")
.str());
}
nanobind::object self = superCls.attr("__new__")(cls, otherAttribute);
return self;
},
nanobind::name("__new__"), nanobind::arg("cls"),
nanobind::arg("cast_from_attr"));
thisClass.attr("__new__") = newCf;
// 'isinstance' method.
def_staticmethod(
"isinstance",
[isaFunction](MlirAttribute other) { return isaFunction(other); },
nanobind::arg("other_attribute"));
def("__repr__", [superCls, captureTypeName](nanobind::object self) {
return nanobind::repr(superCls(self))
.attr("replace")(superCls.attr("__name__"), captureTypeName);
});
if (getTypeIDFunction) {
def_staticmethod("get_static_typeid",
[getTypeIDFunction]() { return getTypeIDFunction(); });
nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
.attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
getTypeIDFunction())(nanobind::cpp_function(
[thisClass = thisClass](const nanobind::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(nanobind::handle scope, const char *typeClassName,
IsAFunctionTy isaFunction,
GetTypeIDFunctionTy getTypeIDFunction = nullptr)
: mlir_type_subclass(
scope, typeClassName, isaFunction,
nanobind::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(nanobind::handle scope, const char *typeClassName,
IsAFunctionTy isaFunction,
const nanobind::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 nanobind 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.
nanobind::object newCf = nanobind::cpp_function(
[superCls, isaFunction, captureTypeName](nanobind::object cls,
nanobind::object otherType) {
MlirType rawType = nanobind::cast<MlirType>(otherType);
if (!isaFunction(rawType)) {
auto origRepr =
nanobind::cast<std::string>(nanobind::repr(otherType));
throw std::invalid_argument((llvm::Twine("Cannot cast type to ") +
captureTypeName + " (from " +
origRepr + ")")
.str());
}
nanobind::object self = superCls.attr("__new__")(cls, otherType);
return self;
},
nanobind::name("__new__"), nanobind::arg("cls"),
nanobind::arg("cast_from_type"));
thisClass.attr("__new__") = newCf;
// 'isinstance' method.
def_staticmethod(
"isinstance",
[isaFunction](MlirType other) { return isaFunction(other); },
nanobind::arg("other_type"));
def("__repr__", [superCls, captureTypeName](nanobind::object self) {
return nanobind::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(); });
nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
.attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
getTypeIDFunction())(nanobind::cpp_function(
[thisClass = thisClass](const nanobind::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(nanobind::handle scope, const char *valueClassName,
IsAFunctionTy isaFunction)
: mlir_value_subclass(
scope, valueClassName, isaFunction,
nanobind::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(nanobind::handle scope, const char *valueClassName,
IsAFunctionTy isaFunction,
const nanobind::object &superCls)
: pure_subclass(scope, valueClassName, superCls) {
// Casting constructor. Note that it hard, if not impossible, to properly
// call chain to parent `__init__` in nanobind 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.
nanobind::object newCf = nanobind::cpp_function(
[superCls, isaFunction, captureValueName](nanobind::object cls,
nanobind::object otherValue) {
MlirValue rawValue = nanobind::cast<MlirValue>(otherValue);
if (!isaFunction(rawValue)) {
auto origRepr =
nanobind::cast<std::string>(nanobind::repr(otherValue));
throw std::invalid_argument((llvm::Twine("Cannot cast value to ") +
captureValueName + " (from " +
origRepr + ")")
.str());
}
nanobind::object self = superCls.attr("__new__")(cls, otherValue);
return self;
},
nanobind::name("__new__"), nanobind::arg("cls"),
nanobind::arg("cast_from_value"));
thisClass.attr("__new__") = newCf;
// 'isinstance' method.
def_staticmethod(
"isinstance",
[isaFunction](MlirValue other) { return isaFunction(other); },
nanobind::arg("other_value"));
}
};
} // namespace nanobind_adaptors
/// RAII scope intercepting all diagnostics into a string. The message must be
/// checked before this goes out of scope.
class CollectDiagnosticsToStringScope {
public:
explicit CollectDiagnosticsToStringScope(MlirContext ctx) : context(ctx) {
handlerID = mlirContextAttachDiagnosticHandler(ctx, &handler, &errorMessage,
/*deleteUserData=*/nullptr);
}
~CollectDiagnosticsToStringScope() {
assert(errorMessage.empty() && "unchecked error message");
mlirContextDetachDiagnosticHandler(context, handlerID);
}
[[nodiscard]] std::string takeMessage() { return std::move(errorMessage); }
private:
static MlirLogicalResult handler(MlirDiagnostic diag, void *data) {
auto printer = +[](MlirStringRef message, void *data) {
*static_cast<std::string *>(data) +=
llvm::StringRef(message.data, message.length);
};
MlirLocation loc = mlirDiagnosticGetLocation(diag);
*static_cast<std::string *>(data) += "at ";
mlirLocationPrint(loc, printer, data);
*static_cast<std::string *>(data) += ": ";
mlirDiagnosticPrint(diag, printer, data);
return mlirLogicalResultSuccess();
}
MlirContext context;
MlirDiagnosticHandlerID handlerID;
std::string errorMessage = "";
};
} // namespace python
} // namespace mlir
#endif // MLIR_BINDINGS_PYTHON_NANOBINDADAPTORS_H

43
mlir/include/mlir/Bindings/Python/PybindAdaptors.h
View File

@ -1,4 +1,4 @@
//===- PybindAdaptors.h - Adaptors for interop with MLIR APIs -------------===//
//===- 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.
@ -6,9 +6,10 @@
//
//===----------------------------------------------------------------------===//
// This file contains adaptors for clients of the core MLIR Python APIs to
// interop via MLIR CAPI types. 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).
// 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
@ -611,40 +612,6 @@ public:
} // namespace adaptors
/// RAII scope intercepting all diagnostics into a string. The message must be
/// checked before this goes out of scope.
class CollectDiagnosticsToStringScope {
public:
explicit CollectDiagnosticsToStringScope(MlirContext ctx) : context(ctx) {
handlerID = mlirContextAttachDiagnosticHandler(ctx, &handler, &errorMessage,
/*deleteUserData=*/nullptr);
}
~CollectDiagnosticsToStringScope() {
assert(errorMessage.empty() && "unchecked error message");
mlirContextDetachDiagnosticHandler(context, handlerID);
}
[[nodiscard]] std::string takeMessage() { return std::move(errorMessage); }
private:
static MlirLogicalResult handler(MlirDiagnostic diag, void *data) {
auto printer = +[](MlirStringRef message, void *data) {
*static_cast<std::string *>(data) +=
llvm::StringRef(message.data, message.length);
};
MlirLocation loc = mlirDiagnosticGetLocation(diag);
*static_cast<std::string *>(data) += "at ";
mlirLocationPrint(loc, printer, data);
*static_cast<std::string *>(data) += ": ";
mlirDiagnosticPrint(diag, printer, data);
return mlirLogicalResultSuccess();
}
MlirContext context;
MlirDiagnosticHandlerID handlerID;
std::string errorMessage = "";
};
} // namespace python
} // namespace mlir

4
mlir/lib/Bindings/Python/DialectLLVM.cpp
View File

@ -6,11 +6,13 @@
//
//===----------------------------------------------------------------------===//
#include <string>
#include "mlir-c/Dialect/LLVM.h"
#include "mlir-c/IR.h"
#include "mlir-c/Support.h"
#include "mlir/Bindings/Python/Diagnostics.h"
#include "mlir/Bindings/Python/PybindAdaptors.h"
#include <string>
namespace py = pybind11;
using namespace llvm;

7
mlir/lib/Bindings/Python/TransformInterpreter.cpp
View File

@ -10,14 +10,15 @@
//
//===----------------------------------------------------------------------===//
#include <pybind11/detail/common.h>
#include <pybind11/pybind11.h>
#include "mlir-c/Dialect/Transform/Interpreter.h"
#include "mlir-c/IR.h"
#include "mlir-c/Support.h"
#include "mlir/Bindings/Python/Diagnostics.h"
#include "mlir/Bindings/Python/PybindAdaptors.h"
#include <pybind11/detail/common.h>
#include <pybind11/pybind11.h>
namespace py = pybind11;
namespace {

23
mlir/python/CMakeLists.txt
View File

@ -683,7 +683,9 @@ if(MLIR_INCLUDE_TESTS)
MLIRPythonTestSources.Dialects.PythonTest
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
ADD_TO_PARENT MLIRPythonTestSources.Dialects
SOURCES dialects/python_test.py)
SOURCES
dialects/python_test.py
)
set(LLVM_TARGET_DEFINITIONS
"${MLIR_MAIN_SRC_DIR}/test/python/python_test_ops.td")
mlir_tablegen(
@ -697,12 +699,25 @@ if(MLIR_INCLUDE_TESTS)
ADD_TO_PARENT MLIRPythonTestSources.Dialects.PythonTest
SOURCES "dialects/_python_test_ops_gen.py")
declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtension
MODULE_NAME _mlirPythonTest
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
PythonTestModule.cpp
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
LLVMSupport
EMBED_CAPI_LINK_LIBS

17
mlir/python/mlir/dialects/python_test.py
View File

@ -3,15 +3,14 @@
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from ._python_test_ops_gen import *
from .._mlir_libs._mlirPythonTest import (
TestAttr,
TestType,
TestTensorValue,
TestIntegerRankedTensorType,
)
def register_python_test_dialect(registry):
from .._mlir_libs import _mlirPythonTest
def register_python_test_dialect(registry, use_nanobind):
if use_nanobind:
from .._mlir_libs import _mlirPythonTestNanobind
_mlirPythonTest.register_dialect(registry)
_mlirPythonTestNanobind.register_dialect(registry)
else:
from .._mlir_libs import _mlirPythonTestPybind11
_mlirPythonTestPybind11.register_dialect(registry)

1
mlir/python/requirements.txt
View File

@ -1,3 +1,4 @@
nanobind>=2.0, <3.0
numpy>=1.19.5, <=2.1.2
pybind11>=2.10.0, <=2.13.6
PyYAML>=5.4.0, <=6.0.1

59
mlir/test/python/dialects/python_test.py
View File

@ -1,12 +1,33 @@
# RUN: %PYTHON %s | FileCheck %s
# RUN: %PYTHON %s pybind11 | FileCheck %s
# RUN: %PYTHON %s nanobind | FileCheck %s
import sys
from mlir.ir import *
import mlir.dialects.func as func
import mlir.dialects.python_test as test
import mlir.dialects.tensor as tensor
import mlir.dialects.arith as arith
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):
@ -308,7 +329,7 @@ def testOptionalOperandOp():
@run
def testCustomAttribute():
with Context() as ctx, Location.unknown():
a = test.TestAttr.get()
a = TestAttr.get()
# CHECK: #python_test.test_attr
print(a)
@ -325,11 +346,11 @@ def testCustomAttribute():
print(repr(op2.test_attr))
# The following cast must not assert.
b = test.TestAttr(a)
b = TestAttr(a)
unit = UnitAttr.get()
try:
test.TestAttr(unit)
TestAttr(unit)
except ValueError as e:
assert "Cannot cast attribute to TestAttr" in str(e)
else:
@ -338,7 +359,7 @@ def testCustomAttribute():
# The following must trigger a TypeError from our adaptors and must not
# crash.
try:
test.TestAttr(42)
TestAttr(42)
except TypeError as e:
assert "Expected an MLIR object" in str(e)
else:
@ -347,7 +368,7 @@ def testCustomAttribute():
# The following must trigger a TypeError from pybind (therefore, not
# checking its message) and must not crash.
try:
test.TestAttr(42, 56)
TestAttr(42, 56)
except TypeError:
pass
else:
@ -357,12 +378,12 @@ def testCustomAttribute():
@run
def testCustomType():
with Context() as ctx:
a = test.TestType.get()
a = TestType.get()
# CHECK: !python_test.test_type
print(a)
# The following cast must not assert.
b = test.TestType(a)
b = TestType(a)
# Instance custom types should have typeids
assert isinstance(b.typeid, TypeID)
# Subclasses of ir.Type should not have a static_typeid
@ -374,7 +395,7 @@ def testCustomType():
i8 = IntegerType.get_signless(8)
try:
test.TestType(i8)
TestType(i8)
except ValueError as e:
assert "Cannot cast type to TestType" in str(e)
else:
@ -383,7 +404,7 @@ def testCustomType():
# The following must trigger a TypeError from our adaptors and must not
# crash.
try:
test.TestType(42)
TestType(42)
except TypeError as e:
assert "Expected an MLIR object" in str(e)
else:
@ -392,7 +413,7 @@ def testCustomType():
# The following must trigger a TypeError from pybind (therefore, not
# checking its message) and must not crash.
try:
test.TestType(42, 56)
TestType(42, 56)
except TypeError:
pass
else:
@ -405,7 +426,7 @@ def testTensorValue():
with Context() as ctx, Location.unknown():
i8 = IntegerType.get_signless(8)
class Tensor(test.TestTensorValue):
class Tensor(TestTensorValue):
def __str__(self):
return super().__str__().replace("Value", "Tensor")
@ -425,9 +446,9 @@ def testTensorValue():
# Classes of custom types that inherit from concrete types should have
# static_typeid
assert isinstance(test.TestIntegerRankedTensorType.static_typeid, TypeID)
assert isinstance(TestIntegerRankedTensorType.static_typeid, TypeID)
# And it should be equal to the in-tree concrete type
assert test.TestIntegerRankedTensorType.static_typeid == t.type.typeid
assert TestIntegerRankedTensorType.static_typeid == t.type.typeid
d = tensor.EmptyOp([1, 2, 3], IntegerType.get_signless(5)).result
# CHECK: Value(%{{.*}} = tensor.empty() : tensor<1x2x3xi5>)
@ -491,7 +512,7 @@ def inferReturnTypeComponents():
@run
def testCustomTypeTypeCaster():
with Context() as ctx, Location.unknown():
a = test.TestType.get()
a = TestType.get()
assert a.typeid is not None
b = Type.parse("!python_test.test_type")
@ -500,7 +521,7 @@ def testCustomTypeTypeCaster():
# CHECK: TestType(!python_test.test_type)
print(repr(b))
c = test.TestIntegerRankedTensorType.get([10, 10], 5)
c = TestIntegerRankedTensorType.get([10, 10], 5)
# CHECK: tensor<10x10xi5>
print(c)
# CHECK: TestIntegerRankedTensorType(tensor<10x10xi5>)
@ -511,7 +532,7 @@ def testCustomTypeTypeCaster():
@register_type_caster(c.typeid)
def type_caster(pytype):
return test.TestIntegerRankedTensorType(pytype)
return TestIntegerRankedTensorType(pytype)
except RuntimeError as e:
print(e)
@ -530,7 +551,7 @@ def testCustomTypeTypeCaster():
@register_type_caster(c.typeid, replace=True)
def type_caster(pytype):
return test.TestIntegerRankedTensorType(pytype)
return TestIntegerRankedTensorType(pytype)
d = tensor.EmptyOp([10, 10], IntegerType.get_signless(5)).result
# CHECK: tensor<10x10xi5>

3
mlir/test/python/lib/CMakeLists.txt
View File

@ -1,7 +1,8 @@
set(LLVM_OPTIONAL_SOURCES
PythonTestCAPI.cpp
PythonTestDialect.cpp
PythonTestModule.cpp
PythonTestModulePybind11.cpp
PythonTestModuleNanobind.cpp
)
add_mlir_library(MLIRPythonTestDialect

121
mlir/test/python/lib/PythonTestModuleNanobind.cpp Normal file
View File

@ -0,0 +1,121 @@
//===- PythonTestModuleNanobind.cpp - PythonTest dialect extension --------===//
//
// 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 nanobind edition of the PythonTest dialect module.
//===----------------------------------------------------------------------===//
#include <nanobind/nanobind.h>
#include <nanobind/stl/vector.h>
#include "PythonTestCAPI.h"
#include "mlir-c/BuiltinAttributes.h"
#include "mlir-c/BuiltinTypes.h"
#include "mlir-c/IR.h"
#include "mlir/Bindings/Python/NanobindAdaptors.h"
namespace nb = nanobind;
using namespace mlir::python::nanobind_adaptors;
static bool mlirTypeIsARankedIntegerTensor(MlirType t) {
return mlirTypeIsARankedTensor(t) &&
mlirTypeIsAInteger(mlirShapedTypeGetElementType(t));
}
NB_MODULE(_mlirPythonTestNanobind, m) {
m.def(
"register_python_test_dialect",
[](MlirContext context, bool load) {
MlirDialectHandle pythonTestDialect =
mlirGetDialectHandle__python_test__();
mlirDialectHandleRegisterDialect(pythonTestDialect, context);
if (load) {
mlirDialectHandleLoadDialect(pythonTestDialect, context);
}
},
nb::arg("context"), nb::arg("load") = true);
m.def(
"register_dialect",
[](MlirDialectRegistry registry) {
MlirDialectHandle pythonTestDialect =
mlirGetDialectHandle__python_test__();
mlirDialectHandleInsertDialect(pythonTestDialect, registry);
},
nb::arg("registry"));
mlir_attribute_subclass(m, "TestAttr",
mlirAttributeIsAPythonTestTestAttribute,
mlirPythonTestTestAttributeGetTypeID)
.def_classmethod(
"get",
[](const nb::object &cls, MlirContext ctx) {
return cls(mlirPythonTestTestAttributeGet(ctx));
},
nb::arg("cls"), nb::arg("context").none() = nb::none());
mlir_type_subclass(m, "TestType", mlirTypeIsAPythonTestTestType,
mlirPythonTestTestTypeGetTypeID)
.def_classmethod(
"get",
[](const nb::object &cls, MlirContext ctx) {
return cls(mlirPythonTestTestTypeGet(ctx));
},
nb::arg("cls"), nb::arg("context").none() = nb::none());
auto typeCls =
mlir_type_subclass(m, "TestIntegerRankedTensorType",
mlirTypeIsARankedIntegerTensor,
nb::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
.attr("RankedTensorType"))
.def_classmethod(
"get",
[](const nb::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));
},
nb::arg("cls"), nb::arg("shape"), nb::arg("width"),
nb::arg("context").none() = nb::none());
assert(nb::hasattr(typeCls.get_class(), "static_typeid") &&
"TestIntegerRankedTensorType has no static_typeid");
MlirTypeID mlirRankedTensorTypeID = mlirRankedTensorTypeGetTypeID();
nb::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
.attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
mlirRankedTensorTypeID, nb::arg("replace") = true)(
nanobind::cpp_function([typeCls](const nb::object &mlirType) {
return typeCls.get_class()(mlirType);
}));
auto valueCls = mlir_value_subclass(m, "TestTensorValue",
mlirTypeIsAPythonTestTestTensorValue)
.def("is_null", [](MlirValue &self) {
return mlirValueIsNull(self);
});
nb::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
.attr(MLIR_PYTHON_CAPI_VALUE_CASTER_REGISTER_ATTR)(
mlirRankedTensorTypeID)(
nanobind::cpp_function([valueCls](const nb::object &valueObj) {
nb::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;
}));
}

4
mlir/test/python/lib/PythonTestModule.cpp → mlir/test/python/lib/PythonTestModulePybind11.cpp
View File

@ -5,6 +5,8 @@
// 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"
@ -21,7 +23,7 @@ static bool mlirTypeIsARankedIntegerTensor(MlirType t) {
mlirTypeIsAInteger(mlirShapedTypeGetElementType(t));
}
PYBIND11_MODULE(_mlirPythonTest, m) {
PYBIND11_MODULE(_mlirPythonTestPybind11, m) {
m.def(
"register_python_test_dialect",
[](MlirContext context, bool load) {

18
utils/bazel/WORKSPACE
View File

@ -148,6 +148,24 @@ maybe(
url = "https://github.com/pybind/pybind11/archive/v2.10.3.zip",
)
maybe(
http_archive,
name = "robin_map",
strip_prefix = "robin-map-1.3.0",
sha256 = "a8424ad3b0affd4c57ed26f0f3d8a29604f0e1f2ef2089f497f614b1c94c7236",
build_file = "@llvm-raw//utils/bazel/third_party_build:robin_map.BUILD",
url = "https://github.com/Tessil/robin-map/archive/refs/tags/v1.3.0.tar.gz",
)
maybe(
http_archive,
name = "nanobind",
build_file = "@llvm-raw//utils/bazel/third_party_build:nanobind.BUILD",
sha256 = "bfbfc7e5759f1669e4ddb48752b1ddc5647d1430e94614d6f8626df1d508e65a",
strip_prefix = "nanobind-2.2.0",
url = "https://github.com/wjakob/nanobind/archive/refs/tags/v2.2.0.tar.gz",
)
load("@rules_python//python:repositories.bzl", "py_repositories", "python_register_toolchains")
py_repositories()

50
utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
View File

@ -932,7 +932,6 @@ exports_files(
filegroup(
name = "MLIRBindingsPythonHeaderFiles",
srcs = glob([
"lib/Bindings/Python/*.h",
"include/mlir-c/Bindings/Python/*.h",
"include/mlir/Bindings/Python/*.h",
]),
@ -942,12 +941,10 @@ cc_library(
name = "MLIRBindingsPythonHeaders",
includes = [
"include",
"lib/Bindings/Python",
],
textual_hdrs = [":MLIRBindingsPythonHeaderFiles"],
deps = [
":CAPIIRHeaders",
":CAPITransformsHeaders",
"@pybind11",
"@rules_python//python/cc:current_py_cc_headers",
],
@ -957,17 +954,41 @@ cc_library(
name = "MLIRBindingsPythonHeadersAndDeps",
includes = [
"include",
"lib/Bindings/Python",
],
textual_hdrs = [":MLIRBindingsPythonHeaderFiles"],
deps = [
":CAPIIR",
":CAPITransforms",
"@pybind11",
"@rules_python//python/cc:current_py_cc_headers",
],
)
cc_library(
name = "MLIRBindingsPythonNanobindHeaders",
includes = [
"include",
],
textual_hdrs = [":MLIRBindingsPythonHeaderFiles"],
deps = [
":CAPIIRHeaders",
"@nanobind",
"@rules_python//python/cc:current_py_cc_headers",
],
)
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.
PYBIND11_COPTS = [
"-fexceptions",
@ -993,16 +1014,25 @@ filegroup(
],
)
filegroup(
name = "MLIRBindingsPythonCoreHeaders",
srcs = glob([
"lib/Bindings/Python/*.h",
]),
)
cc_library(
name = "MLIRBindingsPythonCore",
srcs = [":MLIRBindingsPythonSourceFiles"],
copts = PYBIND11_COPTS,
features = PYBIND11_FEATURES,
textual_hdrs = [":MLIRBindingsPythonCoreHeaders"],
deps = [
":CAPIAsync",
":CAPIDebug",
":CAPIIR",
":CAPIInterfaces",
":CAPITransforms",
":MLIRBindingsPythonHeadersAndDeps",
":Support",
":config",
@ -1017,10 +1047,12 @@ cc_library(
srcs = [":MLIRBindingsPythonSourceFiles"],
copts = PYBIND11_COPTS,
features = PYBIND11_FEATURES,
textual_hdrs = [":MLIRBindingsPythonCoreHeaders"],
deps = [
":CAPIAsyncHeaders",
":CAPIDebugHeaders",
":CAPIIRHeaders",
":CAPITransformsHeaders",
":MLIRBindingsPythonHeaders",
":Support",
":config",
@ -1050,6 +1082,9 @@ cc_binary(
# These flags are needed for pybind11 to work.
copts = PYBIND11_COPTS,
features = PYBIND11_FEATURES,
includes = [
"lib/Bindings/Python",
],
linkshared = 1,
linkstatic = 0,
deps = [
@ -1063,6 +1098,9 @@ cc_binary(
srcs = ["lib/Bindings/Python/DialectLinalg.cpp"],
copts = PYBIND11_COPTS,
features = PYBIND11_FEATURES,
includes = [
"lib/Bindings/Python",
],
linkshared = 1,
linkstatic = 0,
deps = [
@ -8448,9 +8486,9 @@ cc_library(
hdrs = ["include/mlir/Conversion/ConvertToLLVM/ToLLVMPass.h"],
includes = ["include"],
deps = [
":Analysis",
":ConversionPassIncGen",
":ConvertToLLVMInterface",
":Analysis",
":IR",
":LLVMCommonConversion",
":LLVMDialect",

25
utils/bazel/third_party_build/nanobind.BUILD Normal file
View File

@ -0,0 +1,25 @@
cc_library(
name = "nanobind",
srcs = glob(
[
"src/*.cpp",
],
exclude = ["src/nb_combined.cpp"],
),
defines = [
"NB_BUILD=1",
"NB_SHARED=1",
],
includes = ["include"],
textual_hdrs = glob(
[
"include/**/*.h",
"src/*.h",
],
),
visibility = ["//visibility:public"],
deps = [
"@robin_map",
"@rules_python//python/cc:current_py_cc_headers",
],
)

12
utils/bazel/third_party_build/robin_map.BUILD Normal file
View File

@ -0,0 +1,12 @@
cc_library(
name = "robin_map",
hdrs = [
"include/tsl/robin_growth_policy.h",
"include/tsl/robin_hash.h",
"include/tsl/robin_map.h",
"include/tsl/robin_set.h",
],
includes = ["."],
strip_include_prefix = "include",
visibility = ["//visibility:public"],
)