[Mlir-commits] [mlir] 28063a2 - [mlir][sparse] refactored python setup of sparse compiler
Aart Bik
llvmlistbot at llvm.org
Tue Apr 12 11:58:48 PDT 2022
Author: Aart Bik
Date: 2022-04-12T11:58:41-07:00
New Revision: 28063a281b18a9a94a7d56815db729b6c1a087ea
URL: https://github.com/llvm/llvm-project/commit/28063a281b18a9a94a7d56815db729b6c1a087ea
DIFF: https://github.com/llvm/llvm-project/commit/28063a281b18a9a94a7d56815db729b6c1a087ea.diff
LOG: [mlir][sparse] refactored python setup of sparse compiler
Reviewed By: bixia
Differential Revision: https://reviews.llvm.org/D123419
Added:
Modified:
mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py
mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py
mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py
mlir/test/Integration/Dialect/SparseTensor/python/test_output.py
mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py
mlir/test/Integration/Dialect/SparseTensor/python/tools/sparse_compiler.py
Removed:
################################################################################
diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py
index 353cd58d23e83..d8dde93e96d95 100644
--- 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 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')
diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py
index 4c51e19c31aed..5fda4ee32e1f4 100644
--- 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 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')
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
index 1cc79c43f728a..91de742a83cad 100644
--- 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
index 5e2210b2d81ba..324561cd725c9 100644
--- 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 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
diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py
index 632619581f218..76dfd3cf145fc 100644
--- 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 @@ 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')
diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/tools/sparse_compiler.py b/mlir/test/Integration/Dialect/SparseTensor/python/tools/sparse_compiler.py
index 47b145ff3cb3c..82c40625bcc44 100644
--- 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)
More information about the Mlir-commits
mailing list