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

30 lines
1000 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.
Update test configuration to import mlir from LLVM install location. * Also adds two lit tests to verify that all of our extensions load without fireworks, which is a good indication that the shared library deps are sane. * Bumps llvm-project version to use D89167.
2020-10-10 07:21:01 +08:00
# RUN: %PYTHON %s | FileCheck %s
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
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
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
t0 = torch.randn((1,4))
t1 = torch.randn((4,1))
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
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()
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
with mb.capture_function("foobar", [t0, t1]) as f:
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
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
f.returns([result])
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
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 {
Set PyTorch captured function return type. * Resolves various TODOs that required an LLVM change/bump. * Bumps LLVM to 4aa217160e5f06a96c6effc4950c3b402374de58
2020-10-08 01:14:34 +08:00
# CHECK: func @foobar(%arg0: !numpy.ndarray<[1,4]:f32>, %arg1: !numpy.ndarray<[4,1]:f32>) -> !numpy.ndarray<[4,4]:f32> {
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
# CHECK: %c1_i64 = constant 1 : i64
# CHECK: %0 = torch.kernel_call "aten::add" %arg0, %arg1, %c1_i64 : (!numpy.ndarray<[1,4]:f32>, !numpy.ndarray<[4,1]:f32>, i64) -> !numpy.ndarray<[4,4]:f32>
# CHECK: return %0 : !numpy.ndarray<[4,4]:f32>
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)
Implement torch.kernel_call capture. * Had to stop short of modifying the function return signature because of a missing C-API upstream. * Committing here is good enough for a test and will resolve the various TODOs about upstream APIs next.
2020-10-06 14:21:21 +08:00
for line in f.get_debug_log(): print(line)