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/test/c10_dispatch/module_builder.py

27 lines
716 B
Python
Raw Normal View History

Add boilerplate to do device capture (pytorch 1.6). * Uses the new dispatcher API. * Just prints to the console for the moment when an op is captured. * Executes the op through the existing implementation.
2020-09-26 09:13:16 +08:00
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See frontends/pytorch/LICENSE for license information.
# RUN: python %s | FileCheck %s
Add capture function arguments. * Adds at::Tensor -> MlirValue tracking. * Adds conversions for tensor and scalar types to MLIR types. * Adds npcomp C APIs for constructing custom types. * Reworks pybind include so as to get Torch pybind helpers (needed to pass at::Tensor type from Python->C++).
2020-10-02 09:59:58 +08:00
# TODO: Once stabilized, expand tests to include all argument dtypes.
Add boilerplate to do device capture (pytorch 1.6). * Uses the new dispatcher API. * Just prints to the console for the moment when an op is captured. * Executes the op through the existing implementation.
2020-09-26 09:13:16 +08:00
import torch
import _torch_mlir as m
t0 = torch.randn((4,4))
t1 = torch.randn((4,4))
Add public API for constructing a module/function to capture PyTorch ops. * Uses the MLIR-C API since that will save us a lot of grief down the road (i.e. will give PyTorch and libMLIR/libNPCOMP the ability to skew version-wise). * Quite a few TODOs and not yet populating the function in any way.
2020-09-29 09:36:00 +08:00
mb = m.ModuleBuilder()
Add capture function arguments. * Adds at::Tensor -> MlirValue tracking. * Adds conversions for tensor and scalar types to MLIR types. * Adds npcomp C APIs for constructing custom types. * Reworks pybind include so as to get Torch pybind helpers (needed to pass at::Tensor type from Python->C++).
2020-10-02 09:59:58 +08:00
with mb.capture_function("foobar", [t0, t1]) as c:
Add boilerplate to do device capture (pytorch 1.6). * Uses the new dispatcher API. * Just prints to the console for the moment when an op is captured. * Executes the op through the existing implementation.
2020-09-26 09:13:16 +08:00
result = t0 + t1
Add public API for constructing a module/function to capture PyTorch ops. * Uses the MLIR-C API since that will save us a lot of grief down the road (i.e. will give PyTorch and libMLIR/libNPCOMP the ability to skew version-wise). * Quite a few TODOs and not yet populating the function in any way.
2020-09-29 09:36:00 +08:00
# CHECK: module {
Add capture function arguments. * Adds at::Tensor -> MlirValue tracking. * Adds conversions for tensor and scalar types to MLIR types. * Adds npcomp C APIs for constructing custom types. * Reworks pybind include so as to get Torch pybind helpers (needed to pass at::Tensor type from Python->C++).
2020-10-02 09:59:58 +08:00
# CHECK: func @foobar(%arg0: !numpy.ndarray<[4,4]:f32>, %arg1: !numpy.ndarray<[4,4]:f32>) {
# CHECK: return
Add public API for constructing a module/function to capture PyTorch ops. * Uses the MLIR-C API since that will save us a lot of grief down the road (i.e. will give PyTorch and libMLIR/libNPCOMP the ability to skew version-wise). * Quite a few TODOs and not yet populating the function in any way.
2020-09-29 09:36:00 +08:00
# CHECK: }
# CHECK: }
print(mb)
Add boilerplate to do device capture (pytorch 1.6). * Uses the new dispatcher API. * Just prints to the console for the moment when an op is captured. * Executes the op through the existing implementation.
2020-09-26 09:13:16 +08:00
# CHECK: CAPTURE: aten::add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> (Tensor)
Add public API for constructing a module/function to capture PyTorch ops. * Uses the MLIR-C API since that will save us a lot of grief down the road (i.e. will give PyTorch and libMLIR/libNPCOMP the ability to skew version-wise). * Quite a few TODOs and not yet populating the function in any way.
2020-09-29 09:36:00 +08:00
for line in c.get_debug_log(): print(line)