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

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D123419

Author: aartbik

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

@@ -8,7 +8,6 @@
 
 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
@@ -69,17 +68,14 @@ def boilerplate(attr: st.EncodingAttr):
 """
 
 
-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],
@@ -156,8 +152,9 @@ def main():
                   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')

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

@@ -8,7 +8,6 @@
 
 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
@@ -69,17 +68,14 @@ def boilerplate(attr: st.EncodingAttr):
 """
 
 
-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(
@@ -140,13 +136,14 @@ def main():
         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')

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

@@ -7,7 +7,6 @@
 
 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
@@ -61,10 +60,10 @@
 
 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(

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

@@ -6,7 +6,6 @@
 import sys
 import tempfile
 
-from mlir import execution_engine
 from mlir import ir
 from mlir import runtime as rt
 
@@ -49,13 +48,10 @@ def expected():
 """
 
 
-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:
@@ -88,12 +84,13 @@ def main():
         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

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

@@ -13,7 +13,6 @@
 
 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
@@ -139,15 +138,13 @@ def writeTo(self, filename):
       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:
@@ -194,7 +191,8 @@ def main():
         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
@@ -230,9 +228,8 @@ def main():
     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')

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

@@ -5,15 +5,35 @@
 #  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)