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

153 lines
6.5 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 prim::device and handle derefining for prim::CallMethod
2021-03-11 09:25:39 +08:00
# CHECK-SAME: %[[ARG:.*]]: !torch.optional<i64>) -> i64 {
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: %[[NONE:.*]] = basicpy.singleton : !basicpy.NoneType
# CHECK: %[[C3:.*]] = constant 3 : i64
Add prim::device and handle derefining for prim::CallMethod
2021-03-11 09:25:39 +08:00
# CHECK: %[[IS_NONE:.*]] = torch.kernel_call "aten::__is__" %[[ARG]], %[[NONE]] : (!torch.optional<i64>, !basicpy.NoneType) -> !basicpy.BoolType
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: %[[COND:.*]] = basicpy.bool_cast %[[IS_NONE]] : !basicpy.BoolType -> i1
# CHECK: %[[RESULT:.*]] = scf.if %[[COND]] -> (i64) {
# CHECK: scf.yield %[[C3]] : i64
# CHECK: } else {
Add prim::device and handle derefining for prim::CallMethod
2021-03-11 09:25:39 +08:00
# CHECK: %[[CASTED:.*]] = torch.prim.unchecked_cast %[[ARG]] : !torch.optional<i64> -> i64
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: 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
Miscellaneous changes while trying to work on ResNet18 - Move frontend lowering pipelines to c++ (this helps with reproducing failures in npcomp-opt) - Add debugging printouts when compilation fails on RefBackendTestConfig The experience now when a test fails during MLIR lowering is now like this: ``` NPCOMP TorchScript Object Graph IR -> NPCOMP Backend IR lowering failed with the following diagnostics: failed to legalize operation 'torch.global_slot' Module does not conform to npcomp's backend contract. See dialect conversion legality information above. Error can be reproduced with: $ npcomp-opt -torchscript-to-npcomp-backend-pipeline /tmp/ResNet18Module.mlir ``` And when TorchScript->MLIR import fails it looks like this: ``` PyTorch TorchScript module -> NPCOMP Object Graph IR import failed with the following diagnostics: unhandled prim operation: %18 : int = prim::min(%17) # /usr/local/google/home/silvasean/.local/lib/python3.9/site-packages/torch/nn/functional.py:4532:4 ``` Also, - Add `--filter=<regex>` to e2e test harness to filter tests. - Add a few prim ops that were needed to import ResNet18 - Fix torch.prim.Loop.condition assemblyFormat (it previously would not round-trip in the case of no loop-carried variables)
2021-04-22 06:07:15 +08:00
# CHECK-LABEL: func @__torch__.prim_layout(
# CHECK-SAME: %[[ARG:.*]]: !numpy.ndarray<*:!numpy.any_dtype>) -> i64 {
# CHECK: %[[RET:.*]] = torch.prim.layout %[[ARG]] : !numpy.ndarray<*:!numpy.any_dtype> -> i64
# CHECK: return %[[RET]] : i64
@mb.import_function
@torch.jit.script
def prim_layout(x):
return x.layout
Add prim::device and handle derefining for prim::CallMethod
2021-03-11 09:25:39 +08:00
# CHECK-LABEL: func @__torch__.prim_device(
# CHECK-SAME: %[[ARG:.*]]: !numpy.ndarray<*:!numpy.any_dtype>) -> !torch.Device {
# CHECK: %[[RET:.*]] = torch.prim.device %[[ARG]] : !numpy.ndarray<*:!numpy.any_dtype> -> !torch.Device
# CHECK: return %[[RET]] : !torch.Device
@mb.import_function
@torch.jit.script
def prim_device(x):
return x.device
Miscellaneous changes while trying to work on ResNet18 - Move frontend lowering pipelines to c++ (this helps with reproducing failures in npcomp-opt) - Add debugging printouts when compilation fails on RefBackendTestConfig The experience now when a test fails during MLIR lowering is now like this: ``` NPCOMP TorchScript Object Graph IR -> NPCOMP Backend IR lowering failed with the following diagnostics: failed to legalize operation 'torch.global_slot' Module does not conform to npcomp's backend contract. See dialect conversion legality information above. Error can be reproduced with: $ npcomp-opt -torchscript-to-npcomp-backend-pipeline /tmp/ResNet18Module.mlir ``` And when TorchScript->MLIR import fails it looks like this: ``` PyTorch TorchScript module -> NPCOMP Object Graph IR import failed with the following diagnostics: unhandled prim operation: %18 : int = prim::min(%17) # /usr/local/google/home/silvasean/.local/lib/python3.9/site-packages/torch/nn/functional.py:4532:4 ``` Also, - Add `--filter=<regex>` to e2e test harness to filter tests. - Add a few prim ops that were needed to import ResNet18 - Fix torch.prim.Loop.condition assemblyFormat (it previously would not round-trip in the case of no loop-carried variables)
2021-04-22 06:07:15 +08:00
# CHECK-LABEL: func @__torch__.prim_min(
# CHECK-SAME: %[[ARG:.*]]: i64) -> !basicpy.TupleType {
# CHECK: %[[SINGLETON:.*]] = basicpy.build_list %[[ARG]] : (i64) -> !basicpy.ListType
# CHECK: %[[MIN1:.*]] = torch.prim.min %[[SINGLETON]] : !basicpy.ListType -> i64
# CHECK: %[[MIN2:.*]] = torch.prim.min %[[ARG]], %[[ARG]] : i64, i64 -> i64
# CHECK: %[[ARG_3_TIMES:.*]] = basicpy.build_list %[[ARG]], %[[ARG]], %[[ARG]] : (i64, i64, i64) -> !basicpy.ListType
# CHECK: %[[MIN3:.*]] = torch.prim.min %[[ARG_3_TIMES]] : !basicpy.ListType -> i64
# CHECK: %[[RET:.*]] = basicpy.build_tuple %[[MIN1]], %[[MIN2]], %[[MIN3]] : (i64, i64, i64) -> !basicpy.TupleType
# CHECK: return %[[RET]] : !basicpy.TupleType
@mb.import_function
@torch.jit.script
def prim_min(x: int):
return min(x), min(x,x), min(x, x, x)
# CHECK-LABEL: func @__torch__.prim_max(
# CHECK-SAME: %[[ARG:.*]]: i64) -> !basicpy.TupleType {
# CHECK: %[[SINGLETON:.*]] = basicpy.build_list %[[ARG]] : (i64) -> !basicpy.ListType
# CHECK: %[[MAX1:.*]] = torch.prim.max %[[SINGLETON]] : !basicpy.ListType -> i64
# CHECK: %[[MAX2:.*]] = torch.prim.max %[[ARG]], %[[ARG]] : i64, i64 -> i64
# CHECK: %[[ARG_3_TIMES:.*]] = basicpy.build_list %[[ARG]], %[[ARG]], %[[ARG]] : (i64, i64, i64) -> !basicpy.ListType
# CHECK: %[[MAX3:.*]] = torch.prim.max %[[ARG_3_TIMES]] : !basicpy.ListType -> i64
# CHECK: %[[RET:.*]] = basicpy.build_tuple %[[MAX1]], %[[MAX2]], %[[MAX3]] : (i64, i64, i64) -> !basicpy.TupleType
# CHECK: return %[[RET]] : !basicpy.TupleType
@mb.import_function
@torch.jit.script
def prim_max(x: int):
return max(x), max(x,x), max(x, x, x)
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()