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/tanh_out_e2e.py

34 lines
810 B
Python
Raw Normal View History

Add support for "trailing_" and "out" variants of various ops. We already had the `promoteTrailingOutTensor` flag, but weren't using it. A inplaceVariantKernelName flag needed to be added. This change is a little dissatisfying, as the conversions done by the RecognizeKernelsPass are currently non-orthogonal. In particular, `kDropResultAndAliasArg0` probably won't work as intended if mixed with these (we probably need to promote kDropResultAndAliasArg0 to not be an arg-level thing anyway, as we have done with promoteTrailingOutTensor). This involved adding a new op `numpy.overwrite_array`. ``` numpy.overwrite_array %arg2 overwrites %arg0 : tensor<2x3xf32>, !numpy.ndarray<[2,3]:f32> ``` This models the destructive update behavior. Note that in the above op, we cannot simply RAUW %arg0 with a suitably conveted %arg2 (for example, %arg0 might have uses that are not dominated by %arg2, or might have an alias relation with some other array in the program). In general, we need a pass analogous to "SSA-formation" which knows how to see through these to uncover an underlying tensor program. Also, add tanh_out_e2e.py/div_inplace_e2e.py and fix some bitrot in refjit.py which is my running example I'm trying to get working.
2021-03-19 04:13:40 +08:00
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See frontends/pytorch/LICENSE for license information.
import torch
import torch_mlir
import npcomp
Add support for compiling through IREE. Recommended review order: - Changes in frontends/pytorch/examples/ - Changes in python/npcomp/compiler/pytorch/backend/ - Boilerplate for the `npcomp-iree-backend-lower-linkage` pass. This change separates out a `npcomp.compiler.pytorch.backend.frontend_lowering` module that does the common lowering for all backends. The individual compiler backends `npcomp.compiler.pytorch.backend.{refjit,iree}` now accept a loosely defined "TCP + scalar code" IR mix that will be formalized in the future as the interface to codegen backends. This also required adding a small pass `npcomp-iree-backend-lower-linkage` which adds `iree.module.export` onto functions, and layering that into the frontend flow. The pass doesn't require a C++-level dependency on IREE, which is nice for now. TBD how we are going to handle lists (we hope we can get away with sneakerneting some td files and relying on loose IR compatibility). Running through IREE requires the ability to import `iree.compiler` and `iree.runtime`, which can be obtained as follows: ``` python3 -m pip install iree-compiler-snapshot iree-runtime-snapshot -f https://github.com/google/iree/releases/tag/snapshot-20210406.200 PYTHONPATH="${PYTHONPATH}:${MY_IREE_BUILD}/bindings/python/" ``` This patch makes it painfully clear that we don't have any e2e testing harness to really plug into, and also don't have a usable Python API to our compiler stack (something usable in a jupyter notebook). That will be addressed in subsequent commits. We've been flying by the seat of our pants with this `examples` directory that isn't subject to any kind of testing or real usability concerns.
2021-04-09 04:05:16 +08:00
from npcomp.compiler.pytorch.backend import refjit, frontend_lowering
Add support for "trailing_" and "out" variants of various ops. We already had the `promoteTrailingOutTensor` flag, but weren't using it. A inplaceVariantKernelName flag needed to be added. This change is a little dissatisfying, as the conversions done by the RecognizeKernelsPass are currently non-orthogonal. In particular, `kDropResultAndAliasArg0` probably won't work as intended if mixed with these (we probably need to promote kDropResultAndAliasArg0 to not be an arg-level thing anyway, as we have done with promoteTrailingOutTensor). This involved adding a new op `numpy.overwrite_array`. ``` numpy.overwrite_array %arg2 overwrites %arg0 : tensor<2x3xf32>, !numpy.ndarray<[2,3]:f32> ``` This models the destructive update behavior. Note that in the above op, we cannot simply RAUW %arg0 with a suitably conveted %arg2 (for example, %arg0 might have uses that are not dominated by %arg2, or might have an alias relation with some other array in the program). In general, we need a pass analogous to "SSA-formation" which knows how to see through these to uncover an underlying tensor program. Also, add tanh_out_e2e.py/div_inplace_e2e.py and fix some bitrot in refjit.py which is my running example I'm trying to get working.
2021-03-19 04:13:40 +08:00
from npcomp.compiler.utils import logging
import test_utils
logging.enable()
torch.manual_seed(0)
arg0 = torch.ones(2, 2)
def fun(a):
z = torch.zeros(2, 2)
torch.tanh(a, out=z)
return z
mb = torch_mlir.ModuleBuilder()
with mb.capture_function("test", [arg0]) as f:
f.returns([fun(arg0)])
backend = refjit.CompilerBackend()
Add support for compiling through IREE. Recommended review order: - Changes in frontends/pytorch/examples/ - Changes in python/npcomp/compiler/pytorch/backend/ - Boilerplate for the `npcomp-iree-backend-lower-linkage` pass. This change separates out a `npcomp.compiler.pytorch.backend.frontend_lowering` module that does the common lowering for all backends. The individual compiler backends `npcomp.compiler.pytorch.backend.{refjit,iree}` now accept a loosely defined "TCP + scalar code" IR mix that will be formalized in the future as the interface to codegen backends. This also required adding a small pass `npcomp-iree-backend-lower-linkage` which adds `iree.module.export` onto functions, and layering that into the frontend flow. The pass doesn't require a C++-level dependency on IREE, which is nice for now. TBD how we are going to handle lists (we hope we can get away with sneakerneting some td files and relying on loose IR compatibility). Running through IREE requires the ability to import `iree.compiler` and `iree.runtime`, which can be obtained as follows: ``` python3 -m pip install iree-compiler-snapshot iree-runtime-snapshot -f https://github.com/google/iree/releases/tag/snapshot-20210406.200 PYTHONPATH="${PYTHONPATH}:${MY_IREE_BUILD}/bindings/python/" ``` This patch makes it painfully clear that we don't have any e2e testing harness to really plug into, and also don't have a usable Python API to our compiler stack (something usable in a jupyter notebook). That will be addressed in subsequent commits. We've been flying by the seat of our pants with this `examples` directory that isn't subject to any kind of testing or real usability concerns.
2021-04-09 04:05:16 +08:00
jit_module = backend.load(backend.compile(frontend_lowering.lower_module(mb.module)))
Add support for "trailing_" and "out" variants of various ops. We already had the `promoteTrailingOutTensor` flag, but weren't using it. A inplaceVariantKernelName flag needed to be added. This change is a little dissatisfying, as the conversions done by the RecognizeKernelsPass are currently non-orthogonal. In particular, `kDropResultAndAliasArg0` probably won't work as intended if mixed with these (we probably need to promote kDropResultAndAliasArg0 to not be an arg-level thing anyway, as we have done with promoteTrailingOutTensor). This involved adding a new op `numpy.overwrite_array`. ``` numpy.overwrite_array %arg2 overwrites %arg0 : tensor<2x3xf32>, !numpy.ndarray<[2,3]:f32> ``` This models the destructive update behavior. Note that in the above op, we cannot simply RAUW %arg0 with a suitably conveted %arg2 (for example, %arg0 might have uses that are not dominated by %arg2, or might have an alias relation with some other array in the program). In general, we need a pass analogous to "SSA-formation" which knows how to see through these to uncover an underlying tensor program. Also, add tanh_out_e2e.py/div_inplace_e2e.py and fix some bitrot in refjit.py which is my running example I'm trying to get working.
2021-03-19 04:13:40 +08:00
test_utils.compare_outputs(torch.mm, jit_module.test, arg0, arg1)
test_utils.compare_outputs(torch.mm, jit_module.test, arg0 + 1, arg1 + 1)