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

110 lines
3.4 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 io
import subprocess
from mlir.ir import *
from mlir.passmanager import *
from npcomp.compiler.generic.backend import iree as iree_backend
from npcomp.compiler.utils import logging
__all__ = [
"is_enabled",
"CompilerBackend",
]
FRONTEND_PASSES = (
"func(basicpy-type-inference)",
"func(convert-basicpy-to-std)",
"func(canonicalize)",
"func(convert-scf-to-std)",
)
_ireert = None
_cached_config = None
def _get_iree():
"""Dynamically resolves the iree backend module."""
global _ireert
try:
from pyiree import rt as imported_rt
except ImportError:
raise ImportError("IREE runtime library not found (pyiree.rt)")
_ireert = imported_rt
return _ireert
def is_enabled() -> bool:
"""Returns whether the backend is enabled for the current build."""
try:
_get_iree()
return True
except ImportError:
return False
class CompilerBackend:
"""Main entry-point for the backend."""
def __init__(self):
super().__init__()
self._ireert = _get_iree()
self._debug = logging.debug_enabled()
def compile(self, imported_module: Module):
"""Compiles an imported module.
Args:
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.
"""
with imported_module.context:
# Frontend.
if self._debug:
logging.debug("Input IR:\n{}", imported_module)
assert (
imported_module.operation.verify()), "Imported module does not verify"
# Frontend.
pm = PassManager.parse(",".join(FRONTEND_PASSES))
pm.run(imported_module)
if self._debug:
logging.debug("Frontend IR:{}", imported_module)
# TODO: There should be some common utility for invoking backend processes
# safely (and have options like saving temps, etc).
args = [
iree_backend.get_translate_exe(), "--iree-mlir-to-vm-bytecode-module"
]
p = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
imported_module.operation.print(binary=True,
enable_debug_info=True,
file=p.stdin)
out, err = p.communicate()
return out
def load(self, vm_blob):
"""Loads a compiled artifact into the runtime.
This is meant as a simple mechanism for testing and is not optimized or
highly parameterized. It loads a compiled result into a new runtime
instance and returns an object that exposes a python function for each
public function compiled in the imported_ir_module that was compiled.
"""
ireert = self._ireert
m = ireert.VmModule.from_flatbuffer(vm_blob)
global _cached_config
if not _cached_config:
# TODO: Need to make the configuration more flexible.
_cached_config = ireert.Config(driver_name="vmla")
ctx = ireert.SystemContext(config=_cached_config)
ctx.add_module(m)
# TODO: The implicit tying of the 'module' name has got to go.
return ctx.modules.module