diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py --- a/mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py +++ b/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 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 @@ 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') diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py --- a/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py +++ b/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 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 @@ 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') diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py --- a/mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py +++ b/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( diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_output.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_output.py --- a/mlir/test/Integration/Dialect/SparseTensor/python/test_output.py +++ b/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 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 @@ 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 diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py --- a/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py +++ b/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 @@ 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 @@ 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 @@ 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') diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/tools/sparse_compiler.py b/mlir/test/Integration/Dialect/SparseTensor/python/tools/sparse_compiler.py --- a/mlir/test/Integration/Dialect/SparseTensor/python/tools/sparse_compiler.py +++ b/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)