[mlir][sparse] refactored python setup of sparse compiler

[bixia]

Nice!

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 sys
 
 from mlir import ir
 from mlir import runtime as rt
-from mlir import execution_engine
 
 from mlir.dialects import sparse_tensor as st
 from mlir.dialects import builtin
 from mlir.dialects import func
 from mlir.dialects.linalg.opdsl import lang as dsl
 
 _SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(_SCRIPT_PATH)
%0 = call @sddmm(%a, %b, %s, %c) : (tensor<8x8xf64>,
                                       tensor<8x8xf64>,
                                       tensor<8x8xf64, {attr}>,
                                       tensor<8x8xf64>) -> tensor<8x8xf64>
   return %0 : tensor<8x8xf64>
 }}
 """
 
 
-def build_compile_and_run_SDDMMM(attr: st.EncodingAttr, opt: str,
-                                 support_lib: str, compiler):
+def build_compile_and_run_SDDMMM(attr: st.EncodingAttr, compiler):
   # Build.
   module = build_SDDMM(attr)
   func = str(module.operation.regions[0].blocks[0].operations[0].operation)
   module = ir.Module.parse(func + boilerplate(attr))
 
   # Compile.
-  compiler(module)
-  engine = execution_engine.ExecutionEngine(
-      module, opt_level=0, shared_libs=[support_lib])
+  engine = compiler.compile_and_jit(module)
 
   # Set up numpy input and buffer for output.
   a = np.array([[1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1, 8.1],
                 [1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2, 8.2],
                 [1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3, 8.3],
                 [1.4, 2.4, 3.4, 4.4, 5.4, 6.4, 7.4, 8.4],
                 [1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5],
                 [1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6, 8.6],
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 = sparse_compiler.SparseCompiler(options=opt)
-                  build_compile_and_run_SDDMMM(attr, opt, support_lib, compiler)
+                  compiler = sparse_compiler.SparseCompiler(
+                      options=opt, opt_level=0, shared_libs=[support_lib])
+                  build_compile_and_run_SDDMMM(attr, 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 sys
 
 from mlir import ir
 from mlir import runtime as rt
-from mlir import execution_engine
 
 from mlir.dialects import sparse_tensor as st
 from mlir.dialects import builtin
 from mlir.dialects import func
 from mlir.dialects.linalg.opdsl import lang as dsl
 
 _SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(_SCRIPT_PATH)
func @main(%ad: tensor<3x4xf64>, %b: tensor<4x2xf64>, %c: tensor<3x2xf64>) -> tensor<3x2xf64>
   %0 = call @spMxM(%a, %b, %c) : (tensor<3x4xf64, {attr}>,
                                   tensor<4x2xf64>,
                                   tensor<3x2xf64>) -> tensor<3x2xf64>
   return %0 : tensor<3x2xf64>
 }}
 """
 
 
-def build_compile_and_run_SpMM(attr: st.EncodingAttr, support_lib: str,
-                               compiler):
+def build_compile_and_run_SpMM(attr: st.EncodingAttr, compiler):
   # Build.
   module = build_SpMM(attr)
   func = str(module.operation.regions[0].blocks[0].operations[0].operation)
   module = ir.Module.parse(func + boilerplate(attr))
 
   # Compile.
-  compiler(module)
-  engine = execution_engine.ExecutionEngine(
-      module, opt_level=0, shared_libs=[support_lib])
+  engine = compiler.compile_and_jit(module)
 
   # Set up numpy input and buffer for output.
   a = np.array(
       [[1.1, 0.0, 0.0, 1.4], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 3.3, 0.0]],
       np.float64)
   b = np.array([[1.0, 2.0], [4.0, 3.0], [5.0, 6.0], [8.0, 7.0]], np.float64)
   c = np.zeros((3, 2), np.float64)
 
levels = [[st.DimLevelType.dense, st.DimLevelType.dense],
               [st.DimLevelType.dense, st.DimLevelType.compressed],
               [st.DimLevelType.compressed, st.DimLevelType.dense],
               [st.DimLevelType.compressed, st.DimLevelType.compressed]]
     orderings = [
         ir.AffineMap.get_permutation([0, 1]),
         ir.AffineMap.get_permutation([1, 0])
     ]
     bitwidths = [0]
