mirror of https://github.com/llvm/torch-mlir
130 lines
4.3 KiB
Python
130 lines
4.3 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
|
|
|
|
import torch
|
|
|
|
from mlir.ir import *
|
|
from mlir.passmanager import *
|
|
from npcomp.compiler.generic.backend import refjit as refjit_backend
|
|
from npcomp.compiler.utils import logging
|
|
|
|
__all__ = [
|
|
"is_enabled",
|
|
"CompilerBackend",
|
|
]
|
|
|
|
# The set of passes that lowers from a TorchScript object graph representation
|
|
# to a module semantics where symbols correspond to dotted paths into the
|
|
# module.
|
|
OBJECT_GRAPH_LOWERING_PASSES = (
|
|
"torch-globalize-pipeline",
|
|
# symbol-dce is currently needed for correctness, as we don't have a lowering
|
|
# in the backend for torch.global_slot's.
|
|
# Torch usually inserts a few unused global slots that are otherwise
|
|
# bothersome because we don't currently have a lowering for them.
|
|
# TODO: Support global slots in backends.
|
|
"symbol-dce",
|
|
)
|
|
|
|
TORCH_TO_TCF_PASSES = (
|
|
"func(aten-recognize-kernels)",
|
|
"func(convert-aten-to-tcf)",
|
|
"numpy-public-functions-to-tensor",
|
|
"canonicalize",
|
|
)
|
|
|
|
# Re-export.
|
|
is_enabled = refjit_backend.is_enabled
|
|
|
|
|
|
class TorchJitModuleInvoker(refjit_backend.JitModuleInvoker):
|
|
"""Allows torch.Tensor inputs to be passed to module invocations."""
|
|
|
|
def __getitem__(self, function_name: str):
|
|
numpy_invoke = super().__getitem__(function_name)
|
|
|
|
def invoke(*args):
|
|
args = tuple(
|
|
arg.numpy() if isinstance(arg, torch.Tensor) else arg for arg in args)
|
|
return numpy_invoke(*args)
|
|
|
|
return invoke
|
|
|
|
|
|
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, imported_module: Module):
|
|
"""Compiles an imported module, with a flat list of functions.
|
|
|
|
Args:
|
|
imported_module: The MLIR module consisting of funcs in the torch
|
|
dialect.
|
|
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 as context:
|
|
if self._debug:
|
|
logging.debug("Initial PyTorch IR:\n{}", imported_module)
|
|
|
|
# Frontend.
|
|
pipeline_str = ",".join(TORCH_TO_TCF_PASSES)
|
|
if self._debug:
|
|
logging.debug("Running Torch->TCF pipeline '{}'", pipeline_str)
|
|
pm = PassManager.parse(pipeline_str)
|
|
pm.run(imported_module)
|
|
if self._debug:
|
|
logging.debug("TCF IR:\n{}", 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:\n{}", imported_module)
|
|
|
|
jit_module = self._refjit.JITModule.from_compiled_module(
|
|
imported_module, refjit_backend.get_runtime_libs())
|
|
return jit_module
|
|
|
|
def compile_object_graph(self, imported_module: Module):
|
|
"""Compiles an imported module, with TorchScript object graph semantics.
|
|
|
|
Args:
|
|
imported_module: The MLIR module consisting of IR as imported by the
|
|
torch_mlir.import_module
|
|
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 as context:
|
|
if self._debug:
|
|
logging.debug("Initial PyTorch object graph IR:\n{}", imported_module)
|
|
|
|
# Frontend.
|
|
pipeline_str = ",".join(OBJECT_GRAPH_LOWERING_PASSES)
|
|
if self._debug:
|
|
logging.debug(
|
|
"Running Torch object graph lowering pipeline '{}'", pipeline_str)
|
|
pm = PassManager.parse(pipeline_str)
|
|
pm.run(imported_module)
|
|
return self.compile(imported_module)
|
|
|
|
def load(self, jit_module) -> TorchJitModuleInvoker:
|
|
"""Loads a compiled artifact into the runtime."""
|
|
return TorchJitModuleInvoker(jit_module)
|