2021-04-07 06:58:37 +08:00
|
|
|
# -*- Python -*-
|
|
|
|
|
# This file is licensed under a pytorch-style license
|
|
|
|
|
# See frontends/pytorch/LICENSE for license information.
|
|
|
|
|
|
|
|
|
|
import typing
|
|
|
|
|
|
|
|
|
|
import torch
|
|
|
|
|
import torch_mlir
|
|
|
|
|
|
|
|
|
|
import npcomp
|
2021-04-09 04:05:16 +08:00
|
|
|
from npcomp.compiler.pytorch.backend import refjit, frontend_lowering
|
2021-04-07 06:58:37 +08:00
|
|
|
from npcomp.compiler.utils import logging
|
|
|
|
|
|
|
|
|
|
import test_utils
|
|
|
|
|
|
|
|
|
|
#logging.enable()
|
|
|
|
|
|
|
|
|
|
# RUN: %PYTHON %s | npcomp-opt | FileCheck %s
|
|
|
|
|
|
|
|
|
|
mb = torch_mlir.ModuleBuilder()
|
|
|
|
|
|
|
|
|
|
class TestModule(torch.nn.Module):
|
|
|
|
|
def __init__(self):
|
|
|
|
|
super().__init__()
|
|
|
|
|
def forward(self, x):
|
|
|
|
|
return torch.tanh(x)
|
|
|
|
|
|
|
|
|
|
test_module = TestModule()
|
|
|
|
|
class_annotator = torch_mlir.ClassAnnotator()
|
|
|
|
|
recursivescriptmodule = torch.jit.script(test_module)
|
|
|
|
|
torch.jit.save(recursivescriptmodule, '/tmp/foo.pt')
|
|
|
|
|
|
|
|
|
|
class_annotator.exportNone(recursivescriptmodule._c._type())
|
|
|
|
|
class_annotator.exportPath(recursivescriptmodule._c._type(), ['forward'])
|
Introduce `!torch.tensor` / `!torch.vtensor` types.
This removes our reliance on the numpy dialect and avoids our off-label
use of the builtin tnesor type for modeling unknown dtypes. The
`!torch.vtensor` (`ValueTensorType`) type is a value-semantic tensor.
The `!torch.tensor` (`NonValueTensorType`) type is a non-value-semantic
tensor. The new types look as follows syntactically:
```
// Least-static-information, non-value-semantic tensor.
!torch.tensor
// Explicit form of least-static-information variant.
!torch.tensor<*,unk>
// Least-static-information, value-semantic tensor.
!torch.vtensor
// Explicit form of least-static-information variant.
!torch.vtensor<*,unk>
// Fixed-set of allowable element types, with first-class support for
// Torch's frontend signedness semantics.
!torch.tensor<*,si32>
// First-class support for unknown dtypes.
!torch.tensor<[?,?,?],unk>
// Standard MLIR representation of `?` for unknown dimensions.
!torch.tensor<[?,2,?,4],unk>
// Statically shaped / dtyped example.
!torch.vtensor<[1,2,3,4],f32>
```
This required fairly significant changes throughout the compiler, but
overall it is a big cleanup. We now have a much clearer layering of "the
Torch frontend lowering" vs "lowering to std + linalg + etc.".
At the C++ level, there is `ValueTensorType`, `NonValueTensorType`.
We also have a helper `BaseTensorType` (kind of like ShapedType) which
interoperates with those two.
Included changes:
- New `torch.tensor(dense<0.0> : tensor<5xf32>) : !torch.tensor` op for
creating torch tensor literals in the frontend.
- Consistently use signedness for the types (except i1 which I didn't
touch -- we need to sort out the situation with !basicpy.BoolType
there anyway so will be attending to that soon)
- Frontend can annotate whether an argument to the function has value
semantics. We currently require this, as our backend contract does not
currently allow us to even model the non-value-semantic case. Before,
the value-semantic assumption was randomly injected in the middle of
the pass pipeline.
- Move ArrayToTensor (now called MaximizeValueSemantics) and
RefinePublicReturn passes to torch dialect.
- The TorchToStd and TorchToLinalg passes are now type conversions from
`!torch.vtensor` to `tensor` and use the dialect conversion infra.
The overall conversion pipeline is set up following the best practices
of the "Type Conversions the Not-So-Hard Way" talk. This required
introducing `torch-func-builtin-tensorize` and
`torch-finalizing-builtin-tensorize` passes analogous to the upstream
bufferization passes with the corresponding names (mostly just
copypasta from there).
- Misc Torch-level canonicalizations -- we now cleanly layer the
lowering to std later in the pipeline, so we are gradually lessening
our reliance on random std constant folding before we get to that
point.
Recommended review order:
- New types in TorchTypes.td/TorchTypes.h/TorchDialect.cpp
- New ops in TorchOps.td / TorchOps.cpp
- Less important / more mechanical stuff
- Frontend changes.
- Pass changes/additions in `Torch/Transforms` and `Conversion/`
2021-05-21 08:07:18 +08:00
|
|
|
class_annotator.annotateArgs(recursivescriptmodule._c._type(), ['forward'], [
|
2021-04-07 06:58:37 +08:00
|
|
|
None,
|
|
|
|
|
([2, 3, -1], torch.float32)
|
|
|
|
|
])
|
|
|
|
|
# TODO: Automatically handle unpacking Python class RecursiveScriptModule into the underlying ScriptModule.
|
|
|
|
|
mb.import_module(recursivescriptmodule._c, class_annotator)
|
|
|
|
|
#mb.module.operation.print()
|
|
|
|
|
|
|
|
|
|
backend = refjit.CompilerBackend()
|
2021-04-09 04:05:16 +08:00
|
|
|
compiled = backend.compile(frontend_lowering.lower_object_graph(mb.module))
|
2021-04-07 06:58:37 +08:00
|
|
|
jit_module = backend.load(compiled)
|
|
|
|
|
|
|
|
|
|
torch.manual_seed(0)
|
|
|
|
|
input = torch.rand(2, 3, 1)
|
|
|
|
|
test_utils.compare_outputs(test_module.forward, jit_module.forward, input)
|