[mlir][sparse][capi][python] add sparse tensor passes

[aartbik]

Summary

First set of "boilerplate" to get sparse tensor
passes available through CAPI and Python.

[stellaraccident]

mlir/test/python/dialects/sparse_tensor/passes.py

15

Generally, add a "return f" to the above run() function. Then use a decorator like:

@run
def testSparseTensorPass():

And put your CHECK-LABEL above the @run and the CHECK above the print.

[aartbik]

mlir/test/python/dialects/sparse_tensor/passes.py

15

Ah, I did not see that style in your test file. Changed.
Thanks!

[aartbik]

use @run

mlir/include/mlir-c/Dialect/SparseTensor.h

 /// Returns the index bit width in a sparse_tensor.encoding attribute.
 MLIR_CAPI_EXPORTED int
 mlirSparseTensorEncodingAttrGetIndexBitWidth(MlirAttribute attr);
 
 #ifdef __cplusplus
 }
 #endif
 
+#include "mlir/Dialect/SparseTensor/Transforms/Passes.capi.h.inc"
+
 #endif // MLIR_C_DIALECT_SPARSE_TENSOR_H

mlir/include/mlir/Dialect/SparseTensor/Transforms/CMakeLists.txt

 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)

mlir/lib/Bindings/Python/CMakeLists.txt

 add_mlir_python_extension(MLIRAsyncPassesBindingsPythonExtension _mlirAsyncPasses
   INSTALL_DIR
     python
   SOURCES
     AsyncPasses.cpp
 )
 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
   SOURCES
     GPUPasses.cpp
 )
 add_dependencies(MLIRBindingsPythonExtension MLIRGPUPassesBindingsPythonExtension)
 
 add_mlir_python_extension(MLIRLinalgPassesBindingsPythonExtension _mlirLinalgPasses
   INSTALL_DIR
     python
   SOURCES
     LinalgPasses.cpp
 )
 add_dependencies(MLIRBindingsPythonExtension MLIRLinalgPassesBindingsPythonExtension)

mlir/lib/Bindings/Python/SparseTensorPasses.cpp

+//===- 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();
+}

mlir/lib/CAPI/Dialect/CMakeLists.txt

add_mlir_public_c_api_library(MLIRCAPIShape
   PARTIAL_SOURCES_INTENDED
   LINK_LIBS PUBLIC
   MLIRCAPIIR
   MLIRShape
 )
 
 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
   Standard.cpp
 
   PARTIAL_SOURCES_INTENDED
   LINK_LIBS PUBLIC
   MLIRCAPIIR

mlir/lib/CAPI/Dialect/SparseTensorPasses.cpp

+//===- 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

mlir/python/mlir/dialects/sparse_tensor.py

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

mlir/test/python/dialects/sparse_tensor/passes.py

+# 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')