[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.

196

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
+
+from . import mlir_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"
 # The default supporting library if the environment variable is not provided.
 _DEFAULT_SUPPORTLIB = "libmlir_c_runner_utils.so"
 
 # The JIT compiler optimization level.
 _OPT_LEVEL = 2
 # The entry point to the JIT compiled program.
 _ENTRY_NAME = "main"
 
 
 @functools.lru_cache()
 def _get_support_lib_name() -> str:
   """Gets the string name for the supporting C shared library."""
   return os.getenv(_SUPPORTLIB_ENV_VAR, _DEFAULT_SUPPORTLIB)
 
 
[e-mail]_cache()
+def _get_sparse_compiler() -> mlir_sparse_compiler.SparseCompiler:
+  """Gets the MLIR sparse compiler with default setting."""
+  return mlir_sparse_compiler.SparseCompiler(
+      options="", opt_level=_OPT_LEVEL, shared_libs=[_get_support_lib_name()])
+
+
 def _record_support_funcs(
     ty: np.dtype, to_func: _SupportFunc, from_func: _SupportFunc,
     ty_to_funcs: Dict[np.dtype, Tuple[_SupportFunc, _SupportFunc]]) -> None:
   """Records the two supporting functions for a given data type."""
   to_func.restype = ctypes.c_void_p
   from_func.restype = ctypes.c_void_p
   ty_to_funcs[ty] = (to_func, from_func)
 
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()])
+  return _get_sparse_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.

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_sparse_compiler.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 sparse compiler class. It is copied from
+#  test/Integration/Dialect/SparseTensor/python/ until we have a better
+#  solution.
+
+from mlir import all_passes_registration
+from mlir import execution_engine
+from mlir import ir
+from mlir import passmanager
+from typing import Sequence
+
+
+class SparseCompiler:
+  """Sparse compiler class for compiling and building MLIR modules."""
+
+  def __init__(self, options: str, opt_level: int, shared_libs: Sequence[str]):
+    pipeline = f'sparse-compiler{{{options} reassociate-fp-reductions=1 enable-index-optimizations=1}}'
+    self.pipeline = pipeline
+    self.opt_level = opt_level
+    self.shared_libs = shared_libs
+
+  def __call__(self, module: ir.Module):
+    """Convenience application method."""
+    self.compile(module)
+
+  def compile(self, module: ir.Module):
+    """Compiles the module by invoking the sparse copmiler pipeline."""
+    passmanager.PassManager.parse(self.pipeline).run(module)
+
+  def jit(self, module: ir.Module) -> execution_engine.ExecutionEngine:
+    """Wraps the module in a JIT execution engine."""
+    return execution_engine.ExecutionEngine(
+        module, opt_level=self.opt_level, shared_libs=self.shared_libs)
+
+  def compile_and_jit(self,
+                      module: ir.Module) -> execution_engine.ExecutionEngine:
+    """Compiles and jits the module."""
+    self.compile(module)
+    return self.jit(module)