1#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
2#  See https://llvm.org/LICENSE.txt for license information.
3#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
4
5#  This file contains the sparse compiler class. It is copied from
6#  test/Integration/Dialect/SparseTensor/python/ until we have a better
7#  solution.
8
9from mlir import execution_engine
10from mlir import ir
11from mlir import passmanager
12from typing import Sequence
13
14
15class SparseCompiler:
16  """Sparse compiler class for compiling and building MLIR modules."""
17
18  def __init__(self, options: str, opt_level: int, shared_libs: Sequence[str]):
19    pipeline = f'sparse-compiler{{{options} reassociate-fp-reductions=1 enable-index-optimizations=1}}'
20    self.pipeline = pipeline
21    self.opt_level = opt_level
22    self.shared_libs = shared_libs
23
24  def __call__(self, module: ir.Module):
25    """Convenience application method."""
26    self.compile(module)
27
28  def compile(self, module: ir.Module):
29    """Compiles the module by invoking the sparse copmiler pipeline."""
30    passmanager.PassManager.parse(self.pipeline).run(module)
31
32  def jit(self, module: ir.Module) -> execution_engine.ExecutionEngine:
33    """Wraps the module in a JIT execution engine."""
34    return execution_engine.ExecutionEngine(
35        module, opt_level=self.opt_level, shared_libs=self.shared_libs)
36
37  def compile_and_jit(self,
38                      module: ir.Module) -> execution_engine.ExecutionEngine:
39    """Compiles and jits the module."""
40    self.compile(module)
41    return self.jit(module)
42