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/if.py

55 lines
2.1 KiB
Python
Raw Normal View History

Add support for prim::GetAttr/SetAttr/CallMethod/If This required some invasive surgery to graph_importer.h/cpp, specifically moving most of it into node_importer.h/cpp and relayering it. The abstraction that it had didn't work well in the recursive setting that happens with prim::If. The key observation is that torch::jit::Graph doesn't really correspond directly to anything on the MLIR side. It's a weird combination of a context, builder, and function and just holds a `torch::jit::Block`. It is `torch::jit::Node` and `torch::jit::Block` which form the recursive structure analogous to MLIR's operation/region/block. So node_importer.h/cpp makes sense as a core building block. As part of doing this, I did venture a bit into the AcapController code, and realize now that there is functionality duplicated there with the ivalue importer. Will refactor that soon.
2021-02-02 09:59:42 +08:00
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See frontends/pytorch/LICENSE for license information.
import torch
import torch_mlir
# 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: @__torch__.prim_If(
# CHECK-SAME: %[[B:.*]]: !basicpy.BoolType,
# CHECK-SAME: %[[I:.*]]: i64) -> i64 {
Add support for prim::GetAttr/SetAttr/CallMethod/If This required some invasive surgery to graph_importer.h/cpp, specifically moving most of it into node_importer.h/cpp and relayering it. The abstraction that it had didn't work well in the recursive setting that happens with prim::If. The key observation is that torch::jit::Graph doesn't really correspond directly to anything on the MLIR side. It's a weird combination of a context, builder, and function and just holds a `torch::jit::Block`. It is `torch::jit::Node` and `torch::jit::Block` which form the recursive structure analogous to MLIR's operation/region/block. So node_importer.h/cpp makes sense as a core building block. As part of doing this, I did venture a bit into the AcapController code, and realize now that there is functionality duplicated there with the ivalue importer. Will refactor that soon.
2021-02-02 09:59:42 +08:00
@mb.import_function
@torch.jit.script
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
def prim_If(b: bool, i: int):
Add support for prim::GetAttr/SetAttr/CallMethod/If This required some invasive surgery to graph_importer.h/cpp, specifically moving most of it into node_importer.h/cpp and relayering it. The abstraction that it had didn't work well in the recursive setting that happens with prim::If. The key observation is that torch::jit::Graph doesn't really correspond directly to anything on the MLIR side. It's a weird combination of a context, builder, and function and just holds a `torch::jit::Block`. It is `torch::jit::Node` and `torch::jit::Block` which form the recursive structure analogous to MLIR's operation/region/block. So node_importer.h/cpp makes sense as a core building block. As part of doing this, I did venture a bit into the AcapController code, and realize now that there is functionality duplicated there with the ivalue importer. Will refactor that soon.
2021-02-02 09:59:42 +08:00
# CHECK: %[[I1:.*]] = basicpy.bool_cast %[[B]] : !basicpy.BoolType -> i1
# CHECK: %[[RES:.*]] = scf.if %[[I1]] -> (i64) {
# CHECK: %[[ADD:.*]] = torch.kernel_call "aten::add" %[[I]], %[[I]]
# CHECK: scf.yield %[[ADD]] : i64
# CHECK: } else {
# CHECK: %[[MUL:.*]] = torch.kernel_call "aten::mul" %[[I]], %[[I]]
# CHECK: scf.yield %[[MUL]] : i64
# CHECK: }
# CHECK: return %[[RES:.*]] : i64
if b:
return i + i
else:
return i * i
# elif is modeled as a nested if, so no need to specially test it here.
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_If_derefine(
# CHECK-SAME: %[[B:.*]]: !basicpy.BoolType,
# CHECK-SAME: %[[I:.*]]: i64) -> !torch.optional<i64> {
# CHECK: %[[NONE:.*]] = basicpy.singleton : !basicpy.NoneType
# CHECK: %[[PRED:.*]] = basicpy.bool_cast %[[B]] : !basicpy.BoolType -> i1
# CHECK: %[[RES:.*]] = scf.if %[[PRED]] -> (!torch.optional<i64>) {
# CHECK: %[[NONE_DEREFINED:.*]] = torch.derefine %[[NONE]] : !basicpy.NoneType -> !torch.optional<i64>
# CHECK: scf.yield %[[NONE_DEREFINED]] : !torch.optional<i64>
# CHECK: } else {
# CHECK: %[[I_DEREFINED:.*]] = torch.derefine %[[I]] : i64 -> !torch.optional<i64>
# CHECK: scf.yield %[[I_DEREFINED]] : !torch.optional<i64>
# CHECK: }
# CHECK: return %[[RES:.*]] : !torch.optional<i64>
@mb.import_function
@torch.jit.script
def prim_If_derefine(b: bool, i: int):
if b:
return None
return i
Add support for prim::GetAttr/SetAttr/CallMethod/If This required some invasive surgery to graph_importer.h/cpp, specifically moving most of it into node_importer.h/cpp and relayering it. The abstraction that it had didn't work well in the recursive setting that happens with prim::If. The key observation is that torch::jit::Graph doesn't really correspond directly to anything on the MLIR side. It's a weird combination of a context, builder, and function and just holds a `torch::jit::Block`. It is `torch::jit::Node` and `torch::jit::Block` which form the recursive structure analogous to MLIR's operation/region/block. So node_importer.h/cpp makes sense as a core building block. As part of doing this, I did venture a bit into the AcapController code, and realize now that there is functionality duplicated there with the ivalue importer. Will refactor that soon.
2021-02-02 09:59:42 +08:00
mb.module.operation.print()
print()