From 58d12332a4730226d0a640136f06f3bd0861f1a5 Mon Sep 17 00:00:00 2001 From: Aart Bik Date: Wed, 12 May 2021 14:51:16 -0700 Subject: [PATCH] [mlir][sparse][capi][python] add sparse tensor passes First set of "boilerplate" to get sparse tensor passes available through CAPI and Python. Reviewed By: stellaraccident Differential Revision: https://reviews.llvm.org/D102362 --- mlir/include/mlir-c/Dialect/SparseTensor.h | 2 ++ .../SparseTensor/Transforms/CMakeLists.txt | 2 ++ mlir/lib/Bindings/Python/CMakeLists.txt | 8 ++++++ .../Bindings/Python/SparseTensorPasses.cpp | 22 ++++++++++++++++ mlir/lib/CAPI/Dialect/CMakeLists.txt | 2 ++ mlir/lib/CAPI/Dialect/SparseTensorPasses.cpp | 26 +++++++++++++++++++ mlir/python/mlir/dialects/sparse_tensor.py | 5 ++++ .../python/dialects/sparse_tensor/passes.py | 22 ++++++++++++++++ 8 files changed, 89 insertions(+) create mode 100644 mlir/lib/Bindings/Python/SparseTensorPasses.cpp create mode 100644 mlir/lib/CAPI/Dialect/SparseTensorPasses.cpp create mode 100644 mlir/test/python/dialects/sparse_tensor/passes.py diff --git a/mlir/include/mlir-c/Dialect/SparseTensor.h b/mlir/include/mlir-c/Dialect/SparseTensor.h index 2615a1655c15..16d932c16bdb 100644 --- a/mlir/include/mlir-c/Dialect/SparseTensor.h +++ b/mlir/include/mlir-c/Dialect/SparseTensor.h @@ -74,4 +74,6 @@ mlirSparseTensorEncodingAttrGetIndexBitWidth(MlirAttribute attr); } #endif +#include "mlir/Dialect/SparseTensor/Transforms/Passes.capi.h.inc" + #endif // MLIR_C_DIALECT_SPARSE_TENSOR_H diff --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/CMakeLists.txt b/mlir/include/mlir/Dialect/SparseTensor/Transforms/CMakeLists.txt index 2d9eaa567c8f..852b76ad12dd 100644 --- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/CMakeLists.txt +++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/CMakeLists.txt @@ -1,5 +1,7 @@ set(LLVM_TARGET_DEFINITIONS Passes.td) mlir_tablegen(Passes.h.inc -gen-pass-decls -name SparseTensor) +mlir_tablegen(Passes.capi.h.inc -gen-pass-capi-header --prefix SparseTensor) +mlir_tablegen(Passes.capi.cpp.inc -gen-pass-capi-impl --prefix SparseTensor) add_public_tablegen_target(MLIRSparseTensorPassIncGen) add_mlir_doc(Passes SparseTensorPasses ./ -gen-pass-doc) diff --git a/mlir/lib/Bindings/Python/CMakeLists.txt b/mlir/lib/Bindings/Python/CMakeLists.txt index 7dc1f64b4f57..575b9dbbd4ed 100644 --- a/mlir/lib/Bindings/Python/CMakeLists.txt +++ b/mlir/lib/Bindings/Python/CMakeLists.txt @@ -33,6 +33,14 @@ add_mlir_python_extension(MLIRAsyncPassesBindingsPythonExtension _mlirAsyncPasse ) add_dependencies(MLIRBindingsPythonExtension MLIRAsyncPassesBindingsPythonExtension) +add_mlir_python_extension(MLIRSparseTensorPassesBindingsPythonExtension _mlirSparseTensorPasses + INSTALL_DIR + python + SOURCES + SparseTensorPasses.cpp +) +add_dependencies(MLIRBindingsPythonExtension MLIRSparseTensorPassesBindingsPythonExtension) + add_mlir_python_extension(MLIRGPUPassesBindingsPythonExtension _mlirGPUPasses INSTALL_DIR python diff --git a/mlir/lib/Bindings/Python/SparseTensorPasses.cpp b/mlir/lib/Bindings/Python/SparseTensorPasses.cpp new file mode 100644 index 000000000000..2a8e2b802df9 --- /dev/null +++ b/mlir/lib/Bindings/Python/SparseTensorPasses.cpp @@ -0,0 +1,22 @@ +//===- SparseTensorPasses.cpp - Pybind module for the SparseTensor passes -===// +// +// 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 + +// ----------------------------------------------------------------------------- +// Module initialization. +// ----------------------------------------------------------------------------- + +PYBIND11_MODULE(_mlirSparseTensorPasses, m) { + m.doc() = "MLIR SparseTensor Dialect Passes"; + + // Register all SparseTensor passes on load. + mlirRegisterSparseTensorPasses(); +} diff --git a/mlir/lib/CAPI/Dialect/CMakeLists.txt b/mlir/lib/CAPI/Dialect/CMakeLists.txt index 69371fd5718f..053fce30d5dd 100644 --- a/mlir/lib/CAPI/Dialect/CMakeLists.txt +++ b/mlir/lib/CAPI/Dialect/CMakeLists.txt @@ -62,11 +62,13 @@ add_mlir_public_c_api_library(MLIRCAPIShape add_mlir_public_c_api_library(MLIRCAPISparseTensor SparseTensor.cpp + SparseTensorPasses.cpp PARTIAL_SOURCES_INTENDED LINK_LIBS PUBLIC MLIRCAPIIR MLIRSparseTensor + MLIRSparseTensorTransforms ) add_mlir_public_c_api_library(MLIRCAPIStandard diff --git a/mlir/lib/CAPI/Dialect/SparseTensorPasses.cpp b/mlir/lib/CAPI/Dialect/SparseTensorPasses.cpp new file mode 100644 index 000000000000..5b2ba4ca7aa3 --- /dev/null +++ b/mlir/lib/CAPI/Dialect/SparseTensorPasses.cpp @@ -0,0 +1,26 @@ +//===- SparseTensorPasses.cpp - C API for SparseTensor Dialect Passes -----===// +// +// 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/CAPI/Pass.h" +#include "mlir/Dialect/SparseTensor/Transforms/Passes.h" +#include "mlir/Pass/Pass.h" + +// Must include the declarations as they carry important visibility attributes. +#include "mlir/Dialect/SparseTensor/Transforms/Passes.capi.h.inc" + +using namespace mlir; + +#ifdef __cplusplus +extern "C" { +#endif + +#include "mlir/Dialect/SparseTensor/Transforms/Passes.capi.cpp.inc" + +#ifdef __cplusplus +} +#endif diff --git a/mlir/python/mlir/dialects/sparse_tensor.py b/mlir/python/mlir/dialects/sparse_tensor.py index 687d1e6cdbc7..59fd86021dc6 100644 --- a/mlir/python/mlir/dialects/sparse_tensor.py +++ b/mlir/python/mlir/dialects/sparse_tensor.py @@ -3,5 +3,10 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception from .._cext_loader import _reexport_cext +from .._cext_loader import _load_extension + _reexport_cext("dialects.sparse_tensor", __name__) +_cextSparseTensorPasses = _load_extension("_mlirSparseTensorPasses") + del _reexport_cext +del _load_extension diff --git a/mlir/test/python/dialects/sparse_tensor/passes.py b/mlir/test/python/dialects/sparse_tensor/passes.py new file mode 100644 index 000000000000..f3510c5ce816 --- /dev/null +++ b/mlir/test/python/dialects/sparse_tensor/passes.py @@ -0,0 +1,22 @@ +# RUN: %PYTHON %s | FileCheck %s + +from mlir.ir import * +from mlir.passmanager import * + +from mlir.dialects import sparse_tensor as st + + +def run(f): + print('\nTEST:', f.__name__) + f() + return f + + +# CHECK-LABEL: TEST: testSparseTensorPass +@run +def testSparseTensorPass(): + with Context() as context: + PassManager.parse('sparsification') + PassManager.parse('sparse-tensor-conversion') + # CHECK: SUCCESS + print('SUCCESS')