torch-mlir/python/npcomp/compiler/numpy/backend/refjit.py

80 lines
2.5 KiB
Python
Raw Normal View History

# 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
import os
from mlir.ir import *
from mlir.passmanager import *
from _npcomp import register_dialects
from _npcomp import mlir as legacy_mlir
from npcomp.compiler.generic.backend import refjit as refjit_backend
from npcomp.compiler.utils import logging
__all__ = [
"is_enabled",
"CompilerBackend",
]
FRONTEND_PASSES = (
"func(npcomp-cpa-type-inference)",
"numpy-public-functions-to-tensor",
"func(convert-numpy-to-tcf)",
"func(convert-scf-to-std)",
"func(canonicalize)",
"func(tcf-shape-refinement)",
)
# Re-export.
is_enabled = refjit_backend.is_enabled
class CompilerBackend:
"""Main entry-point for the backend."""
def __init__(self):
super().__init__()
self._refjit = refjit_backend.get_refjit()
self._debug = logging.debug_enabled()
def compile(self, legacy_imported_ir_module: legacy_mlir.ir.ModuleOp):
"""Compiles an imported module.
Args:
legacy_imported_ir_module: The MLIR module as imported from the
ImportFrontend.
Returns:
An opaque, backend specific module object that can be passed to load.
The object may actually be something more specific to the backend (i.e.
for IREE, it is a serialized VM flatbuffer) but the contract is that
it is operated on by methods on this class.
"""
# TODO: Once transitioned to new Python API, don't reparse the module.
with Context() as context:
register_dialects(context)
imported_module = Module.parse(legacy_imported_ir_module.to_asm())
# Frontend.
pm = PassManager.parse(",".join(FRONTEND_PASSES))
pm.run(imported_module)
if self._debug:
logging.debug("Frontend IR:{}", imported_module)
# Backend.
# Note that this is a separate pass manager purely to aid in debugging.
pm = PassManager()
self._refjit.build_backend_compilation_pipeline(pm)
pm.run(imported_module)
if self._debug:
logging.debug("Backend IR:{}", imported_module)
jit_module = self._refjit.JITModule.from_compiled_module(
imported_module, refjit_backend.get_runtime_libs())
return jit_module
def load(self, jit_module):
"""Loads a compiled artifact into the runtime.
Since this is a JIT instead of an AOT compiler,
"""
return refjit_backend.JitModuleInvoker(jit_module)