[Mlir-commits] [mlir] [MLIR][Python] enable precise registration (PR #160742)
Maksim Levental
llvmlistbot at llvm.org
Fri Sep 26 22:31:46 PDT 2025
https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/160742
>From 17d76f802f323ce0b909064580921021724e007d Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Thu, 25 Sep 2025 09:40:53 -0700
Subject: [PATCH 1/2] [MLIR][Python] enable precise registration
---
mlir/cmake/modules/AddMLIRPython.cmake | 14 +++++++++++---
mlir/python/CMakeLists.txt | 24 +++++++++++++++++++++---
mlir/python/mlir/_mlir_libs/_capi.py.in | 8 ++++++++
mlir/test/python/ir/capi.py | 6 ++++++
4 files changed, 46 insertions(+), 6 deletions(-)
create mode 100644 mlir/python/mlir/_mlir_libs/_capi.py.in
create mode 100644 mlir/test/python/ir/capi.py
diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index 208cbdd1dd535..d8b6d493f985c 100644
--- a/mlir/cmake/modules/AddMLIRPython.cmake
+++ b/mlir/cmake/modules/AddMLIRPython.cmake
@@ -23,11 +23,14 @@
# grouping. Source groupings form a DAG.
# SOURCES: List of specific source files relative to ROOT_DIR to include.
# SOURCES_GLOB: List of glob patterns relative to ROOT_DIR to include.
+# 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.
function(declare_mlir_python_sources name)
cmake_parse_arguments(ARG
""
"ROOT_DIR;ADD_TO_PARENT"
- "SOURCES;SOURCES_GLOB"
+ "SOURCES;SOURCES_GLOB;EMBED_CAPI_LINK_LIBS"
${ARGN})
if(NOT ARG_ROOT_DIR)
@@ -53,9 +56,10 @@ function(declare_mlir_python_sources name)
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_DEPENDS"
+ EXPORT_PROPERTIES "mlir_python_SOURCES_TYPE;mlir_python_DEPENDS;mlir_python_EMBED_CAPI_LINK_LIBS"
mlir_python_SOURCES_TYPE pure
mlir_python_DEPENDS ""
+ mlir_python_EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
)
# Use the interface include directories and sources on the target to carry the
@@ -374,6 +378,9 @@ 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.
+# 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.
#
# 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
@@ -383,7 +390,7 @@ function(declare_mlir_dialect_python_bindings)
cmake_parse_arguments(ARG
"GEN_ENUM_BINDINGS"
"ROOT_DIR;ADD_TO_PARENT;TD_FILE;DIALECT_NAME"
- "SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE"
+ "SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE;EMBED_CAPI_LINK_LIBS"
${ARGN})
# Sources.
set(_dialect_target "${ARG_ADD_TO_PARENT}.${ARG_DIALECT_NAME}")
@@ -424,6 +431,7 @@ function(declare_mlir_dialect_python_bindings)
ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
ADD_TO_PARENT "${_dialect_target}"
SOURCES ${_sources}
+ EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
)
endif()
endfunction()
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 9f5246de6bda0..368f344841ed1 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -326,7 +326,10 @@ declare_mlir_dialect_python_bindings(
SOURCES
dialects/arith.py
DIALECT_NAME arith
- GEN_ENUM_BINDINGS)
+ GEN_ENUM_BINDINGS
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIArith
+ )
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -846,8 +849,20 @@ endif()
# once ready.
################################################################################
+set(MLIR_PYTHON_CAPI_DYLIB_NAME MLIRPythonCAPI)
+configure_file(
+ "${CMAKE_CURRENT_LIST_DIR}/mlir/_mlir_libs/_capi.py.in"
+ "${CMAKE_CURRENT_BINARY_DIR}/_mlir_libs/_capi.py"
+ @ONLY
+)
+declare_mlir_python_sources(
+ MLIRPythonCAPICTypesBinding
+ ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
+ SOURCES _mlir_libs/_capi.py
+)
+
set(MLIRPythonModules_ROOT_PREFIX "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}")
-add_mlir_python_common_capi_library(MLIRPythonCAPI
+add_mlir_python_common_capi_library(${MLIR_PYTHON_CAPI_DYLIB_NAME}
INSTALL_COMPONENT MLIRPythonModules
INSTALL_DESTINATION "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/_mlir_libs"
OUTPUT_DIRECTORY "${MLIRPythonModules_ROOT_PREFIX}/_mlir_libs"
@@ -963,7 +978,10 @@ endif()
# This must come last.
################################################################################
-set(_declared_sources MLIRPythonSources MLIRPythonExtension.RegisterEverything)
+set(_declared_sources
+ MLIRPythonSources
+ MLIRPythonExtension.RegisterEverything
+ MLIRPythonCAPICTypesBinding)
if(NOT CMAKE_CROSSCOMPILING)
list(APPEND _declared_sources MLIRPythonExtension.Core.type_stub_gen)
endif()
diff --git a/mlir/python/mlir/_mlir_libs/_capi.py.in b/mlir/python/mlir/_mlir_libs/_capi.py.in
new file mode 100644
index 0000000000000..9568845e67de9
--- /dev/null
+++ b/mlir/python/mlir/_mlir_libs/_capi.py.in
@@ -0,0 +1,8 @@
+# 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
+
+import ctypes
+from pathlib import Path
+
+_capi = ctypes.CDLL(str(Path(__file__).parent / "@CMAKE_SHARED_LIBRARY_PREFIX@@MLIR_PYTHON_CAPI_DYLIB_NAME@@CMAKE_SHARED_LIBRARY_SUFFIX@"))
\ No newline at end of file
diff --git a/mlir/test/python/ir/capi.py b/mlir/test/python/ir/capi.py
new file mode 100644
index 0000000000000..d60fbd820f91e
--- /dev/null
+++ b/mlir/test/python/ir/capi.py
@@ -0,0 +1,6 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir._mlir_libs._capi import _capi
+
+print("success")
+# CHECK: success
\ No newline at end of file
>From d3463b23ae9468931cb3b36dafefedbb054e07b5 Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Thu, 25 Sep 2025 14:59:58 -0700
Subject: [PATCH 2/2] [MLIR][Python] enable precise registration
---
mlir/cmake/modules/AddMLIRPython.cmake | 6 +-
mlir/include/mlir-c/Bindings/Python/Interop.h | 9 +
mlir/include/mlir-c/Dialect/Affine.h | 36 +++
mlir/include/mlir-c/Dialect/Bufferization.h | 37 +++
mlir/include/mlir-c/Dialect/Builtin.h | 33 +++
mlir/include/mlir-c/Dialect/Complex.h | 33 +++
mlir/include/mlir-c/Dialect/GPU.h | 3 +
mlir/include/mlir-c/Dialect/MemRef.h | 3 +
mlir/include/mlir-c/Dialect/SMT.h | 3 +
mlir/include/mlir-c/Dialect/SparseTensor.h | 3 +
mlir/include/mlir-c/Dialect/Tensor.h | 3 +
mlir/include/mlir-c/Dialect/Tosa.h | 33 +++
mlir/include/mlir-c/Dialect/Transform.h | 13 +
mlir/include/mlir-c/Dialect/UB.h | 33 +++
mlir/include/mlir-c/Dialect/Vector.h | 3 +
mlir/include/mlir-c/IR.h | 18 +-
mlir/lib/Bindings/Python/IRCore.cpp | 34 ++-
mlir/lib/CAPI/Dialect/Affine.cpp | 14 +
mlir/lib/CAPI/Dialect/Bufferization.cpp | 14 +
mlir/lib/CAPI/Dialect/Builtin.cpp | 13 +
mlir/lib/CAPI/Dialect/CMakeLists.txt | 120 ++++++++
mlir/lib/CAPI/Dialect/Complex.cpp | 14 +
mlir/lib/CAPI/Dialect/Tosa.cpp | 13 +
mlir/lib/CAPI/Dialect/Transform.cpp | 15 +
mlir/lib/CAPI/Dialect/TransformAffine.cpp | 16 ++
.../CAPI/Dialect/TransformBufferization.cpp | 16 ++
mlir/lib/CAPI/Dialect/TransformGPU.cpp | 15 +
mlir/lib/CAPI/Dialect/TransformMemRef.cpp | 15 +
mlir/lib/CAPI/Dialect/TransformSMT.cpp | 15 +
.../CAPI/Dialect/TransformSparseTensor.cpp | 16 ++
mlir/lib/CAPI/Dialect/TransformTensor.cpp | 15 +
mlir/lib/CAPI/Dialect/TransformVector.cpp | 15 +
mlir/lib/CAPI/Dialect/UB.cpp | 13 +
mlir/lib/CAPI/IR/IR.cpp | 12 +
mlir/python/CMakeLists.txt | 262 +++++++++++++-----
mlir/python/mlir/_mlir_libs/_capi.py.in | 45 ++-
mlir/python/mlir/dialects/QuantOps.td | 14 +
mlir/python/mlir/dialects/quant.py | 1 +
.../mlir/dialects/transform/bufferization.py | 9 +
mlir/python/mlir/dialects/transform/gpu.py | 10 +
mlir/python/mlir/dialects/transform/loop.py | 9 +
mlir/python/mlir/dialects/transform/memref.py | 9 +
mlir/python/mlir/dialects/transform/smt.py | 10 +
.../mlir/dialects/transform/sparse_tensor.py | 8 +
mlir/python/mlir/dialects/transform/tensor.py | 9 +
mlir/python/mlir/dialects/transform/tune.py | 9 +
mlir/python/mlir/dialects/transform/vector.py | 8 +
mlir/test/python/dialects/gpu/dialect.py | 1 +
mlir/test/python/dialects/irdl.py | 1 +
mlir/test/python/dialects/memref.py | 1 +
mlir/test/python/integration/dialects/pdl.py | 2 +-
mlir/test/python/ir/builtin_types.py | 2 +-
mlir/test/python/ir/capi.py | 39 ++-
mlir/test/python/ir/diagnostic_handler.py | 1 +
mlir/test/python/ir/dialects.py | 1 +
mlir/test/python/ir/exception.py | 1 +
mlir/test/python/ir/insertion_point.py | 1 +
mlir/test/python/ir/module.py | 1 +
mlir/test/python/ir/operation.py | 2 +-
mlir/test/python/ir/symbol_table.py | 1 +
mlir/test/python/multithreaded_tests.py | 2 +-
mlir/test/python/python_pass.py | 2 +-
mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp | 7 +-
63 files changed, 1035 insertions(+), 87 deletions(-)
create mode 100644 mlir/include/mlir-c/Dialect/Affine.h
create mode 100644 mlir/include/mlir-c/Dialect/Bufferization.h
create mode 100644 mlir/include/mlir-c/Dialect/Builtin.h
create mode 100644 mlir/include/mlir-c/Dialect/Complex.h
create mode 100644 mlir/include/mlir-c/Dialect/Tosa.h
create mode 100644 mlir/include/mlir-c/Dialect/UB.h
create mode 100644 mlir/lib/CAPI/Dialect/Affine.cpp
create mode 100644 mlir/lib/CAPI/Dialect/Bufferization.cpp
create mode 100644 mlir/lib/CAPI/Dialect/Builtin.cpp
create mode 100644 mlir/lib/CAPI/Dialect/Complex.cpp
create mode 100644 mlir/lib/CAPI/Dialect/Tosa.cpp
create mode 100644 mlir/lib/CAPI/Dialect/TransformAffine.cpp
create mode 100644 mlir/lib/CAPI/Dialect/TransformBufferization.cpp
create mode 100644 mlir/lib/CAPI/Dialect/TransformGPU.cpp
create mode 100644 mlir/lib/CAPI/Dialect/TransformMemRef.cpp
create mode 100644 mlir/lib/CAPI/Dialect/TransformSMT.cpp
create mode 100644 mlir/lib/CAPI/Dialect/TransformSparseTensor.cpp
create mode 100644 mlir/lib/CAPI/Dialect/TransformTensor.cpp
create mode 100644 mlir/lib/CAPI/Dialect/TransformVector.cpp
create mode 100644 mlir/lib/CAPI/Dialect/UB.cpp
create mode 100644 mlir/python/mlir/dialects/QuantOps.td
diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index d8b6d493f985c..ceebc931e60dc 100644
--- a/mlir/cmake/modules/AddMLIRPython.cmake
+++ b/mlir/cmake/modules/AddMLIRPython.cmake
@@ -457,11 +457,14 @@ 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.
+# 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.
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"
- "SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE"
+ "SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE;EMBED_CAPI_LINK_LIBS"
${ARGN})
# Source files.
set(_extension_target "${ARG_ADD_TO_PARENT}.${ARG_EXTENSION_NAME}")
@@ -503,6 +506,7 @@ function(declare_mlir_dialect_extension_python_bindings)
ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
ADD_TO_PARENT "${_extension_target}"
SOURCES ${_sources}
+ EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
)
endif()
endfunction()
diff --git a/mlir/include/mlir-c/Bindings/Python/Interop.h b/mlir/include/mlir-c/Bindings/Python/Interop.h
index a33190c380d37..89559da689017 100644
--- a/mlir/include/mlir-c/Bindings/Python/Interop.h
+++ b/mlir/include/mlir-c/Bindings/Python/Interop.h
@@ -84,6 +84,8 @@
#define MLIR_PYTHON_CAPSULE_VALUE MAKE_MLIR_PYTHON_QUALNAME("ir.Value._CAPIPtr")
#define MLIR_PYTHON_CAPSULE_TYPEID \
MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID._CAPIPtr")
+#define MLIR_PYTHON_CAPSULE_DIALECT_HANDLE \
+ MAKE_MLIR_PYTHON_QUALNAME("ir.DialectHandle._CAPIPtr")
/** Attribute on MLIR Python objects that expose their C-API pointer.
* This will be a type-specific capsule created as per one of the helpers
@@ -457,6 +459,13 @@ static inline MlirValue mlirPythonCapsuleToValue(PyObject *capsule) {
return value;
}
+static inline MlirDialectHandle
+mlirPythonCapsuleToDialectHandle(PyObject *capsule) {
+ void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_DIALECT_HANDLE);
+ MlirDialectHandle handle = {ptr};
+ return handle;
+}
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir-c/Dialect/Affine.h b/mlir/include/mlir-c/Dialect/Affine.h
new file mode 100644
index 0000000000000..b2bf5aad44de9
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/Affine.h
@@ -0,0 +1,36 @@
+//===-- mlir-c/Dialect/Affine.h - C API for Affine dialect --------*- C -*-===//
+//
+// 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 header declares the C interface for registering and accessing the
+// Affine dialect. A dialect should be registered with a context to make it
+// available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_AFFINE_H
+#define MLIR_C_DIALECT_AFFINE_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Affine, affine);
+
+MLIR_CAPI_EXPORTED void
+mlirAffineRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_AFFINE_H
diff --git a/mlir/include/mlir-c/Dialect/Bufferization.h b/mlir/include/mlir-c/Dialect/Bufferization.h
new file mode 100644
index 0000000000000..dee4abccda4f3
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/Bufferization.h
@@ -0,0 +1,37 @@
+//===-- mlir-c/Dialect/Bufferization.h - C API for Bufferization dialect
+//----------*- C -*-===//
+//
+// 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 header declares the C interface for registering and accessing the
+// Bufferization dialect. A dialect should be registered with a context to make
+// it available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_BUFFERIZATION_H
+#define MLIR_C_DIALECT_BUFFERIZATION_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Bufferization, bufferization);
+
+MLIR_CAPI_EXPORTED void mlirBufferizationRegisterTransformDialectExtension(
+ MlirDialectRegistry registry);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_BUFFERIZATION_H
diff --git a/mlir/include/mlir-c/Dialect/Builtin.h b/mlir/include/mlir-c/Dialect/Builtin.h
new file mode 100644
index 0000000000000..c5d958249b36f
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/Builtin.h
@@ -0,0 +1,33 @@
+//===-- mlir-c/Dialect/Builtin.h - C API for Builtin dialect ------*- C -*-===//
+//
+// 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 header declares the C interface for registering and accessing the
+// Builtin dialect. A dialect should be registered with a context to make it
+// available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_BUILTIN_H
+#define MLIR_C_DIALECT_BUILTIN_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Builtin, builtin);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_BUILTIN_H
diff --git a/mlir/include/mlir-c/Dialect/Complex.h b/mlir/include/mlir-c/Dialect/Complex.h
new file mode 100644
index 0000000000000..e51a67346a6ee
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/Complex.h
@@ -0,0 +1,33 @@
+//===-- mlir-c/Dialect/Complex.h - C API for Complex dialect ------*- C -*-===//
+//
+// 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 header declares the C interface for registering and accessing the
+// Complex dialect. A dialect should be registered with a context to make it
+// available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_COMPLEX_H
+#define MLIR_C_DIALECT_COMPLEX_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Complex, complex);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_COMPLEX_H
diff --git a/mlir/include/mlir-c/Dialect/GPU.h b/mlir/include/mlir-c/Dialect/GPU.h
index 321c1122c3370..44d6308c270b2 100644
--- a/mlir/include/mlir-c/Dialect/GPU.h
+++ b/mlir/include/mlir-c/Dialect/GPU.h
@@ -63,6 +63,9 @@ mlirGPUObjectAttrHasKernels(MlirAttribute mlirObjectAttr);
MLIR_CAPI_EXPORTED MlirAttribute
mlirGPUObjectAttrGetKernels(MlirAttribute mlirObjectAttr);
+MLIR_CAPI_EXPORTED void
+mlirGPURegisterTransformDialectExtension(MlirDialectRegistry registry);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir-c/Dialect/MemRef.h b/mlir/include/mlir-c/Dialect/MemRef.h
index 087a4b3f85b1d..38184404d6853 100644
--- a/mlir/include/mlir-c/Dialect/MemRef.h
+++ b/mlir/include/mlir-c/Dialect/MemRef.h
@@ -26,6 +26,9 @@ extern "C" {
MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(MemRef, memref);
+MLIR_CAPI_EXPORTED void
+mlirMemRefRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir-c/Dialect/SMT.h b/mlir/include/mlir-c/Dialect/SMT.h
index 0ad64746f148b..e6d8e14efed6e 100644
--- a/mlir/include/mlir-c/Dialect/SMT.h
+++ b/mlir/include/mlir-c/Dialect/SMT.h
@@ -104,6 +104,9 @@ mlirSMTAttrGetBVCmpPredicate(MlirContext ctx, MlirStringRef str);
MLIR_CAPI_EXPORTED MlirAttribute mlirSMTAttrGetIntPredicate(MlirContext ctx,
MlirStringRef str);
+MLIR_CAPI_EXPORTED void
+mlirSMTRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir-c/Dialect/SparseTensor.h b/mlir/include/mlir-c/Dialect/SparseTensor.h
index c816c1b58690e..0701211375a48 100644
--- a/mlir/include/mlir-c/Dialect/SparseTensor.h
+++ b/mlir/include/mlir-c/Dialect/SparseTensor.h
@@ -107,6 +107,9 @@ mlirSparseTensorEncodingAttrBuildLvlType(
const enum MlirSparseTensorLevelPropertyNondefault *properties,
unsigned propSize, unsigned n, unsigned m);
+MLIR_CAPI_EXPORTED void
+mlirSparseTensorRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir-c/Dialect/Tensor.h b/mlir/include/mlir-c/Dialect/Tensor.h
index 74cbc5a6f0fe1..5891ade3bf3fd 100644
--- a/mlir/include/mlir-c/Dialect/Tensor.h
+++ b/mlir/include/mlir-c/Dialect/Tensor.h
@@ -18,6 +18,9 @@ extern "C" {
MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Tensor, tensor);
+MLIR_CAPI_EXPORTED void
+mlirTensorRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir-c/Dialect/Tosa.h b/mlir/include/mlir-c/Dialect/Tosa.h
new file mode 100644
index 0000000000000..3dc1d2a042fb2
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/Tosa.h
@@ -0,0 +1,33 @@
+//===-- mlir-c/Dialect/Tosa.h - C API for Tosa dialect ----------*- C -*-===//
+//
+// 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 header declares the C interface for registering and accessing the
+// Tosa dialect. A dialect should be registered with a context to make it
+// available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_TOSA_H
+#define MLIR_C_DIALECT_TOSA_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Tosa, tosa);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_TOSA_H
diff --git a/mlir/include/mlir-c/Dialect/Transform.h b/mlir/include/mlir-c/Dialect/Transform.h
index 02c99b5921882..1b5232c106edd 100644
--- a/mlir/include/mlir-c/Dialect/Transform.h
+++ b/mlir/include/mlir-c/Dialect/Transform.h
@@ -76,6 +76,19 @@ MLIR_CAPI_EXPORTED MlirType mlirTransformParamTypeGet(MlirContext ctx,
MLIR_CAPI_EXPORTED MlirType mlirTransformParamTypeGetType(MlirType type);
+//===---------------------------------------------------------------------===//
+// Extension Registration
+//===---------------------------------------------------------------------===//
+
+MLIR_CAPI_EXPORTED void
+mlirDebugRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
+MLIR_CAPI_EXPORTED void
+mlirLoopRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
+MLIR_CAPI_EXPORTED void
+mlirTuneRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir-c/Dialect/UB.h b/mlir/include/mlir-c/Dialect/UB.h
new file mode 100644
index 0000000000000..74159f0c705de
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/UB.h
@@ -0,0 +1,33 @@
+//===-- mlir-c/Dialect/UB.h - C API for UB dialect ----------------*- C -*-===//
+//
+// 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 header declares the C interface for registering and accessing the
+// UB dialect. A dialect should be registered with a context to make it
+// available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_UB_H
+#define MLIR_C_DIALECT_UB_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(UB, ub);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_UB_H
diff --git a/mlir/include/mlir-c/Dialect/Vector.h b/mlir/include/mlir-c/Dialect/Vector.h
index 6256c82d172ce..b7248b7985057 100644
--- a/mlir/include/mlir-c/Dialect/Vector.h
+++ b/mlir/include/mlir-c/Dialect/Vector.h
@@ -26,6 +26,9 @@ extern "C" {
MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Vector, vector);
+MLIR_CAPI_EXPORTED void
+mlirVectorRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir-c/IR.h b/mlir/include/mlir-c/IR.h
index 061d7620ba077..55cc86accb8a0 100644
--- a/mlir/include/mlir-c/IR.h
+++ b/mlir/include/mlir-c/IR.h
@@ -66,6 +66,7 @@ DEFINE_C_API_STRUCT(MlirLocation, const void);
DEFINE_C_API_STRUCT(MlirModule, const void);
DEFINE_C_API_STRUCT(MlirType, const void);
DEFINE_C_API_STRUCT(MlirValue, const void);
+DEFINE_C_API_STRUCT(MlirDialectHandle, const void);
#undef DEFINE_C_API_STRUCT
@@ -207,11 +208,6 @@ MLIR_CAPI_EXPORTED MlirStringRef mlirDialectGetNamespace(MlirDialect dialect);
// registration schemes.
//===----------------------------------------------------------------------===//
-struct MlirDialectHandle {
- const void *ptr;
-};
-typedef struct MlirDialectHandle MlirDialectHandle;
-
#define MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Name, Namespace) \
MLIR_CAPI_EXPORTED MlirDialectHandle mlirGetDialectHandle__##Namespace##__( \
void)
@@ -233,6 +229,11 @@ MLIR_CAPI_EXPORTED void mlirDialectHandleRegisterDialect(MlirDialectHandle,
MLIR_CAPI_EXPORTED MlirDialect mlirDialectHandleLoadDialect(MlirDialectHandle,
MlirContext);
+/// Checks if the dialect handle is null.
+static inline bool mlirDialectHandleIsNull(MlirDialectHandle handle) {
+ return !handle.ptr;
+}
+
//===----------------------------------------------------------------------===//
// DialectRegistry API.
//===----------------------------------------------------------------------===//
@@ -249,6 +250,13 @@ static inline bool mlirDialectRegistryIsNull(MlirDialectRegistry registry) {
MLIR_CAPI_EXPORTED void
mlirDialectRegistryDestroy(MlirDialectRegistry registry);
+MLIR_CAPI_EXPORTED int64_t
+mlirDialectRegistryGetNumDialectNames(MlirDialectRegistry registry);
+
+MLIR_CAPI_EXPORTED void
+mlirDialectRegistryGetDialectNames(MlirDialectRegistry registry,
+ MlirStringRef *dialectNames);
+
//===----------------------------------------------------------------------===//
// Location API.
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 83a8757bb72c7..c3d28733c2cee 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2897,6 +2897,14 @@ maybeGetTracebackLocation(const std::optional<PyLocation> &location) {
// Populates the core exports of the 'ir' submodule.
//------------------------------------------------------------------------------
+MlirDialectHandle createMlirDialectHandleFromCapsule(nb::object capsule) {
+ MlirDialectHandle rawRegistry =
+ mlirPythonCapsuleToDialectHandle(capsule.ptr());
+ if (mlirDialectHandleIsNull(rawRegistry))
+ throw nb::python_error();
+ return rawRegistry;
+}
+
void mlir::python::populateIRCore(nb::module_ &m) {
// disable leak warnings which tend to be false positives.
nb::set_leak_warnings(false);
@@ -3126,6 +3134,14 @@ void mlir::python::populateIRCore(nb::module_ &m) {
},
nb::sig("def __repr__(self) -> str"));
+ //----------------------------------------------------------------------------
+ // Mapping of MlirDialectHandle
+ //----------------------------------------------------------------------------
+
+ nb::class_<MlirDialectHandle>(m, "DialectHandle")
+ .def_static(MLIR_PYTHON_CAPI_FACTORY_ATTR,
+ &createMlirDialectHandleFromCapsule);
+
//----------------------------------------------------------------------------
// Mapping of PyDialectRegistry
//----------------------------------------------------------------------------
@@ -3133,7 +3149,23 @@ void mlir::python::populateIRCore(nb::module_ &m) {
.def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyDialectRegistry::getCapsule)
.def_static(MLIR_PYTHON_CAPI_FACTORY_ATTR,
&PyDialectRegistry::createFromCapsule)
- .def(nb::init<>());
+ .def(nb::init<>())
+ .def("insert_dialect",
+ [](PyDialectRegistry &self, MlirDialectHandle handle) {
+ mlirDialectHandleInsertDialect(handle, self.get());
+ })
+ .def("insert_dialect",
+ [](PyDialectRegistry &self, intptr_t ptr) {
+ mlirDialectHandleInsertDialect(
+ {reinterpret_cast<const void *>(ptr)}, self.get());
+ })
+ .def_prop_ro("dialect_names", [](PyDialectRegistry &self) {
+ int64_t numDialectNames =
+ mlirDialectRegistryGetNumDialectNames(self.get());
+ std::vector<MlirStringRef> dialectNames(numDialectNames);
+ mlirDialectRegistryGetDialectNames(self.get(), dialectNames.data());
+ return dialectNames;
+ });
//----------------------------------------------------------------------------
// Mapping of Location
diff --git a/mlir/lib/CAPI/Dialect/Affine.cpp b/mlir/lib/CAPI/Dialect/Affine.cpp
new file mode 100644
index 0000000000000..b796523390a29
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/Affine.cpp
@@ -0,0 +1,14 @@
+//===- Affine.cpp - C Interface for Affine dialect ------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Affine.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Affine, affine,
+ mlir::affine::AffineDialect)
diff --git a/mlir/lib/CAPI/Dialect/Bufferization.cpp b/mlir/lib/CAPI/Dialect/Bufferization.cpp
new file mode 100644
index 0000000000000..da50dfaab44b6
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/Bufferization.cpp
@@ -0,0 +1,14 @@
+//===- Bufferization.cpp - C Interface for Bufferization dialect ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Bufferization.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Bufferization, bufferization,
+ mlir::bufferization::BufferizationDialect)
diff --git a/mlir/lib/CAPI/Dialect/Builtin.cpp b/mlir/lib/CAPI/Dialect/Builtin.cpp
new file mode 100644
index 0000000000000..d095daa294c56
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/Builtin.cpp
@@ -0,0 +1,13 @@
+//===- Builtin.cpp - C Interface for Builtin dialect ----------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Builtin.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/IR/BuiltinDialect.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Builtin, builtin, mlir::BuiltinDialect)
diff --git a/mlir/lib/CAPI/Dialect/CMakeLists.txt b/mlir/lib/CAPI/Dialect/CMakeLists.txt
index bb1fdf8be3c8f..44928f4c86f77 100644
--- a/mlir/lib/CAPI/Dialect/CMakeLists.txt
+++ b/mlir/lib/CAPI/Dialect/CMakeLists.txt
@@ -7,6 +7,15 @@ add_mlir_upstream_c_api_library(MLIRCAPIAMDGPU
MLIRAMDGPUDialect
)
+add_mlir_upstream_c_api_library(MLIRCAPIAffine
+ Affine.cpp
+
+ PARTIAL_SOURCES_INTENDED
+ LINK_LIBS PUBLIC
+ MLIRCAPIIR
+ MLIRAffineDialect
+)
+
add_mlir_upstream_c_api_library(MLIRCAPIArith
Arith.cpp
@@ -16,6 +25,23 @@ add_mlir_upstream_c_api_library(MLIRCAPIArith
MLIRArithDialect
)
+add_mlir_upstream_c_api_library(MLIRCAPIBufferization
+ Bufferization.cpp
+
+ PARTIAL_SOURCES_INTENDED
+ LINK_LIBS PUBLIC
+ MLIRCAPIIR
+ MLIRBufferizationDialect
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPIBuiltin
+ Builtin.cpp
+
+ PARTIAL_SOURCES_INTENDED
+ LINK_LIBS PUBLIC
+ MLIRCAPIIR
+)
+
add_mlir_upstream_c_api_library(MLIRCAPIAsync
Async.cpp
AsyncPasses.cpp
@@ -31,6 +57,15 @@ add_mlir_upstream_c_api_library(MLIRCAPIAsync
MLIRPass
)
+add_mlir_upstream_c_api_library(MLIRCAPIComplex
+ Complex.cpp
+
+ PARTIAL_SOURCES_INTENDED
+ LINK_LIBS PUBLIC
+ MLIRCAPIIR
+ MLIRComplexDialect
+)
+
add_mlir_upstream_c_api_library(MLIRCAPIControlFlow
ControlFlow.cpp
@@ -223,6 +258,9 @@ add_mlir_upstream_c_api_library(MLIRCAPITransformDialect
LINK_LIBS PUBLIC
MLIRCAPIIR
MLIRTransformDialect
+ MLIRTransformDebugExtension
+ MLIRTransformLoopExtension
+ MLIRTransformTuneExtension
)
add_mlir_upstream_c_api_library(MLIRCAPITransformDialectTransforms
@@ -234,6 +272,70 @@ add_mlir_upstream_c_api_library(MLIRCAPITransformDialectTransforms
MLIRTransformDialectTransforms
)
+add_mlir_upstream_c_api_library(MLIRCAPIAffineTransformDialectExtension
+ TransformAffine.cpp
+
+ PARTIAL_SOURCES_INTENDED
+ LINK_LIBS PUBLIC
+ MLIRAffineTransformOps
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPIBufferizationTransformDialectExtension
+ TransformBufferization.cpp
+
+ PARTIAL_SOURCES_INTENDED
+ LINK_LIBS PUBLIC
+ MLIRBufferizationTransformOps
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPIGPUTransformDialectExtension
+ TransformGPU.cpp
+
+ PARTIAL_SOURCES_INTENDED
+ LINK_LIBS PUBLIC
+ MLIRGPUTransformOps
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPIMemRefTransformDialectExtension
+ TransformMemRef.cpp
+
+ PARTIAL_SOURCES_INTENDED
+ LINK_LIBS PUBLIC
+ MLIRMemRefTransformOps
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPISMTTransformDialectExtension
+ TransformSMT.cpp
+
+ PARTIAL_SOURCES_INTENDED
+ LINK_LIBS PUBLIC
+ MLIRTransformSMTExtension
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPISparseTensorTransformDialectExtension
+ TransformSparseTensor.cpp
+
+ PARTIAL_SOURCES_INTENDED
+ LINK_LIBS PUBLIC
+ MLIRSparseTensorTransformOps
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPITensorTransformDialectExtension
+ TransformTensor.cpp
+
+ PARTIAL_SOURCES_INTENDED
+ LINK_LIBS PUBLIC
+ MLIRTensorTransformOps
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPIVectorTransformDialectExtension
+ TransformVector.cpp
+
+ PARTIAL_SOURCES_INTENDED
+ LINK_LIBS PUBLIC
+ MLIRVectorTransformOps
+)
+
add_mlir_upstream_c_api_library(MLIRCAPIQuant
Quant.cpp
@@ -278,3 +380,21 @@ add_mlir_upstream_c_api_library(MLIRCAPISMT
MLIRCAPIIR
MLIRSMT
)
+
+add_mlir_upstream_c_api_library(MLIRCAPITosa
+ Tosa.cpp
+
+ PARTIAL_SOURCES_INTENDED
+ LINK_LIBS PUBLIC
+ MLIRCAPIIR
+ MLIRTosaDialect
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPIUB
+ UB.cpp
+
+ PARTIAL_SOURCES_INTENDED
+ LINK_LIBS PUBLIC
+ MLIRCAPIIR
+ MLIRUBDialect
+)
diff --git a/mlir/lib/CAPI/Dialect/Complex.cpp b/mlir/lib/CAPI/Dialect/Complex.cpp
new file mode 100644
index 0000000000000..7063028e5d640
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/Complex.cpp
@@ -0,0 +1,14 @@
+//===- Complex.cpp - C Interface for Complex dialect ----------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Complex.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Complex/IR/Complex.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Complex, complex,
+ mlir::complex::ComplexDialect)
diff --git a/mlir/lib/CAPI/Dialect/Tosa.cpp b/mlir/lib/CAPI/Dialect/Tosa.cpp
new file mode 100644
index 0000000000000..357717a2ed5a5
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/Tosa.cpp
@@ -0,0 +1,13 @@
+//===- Tosa.cpp - C Interface for Tosa dialect ----------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Tosa.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Tosa, tosa, mlir::tosa::TosaDialect)
diff --git a/mlir/lib/CAPI/Dialect/Transform.cpp b/mlir/lib/CAPI/Dialect/Transform.cpp
index 5fd773572bd3c..8942d0a00cde3 100644
--- a/mlir/lib/CAPI/Dialect/Transform.cpp
+++ b/mlir/lib/CAPI/Dialect/Transform.cpp
@@ -9,8 +9,11 @@
#include "mlir-c/Dialect/Transform.h"
#include "mlir-c/Support.h"
#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Transform/DebugExtension/DebugExtension.h"
#include "mlir/Dialect/Transform/IR/TransformDialect.h"
#include "mlir/Dialect/Transform/IR/TransformTypes.h"
+#include "mlir/Dialect/Transform/LoopExtension/LoopExtension.h"
+#include "mlir/Dialect/Transform/TuneExtension/TuneExtension.h"
using namespace mlir;
@@ -106,3 +109,15 @@ MlirType mlirTransformParamTypeGet(MlirContext ctx, MlirType type) {
MlirType mlirTransformParamTypeGetType(MlirType type) {
return wrap(cast<transform::ParamType>(unwrap(type)).getType());
}
+
+void mlirDebugRegisterTransformDialectExtension(MlirDialectRegistry registry) {
+ mlir::transform::registerDebugExtension(*unwrap(registry));
+}
+
+void mlirLoopRegisterTransformDialectExtension(MlirDialectRegistry registry) {
+ mlir::transform::registerLoopExtension(*unwrap(registry));
+}
+
+void mlirTuneRegisterTransformDialectExtension(MlirDialectRegistry registry) {
+ mlir::transform::registerTuneExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/TransformAffine.cpp b/mlir/lib/CAPI/Dialect/TransformAffine.cpp
new file mode 100644
index 0000000000000..6fd7744e8c2c9
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/TransformAffine.cpp
@@ -0,0 +1,16 @@
+//===- TransformAffine.cpp - C Interface for Transform affine 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Affine.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.h"
+
+MLIR_CAPI_EXPORTED void
+mlirAffineRegisterTransformDialectExtension(MlirDialectRegistry registry) {
+ mlir::affine::registerTransformDialectExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/TransformBufferization.cpp b/mlir/lib/CAPI/Dialect/TransformBufferization.cpp
new file mode 100644
index 0000000000000..f1bd61be10880
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/TransformBufferization.cpp
@@ -0,0 +1,16 @@
+//==- TransformBufferization.cpp - C Interface for bufferization 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Bufferization.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h"
+
+void mlirBufferizationRegisterTransformDialectExtension(
+ MlirDialectRegistry registry) {
+ mlir::bufferization::registerTransformDialectExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/TransformGPU.cpp b/mlir/lib/CAPI/Dialect/TransformGPU.cpp
new file mode 100644
index 0000000000000..51187a28fa752
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/TransformGPU.cpp
@@ -0,0 +1,15 @@
+//===- TransformGPU.cpp - C Interface for Transform GPU 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/GPU.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h"
+
+void mlirGPURegisterTransformDialectExtension(MlirDialectRegistry registry) {
+ mlir::gpu::registerTransformDialectExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/TransformMemRef.cpp b/mlir/lib/CAPI/Dialect/TransformMemRef.cpp
new file mode 100644
index 0000000000000..9b80295748514
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/TransformMemRef.cpp
@@ -0,0 +1,15 @@
+//===- TransformMemRef.cpp - C Interface for Transform MemRef 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/MemRef.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.h"
+
+void mlirMemRefRegisterTransformDialectExtension(MlirDialectRegistry registry) {
+ mlir::memref::registerTransformDialectExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/TransformSMT.cpp b/mlir/lib/CAPI/Dialect/TransformSMT.cpp
new file mode 100644
index 0000000000000..449129dac99e5
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/TransformSMT.cpp
@@ -0,0 +1,15 @@
+//===- TransformSMT.cpp - C Interface for Transform SMT 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/SMT.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Transform/SMTExtension/SMTExtension.h"
+
+void mlirSMTRegisterTransformDialectExtension(MlirDialectRegistry registry) {
+ mlir::transform::registerSMTExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/TransformSparseTensor.cpp b/mlir/lib/CAPI/Dialect/TransformSparseTensor.cpp
new file mode 100644
index 0000000000000..3e5b91099a6fb
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/TransformSparseTensor.cpp
@@ -0,0 +1,16 @@
+//===- TransformSparseTensor.cpp - C Interface for SparseTensor 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/SparseTensor.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.h"
+
+void mlirSparseTensorRegisterTransformDialectExtension(
+ MlirDialectRegistry registry) {
+ mlir::sparse_tensor::registerTransformDialectExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/TransformTensor.cpp b/mlir/lib/CAPI/Dialect/TransformTensor.cpp
new file mode 100644
index 0000000000000..fc7a30b5b237e
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/TransformTensor.cpp
@@ -0,0 +1,15 @@
+//===- TransformTensor.cpp - C Interface for Transform tensor extensio ----===//
+//
+// 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 "mlir-c/Dialect/Tensor.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h"
+
+void mlirTensorRegisterTransformDialectExtension(MlirDialectRegistry registry) {
+ mlir::tensor::registerTransformDialectExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/TransformVector.cpp b/mlir/lib/CAPI/Dialect/TransformVector.cpp
new file mode 100644
index 0000000000000..c2cf9057f16e6
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/TransformVector.cpp
@@ -0,0 +1,15 @@
+//===- TransformVector.cpp - C Interface for Transform Vector 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Vector.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.h"
+
+void mlirVectorRegisterTransformDialectExtension(MlirDialectRegistry registry) {
+ mlir::vector::registerTransformDialectExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/UB.cpp b/mlir/lib/CAPI/Dialect/UB.cpp
new file mode 100644
index 0000000000000..de989237159c4
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/UB.cpp
@@ -0,0 +1,13 @@
+//===- Ub.cpp - C Interface for UB dialect --------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/UB.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(UB, ub, mlir::ub::UBDialect)
diff --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp
index e9844a7cc1909..a81e2a14e5255 100644
--- a/mlir/lib/CAPI/IR/IR.cpp
+++ b/mlir/lib/CAPI/IR/IR.cpp
@@ -150,6 +150,18 @@ void mlirDialectRegistryDestroy(MlirDialectRegistry registry) {
delete unwrap(registry);
}
+int64_t mlirDialectRegistryGetNumDialectNames(MlirDialectRegistry registry) {
+ auto dialectNames = unwrap(registry)->getDialectNames();
+ return std::distance(dialectNames.begin(), dialectNames.end());
+}
+
+void mlirDialectRegistryGetDialectNames(MlirDialectRegistry registry,
+ MlirStringRef *dialectNames) {
+ for (auto [i, location] :
+ llvm::enumerate(unwrap(registry)->getDialectNames()))
+ dialectNames[i] = wrap(location);
+}
+
//===----------------------------------------------------------------------===//
// AsmState API.
//===----------------------------------------------------------------------===//
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 368f344841ed1..c6e38592d17d7 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -64,7 +64,10 @@ declare_mlir_dialect_python_bindings(
SOURCES
dialects/affine.py
DIALECT_NAME affine
- GEN_ENUM_BINDINGS)
+ GEN_ENUM_BINDINGS
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIAffine
+ )
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -73,14 +76,20 @@ declare_mlir_dialect_python_bindings(
SOURCES
dialects/amdgpu.py
DIALECT_NAME amdgpu
- GEN_ENUM_BINDINGS)
+ GEN_ENUM_BINDINGS
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIAMDGPU
+ )
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
TD_FILE dialects/AsyncOps.td
SOURCES_GLOB dialects/async_dialect/*.py
- DIALECT_NAME async)
+ DIALECT_NAME async
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIAsync
+ )
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -91,6 +100,8 @@ declare_mlir_dialect_python_bindings(
DIALECT_NAME bufferization
GEN_ENUM_BINDINGS_TD_FILE
"../../include/mlir/Dialect/Bufferization/IR/BufferizationEnums.td"
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIBufferization
)
declare_mlir_dialect_python_bindings(
@@ -99,7 +110,10 @@ declare_mlir_dialect_python_bindings(
TD_FILE dialects/BuiltinOps.td
SOURCES
dialects/builtin.py
- DIALECT_NAME builtin)
+ DIALECT_NAME builtin
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIBuiltin
+ )
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -107,7 +121,10 @@ declare_mlir_dialect_python_bindings(
TD_FILE dialects/ComplexOps.td
SOURCES
dialects/complex.py
- DIALECT_NAME complex)
+ DIALECT_NAME complex
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIComplex
+ )
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -116,7 +133,10 @@ declare_mlir_dialect_python_bindings(
SOURCES
dialects/index.py
DIALECT_NAME index
- GEN_ENUM_BINDINGS)
+ GEN_ENUM_BINDINGS
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIIndex
+ )
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -124,7 +144,10 @@ declare_mlir_dialect_python_bindings(
TD_FILE dialects/ControlFlowOps.td
SOURCES
dialects/cf.py
- DIALECT_NAME cf)
+ DIALECT_NAME cf
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIControlFlow
+ )
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -132,7 +155,10 @@ declare_mlir_dialect_python_bindings(
TD_FILE dialects/FuncOps.td
SOURCES
dialects/func.py
- DIALECT_NAME func)
+ DIALECT_NAME func
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIFunc
+ )
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -140,7 +166,10 @@ declare_mlir_dialect_python_bindings(
TD_FILE dialects/GPUOps.td
SOURCES_GLOB dialects/gpu/*.py
DIALECT_NAME gpu
- GEN_ENUM_BINDINGS)
+ GEN_ENUM_BINDINGS
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIGPU
+ )
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -151,7 +180,10 @@ declare_mlir_dialect_python_bindings(
dialects/linalg/*.py
DIALECT_NAME linalg
DEPENDS LinalgOdsGen
- GEN_ENUM_BINDINGS)
+ GEN_ENUM_BINDINGS
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPILinalg
+ )
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -160,11 +192,14 @@ declare_mlir_dialect_python_bindings(
SOURCES
dialects/llvm.py
DIALECT_NAME llvm
- GEN_ENUM_BINDINGS)
+ GEN_ENUM_BINDINGS
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPILLVM
+ )
declare_mlir_dialect_extension_python_bindings(
-ADD_TO_PARENT MLIRPythonSources.Dialects
-ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+ ADD_TO_PARENT MLIRPythonSources.Dialects
+ ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
TD_FILE dialects/TransformPDLExtensionOps.td
SOURCES
dialects/transform/pdl.py
@@ -172,8 +207,8 @@ ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
EXTENSION_NAME transform_pdl_extension)
declare_mlir_dialect_extension_python_bindings(
-ADD_TO_PARENT MLIRPythonSources.Dialects
-ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+ ADD_TO_PARENT MLIRPythonSources.Dialects
+ ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
TD_FILE dialects/TransformSMTExtensionOps.td
SOURCES
dialects/transform/smt.py
@@ -181,8 +216,8 @@ ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
EXTENSION_NAME transform_smt_extension)
declare_mlir_dialect_extension_python_bindings(
-ADD_TO_PARENT MLIRPythonSources.Dialects
-ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+ ADD_TO_PARENT MLIRPythonSources.Dialects
+ ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
TD_FILE dialects/TransformDebugExtensionOps.td
SOURCES
dialects/transform/debug.py
@@ -190,8 +225,8 @@ ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
EXTENSION_NAME transform_debug_extension)
declare_mlir_dialect_extension_python_bindings(
-ADD_TO_PARENT MLIRPythonSources.Dialects
-ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+ ADD_TO_PARENT MLIRPythonSources.Dialects
+ ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
TD_FILE dialects/TransformTuneExtensionOps.td
SOURCES
dialects/transform/tune.py
@@ -208,7 +243,9 @@ declare_mlir_dialect_python_bindings(
DIALECT_NAME transform
GEN_ENUM_BINDINGS_TD_FILE
"../../include/mlir/Dialect/Transform/IR/TransformAttrs.td"
-)
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPITransformDialect
+ )
declare_mlir_python_sources(
MLIRPythonSources.Dialects.transform.extras
@@ -232,7 +269,11 @@ declare_mlir_dialect_extension_python_bindings(
SOURCES
dialects/transform/bufferization.py
DIALECT_NAME transform
- EXTENSION_NAME bufferization_transform)
+ EXTENSION_NAME bufferization_transform
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPITransformDialect
+ MLIRCAPIBufferizationTransformDialectExtension
+ )
declare_mlir_dialect_extension_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -241,7 +282,11 @@ declare_mlir_dialect_extension_python_bindings(
SOURCES
dialects/transform/gpu.py
DIALECT_NAME transform
- EXTENSION_NAME gpu_transform)
+ EXTENSION_NAME gpu_transform
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPITransformDialect
+ MLIRCAPIGPUTransformDialectExtension
+)
declare_mlir_dialect_extension_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -250,7 +295,10 @@ declare_mlir_dialect_extension_python_bindings(
SOURCES
dialects/transform/loop.py
DIALECT_NAME transform
- EXTENSION_NAME loop_transform)
+ EXTENSION_NAME loop_transform
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPITransformDialect
+)
declare_mlir_dialect_extension_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -259,7 +307,11 @@ declare_mlir_dialect_extension_python_bindings(
SOURCES
dialects/transform/memref.py
DIALECT_NAME transform
- EXTENSION_NAME memref_transform)
+ EXTENSION_NAME memref_transform
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPITransformDialect
+ MLIRCAPIMemRefTransformDialectExtension
+)
declare_mlir_dialect_extension_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -289,7 +341,11 @@ declare_mlir_dialect_extension_python_bindings(
SOURCES
dialects/transform/sparse_tensor.py
DIALECT_NAME transform
- EXTENSION_NAME sparse_tensor_transform)
+ EXTENSION_NAME sparse_tensor_transform
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPITransformDialect
+ MLIRCAPISparseTensorTransformDialectExtension
+)
declare_mlir_dialect_extension_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -298,7 +354,11 @@ declare_mlir_dialect_extension_python_bindings(
SOURCES
dialects/transform/tensor.py
DIALECT_NAME transform
- EXTENSION_NAME tensor_transform)
+ EXTENSION_NAME tensor_transform
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPITransformDialect
+ MLIRCAPITensorTransformDialectExtension
+)
declare_mlir_dialect_extension_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -310,6 +370,9 @@ declare_mlir_dialect_extension_python_bindings(
EXTENSION_NAME vector_transform
GEN_ENUM_BINDINGS_TD_FILE
"../../include/mlir/Dialect/Vector/Transforms/VectorTransformsBase.td"
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPITransformDialect
+ MLIRCAPIVectorTransformDialectExtension
)
declare_mlir_dialect_python_bindings(
@@ -317,7 +380,10 @@ declare_mlir_dialect_python_bindings(
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
TD_FILE dialects/MathOps.td
SOURCES dialects/math.py
- DIALECT_NAME math)
+ DIALECT_NAME math
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIMath
+)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -329,7 +395,7 @@ declare_mlir_dialect_python_bindings(
GEN_ENUM_BINDINGS
EMBED_CAPI_LINK_LIBS
MLIRCAPIArith
- )
+)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -337,7 +403,10 @@ declare_mlir_dialect_python_bindings(
TD_FILE dialects/MemRefOps.td
SOURCES
dialects/memref.py
- DIALECT_NAME memref)
+ DIALECT_NAME memref
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIMemRef
+)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -345,7 +414,10 @@ declare_mlir_dialect_python_bindings(
TD_FILE dialects/MLProgramOps.td
SOURCES
dialects/ml_program.py
- DIALECT_NAME ml_program)
+ DIALECT_NAME ml_program
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIMLProgram
+)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -354,7 +426,10 @@ declare_mlir_dialect_python_bindings(
SOURCES
dialects/nvgpu.py
DIALECT_NAME nvgpu
- GEN_ENUM_BINDINGS)
+ GEN_ENUM_BINDINGS
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPINVGPU
+)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -363,7 +438,10 @@ declare_mlir_dialect_python_bindings(
SOURCES
dialects/nvvm.py
DIALECT_NAME nvvm
- GEN_ENUM_BINDINGS)
+ GEN_ENUM_BINDINGS
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPINVVM
+)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -371,16 +449,23 @@ declare_mlir_dialect_python_bindings(
TD_FILE dialects/ROCDLOps.td
SOURCES
dialects/rocdl.py
- DIALECT_NAME rocdl)
+ DIALECT_NAME rocdl
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIROCDL
+)
-declare_mlir_python_sources(
- MLIRPythonSources.Dialects.quant
+declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+ TD_FILE dialects/QuantOps.td
GEN_ENUM_BINDINGS
SOURCES
dialects/quant.py
- _mlir_libs/_mlir/dialects/quant.pyi)
+ _mlir_libs/_mlir/dialects/quant.pyi
+ DIALECT_NAME quant
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIQuant
+)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -388,7 +473,10 @@ declare_mlir_dialect_python_bindings(
TD_FILE dialects/EmitC.td
SOURCES
dialects/emitc.py
- DIALECT_NAME emitc)
+ DIALECT_NAME emitc
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIEmitC
+)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -397,7 +485,10 @@ declare_mlir_dialect_python_bindings(
SOURCES
dialects/pdl.py
_mlir_libs/_mlir/dialects/pdl.pyi
- DIALECT_NAME pdl)
+ DIALECT_NAME pdl
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIPDL
+)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -406,7 +497,10 @@ declare_mlir_dialect_python_bindings(
SOURCES
dialects/openmp.py
DIALECT_NAME omp
- DEPENDS omp_common_td)
+ DEPENDS omp_common_td
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIOpenMP
+)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -414,14 +508,20 @@ declare_mlir_dialect_python_bindings(
TD_FILE dialects/SCFOps.td
SOURCES
dialects/scf.py
- DIALECT_NAME scf)
+ DIALECT_NAME scf
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPISCF
+)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
TD_FILE dialects/ShapeOps.td
SOURCES dialects/shape.py
- DIALECT_NAME shape)
+ DIALECT_NAME shape
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIShape
+)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -431,6 +531,8 @@ declare_mlir_dialect_python_bindings(
DIALECT_NAME sparse_tensor
GEN_ENUM_BINDINGS_TD_FILE
"../../include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td"
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPISparseTensor
)
declare_mlir_dialect_python_bindings(
@@ -440,14 +542,20 @@ declare_mlir_dialect_python_bindings(
GEN_ENUM_BINDINGS
SOURCES
dialects/smt.py
- DIALECT_NAME smt)
+ DIALECT_NAME smt
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPISMT
+)
declare_mlir_dialect_python_bindings(
- ADD_TO_PARENT MLIRPythonSources.Dialects
- ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
- TD_FILE dialects/SPIRVOps.td
- SOURCES dialects/spirv.py
- DIALECT_NAME spirv)
+ ADD_TO_PARENT MLIRPythonSources.Dialects
+ ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+ TD_FILE dialects/SPIRVOps.td
+ SOURCES dialects/spirv.py
+ DIALECT_NAME spirv
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPISPIRV
+)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -455,7 +563,10 @@ declare_mlir_dialect_python_bindings(
TD_FILE dialects/TensorOps.td
SOURCES
dialects/tensor.py
- DIALECT_NAME tensor)
+ DIALECT_NAME tensor
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPITensor
+)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -463,6 +574,8 @@ declare_mlir_dialect_python_bindings(
TD_FILE dialects/TosaOps.td
SOURCES dialects/tosa.py
DIALECT_NAME tosa
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPITosa
)
declare_mlir_dialect_python_bindings(
@@ -471,6 +584,8 @@ declare_mlir_dialect_python_bindings(
TD_FILE dialects/UBOps.td
SOURCES dialects/ub.py
DIALECT_NAME ub
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIUB
)
declare_mlir_dialect_python_bindings(
@@ -480,7 +595,10 @@ declare_mlir_dialect_python_bindings(
SOURCES dialects/vector.py
DIALECT_NAME vector
GEN_ENUM_BINDINGS_TD_FILE
- "dialects/VectorAttributes.td")
+ "dialects/VectorAttributes.td"
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIVector
+)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -489,6 +607,8 @@ declare_mlir_dialect_python_bindings(
SOURCES dialects/irdl.py
DIALECT_NAME irdl
GEN_ENUM_BINDINGS
+ EMBED_CAPI_LINK_LIBS
+ MLIRCAPIIRDL
)
################################################################################
@@ -528,31 +648,32 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
MLIRCAPIDebug
MLIRCAPIIR
MLIRCAPIInterfaces
+ MLIRCAPITransforms
+ MLIRCAPIBuiltin
# Dialects
MLIRCAPIFunc
)
-# This extension exposes an API to register all dialects, extensions, and passes
-# packaged in upstream MLIR and it is used for the upstream "mlir" Python
-# package. Downstreams will likely want to provide their own and not depend
-# on this one, since it links in the world.
-# Note that this is not added to any top-level source target for transitive
-# inclusion: It must be included explicitly by downstreams if desired. Note that
-# this has a very large impact on what gets built/packaged.
-declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
- MODULE_NAME _mlirRegisterEverything
- ROOT_DIR "${PYTHON_SOURCE_DIR}"
- PYTHON_BINDINGS_LIBRARY nanobind
- SOURCES
- RegisterEverything.cpp
- PRIVATE_LINK_LIBS
- LLVMSupport
- EMBED_CAPI_LINK_LIBS
- MLIRCAPIConversion
- MLIRCAPITransforms
- MLIRCAPIRegisterEverything
-)
+## This extension exposes an API to register all dialects, extensions, and passes
+## packaged in upstream MLIR and it is used for the upstream "mlir" Python
+## package. Downstreams will likely want to provide their own and not depend
+## on this one, since it links in the world.
+## Note that this is not added to any top-level source target for transitive
+## inclusion: It must be included explicitly by downstreams if desired. Note that
+## this has a very large impact on what gets built/packaged.
+#declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
+# MODULE_NAME _mlirRegisterEverything
+# ROOT_DIR "${PYTHON_SOURCE_DIR}"
+# PYTHON_BINDINGS_LIBRARY nanobind
+# SOURCES
+# RegisterEverything.cpp
+# PRIVATE_LINK_LIBS
+# LLVMSupport
+# MLIRCAPIConversion
+# MLIRCAPITransforms
+# MLIRCAPIRegisterEverything
+#)
declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
MODULE_NAME _mlirDialectsLinalg
@@ -749,6 +870,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
MLIRCAPIIR
MLIRCAPISMT
MLIRCAPIExportSMTLIB
+ MLIRCAPISMTTransformDialectExtension
)
declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
@@ -871,7 +993,6 @@ add_mlir_python_common_capi_library(${MLIR_PYTHON_CAPI_DYLIB_NAME}
MLIRPythonCAPI.HeaderSources
DECLARED_SOURCES
MLIRPythonSources
- MLIRPythonExtension.RegisterEverything
${_ADDL_TEST_SOURCES}
)
@@ -980,7 +1101,6 @@ endif()
set(_declared_sources
MLIRPythonSources
- MLIRPythonExtension.RegisterEverything
MLIRPythonCAPICTypesBinding)
if(NOT CMAKE_CROSSCOMPILING)
list(APPEND _declared_sources MLIRPythonExtension.Core.type_stub_gen)
diff --git a/mlir/python/mlir/_mlir_libs/_capi.py.in b/mlir/python/mlir/_mlir_libs/_capi.py.in
index 9568845e67de9..ef6e0ea973cac 100644
--- a/mlir/python/mlir/_mlir_libs/_capi.py.in
+++ b/mlir/python/mlir/_mlir_libs/_capi.py.in
@@ -5,4 +5,47 @@
import ctypes
from pathlib import Path
-_capi = ctypes.CDLL(str(Path(__file__).parent / "@CMAKE_SHARED_LIBRARY_PREFIX@@MLIR_PYTHON_CAPI_DYLIB_NAME@@CMAKE_SHARED_LIBRARY_SUFFIX@"))
\ No newline at end of file
+_capi = ctypes.CDLL(
+ str(
+ Path(__file__).parent
+ / "@CMAKE_SHARED_LIBRARY_PREFIX@@MLIR_PYTHON_CAPI_DYLIB_NAME@@CMAKE_SHARED_LIBRARY_SUFFIX@"
+ )
+)
+
+PyCapsule_New = ctypes.pythonapi.PyCapsule_New
+PyCapsule_New.restype = ctypes.py_object
+PyCapsule_New.argtypes = ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p
+
+PyCapsule_GetPointer = ctypes.pythonapi.PyCapsule_GetPointer
+PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p]
+PyCapsule_GetPointer.restype = ctypes.c_void_p
+
+MLIR_PYTHON_CAPSULE_DIALECT_HANDLE = (
+ "@MLIR_PYTHON_PACKAGE_PREFIX at .ir.DialectHandle._CAPIPtr"
+).encode()
+
+MLIR_PYTHON_CAPSULE_DIALECT_REGISTRY_HANDLE = (
+ "@MLIR_PYTHON_PACKAGE_PREFIX at .ir.DialectRegistry._CAPIPtr"
+).encode()
+
+
+def register_dialect(dialect_handle_capi_name, dialect_registry):
+ if not hasattr(_capi, dialect_handle_capi_name):
+ raise RuntimeError(f"missing {dialect_handle_capi_name} API")
+ dialect_handle_capi = getattr(_capi, dialect_handle_capi_name)
+ dialect_handle_capi.argtypes = []
+ dialect_handle_capi.restype = ctypes.c_void_p
+ handle = dialect_handle_capi()
+ dialect_registry.insert_dialect(handle)
+
+
+def register_transform_dialect_extension(registration_capi_name, dialect_registry):
+ if not hasattr(_capi, registration_capi_name):
+ raise RuntimeError(f"missing {registration_capi_name} API")
+ registration_capi = getattr(_capi, registration_capi_name)
+ registration_capi.argtypes = [ctypes.c_void_p]
+ registration_capi(
+ PyCapsule_GetPointer(
+ dialect_registry._CAPIPtr, MLIR_PYTHON_CAPSULE_DIALECT_REGISTRY_HANDLE
+ )
+ )
diff --git a/mlir/python/mlir/dialects/QuantOps.td b/mlir/python/mlir/dialects/QuantOps.td
new file mode 100644
index 0000000000000..46385fd00ac0c
--- /dev/null
+++ b/mlir/python/mlir/dialects/QuantOps.td
@@ -0,0 +1,14 @@
+//===-- QuantOps.td - Entry point for QuantOps bind --------*- tablegen -*-===//
+//
+// 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 PYTHON_BINDINGS_QUANT_OPS
+#define PYTHON_BINDINGS_QUANT_OPS
+
+include "mlir/Dialect/Quant/IR/QuantOps.td"
+
+#endif
diff --git a/mlir/python/mlir/dialects/quant.py b/mlir/python/mlir/dialects/quant.py
index bf1fc5f2de378..7a7273d8e26be 100644
--- a/mlir/python/mlir/dialects/quant.py
+++ b/mlir/python/mlir/dialects/quant.py
@@ -2,4 +2,5 @@
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+from ._quant_ops_gen import *
from .._mlir_libs._mlirDialectsQuant import *
diff --git a/mlir/python/mlir/dialects/transform/bufferization.py b/mlir/python/mlir/dialects/transform/bufferization.py
index 485a8a36b6305..228c9dab52b8d 100644
--- a/mlir/python/mlir/dialects/transform/bufferization.py
+++ b/mlir/python/mlir/dialects/transform/bufferization.py
@@ -15,6 +15,15 @@
from enum import Enum
from typing import Optional, overload, Union
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+ register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+ "mlirBufferizationRegisterTransformDialectExtension", _get_dialect_registry()
+)
+
@_ods_cext.register_operation(_Dialect, replace=True)
class EmptyTensorToAllocTensorOp(EmptyTensorToAllocTensorOp):
diff --git a/mlir/python/mlir/dialects/transform/gpu.py b/mlir/python/mlir/dialects/transform/gpu.py
index 00cf0840eeae9..b35004cb4a09a 100644
--- a/mlir/python/mlir/dialects/transform/gpu.py
+++ b/mlir/python/mlir/dialects/transform/gpu.py
@@ -15,6 +15,16 @@
from typing import Optional, Sequence, Union, overload
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+ register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+ "mlirGPURegisterTransformDialectExtension", _get_dialect_registry()
+)
+
+
@_ods_cext.register_operation(_Dialect, replace=True)
class MapForallToBlocks(MapForallToBlocks):
"""Specialization for MapForallToBlocks class."""
diff --git a/mlir/python/mlir/dialects/transform/loop.py b/mlir/python/mlir/dialects/transform/loop.py
index c4770b1c4067e..3311ae47b4fa9 100644
--- a/mlir/python/mlir/dialects/transform/loop.py
+++ b/mlir/python/mlir/dialects/transform/loop.py
@@ -16,6 +16,15 @@
from typing import Optional, Union
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+ register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+ "mlirLoopRegisterTransformDialectExtension", _get_dialect_registry()
+)
+
@_ods_cext.register_operation(_Dialect, replace=True)
class LoopOutlineOp(LoopOutlineOp):
diff --git a/mlir/python/mlir/dialects/transform/memref.py b/mlir/python/mlir/dialects/transform/memref.py
index 56ea61eb817f8..f45d8054bfa58 100644
--- a/mlir/python/mlir/dialects/transform/memref.py
+++ b/mlir/python/mlir/dialects/transform/memref.py
@@ -14,6 +14,15 @@
from typing import Optional, overload, Union
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+ register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+ "mlirMemRefRegisterTransformDialectExtension", _get_dialect_registry()
+)
+
@_ods_cext.register_operation(_Dialect, replace=True)
class MemRefAllocaToGlobalOp(MemRefAllocaToGlobalOp):
diff --git a/mlir/python/mlir/dialects/transform/smt.py b/mlir/python/mlir/dialects/transform/smt.py
index 1f0b7f066118c..a4e8eb96d388c 100644
--- a/mlir/python/mlir/dialects/transform/smt.py
+++ b/mlir/python/mlir/dialects/transform/smt.py
@@ -15,6 +15,16 @@
raise RuntimeError("Error loading imports from extension module") from e
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+ register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+ "mlirSMTRegisterTransformDialectExtension", _get_dialect_registry()
+)
+
+
@_ods_cext.register_operation(_Dialect, replace=True)
class ConstrainParamsOp(ConstrainParamsOp):
def __init__(
diff --git a/mlir/python/mlir/dialects/transform/sparse_tensor.py b/mlir/python/mlir/dialects/transform/sparse_tensor.py
index 8b33270dc23a1..efaf438e07b83 100644
--- a/mlir/python/mlir/dialects/transform/sparse_tensor.py
+++ b/mlir/python/mlir/dialects/transform/sparse_tensor.py
@@ -3,3 +3,11 @@
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from .._sparse_tensor_transform_ops_gen import *
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+ register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+ "mlirSparseTensorRegisterTransformDialectExtension", _get_dialect_registry()
+)
diff --git a/mlir/python/mlir/dialects/transform/tensor.py b/mlir/python/mlir/dialects/transform/tensor.py
index 4eb30398f0872..904f57827146d 100644
--- a/mlir/python/mlir/dialects/transform/tensor.py
+++ b/mlir/python/mlir/dialects/transform/tensor.py
@@ -14,6 +14,15 @@
from typing import Optional, overload, Union
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+ register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+ "mlirTensorRegisterTransformDialectExtension", _get_dialect_registry()
+)
+
@_ods_cext.register_operation(_Dialect, replace=True)
class MakeLoopIndependentOp(MakeLoopIndependentOp):
diff --git a/mlir/python/mlir/dialects/transform/tune.py b/mlir/python/mlir/dialects/transform/tune.py
index f63f88a382422..6eadd9b4e9d5d 100644
--- a/mlir/python/mlir/dialects/transform/tune.py
+++ b/mlir/python/mlir/dialects/transform/tune.py
@@ -25,6 +25,15 @@
from typing import Union
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+ register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+ "mlirTuneRegisterTransformDialectExtension", _get_dialect_registry()
+)
+
@_ods_cext.register_operation(_Dialect, replace=True)
class KnobOp(KnobOp):
diff --git a/mlir/python/mlir/dialects/transform/vector.py b/mlir/python/mlir/dialects/transform/vector.py
index af2435cb26cc4..1d7620188da83 100644
--- a/mlir/python/mlir/dialects/transform/vector.py
+++ b/mlir/python/mlir/dialects/transform/vector.py
@@ -4,3 +4,11 @@
from .._vector_transform_enum_gen import *
from .._vector_transform_ops_gen import *
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+ register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+ "mlirVectorRegisterTransformDialectExtension", _get_dialect_registry()
+)
diff --git a/mlir/test/python/dialects/gpu/dialect.py b/mlir/test/python/dialects/gpu/dialect.py
index 26ee9f34cb332..3517b6519b93e 100644
--- a/mlir/test/python/dialects/gpu/dialect.py
+++ b/mlir/test/python/dialects/gpu/dialect.py
@@ -2,6 +2,7 @@
from mlir.ir import *
import mlir.dialects.gpu as gpu
+import mlir.dialects.nvvm as nvvm
import mlir.dialects.gpu.passes
from mlir.passmanager import *
diff --git a/mlir/test/python/dialects/irdl.py b/mlir/test/python/dialects/irdl.py
index ed62db9b69968..4dbcdd6d0cef2 100644
--- a/mlir/test/python/dialects/irdl.py
+++ b/mlir/test/python/dialects/irdl.py
@@ -2,6 +2,7 @@
from mlir.ir import *
from mlir.dialects.irdl import *
+import mlir.dialects.arith
import sys
diff --git a/mlir/test/python/dialects/memref.py b/mlir/test/python/dialects/memref.py
index b91fdc367cf30..0967e7e5d070c 100644
--- a/mlir/test/python/dialects/memref.py
+++ b/mlir/test/python/dialects/memref.py
@@ -2,6 +2,7 @@
import mlir.dialects.arith as arith
import mlir.dialects.memref as memref
+import mlir.dialects.func
import mlir.extras.types as T
from mlir.dialects.memref import _infer_memref_subview_result_type
from mlir.ir import *
diff --git a/mlir/test/python/integration/dialects/pdl.py b/mlir/test/python/integration/dialects/pdl.py
index c8e6197e03842..5876ecfa41a5a 100644
--- a/mlir/test/python/integration/dialects/pdl.py
+++ b/mlir/test/python/integration/dialects/pdl.py
@@ -1,6 +1,6 @@
# RUN: %PYTHON %s 2>&1 | FileCheck %s
-from mlir.dialects import arith, func, pdl
+from mlir.dialects import arith, func, pdl, irdl
from mlir.dialects.builtin import module
from mlir.ir import *
from mlir.rewrite import *
diff --git a/mlir/test/python/ir/builtin_types.py b/mlir/test/python/ir/builtin_types.py
index b42bfd9bc6587..22a68b5d76b97 100644
--- a/mlir/test/python/ir/builtin_types.py
+++ b/mlir/test/python/ir/builtin_types.py
@@ -2,7 +2,7 @@
import gc
from mlir.ir import *
-from mlir.dialects import arith, tensor, func, memref
+from mlir.dialects import arith, tensor, func, memref, transform, scf
import mlir.extras.types as T
diff --git a/mlir/test/python/ir/capi.py b/mlir/test/python/ir/capi.py
index d60fbd820f91e..934fab91f05a0 100644
--- a/mlir/test/python/ir/capi.py
+++ b/mlir/test/python/ir/capi.py
@@ -1,6 +1,41 @@
# RUN: %PYTHON %s | FileCheck %s
-from mlir._mlir_libs._capi import _capi
+import ctypes
+
+from mlir._mlir_libs import get_dialect_registry
+from mlir._mlir_libs._capi import (
+ _capi,
+ PyCapsule_New,
+ MLIR_PYTHON_CAPSULE_DIALECT_HANDLE,
+)
+from mlir.ir import DialectHandle
print("success")
-# CHECK: success
\ No newline at end of file
+# CHECK: success
+
+
+if not hasattr(_capi, "mlirGetDialectHandle__arith__"):
+ raise Exception("missing API")
+_capi.mlirGetDialectHandle__arith__.argtypes = []
+_capi.mlirGetDialectHandle__arith__.restype = ctypes.c_void_p
+
+if not hasattr(_capi, "mlirGetDialectHandle__quant__"):
+ raise Exception("missing API")
+_capi.mlirGetDialectHandle__quant__.argtypes = []
+_capi.mlirGetDialectHandle__quant__.restype = ctypes.c_void_p
+
+dialect_registry = get_dialect_registry()
+# CHECK: ['builtin']
+print(dialect_registry.dialect_names)
+
+arith_handle = _capi.mlirGetDialectHandle__arith__()
+dialect_registry.insert_dialect(arith_handle)
+# CHECK: ['arith', 'builtin']
+print(dialect_registry.dialect_names)
+
+quant_handle = _capi.mlirGetDialectHandle__quant__()
+capsule = PyCapsule_New(quant_handle, MLIR_PYTHON_CAPSULE_DIALECT_HANDLE, None)
+dialect_handle = DialectHandle._CAPICreate(capsule)
+dialect_registry.insert_dialect(dialect_handle)
+# CHECK: ['arith', 'builtin', 'quant']
+print(dialect_registry.dialect_names)
diff --git a/mlir/test/python/ir/diagnostic_handler.py b/mlir/test/python/ir/diagnostic_handler.py
index 6d273e5092e42..f06090d982476 100644
--- a/mlir/test/python/ir/diagnostic_handler.py
+++ b/mlir/test/python/ir/diagnostic_handler.py
@@ -5,6 +5,7 @@
from mlir._mlir_libs._mlirPythonTestNanobind import (
test_diagnostics_with_errors_and_notes,
)
+import mlir.dialects.arith
def run(f):
diff --git a/mlir/test/python/ir/dialects.py b/mlir/test/python/ir/dialects.py
index 5a2ed684d298b..e4fa28db28272 100644
--- a/mlir/test/python/ir/dialects.py
+++ b/mlir/test/python/ir/dialects.py
@@ -4,6 +4,7 @@
import sys
from mlir.ir import *
from mlir.dialects._ods_common import _cext
+from mlir.dialects import func, arith
def run(f):
diff --git a/mlir/test/python/ir/exception.py b/mlir/test/python/ir/exception.py
index 74085cd349643..bcf95bf43acc3 100644
--- a/mlir/test/python/ir/exception.py
+++ b/mlir/test/python/ir/exception.py
@@ -2,6 +2,7 @@
import gc
from mlir.ir import *
+import mlir.dialects.func
def run(f):
diff --git a/mlir/test/python/ir/insertion_point.py b/mlir/test/python/ir/insertion_point.py
index f48beb25f04b2..74223ce18599e 100644
--- a/mlir/test/python/ir/insertion_point.py
+++ b/mlir/test/python/ir/insertion_point.py
@@ -2,6 +2,7 @@
import gc
from mlir.ir import *
+import mlir.dialects.func
def run(f):
diff --git a/mlir/test/python/ir/module.py b/mlir/test/python/ir/module.py
index 33959bea9ffb6..c9bb8fc309cc2 100644
--- a/mlir/test/python/ir/module.py
+++ b/mlir/test/python/ir/module.py
@@ -3,6 +3,7 @@
import gc
from tempfile import NamedTemporaryFile
from mlir.ir import *
+import mlir.dialects.func
def run(f):
diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
index 4a3625c953d52..2b37f79251af5 100644
--- a/mlir/test/python/ir/operation.py
+++ b/mlir/test/python/ir/operation.py
@@ -6,7 +6,7 @@
from tempfile import NamedTemporaryFile
from mlir.ir import *
from mlir.dialects.builtin import ModuleOp
-from mlir.dialects import arith
+from mlir.dialects import arith, func
from mlir.dialects._ods_common import _cext
diff --git a/mlir/test/python/ir/symbol_table.py b/mlir/test/python/ir/symbol_table.py
index 99d5fadfea10a..81efe3725f911 100644
--- a/mlir/test/python/ir/symbol_table.py
+++ b/mlir/test/python/ir/symbol_table.py
@@ -4,6 +4,7 @@
import io
import itertools
from mlir.ir import *
+import mlir.dialects.func
def run(f):
diff --git a/mlir/test/python/multithreaded_tests.py b/mlir/test/python/multithreaded_tests.py
index 6e1a668346872..513ec90fddfeb 100644
--- a/mlir/test/python/multithreaded_tests.py
+++ b/mlir/test/python/multithreaded_tests.py
@@ -50,7 +50,7 @@
from typing import Optional, List
import mlir.dialects.arith as arith
-from mlir.dialects import transform
+from mlir.dialects import transform, func, memref
from mlir.ir import Context, Location, Module, IntegerType, InsertionPoint
diff --git a/mlir/test/python/python_pass.py b/mlir/test/python/python_pass.py
index 50c42102f66d3..7d47cf3684f5c 100644
--- a/mlir/test/python/python_pass.py
+++ b/mlir/test/python/python_pass.py
@@ -4,7 +4,7 @@
from mlir.ir import *
from mlir.passmanager import *
from mlir.dialects.builtin import ModuleOp
-from mlir.dialects import pdl
+from mlir.dialects import pdl, func, arith
from mlir.rewrite import *
diff --git a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
index 0172b3fa38a6b..02f9434e4129b 100644
--- a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
@@ -44,6 +44,11 @@ from ._ods_common import (
_ods_ir = _ods_cext.ir
_ods_cext.globals.register_traceback_file_exclusion(__file__)
+from .._mlir_libs import get_dialect_registry as _get_dialect_registry
+from .._mlir_libs._capi import register_dialect as _register_dialect
+
+_register_dialect("mlirGetDialectHandle__{0}__", _get_dialect_registry())
+
import builtins
from typing import Sequence as _Sequence, Union as _Union, Optional as _Optional
@@ -1191,7 +1196,7 @@ static bool emitAllOps(const RecordKeeper &records, raw_ostream &os) {
if (clDialectName.empty())
llvm::PrintFatalError("dialect name not provided");
- os << fileHeader;
+ os << formatv(fileHeader, clDialectName.getValue());
if (!clDialectExtensionName.empty())
os << formatv(dialectExtensionTemplate, clDialectName.getValue());
else
More information about the Mlir-commits
mailing list