[mlir][sparse] refactor sparse compiler pipeline to single place

[aartbik]

rebased with main head

[bixia]

Nice!

mlir/test/Integration/Dialect/SparseTensor/python/tools/sparse_compiler.py

7

Shall we write this as
from mlir import all_passes_registration
so it is consistent with the following two lines?

If we will make this change, we can also move this line to the same block with the following two lines.

[aartbik]

mlir/test/Integration/Dialect/SparseTensor/python/tools/sparse_compiler.py

7

Yes, looks nicer. Done.

[aartbik]

addressed comment

mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py

 # RUN: SUPPORT_LIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext \
 # RUN:   %PYTHON %s | FileCheck %s
 
 import ctypes
 import numpy as np
 import os
-
-import mlir.all_passes_registration
+import sys
 
 from mlir import ir
 from mlir import runtime as rt
 from mlir import execution_engine
-from mlir import passmanager
 
 from mlir.dialects import sparse_tensor as st
 from mlir.dialects import builtin
 from mlir.dialects.linalg.opdsl import lang as dsl
 
+_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(_SCRIPT_PATH)
+from tools import sparse_compiler
 
 @dsl.linalg_structured_op
 def sddmm_dsl(
     A=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.K),
     B=dsl.TensorDef(dsl.T, dsl.S.K, dsl.S.N),
     S=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.N),
     C=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.N, output=True)):
   C[dsl.D.m,
def build_compile_and_run_SDDMMM(attr: st.EncodingAttr, opt: str,
   expected[4, 1] = 3.0 * full_matmul[4, 1]
   c = rt.ranked_memref_to_numpy(mem_out[0])
   if np.allclose(c, expected):
     pass
   else:
     quit(f'FAILURE')
 
 
-class SparseCompiler:
-  """Sparse compiler passes."""
-
-  def __init__(self, options: str):
-    pipeline = (
-        f'sparse-compiler{{{options} reassociate-fp-reductions=1 enable-index-optimizations=1}}')
-    self.pipeline = pipeline
-
-  def __call__(self, module: ir.Module):
-    passmanager.PassManager.parse(self.pipeline).run(module)
-
-
 def main():
   support_lib = os.getenv('SUPPORT_LIB')
   assert support_lib is not None, 'SUPPORT_LIB is undefined'
   if not os.path.exists(support_lib):
     raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT),
                             support_lib)
 
   # CHECK-LABEL: TEST: testSDDMMM
for level in levels:
             for par in [0]:
               for vec in [0, 1]:
                 for e in [True]:
                   vl = 1 if vec == 0 else 16
                   attr = st.EncodingAttr.get(level, ordering, pwidth, iwidth)
                   opt = (f'parallelization-strategy={par} '
                          f'vectorization-strategy={vec} '
                          f'vl={vl} enable-simd-index32={e}')
-                  compiler = SparseCompiler(options=opt)
+                  compiler = sparse_compiler.SparseCompiler(options=opt)
                   build_compile_and_run_SDDMMM(attr, opt, support_lib, compiler)
                   count = count + 1
   # CHECK: Passed 16 tests
   print('Passed ', count, 'tests')
 
 
 if __name__ == '__main__':
   main()

mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py

 # RUN: SUPPORT_LIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext \
 # RUN:   %PYTHON %s | FileCheck %s
 
 import ctypes
 import numpy as np
 import os
-
-import mlir.all_passes_registration
+import sys
 
 from mlir import ir
 from mlir import runtime as rt
 from mlir import execution_engine
-from mlir import passmanager
 
 from mlir.dialects import sparse_tensor as st
 from mlir.dialects import builtin
 from mlir.dialects.linalg.opdsl import lang as dsl
 
+_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(_SCRIPT_PATH)
+from tools import sparse_compiler
 
 @dsl.linalg_structured_op
 def matmul_dsl(
     A=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.K),
     B=dsl.TensorDef(dsl.T, dsl.S.K, dsl.S.N),
     C=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.N, output=True)):
   C[dsl.D.m, dsl.D.n] += A[dsl.D.m, dsl.D.k] * B[dsl.D.k, dsl.D.n]
 
