From 17d76f802f323ce0b909064580921021724e007d Mon Sep 17 00:00:00 2001 From: makslevental Date: Thu, 25 Sep 2025 09:40:53 -0700 Subject: [PATCH 1/4] [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 b02964f047538eca9f1b1a10ee3e6ad6cf20fc2b Mon Sep 17 00:00:00 2001 From: makslevental Date: Thu, 25 Sep 2025 14:59:58 -0700 Subject: [PATCH 2/4] [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 | 36 +++ 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/Linalg.h | 3 + mlir/include/mlir-c/Dialect/MemRef.h | 3 + mlir/include/mlir-c/Dialect/NVGPU.h | 3 + mlir/include/mlir-c/Dialect/PDL.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 | 144 ++++++++++ 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 + .../lib/CAPI/Dialect/TransformInterpreter.cpp | 3 - mlir/lib/CAPI/Dialect/TransformLinalg.cpp | 15 + mlir/lib/CAPI/Dialect/TransformMemRef.cpp | 15 + mlir/lib/CAPI/Dialect/TransformNVGPU.cpp | 15 + mlir/lib/CAPI/Dialect/TransformPDL.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 + .../lib/Dialect/Transform/IR/TransformOps.cpp | 3 +- mlir/python/CMakeLists.txt | 265 +++++++++++++----- mlir/python/mlir/_mlir_libs/_capi.py.in | 44 ++- 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/debug.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/nvgpu.py | 9 + mlir/python/mlir/dialects/transform/pdl.py | 10 + mlir/python/mlir/dialects/transform/smt.py | 10 + .../mlir/dialects/transform/sparse_tensor.py | 8 + .../mlir/dialects/transform/structured.py | 9 + 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/dialects/transform.py | 1 + .../dialects/transform_bufferization_ext.py | 1 + .../python/dialects/transform_debug_ext.py | 1 + mlir/test/python/dialects/transform_extras.py | 3 +- .../test/python/dialects/transform_gpu_ext.py | 1 + .../python/dialects/transform_loop_ext.py | 1 + .../python/dialects/transform_memref_ext.py | 1 + .../python/dialects/transform_nvgpu_ext.py | 1 + .../test/python/dialects/transform_smt_ext.py | 1 + .../dialects/transform_sparse_tensor_ext.py | 1 + .../dialects/transform_structured_ext.py | 1 + .../python/dialects/transform_tensor_ext.py | 1 + .../python/dialects/transform_tune_ext.py | 1 + .../python/dialects/transform_vector_ext.py | 1 + mlir/test/python/integration/dialects/pdl.py | 2 +- .../python/integration/dialects/transform.py | 1 + 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/pass_manager.py | 1 + mlir/test/python/python_pass.py | 2 +- mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp | 9 +- 91 files changed, 1168 insertions(+), 96 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/TransformLinalg.cpp create mode 100644 mlir/lib/CAPI/Dialect/TransformMemRef.cpp create mode 100644 mlir/lib/CAPI/Dialect/TransformNVGPU.cpp create mode 100644 mlir/lib/CAPI/Dialect/TransformPDL.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..41af7a294eb5c --- /dev/null +++ b/mlir/include/mlir-c/Dialect/Bufferization.h @@ -0,0 +1,36 @@ +//===-- mlir-c/Dialect/Bufferization.h - C API 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 +// +//===----------------------------------------------------------------------===// +// +// 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/Linalg.h b/mlir/include/mlir-c/Dialect/Linalg.h index 339e63d667c5e..a3fa3a93a70c7 100644 --- a/mlir/include/mlir-c/Dialect/Linalg.h +++ b/mlir/include/mlir-c/Dialect/Linalg.h @@ -55,6 +55,9 @@ mlirLinalgGetIndexingMapsAttribute(MlirOperation op); MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Linalg, linalg); +MLIR_CAPI_EXPORTED void +mlirLinalgRegisterTransformDialectExtension(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/NVGPU.h b/mlir/include/mlir-c/Dialect/NVGPU.h index e58015a4a3421..89783d3758f9a 100644 --- a/mlir/include/mlir-c/Dialect/NVGPU.h +++ b/mlir/include/mlir-c/Dialect/NVGPU.h @@ -29,6 +29,9 @@ MLIR_CAPI_EXPORTED MlirType mlirNVGPUTensorMapDescriptorTypeGet( MlirContext ctx, MlirType tensorMemrefType, int swizzle, int l2promo, int oobFill, int interleave); +MLIR_CAPI_EXPORTED void +mlirNVGPURegisterTransformDialectExtension(MlirDialectRegistry registry); + #ifdef __cplusplus } #endif diff --git a/mlir/include/mlir-c/Dialect/PDL.h b/mlir/include/mlir-c/Dialect/PDL.h index 6ad2e2da62d87..eb611964e278b 100644 --- a/mlir/include/mlir-c/Dialect/PDL.h +++ b/mlir/include/mlir-c/Dialect/PDL.h @@ -66,6 +66,9 @@ MLIR_CAPI_EXPORTED bool mlirTypeIsAPDLValueType(MlirType type); MLIR_CAPI_EXPORTED MlirType mlirPDLValueTypeGet(MlirContext ctx); +MLIR_CAPI_EXPORTED void +mlirPDLRegisterTransformDialectExtension(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..ed55577996604 --- /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 &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_(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(ptr)}, self.get()); + }) + .def_prop_ro("dialect_names", [](PyDialectRegistry &self) { + int64_t numDialectNames = + mlirDialectRegistryGetNumDialectNames(self.get()); + std::vector 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..8b1321d0c125a 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,94 @@ 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(MLIRCAPILinalgTransformDialectExtension + TransformLinalg.cpp + + PARTIAL_SOURCES_INTENDED + LINK_LIBS PUBLIC + MLIRLinalgTransformOps +) + +add_mlir_upstream_c_api_library(MLIRCAPIMemRefTransformDialectExtension + TransformMemRef.cpp + + PARTIAL_SOURCES_INTENDED + LINK_LIBS PUBLIC + MLIRMemRefTransformOps +) + +add_mlir_upstream_c_api_library(MLIRCAPINVGPUTransformDialectExtension + TransformNVGPU.cpp + + PARTIAL_SOURCES_INTENDED + LINK_LIBS PUBLIC + MLIRNVGPUTransformOps +) + +add_mlir_upstream_c_api_library(MLIRCAPIPDLTransformDialectExtension + TransformPDL.cpp + + PARTIAL_SOURCES_INTENDED + LINK_LIBS PUBLIC + MLIRTransformPDLExtension +) + +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 +404,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(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/TransformInterpreter.cpp b/mlir/lib/CAPI/Dialect/TransformInterpreter.cpp index 145455e1c1b3d..05648de959e7a 100644 --- a/mlir/lib/CAPI/Dialect/TransformInterpreter.cpp +++ b/mlir/lib/CAPI/Dialect/TransformInterpreter.cpp @@ -23,8 +23,6 @@ using namespace mlir; DEFINE_C_API_PTR_METHODS(MlirTransformOptions, transform::TransformOptions) -extern "C" { - MlirTransformOptions mlirTransformOptionsCreate() { return wrap(new transform::TransformOptions); } @@ -80,4 +78,3 @@ MlirLogicalResult mlirMergeSymbolsIntoFromClone(MlirOperation target, unwrap(target), std::move(otherOwning)); return wrap(result); } -} diff --git a/mlir/lib/CAPI/Dialect/TransformLinalg.cpp b/mlir/lib/CAPI/Dialect/TransformLinalg.cpp new file mode 100644 index 0000000000000..5e62a80011a1b --- /dev/null +++ b/mlir/lib/CAPI/Dialect/TransformLinalg.cpp @@ -0,0 +1,15 @@ +//===- TransformLinalg.cpp - C Interface for Transform Linalg 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/Linalg.h" +#include "mlir/CAPI/Registration.h" +#include "mlir/Dialect/Linalg/TransformOps/DialectExtension.h" + +void mlirLinalgRegisterTransformDialectExtension(MlirDialectRegistry registry) { + mlir::linalg::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/TransformNVGPU.cpp b/mlir/lib/CAPI/Dialect/TransformNVGPU.cpp new file mode 100644 index 0000000000000..9159dc319aeb1 --- /dev/null +++ b/mlir/lib/CAPI/Dialect/TransformNVGPU.cpp @@ -0,0 +1,15 @@ +//===- TransformNVGPU.cpp - C Interface for Transform NVGPU 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/NVGPU.h" +#include "mlir/CAPI/Registration.h" +#include "mlir/Dialect/NVGPU/TransformOps/NVGPUTransformOps.h" + +void mlirNVGPURegisterTransformDialectExtension(MlirDialectRegistry registry) { + mlir::nvgpu::registerTransformDialectExtension(*unwrap(registry)); +} diff --git a/mlir/lib/CAPI/Dialect/TransformPDL.cpp b/mlir/lib/CAPI/Dialect/TransformPDL.cpp new file mode 100644 index 0000000000000..550bcabee55ec --- /dev/null +++ b/mlir/lib/CAPI/Dialect/TransformPDL.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/PDL.h" +#include "mlir/CAPI/Registration.h" +#include "mlir/Dialect/Transform/PDLExtension/PDLExtension.h" + +void mlirPDLRegisterTransformDialectExtension(MlirDialectRegistry registry) { + mlir::transform::registerPDLExtension(*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..027d37ea34e79 --- /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/lib/Dialect/Transform/IR/TransformOps.cpp b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp index 132ed815c354e..d8e54bf060bd6 100644 --- a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp +++ b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp @@ -2928,7 +2928,8 @@ LogicalResult transform::SequenceOp::verify() { InFlightDiagnostic diag = emitOpError() << "expected children ops to implement TransformOpInterface"; - diag.attachNote(child.getLoc()) << "op without interface"; + diag.attachNote(child.getLoc()) + << "op without interface: " << child.getName(); return diag; } diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt index 368f344841ed1..082e4ebac43ee 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,29 +192,40 @@ 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 DIALECT_NAME transform - EXTENSION_NAME transform_pdl_extension) + EXTENSION_NAME transform_pdl_extension + EMBED_CAPI_LINK_LIBS + MLIRCAPITransformDialect + MLIRCAPIPDLTransformDialectExtension +) 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 DIALECT_NAME transform - EXTENSION_NAME transform_smt_extension) + EXTENSION_NAME transform_smt_extension + EMBED_CAPI_LINK_LIBS + MLIRCAPITransformDialect + MLIRCAPISMTTransformDialectExtension +) 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 +233,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 +251,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 +277,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 +290,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 +303,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 +315,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 @@ -268,7 +328,11 @@ declare_mlir_dialect_extension_python_bindings( SOURCES dialects/transform/nvgpu.py DIALECT_NAME transform - EXTENSION_NAME nvgpu_transform) + EXTENSION_NAME nvgpu_transform + EMBED_CAPI_LINK_LIBS + MLIRCAPITransformDialect + MLIRCAPINVGPUTransformDialectExtension +) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects @@ -280,6 +344,9 @@ declare_mlir_dialect_extension_python_bindings( EXTENSION_NAME structured_transform GEN_ENUM_BINDINGS_TD_FILE "../../include/mlir/Dialect/Linalg/TransformOps/LinalgTransformEnums.td" + EMBED_CAPI_LINK_LIBS + MLIRCAPITransformDialect + MLIRCAPILinalgTransformDialectExtension ) declare_mlir_dialect_extension_python_bindings( @@ -289,7 +356,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 +369,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 +385,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 +395,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 +410,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 +418,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 +429,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 +441,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 +453,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 +464,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 +488,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 +500,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 +512,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 +523,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 +546,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 +557,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 +578,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 +589,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 +599,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 +610,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 +622,8 @@ declare_mlir_dialect_python_bindings( SOURCES dialects/irdl.py DIALECT_NAME irdl GEN_ENUM_BINDINGS + EMBED_CAPI_LINK_LIBS + MLIRCAPIIRDL ) ################################################################################ @@ -528,32 +663,13 @@ 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 -) - declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind MODULE_NAME _mlirDialectsLinalg ADD_TO_PARENT MLIRPythonSources.Dialects.linalg @@ -749,6 +865,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind MLIRCAPIIR MLIRCAPISMT MLIRCAPIExportSMTLIB + MLIRCAPISMTTransformDialectExtension ) declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses @@ -871,7 +988,6 @@ add_mlir_python_common_capi_library(${MLIR_PYTHON_CAPI_DYLIB_NAME} MLIRPythonCAPI.HeaderSources DECLARED_SOURCES MLIRPythonSources - MLIRPythonExtension.RegisterEverything ${_ADDL_TEST_SOURCES} ) @@ -980,7 +1096,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..2313ecc4f9823 100644 --- a/mlir/python/mlir/_mlir_libs/_capi.py.in +++ b/mlir/python/mlir/_mlir_libs/_capi.py.in @@ -5,4 +5,46 @@ 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@.ir.DialectHandle._CAPIPtr" +).encode() + +MLIR_PYTHON_CAPSULE_DIALECT_REGISTRY = ( + "@MLIR_PYTHON_PACKAGE_PREFIX@.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] + dialect_registry_ptr = PyCapsule_GetPointer( + dialect_registry._CAPIPtr, MLIR_PYTHON_CAPSULE_DIALECT_REGISTRY + ) + registration_capi(dialect_registry_ptr) 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/debug.py b/mlir/python/mlir/dialects/transform/debug.py index f7c04268dc03d..353e4f07c883a 100644 --- a/mlir/python/mlir/dialects/transform/debug.py +++ b/mlir/python/mlir/dialects/transform/debug.py @@ -15,6 +15,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( + "mlirDebugRegisterTransformDialectExtension", _get_dialect_registry() +) + @_ods_cext.register_operation(_Dialect, replace=True) class EmitParamAsRemarkOp(EmitParamAsRemarkOp): 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/nvgpu.py b/mlir/python/mlir/dialects/transform/nvgpu.py index 74ba4c9aeb6c1..3d6427d1b1ed9 100644 --- a/mlir/python/mlir/dialects/transform/nvgpu.py +++ b/mlir/python/mlir/dialects/transform/nvgpu.py @@ -3,3 +3,12 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception from .._nvgpu_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( + "mlirNVGPURegisterTransformDialectExtension", _get_dialect_registry() +) diff --git a/mlir/python/mlir/dialects/transform/pdl.py b/mlir/python/mlir/dialects/transform/pdl.py index bb5fa7ffd3065..15b0c37837902 100644 --- a/mlir/python/mlir/dialects/transform/pdl.py +++ b/mlir/python/mlir/dialects/transform/pdl.py @@ -18,6 +18,16 @@ 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( + "mlirPDLRegisterTransformDialectExtension", _get_dialect_registry() +) + + @_ods_cext.register_operation(_Dialect, replace=True) class PDLMatchOp(PDLMatchOp): def __init__( 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/structured.py b/mlir/python/mlir/dialects/transform/structured.py index bf40cc532065d..85833d32ba55e 100644 --- a/mlir/python/mlir/dialects/transform/structured.py +++ b/mlir/python/mlir/dialects/transform/structured.py @@ -28,6 +28,15 @@ from typing import List, 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( + "mlirLinalgRegisterTransformDialectExtension", _get_dialect_registry() +) + @_ods_cext.register_operation(_Dialect, replace=True) class BufferizeToAllocationOp(BufferizeToAllocationOp): 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/dialects/transform.py b/mlir/test/python/dialects/transform.py index 6c5e4e5505b1c..3477539b8d2f0 100644 --- a/mlir/test/python/dialects/transform.py +++ b/mlir/test/python/dialects/transform.py @@ -12,6 +12,7 @@ def run(f): print("\nTEST:", f.__name__) f(module) print(module) + module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_bufferization_ext.py b/mlir/test/python/dialects/transform_bufferization_ext.py index 03b16c324a407..39014735aa8fb 100644 --- a/mlir/test/python/dialects/transform_bufferization_ext.py +++ b/mlir/test/python/dialects/transform_bufferization_ext.py @@ -13,6 +13,7 @@ def run(f): print("\nTEST:", f.__name__) f() print(module) + module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_debug_ext.py b/mlir/test/python/dialects/transform_debug_ext.py index 2dfdaed343865..0826ed14d199b 100644 --- a/mlir/test/python/dialects/transform_debug_ext.py +++ b/mlir/test/python/dialects/transform_debug_ext.py @@ -19,6 +19,7 @@ def run(f): f(sequence.bodyTarget) transform.YieldOp() print(module) + module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_extras.py b/mlir/test/python/dialects/transform_extras.py index ea47f170cb632..3aeac681f8d49 100644 --- a/mlir/test/python/dialects/transform_extras.py +++ b/mlir/test/python/dialects/transform_extras.py @@ -2,7 +2,7 @@ from typing import Callable from mlir import ir -from mlir.dialects import scf, pdl +from mlir.dialects import scf, pdl, arith from mlir.dialects.transform import ( structured, get_parent_op, @@ -30,6 +30,7 @@ def construct_and_print_in_module(f): with ir.InsertionPoint(module.body): f() print(module) + module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_gpu_ext.py b/mlir/test/python/dialects/transform_gpu_ext.py index 17f5a6d38f8c1..5771ebd13ab77 100644 --- a/mlir/test/python/dialects/transform_gpu_ext.py +++ b/mlir/test/python/dialects/transform_gpu_ext.py @@ -19,6 +19,7 @@ def run(f): transform.YieldOp() print("\nTEST:", f.__name__) print(module) + module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_loop_ext.py b/mlir/test/python/dialects/transform_loop_ext.py index 430b33fba04c7..50c93fe9f46d7 100644 --- a/mlir/test/python/dialects/transform_loop_ext.py +++ b/mlir/test/python/dialects/transform_loop_ext.py @@ -13,6 +13,7 @@ def run(f): print("\nTEST:", f.__name__) f() print(module) + module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_memref_ext.py b/mlir/test/python/dialects/transform_memref_ext.py index e7d871c9eac8c..55eceb57f97c8 100644 --- a/mlir/test/python/dialects/transform_memref_ext.py +++ b/mlir/test/python/dialects/transform_memref_ext.py @@ -13,6 +13,7 @@ def run(f): print("\nTEST:", f.__name__) f() print(module) + module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_nvgpu_ext.py b/mlir/test/python/dialects/transform_nvgpu_ext.py index 1de8b25bab7a5..45672ff3b58e4 100644 --- a/mlir/test/python/dialects/transform_nvgpu_ext.py +++ b/mlir/test/python/dialects/transform_nvgpu_ext.py @@ -12,6 +12,7 @@ def run(f): print("\nTEST:", f.__name__) f() print(module) + module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_smt_ext.py b/mlir/test/python/dialects/transform_smt_ext.py index 3692fd92344a6..0593421c7713d 100644 --- a/mlir/test/python/dialects/transform_smt_ext.py +++ b/mlir/test/python/dialects/transform_smt_ext.py @@ -19,6 +19,7 @@ def run(f): f(sequence.bodyTarget) transform.YieldOp() print(module) + module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_sparse_tensor_ext.py b/mlir/test/python/dialects/transform_sparse_tensor_ext.py index e11cc6bf1e074..b22f27ab6494b 100644 --- a/mlir/test/python/dialects/transform_sparse_tensor_ext.py +++ b/mlir/test/python/dialects/transform_sparse_tensor_ext.py @@ -19,6 +19,7 @@ def run(f): transform.YieldOp() print("\nTEST:", f.__name__) print(module) + module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_structured_ext.py b/mlir/test/python/dialects/transform_structured_ext.py index 8785d6d360074..ee15ac43e3d1a 100644 --- a/mlir/test/python/dialects/transform_structured_ext.py +++ b/mlir/test/python/dialects/transform_structured_ext.py @@ -19,6 +19,7 @@ def run(f): f() module.operation.verify() print(module) + module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_tensor_ext.py b/mlir/test/python/dialects/transform_tensor_ext.py index a2e7aa242b9da..9dde72f52a8fc 100644 --- a/mlir/test/python/dialects/transform_tensor_ext.py +++ b/mlir/test/python/dialects/transform_tensor_ext.py @@ -19,6 +19,7 @@ def run(f): f(sequence.bodyTarget) transform.YieldOp() print(module) + module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_tune_ext.py b/mlir/test/python/dialects/transform_tune_ext.py index dfb93594bca52..f8db151e10703 100644 --- a/mlir/test/python/dialects/transform_tune_ext.py +++ b/mlir/test/python/dialects/transform_tune_ext.py @@ -19,6 +19,7 @@ def run(f): f(sequence.bodyTarget) transform.YieldOp() print(module) + module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_vector_ext.py b/mlir/test/python/dialects/transform_vector_ext.py index 0cd9333dc1218..0e841f2895b55 100644 --- a/mlir/test/python/dialects/transform_vector_ext.py +++ b/mlir/test/python/dialects/transform_vector_ext.py @@ -21,6 +21,7 @@ def run_apply_patterns(f): transform.YieldOp() print("\nTEST:", f.__name__) print(module) + module.operation.verify() return f 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/integration/dialects/transform.py b/mlir/test/python/integration/dialects/transform.py index 303274a8f8828..f57eee7d81819 100644 --- a/mlir/test/python/integration/dialects/transform.py +++ b/mlir/test/python/integration/dialects/transform.py @@ -24,6 +24,7 @@ def construct_and_print_in_module(f): module = f(module) if module is not None: print(module) + module.operation.verify() return f 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..5ac077830cf02 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, scf, cf 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/pass_manager.py b/mlir/test/python/pass_manager.py index 5f92f5b52a09a..db2915246d30f 100644 --- a/mlir/test/python/pass_manager.py +++ b/mlir/test/python/pass_manager.py @@ -5,6 +5,7 @@ from mlir.passmanager import * from mlir.dialects.func import FuncOp from mlir.dialects.builtin import ModuleOp +import mlir.dialects.arith # Log everything to stderr and flush so that we have a unified stream to match 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..6c98d3ed8d0a9 100644 --- a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp +++ b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp @@ -44,6 +44,13 @@ 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 + +_dialect_registry = _get_dialect_registry() +if "{0}" not in _dialect_registry.dialect_names: + _register_dialect("mlirGetDialectHandle__{0}__", _dialect_registry) + import builtins from typing import Sequence as _Sequence, Union as _Union, Optional as _Optional @@ -1191,7 +1198,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 From e85d9c80c4dbc77b22a98f6d3ace5af36c35aacc Mon Sep 17 00:00:00 2001 From: Maksim Levental Date: Sat, 27 Sep 2025 03:58:37 -0400 Subject: [PATCH 3/4] factor registereverything --- mlir/include/mlir-c/RegisterAllExtensions.h | 26 ++++++++++++++++ .../mlir-c/RegisterAllLLVMTranslations.h | 26 ++++++++++++++++ mlir/include/mlir-c/RegisterAllPasses.h | 26 ++++++++++++++++ mlir/include/mlir-c/RegisterEverything.h | 4 +-- .../Bindings/Python/RegisterEverything.cpp | 11 ++++--- .../CAPI/RegisterEverything/CMakeLists.txt | 30 +++++++++++++++++++ .../RegisterAllExtensions.cpp | 16 ++++++++++ .../RegisterAllLLVMTranslations.cpp | 20 +++++++++++++ .../RegisterEverything/RegisterAllPasses.cpp | 12 ++++++++ mlir/python/CMakeLists.txt | 25 ++++++++++++++++ mlir/python/mlir/_mlir_libs/__init__.py | 3 ++ 11 files changed, 193 insertions(+), 6 deletions(-) create mode 100644 mlir/include/mlir-c/RegisterAllExtensions.h create mode 100644 mlir/include/mlir-c/RegisterAllLLVMTranslations.h create mode 100644 mlir/include/mlir-c/RegisterAllPasses.h create mode 100644 mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp create mode 100644 mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp create mode 100644 mlir/lib/CAPI/RegisterEverything/RegisterAllPasses.cpp diff --git a/mlir/include/mlir-c/RegisterAllExtensions.h b/mlir/include/mlir-c/RegisterAllExtensions.h new file mode 100644 index 0000000000000..7238334ed5b8f --- /dev/null +++ b/mlir/include/mlir-c/RegisterAllExtensions.h @@ -0,0 +1,26 @@ +//===-- mlir-c/RegisterAllExtensions.h - Register all MLIR Extensions ---*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM +// Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_C_REGISTER_Extensions_H +#define MLIR_C_REGISTER_Extensions_H + +#include "mlir-c/IR.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/// Register all compiler Extensions of MLIR. +MLIR_CAPI_EXPORTED void mlirRegisterAllExtensions(MlirDialectRegistry registry); + +#ifdef __cplusplus +} +#endif + +#endif // MLIR_C_REGISTER_Extensions_H diff --git a/mlir/include/mlir-c/RegisterAllLLVMTranslations.h b/mlir/include/mlir-c/RegisterAllLLVMTranslations.h new file mode 100644 index 0000000000000..b2be041e8a93a --- /dev/null +++ b/mlir/include/mlir-c/RegisterAllLLVMTranslations.h @@ -0,0 +1,26 @@ +//===-- mlir-c/RegisterAllLLVMTranslations.h - Register all LLVM translations // +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM +// Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_C_REGISTER_LLVM_TRANSLATIONS_H +#define MLIR_C_REGISTER_LLVM_TRANSLATIONS_H + +#include "mlir-c/IR.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/// Register all LLVM translations of MLIR. +MLIR_CAPI_EXPORTED void mlirRegisterAllLLVMTranslations(MlirContext context); + +#ifdef __cplusplus +} +#endif + +#endif // MLIR_C_REGISTER_LLVM_TRANSLATIONS_H diff --git a/mlir/include/mlir-c/RegisterAllPasses.h b/mlir/include/mlir-c/RegisterAllPasses.h new file mode 100644 index 0000000000000..797b865885682 --- /dev/null +++ b/mlir/include/mlir-c/RegisterAllPasses.h @@ -0,0 +1,26 @@ +//===-- mlir-c/RegisterAllPasses.h - Register all MLIR Pass --*- 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 +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_C_REGISTER_PASSES_H +#define MLIR_C_REGISTER_PASSES_H + +#include "mlir-c/IR.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/// Register all compiler passes of MLIR. +MLIR_CAPI_EXPORTED void mlirRegisterAllPasses(void); + +#ifdef __cplusplus +} +#endif + +#endif // MLIR_C_REGISTER_PASSES_H diff --git a/mlir/include/mlir-c/RegisterEverything.h b/mlir/include/mlir-c/RegisterEverything.h index ea2ea86449727..e759f6b188778 100644 --- a/mlir/include/mlir-c/RegisterEverything.h +++ b/mlir/include/mlir-c/RegisterEverything.h @@ -28,8 +28,8 @@ MLIR_CAPI_EXPORTED void mlirRegisterAllDialects(MlirDialectRegistry registry); /// Register all translations to LLVM IR for dialects that can support it. MLIR_CAPI_EXPORTED void mlirRegisterAllLLVMTranslations(MlirContext context); -/// Register all compiler passes of MLIR. -MLIR_CAPI_EXPORTED void mlirRegisterAllPasses(void); +// /// Register all compiler passes of MLIR. +// MLIR_CAPI_EXPORTED void mlirRegisterAllPasses(void); #ifdef __cplusplus } diff --git a/mlir/lib/Bindings/Python/RegisterEverything.cpp b/mlir/lib/Bindings/Python/RegisterEverything.cpp index 3edcb099c0a24..d773c15fbfc30 100644 --- a/mlir/lib/Bindings/Python/RegisterEverything.cpp +++ b/mlir/lib/Bindings/Python/RegisterEverything.cpp @@ -6,15 +6,18 @@ // //===----------------------------------------------------------------------===// -#include "mlir-c/RegisterEverything.h" +#include "mlir-c/RegisterAllExtensions.h" +#include "mlir-c/RegisterAllLLVMTranslations.h" +#include "mlir-c/RegisterAllPasses.h" #include "mlir/Bindings/Python/Nanobind.h" #include "mlir/Bindings/Python/NanobindAdaptors.h" NB_MODULE(_mlirRegisterEverything, m) { - m.doc() = "MLIR All Upstream Dialects, Translations and Passes Registration"; + m.doc() = + "MLIR All Upstream Extensions, Translations and Passes Registration"; - m.def("register_dialects", [](MlirDialectRegistry registry) { - mlirRegisterAllDialects(registry); + m.def("register_extensions", [](MlirDialectRegistry registry) { + mlirRegisterAllExtensions(registry); }); m.def("register_llvm_translations", [](MlirContext context) { mlirRegisterAllLLVMTranslations(context); }); diff --git a/mlir/lib/CAPI/RegisterEverything/CMakeLists.txt b/mlir/lib/CAPI/RegisterEverything/CMakeLists.txt index ccda668ece3d4..2ac01bc4c84f3 100644 --- a/mlir/lib/CAPI/RegisterEverything/CMakeLists.txt +++ b/mlir/lib/CAPI/RegisterEverything/CMakeLists.txt @@ -2,6 +2,7 @@ get_property(translation_libs GLOBAL PROPERTY MLIR_TRANSLATION_LIBS) add_mlir_upstream_c_api_library(MLIRCAPIRegisterEverything RegisterEverything.cpp + PARTIAL_SOURCES_INTENDED LINK_LIBS PUBLIC ${translation_libs} @@ -14,3 +15,32 @@ add_mlir_upstream_c_api_library(MLIRCAPIRegisterEverything MLIRRegisterAllExtensions MLIRRegisterAllPasses ) + +add_mlir_upstream_c_api_library(MLIRCAPIRegisterAllPasses + RegisterAllPasses.cpp + PARTIAL_SOURCES_INTENDED + + LINK_LIBS PUBLIC + MLIRCAPIIR + MLIRRegisterAllPasses +) + +add_mlir_upstream_c_api_library(MLIRCAPIRegisterAllLLVMTranslations + RegisterAllLLVMTranslations.cpp + PARTIAL_SOURCES_INTENDED + + LINK_LIBS PUBLIC + MLIRCAPIIR + ${translation_libs} + MLIRBuiltinToLLVMIRTranslation + MLIRLLVMToLLVMIRTranslation +) + +add_mlir_upstream_c_api_library(MLIRCAPIRegisterAllExtensions + RegisterAllExtensions.cpp + PARTIAL_SOURCES_INTENDED + + LINK_LIBS PUBLIC + MLIRCAPIIR + MLIRRegisterAllExtensions +) \ No newline at end of file diff --git a/mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp b/mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp new file mode 100644 index 0000000000000..9cc9fd394e728 --- /dev/null +++ b/mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp @@ -0,0 +1,16 @@ +//===- RegisterAllExtensions.cpp - Register all MLIR entities +//-----------------===// +// +// 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/RegisterAllExtensions.h" +#include "mlir/CAPI/IR.h" +#include "mlir/InitAllExtensions.h" + +void mlirRegisterAllExtensions(MlirDialectRegistry registry) { + mlir::registerAllExtensions(*unwrap(registry)); +} diff --git a/mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp b/mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp new file mode 100644 index 0000000000000..c5dc64ca17834 --- /dev/null +++ b/mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp @@ -0,0 +1,20 @@ +//===- RegisterAllPasses.cpp - Register all MLIR entities -----------------===// +// +// 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/RegisterAllLLVMTranslations.h" +#include "mlir/CAPI/IR.h" +#include "mlir/Target/LLVMIR/Dialect/All.h" +#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h" +#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h" + +void mlirRegisterAllLLVMTranslations(MlirContext context) { + auto &ctx = *unwrap(context); + mlir::DialectRegistry registry; + mlir::registerAllToLLVMIRTranslations(registry); + ctx.appendDialectRegistry(registry); +} diff --git a/mlir/lib/CAPI/RegisterEverything/RegisterAllPasses.cpp b/mlir/lib/CAPI/RegisterEverything/RegisterAllPasses.cpp new file mode 100644 index 0000000000000..159dff0b86e44 --- /dev/null +++ b/mlir/lib/CAPI/RegisterEverything/RegisterAllPasses.cpp @@ -0,0 +1,12 @@ +//===- RegisterAllPasses.cpp - Register all MLIR entities -----------------===// +// +// 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/RegisterAllPasses.h" +#include "mlir/InitAllPasses.h" + +void mlirRegisterAllPasses() { mlir::registerAllPasses(); } diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt index 082e4ebac43ee..9a45c5026f60e 100644 --- a/mlir/python/CMakeLists.txt +++ b/mlir/python/CMakeLists.txt @@ -670,6 +670,29 @@ declare_mlir_python_extension(MLIRPythonExtension.Core 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 + MLIRCAPIRegisterAllPasses + MLIRCAPIRegisterAllExtensions + MLIRCAPIRegisterAllLLVMTranslations +) + declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind MODULE_NAME _mlirDialectsLinalg ADD_TO_PARENT MLIRPythonSources.Dialects.linalg @@ -988,6 +1011,7 @@ add_mlir_python_common_capi_library(${MLIR_PYTHON_CAPI_DYLIB_NAME} MLIRPythonCAPI.HeaderSources DECLARED_SOURCES MLIRPythonSources + MLIRPythonExtension.RegisterEverything ${_ADDL_TEST_SOURCES} ) @@ -1096,6 +1120,7 @@ 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/__init__.py b/mlir/python/mlir/_mlir_libs/__init__.py index 63244212ba42c..2c36397d2cf8b 100644 --- a/mlir/python/mlir/_mlir_libs/__init__.py +++ b/mlir/python/mlir/_mlir_libs/__init__.py @@ -123,6 +123,9 @@ def process_initializer_module(module_name): if hasattr(m, "register_dialects"): logger.debug("Registering dialects from initializer %r", m) m.register_dialects(get_dialect_registry()) + if hasattr(m, "register_extensions"): + logger.debug("Registering extensions from initializer %r", m) + m.register_extensions(get_dialect_registry()) if hasattr(m, "context_init_hook"): logger.debug("Adding context init hook from %r", m) post_init_hooks.append(m.context_init_hook) From 7179ffae6e552b5ff7d9fe7c4ad3f66e8b8b3f33 Mon Sep 17 00:00:00 2001 From: Maksim Levental Date: Sat, 27 Sep 2025 04:06:35 -0400 Subject: [PATCH 4/4] undo transform extension stuff --- mlir/cmake/modules/AddMLIRPython.cmake | 6 +- mlir/examples/standalone/pyproject.toml | 5 + .../examples/standalone/python/CMakeLists.txt | 1 + mlir/include/mlir-c/Dialect/Transform.h | 13 --- mlir/include/mlir-c/RegisterAllExtensions.h | 6 +- .../mlir-c/RegisterAllExternalModels.h | 27 +++++ mlir/include/mlir-c/RegisterEverything.h | 4 +- mlir/lib/Bindings/Python/IRCore.cpp | 8 ++ .../Bindings/Python/RegisterEverything.cpp | 4 + mlir/lib/CAPI/Dialect/CMakeLists.txt | 91 ---------------- 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/TransformLinalg.cpp | 15 --- mlir/lib/CAPI/Dialect/TransformMemRef.cpp | 15 --- mlir/lib/CAPI/Dialect/TransformNVGPU.cpp | 15 --- mlir/lib/CAPI/Dialect/TransformPDL.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 --- .../CAPI/RegisterEverything/CMakeLists.txt | 11 +- .../RegisterAllExtensions.cpp | 10 +- .../RegisterAllExternalModels.cpp | 101 ++++++++++++++++++ .../RegisterAllLLVMTranslations.cpp | 2 +- .../RegisterEverything/RegisterAllPasses.cpp | 2 +- mlir/python/CMakeLists.txt | 79 +++----------- mlir/python/mlir/_mlir_libs/__init__.py | 3 + .../mlir/_mlir_libs/{_capi.py.in => _capi.py} | 50 ++++----- .../mlir/dialects/transform/bufferization.py | 9 -- mlir/python/mlir/dialects/transform/debug.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/nvgpu.py | 9 -- mlir/python/mlir/dialects/transform/pdl.py | 10 -- mlir/python/mlir/dialects/transform/smt.py | 10 -- .../mlir/dialects/transform/sparse_tensor.py | 8 -- .../mlir/dialects/transform/structured.py | 9 -- mlir/python/mlir/dialects/transform/tensor.py | 9 -- mlir/python/mlir/dialects/transform/tune.py | 9 -- mlir/python/mlir/dialects/transform/vector.py | 8 -- .../dialects/gpu/module-to-binary-nvvm.py | 2 +- .../dialects/gpu/module-to-binary-rocdl.py | 2 +- mlir/test/python/dialects/transform.py | 1 - .../dialects/transform_bufferization_ext.py | 1 - .../python/dialects/transform_debug_ext.py | 1 - mlir/test/python/dialects/transform_extras.py | 1 - .../test/python/dialects/transform_gpu_ext.py | 1 - .../python/dialects/transform_loop_ext.py | 1 - .../python/dialects/transform_memref_ext.py | 1 - .../python/dialects/transform_nvgpu_ext.py | 1 - .../test/python/dialects/transform_smt_ext.py | 1 - .../dialects/transform_sparse_tensor_ext.py | 1 - .../dialects/transform_structured_ext.py | 1 - .../python/dialects/transform_tensor_ext.py | 1 - .../python/dialects/transform_tune_ext.py | 1 - .../python/dialects/transform_vector_ext.py | 1 - mlir/test/python/execution_engine.py | 1 + mlir/test/python/global_constructors.py | 1 + 61 files changed, 215 insertions(+), 529 deletions(-) create mode 100644 mlir/include/mlir-c/RegisterAllExternalModels.h delete mode 100644 mlir/lib/CAPI/Dialect/TransformAffine.cpp delete mode 100644 mlir/lib/CAPI/Dialect/TransformBufferization.cpp delete mode 100644 mlir/lib/CAPI/Dialect/TransformGPU.cpp delete mode 100644 mlir/lib/CAPI/Dialect/TransformLinalg.cpp delete mode 100644 mlir/lib/CAPI/Dialect/TransformMemRef.cpp delete mode 100644 mlir/lib/CAPI/Dialect/TransformNVGPU.cpp delete mode 100644 mlir/lib/CAPI/Dialect/TransformPDL.cpp delete mode 100644 mlir/lib/CAPI/Dialect/TransformSMT.cpp delete mode 100644 mlir/lib/CAPI/Dialect/TransformSparseTensor.cpp delete mode 100644 mlir/lib/CAPI/Dialect/TransformTensor.cpp delete mode 100644 mlir/lib/CAPI/Dialect/TransformVector.cpp create mode 100644 mlir/lib/CAPI/RegisterEverything/RegisterAllExternalModels.cpp rename mlir/python/mlir/_mlir_libs/{_capi.py.in => _capi.py} (53%) diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake index ceebc931e60dc..d8b6d493f985c 100644 --- a/mlir/cmake/modules/AddMLIRPython.cmake +++ b/mlir/cmake/modules/AddMLIRPython.cmake @@ -457,14 +457,11 @@ 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;EMBED_CAPI_LINK_LIBS" + "SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE" ${ARGN}) # Source files. set(_extension_target "${ARG_ADD_TO_PARENT}.${ARG_EXTENSION_NAME}") @@ -506,7 +503,6 @@ 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/examples/standalone/pyproject.toml b/mlir/examples/standalone/pyproject.toml index 5a1e6e86513c3..90535030df817 100644 --- a/mlir/examples/standalone/pyproject.toml +++ b/mlir/examples/standalone/pyproject.toml @@ -37,6 +37,8 @@ cmake.source-dir = "." # This is for installing/distributing the python bindings target and only the python bindings target. build.targets = ["StandalonePythonModules"] install.components = ["StandalonePythonModules"] +# The default is true but make it explicit to highlight that this option exists (turn off for debug symbols). +install.strip = true [tool.scikit-build.cmake.define] # Optional @@ -51,6 +53,9 @@ LLVM_USE_LINKER = { env = "LLVM_USE_LINKER", default = "" } CMAKE_VISIBILITY_INLINES_HIDDEN = "ON" CMAKE_C_VISIBILITY_PRESET = "hidden" CMAKE_CXX_VISIBILITY_PRESET = "hidden" +# Disables generation of "version soname" (i.e. libFoo.so.), +# which causes pure duplication of various shlibs for Python wheels. +CMAKE_PLATFORM_NO_VERSIONED_SONAME = "ON" # Non-optional (alternatively you could use CMAKE_PREFIX_PATH here). MLIR_DIR = { env = "MLIR_DIR", default = "" } diff --git a/mlir/examples/standalone/python/CMakeLists.txt b/mlir/examples/standalone/python/CMakeLists.txt index 905c944939756..505add5232087 100644 --- a/mlir/examples/standalone/python/CMakeLists.txt +++ b/mlir/examples/standalone/python/CMakeLists.txt @@ -138,6 +138,7 @@ set(_declared_sources StandalonePythonSources MLIRPythonSources.Core MLIRPythonSources.Dialects.builtin + MLIRPythonSources.Dialects.arith ) # For an external projects build, the MLIRPythonExtension.Core.type_stub_gen # target already exists and can just be added to DECLARED_SOURCES. diff --git a/mlir/include/mlir-c/Dialect/Transform.h b/mlir/include/mlir-c/Dialect/Transform.h index 1b5232c106edd..02c99b5921882 100644 --- a/mlir/include/mlir-c/Dialect/Transform.h +++ b/mlir/include/mlir-c/Dialect/Transform.h @@ -76,19 +76,6 @@ 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/RegisterAllExtensions.h b/mlir/include/mlir-c/RegisterAllExtensions.h index 7238334ed5b8f..dc2dcfbb28098 100644 --- a/mlir/include/mlir-c/RegisterAllExtensions.h +++ b/mlir/include/mlir-c/RegisterAllExtensions.h @@ -7,8 +7,8 @@ // //===----------------------------------------------------------------------===// -#ifndef MLIR_C_REGISTER_Extensions_H -#define MLIR_C_REGISTER_Extensions_H +#ifndef MLIR_C_REGISTER_EXTENSIONS_H +#define MLIR_C_REGISTER_EXTENSIONS_H #include "mlir-c/IR.h" @@ -23,4 +23,4 @@ MLIR_CAPI_EXPORTED void mlirRegisterAllExtensions(MlirDialectRegistry registry); } #endif -#endif // MLIR_C_REGISTER_Extensions_H +#endif // MLIR_C_REGISTER_EXTENSIONS_H diff --git a/mlir/include/mlir-c/RegisterAllExternalModels.h b/mlir/include/mlir-c/RegisterAllExternalModels.h new file mode 100644 index 0000000000000..baa8a742cd019 --- /dev/null +++ b/mlir/include/mlir-c/RegisterAllExternalModels.h @@ -0,0 +1,27 @@ +//===-- mlir-c/RegisterAllExternalModels.h - Register all MLIR model ----*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM +// Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_C_REGISTER_EXTERNAL_MODELS_H +#define MLIR_C_REGISTER_EXTERNAL_MODELS_H + +#include "mlir-c/IR.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/// Register all compiler External models of MLIR. +MLIR_CAPI_EXPORTED void +mlirRegisterAllExternalModels(MlirDialectRegistry registry); + +#ifdef __cplusplus +} +#endif + +#endif // MLIR_C_REGISTER_EXTERNAL_MODELS_H diff --git a/mlir/include/mlir-c/RegisterEverything.h b/mlir/include/mlir-c/RegisterEverything.h index e759f6b188778..ea2ea86449727 100644 --- a/mlir/include/mlir-c/RegisterEverything.h +++ b/mlir/include/mlir-c/RegisterEverything.h @@ -28,8 +28,8 @@ MLIR_CAPI_EXPORTED void mlirRegisterAllDialects(MlirDialectRegistry registry); /// Register all translations to LLVM IR for dialects that can support it. MLIR_CAPI_EXPORTED void mlirRegisterAllLLVMTranslations(MlirContext context); -// /// Register all compiler passes of MLIR. -// MLIR_CAPI_EXPORTED void mlirRegisterAllPasses(void); +/// Register all compiler passes of MLIR. +MLIR_CAPI_EXPORTED void mlirRegisterAllPasses(void); #ifdef __cplusplus } diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp index c3d28733c2cee..21e60d5550a51 100644 --- a/mlir/lib/Bindings/Python/IRCore.cpp +++ b/mlir/lib/Bindings/Python/IRCore.cpp @@ -3139,6 +3139,10 @@ void mlir::python::populateIRCore(nb::module_ &m) { //---------------------------------------------------------------------------- nb::class_(m, "DialectHandle") + .def_prop_ro_static( + "_capsule_name", + [](nb::handle &) { return MLIR_PYTHON_CAPSULE_DIALECT_HANDLE; }, + nb::sig("def _capsule_name(/) -> str")) .def_static(MLIR_PYTHON_CAPI_FACTORY_ATTR, &createMlirDialectHandleFromCapsule); @@ -3147,6 +3151,10 @@ void mlir::python::populateIRCore(nb::module_ &m) { //---------------------------------------------------------------------------- nb::class_(m, "DialectRegistry") .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyDialectRegistry::getCapsule) + .def_prop_ro_static( + "_capsule_name", + [](nb::handle &) { return MLIR_PYTHON_CAPSULE_DIALECT_REGISTRY; }, + nb::sig("def _capsule_name(/) -> str")) .def_static(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyDialectRegistry::createFromCapsule) .def(nb::init<>()) diff --git a/mlir/lib/Bindings/Python/RegisterEverything.cpp b/mlir/lib/Bindings/Python/RegisterEverything.cpp index d773c15fbfc30..e825bceace87c 100644 --- a/mlir/lib/Bindings/Python/RegisterEverything.cpp +++ b/mlir/lib/Bindings/Python/RegisterEverything.cpp @@ -7,6 +7,7 @@ //===----------------------------------------------------------------------===// #include "mlir-c/RegisterAllExtensions.h" +#include "mlir-c/RegisterAllExternalModels.h" #include "mlir-c/RegisterAllLLVMTranslations.h" #include "mlir-c/RegisterAllPasses.h" #include "mlir/Bindings/Python/Nanobind.h" @@ -16,6 +17,9 @@ NB_MODULE(_mlirRegisterEverything, m) { m.doc() = "MLIR All Upstream Extensions, Translations and Passes Registration"; + m.def("register_external_models", [](MlirDialectRegistry registry) { + mlirRegisterAllExternalModels(registry); + }); m.def("register_extensions", [](MlirDialectRegistry registry) { mlirRegisterAllExtensions(registry); }); diff --git a/mlir/lib/CAPI/Dialect/CMakeLists.txt b/mlir/lib/CAPI/Dialect/CMakeLists.txt index 8b1321d0c125a..1deafe41a4b7a 100644 --- a/mlir/lib/CAPI/Dialect/CMakeLists.txt +++ b/mlir/lib/CAPI/Dialect/CMakeLists.txt @@ -258,9 +258,6 @@ add_mlir_upstream_c_api_library(MLIRCAPITransformDialect LINK_LIBS PUBLIC MLIRCAPIIR MLIRTransformDialect - MLIRTransformDebugExtension - MLIRTransformLoopExtension - MLIRTransformTuneExtension ) add_mlir_upstream_c_api_library(MLIRCAPITransformDialectTransforms @@ -272,94 +269,6 @@ 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(MLIRCAPILinalgTransformDialectExtension - TransformLinalg.cpp - - PARTIAL_SOURCES_INTENDED - LINK_LIBS PUBLIC - MLIRLinalgTransformOps -) - -add_mlir_upstream_c_api_library(MLIRCAPIMemRefTransformDialectExtension - TransformMemRef.cpp - - PARTIAL_SOURCES_INTENDED - LINK_LIBS PUBLIC - MLIRMemRefTransformOps -) - -add_mlir_upstream_c_api_library(MLIRCAPINVGPUTransformDialectExtension - TransformNVGPU.cpp - - PARTIAL_SOURCES_INTENDED - LINK_LIBS PUBLIC - MLIRNVGPUTransformOps -) - -add_mlir_upstream_c_api_library(MLIRCAPIPDLTransformDialectExtension - TransformPDL.cpp - - PARTIAL_SOURCES_INTENDED - LINK_LIBS PUBLIC - MLIRTransformPDLExtension -) - -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 diff --git a/mlir/lib/CAPI/Dialect/Transform.cpp b/mlir/lib/CAPI/Dialect/Transform.cpp index 8942d0a00cde3..5fd773572bd3c 100644 --- a/mlir/lib/CAPI/Dialect/Transform.cpp +++ b/mlir/lib/CAPI/Dialect/Transform.cpp @@ -9,11 +9,8 @@ #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; @@ -109,15 +106,3 @@ MlirType mlirTransformParamTypeGet(MlirContext ctx, MlirType type) { MlirType mlirTransformParamTypeGetType(MlirType type) { return wrap(cast(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 deleted file mode 100644 index 6fd7744e8c2c9..0000000000000 --- a/mlir/lib/CAPI/Dialect/TransformAffine.cpp +++ /dev/null @@ -1,16 +0,0 @@ -//===- 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 deleted file mode 100644 index f1bd61be10880..0000000000000 --- a/mlir/lib/CAPI/Dialect/TransformBufferization.cpp +++ /dev/null @@ -1,16 +0,0 @@ -//==- 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 deleted file mode 100644 index 51187a28fa752..0000000000000 --- a/mlir/lib/CAPI/Dialect/TransformGPU.cpp +++ /dev/null @@ -1,15 +0,0 @@ -//===- 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/TransformLinalg.cpp b/mlir/lib/CAPI/Dialect/TransformLinalg.cpp deleted file mode 100644 index 5e62a80011a1b..0000000000000 --- a/mlir/lib/CAPI/Dialect/TransformLinalg.cpp +++ /dev/null @@ -1,15 +0,0 @@ -//===- TransformLinalg.cpp - C Interface for Transform Linalg 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/Linalg.h" -#include "mlir/CAPI/Registration.h" -#include "mlir/Dialect/Linalg/TransformOps/DialectExtension.h" - -void mlirLinalgRegisterTransformDialectExtension(MlirDialectRegistry registry) { - mlir::linalg::registerTransformDialectExtension(*unwrap(registry)); -} diff --git a/mlir/lib/CAPI/Dialect/TransformMemRef.cpp b/mlir/lib/CAPI/Dialect/TransformMemRef.cpp deleted file mode 100644 index 9b80295748514..0000000000000 --- a/mlir/lib/CAPI/Dialect/TransformMemRef.cpp +++ /dev/null @@ -1,15 +0,0 @@ -//===- 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/TransformNVGPU.cpp b/mlir/lib/CAPI/Dialect/TransformNVGPU.cpp deleted file mode 100644 index 9159dc319aeb1..0000000000000 --- a/mlir/lib/CAPI/Dialect/TransformNVGPU.cpp +++ /dev/null @@ -1,15 +0,0 @@ -//===- TransformNVGPU.cpp - C Interface for Transform NVGPU 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/NVGPU.h" -#include "mlir/CAPI/Registration.h" -#include "mlir/Dialect/NVGPU/TransformOps/NVGPUTransformOps.h" - -void mlirNVGPURegisterTransformDialectExtension(MlirDialectRegistry registry) { - mlir::nvgpu::registerTransformDialectExtension(*unwrap(registry)); -} diff --git a/mlir/lib/CAPI/Dialect/TransformPDL.cpp b/mlir/lib/CAPI/Dialect/TransformPDL.cpp deleted file mode 100644 index 550bcabee55ec..0000000000000 --- a/mlir/lib/CAPI/Dialect/TransformPDL.cpp +++ /dev/null @@ -1,15 +0,0 @@ -//===- 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/PDL.h" -#include "mlir/CAPI/Registration.h" -#include "mlir/Dialect/Transform/PDLExtension/PDLExtension.h" - -void mlirPDLRegisterTransformDialectExtension(MlirDialectRegistry registry) { - mlir::transform::registerPDLExtension(*unwrap(registry)); -} diff --git a/mlir/lib/CAPI/Dialect/TransformSMT.cpp b/mlir/lib/CAPI/Dialect/TransformSMT.cpp deleted file mode 100644 index 449129dac99e5..0000000000000 --- a/mlir/lib/CAPI/Dialect/TransformSMT.cpp +++ /dev/null @@ -1,15 +0,0 @@ -//===- 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 deleted file mode 100644 index 3e5b91099a6fb..0000000000000 --- a/mlir/lib/CAPI/Dialect/TransformSparseTensor.cpp +++ /dev/null @@ -1,16 +0,0 @@ -//===- 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 deleted file mode 100644 index fc7a30b5b237e..0000000000000 --- a/mlir/lib/CAPI/Dialect/TransformTensor.cpp +++ /dev/null @@ -1,15 +0,0 @@ -//===- 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 deleted file mode 100644 index 027d37ea34e79..0000000000000 --- a/mlir/lib/CAPI/Dialect/TransformVector.cpp +++ /dev/null @@ -1,15 +0,0 @@ -//===- 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/RegisterEverything/CMakeLists.txt b/mlir/lib/CAPI/RegisterEverything/CMakeLists.txt index 2ac01bc4c84f3..473c3844538c5 100644 --- a/mlir/lib/CAPI/RegisterEverything/CMakeLists.txt +++ b/mlir/lib/CAPI/RegisterEverything/CMakeLists.txt @@ -43,4 +43,13 @@ add_mlir_upstream_c_api_library(MLIRCAPIRegisterAllExtensions LINK_LIBS PUBLIC MLIRCAPIIR MLIRRegisterAllExtensions -) \ No newline at end of file +) + +add_mlir_upstream_c_api_library(MLIRCAPIRegisterAllExternalModels + RegisterAllExternalModels.cpp + PARTIAL_SOURCES_INTENDED + + LINK_LIBS PUBLIC + MLIRCAPIIR + MLIRRegisterAllDialects +) diff --git a/mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp b/mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp index 9cc9fd394e728..9be46bfff012f 100644 --- a/mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp +++ b/mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp @@ -1,5 +1,4 @@ -//===- RegisterAllExtensions.cpp - Register all MLIR entities -//-----------------===// +//===- RegisterAllExtensions.cpp - Register all MLIR extensions -----------===// // // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // See https://llvm.org/LICENSE.txt for license information. @@ -11,6 +10,9 @@ #include "mlir/CAPI/IR.h" #include "mlir/InitAllExtensions.h" -void mlirRegisterAllExtensions(MlirDialectRegistry registry) { - mlir::registerAllExtensions(*unwrap(registry)); +using namespace mlir; + +void mlirRegisterAllExtensions(MlirDialectRegistry mlirRegistry) { + mlir::DialectRegistry *registry = unwrap(mlirRegistry); + mlir::registerAllExtensions(*registry); } diff --git a/mlir/lib/CAPI/RegisterEverything/RegisterAllExternalModels.cpp b/mlir/lib/CAPI/RegisterEverything/RegisterAllExternalModels.cpp new file mode 100644 index 0000000000000..cee3d717ccb33 --- /dev/null +++ b/mlir/lib/CAPI/RegisterEverything/RegisterAllExternalModels.cpp @@ -0,0 +1,101 @@ +//===- RegisterAllExtensions.cpp - Register all MLIR extensions -----------===// +// +// 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/RegisterAllExternalModels.h" +#include "mlir/CAPI/IR.h" +#include "mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h" +#include "mlir/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.h" +#include "mlir/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.h" +#include "mlir/Dialect/Arith/Transforms/BufferViewFlowOpInterfaceImpl.h" +#include "mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h" +#include "mlir/Dialect/Arith/Transforms/ShardingInterfaceImpl.h" +#include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h" +#include "mlir/Dialect/ControlFlow/Transforms/BufferDeallocationOpInterfaceImpl.h" +#include "mlir/Dialect/ControlFlow/Transforms/BufferizableOpInterfaceImpl.h" +#include "mlir/Dialect/GPU/IR/ValueBoundsOpInterfaceImpl.h" +#include "mlir/Dialect/GPU/Transforms/BufferDeallocationOpInterfaceImpl.h" +#include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h" +#include "mlir/Dialect/Linalg/Transforms/AllInterfaces.h" +#include "mlir/Dialect/Linalg/Transforms/RuntimeOpVerification.h" +#include "mlir/Dialect/MLProgram/Transforms/BufferizableOpInterfaceImpl.h" +#include "mlir/Dialect/MemRef/IR/MemRefMemorySlot.h" +#include "mlir/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.h" +#include "mlir/Dialect/MemRef/Transforms/AllocationOpInterfaceImpl.h" +#include "mlir/Dialect/MemRef/Transforms/BufferViewFlowOpInterfaceImpl.h" +#include "mlir/Dialect/MemRef/Transforms/RuntimeOpVerification.h" +#include "mlir/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.h" +#include "mlir/Dialect/SCF/Transforms/BufferDeallocationOpInterfaceImpl.h" +#include "mlir/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.h" +#include "mlir/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.h" +#include "mlir/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.h" +#include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h" +#include "mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h" +#include "mlir/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.h" +#include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h" +#include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h" +#include "mlir/Dialect/Tensor/Transforms/RuntimeOpVerification.h" +#include "mlir/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.h" +#include "mlir/Dialect/Tosa/IR/ShardingInterfaceImpl.h" +#include "mlir/Dialect/Vector/IR/ValueBoundsOpInterfaceImpl.h" +#include "mlir/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.h" +#include "mlir/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.h" +#include "mlir/Interfaces/CastInterfaces.h" +#include "mlir/Target/LLVM/NVVM/Target.h" +#include "mlir/Target/LLVM/ROCDL/Target.h" +#include "mlir/Target/LLVM/XeVM/Target.h" +#include "mlir/Target/SPIRV/Target.h" + +using namespace mlir; + +void mlirRegisterAllExternalModels(MlirDialectRegistry mlirRegistry) { + mlir::DialectRegistry *registry = unwrap(mlirRegistry); + // Register all external models. + affine::registerValueBoundsOpInterfaceExternalModels(*registry); + arith::registerBufferDeallocationOpInterfaceExternalModels(*registry); + arith::registerBufferizableOpInterfaceExternalModels(*registry); + arith::registerBufferViewFlowOpInterfaceExternalModels(*registry); + arith::registerShardingInterfaceExternalModels(*registry); + arith::registerValueBoundsOpInterfaceExternalModels(*registry); + bufferization::func_ext::registerBufferizableOpInterfaceExternalModels( + *registry); + builtin::registerCastOpInterfaceExternalModels(*registry); + cf::registerBufferizableOpInterfaceExternalModels(*registry); + cf::registerBufferDeallocationOpInterfaceExternalModels(*registry); + gpu::registerBufferDeallocationOpInterfaceExternalModels(*registry); + gpu::registerValueBoundsOpInterfaceExternalModels(*registry); + LLVM::registerInlinerInterface(*registry); + NVVM::registerInlinerInterface(*registry); + linalg::registerAllDialectInterfaceImplementations(*registry); + linalg::registerRuntimeVerifiableOpInterfaceExternalModels(*registry); + memref::registerAllocationOpInterfaceExternalModels(*registry); + memref::registerBufferViewFlowOpInterfaceExternalModels(*registry); + memref::registerRuntimeVerifiableOpInterfaceExternalModels(*registry); + memref::registerValueBoundsOpInterfaceExternalModels(*registry); + memref::registerMemorySlotExternalModels(*registry); + ml_program::registerBufferizableOpInterfaceExternalModels(*registry); + scf::registerBufferDeallocationOpInterfaceExternalModels(*registry); + scf::registerBufferizableOpInterfaceExternalModels(*registry); + scf::registerValueBoundsOpInterfaceExternalModels(*registry); + shape::registerBufferizableOpInterfaceExternalModels(*registry); + sparse_tensor::registerBufferizableOpInterfaceExternalModels(*registry); + tensor::registerBufferizableOpInterfaceExternalModels(*registry); + tensor::registerFindPayloadReplacementOpInterfaceExternalModels(*registry); + tensor::registerInferTypeOpInterfaceExternalModels(*registry); + tensor::registerRuntimeVerifiableOpInterfaceExternalModels(*registry); + tensor::registerSubsetOpInterfaceExternalModels(*registry); + tensor::registerTilingInterfaceExternalModels(*registry); + tensor::registerValueBoundsOpInterfaceExternalModels(*registry); + tosa::registerShardingInterfaceExternalModels(*registry); + vector::registerBufferizableOpInterfaceExternalModels(*registry); + vector::registerSubsetOpInterfaceExternalModels(*registry); + vector::registerValueBoundsOpInterfaceExternalModels(*registry); + NVVM::registerNVVMTargetInterfaceExternalModels(*registry); + ROCDL::registerROCDLTargetInterfaceExternalModels(*registry); + spirv::registerSPIRVTargetInterfaceExternalModels(*registry); + xevm::registerXeVMTargetInterfaceExternalModels(*registry); +} diff --git a/mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp b/mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp index c5dc64ca17834..0eeeaa279aef7 100644 --- a/mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp +++ b/mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp @@ -1,4 +1,4 @@ -//===- RegisterAllPasses.cpp - Register all MLIR entities -----------------===// +//===- RegisterAllPasses.cpp - Register all MLIR LLVM translations --------===// // // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // See https://llvm.org/LICENSE.txt for license information. diff --git a/mlir/lib/CAPI/RegisterEverything/RegisterAllPasses.cpp b/mlir/lib/CAPI/RegisterEverything/RegisterAllPasses.cpp index 159dff0b86e44..88a189cf73605 100644 --- a/mlir/lib/CAPI/RegisterEverything/RegisterAllPasses.cpp +++ b/mlir/lib/CAPI/RegisterEverything/RegisterAllPasses.cpp @@ -1,4 +1,4 @@ -//===- RegisterAllPasses.cpp - Register all MLIR entities -----------------===// +//===- RegisterAllPasses.cpp - Register all MLIR passes -------------------===// // // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // See https://llvm.org/LICENSE.txt for license information. diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt index 9a45c5026f60e..de736c77ebdd2 100644 --- a/mlir/python/CMakeLists.txt +++ b/mlir/python/CMakeLists.txt @@ -23,6 +23,7 @@ declare_mlir_python_sources(MLIRPythonSources.Core.Python ADD_TO_PARENT MLIRPythonSources.Core SOURCES _mlir_libs/__init__.py + _mlir_libs/_capi.py _mlir_libs/_mlir/py.typed ir.py passmanager.py @@ -204,11 +205,7 @@ declare_mlir_dialect_extension_python_bindings( SOURCES dialects/transform/pdl.py DIALECT_NAME transform - EXTENSION_NAME transform_pdl_extension - EMBED_CAPI_LINK_LIBS - MLIRCAPITransformDialect - MLIRCAPIPDLTransformDialectExtension -) + EXTENSION_NAME transform_pdl_extension) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects @@ -217,11 +214,7 @@ declare_mlir_dialect_extension_python_bindings( SOURCES dialects/transform/smt.py DIALECT_NAME transform - EXTENSION_NAME transform_smt_extension - EMBED_CAPI_LINK_LIBS - MLIRCAPITransformDialect - MLIRCAPISMTTransformDialectExtension -) + EXTENSION_NAME transform_smt_extension) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects @@ -277,11 +270,7 @@ declare_mlir_dialect_extension_python_bindings( SOURCES dialects/transform/bufferization.py DIALECT_NAME transform - EXTENSION_NAME bufferization_transform - EMBED_CAPI_LINK_LIBS - MLIRCAPITransformDialect - MLIRCAPIBufferizationTransformDialectExtension - ) + EXTENSION_NAME bufferization_transform) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects @@ -290,11 +279,7 @@ declare_mlir_dialect_extension_python_bindings( SOURCES dialects/transform/gpu.py DIALECT_NAME transform - EXTENSION_NAME gpu_transform - EMBED_CAPI_LINK_LIBS - MLIRCAPITransformDialect - MLIRCAPIGPUTransformDialectExtension -) + EXTENSION_NAME gpu_transform) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects @@ -303,10 +288,7 @@ declare_mlir_dialect_extension_python_bindings( SOURCES dialects/transform/loop.py DIALECT_NAME transform - EXTENSION_NAME loop_transform - EMBED_CAPI_LINK_LIBS - MLIRCAPITransformDialect -) + EXTENSION_NAME loop_transform) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects @@ -315,11 +297,7 @@ declare_mlir_dialect_extension_python_bindings( SOURCES dialects/transform/memref.py DIALECT_NAME transform - EXTENSION_NAME memref_transform - EMBED_CAPI_LINK_LIBS - MLIRCAPITransformDialect - MLIRCAPIMemRefTransformDialectExtension -) + EXTENSION_NAME memref_transform) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects @@ -328,11 +306,7 @@ declare_mlir_dialect_extension_python_bindings( SOURCES dialects/transform/nvgpu.py DIALECT_NAME transform - EXTENSION_NAME nvgpu_transform - EMBED_CAPI_LINK_LIBS - MLIRCAPITransformDialect - MLIRCAPINVGPUTransformDialectExtension -) + EXTENSION_NAME nvgpu_transform) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects @@ -344,9 +318,6 @@ declare_mlir_dialect_extension_python_bindings( EXTENSION_NAME structured_transform GEN_ENUM_BINDINGS_TD_FILE "../../include/mlir/Dialect/Linalg/TransformOps/LinalgTransformEnums.td" - EMBED_CAPI_LINK_LIBS - MLIRCAPITransformDialect - MLIRCAPILinalgTransformDialectExtension ) declare_mlir_dialect_extension_python_bindings( @@ -356,11 +327,7 @@ declare_mlir_dialect_extension_python_bindings( SOURCES dialects/transform/sparse_tensor.py DIALECT_NAME transform - EXTENSION_NAME sparse_tensor_transform - EMBED_CAPI_LINK_LIBS - MLIRCAPITransformDialect - MLIRCAPISparseTensorTransformDialectExtension -) + EXTENSION_NAME sparse_tensor_transform) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects @@ -369,11 +336,7 @@ declare_mlir_dialect_extension_python_bindings( SOURCES dialects/transform/tensor.py DIALECT_NAME transform - EXTENSION_NAME tensor_transform - EMBED_CAPI_LINK_LIBS - MLIRCAPITransformDialect - MLIRCAPITensorTransformDialectExtension -) + EXTENSION_NAME tensor_transform) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects @@ -385,9 +348,6 @@ 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( @@ -691,6 +651,7 @@ declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything MLIRCAPIRegisterAllPasses MLIRCAPIRegisterAllExtensions MLIRCAPIRegisterAllLLVMTranslations + MLIRCAPIRegisterAllExternalModels ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind @@ -888,7 +849,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind MLIRCAPIIR MLIRCAPISMT MLIRCAPIExportSMTLIB - MLIRCAPISMTTransformDialectExtension ) declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses @@ -989,20 +949,8 @@ 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(${MLIR_PYTHON_CAPI_DYLIB_NAME} +add_mlir_python_common_capi_library(MLIRPythonCAPI INSTALL_COMPONENT MLIRPythonModules INSTALL_DESTINATION "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/_mlir_libs" OUTPUT_DIRECTORY "${MLIRPythonModules_ROOT_PREFIX}/_mlir_libs" @@ -1120,8 +1068,7 @@ endif() set(_declared_sources MLIRPythonSources - MLIRPythonExtension.RegisterEverything - MLIRPythonCAPICTypesBinding) + MLIRPythonExtension.RegisterEverything) if(NOT CMAKE_CROSSCOMPILING) list(APPEND _declared_sources MLIRPythonExtension.Core.type_stub_gen) endif() diff --git a/mlir/python/mlir/_mlir_libs/__init__.py b/mlir/python/mlir/_mlir_libs/__init__.py index 2c36397d2cf8b..e36ed636b619d 100644 --- a/mlir/python/mlir/_mlir_libs/__init__.py +++ b/mlir/python/mlir/_mlir_libs/__init__.py @@ -126,6 +126,9 @@ def process_initializer_module(module_name): if hasattr(m, "register_extensions"): logger.debug("Registering extensions from initializer %r", m) m.register_extensions(get_dialect_registry()) + if hasattr(m, "register_external_models"): + logger.debug("Registering external models from initializer %r", m) + m.register_external_models(get_dialect_registry()) if hasattr(m, "context_init_hook"): logger.debug("Adding context init hook from %r", m) post_init_hooks.append(m.context_init_hook) diff --git a/mlir/python/mlir/_mlir_libs/_capi.py.in b/mlir/python/mlir/_mlir_libs/_capi.py similarity index 53% rename from mlir/python/mlir/_mlir_libs/_capi.py.in rename to mlir/python/mlir/_mlir_libs/_capi.py index 2313ecc4f9823..3d7836daa1c8b 100644 --- a/mlir/python/mlir/_mlir_libs/_capi.py.in +++ b/mlir/python/mlir/_mlir_libs/_capi.py @@ -3,14 +3,33 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception import ctypes +import platform from pathlib import Path -_capi = ctypes.CDLL( - str( - Path(__file__).parent - / "@CMAKE_SHARED_LIBRARY_PREFIX@@MLIR_PYTHON_CAPI_DYLIB_NAME@@CMAKE_SHARED_LIBRARY_SUFFIX@" - ) -) +from . import _mlir, get_dialect_registry as _get_dialect_registry + +_get_dialect_registry() + +MLIR_PYTHON_CAPSULE_DIALECT_HANDLE = _mlir.ir.DialectHandle._capsule_name.encode() + +MLIR_PYTHON_CAPSULE_DIALECT_REGISTRY = _mlir.ir.DialectRegistry._capsule_name.encode() + +if platform.system() == "Windows": + _ext_suffix = "dll" +elif platform.system() == "Darwin": + _ext_suffix = "dylib" +else: + _ext_suffix = "so" + +for fp in Path(__file__).parent.glob(f"*.{_ext_suffix}"): + if "CAPI" in fp.name: + _capi_dylib = fp + break +else: + raise ValueError("Couldn't find CAPI dylib") + + +_capi = ctypes.CDLL(str(Path(__file__).parent / _capi_dylib)) PyCapsule_New = ctypes.pythonapi.PyCapsule_New PyCapsule_New.restype = ctypes.py_object @@ -20,14 +39,6 @@ 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@.ir.DialectHandle._CAPIPtr" -).encode() - -MLIR_PYTHON_CAPSULE_DIALECT_REGISTRY = ( - "@MLIR_PYTHON_PACKAGE_PREFIX@.ir.DialectRegistry._CAPIPtr" -).encode() - def register_dialect(dialect_handle_capi_name, dialect_registry): if not hasattr(_capi, dialect_handle_capi_name): @@ -37,14 +48,3 @@ def register_dialect(dialect_handle_capi_name, dialect_registry): 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] - dialect_registry_ptr = PyCapsule_GetPointer( - dialect_registry._CAPIPtr, MLIR_PYTHON_CAPSULE_DIALECT_REGISTRY - ) - registration_capi(dialect_registry_ptr) diff --git a/mlir/python/mlir/dialects/transform/bufferization.py b/mlir/python/mlir/dialects/transform/bufferization.py index 228c9dab52b8d..485a8a36b6305 100644 --- a/mlir/python/mlir/dialects/transform/bufferization.py +++ b/mlir/python/mlir/dialects/transform/bufferization.py @@ -15,15 +15,6 @@ 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/debug.py b/mlir/python/mlir/dialects/transform/debug.py index 353e4f07c883a..f7c04268dc03d 100644 --- a/mlir/python/mlir/dialects/transform/debug.py +++ b/mlir/python/mlir/dialects/transform/debug.py @@ -15,15 +15,6 @@ 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( - "mlirDebugRegisterTransformDialectExtension", _get_dialect_registry() -) - @_ods_cext.register_operation(_Dialect, replace=True) class EmitParamAsRemarkOp(EmitParamAsRemarkOp): diff --git a/mlir/python/mlir/dialects/transform/gpu.py b/mlir/python/mlir/dialects/transform/gpu.py index b35004cb4a09a..00cf0840eeae9 100644 --- a/mlir/python/mlir/dialects/transform/gpu.py +++ b/mlir/python/mlir/dialects/transform/gpu.py @@ -15,16 +15,6 @@ 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 3311ae47b4fa9..c4770b1c4067e 100644 --- a/mlir/python/mlir/dialects/transform/loop.py +++ b/mlir/python/mlir/dialects/transform/loop.py @@ -16,15 +16,6 @@ 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 f45d8054bfa58..56ea61eb817f8 100644 --- a/mlir/python/mlir/dialects/transform/memref.py +++ b/mlir/python/mlir/dialects/transform/memref.py @@ -14,15 +14,6 @@ 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/nvgpu.py b/mlir/python/mlir/dialects/transform/nvgpu.py index 3d6427d1b1ed9..74ba4c9aeb6c1 100644 --- a/mlir/python/mlir/dialects/transform/nvgpu.py +++ b/mlir/python/mlir/dialects/transform/nvgpu.py @@ -3,12 +3,3 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception from .._nvgpu_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( - "mlirNVGPURegisterTransformDialectExtension", _get_dialect_registry() -) diff --git a/mlir/python/mlir/dialects/transform/pdl.py b/mlir/python/mlir/dialects/transform/pdl.py index 15b0c37837902..bb5fa7ffd3065 100644 --- a/mlir/python/mlir/dialects/transform/pdl.py +++ b/mlir/python/mlir/dialects/transform/pdl.py @@ -18,16 +18,6 @@ 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( - "mlirPDLRegisterTransformDialectExtension", _get_dialect_registry() -) - - @_ods_cext.register_operation(_Dialect, replace=True) class PDLMatchOp(PDLMatchOp): def __init__( diff --git a/mlir/python/mlir/dialects/transform/smt.py b/mlir/python/mlir/dialects/transform/smt.py index a4e8eb96d388c..1f0b7f066118c 100644 --- a/mlir/python/mlir/dialects/transform/smt.py +++ b/mlir/python/mlir/dialects/transform/smt.py @@ -15,16 +15,6 @@ 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 efaf438e07b83..8b33270dc23a1 100644 --- a/mlir/python/mlir/dialects/transform/sparse_tensor.py +++ b/mlir/python/mlir/dialects/transform/sparse_tensor.py @@ -3,11 +3,3 @@ # 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/structured.py b/mlir/python/mlir/dialects/transform/structured.py index 85833d32ba55e..bf40cc532065d 100644 --- a/mlir/python/mlir/dialects/transform/structured.py +++ b/mlir/python/mlir/dialects/transform/structured.py @@ -28,15 +28,6 @@ from typing import List, 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( - "mlirLinalgRegisterTransformDialectExtension", _get_dialect_registry() -) - @_ods_cext.register_operation(_Dialect, replace=True) class BufferizeToAllocationOp(BufferizeToAllocationOp): diff --git a/mlir/python/mlir/dialects/transform/tensor.py b/mlir/python/mlir/dialects/transform/tensor.py index 904f57827146d..4eb30398f0872 100644 --- a/mlir/python/mlir/dialects/transform/tensor.py +++ b/mlir/python/mlir/dialects/transform/tensor.py @@ -14,15 +14,6 @@ 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 6eadd9b4e9d5d..f63f88a382422 100644 --- a/mlir/python/mlir/dialects/transform/tune.py +++ b/mlir/python/mlir/dialects/transform/tune.py @@ -25,15 +25,6 @@ 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 1d7620188da83..af2435cb26cc4 100644 --- a/mlir/python/mlir/dialects/transform/vector.py +++ b/mlir/python/mlir/dialects/transform/vector.py @@ -4,11 +4,3 @@ 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/module-to-binary-nvvm.py b/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py index 96e3f6efcc5d3..cda570fceb64b 100644 --- a/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py +++ b/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py @@ -2,7 +2,7 @@ # RUN: %PYTHON %s | FileCheck %s from mlir.ir import * -import mlir.dialects.gpu as gpu +from mlir.dialects import gpu, llvm, nvvm import mlir.dialects.gpu.passes from mlir.passmanager import * diff --git a/mlir/test/python/dialects/gpu/module-to-binary-rocdl.py b/mlir/test/python/dialects/gpu/module-to-binary-rocdl.py index f5ee2856f3950..6e8c13b8dea3c 100644 --- a/mlir/test/python/dialects/gpu/module-to-binary-rocdl.py +++ b/mlir/test/python/dialects/gpu/module-to-binary-rocdl.py @@ -2,7 +2,7 @@ # RUN: %PYTHON %s | FileCheck %s from mlir.ir import * -import mlir.dialects.gpu as gpu +from mlir.dialects import gpu, llvm, rocdl import mlir.dialects.gpu.passes from mlir.passmanager import * diff --git a/mlir/test/python/dialects/transform.py b/mlir/test/python/dialects/transform.py index 3477539b8d2f0..6c5e4e5505b1c 100644 --- a/mlir/test/python/dialects/transform.py +++ b/mlir/test/python/dialects/transform.py @@ -12,7 +12,6 @@ def run(f): print("\nTEST:", f.__name__) f(module) print(module) - module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_bufferization_ext.py b/mlir/test/python/dialects/transform_bufferization_ext.py index 39014735aa8fb..03b16c324a407 100644 --- a/mlir/test/python/dialects/transform_bufferization_ext.py +++ b/mlir/test/python/dialects/transform_bufferization_ext.py @@ -13,7 +13,6 @@ def run(f): print("\nTEST:", f.__name__) f() print(module) - module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_debug_ext.py b/mlir/test/python/dialects/transform_debug_ext.py index 0826ed14d199b..2dfdaed343865 100644 --- a/mlir/test/python/dialects/transform_debug_ext.py +++ b/mlir/test/python/dialects/transform_debug_ext.py @@ -19,7 +19,6 @@ def run(f): f(sequence.bodyTarget) transform.YieldOp() print(module) - module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_extras.py b/mlir/test/python/dialects/transform_extras.py index 3aeac681f8d49..c52d5b41e11b6 100644 --- a/mlir/test/python/dialects/transform_extras.py +++ b/mlir/test/python/dialects/transform_extras.py @@ -30,7 +30,6 @@ def construct_and_print_in_module(f): with ir.InsertionPoint(module.body): f() print(module) - module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_gpu_ext.py b/mlir/test/python/dialects/transform_gpu_ext.py index 5771ebd13ab77..17f5a6d38f8c1 100644 --- a/mlir/test/python/dialects/transform_gpu_ext.py +++ b/mlir/test/python/dialects/transform_gpu_ext.py @@ -19,7 +19,6 @@ def run(f): transform.YieldOp() print("\nTEST:", f.__name__) print(module) - module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_loop_ext.py b/mlir/test/python/dialects/transform_loop_ext.py index 50c93fe9f46d7..430b33fba04c7 100644 --- a/mlir/test/python/dialects/transform_loop_ext.py +++ b/mlir/test/python/dialects/transform_loop_ext.py @@ -13,7 +13,6 @@ def run(f): print("\nTEST:", f.__name__) f() print(module) - module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_memref_ext.py b/mlir/test/python/dialects/transform_memref_ext.py index 55eceb57f97c8..e7d871c9eac8c 100644 --- a/mlir/test/python/dialects/transform_memref_ext.py +++ b/mlir/test/python/dialects/transform_memref_ext.py @@ -13,7 +13,6 @@ def run(f): print("\nTEST:", f.__name__) f() print(module) - module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_nvgpu_ext.py b/mlir/test/python/dialects/transform_nvgpu_ext.py index 45672ff3b58e4..1de8b25bab7a5 100644 --- a/mlir/test/python/dialects/transform_nvgpu_ext.py +++ b/mlir/test/python/dialects/transform_nvgpu_ext.py @@ -12,7 +12,6 @@ def run(f): print("\nTEST:", f.__name__) f() print(module) - module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_smt_ext.py b/mlir/test/python/dialects/transform_smt_ext.py index 0593421c7713d..3692fd92344a6 100644 --- a/mlir/test/python/dialects/transform_smt_ext.py +++ b/mlir/test/python/dialects/transform_smt_ext.py @@ -19,7 +19,6 @@ def run(f): f(sequence.bodyTarget) transform.YieldOp() print(module) - module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_sparse_tensor_ext.py b/mlir/test/python/dialects/transform_sparse_tensor_ext.py index b22f27ab6494b..e11cc6bf1e074 100644 --- a/mlir/test/python/dialects/transform_sparse_tensor_ext.py +++ b/mlir/test/python/dialects/transform_sparse_tensor_ext.py @@ -19,7 +19,6 @@ def run(f): transform.YieldOp() print("\nTEST:", f.__name__) print(module) - module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_structured_ext.py b/mlir/test/python/dialects/transform_structured_ext.py index ee15ac43e3d1a..8785d6d360074 100644 --- a/mlir/test/python/dialects/transform_structured_ext.py +++ b/mlir/test/python/dialects/transform_structured_ext.py @@ -19,7 +19,6 @@ def run(f): f() module.operation.verify() print(module) - module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_tensor_ext.py b/mlir/test/python/dialects/transform_tensor_ext.py index 9dde72f52a8fc..a2e7aa242b9da 100644 --- a/mlir/test/python/dialects/transform_tensor_ext.py +++ b/mlir/test/python/dialects/transform_tensor_ext.py @@ -19,7 +19,6 @@ def run(f): f(sequence.bodyTarget) transform.YieldOp() print(module) - module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_tune_ext.py b/mlir/test/python/dialects/transform_tune_ext.py index f8db151e10703..dfb93594bca52 100644 --- a/mlir/test/python/dialects/transform_tune_ext.py +++ b/mlir/test/python/dialects/transform_tune_ext.py @@ -19,7 +19,6 @@ def run(f): f(sequence.bodyTarget) transform.YieldOp() print(module) - module.operation.verify() return f diff --git a/mlir/test/python/dialects/transform_vector_ext.py b/mlir/test/python/dialects/transform_vector_ext.py index 0e841f2895b55..0cd9333dc1218 100644 --- a/mlir/test/python/dialects/transform_vector_ext.py +++ b/mlir/test/python/dialects/transform_vector_ext.py @@ -21,7 +21,6 @@ def run_apply_patterns(f): transform.YieldOp() print("\nTEST:", f.__name__) print(module) - module.operation.verify() return f diff --git a/mlir/test/python/execution_engine.py b/mlir/test/python/execution_engine.py index d569fcef32bfd..f93337c5cbb6a 100644 --- a/mlir/test/python/execution_engine.py +++ b/mlir/test/python/execution_engine.py @@ -5,6 +5,7 @@ from mlir.passmanager import * from mlir.execution_engine import * from mlir.runtime import * +from mlir.dialects import func, memref, arith, complex, cf, llvm try: from ml_dtypes import bfloat16, float8_e5m2 diff --git a/mlir/test/python/global_constructors.py b/mlir/test/python/global_constructors.py index 5020c00344a33..a11045bc8526c 100644 --- a/mlir/test/python/global_constructors.py +++ b/mlir/test/python/global_constructors.py @@ -6,6 +6,7 @@ from mlir.passmanager import * from mlir.execution_engine import * from mlir.runtime import * +from mlir.dialects import func, llvm # Log everything to stderr and flush so that we have a unified stream to match