[Mlir-commits] [mlir] 58d1233 - [mlir][sparse][capi][python] add sparse tensor passes
Aart Bik
llvmlistbot at llvm.org
Wed May 12 16:41:08 PDT 2021
Author: Aart Bik
Date: 2021-05-12T16:40:50-07:00
New Revision: 58d12332a4730226d0a640136f06f3bd0861f1a5
URL: https://github.com/llvm/llvm-project/commit/58d12332a4730226d0a640136f06f3bd0861f1a5
DIFF: https://github.com/llvm/llvm-project/commit/58d12332a4730226d0a640136f06f3bd0861f1a5.diff
LOG: [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
Added:
mlir/lib/Bindings/Python/SparseTensorPasses.cpp
mlir/lib/CAPI/Dialect/SparseTensorPasses.cpp
mlir/test/python/dialects/sparse_tensor/passes.py
Modified:
mlir/include/mlir-c/Dialect/SparseTensor.h
mlir/include/mlir/Dialect/SparseTensor/Transforms/CMakeLists.txt
mlir/lib/Bindings/Python/CMakeLists.txt
mlir/lib/CAPI/Dialect/CMakeLists.txt
mlir/python/mlir/dialects/sparse_tensor.py
Removed:
################################################################################
diff --git a/mlir/include/mlir-c/Dialect/SparseTensor.h b/mlir/include/mlir-c/Dialect/SparseTensor.h
index 2615a1655c15b..16d932c16bdbc 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 2d9eaa567c8f7..852b76ad12dde 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 7dc1f64b4f57e..575b9dbbd4ed7 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 0000000000000..2a8e2b802df9c
--- /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 <pybind11/pybind11.h>
+
+// -----------------------------------------------------------------------------
+// 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 69371fd5718f9..053fce30d5ddc 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 0000000000000..5b2ba4ca7aa31
--- /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 687d1e6cdbc75..59fd86021dc6e 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 0000000000000..f3510c5ce8169
--- /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
+ at run
+def testSparseTensorPass():
+ with Context() as context:
+ PassManager.parse('sparsification')
+ PassManager.parse('sparse-tensor-conversion')
+ # CHECK: SUCCESS
+ print('SUCCESS')
More information about the Mlir-commits
mailing list