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/builder/get_registered_ops.py

12 lines
672 B
Python
Raw Normal View History

Move tests around so we can have dedicated tests for the c10 dispatcher. * Adds a trivial missing test for _torch_mlir.c10.get_registered_ops() * Disables the regression tests for now on c10 (until implemented).
2020-09-25 09:26:51 +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
Move tests around so we can have dedicated tests for the c10 dispatcher. * Adds a trivial missing test for _torch_mlir.c10.get_registered_ops() * Disables the regression tests for now on c10 (until implemented).
2020-09-25 09:26:51 +08:00
import _torch_mlir
# This check is just for a built-in op that is unlikely to change (and is
# otherwise insignificant).
Add recognition/folder/lowering for aten::__is__, aten::ne.int, and aten::dim Interestingly, TorchScript has its own op (`torch::jit::Operator`) registry separate from the dispatcher (it is a superset of the dispatcher). This is where the "prim" ops and some "aten" ops (that should probably be renamed to "prim") live. In particular, `aten::__is__` is in that latter category of "aten but really prim". This registry is also the source of truth for what the TorchScript interpreter calls into when it executes. The bulk of the "not part of the dispatcher" ops live in https://github.com/pytorch/pytorch/blob/09feb5f579d40b550559360870907963e20c6954/torch/csrc/jit/runtime/register_prim_ops.cpp#L82 And the registry itself lives in: https://github.com/pytorch/pytorch/blob/09feb5f579d40b550559360870907963e20c6954/torch/csrc/jit/runtime/operator.cpp#L196 This fold further reduces the IR of ResNet by folding away some more not-taken branches. These not-taken branches in ResNet require first-class handling of the list type which we don't yet have on any backend.
2021-04-28 06:15:50 +08:00
# CHECK: {'name': ('aten::mul', 'Tensor'), 'is_c10_op': True, 'is_vararg': False, 'is_varret': False, 'is_mutable': False, 'arguments': [{'name': 'self', 'type': 'Tensor', 'pytype': 'Tensor'}, {'name': 'other', 'type': 'Tensor', 'pytype': 'Tensor'}], 'returns': [{'name': '', 'type': 'Tensor', 'pytype': '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
print('\n\n'.join([repr(r) for r in _torch_mlir.get_registered_ops()]))