+    compiler = sparse_compiler.SparseCompiler(
+        options=opt, opt_level=0, shared_libs=[support_lib])
     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 = sparse_compiler.SparseCompiler(options=opt)
-            build_compile_and_run_SpMM(attr, support_lib, compiler)
+            build_compile_and_run_SpMM(attr, 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
 
 from mlir import ir
 from mlir import runtime as rt
-from mlir import execution_engine
 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
func @main(%ad: tensor<3x4xf64>, %bd: tensor<3x4xf64>) -> tensor<3x4xf64, #DCSR>
   %0 = call @sparse_add_elt(%a, %b) : (tensor<3x4xf64, #DCSR>, tensor<3x4xf64, #DCSR>) -> tensor<3x4xf64, #DCSR>
   return %0 : tensor<3x4xf64, #DCSR>
 }
 """
 
 
 def _run_test(support_lib, kernel):
   """Compiles, runs and checks results."""
+  compiler = sparse_compiler.SparseCompiler(
+      options='', opt_level=2, shared_libs=[support_lib])
   module = ir.Module.parse(kernel)
-  sparse_compiler.SparseCompiler(options='')(module)
-  engine = execution_engine.ExecutionEngine(
-      module, opt_level=0, shared_libs=[support_lib])
+  engine = compiler.compile_and_jit(module)
 
   # 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(
       [[1.1, 0.0, 0.0, 2.8], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]],
       np.float64)

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
 
-from mlir import execution_engine
 from mlir import ir
 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)

 1 1 1
 1 10 3
 2 2 2
 5 5 5
 10 1 4
 """
 
 
-def build_compile_and_run_output(attr: st.EncodingAttr, support_lib: str,
-                                 compiler):
+def build_compile_and_run_output(attr: st.EncodingAttr, compiler):
   # Build and Compile.
   module = ir.Module.parse(boilerplate(attr))
-  compiler(module)
-  engine = execution_engine.ExecutionEngine(
-      module, opt_level=0, shared_libs=[support_lib])
+  engine = compiler.compile_and_jit(module)
 
   # Invoke the kernel and compare output.
   with tempfile.TemporaryDirectory() as test_dir:
     out = os.path.join(test_dir, 'out.tns')
     buf = out.encode('utf-8')
     mem_a = ctypes.pointer(ctypes.pointer(ctypes.create_string_buffer(buf)))
     engine.invoke('main', mem_a)
 
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]
+    compiler = sparse_compiler.SparseCompiler(
+        options='', opt_level=2, shared_libs=[support_lib])
     for level in levels:
       for ordering in orderings:
         for bwidth in bitwidths:
           attr = st.EncodingAttr.get(level, ordering, bwidth, bwidth)
-          compiler = sparse_compiler.SparseCompiler(options='')
-          build_compile_and_run_output(attr, support_lib, compiler)
+          build_compile_and_run_output(attr, 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
 
 from mlir import ir
 from mlir import runtime as rt
-from mlir.execution_engine import ExecutionEngine
 
 from mlir.dialects import builtin
 from mlir.dialects import func
 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
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, support_lib: str):
+  def compile(self, compiler):
     """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'
-    compiler(self._module)
-    self._engine = ExecutionEngine(
-        self._module, opt_level=0, shared_libs=[support_lib])
+    self._engine = compiler.compile_and_jit(self._module)
     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(
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 = sparse_compiler.SparseCompiler(options=sparsification_options)
+    compiler = sparse_compiler.SparseCompiler(
+        options=sparsification_options, opt_level=0, shared_libs=[support_lib])
     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():
     #
     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, support_lib).writeTo(
-                    sys.argv[2] if len(sys.argv) > 2 else None).run(np_arg0))
+            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))
     # 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.
 
 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 definition."""
+  """Sparse compiler class for compiling and building MLIR modules."""
 
-  def __init__(self, options: str):
+  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)