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

33 lines
918 B
Python
Raw Normal View History

Add cos_e2e.py, test_utils and support for tensor inputs (#134)
2020-11-25 11:02:50 +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 cos_e2e.py, test_utils and support for tensor inputs (#134)
2020-11-25 11:02:50 +08:00
from npcomp.compiler.utils import logging
import test_utils
logging.enable()
torch.manual_seed(0)
input = torch.rand(2, 3)
mb = torch_mlir.ModuleBuilder()
with mb.capture_function("cos", [input]) as f:
result = torch.cos(input)
f.returns([result])
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 cos_e2e.py, test_utils and support for tensor inputs (#134)
2020-11-25 11:02:50 +08:00
logging.debug(f"Executing jit_module.cos")
test_utils.compare_outputs(torch.cos, jit_module.cos, input)
# This fails because ModuleBuilder represents torch.cos with a constant:
# https://github.com/llvm/mlir-npcomp/issues/135
test_utils.compare_outputs(torch.cos, jit_module.cos, input + 1)