49 lines · python
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-exception4 5# This file contains the Sparsifier class.6 7from mlir import execution_engine8from mlir import ir9from mlir import passmanager10from typing import Sequence11 12 13class Sparsifier:14 """Sparsifier class for compiling and building MLIR modules."""15 16 def __init__(17 self,18 extras: str,19 options: str,20 opt_level: int,21 shared_libs: Sequence[str],22 ):23 pipeline = (24 f"builtin.module({extras}sparsifier{{{options} reassociate-fp-reductions=1"25 " enable-index-optimizations=1})"26 )27 self.pipeline = pipeline28 self.opt_level = opt_level29 self.shared_libs = shared_libs30 31 def __call__(self, module: ir.Module):32 """Convenience application method."""33 self.compile(module)34 35 def compile(self, module: ir.Module):36 """Compiles the module by invoking the sparsifier pipeline."""37 passmanager.PassManager.parse(self.pipeline).run(module.operation)38 39 def jit(self, module: ir.Module) -> execution_engine.ExecutionEngine:40 """Wraps the module in a JIT execution engine."""41 return execution_engine.ExecutionEngine(42 module, opt_level=self.opt_level, shared_libs=self.shared_libs43 )44 45 def compile_and_jit(self, module: ir.Module) -> execution_engine.ExecutionEngine:46 """Compiles and jits the module."""47 self.compile(module)48 return self.jit(module)49