[Mlir-commits] [mlir] [MLIR][Python] Extend bindings for external projects without duplication (PR #173241)
Sylvain Noiry
llvmlistbot at llvm.org
Mon Dec 22 03:21:02 PST 2025
https://github.com/ElectrikSpace created https://github.com/llvm/llvm-project/pull/173241
This patch allow to creation of Python binding extensions for external projects that still rely on the main MLIR package, but without duplication of code. It's now possible to mix imports from mlir and one or multiple external projects.
It contains two points:
1) Add a `MLIR_PYTHON_PACKAGE` option to `declare_mlir_dialect_python_bindings` and `declare_mlir_dialect_extension_python_bindings` to create bindings as and extension to a main (usually upstream) mlir package. This changes the import paths in Python files generated by `mlir-tblgen`. Previous relative paths required to duplicate the installation of some files.
2) MLIRPythonCAPI is now installed to be linked dynamically in add_mlir_python_modules. For an external project, the relationship looks like:
```
*cpython*.so --------> libMLIRPythonCAPI.so
|------> libMyExtensionsPythonCAPI.so
```
This prevents mangling issues due to code duplication. However it is orthogonal and does not resolve #108253.
I didn't change the Standalone example because this approach would introduce a dependency to the mlir package.
This is my first PR in llvm. Comments are welcomed!
>From fff6f57b470756ce585bd79e156daec53ad7306e Mon Sep 17 00:00:00 2001
From: Sylvain Noiry <snoiry at kalrayinc.com>
Date: Mon, 22 Dec 2025 10:25:52 +0100
Subject: [PATCH] [MLIR][Python] Extend bindings without duplication
Add a MLIR_PYTHON_PACKAGE option to declare_mlir_dialect_python_bindings
and declare_mlir_dialect_extension_python_bindings to create bindings
as and extension to a main (usually upstream) mlir package. This changes
the import paths in Python files generated by mlir-tblgen.
MLIRPythonCAPI is now installed and can be linked dynamically in
add_mlir_python_modules to prevent mangling issues due to code
duplication.
---
mlir/cmake/modules/AddMLIRPython.cmake | 25 ++++++++++++++-----
.../examples/standalone/python/CMakeLists.txt | 3 ++-
.../test/mlir-tblgen/enums-python-bindings.td | 4 +--
.../mlir-tblgen/EnumPythonBindingGen.cpp | 10 +++++---
mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp | 23 ++++++++++++-----
5 files changed, 47 insertions(+), 18 deletions(-)
diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index ca90151e76268..2d86a9a9d0bc4 100644
--- a/mlir/cmake/modules/AddMLIRPython.cmake
+++ b/mlir/cmake/modules/AddMLIRPython.cmake
@@ -389,6 +389,8 @@ endfunction()
# This file is where the *EnumAttrs are defined, not where the *Enums are defined.
# **WARNING**: This arg will shortly be removed when the just-below TODO is satisfied. Use at your
# risk.
+# MLIR_PYTHON_PACKAGE: Optional name of the current main MLIR package.
+# It can be used to build extensions against a main package.
#
# TODO: Right now `TD_FILE` can't be the actual dialect tablegen file, since we
# use its path to determine where to place the generated python file. If
@@ -397,7 +399,7 @@ endfunction()
function(declare_mlir_dialect_python_bindings)
cmake_parse_arguments(ARG
"GEN_ENUM_BINDINGS"
- "ROOT_DIR;ADD_TO_PARENT;TD_FILE;DIALECT_NAME"
+ "ROOT_DIR;ADD_TO_PARENT;TD_FILE;DIALECT_NAME;MLIR_PYTHON_PACKAGE"
"SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE"
${ARGN})
# Sources.
@@ -417,8 +419,12 @@ function(declare_mlir_dialect_python_bindings)
file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/${relative_td_directory}")
set(dialect_filename "${relative_td_directory}/_${ARG_DIALECT_NAME}_ops_gen.py")
set(LLVM_TARGET_DEFINITIONS ${td_file})
+ if(NOT DEFINED ARG_MLIR_PYTHON_PACKAGE)
+ set(ARG_MLIR_PYTHON_PACKAGE "mlir")
+ endif()
mlir_tablegen("${dialect_filename}"
-gen-python-op-bindings -bind-dialect=${ARG_DIALECT_NAME}
+ -python-package-prefix=${ARG_MLIR_PYTHON_PACKAGE}
DEPENDS ${ARG_DEPENDS}
)
add_public_tablegen_target(${tblgen_target})
@@ -430,7 +436,8 @@ function(declare_mlir_dialect_python_bindings)
set(LLVM_TARGET_DEFINITIONS ${td_file})
endif()
set(enum_filename "${relative_td_directory}/_${ARG_DIALECT_NAME}_enum_gen.py")
- mlir_tablegen(${enum_filename} -gen-python-enum-bindings)
+ mlir_tablegen(${enum_filename} -gen-python-enum-bindings
+ -python-package-prefix=${ARG_MLIR_PYTHON_PACKAGE})
list(APPEND _sources ${enum_filename})
endif()
@@ -464,10 +471,12 @@ endfunction()
# This file is where the *Attrs are defined, not where the *Enums are defined.
# **WARNING**: This arg will shortly be removed when the TODO for
# declare_mlir_dialect_python_bindings is satisfied. Use at your risk.
+# MLIR_PYTHON_PACKAGE: Optional name of the current main MLIR package.
+# It can be used to build extensions against a main package.
function(declare_mlir_dialect_extension_python_bindings)
cmake_parse_arguments(ARG
"GEN_ENUM_BINDINGS"
- "ROOT_DIR;ADD_TO_PARENT;TD_FILE;DIALECT_NAME;EXTENSION_NAME"
+ "ROOT_DIR;ADD_TO_PARENT;TD_FILE;DIALECT_NAME;EXTENSION_NAME;MLIR_PYTHON_PACKAGE"
"SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE"
${ARGN})
# Source files.
@@ -487,9 +496,13 @@ function(declare_mlir_dialect_extension_python_bindings)
file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/${relative_td_directory}")
set(output_filename "${relative_td_directory}/_${ARG_EXTENSION_NAME}_ops_gen.py")
set(LLVM_TARGET_DEFINITIONS ${td_file})
+ if(NOT DEFINED ARG_MLIR_PYTHON_PACKAGE)
+ set(ARG_MLIR_PYTHON_PACKAGE "mlir")
+ endif()
mlir_tablegen("${output_filename}" -gen-python-op-bindings
-bind-dialect=${ARG_DIALECT_NAME}
- -dialect-extension=${ARG_EXTENSION_NAME})
+ -dialect-extension=${ARG_EXTENSION_NAME}
+ -python-package-prefix=${ARG_MLIR_PYTHON_PACKAGE})
add_public_tablegen_target(${tblgen_target})
if(ARG_DEPENDS)
add_dependencies(${tblgen_target} ${ARG_DEPENDS})
@@ -502,7 +515,8 @@ function(declare_mlir_dialect_extension_python_bindings)
set(LLVM_TARGET_DEFINITIONS ${td_file})
endif()
set(enum_filename "${relative_td_directory}/_${ARG_EXTENSION_NAME}_enum_gen.py")
- mlir_tablegen(${enum_filename} -gen-python-enum-bindings)
+ mlir_tablegen(${enum_filename} -gen-python-enum-bindings
+ -python-package-prefix=${ARG_MLIR_PYTHON_PACKAGE})
list(APPEND _sources ${enum_filename})
endif()
@@ -601,7 +615,6 @@ function(add_mlir_python_common_capi_library name)
# Generate the aggregate .so that everything depends on.
add_mlir_aggregate(${name}
SHARED
- DISABLE_INSTALL
EMBED_LIBS ${_embed_libs}
)
diff --git a/mlir/examples/standalone/python/CMakeLists.txt b/mlir/examples/standalone/python/CMakeLists.txt
index edaedf18cc843..91373da0e9377 100644
--- a/mlir/examples/standalone/python/CMakeLists.txt
+++ b/mlir/examples/standalone/python/CMakeLists.txt
@@ -18,7 +18,8 @@ declare_mlir_dialect_python_bindings(
SOURCES
dialects/standalone_nanobind.py
_mlir_libs/_standaloneDialectsNanobind/py.typed
- DIALECT_NAME standalone)
+ DIALECT_NAME standalone
+ MLIR_PYTHON_PACKAGE "${MLIR_PYTHON_PACKAGE_PREFIX}")
declare_mlir_python_extension(StandalonePythonSources.NanobindExtension
MODULE_NAME _standaloneDialectsNanobind
diff --git a/mlir/test/mlir-tblgen/enums-python-bindings.td b/mlir/test/mlir-tblgen/enums-python-bindings.td
index cd23b6a2effb9..5dd002ca21bd3 100644
--- a/mlir/test/mlir-tblgen/enums-python-bindings.td
+++ b/mlir/test/mlir-tblgen/enums-python-bindings.td
@@ -10,8 +10,8 @@ def Test_Dialect : Dialect {
// CHECK: Autogenerated by mlir-tblgen; don't manually edit.
// CHECK: from enum import IntEnum, auto, IntFlag
-// CHECK: from ._ods_common import _cext as _ods_cext
-// CHECK: from ..ir import register_attribute_builder
+// CHECK: from mlir.dialects._ods_common import _cext as _ods_cext
+// CHECK: from mlir.ir import register_attribute_builder
// CHECK: _ods_ir = _ods_cext.ir
def One : I32EnumAttrCase<"CaseOne", 1, "one">;
diff --git a/mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp b/mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp
index acc9b61d7121c..31211a2094f06 100644
--- a/mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp
+++ b/mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp
@@ -17,6 +17,7 @@
#include "mlir/TableGen/Dialect.h"
#include "mlir/TableGen/EnumInfo.h"
#include "mlir/TableGen/GenInfo.h"
+#include "llvm/Support/CommandLine.h"
#include "llvm/Support/FormatVariadic.h"
#include "llvm/TableGen/Record.h"
@@ -27,16 +28,19 @@ using llvm::Record;
using llvm::RecordKeeper;
/// File header and includes.
+/// {0} is the Python package prefix.
constexpr const char *fileHeader = R"Py(
# Autogenerated by mlir-tblgen; don't manually edit.
from enum import IntEnum, auto, IntFlag
-from ._ods_common import _cext as _ods_cext
-from ..ir import register_attribute_builder
+from {0}.dialects._ods_common import _cext as _ods_cext
+from {0}.ir import register_attribute_builder
_ods_ir = _ods_cext.ir
)Py";
+extern llvm::cl::opt<std::string> clPythonPackagePrefix;
+
/// Makes enum case name Python-compatible, i.e. UPPER_SNAKE_CASE.
static std::string makePythonEnumCaseName(StringRef name) {
if (isPythonReserved(name.str()))
@@ -122,7 +126,7 @@ static bool emitDialectEnumAttributeBuilder(StringRef attrDefName,
/// Emits Python bindings for all enums in the record keeper. Returns
/// `false` on success, `true` on failure.
static bool emitPythonEnums(const RecordKeeper &records, raw_ostream &os) {
- os << fileHeader;
+ os << formatv(fileHeader, clPythonPackagePrefix);
for (const Record *it :
records.getAllDerivedDefinitionsIfDefined("EnumInfo")) {
EnumInfo enumInfo(*it);
diff --git a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
index 2c33f4efac3ac..fd3a82427352d 100644
--- a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
@@ -30,12 +30,12 @@ using llvm::Record;
using llvm::RecordKeeper;
/// File header and includes.
-/// {0} is the dialect namespace.
+/// {0} is the Python package prefix.
constexpr const char *fileHeader = R"Py(
# Autogenerated by mlir-tblgen; don't manually edit.
-from ._ods_common import _cext as _ods_cext
-from ._ods_common import (
+from {0}.dialects._ods_common import _cext as _ods_cext
+from {0}.dialects._ods_common import (
equally_sized_accessor as _ods_equally_sized_accessor,
get_default_loc_context as _ods_get_default_loc_context,
get_op_results_or_values as _get_op_results_or_values,
@@ -51,6 +51,7 @@ from typing import Sequence as _Sequence, Union as _Union, Optional as _Optional
/// Template for dialect class:
/// {0} is the dialect namespace.
+/// {1} is the Python package prefix.
constexpr const char *dialectClassTemplate = R"Py(
@_ods_cext.register_dialect
class _Dialect(_ods_ir.Dialect):
@@ -58,7 +59,7 @@ class _Dialect(_ods_ir.Dialect):
)Py";
constexpr const char *dialectExtensionTemplate = R"Py(
-from ._{0}_ops_gen import _Dialect
+from {1}.dialects._{0}_ops_gen import _Dialect
)Py";
/// Template for operation class:
@@ -293,6 +294,15 @@ def {0}({2}) -> _Union[_ods_ir.OpResult, _ods_ir.OpResultList, {1}]:
return results if len(results) > 1 else (results[0] if len(results) == 1 else op)
)Py";
+static llvm::cl::OptionCategory
+ clPythonBindingCat("Options for -gen-python-op-bindings and "
+ "-gen-python-enum-bindings");
+
+llvm::cl::opt<std::string>
+ clPythonPackagePrefix("python-package-prefix",
+ llvm::cl::desc("The prefix of the MLIR Python package"),
+ llvm::cl::init("mlir"), llvm::cl::cat(clPythonBindingCat));
+
static llvm::cl::OptionCategory
clOpPythonBindingCat("Options for -gen-python-op-bindings");
@@ -1222,9 +1232,10 @@ static bool emitAllOps(const RecordKeeper &records, raw_ostream &os) {
if (clDialectName.empty())
llvm::PrintFatalError("dialect name not provided");
- os << fileHeader;
+ os << formatv(fileHeader, clPythonPackagePrefix.getValue());
if (!clDialectExtensionName.empty())
- os << formatv(dialectExtensionTemplate, clDialectName.getValue());
+ os << formatv(dialectExtensionTemplate, clDialectName.getValue(),
+ clPythonPackagePrefix.getValue());
else
os << formatv(dialectClassTemplate, clDialectName.getValue());
More information about the Mlir-commits
mailing list