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

73 lines
3.5 KiB
Python
Raw Normal View History

Add support for prim::Loop op. This is a funny one. It combines a `for` and `while` loop in one op. We will need to write some conversions to `scf`.
2021-03-02 07:00:32 +08:00
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See frontends/pytorch/LICENSE for license information.
import torch
import torch_mlir
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
import typing
Add support for prim::Loop op. This is a funny one. It combines a `for` and `while` loop in one op. We will need to write some conversions to `scf`.
2021-03-02 07:00:32 +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_Loop_forlike(
Add a !torch.float type. This removes the dependence of the `torch` dialect on the low-level builtin types. Now the `torch` dialect is a standalone layer, suitable for targeting from higher-level Python abstractions without any premature lowering to primitive types.
2021-06-17 23:24:31 +08:00
# CHECK-SAME: %[[MAX_ITERATIONS:.*]]: !torch.int) -> !torch.float {
Add `!torch.bool` type. This finishes removing the dependence on the basicpy dialect! Changes: - Add `!torch.bool` type and replace use of `!basicpy.BoolType` in Torch-related code. - Rename BuiltinTensorize to BackendTypeConversion since now it handles bool conversions (and, when we add !torch.int and !torch.float, it will handle those as well), and generalize the related utilities (I also moved them to Torch/Transforms since they aren't really part of Torch/IR). - Add `torch.to_i1` and `torch.from_i1` ops for materializations - [cleanup] Reorganize `torch.constant.*` ops in TorchOps.td - Remove dependency of `torch` dialect on `basicpy` dialect and also `std` dialect. For `std`, we use some call related ops, but the `torch` dialect itself never produces them (we have passes that do though). This is fairly mechanical. Recommended review order: - New stuff in Torch/IR - New BuiltinTypeConversion files. - Mechnical fixups elsewhere.
2021-06-16 07:47:53 +08:00
# CHECK: %[[BOOL_TRUE:.*]] = torch.constant.bool true
Add `torch.constant.int` and `torch.constant.float`. - This removes reliance on basicpy.numeric_constant. - Also, add OpAsmOpInterface to the `torch.constant.none` and `torch.constant.str` ops.
2021-06-16 03:42:51 +08:00
# CHECK: %[[F_INIT:.*]] = torch.constant.float 0.000000e+00
Add support for prim::Loop op. This is a funny one. It combines a `for` and `while` loop in one op. We will need to write some conversions to `scf`.
2021-03-02 07:00:32 +08:00
# CHECK: %[[RESULTS:.*]] = torch.prim.Loop %[[MAX_ITERATIONS]], %[[BOOL_TRUE]], init(%[[F_INIT]]) {
Add a !torch.float type. This removes the dependence of the `torch` dialect on the low-level builtin types. Now the `torch` dialect is a standalone layer, suitable for targeting from higher-level Python abstractions without any premature lowering to primitive types.
2021-06-17 23:24:31 +08:00
# CHECK: ^bb0(%[[IV:.*]]: !torch.int, %[[F_ITER:.*]]: !torch.float):
# CHECK: %[[F_NEXT:.*]] = torch.aten.add.float_int %[[F_ITER]], %[[IV]] : !torch.float, !torch.int -> !torch.float
# CHECK: torch.prim.Loop.condition %[[BOOL_TRUE]], iter(%[[F_NEXT]] : !torch.float)
# CHECK: } : (!torch.int, !torch.bool, !torch.float) -> !torch.float
# CHECK: return %[[RESULTS:.*]] : !torch.float
Add support for prim::Loop op. This is a funny one. It combines a `for` and `while` loop in one op. We will need to write some conversions to `scf`.
2021-03-02 07:00:32 +08:00
@mb.import_function
@torch.jit.script
def prim_Loop_forlike(n: int):
f = 0.0
for i in range(n):
f += i
return f
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_Loop_whilelike(
Add a !torch.float type. This removes the dependence of the `torch` dialect on the low-level builtin types. Now the `torch` dialect is a standalone layer, suitable for targeting from higher-level Python abstractions without any premature lowering to primitive types.
2021-06-17 23:24:31 +08:00
# CHECK-SAME: %[[VAL_0:.*]]: !torch.int) -> !torch.float {
Add `torch.constant.int` and `torch.constant.float`. - This removes reliance on basicpy.numeric_constant. - Also, add OpAsmOpInterface to the `torch.constant.none` and `torch.constant.str` ops.
2021-06-16 03:42:51 +08:00
# CHECK: %[[F_INIT:.*]] = torch.constant.float 3.200000e+00
Add `!torch.int` type. This replaces the ad-hoc use of `i64` throughout the Torch layer, and helps to keep it crystal clear the distinction between `!torch.int` (which is modeling the Python `int` type) and the various types that serve as dtypes of tensors, which are a totally different type universe. Changes: - `!torch.int` type and C bindings. - Change `torch.constant.int` parser to not need the `: i64` at the end. - `m_TorchConstantInt` matcher to aid with matching constants. - BackendTypeConversion changes for `!torch.int` -> `i64` type conversion. - Refactor finalizing patterns in FinalizingBackendTypeConversionPass (they were getting very repetitive). - Mechanical rewriting of `!torch.int` to `i64` in all the tests, and `AnyTorchIntType` to `Torch_IntType` in the `.td` files.
2021-06-17 06:53:15 +08:00
# CHECK: %[[MAX_ITERATIONS:.*]] = torch.constant.int 9223372036854775807
Add a !torch.float type. This removes the dependence of the `torch` dialect on the low-level builtin types. Now the `torch` dialect is a standalone layer, suitable for targeting from higher-level Python abstractions without any premature lowering to primitive types.
2021-06-17 23:24:31 +08:00
# CHECK: %[[COND_INIT:.*]] = torch.aten.lt.float_int %[[F_INIT]], %[[VAL_0]] : !torch.float, !torch.int -> !torch.bool
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: %[[RET:.*]] = torch.prim.Loop %[[MAX_ITERATIONS]], %[[COND_INIT]], init(%[[F_INIT]]) {
Add a !torch.float type. This removes the dependence of the `torch` dialect on the low-level builtin types. Now the `torch` dialect is a standalone layer, suitable for targeting from higher-level Python abstractions without any premature lowering to primitive types.
2021-06-17 23:24:31 +08:00
# CHECK: ^bb0(%[[F_ITER:.*]]: !torch.int, %[[F_ITER:.*]]: !torch.float):
# CHECK: %[[F_NEXT:.*]] = torch.aten.mul.float %[[F_ITER]], %[[F_ITER]] : !torch.float, !torch.float -> !torch.float
# CHECK: %[[COND_ITER:.*]] = torch.aten.lt.float_int %[[F_NEXT]], %[[VAL_0]] : !torch.float, !torch.int -> !torch.bool
# CHECK: torch.prim.Loop.condition %[[COND_ITER]], iter(%[[F_NEXT]] : !torch.float)
# CHECK: } : (!torch.int, !torch.bool, !torch.float) -> !torch.float
# CHECK: return %[[RET:.*]] : !torch.float
Add support for prim::Loop op. This is a funny one. It combines a `for` and `while` loop in one op. We will need to write some conversions to `scf`.
2021-03-02 07:00:32 +08:00
@mb.import_function
@torch.jit.script
def prim_Loop_whilelike(n: int):
f = 3.2
while f < n:
f = f * f
return f
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_Loop_derefine(
Add `!torch.int` type. This replaces the ad-hoc use of `i64` throughout the Torch layer, and helps to keep it crystal clear the distinction between `!torch.int` (which is modeling the Python `int` type) and the various types that serve as dtypes of tensors, which are a totally different type universe. Changes: - `!torch.int` type and C bindings. - Change `torch.constant.int` parser to not need the `: i64` at the end. - `m_TorchConstantInt` matcher to aid with matching constants. - BackendTypeConversion changes for `!torch.int` -> `i64` type conversion. - Refactor finalizing patterns in FinalizingBackendTypeConversionPass (they were getting very repetitive). - Mechanical rewriting of `!torch.int` to `i64` in all the tests, and `AnyTorchIntType` to `Torch_IntType` in the `.td` files.
2021-06-17 06:53:15 +08:00
# CHECK-SAME: %[[ARG:.*]]: !torch.int) -> !torch.optional<!torch.int> {
Add `!torch.bool` type. This finishes removing the dependence on the basicpy dialect! Changes: - Add `!torch.bool` type and replace use of `!basicpy.BoolType` in Torch-related code. - Rename BuiltinTensorize to BackendTypeConversion since now it handles bool conversions (and, when we add !torch.int and !torch.float, it will handle those as well), and generalize the related utilities (I also moved them to Torch/Transforms since they aren't really part of Torch/IR). - Add `torch.to_i1` and `torch.from_i1` ops for materializations - [cleanup] Reorganize `torch.constant.*` ops in TorchOps.td - Remove dependency of `torch` dialect on `basicpy` dialect and also `std` dialect. For `std`, we use some call related ops, but the `torch` dialect itself never produces them (we have passes that do though). This is fairly mechanical. Recommended review order: - New stuff in Torch/IR - New BuiltinTypeConversion files. - Mechnical fixups elsewhere.
2021-06-16 07:47:53 +08:00
# CHECK: %[[TRUE:.*]] = torch.constant.bool true
Introduce native `!torch.none` type. - Add `torch.constant.none` op to construct it (naming is chosen to be analogous to Torch's representation of a prim::Constant with NoneType, rather than using the "singleton" terminology of Basicpy).
2021-06-15 02:36:10 +08:00
# CHECK: %[[NONE:.*]] = torch.constant.none
Add `!torch.int` type. This replaces the ad-hoc use of `i64` throughout the Torch layer, and helps to keep it crystal clear the distinction between `!torch.int` (which is modeling the Python `int` type) and the various types that serve as dtypes of tensors, which are a totally different type universe. Changes: - `!torch.int` type and C bindings. - Change `torch.constant.int` parser to not need the `: i64` at the end. - `m_TorchConstantInt` matcher to aid with matching constants. - BackendTypeConversion changes for `!torch.int` -> `i64` type conversion. - Refactor finalizing patterns in FinalizingBackendTypeConversionPass (they were getting very repetitive). - Mechanical rewriting of `!torch.int` to `i64` in all the tests, and `AnyTorchIntType` to `Torch_IntType` in the `.td` files.
2021-06-17 06:53:15 +08:00
# CHECK: %[[NONE_DEREFINED:.*]] = torch.derefine %[[NONE]] : !torch.none to !torch.optional<!torch.int>
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: %[[RET:.*]] = torch.prim.Loop %[[ARG]], %[[TRUE]], init(%[[NONE_DEREFINED]]) {
Add `!torch.int` type. This replaces the ad-hoc use of `i64` throughout the Torch layer, and helps to keep it crystal clear the distinction between `!torch.int` (which is modeling the Python `int` type) and the various types that serve as dtypes of tensors, which are a totally different type universe. Changes: - `!torch.int` type and C bindings. - Change `torch.constant.int` parser to not need the `: i64` at the end. - `m_TorchConstantInt` matcher to aid with matching constants. - BackendTypeConversion changes for `!torch.int` -> `i64` type conversion. - Refactor finalizing patterns in FinalizingBackendTypeConversionPass (they were getting very repetitive). - Mechanical rewriting of `!torch.int` to `i64` in all the tests, and `AnyTorchIntType` to `Torch_IntType` in the `.td` files.
2021-06-17 06:53:15 +08:00
# CHECK: ^bb0(%[[IV:.*]]: !torch.int, %[[X_ITER:.*]]: !torch.optional<!torch.int>):
# CHECK: %[[X_NEXT:.*]] = torch.derefine %[[ARG]] : !torch.int to !torch.optional<!torch.int>
# CHECK: torch.prim.Loop.condition %[[TRUE]], iter(%[[X_NEXT]] : !torch.optional<!torch.int>)
# CHECK: } : (!torch.int, !torch.bool, !torch.optional<!torch.int>) -> !torch.optional<!torch.int>
# CHECK: return %[[RET:.*]] : !torch.optional<!torch.int>
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
@mb.import_function
@torch.jit.script
def prim_Loop_derefine(n: int):
x: typing.Optional[int] = None
for i in range(n):
x = n
return x
Add support for prim::Loop op. This is a funny one. It combines a `for` and `while` loop in one op. We will need to write some conversions to `scf`.
2021-03-02 07:00:32 +08:00
mb.module.operation.print()
print()