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

107 lines
4.0 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()
Properly model "derefinement". In terms of IR structure, TorchScript allows types to vary in many circumstances where MLIR requires pointer-identical types. In particular, it is valid to pass any subtype in place of a type. For example, if an `Optional[int]` is required somewhere in the IR, it is legal to pass a value of just `int` (but not the other way around; see `torch.prim.unchecked_cast`). In effect, every *use* can have a different type. We introduce a new op `torch.derefine` that models that impedance mismatch. This op allows casting a value from one type to a type that it is a subtype of to model this behavior. Recommended review order: - TorchOps.td for new torch.derefine (and updated docs for `torch.prim.unchecked_cast`) - new test code in if.py, loop.py, function-derefine.py - new code in node_importer.cpp for handling derefinement insertion - function_importer.cpp and utils changes in torch_to_mlir_utils.cpp Properly handling derefinement on function boundaries required relayering the code so that graph_importer.cpp/.h is now function_importer.cpp/.h because only the `torch::jit::Function` (actually the `c10::FunctionSchema` it holds) knows the derefined types that are actually needed at the boundary (see `function-derefine.py` for a test). Annoyingly, this churns all the functions which are now prefixed with `__torch__.` but that is more correct anyway (that is their linkage name in the `torch::jit::CompilationUnit`; the previous `mb.import_function` was actually buggy in the case of functions calling each other as it would reference their unqualified name). With this change, we can import `resnet18` from `torchvision` :) IR: https://gist.github.com/silvasean/6426a5272d8a6c7caae533fce05ab704
2021-03-02 09:24:15 +08:00
# CHECK-LABEL: func @__torch__.prim_NumToTensor(
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
# 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 model "derefinement". In terms of IR structure, TorchScript allows types to vary in many circumstances where MLIR requires pointer-identical types. In particular, it is valid to pass any subtype in place of a type. For example, if an `Optional[int]` is required somewhere in the IR, it is legal to pass a value of just `int` (but not the other way around; see `torch.prim.unchecked_cast`). In effect, every *use* can have a different type. We introduce a new op `torch.derefine` that models that impedance mismatch. This op allows casting a value from one type to a type that it is a subtype of to model this behavior. Recommended review order: - TorchOps.td for new torch.derefine (and updated docs for `torch.prim.unchecked_cast`) - new test code in if.py, loop.py, function-derefine.py - new code in node_importer.cpp for handling derefinement insertion - function_importer.cpp and utils changes in torch_to_mlir_utils.cpp Properly handling derefinement on function boundaries required relayering the code so that graph_importer.cpp/.h is now function_importer.cpp/.h because only the `torch::jit::Function` (actually the `c10::FunctionSchema` it holds) knows the derefined types that are actually needed at the boundary (see `function-derefine.py` for a test). Annoyingly, this churns all the functions which are now prefixed with `__torch__.` but that is more correct anyway (that is their linkage name in the `torch::jit::CompilationUnit`; the previous `mb.import_function` was actually buggy in the case of functions calling each other as it would reference their unqualified name). With this change, we can import `resnet18` from `torchvision` :) IR: https://gist.github.com/silvasean/6426a5272d8a6c7caae533fce05ab704
2021-03-02 09:24:15 +08:00
# CHECK-LABEL: func @__torch__.prim_Print(
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-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)
Properly model "derefinement". In terms of IR structure, TorchScript allows types to vary in many circumstances where MLIR requires pointer-identical types. In particular, it is valid to pass any subtype in place of a type. For example, if an `Optional[int]` is required somewhere in the IR, it is legal to pass a value of just `int` (but not the other way around; see `torch.prim.unchecked_cast`). In effect, every *use* can have a different type. We introduce a new op `torch.derefine` that models that impedance mismatch. This op allows casting a value from one type to a type that it is a subtype of to model this behavior. Recommended review order: - TorchOps.td for new torch.derefine (and updated docs for `torch.prim.unchecked_cast`) - new test code in if.py, loop.py, function-derefine.py - new code in node_importer.cpp for handling derefinement insertion - function_importer.cpp and utils changes in torch_to_mlir_utils.cpp Properly handling derefinement on function boundaries required relayering the code so that graph_importer.cpp/.h is now function_importer.cpp/.h because only the `torch::jit::Function` (actually the `c10::FunctionSchema` it holds) knows the derefined types that are actually needed at the boundary (see `function-derefine.py` for a test). Annoyingly, this churns all the functions which are now prefixed with `__torch__.` but that is more correct anyway (that is their linkage name in the `torch::jit::CompilationUnit`; the previous `mb.import_function` was actually buggy in the case of functions calling each other as it would reference their unqualified name). With this change, we can import `resnet18` from `torchvision` :) IR: https://gist.github.com/silvasean/6426a5272d8a6c7caae533fce05ab704
2021-03-02 09:24:15 +08:00
# CHECK-LABEL: func @__torch__.prim_RaiseException() -> !basicpy.NoneType {
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: %[[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")
Properly model "derefinement". In terms of IR structure, TorchScript allows types to vary in many circumstances where MLIR requires pointer-identical types. In particular, it is valid to pass any subtype in place of a type. For example, if an `Optional[int]` is required somewhere in the IR, it is legal to pass a value of just `int` (but not the other way around; see `torch.prim.unchecked_cast`). In effect, every *use* can have a different type. We introduce a new op `torch.derefine` that models that impedance mismatch. This op allows casting a value from one type to a type that it is a subtype of to model this behavior. Recommended review order: - TorchOps.td for new torch.derefine (and updated docs for `torch.prim.unchecked_cast`) - new test code in if.py, loop.py, function-derefine.py - new code in node_importer.cpp for handling derefinement insertion - function_importer.cpp and utils changes in torch_to_mlir_utils.cpp Properly handling derefinement on function boundaries required relayering the code so that graph_importer.cpp/.h is now function_importer.cpp/.h because only the `torch::jit::Function` (actually the `c10::FunctionSchema` it holds) knows the derefined types that are actually needed at the boundary (see `function-derefine.py` for a test). Annoyingly, this churns all the functions which are now prefixed with `__torch__.` but that is more correct anyway (that is their linkage name in the `torch::jit::CompilationUnit`; the previous `mb.import_function` was actually buggy in the case of functions calling each other as it would reference their unqualified name). With this change, we can import `resnet18` from `torchvision` :) IR: https://gist.github.com/silvasean/6426a5272d8a6c7caae533fce05ab704
2021-03-02 09:24:15 +08:00
# CHECK-LABEL: func @__torch__.prim_unchecked_cast(
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-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
Properly model "derefinement". In terms of IR structure, TorchScript allows types to vary in many circumstances where MLIR requires pointer-identical types. In particular, it is valid to pass any subtype in place of a type. For example, if an `Optional[int]` is required somewhere in the IR, it is legal to pass a value of just `int` (but not the other way around; see `torch.prim.unchecked_cast`). In effect, every *use* can have a different type. We introduce a new op `torch.derefine` that models that impedance mismatch. This op allows casting a value from one type to a type that it is a subtype of to model this behavior. Recommended review order: - TorchOps.td for new torch.derefine (and updated docs for `torch.prim.unchecked_cast`) - new test code in if.py, loop.py, function-derefine.py - new code in node_importer.cpp for handling derefinement insertion - function_importer.cpp and utils changes in torch_to_mlir_utils.cpp Properly handling derefinement on function boundaries required relayering the code so that graph_importer.cpp/.h is now function_importer.cpp/.h because only the `torch::jit::Function` (actually the `c10::FunctionSchema` it holds) knows the derefined types that are actually needed at the boundary (see `function-derefine.py` for a test). Annoyingly, this churns all the functions which are now prefixed with `__torch__.` but that is more correct anyway (that is their linkage name in the `torch::jit::CompilationUnit`; the previous `mb.import_function` was actually buggy in the case of functions calling each other as it would reference their unqualified name). With this change, we can import `resnet18` from `torchvision` :) IR: https://gist.github.com/silvasean/6426a5272d8a6c7caae533fce05ab704
2021-03-02 09:24:15 +08:00
# CHECK-LABEL: func @__torch__.prim_TupleUnpack(
[prim] Add support for List|TupleUnpack
2021-03-03 06:39:48 +08:00
# 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
[chore] Make variable names in prim.py more clear
2021-03-03 06:42:21 +08:00
def prim_TupleUnpack(tup: typing.Tuple[int, int]):
val, _ = tup
[prim] Add support for List|TupleUnpack
2021-03-03 06:39:48 +08:00
return val
Properly model "derefinement". In terms of IR structure, TorchScript allows types to vary in many circumstances where MLIR requires pointer-identical types. In particular, it is valid to pass any subtype in place of a type. For example, if an `Optional[int]` is required somewhere in the IR, it is legal to pass a value of just `int` (but not the other way around; see `torch.prim.unchecked_cast`). In effect, every *use* can have a different type. We introduce a new op `torch.derefine` that models that impedance mismatch. This op allows casting a value from one type to a type that it is a subtype of to model this behavior. Recommended review order: - TorchOps.td for new torch.derefine (and updated docs for `torch.prim.unchecked_cast`) - new test code in if.py, loop.py, function-derefine.py - new code in node_importer.cpp for handling derefinement insertion - function_importer.cpp and utils changes in torch_to_mlir_utils.cpp Properly handling derefinement on function boundaries required relayering the code so that graph_importer.cpp/.h is now function_importer.cpp/.h because only the `torch::jit::Function` (actually the `c10::FunctionSchema` it holds) knows the derefined types that are actually needed at the boundary (see `function-derefine.py` for a test). Annoyingly, this churns all the functions which are now prefixed with `__torch__.` but that is more correct anyway (that is their linkage name in the `torch::jit::CompilationUnit`; the previous `mb.import_function` was actually buggy in the case of functions calling each other as it would reference their unqualified name). With this change, we can import `resnet18` from `torchvision` :) IR: https://gist.github.com/silvasean/6426a5272d8a6c7caae533fce05ab704
2021-03-02 09:24:15 +08:00
# CHECK-LABEL: func @__torch__.prim_TupleIndex(
[prim] Add TupleIndex support I could not find a corresponding ListIndex in prim, which seems to translate to a __get_attr__ under the hood. I think the reason a tuple Index op can exist is because Tuple's are supposed to be frozen, where List operands can be mutable.
2021-03-03 07:26:03 +08:00
# CHECK-SAME: %[[ARG:.*]]: !basicpy.TupleType) -> i64 {
# CHECK: %[[RET:.*]] = torch.prim.TupleIndex %[[ARG]], %[[IDX:.*]] : !basicpy.TupleType, i64 -> i64
# CHECK: return %[[RET]] : i64
@mb.import_function
@torch.jit.script
def prim_TupleIndex(tup: typing.Tuple[int, int]):
return tup[0]
Properly model "derefinement". In terms of IR structure, TorchScript allows types to vary in many circumstances where MLIR requires pointer-identical types. In particular, it is valid to pass any subtype in place of a type. For example, if an `Optional[int]` is required somewhere in the IR, it is legal to pass a value of just `int` (but not the other way around; see `torch.prim.unchecked_cast`). In effect, every *use* can have a different type. We introduce a new op `torch.derefine` that models that impedance mismatch. This op allows casting a value from one type to a type that it is a subtype of to model this behavior. Recommended review order: - TorchOps.td for new torch.derefine (and updated docs for `torch.prim.unchecked_cast`) - new test code in if.py, loop.py, function-derefine.py - new code in node_importer.cpp for handling derefinement insertion - function_importer.cpp and utils changes in torch_to_mlir_utils.cpp Properly handling derefinement on function boundaries required relayering the code so that graph_importer.cpp/.h is now function_importer.cpp/.h because only the `torch::jit::Function` (actually the `c10::FunctionSchema` it holds) knows the derefined types that are actually needed at the boundary (see `function-derefine.py` for a test). Annoyingly, this churns all the functions which are now prefixed with `__torch__.` but that is more correct anyway (that is their linkage name in the `torch::jit::CompilationUnit`; the previous `mb.import_function` was actually buggy in the case of functions calling each other as it would reference their unqualified name). With this change, we can import `resnet18` from `torchvision` :) IR: https://gist.github.com/silvasean/6426a5272d8a6c7caae533fce05ab704
2021-03-02 09:24:15 +08:00
# CHECK-LABEL: func @__torch__.prim_ListUnpack(
[prim] Add support for List|TupleUnpack
2021-03-03 06:39:48 +08:00
# 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
[chore] Make variable names in prim.py more clear
2021-03-03 06:42:21 +08:00
def prim_ListUnpack(l: typing.List[int]):
_, val, _ = l
[prim] Add support for List|TupleUnpack
2021-03-03 06:39:48 +08:00
return val
prim::dtype op
2021-03-11 08:41:18 +08:00
# CHECK-LABEL: func @__torch__.prim_dtype(
# CHECK-SAME: %[[ARG:.*]]: !numpy.ndarray<*:!numpy.any_dtype>) -> i64 {
# CHECK: %[[RET:.*]] = torch.prim.dtype %[[ARG]] : !numpy.ndarray<*:!numpy.any_dtype> -> i64
# CHECK: return %[[RET]] : i64
@mb.import_function
@torch.jit.script
def prim_dtype(x):
return x.dtype
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()