def build_compile_and_run_SpMM(attr: st.EncodingAttr, support_lib: str,
   expected = np.matmul(a, b);
   c = rt.ranked_memref_to_numpy(mem_out[0])
   if np.allclose(c, expected):
     pass
   else:
     quit(f'FAILURE')
 
 
-class SparseCompiler:
-  """Sparse compiler passes."""
-
-  def __init__(self, options: str):
-    pipeline = (
-        f'sparse-compiler{{{options} reassociate-fp-reductions=1 enable-index-optimizations=1}}')
-    self.pipeline = pipeline
-
-  def __call__(self, module: ir.Module):
-    passmanager.PassManager.parse(self.pipeline).run(module)
-
-
 def main():
   support_lib = os.getenv('SUPPORT_LIB')
   assert support_lib is not None, 'SUPPORT_LIB is undefined'
   if not os.path.exists(support_lib):
     raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), support_lib)
 
   # CHECK-LABEL: TEST: testSpMM
   print('\nTEST: testSpMM')
orderings = [
         ir.AffineMap.get_permutation([1, 0])
     ]
     bitwidths = [0]
     for level in levels:
       for ordering in orderings:
         for pwidth in bitwidths:
           for iwidth in bitwidths:
             attr = st.EncodingAttr.get(level, ordering, pwidth, iwidth)
-            compiler = SparseCompiler(options=opt)
+            compiler = sparse_compiler.SparseCompiler(options=opt)
             build_compile_and_run_SpMM(attr, support_lib, compiler)
             count = count + 1
     # CHECK: Passed 8 tests
     print('Passed ', count, 'tests')
 
 
 if __name__ == '__main__':
   main()

mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py

 # RUN: SUPPORT_LIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
 
 import ctypes
 import numpy as np
 import os
 import sys
 
-import mlir.all_passes_registration
-
 from mlir import ir
 from mlir import runtime as rt
 from mlir import execution_engine
-from mlir import passmanager
 from mlir.dialects import sparse_tensor as st
 from mlir.dialects import builtin
 from mlir.dialects.linalg.opdsl import lang as dsl
 
 _SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(_SCRIPT_PATH)
 from tools import np_to_sparse_tensor as test_tools
+from tools import sparse_compiler
 
 # TODO: Use linalg_structured_op to generate the kernel after making it to
 # handle sparse tensor outputs.
 _KERNEL_STR = """
 #DCSR = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed" ]
 }>
 
func @main(%ad: tensor<3x4xf64>, %bd: tensor<3x4xf64>) -> tensor<3x4xf64, #DCSR>
   %a = sparse_tensor.convert %ad : tensor<3x4xf64> to tensor<3x4xf64, #DCSR>
   %b = sparse_tensor.convert %bd : tensor<3x4xf64> to tensor<3x4xf64, #DCSR>
   %0 = call @sparse_add_elt(%a, %b) : (tensor<3x4xf64, #DCSR>, tensor<3x4xf64, #DCSR>) -> tensor<3x4xf64, #DCSR>
   return %0 : tensor<3x4xf64, #DCSR>
 }
 """
 
 
-class _SparseCompiler:
-  """Sparse compiler passes."""
-
-  def __init__(self):
-    self.pipeline = (
-        f'sparse-compiler{{reassociate-fp-reductions=1 enable-index-optimizations=1}}')
-
-  def __call__(self, module: ir.Module):
-    passmanager.PassManager.parse(self.pipeline).run(module)
-
-
 def _run_test(support_lib, kernel):
   """Compiles, runs and checks results."""
   module = ir.Module.parse(kernel)
