mirror of https://github.com/llvm/torch-mlir
80 lines
2.5 KiB
Python
80 lines
2.5 KiB
Python
# 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)
|