mirror of https://github.com/llvm/torch-mlir
418 lines
17 KiB
Python
418 lines
17 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
|
|
# Also available under a BSD-style license. See LICENSE.
|
|
|
|
from typing import Optional, Sequence, Union, List, Dict, Tuple
|
|
from enum import Enum
|
|
|
|
import sys
|
|
from io import StringIO
|
|
|
|
from torch._functorch.compile_utils import strip_overloads
|
|
import torch
|
|
|
|
from torch_mlir.passmanager import PassManager
|
|
from .compiler_utils import run_pipeline_with_repro_report
|
|
from torch_mlir.dialects.torch.importer.jit_ir import ClassAnnotator, ImportOptions, ModuleBuilder
|
|
|
|
|
|
class OutputType(Enum):
|
|
"""The kind of output that `torch_mlir.compile` can produce.
|
|
|
|
In MLIR terminology, this describes the mix of dialects that will be
|
|
produced by the conversion process.
|
|
|
|
In user-facing API's, this type can always be passed interchangeably with an
|
|
appropriate string specifying the output type. The allowed strings are
|
|
the set of enum vales, allowed to be case insensitive and with `-` allowed
|
|
in place of `_`. The `OutputType.get` static method can be used to convert
|
|
from a string to an `OutputType` instance.
|
|
"""
|
|
|
|
# This output type consists of `torch` dialect ops that have been converted
|
|
# maximally to value semantics, decomposed, and shapes have been inferred.
|
|
TORCH = "torch"
|
|
|
|
# The output type contains a mix of `linalg`-on-tensors ops, `scf`, and
|
|
# `arith` ops (and also `math` and `tm_tensor`). It can be thought of
|
|
# as taking the `TORCH` output type and lowering it so that tensor
|
|
# computations are done with `linalg`-on-tensors ops.
|
|
LINALG_ON_TENSORS = "linalg-on-tensors"
|
|
|
|
# This output type consists of `tosa` dialect ops. It can be thought of
|
|
# as taking the `TORCH` output type and lowering it to TOSA.
|
|
TOSA = "tosa"
|
|
|
|
# This output type consists of `mhlo` dialect ops. It can be thought of
|
|
# as taking the `TORCH` output type and lowering it to MHLO.
|
|
MHLO = "mhlo"
|
|
|
|
# Raw output of the JIT IR importer. This is not expected to be useful
|
|
# for end-users, but can be convenient for development or reporting bugs.
|
|
RAW = "raw"
|
|
|
|
@staticmethod
|
|
def get(spec: Union[str, "OutputType"]) -> "OutputType":
|
|
"""Gets an OutputType from allowed way to specify one.
|
|
|
|
Args:
|
|
spec: An OutputType instance or the case-insensitive name of one of the
|
|
enum values.
|
|
Returns:
|
|
An OutputType instance.
|
|
"""
|
|
if isinstance(spec, OutputType):
|
|
return spec
|
|
spec = spec.upper().replace("-", "_")
|
|
if spec not in OutputType.__members__:
|
|
raise ValueError(f"For output_type= argument, expected one of: "
|
|
f"{', '.join(OutputType.__members__.keys())}")
|
|
return OutputType[spec]
|
|
|
|
|
|
class TensorPlaceholder:
|
|
"""A class that represents a formal parameter of a given shape and dtype.
|
|
|
|
This class can be constructed explicitly from a shape and dtype:
|
|
```python
|
|
placeholder = TensorPlaceholder([3, 4], torch.float32)
|
|
```
|
|
|
|
This class can also be constructed from a `torch.Tensor` which is already
|
|
known to be a valid input to the function. In this case, a set of
|
|
dynamic axes are allowed to be specified.
|
|
```python
|
|
placeholder = TensorPlaceholder.like(torch.ones(3, 4), dynamic_axes=[1])
|
|
# Equivalent to `TensorPlaceholder([3, -1], torch.float32)`
|
|
```
|
|
"""
|
|
|
|
def __init__(self, shape: List[int], dtype: torch.dtype):
|
|
"""Create a tensor with shape `shape` and dtype `dtype`.
|
|
|
|
Args:
|
|
shape: The shape of the tensor. A size of `-1` indicates that the
|
|
dimension has an unknown size.
|
|
dtype: The dtype of the tensor.
|
|
"""
|
|
self.shape = shape
|
|
self.dtype = dtype
|
|
|
|
@staticmethod
|
|
def like(tensor: torch.Tensor, dynamic_axes: List[int] = None):
|
|
"""Create a tensor placeholder that is like the given tensor.
|
|
|
|
Args:
|
|
tensor: The tensor to create a placeholder for.
|
|
dynamic_axes: A list of dynamic axes. If specified, the compiled
|
|
module will allow those axes to be any size at runtime.
|
|
"""
|
|
if dynamic_axes is None:
|
|
dynamic_axes = []
|
|
shape = []
|
|
for i, dim in enumerate(tensor.shape):
|
|
if i in dynamic_axes:
|
|
shape.append(-1)
|
|
else:
|
|
shape.append(dim)
|
|
return TensorPlaceholder(shape, tensor.dtype)
|
|
|
|
|
|
_example_arg = Union[TensorPlaceholder, torch.Tensor]
|
|
_example_args_for_one_method = Union[_example_arg, Sequence[_example_arg]]
|
|
_example_args = Union[_example_args_for_one_method, "ExampleArgs"]
|
|
|
|
|
|
class ExampleArgs:
|
|
"""A class representing the example arguments to an nn.Module.
|
|
|
|
In general, an nn.Module may have multiple methods that need to be compiled.
|
|
This requires example args for each method. This class is a lightweight
|
|
wrapper around a dictionary that maps method names to example arguments.
|
|
|
|
In user-facing API's, this type can always be passed interchangeably with a
|
|
single arg or list of args, which normalizes to an ExampleArgs for just
|
|
the `forward` method via the `ExampleArgs.get` static method.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._example_args = {}
|
|
|
|
def add_method(self, method_name: str, example_args: _example_args_for_one_method):
|
|
"""Adds example args for a method.
|
|
|
|
Args:
|
|
method_name: The name of the method. Must have not already been
|
|
added previously as a method.
|
|
example_args: The example args for the method.
|
|
Returns:
|
|
self, for chaining.
|
|
"""
|
|
assert method_name not in self._example_args
|
|
self._example_args[method_name] = ExampleArgs._canonicalize_args(
|
|
example_args)
|
|
return self
|
|
|
|
@staticmethod
|
|
def get(example_args: _example_args) -> "ExampleArgs":
|
|
"""Gets an ExampleArgs from one of the permissible ways to specify one.
|
|
|
|
Args:
|
|
example_args: An ExampleArgs instance or a single arg or list of args.
|
|
Returns:
|
|
An ExampleArgs instance.
|
|
"""
|
|
if isinstance(example_args, ExampleArgs):
|
|
return example_args
|
|
return ExampleArgs().add_method("forward", example_args)
|
|
|
|
@staticmethod
|
|
def _canonicalize_args(example_args: _example_args_for_one_method):
|
|
"""Canonicalize the args for one method into a tuple."""
|
|
if not isinstance(example_args, Sequence):
|
|
example_args = [example_args]
|
|
for arg in example_args:
|
|
if not isinstance(arg, (TensorPlaceholder, torch.Tensor)):
|
|
raise Exception(f"Only Tensor's, TensorPlaceholder's, or sequences of "
|
|
f"Tensor's and TensorPlaceholder's are supported as "
|
|
f"example args for method inputs. "
|
|
f"Got '{arg}'.")
|
|
return tuple(example_args)
|
|
|
|
def _get_methods(self):
|
|
return self._example_args.keys()
|
|
|
|
def _get_for_annotation(self):
|
|
result = {}
|
|
for method_name, example_args in self._example_args.items():
|
|
placeholders = []
|
|
for arg in example_args:
|
|
if isinstance(arg, TensorPlaceholder):
|
|
placeholders.append(arg)
|
|
else:
|
|
assert isinstance(arg, torch.Tensor)
|
|
placeholders.append(TensorPlaceholder.like(arg))
|
|
result[method_name] = placeholders
|
|
return result
|
|
|
|
def _get_for_tracing(
|
|
self,
|
|
use_tracing: bool,
|
|
ignore_traced_shapes: bool,
|
|
) -> Dict[str, Tuple[_example_arg, ...]]:
|
|
result = {}
|
|
for method_name, example_args in self._example_args.items():
|
|
# If we are tracing, then we need to convert any placeholders into
|
|
# concrete values.
|
|
if use_tracing:
|
|
example_args_for_trace = []
|
|
for arg in example_args:
|
|
if isinstance(arg, TensorPlaceholder):
|
|
if not ignore_traced_shapes:
|
|
# To avoid accidental footguns, we require
|
|
# `ignore_traced_shapes` to be true if we're using
|
|
# TensorPlaceholder's, as it falls into the same
|
|
# "hopefully the trace works for different inputs"
|
|
# bucket of concerns.
|
|
raise Exception(
|
|
"TensorPlaceholder can only be used with tracing when `ignore_traced_shapes=True`")
|
|
# For any dynamic dimensions, replace them with "7"
|
|
# arbitrarily. If a user is using dynamic dimensions with
|
|
# tracing, they are walking on thin ice already -- assume
|
|
# they know what they are doing and that their trace is
|
|
# correct for any specific concrete size.
|
|
shape = [s if s != -1 else 7 for s in arg.shape]
|
|
example_args_for_trace.append(
|
|
torch.ones(*shape, dtype=arg.dtype))
|
|
else:
|
|
assert isinstance(arg, torch.Tensor)
|
|
example_args_for_trace.append(arg)
|
|
example_args = tuple(example_args_for_trace)
|
|
result[method_name] = example_args
|
|
return result
|
|
|
|
|
|
# The set of ops that are considered legal for each backend.
|
|
# These are currently quite load-bearing, since different backends might be
|
|
# missing patterns for decomposed forms of certain ops.
|
|
# TODO: Tighten up the definition of these "conditionally legal for backends"
|
|
# ops in the backend contract, and move these lists somewhere deeper in the
|
|
# compiler where each backend can "own" its set of legal ops.
|
|
BACKEND_LEGAL_OPS = {
|
|
OutputType.TOSA: ['torch.aten.flatten.using_ints', 'torch.aten.native_layer_norm', 'torch.aten.linear'],
|
|
OutputType.LINALG_ON_TENSORS: ['torch.aten.flatten.using_ints', ],
|
|
OutputType.MHLO: [],
|
|
}
|
|
|
|
|
|
def compile(model: torch.nn.Module,
|
|
example_args: _example_args,
|
|
output_type: Union[str, "OutputType"] = OutputType.TORCH,
|
|
use_tracing: bool = False,
|
|
ignore_traced_shapes=False,
|
|
backend_legal_ops: Optional[Sequence[str]] = None,
|
|
verbose: bool = False):
|
|
"""Convert a PyTorch model to MLIR.
|
|
|
|
Args:
|
|
model: The PyTorch model to convert.
|
|
example_args: A list of example arguments to use when inferring the
|
|
shapes of the arguments to `forward` method of the model.
|
|
A single tensor is treated as a list of a single tensor.
|
|
A TensorPlaceholder object is also allowed in the place of any
|
|
Tensor. For models with multiple methods, an `ExampleArgs` object
|
|
can be passed.
|
|
output_type: The kind of output to produce. See `OutputType` for more
|
|
details.
|
|
use_tracing: If True, use `torch.jit.trace` to convert the model to
|
|
JIT IR rather than `torch.jit.script`.
|
|
ignore_traced_shapes: If True, ignore the shapes that were observed
|
|
during tracing. This should only be used if one knows that the
|
|
original traced program would result in the same trace (modulo
|
|
shapes) for all shape combinations implied by any
|
|
`TensorPlaceholder`'s used as `example_args`. Also,
|
|
strictly-speaking, this option covers dtypes too, but we just say
|
|
"shapes" to be succinct.
|
|
backend_legal_ops: A list of ops that should be considered legal for
|
|
the backend. An op that is considered legal will not be decomposed.
|
|
This option is only valid with the `"torch"` output type.
|
|
verbose: If true, print extra information about the conversion.
|
|
|
|
Returns:
|
|
An MLIR module that contains the converted model in the specified
|
|
output type.
|
|
"""
|
|
output_type = OutputType.get(output_type)
|
|
example_args = ExampleArgs.get(example_args)
|
|
if ignore_traced_shapes and not use_tracing:
|
|
raise Exception("`ignore_traced_shapes` requires `use_tracing`")
|
|
|
|
# We only allow `backend_legal_ops` to be specified for the `"torch"`
|
|
# output type because the other output types actually invoke their
|
|
# respective backends (Linalg, TOSA, or MHLO), and those backends have
|
|
# very specific requirements about the ops which are legal.
|
|
# See `BACKEND_LEGAL_OPS` for more details.
|
|
if backend_legal_ops is not None:
|
|
if output_type != OutputType.TORCH:
|
|
raise Exception("`backend_legal_ops` is only valid with the "
|
|
"`torch` output type")
|
|
backend_legal_ops = list(sorted(set(backend_legal_ops)))
|
|
else:
|
|
backend_legal_ops = BACKEND_LEGAL_OPS.get(output_type, [])
|
|
|
|
# For FX-based models, automatically strip overloads.
|
|
if isinstance(model, torch.fx.GraphModule):
|
|
strip_overloads(model)
|
|
|
|
# Get the model as JIT IR (TorchScript) for import.
|
|
# TODO: Longer-term, we probably need to split `torch_mlir.compile`.
|
|
# There should be an "acquisition" step that does
|
|
# tracing/scripting/importing from FX/using torchdynamo.export/etc.
|
|
# + any lowering to the backend contract. Then there should be a
|
|
# "backend lowering" step that does the actual lowering to each
|
|
# backend. This separation should be visible at the Python API level, and
|
|
# we can implement a deliberately simplified API like `torch_mlir.compile`
|
|
# on top of those building blocks.
|
|
if isinstance(model, torch.jit._script.RecursiveScriptModule):
|
|
# If the user already converted the model to JIT IR themselves, just
|
|
# do some basic error checking, but take the model as-is.
|
|
for method_name in example_args._get_methods():
|
|
if not hasattr(model, method_name):
|
|
raise Exception(
|
|
f"Model does not have exported method '{method_name}', "
|
|
f"requested in `example_args`. Consider adding "
|
|
f"`@torch.jit.export` to the method definition.")
|
|
scripted = model
|
|
elif use_tracing:
|
|
scripted = torch.jit.trace_module(
|
|
model,
|
|
example_args._get_for_tracing(use_tracing, ignore_traced_shapes)
|
|
)
|
|
else:
|
|
# Make sure that all the methods that the user requested get scripted.
|
|
# By default, PyTorch only scripts the `forward` method and transitive
|
|
# callees.
|
|
for method_name in example_args._get_methods():
|
|
torch.jit.export(getattr(model, method_name).__func__)
|
|
scripted = torch.jit.script(model)
|
|
class_annotator = ClassAnnotator()
|
|
class_annotator.exportNone(scripted._c._type())
|
|
for method_name, example_args in example_args._get_for_annotation().items():
|
|
class_annotator.exportPath(scripted._c._type(), [method_name])
|
|
annotation = [None] # `None` is always the annotation for "self".
|
|
for arg in example_args:
|
|
annotation.append((arg.shape, arg.dtype, True))
|
|
class_annotator.annotateArgs(
|
|
scripted._c._type(), [method_name], annotation)
|
|
|
|
mb = ModuleBuilder()
|
|
import_options = ImportOptions()
|
|
import_options.ignoreExistingTensorShapesAndDtypes = ignore_traced_shapes
|
|
try:
|
|
original_stderr = sys.stderr
|
|
sys.stderr = StringIO()
|
|
# Import the TorchScript module to MLIR
|
|
mb.import_module(scripted._c, class_annotator, import_options)
|
|
except Exception as e:
|
|
raise Exception(f"""
|
|
PyTorch TorchScript module -> torch-mlir Object Graph IR import failed with:
|
|
### Importer C++ Exception:
|
|
{e}
|
|
### Importer Diagnostics:
|
|
{sys.stderr.getvalue()}
|
|
""") from None
|
|
finally:
|
|
sys.stderr = original_stderr
|
|
if output_type == OutputType.RAW:
|
|
return mb.module
|
|
|
|
option_string = "{backend-legal-ops=" + ",".join(backend_legal_ops) + "}"
|
|
run_pipeline_with_repro_report(
|
|
mb.module,
|
|
f"builtin.module(torchscript-module-to-torch-backend-pipeline{option_string})",
|
|
"Lowering TorchScript IR -> Torch Backend IR",
|
|
)
|
|
|
|
if verbose:
|
|
print("\n====================")
|
|
print("Torch Backend IR")
|
|
print(mb.module)
|
|
|
|
if output_type == OutputType.TORCH:
|
|
return mb.module
|
|
|
|
if output_type == OutputType.TOSA:
|
|
run_pipeline_with_repro_report(
|
|
mb.module,
|
|
"builtin.module(torch-backend-to-tosa-backend-pipeline)",
|
|
"Lowering Torch Backend IR -> TOSA Backend IR")
|
|
if verbose:
|
|
print("\n====================")
|
|
print("TOSA Backend IR")
|
|
print(mb.module)
|
|
return mb.module
|
|
|
|
if output_type == OutputType.LINALG_ON_TENSORS:
|
|
run_pipeline_with_repro_report(
|
|
mb.module,
|
|
"builtin.module(torch-backend-to-linalg-on-tensors-backend-pipeline)",
|
|
"Lowering Torch Backend IR -> Linalg-on-Tensors Backend IR")
|
|
if verbose:
|
|
print("\n====================")
|
|
print("LINALG Backend IR")
|
|
print(mb.module)
|
|
return mb.module
|
|
|
|
elif output_type == OutputType.MHLO:
|
|
run_pipeline_with_repro_report(
|
|
mb.module,
|
|
"builtin.module(torch-backend-to-mhlo-backend-pipeline)",
|
|
"Lowering Torch Backend IR -> MHLO Backend IR")
|
|
if verbose:
|
|
print("\n====================")
|
|
print("MHLO Backend IR")
|
|
print(mb.module)
|
|
return mb.module
|
|
raise Exception(f"Unknown OutputType: {output_type}")
|