OSchip
/
llvm-project
mirror of https://github.com/THU-DSP-LAB/llvm-project.git
13
0
Fork 2
Code Issues Packages Projects Releases Wiki Activity

[mlir] Split out Python bindings for dialects into separate libs 

Historically, the bindings for the Linalg dialect were included into the
"core" bindings library because they depended on the C++ implementation
of the "core" bindings. The other dialects followed the pattern. Now
that this dependency is gone, split out each dialect into a separate
Python extension library.

Depends On D116649, D116605

Reviewed By: stellaraccident

Differential Revision: https://reviews.llvm.org/D116662
This commit is contained in:
Alex Zinenko 2022-01-05 11:21:21 +01:00
parent 6e4bbbfcc8
commit 95ddbed9b7
12 changed files with 154 additions and 118 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
mlir/lib/Bindings/Python/DialectLinalg.cpp
View File

@ -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 @@ void mlir::python::populateDialectLinalgSubmodule(py::module m) {
"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);
}

16
mlir/lib/Bindings/Python/DialectQuant.cpp
View File

@ -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 llvm;
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 @@ void mlir::python::populateDialectQuantSubmodule(const py::module &m,
return mlirCalibratedQuantizedTypeGetMax(type);
});
}
PYBIND11_MODULE(_mlirDialectsQuant, m) {
m.doc() = "MLIR Quantization dialect";
populateDialectQuantSubmodule(m);
}

14
mlir/lib/Bindings/Python/DialectSparseTensor.cpp
View File

@ -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 llvm;
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_<MlirSparseTensorDimLevelType>(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 @@ void mlir::python::populateDialectSparseTensorSubmodule(
return mlirSparseTensorEncodingAttrGetIndexBitWidth(self);
});
}
PYBIND11_MODULE(_mlirDialectsSparseTensor, m) {
m.doc() = "MLIR SparseTensor dialect.";
populateDialectSparseTensorSubmodule(m);
}

26
mlir/lib/Bindings/Python/Dialects.h
View File

@ -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 <pybind11/pybind11.h>
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

10
mlir/lib/Bindings/Python/MainModule.cpp
View File

@ -10,7 +10,6 @@
#include "PybindUtils.h"
#include "Dialects.h"
#include "Globals.h"
#include "IRModule.h"
#include "Pass.h"
@ -100,13 +99,4 @@ PYBIND11_MODULE(_mlir, m) {
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);
}

51
mlir/python/CMakeLists.txt
View File

@ -25,8 +25,6 @@ declare_mlir_python_sources(MLIRPythonSources.Core
_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 @@ declare_mlir_python_sources(
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 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
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 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
Pass.cpp
# Headers must be included explicitly so they are installed.
Dialects.h
Globals.h
IRModule.h
Pass.h
@ -219,10 +214,46 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
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

2
mlir/python/mlir/dialects/_linalg_ops_ext.py
View File

@ -6,7 +6,7 @@ try:
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

25
mlir/python/mlir/dialects/linalg/__init__.py
View File

@ -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 @@ from .._linalg_ops_gen import *
# 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 *

3
mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py
View File

@ -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 @@ def emit_named_structured_op(op_config: LinalgStructuredOpConfig, op_name: str,
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.

2
mlir/python/mlir/dialects/quant.py
View File

@ -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 *

2
mlir/python/mlir/dialects/sparse_tensor.py
View File

@ -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

112
utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
View File

@ -647,10 +647,18 @@ cc_library(
],
)
# 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 @@ MLIR_PYTHON_BINDINGS_SOURCES = [
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 @@ cc_library(
":CAPIGPU",
":CAPIIR",
":CAPIInterfaces",
":CAPILinalg",
":CAPIQuant",
":CAPIRegistration",
":CAPISparseTensor",
":MLIRBindingsPythonHeadersAndDeps",
"//llvm:Support",
"@pybind11",
@ -697,15 +695,8 @@ cc_library(
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 @@ cc_library(
":CAPIDebugHeaders",
":CAPIGPUHeaders",
":CAPIIRHeaders",
":CAPILinalgHeaders",
":CAPIQuantHeaders",
":CAPIRegistrationHeaders",
":CAPISparseTensorHeaders",
":MLIRBindingsPythonHeaders",
"//llvm:Support",
"@pybind11",
@ -740,23 +728,10 @@ cc_library(
":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(
],
)
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",