Skip to content

[mlir python] Add nanobind support for standalone dialects. #117922

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Dec 3, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 21 additions & 6 deletions mlir/cmake/modules/AddMLIRPython.cmake
Original file line number Diff line number Diff line change
Expand Up @@ -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})

Expand All @@ -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
Expand Down Expand Up @@ -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}
Expand Down Expand Up @@ -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)
Expand All @@ -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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was imagining this would be an arg (that naturally defaulted to mlir) so that downstream users could set this to whatever they wanted (if they wanted).

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's do that in a follow-on change?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure np

${ARG_SOURCES}
)
endif()

# The extension itself must be compiled with RTTI and exceptions enabled.
# Also, some warning classes triggered by pybind11 are disabled.
Expand Down
39 changes: 39 additions & 0 deletions mlir/cmake/modules/MLIRDetectPythonEnv.cmake
Original file line number Diff line number Diff line change
Expand Up @@ -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}")
Expand All @@ -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()

Expand Down Expand Up @@ -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 changes: 12 additions & 8 deletions mlir/docs/Bindings/Python.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 changes: 18 additions & 4 deletions mlir/examples/standalone/python/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -17,18 +17,32 @@ declare_mlir_dialect_python_bindings(
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir_standalone"
TD_FILE dialects/StandaloneOps.td
SOURCES
dialects/standalone.py
dialects/standalone_pybind11.py
dialects/standalone_nanobind.py
DIALECT_NAME standalone)

declare_mlir_python_extension(StandalonePythonSources.Extension
MODULE_NAME _standaloneDialects

declare_mlir_python_extension(StandalonePythonSources.Pybind11Extension
MODULE_NAME _standaloneDialectsPybind11
ADD_TO_PARENT StandalonePythonSources
SOURCES
StandaloneExtensionPybind11.cpp
EMBED_CAPI_LINK_LIBS
StandaloneCAPI
PYTHON_BINDINGS_LIBRARY pybind11
)

declare_mlir_python_extension(StandalonePythonSources.NanobindExtension
MODULE_NAME _standaloneDialectsNanobind
ADD_TO_PARENT StandalonePythonSources
SOURCES
StandaloneExtension.cpp
StandaloneExtensionNanobind.cpp
EMBED_CAPI_LINK_LIBS
StandaloneCAPI
PYTHON_BINDINGS_LIBRARY nanobind
)


################################################################################
# Common CAPI
################################################################################
Expand Down
35 changes: 35 additions & 0 deletions mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp
Original file line number Diff line number Diff line change
@@ -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);
}
Original file line number Diff line number Diff line change
@@ -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.
Expand All @@ -11,7 +14,7 @@

using namespace mlir::python::adaptors;

PYBIND11_MODULE(_standaloneDialects, m) {
PYBIND11_MODULE(_standaloneDialectsPybind11, m) {
//===--------------------------------------------------------------------===//
// standalone dialect
//===--------------------------------------------------------------------===//
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 *
Original file line number Diff line number Diff line change
@@ -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 changes: 12 additions & 2 deletions mlir/examples/standalone/test/python/smoketest.py
Original file line number Diff line number Diff line change
@@ -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()
Expand Down
59 changes: 59 additions & 0 deletions mlir/include/mlir/Bindings/Python/Diagnostics.h
Original file line number Diff line number Diff line change
@@ -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
Loading
Loading