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(
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'
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
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(
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'
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
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.
"""
-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])
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):
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'
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
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
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:
# ===----------------------------------------------------------------------=== #
-# 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]]
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
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')
--- /dev/null
+# 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 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)
projects / platform / upstream / llvm.git / commitdiff