diff --git a/mlir/lib/Bindings/Python/DialectLinalg.cpp b/mlir/lib/Bindings/Python/DialectLinalg.cpp --- a/mlir/lib/Bindings/Python/DialectLinalg.cpp +++ b/mlir/lib/Bindings/Python/DialectLinalg.cpp @@ -6,14 +6,13 @@ // //===----------------------------------------------------------------------===// -#include "Dialects.h" #include "mlir-c/Dialect/Linalg.h" #include "mlir-c/IR.h" #include "mlir/Bindings/Python/PybindAdaptors.h" namespace py = pybind11; -void mlir::python::populateDialectLinalgSubmodule(py::module m) { +static void populateDialectLinalgSubmodule(py::module m) { m.def( "fill_builtin_region", [](MlirOperation op) { mlirLinalgFillBuiltinNamedOpRegion(op); }, @@ -21,3 +20,9 @@ "Fill the region for `op`, which is assumed to be a builtin named Linalg " "op."); } + +PYBIND11_MODULE(_mlirDialectsLinalg, m) { + m.doc() = "MLIR Linalg dialect."; + + populateDialectLinalgSubmodule(m); +} diff --git a/mlir/lib/Bindings/Python/DialectQuant.cpp b/mlir/lib/Bindings/Python/DialectQuant.cpp --- a/mlir/lib/Bindings/Python/DialectQuant.cpp +++ b/mlir/lib/Bindings/Python/DialectQuant.cpp @@ -6,7 +6,6 @@ // //===----------------------------------------------------------------------===// -#include "Dialects.h" #include "mlir-c/Dialect/Quant.h" #include "mlir-c/IR.h" #include "mlir/Bindings/Python/PybindAdaptors.h" @@ -16,16 +15,13 @@ using namespace mlir; using namespace mlir::python::adaptors; -void mlir::python::populateDialectQuantSubmodule(const py::module &m, - const py::module &irModule) { - auto typeClass = irModule.attr("Type"); - +static void populateDialectQuantSubmodule(const py::module &m) { //===-------------------------------------------------------------------===// // QuantizedType //===-------------------------------------------------------------------===// - auto quantizedType = mlir_type_subclass(m, "QuantizedType", - mlirTypeIsAQuantizedType, typeClass); + auto quantizedType = + mlir_type_subclass(m, "QuantizedType", mlirTypeIsAQuantizedType); quantizedType.def_staticmethod( "default_minimum_for_integer", [](bool isSigned, unsigned integralWidth) { @@ -305,3 +301,9 @@ return mlirCalibratedQuantizedTypeGetMax(type); }); } + +PYBIND11_MODULE(_mlirDialectsQuant, m) { + m.doc() = "MLIR Quantization dialect"; + + populateDialectQuantSubmodule(m); +} \ No newline at end of file diff --git a/mlir/lib/Bindings/Python/DialectSparseTensor.cpp b/mlir/lib/Bindings/Python/DialectSparseTensor.cpp --- a/mlir/lib/Bindings/Python/DialectSparseTensor.cpp +++ b/mlir/lib/Bindings/Python/DialectSparseTensor.cpp @@ -6,7 +6,6 @@ // //===----------------------------------------------------------------------===// -#include "Dialects.h" #include "mlir-c/Dialect/SparseTensor.h" #include "mlir-c/IR.h" #include "mlir/Bindings/Python/PybindAdaptors.h" @@ -16,18 +15,14 @@ using namespace mlir; using namespace mlir::python::adaptors; -void mlir::python::populateDialectSparseTensorSubmodule( - const py::module &m, const py::module &irModule) { - auto attributeClass = irModule.attr("Attribute"); - +static void populateDialectSparseTensorSubmodule(const py::module &m) { py::enum_(m, "DimLevelType", py::module_local()) .value("dense", MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE) .value("compressed", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED) .value("singleton", MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON); mlir_attribute_subclass(m, "EncodingAttr", - mlirAttributeIsASparseTensorEncodingAttr, - attributeClass) + mlirAttributeIsASparseTensorEncodingAttr) .def_classmethod( "get", [](py::object cls, @@ -72,3 +67,8 @@ return mlirSparseTensorEncodingAttrGetIndexBitWidth(self); }); } + +PYBIND11_MODULE(_mlirDialectsSparseTensor, m) { + m.doc() = "MLIR SparseTensor dialect."; + populateDialectSparseTensorSubmodule(m); +} diff --git a/mlir/lib/Bindings/Python/Dialects.h b/mlir/lib/Bindings/Python/Dialects.h deleted file mode 100644 --- a/mlir/lib/Bindings/Python/Dialects.h +++ /dev/null @@ -1,26 +0,0 @@ -//===- Dialects.h - Declaration for dialect submodule factories -----------===// -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// -//===----------------------------------------------------------------------===// - -#ifndef MLIR_BINDINGS_PYTHON_DIALECTS_H -#define MLIR_BINDINGS_PYTHON_DIALECTS_H - -#include - -namespace mlir { -namespace python { - -void populateDialectLinalgSubmodule(pybind11::module m); -void populateDialectSparseTensorSubmodule(const pybind11::module &m, - const pybind11::module &irModule); -void populateDialectQuantSubmodule(const pybind11::module &m, - const pybind11::module &irModule); - -} // namespace python -} // namespace mlir - -#endif // MLIR_BINDINGS_PYTHON_DIALECTS_H diff --git a/mlir/lib/Bindings/Python/MainModule.cpp b/mlir/lib/Bindings/Python/MainModule.cpp --- a/mlir/lib/Bindings/Python/MainModule.cpp +++ b/mlir/lib/Bindings/Python/MainModule.cpp @@ -10,7 +10,6 @@ #include "PybindUtils.h" -#include "Dialects.h" #include "Globals.h" #include "IRModule.h" #include "Pass.h" @@ -100,13 +99,4 @@ auto passModule = m.def_submodule("passmanager", "MLIR Pass Management Bindings"); populatePassManagerSubmodule(passModule); - - // Define and populate dialect submodules. - auto dialectsModule = m.def_submodule("dialects"); - auto linalgModule = dialectsModule.def_submodule("linalg"); - populateDialectLinalgSubmodule(linalgModule); - populateDialectSparseTensorSubmodule( - dialectsModule.def_submodule("sparse_tensor"), irModule); - populateDialectQuantSubmodule(dialectsModule.def_submodule("quant"), - irModule); } diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt --- a/mlir/python/CMakeLists.txt +++ b/mlir/python/CMakeLists.txt @@ -25,8 +25,6 @@ _mlir_libs/_mlir/__init__.pyi _mlir_libs/_mlir/ir.pyi _mlir_libs/_mlir/passmanager.pyi - # TODO: this should be split out into a separate library. - _mlir_libs/_mlir/dialects/quant.pyi ) declare_mlir_python_sources(MLIRPythonSources.ExecutionEngine @@ -122,7 +120,8 @@ ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" SOURCES - dialects/quant.py) + dialects/quant.py + _mlir_libs/_mlir/dialects/quant.pyi) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects @@ -191,9 +190,6 @@ ADD_TO_PARENT MLIRPythonSources.Core ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES - DialectLinalg.cpp # TODO: Break this out. - DialectSparseTensor.cpp # TODO: Break this out. - DialectQuant.cpp # TODO: Break this out. MainModule.cpp IRAffine.cpp IRAttributes.cpp @@ -205,7 +201,6 @@ Pass.cpp # Headers must be included explicitly so they are installed. - Dialects.h Globals.h IRModule.h Pass.h @@ -219,10 +214,46 @@ MLIRCAPIRegistration # TODO: See about dis-aggregating # Dialects - MLIRCAPILinalg # TODO: Remove when above is removed. - MLIRCAPISparseTensor # TODO: Remove when above is removed. MLIRCAPIStandard - MLIRCAPIQuant # TODO: Remove when above is removed. +) + +declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind + MODULE_NAME _mlirDialectsLinalg + ADD_TO_PARENT MLIRPythonSources.Dialects.linalg + ROOT_DIR "${PYTHON_SOURCE_DIR}" + SOURCES + DialectLinalg.cpp + PRIVATE_LINK_LIBS + LLVMSupport + EMBED_CAPI_LINK_LIBS + MLIRCAPIIR + MLIRCAPILinalg +) + +declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind + MODULE_NAME _mlirDialectsQuant + ADD_TO_PARENT MLIRPythonSources.Dialects.quant + ROOT_DIR "${PYTHON_SOURCE_DIR}" + SOURCES + DialectQuant.cpp + PRIVATE_LINK_LIBS + LLVMSupport + EMBED_CAPI_LINK_LIBS + MLIRCAPIIR + MLIRCAPIQuant +) + +declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind + MODULE_NAME _mlirDialectsSparseTensor + ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor + ROOT_DIR "${PYTHON_SOURCE_DIR}" + SOURCES + DialectSparseTensor.cpp + PRIVATE_LINK_LIBS + LLVMSupport + EMBED_CAPI_LINK_LIBS + MLIRCAPIIR + MLIRCAPISparseTensor ) declare_mlir_python_extension(MLIRPythonExtension.AllPassesRegistration diff --git a/mlir/python/mlir/dialects/_linalg_ops_ext.py b/mlir/python/mlir/dialects/_linalg_ops_ext.py --- a/mlir/python/mlir/dialects/_linalg_ops_ext.py +++ b/mlir/python/mlir/dialects/_linalg_ops_ext.py @@ -6,7 +6,7 @@ from typing import Optional, Sequence, Union from ..ir import * from ._ods_common import get_default_loc_context - from .._mlir_libs._mlir.dialects.linalg import fill_builtin_region + from .._mlir_libs._mlirDialectsLinalg import fill_builtin_region except ImportError as e: raise RuntimeError("Error loading imports from extension module") from e diff --git a/mlir/python/mlir/dialects/linalg/__init__.py b/mlir/python/mlir/dialects/linalg/__init__.py --- a/mlir/python/mlir/dialects/linalg/__init__.py +++ b/mlir/python/mlir/dialects/linalg/__init__.py @@ -2,6 +2,9 @@ # See https://llvm.org/LICENSE.txt for license information. # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +# Re-export the objects provided by pybind. +from ..._mlir_libs._mlirDialectsLinalg import * + # These are the backing OpView classes generated from the linalg tablegen # definitions following these steps: # DSL -> YAML -> tblgen -> pytblgen -> build/.../_linalg_ops_gen.py. @@ -15,39 +18,39 @@ # C=TensorDef(U, S.M, S.N, output=True)): # ``` # using the linalg-py eDSL. -# The linalg-py eDSL builds a python representation (PyRepr) that is +# The linalg-py eDSL builds a python representation (PyRepr) that is # used in following ways: # 1. PyRepr -> YAML to generate the C++ and Python .td files. These # then turn into the core C++ Op classes and Python OpView classes -# respectively (made available in _linalg_ops_gen). The generic OpView class +# respectively (made available in _linalg_ops_gen). The generic OpView class # mechanism makes the C++ classes available to python through the CAPI. # PyRepr -> YAML currently occurs before compiler compile time. # The other steps in this category occur at compiler compile time. -# 2. PyRepr -> linalg.core_named_ops calls: piggybacks on the +# 2. PyRepr -> linalg.core_named_ops calls: piggybacks on the # _linalg_ops_gen classes and the OpView mechanism to build IR at # runtime in python: # a. by default, the Named Op Form is emitted, e.g.: # `linalg.matmul(lhs, rhs, outs=[out])` creates the following IR: # ``` -# %1 = linalg.matmul ins(%arg0, %arg1 : tensor<4x16xf32>, tensor<16x8xf32>) +# %1 = linalg.matmul ins(%arg0, %arg1 : tensor<4x16xf32>, tensor<16x8xf32>) # outs(%0 : tensor<4x8xf32>) -# -> tensor<4x8xf32> +# -> tensor<4x8xf32> # ``` # b. by setting emit_generic=True, the Generic Op Form is emitted, e.g.: # `linalg.matmul(lhs, rhs, outs=[out], emit_generic=True)` creates the following IR: # ``` -# %1 = linalg.generic {indexing_maps = [...], iterator_types = [...]} -# ins(%arg0, %arg1 : tensor<4x16xf32>, tensor<16x8xf32>) +# %1 = linalg.generic {indexing_maps = [...], iterator_types = [...]} +# ins(%arg0, %arg1 : tensor<4x16xf32>, tensor<16x8xf32>) # outs(%0 : tensor<4x8xf32>) { -# ^bb0(%arg2: f32, %arg3: f32, %arg4: f32): +# ^bb0(%arg2: f32, %arg3: f32, %arg4: f32): # ... # linalg.yield %3 : f32 -# } -> tensor<4x8xf32> +# } -> tensor<4x8xf32> # ``` # 3. PyRepr -> Runtime Custom Op definitions: directly generates a # linalg.generic form like in 2.b. -# !!!WARNING!!!: if one creates a runtime custom op with the same name +# !!!WARNING!!!: if one creates a runtime custom op with the same name # as an existing core named op, step 2. will likely take precedence. -# TODO: guard against surprises and fail create Runtime Custom Ops with +# TODO: guard against surprises and fail create Runtime Custom Ops with # the same name as existing Core Named Ops. from .opdsl.ops.core_named_ops import * diff --git a/mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py b/mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py --- a/mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py +++ b/mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py @@ -5,7 +5,6 @@ from typing import Dict, List, Sequence, Tuple, Union from .....ir import * -from ....._mlir_libs._mlir.dialects.linalg import fill_builtin_region from .... import linalg from .... import std @@ -173,7 +172,7 @@ f"Unknown named op_name / op_class_name: {op_name} / {op_class_name}") named_op = getattr(linalg, op_class_name)(ins, outs, result_types) - fill_builtin_region(named_op.operation) + linalg.fill_builtin_region(named_op.operation) # Note: mlir-linalg-ods-yaml-gen.cpp uses a special linalg.memoized_indexing_maps # attribute that the non-yaml path does not. The non-yaml path hardcodes the # indexing_maps in C++ directly. diff --git a/mlir/python/mlir/dialects/quant.py b/mlir/python/mlir/dialects/quant.py --- a/mlir/python/mlir/dialects/quant.py +++ b/mlir/python/mlir/dialects/quant.py @@ -2,4 +2,4 @@ # See https://llvm.org/LICENSE.txt for license information. # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -from .._mlir_libs._mlir.dialects.quant import * +from .._mlir_libs._mlirDialectsQuant import * diff --git a/mlir/python/mlir/dialects/sparse_tensor.py b/mlir/python/mlir/dialects/sparse_tensor.py --- a/mlir/python/mlir/dialects/sparse_tensor.py +++ b/mlir/python/mlir/dialects/sparse_tensor.py @@ -3,5 +3,5 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception from ._sparse_tensor_ops_gen import * -from .._mlir_libs._mlir.dialects.sparse_tensor import * +from .._mlir_libs._mlirDialectsSparseTensor import * from .._mlir_libs import _mlirSparseTensorPasses as _cextSparseTensorPasses diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel --- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel +++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel @@ -647,10 +647,18 @@ ], ) +# These flags are needed for pybind11 to work. +PYBIND11_COPTS = [ + "-fexceptions", + "-frtti", +] + +PYBIND11_FEATURES = [ + # Cannot use header_modules (parse_headers feature fails). + "-use_header_modules", +] + MLIR_PYTHON_BINDINGS_SOURCES = [ - "lib/Bindings/Python/DialectLinalg.cpp", - "lib/Bindings/Python/DialectSparseTensor.cpp", - "lib/Bindings/Python/DialectQuant.cpp", "lib/Bindings/Python/IRAffine.cpp", "lib/Bindings/Python/IRAttributes.cpp", "lib/Bindings/Python/IRCore.cpp", @@ -664,15 +672,8 @@ cc_library( name = "MLIRBindingsPythonCore", srcs = MLIR_PYTHON_BINDINGS_SOURCES, - # These flags are needed for pybind11 to work. - copts = [ - "-fexceptions", - "-frtti", - ], - features = [ - # Cannot use header_modules (parse_headers feature fails). - "-use_header_modules", - ], + copts = PYBIND11_COPTS, + features = PYBIND11_FEATURES, tags = [ "manual", # External dependency "nobuildkite", # TODO(gcmn): Add support for this target @@ -683,10 +684,7 @@ ":CAPIGPU", ":CAPIIR", ":CAPIInterfaces", - ":CAPILinalg", - ":CAPIQuant", ":CAPIRegistration", - ":CAPISparseTensor", ":MLIRBindingsPythonHeadersAndDeps", "//llvm:Support", "@pybind11", @@ -697,15 +695,8 @@ cc_library( name = "MLIRBindingsPythonCoreNoCAPI", srcs = MLIR_PYTHON_BINDINGS_SOURCES, - # These flags are needed for pybind11 to work. - copts = [ - "-fexceptions", - "-frtti", - ], - features = [ - # Cannot use header_modules (parse_headers feature fails). - "-use_header_modules", - ], + copts = PYBIND11_COPTS, + features = PYBIND11_FEATURES, tags = [ "manual", # External dependency "nobuildkite", # TODO(gcmn): Add support for this target @@ -715,10 +706,7 @@ ":CAPIDebugHeaders", ":CAPIGPUHeaders", ":CAPIIRHeaders", - ":CAPILinalgHeaders", - ":CAPIQuantHeaders", ":CAPIRegistrationHeaders", - ":CAPISparseTensorHeaders", ":MLIRBindingsPythonHeaders", "//llvm:Support", "@pybind11", @@ -740,23 +728,10 @@ ":CAPIGPUObjects", ":CAPIIRObjects", ":CAPIInterfacesObjects", - ":CAPILinalgObjects", - ":CAPIQuantObjects", ":CAPIRegistrationObjects", - ":CAPISparseTensorObjects", ], ) -PYBIND11_COPTS = [ - "-fexceptions", - "-frtti", -] - -PYBIND11_FEATURES = [ - # Cannot use header_modules (parse_headers feature fails). - "-use_header_modules", -] - # Dynamic library with the MLIR Python extension. cc_binary( name = "_mlir.so", @@ -776,6 +751,63 @@ ], ) +cc_binary( + name = "_mlirDialectsLinalg.so", + srcs = ["lib/Bindings/Python/DialectLinalg.cpp"], + copts = PYBIND11_COPTS, + features = PYBIND11_FEATURES, + linkshared = 1, + linkstatic = 0, + tags = [ + "manual", # External dependency + "nobuildkite", # TODO(gcmn): Add support for this target + ], + deps = [ + ":CAPIIR", + ":CAPILinalg", + ":CAPIRegistration", + ":MLIRBindingsPythonHeadersAndDeps", + ], +) + +cc_binary( + name = "_mlirDialectsQuant.so", + srcs = ["lib/Bindings/Python/DialectQuant.cpp"], + copts = PYBIND11_COPTS, + features = PYBIND11_FEATURES, + linkshared = 1, + linkstatic = 0, + tags = [ + "manual", # External dependency + "nobuildkite", # TODO(gcmn): Add support for this target + ], + deps = [ + ":CAPIIR", + ":CAPIQuant", + ":CAPIRegistration", + ":MLIRBindingsPythonHeadersAndDeps", + ], +) + +cc_binary( + name = "_mlirDialectsSparseTensor.so", + srcs = ["lib/Bindings/Python/DialectSparseTensor.cpp"], + copts = PYBIND11_COPTS, + features = PYBIND11_FEATURES, + linkshared = 1, + linkstatic = 0, + tags = [ + "manual", # External dependency + "nobuildkite", # TODO(gcmn): Add support for this target + ], + deps = [ + ":CAPIIR", + ":CAPISparseTensor", + ":CAPIRegistration", + ":MLIRBindingsPythonHeadersAndDeps", + ], +) + # Dynamic library with the MLIR Conversions Python extension. cc_binary( name = "_mlirConversions.so",