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/node_import/prim.py

89 lines
3.2 KiB
Python
Raw Normal View History

Add support for prim.NumToTensor With this, we can import BERT! ``` pt_util ~/tmp/bert.pt --import --exported-name=forward \ | npcomp-opt -torch-globalize-object-graph -inline -symbol-dce ``` https://gist.github.com/silvasean/fe7735ff5d065cc9216f7b0346d0e977 The test case here is a bit unconventional -- it isn't actually valid Python. To figure out how to generate it I had to go search the PyTorch codebase for "NumToTensor" and work backward. In this case I found this [code](https://github.com/pytorch/pytorch/blob/649760e5f12b90a29559785ac2b3cb132cb3bf84/torch/csrc/jit/frontend/ir_emitter.cpp#L464) which via a wild guess I was able to turn into a test case. In this case it didn't take me too long, but when doing this kind of "add a bunch of trivial stuff to bring up a real model", I'm starting to think that we might skimp on test cases when it's fairly trivial and not obvious how to test with a small test.
2021-02-26 08:35:29 +08:00
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See frontends/pytorch/LICENSE for license information.
[prim] Add support for List|TupleUnpack
2021-03-03 06:39:48 +08:00
import typing
Add support for prim.NumToTensor With this, we can import BERT! ``` pt_util ~/tmp/bert.pt --import --exported-name=forward \ | npcomp-opt -torch-globalize-object-graph -inline -symbol-dce ``` https://gist.github.com/silvasean/fe7735ff5d065cc9216f7b0346d0e977 The test case here is a bit unconventional -- it isn't actually valid Python. To figure out how to generate it I had to go search the PyTorch codebase for "NumToTensor" and work backward. In this case I found this [code](https://github.com/pytorch/pytorch/blob/649760e5f12b90a29559785ac2b3cb132cb3bf84/torch/csrc/jit/frontend/ir_emitter.cpp#L464) which via a wild guess I was able to turn into a test case. In this case it didn't take me too long, but when doing this kind of "add a bunch of trivial stuff to bring up a real model", I'm starting to think that we might skimp on test cases when it's fairly trivial and not obvious how to test with a small test.
2021-02-26 08:35:29 +08:00
import torch
import torch_mlir
Add support for `prim::unchecked_cast`. This arises when casting optionals, which happens a lot especially around handling of default arguments (python `if arg is None` idiom). In this case, the offending code for the model is in max_pool2d: [code link](https://github.com/pytorch/pytorch/blob/b3bf08e67f5cf29cbfd197e3c9e9a2f45e0b5c9d/torch/nn/functional.py#L657)
2021-03-02 07:26:57 +08:00
import typing
Add support for prim.NumToTensor With this, we can import BERT! ``` pt_util ~/tmp/bert.pt --import --exported-name=forward \ | npcomp-opt -torch-globalize-object-graph -inline -symbol-dce ``` https://gist.github.com/silvasean/fe7735ff5d065cc9216f7b0346d0e977 The test case here is a bit unconventional -- it isn't actually valid Python. To figure out how to generate it I had to go search the PyTorch codebase for "NumToTensor" and work backward. In this case I found this [code](https://github.com/pytorch/pytorch/blob/649760e5f12b90a29559785ac2b3cb132cb3bf84/torch/csrc/jit/frontend/ir_emitter.cpp#L464) which via a wild guess I was able to turn into a test case. In this case it didn't take me too long, but when doing this kind of "add a bunch of trivial stuff to bring up a real model", I'm starting to think that we might skimp on test cases when it's fairly trivial and not obvious how to test with a small test.
2021-02-26 08:35:29 +08:00
# RUN: %PYTHON %s | npcomp-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
# CHECK-LABEL: func @prim_NumToTensor(
# CHECK-SAME: %[[ARG:.*]]: i64) -> !numpy.ndarray<*:!numpy.any_dtype> {
# CHECK: %[[RET:.*]] = torch.prim.NumToTensor %[[ARG]] : i64 -> !numpy.ndarray<*:!numpy.any_dtype>
# CHECK: return %[[RET]] : !numpy.ndarray<*:!numpy.any_dtype>
# CHECK: }
@mb.import_function
@torch.jit.script
def prim_NumToTensor(i: int):
return _to_tensor(i)
Properly import the entire torch::jit::CompilationUnit This primarily unlocks proper handling of free functions (that is, functions that are not methods of any torch.nn.Module). Recommended review order: - `ivalue_importer.cpp` + `ivalue_import/functions*.py` - `GlobalizeObjectGraph.cpp` + test case - misc other stuff The `torch::jit::CompilationUnit` is basically a backing store or "context" holding all the possible functions in the program. The previous code was not explicitly accessing this data structure, since it just imported the `torch::jit::Function`'s that it saw attached to methods. Subtly, any time a TorchScript module called into a free function, the free function gets incorporated into the torch::jit::CompilationUnit, but doesn't show up anywhere when dumping the module, except in the curious pattern: ``` %5 : Function = prim::Constant[name="adaptive_avg_pool2d"]() %6 : Tensor = prim::CallFunction(%5, %input.1, %4) ``` That is, calls are indirect calls, and are accessed via `prim::Constant` materializing a function object. Even stranger, the `name` attribute here doesn't really even tell the full story -- it doesn't correspond to anything. It turns out that the c10::FunctionType itself actually holds a pointer to the `torch::jit::Function` in the compilation unit directly (so there is actually no indirection in prim::CallMethod, because any two values of the same FunctionType call the same function!). E.g. when converting the IR to bytecode, the "name" is ignored [code link](https://github.com/pytorch/pytorch/blob/1d6bd157902d4b1347a5d03122d02b407658e263/torch/csrc/jit/runtime/interpreter.cpp#L937). We do import `prim::CallFunction` as a `std.call_indirect` though because it's more braindead to do it that way (it gets canonicalized to a direct call easily).
2021-02-27 08:20:35 +08:00
# CHECK-LABEL: func @prim_Print(
# CHECK-SAME: %[[ARG:.*]]: !numpy.ndarray<*:!numpy.any_dtype>) -> !basicpy.NoneType {
# CHECK: %[[STR:.*]] = basicpy.bytes_constant "x"
# CHECK: torch.prim.Print(%[[STR]], %[[ARG]]) : !basicpy.BytesType, !numpy.ndarray<*:!numpy.any_dtype>
@mb.import_function
@torch.jit.script
def prim_Print(x):
print("x", x)
Add support for prim::RaiseException. Used by resnet18. It seems to originate from a helper `_verify_batch_size`: [code link](https://github.com/pytorch/pytorch/blob/b3bf08e67f5cf29cbfd197e3c9e9a2f45e0b5c9d/torch/nn/functional.py#L2099). I couldn't find a way to test `prim::RaiseException` without also having `prim::Uninitialized`.
2021-03-02 05:47:50 +08:00
# CHECK-LABEL: func @prim_RaiseException() -> !basicpy.NoneType {
# CHECK: %[[ERRORSTR:.*]] = basicpy.bytes_constant "Error"
# CHECK: %[[NONE:.*]] = torch.prim.Uninitialized : !basicpy.NoneType
# CHECK: torch.prim.RaiseException %[[ERRORSTR]]
# CHECK: return %[[NONE]] : !basicpy.NoneType
@mb.import_function
@torch.jit.script
def prim_RaiseException():
raise Exception("Error")
Add support for `prim::unchecked_cast`. This arises when casting optionals, which happens a lot especially around handling of default arguments (python `if arg is None` idiom). In this case, the offending code for the model is in max_pool2d: [code link](https://github.com/pytorch/pytorch/blob/b3bf08e67f5cf29cbfd197e3c9e9a2f45e0b5c9d/torch/nn/functional.py#L657)
2021-03-02 07:26:57 +08:00
# CHECK-LABEL: func @prim_unchecked_cast(
# CHECK-SAME: %[[VAL_0:.*]]: !torch.optional<i64>) -> i64 {
# CHECK: %[[NONE:.*]] = basicpy.singleton : !basicpy.NoneType
# CHECK: %[[C3:.*]] = constant 3 : i64
# CHECK: %[[IS_NONE:.*]] = torch.kernel_call "aten::__is__" %[[VAL_0]], %[[NONE]] : (!torch.optional<i64>, !basicpy.NoneType) -> !basicpy.BoolType
# CHECK: %[[COND:.*]] = basicpy.bool_cast %[[IS_NONE]] : !basicpy.BoolType -> i1
# CHECK: %[[RESULT:.*]] = scf.if %[[COND]] -> (i64) {
# CHECK: scf.yield %[[C3]] : i64
# CHECK: } else {
# CHECK: %[[CASTED:.*]] = torch.prim.unchecked_cast %[[VAL_0]] : !torch.optional<i64> -> i64
# CHECK: scf.yield %[[CASTED]] : i64
# CHECK: }
# CHECK: return %[[RESULT:.*]] : i64
@mb.import_function
@torch.jit.script
def prim_unchecked_cast(i: typing.Optional[int]):
if i is None:
return 3
return i
[prim] Add support for List|TupleUnpack
2021-03-03 06:39:48 +08:00
# CHECK-LABEL: func @prim_TupleUnpack(
# CHECK-SAME: %[[ARG:.*]]: !basicpy.TupleType) -> i64 {
# CHECK: %[[RET:.*]]:2 = torch.prim.TupleUnpack %[[ARG]] : !basicpy.TupleType -> i64, i64
# CHECK: return %[[RET]]#0 : i64
@mb.import_function
@torch.jit.script
def prim_TupleUnpack(lt: typing.Tuple[int, int]):
val, _ = lt
return val
# CHECK-LABEL: func @prim_ListUnpack(
# CHECK-SAME: %[[ARG:.*]]: !basicpy.ListType) -> i64 {
# CHECK: %[[RET:.*]]:3 = torch.prim.ListUnpack %[[ARG]] : !basicpy.ListType -> i64, i64
# CHECK: return %[[RET]]#1 : i64
@mb.import_function
@torch.jit.script
def prim_ListUnpack(lt: typing.List[int]):
_, val, _ = lt
return val
Add support for prim.NumToTensor With this, we can import BERT! ``` pt_util ~/tmp/bert.pt --import --exported-name=forward \ | npcomp-opt -torch-globalize-object-graph -inline -symbol-dce ``` https://gist.github.com/silvasean/fe7735ff5d065cc9216f7b0346d0e977 The test case here is a bit unconventional -- it isn't actually valid Python. To figure out how to generate it I had to go search the PyTorch codebase for "NumToTensor" and work backward. In this case I found this [code](https://github.com/pytorch/pytorch/blob/649760e5f12b90a29559785ac2b3cb132cb3bf84/torch/csrc/jit/frontend/ir_emitter.cpp#L464) which via a wild guess I was able to turn into a test case. In this case it didn't take me too long, but when doing this kind of "add a bunch of trivial stuff to bring up a real model", I'm starting to think that we might skimp on test cases when it's fairly trivial and not obvious how to test with a small test.
2021-02-26 08:35:29 +08:00
mb.module.operation.print()
print()