-  _SparseCompiler()(module)
+  sparse_compiler.SparseCompiler(options='')(module)
   engine = execution_engine.ExecutionEngine(
       module, opt_level=0, shared_libs=[support_lib])
 
   # Set up numpy inputs and buffer for output.
   a = np.array(
       [[1.1, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 6.6, 0.0]],
       np.float64)
   b = np.array(

mlir/test/Integration/Dialect/SparseTensor/python/test_output.py

 # RUN: SUPPORT_LIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext \
 # RUN:   %PYTHON %s | FileCheck %s
 
 import ctypes
 import os
+import sys
 import tempfile
 
-import mlir.all_passes_registration
-
 from mlir import execution_engine
 from mlir import ir
-from mlir import passmanager
 from mlir import runtime as rt
 
 from mlir.dialects import builtin
 from mlir.dialects import sparse_tensor as st
 
+_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(_SCRIPT_PATH)
+from tools import sparse_compiler
 
 # TODO: move more into actual IR building.
 def boilerplate(attr: st.EncodingAttr):
   """Returns boilerplate main method."""
   return f"""
 func @main(%p : !llvm.ptr<i8>) -> () attributes {{ llvm.emit_c_interface }} {{
   %d = arith.constant sparse<[[0, 0], [1, 1], [0, 9], [9, 0], [4, 4]],
                              [1.0, 2.0, 3.0, 4.0, 5.0]> : tensor<10x10xf64>
with tempfile.TemporaryDirectory() as test_dir:
     mem_a = ctypes.pointer(ctypes.pointer(ctypes.create_string_buffer(buf)))
     engine.invoke('main', mem_a)
 
     actual = open(out).read()
     if actual != expected():
       quit('FAILURE')
 
 
-class SparseCompiler:
-  """Sparse compiler passes."""
-
-  def __init__(self):
-    pipeline = (
-        f'sparse-compiler{{reassociate-fp-reductions=1 enable-index-optimizations=1}}')
-    self.pipeline = pipeline
-
-  def __call__(self, module: ir.Module):
-    passmanager.PassManager.parse(self.pipeline).run(module)
-
-
 def main():
   support_lib = os.getenv('SUPPORT_LIB')
   assert support_lib is not None, 'SUPPORT_LIB is undefined'
   if not os.path.exists(support_lib):
     raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT),
                             support_lib)
 
   # CHECK-LABEL: TEST: test_output
   print('\nTEST: test_output')
   count = 0
   with ir.Context() as ctx, ir.Location.unknown():
     # Loop over various sparse types: CSR, DCSR, CSC, DCSC.
     levels = [[st.DimLevelType.dense, st.DimLevelType.compressed],
               [st.DimLevelType.compressed, st.DimLevelType.compressed]]
     orderings = [
         ir.AffineMap.get_permutation([0, 1]),
         ir.AffineMap.get_permutation([1, 0])
     ]
     bitwidths = [8, 16, 32, 64]
     for level in levels:
       for ordering in orderings:
         for bwidth in bitwidths:
           attr = st.EncodingAttr.get(level, ordering, bwidth, bwidth)
-          compiler = SparseCompiler()
+          compiler = sparse_compiler.SparseCompiler(options='')
           build_compile_and_run_output(attr, support_lib, compiler)
           count = count + 1
 
   # CHECK: Passed 16 tests
   print('Passed', count, 'tests')
 
 
 if __name__ == '__main__':
   main()

mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py

 # RUN: SUPPORT_LIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext \
 # RUN:   %PYTHON %s | FileCheck %s
 
 import ctypes
 import errno
 import itertools
 import os
 import sys
+
 from typing import List, Callable
 
 import numpy as np
 
-import mlir.all_passes_registration
-
 from mlir import ir
 from mlir import runtime as rt
 from mlir.execution_engine import ExecutionEngine
-from mlir.passmanager import PassManager
 
 from mlir.dialects import builtin
 from mlir.dialects import std
 from mlir.dialects import sparse_tensor as st
 
+_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(_SCRIPT_PATH)
+from tools import sparse_compiler
+
 # ===----------------------------------------------------------------------=== #
 
 # TODO: move this boilerplate to its own module, so it can be used by
 # other tests and programs.
 class TypeConverter:
   """Converter between NumPy types and MLIR types."""
 
   def __init__(self, context: ir.Context):
if filename is None:
       # Silent no-op, for convenience.
       return self
     if os.path.exists(filename):
       raise FileExistsError(errno.EEXIST, os.strerror(errno.EEXIST), filename)
     with open(filename, 'w') as f:
       f.write(str(self._module))
     return self
 
-  def compile(self, compiler: Callable[[ir.Module], ExecutionEngine]):
+  def compile(self, compiler, support_lib: str):
     """Compile the ir.Module."""
     assert self._module is not None, \
         'StressTest: must call build() before compile()'
     assert self._engine is None, \
         'StressTest: must not call compile() repeatedly'
-    self._engine = compiler(self._module)
+    compiler(self._module)
+    self._engine = ExecutionEngine(
+        self._module, opt_level=0, shared_libs=[support_lib])
     return self
 
   def run(self, np_arg0: np.ndarray) -> np.ndarray:
     """Runs the test on the given numpy array, and returns the resulting
     numpy array."""
     assert self._engine is not None, \
         'StressTest: must call compile() before run()'
     self._assertEqualsRoundtripTp(
         self._tyconv.get_RankedTensorType_of_nparray(np_arg0))
     np_out = np.zeros(np_arg0.shape, dtype=np_arg0.dtype)
     self._assertEqualsRoundtripTp(
         self._tyconv.get_RankedTensorType_of_nparray(np_out))
     mem_arg0 = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(np_arg0)))
     mem_out = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(np_out)))
     self._engine.invoke('main', mem_out, mem_arg0)
     return rt.ranked_memref_to_numpy(mem_out[0])
 
 # ===----------------------------------------------------------------------=== #
 
-# TODO: move this boilerplate to its own module, so it can be used by
-# other tests and programs.
-class SparseCompiler:
-  """Sparse compiler passes."""
-
-  def __init__(self, sparsification_options: str, support_lib: str):
-    self._support_lib = support_lib
-    self._pipeline = (
-        f'sparse-compiler{{{sparsification_options} reassociate-fp-reductions=1 enable-index-optimizations=1}}')
-    # Must be in the scope of a `with ir.Context():`
-    self._passmanager = PassManager.parse(self._pipeline)
-
-  def __call__(self, module: ir.Module) -> ExecutionEngine:
-    self._passmanager.run(module)
-    return ExecutionEngine(module, opt_level=0, shared_libs=[self._support_lib])
-
-# ===----------------------------------------------------------------------=== #
-
 def main():
   """
   USAGE: python3 test_stress.py [raw_module.mlir [compiled_module.mlir]]
 
   The environment variable SUPPORT_LIB must be set to point to the
   libmlir_c_runner_utils shared library.  There are two optional
   arguments, for debugging purposes.  The first argument specifies where
   to write out the raw/generated ir.Module.  The second argument specifies
with ir.Context() as ctx, ir.Location.unknown():
     vec = 0
     vl = 1
     e = False
     sparsification_options = (
         f'parallelization-strategy={par} '
         f'vectorization-strategy={vec} '
         f'vl={vl} '
         f'enable-simd-index32={e}')
-    compiler = SparseCompiler(sparsification_options, support_lib)
+    compiler = sparse_compiler.SparseCompiler(options=sparsification_options)
     f64 = ir.F64Type.get()
     # Be careful about increasing this because
     #     len(types) = 1 + 2^rank * rank! * len(bitwidths)^2
     shape = range(2, 6)
     rank = len(shape)
     # All combinations.
     levels = list(itertools.product(*itertools.repeat(
       [st.DimLevelType.dense, st.DimLevelType.compressed], rank)))
with ir.Context() as ctx, ir.Location.unknown():
     # at least hits all n elements of the `types` set.
     #
     tyconv = TypeConverter(ctx)
     size = 1
     for d in shape:
       size *= d
     np_arg0 = np.arange(size, dtype=tyconv.irtype_to_dtype(f64)).reshape(*shape)
     np_out = (
-        StressTest(tyconv)
-        .build(types)
-        .writeTo(sys.argv[1] if len(sys.argv) > 1 else None)
-        .compile(compiler)
-        .writeTo(sys.argv[2] if len(sys.argv) > 2 else None)
-        .run(np_arg0))
+        StressTest(tyconv).build(types).writeTo(
+            sys.argv[1] if len(sys.argv) > 1 else None).compile(
+                compiler, support_lib).writeTo(
+                    sys.argv[2] if len(sys.argv) > 2 else None).run(np_arg0))
     # CHECK: Passed
     if np.allclose(np_out, np_arg0):
       print('Passed')
     else:
       sys.exit('FAILURE')
 
 if __name__ == '__main__':
   main()

mlir/test/Integration/Dialect/SparseTensor/python/tools/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.
+
+import mlir.all_passes_registration
+
+from mlir import ir
+from mlir import passmanager
+
+
+class SparseCompiler:
+  """Sparse compiler definition."""
+
+  def __init__(self, options: str):
+    pipeline = f'sparse-compiler{{{options} reassociate-fp-reductions=1 enable-index-optimizations=1}}'
+    self.pipeline = pipeline
+
+  def __call__(self, module: ir.Module):
+    passmanager.PassManager.parse(self.pipeline).run(module)