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

30 lines
730 B
Python
Raw Normal View History

Add aten.mm op and "test" it e2e. Note that unlike aten.matmul which has dynamic behavior depending on the argument ranks (can do matrix-matrix, matrix-vector, batch matmul, etc.), aten.mm is just a vanilla matrix multiply, which can be lowered precisely to tcf.matmul. The "test" is really just an example that I stared at while getting my feet wet with this. We probably want something that actually tests this as part of `ninja check-npcomp`.
2020-11-21 07:59:55 +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 cos_e2e.py, test_utils and support for tensor inputs (#134)
2020-11-25 11:02:50 +08:00
from npcomp.compiler.pytorch.backend import refjit
Add aten.mm op and "test" it e2e. Note that unlike aten.matmul which has dynamic behavior depending on the argument ranks (can do matrix-matrix, matrix-vector, batch matmul, etc.), aten.mm is just a vanilla matrix multiply, which can be lowered precisely to tcf.matmul. The "test" is really just an example that I stared at while getting my feet wet with this. We probably want something that actually tests this as part of `ninja check-npcomp`.
2020-11-21 07:59:55 +08:00
from npcomp.compiler.utils import logging
Add cos_e2e.py, test_utils and support for tensor inputs (#134)
2020-11-25 11:02:50 +08:00
import test_utils
Add aten.mm op and "test" it e2e. Note that unlike aten.matmul which has dynamic behavior depending on the argument ranks (can do matrix-matrix, matrix-vector, batch matmul, etc.), aten.mm is just a vanilla matrix multiply, which can be lowered precisely to tcf.matmul. The "test" is really just an example that I stared at while getting my feet wet with this. We probably want something that actually tests this as part of `ninja check-npcomp`.
2020-11-21 07:59:55 +08:00
logging.enable()
torch.manual_seed(0)
lhs = torch.rand(2, 3)
rhs = torch.rand(3, 4)
mb = torch_mlir.ModuleBuilder()
with mb.capture_function("mm", [lhs, rhs]) as f:
result = torch.mm(lhs, rhs)
f.returns([result])
Add cos_e2e.py, test_utils and support for tensor inputs (#134)
2020-11-25 11:02:50 +08:00
backend = refjit.CompilerBackend()
Add aten.mm op and "test" it e2e. Note that unlike aten.matmul which has dynamic behavior depending on the argument ranks (can do matrix-matrix, matrix-vector, batch matmul, etc.), aten.mm is just a vanilla matrix multiply, which can be lowered precisely to tcf.matmul. The "test" is really just an example that I stared at while getting my feet wet with this. We probably want something that actually tests this as part of `ninja check-npcomp`.
2020-11-21 07:59:55 +08:00
jit_module = backend.load(backend.compile(mb.module))
Add cos_e2e.py, test_utils and support for tensor inputs (#134)
2020-11-25 11:02:50 +08:00
test_utils.compare_outputs(torch.mm, jit_module.mm, lhs, rhs)
test_utils.compare_outputs(torch.mm, jit_module.mm, lhs + 1, rhs - 1)