[mlir][sparse][taco] Use the SparseCompiler from python/tools.

[bixia]

Summary

Copy the implementation of SparseCompiler from python/tools to taco/tools until we have a common place to install it. Modify TACO to use this SparseCompiler for compilation and jitting.

[bixia]

Rebase.

[aartbik]

Thanks for following up so quickly with this!

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py

21–25

I am also okay if you simply copy the sparse_compiler.py into mlir_sparse_compiler.mlir of this directory to avoid all this ugliness. It is not that much code duplication, and makes this directory self contained.

195

shall we only set up a sparse_compiler object once, and reuse it between all invocations?
wdyt?

[bixia]

Address review comment.

[bixia]

Removed an import line.

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.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
 
 # This file contains the utilities to process sparse tensor outputs.
 
 from typing import Callable, Dict, Sequence, Tuple
 import ctypes
 import functools
 import numpy as np
 import os
+import sys
 
 # Import MLIR related modules.
 from mlir import all_passes_registration  # Register MLIR compiler passes.
 from mlir import execution_engine
 from mlir import ir
 from mlir import runtime
 from mlir.dialects import sparse_tensor
-from mlir.passmanager import PassManager
+
+_TOOL_PATH = os.path.join(
+    os.path.join(
+        os.path.dirname(
+            os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
+        "python"), "tools")
+sys.path.append(_TOOL_PATH)
+import sparse_compiler
 
 # Type aliases for type annotation.
 _SupportFunc = Callable[..., None]
 _SupportFuncLocator = Callable[[np.dtype], Tuple[_SupportFunc, _SupportFunc]]
 
 # The name for the environment variable that provides the full path for the
 # supporting library.
 _SUPPORTLIB_ENV_VAR = "SUPPORTLIB"
def compile_and_build_engine(
 
   Args:
     module: The MLIR module.
 
   Returns:
     A JIT execution engine for the MLIR module.
 
   """
-  pipeline = f"sparse-compiler"
-  PassManager.parse(pipeline).run(module)
-  return execution_engine.ExecutionEngine(
-      module, opt_level=_OPT_LEVEL, shared_libs=[_get_support_lib_name()])
+  compiler = sparse_compiler.SparseCompiler(
+      options="", opt_level=_OPT_LEVEL, shared_libs=[_get_support_lib_name()])
+  return compiler.compile_and_jit(module)
 
 
 class _SparseTensorDescriptor(ctypes.Structure):
   """A C structure for an MLIR sparse tensor."""
   _fields_ = [
       # A pointer for the MLIR sparse tensor storage.
       ("storage", ctypes.POINTER(ctypes.c_ulonglong)),
       # An MLIR MemRef descriptor for the shape of the sparse tensor.