mirror of https://github.com/llvm/torch-mlir
409 lines
16 KiB
Python
409 lines
16 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, Callable, Iterable
|
|
from enum import Enum
|
|
|
|
import sys
|
|
from io import StringIO
|
|
import tempfile
|
|
import os
|
|
|
|
from torch._functorch.compile_utils import strip_overloads
|
|
import torch
|
|
import torch.fx
|
|
from torch_mlir.dynamo import _get_decomposition_table
|
|
from torch.fx.experimental.proxy_tensor import make_fx
|
|
|
|
from torch_mlir.compiler_utils import (
|
|
run_pipeline_with_repro_report,
|
|
OutputType,
|
|
lower_mlir_module,
|
|
)
|
|
from torch_mlir.jit_ir_importer import ClassAnnotator, ImportOptions, ModuleBuilder
|
|
from torch_mlir.jit_ir_importer.build_tools.library_generator import generate_library
|
|
|
|
|
|
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]
|
|
if len(shape) == 0:
|
|
example_args_for_trace.append(torch.tensor(1))
|
|
else:
|
|
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: [
|
|
"aten.flatten.using_ints",
|
|
"aten.native_layer_norm",
|
|
"aten.linear",
|
|
],
|
|
OutputType.LINALG_ON_TENSORS: [
|
|
"aten.flatten.using_ints",
|
|
"aten.adaptive_avg_pool1d",
|
|
"aten.adaptive_avg_pool2d",
|
|
"aten.unflatten.int",
|
|
],
|
|
OutputType.STABLEHLO: [
|
|
"aten.amax",
|
|
"aten.amin",
|
|
"aten.randn.generator",
|
|
"aten.normal_functional",
|
|
],
|
|
}
|
|
|
|
|
|
def _canon_extra_library(
|
|
extra_library, extra_library_file_name="custom_op_extra_library.mlir"
|
|
):
|
|
if len(extra_library) != 0:
|
|
extra_library_dict = {}
|
|
for library_func in extra_library:
|
|
extra_library_dict[library_func.__name__] = library_func
|
|
mlir_library = generate_library(extra_library_dict)
|
|
|
|
extra_library_file = os.path.join(
|
|
tempfile.gettempdir(), extra_library_file_name
|
|
)
|
|
with open(extra_library_file, "w") as f:
|
|
f.write(mlir_library)
|
|
return extra_library_file
|
|
else:
|
|
return ""
|
|
|
|
|
|
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,
|
|
extra_library: Iterable[Callable] = [],
|
|
verbose: bool = False,
|
|
use_make_fx: bool = False,
|
|
enable_ir_printing: 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.
|
|
extra_library: List of abstract interpretation functions to splice
|
|
into the abstract interpretation library. See
|
|
`docs/adding_abstract_interpretation_functions.md` for more info
|
|
on the format the functions should have.
|
|
verbose: If true, print extra information about the conversion to
|
|
stdout.
|
|
enable_ir_printing: If true, print the IR before and after each pass to
|
|
stderr. This is equivalent to setting MLIR's `-print-ir-after-all`
|
|
flag. Note that this can easily generate many gigabytes of text,
|
|
so make sure to pipe stderr to a file (for example, run
|
|
`python tinymodel.py 2> tinymodel.stderr` on Linux).
|
|
|
|
Returns:
|
|
An MLIR module that contains the converted model in the specified
|
|
output type.
|
|
"""
|
|
extra_library_file_name = _canon_extra_library(extra_library)
|
|
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 STABLEHLO), 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, [])
|
|
|
|
if use_make_fx:
|
|
args = example_args._get_for_tracing(
|
|
use_tracing=True, ignore_traced_shapes=True
|
|
)["forward"]
|
|
model = make_fx(model, decomposition_table=_get_decomposition_table())(*args)
|
|
|
|
# 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 `torchscript.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 `torchscript.compile`
|
|
# on top of those building blocks.
|
|
if isinstance(model, torch.jit.ScriptModule):
|
|
# 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 verbose:
|
|
print("\n====================")
|
|
print("TorchScript RAW IR")
|
|
print(mb.module)
|
|
|
|
if output_type == OutputType.RAW:
|
|
return mb.module
|
|
|
|
option_string = (
|
|
"{backend-legal-ops="
|
|
+ ",".join(backend_legal_ops)
|
|
+ " extra-library="
|
|
+ extra_library_file_name
|
|
+ "}"
|
|
)
|
|
run_pipeline_with_repro_report(
|
|
mb.module,
|
|
f"builtin.module(torchscript-module-to-torch-backend-pipeline{option_string})",
|
|
"Lowering TorchScript IR -> Torch Backend IR",
|
|
enable_ir_printing=enable_ir_printing,
|
|
)
|
|
|
|
return lower_mlir_module(verbose, output_type, mb.module)
|