llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
torch-mlir/frontends/pytorch/examples/torchscript_mm_e2e.py

59 lines
1.7 KiB
Python
Raw Normal View History

Add `aten.mm` to linalg lowering. This is our first op with error semantics, and stresses the system. There are a few design notes of special interest: - RefineTypes.cpp's note about shape inference in the presence of code that dynamically produces and error, and it is provable statically. - ATenToLinalg.cpp's notes about future automation of the ATen->linalg path. - The notes in Passes.td about using low-tech `std.assert` ops instead of `shape.assuming`. Note: Doesn't work on IREE yet due to the `std.assert` op (needs to be lowered to `vm.fail` on the IREE side).
2021-04-09 08:43:41 +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
from npcomp.compiler.pytorch.backend import refjit, frontend_lowering
from npcomp.compiler.utils import logging
import test_utils
logging.enable()
# RUN: %PYTHON %s | npcomp-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
Add npcomp-verify-backend-contract pass. This pass verifies that a given module satisfies the contract that we have for backends. This is phrased as an "allowlist", because we want to keep this interface tight. Also, this gives much better diagnostics than a backend randomly crashing or failing to compile would (though they could still be improved). This was especially painful because if we had `tensor<?x!numpy.any_dtype>` slip through, at some point RefBackend would convert it to a memref type and trip the "verify type invariants" assertion which gives no location or anything and crashed the process, which was very unpleasant. We implement this with the dialect conversion framework, which works reasonably well and was quick to put together and familiar, but is still very "op oriented". We probably want to make this hand-rolled eventually, especially the error reporting (the most useful kind of error for a dialect conversion user is not necessarily the best for this use case). Also, in production, these error will go to users, and need to be surfaced carefully such as "the compiler needs a type annotation on this function parameter" which in general requires some special analysis, wordsmithing, and overall awareness of the e2e use case (such as how much we can lean into certain source locations) to provide a meaningful user-level diagnostic. Also, add `inline` to the current frontend lowering pass pipeline to allow slightly more complicated programs that otherwise would fail on shape inference.
2021-04-13 09:39:53 +08:00
class Submodule(torch.nn.Module):
Add `aten.mm` to linalg lowering. This is our first op with error semantics, and stresses the system. There are a few design notes of special interest: - RefineTypes.cpp's note about shape inference in the presence of code that dynamically produces and error, and it is provable statically. - ATenToLinalg.cpp's notes about future automation of the ATen->linalg path. - The notes in Passes.td about using low-tech `std.assert` ops instead of `shape.assuming`. Note: Doesn't work on IREE yet due to the `std.assert` op (needs to be lowered to `vm.fail` on the IREE side).
2021-04-09 08:43:41 +08:00
def __init__(self):
super().__init__()
def forward(self, lhs, rhs):
return torch.mm(lhs, rhs)
Add npcomp-verify-backend-contract pass. This pass verifies that a given module satisfies the contract that we have for backends. This is phrased as an "allowlist", because we want to keep this interface tight. Also, this gives much better diagnostics than a backend randomly crashing or failing to compile would (though they could still be improved). This was especially painful because if we had `tensor<?x!numpy.any_dtype>` slip through, at some point RefBackend would convert it to a memref type and trip the "verify type invariants" assertion which gives no location or anything and crashed the process, which was very unpleasant. We implement this with the dialect conversion framework, which works reasonably well and was quick to put together and familiar, but is still very "op oriented". We probably want to make this hand-rolled eventually, especially the error reporting (the most useful kind of error for a dialect conversion user is not necessarily the best for this use case). Also, in production, these error will go to users, and need to be surfaced carefully such as "the compiler needs a type annotation on this function parameter" which in general requires some special analysis, wordsmithing, and overall awareness of the e2e use case (such as how much we can lean into certain source locations) to provide a meaningful user-level diagnostic. Also, add `inline` to the current frontend lowering pass pipeline to allow slightly more complicated programs that otherwise would fail on shape inference.
2021-04-13 09:39:53 +08:00
class TestModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.s = Submodule()
def forward(self, lhs, rhs):
return self.s.forward(lhs, rhs)
Add `aten.mm` to linalg lowering. This is our first op with error semantics, and stresses the system. There are a few design notes of special interest: - RefineTypes.cpp's note about shape inference in the presence of code that dynamically produces and error, and it is provable statically. - ATenToLinalg.cpp's notes about future automation of the ATen->linalg path. - The notes in Passes.td about using low-tech `std.assert` ops instead of `shape.assuming`. Note: Doesn't work on IREE yet due to the `std.assert` op (needs to be lowered to `vm.fail` on the IREE side).
2021-04-09 08:43:41 +08:00
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'], [
Add `aten.mm` to linalg lowering. This is our first op with error semantics, and stresses the system. There are a few design notes of special interest: - RefineTypes.cpp's note about shape inference in the presence of code that dynamically produces and error, and it is provable statically. - ATenToLinalg.cpp's notes about future automation of the ATen->linalg path. - The notes in Passes.td about using low-tech `std.assert` ops instead of `shape.assuming`. Note: Doesn't work on IREE yet due to the `std.assert` op (needs to be lowered to `vm.fail` on the IREE side).
2021-04-09 08:43:41 +08:00
None,
([-1, -1], torch.float32),
([-1, -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()
compiled = backend.compile(frontend_lowering.lower_object_graph(mb.module))
jit_module = backend.load(compiled)
torch.manual_seed(0)
lhs = torch.rand(2, 3)
rhs = torch.rand(3, 4)
test_utils.compare_outputs(test_module.forward, jit_module.forward, lhs, rhs)