llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
torch-mlir/test/Dialect/Torch/maximize-value-semantics.mlir

247 lines
15 KiB
MLIR
Raw Normal View History

[torch-mlir earthmoving (1/N)] C/C++ code movement. This creates the `external/torch-mlir` directory as an LLVM_EXTERNAL_PROJECTS-compatible project (analogous to `iree-dialects`) and completes movement/rename of all pure MLIR C/C++ compiler code into there. The next step will be to move all the Python code / code that links/includes PyTorch C++ code (which currently lives in `frontends/pytorch`) into a subdirectory here. I call this "earthmoving" because it is mostly mechanical changes and renames. As a quick summary (we can change this down the road easily) - C++ `mlir::NPCOMP::Torch -> mlir::torch::Torch` - CAPI `npcompTorchListTypeGet -> torchMlirTorchListTypeGet` - preprocessor `#ifndef NPCOMP_ -> #ifndef TORCHMLIR_` - CMake `NPCOMPFoo -> TorchMLIRFoo` The goal of this is to create a standalone project creating a center of mass for entry into the MLIR ecosystem from PyTorch, suitable in scope for eventual inclusion/ownership in PyTorch. The idea is that `external/torch-mlir` will some day be pulled out into its own repository, and then npcomp will simply pull it in as a submodule. Layering-wise, what lives in `torch-mlir` lowers code from PyTorch (currently TorchScript, but TorchFX or pytorch/xla-style tracing are possible extensions) down to what we have been calling the "Torch backend contract" which is cleaned up IR (inlining, simplifcation, conversion to value tensors, ...) entirely in the `torch` dialect. This is the branching off point for further lowering, of which npcomp takes one opinion (outside `torch-mlir` of course!), namely the `TorchConversion` dialect/transforms which lower to IR suitable for IREE and other linalg-on-tensors based lower-level compilers. Summary of changes: - move `{include,lib,test}/Dialect/Torch` into `torch-mlir` - move relevant parts of CAPI into `torch-mlir`. - leave a few things related to the `torch-mlir` Python build commented out, which should be resolved in a subsequent change.
2021-09-10 03:24:10 +08:00
// RUN: torch-mlir-opt -split-input-file -allow-unregistered-dialect %s -torch-maximize-value-semantics | FileCheck %s
Make MaximizeValueSemantics a bit smarter. This adds a pattern to MaximizeValueSemantics which does a simple abstract interpretation within a block, which handles simple cases of `torch.overwrite_tensor`, enough to remove all the unnecessary uses of non-value tensors in ResNet right now. Before/after IR: [gist](https://gist.github.com/silvasean/a3e1ef625b19dfc63579f73cd3b543b6) Also, - Split `torch.copy.tensor` into `torch.copy.to_tensor` and `torch.copy.to_vtensor` which convert between value and non-value semantic tensors. This is a much cleaner factorization as they have very separate use cases and properties (e.g. different side effects) - Remove the various canonicalization patterns they had, which were confusing because they resulted in limited forms of maximizing value semantics throughout the pipeline. We should structure our compilation pipeline such that only MaximizeValueSemantics should be maximizing value semantics. - Adjust pass pipeline to only run MaximizeValueSemantics once. - Make OverwriteTensorOp `$value` always be a value tensor and `$overwritten` be a non-value tensor.
2021-06-19 04:47:47 +08:00
// CHECK-LABEL: func @torch.copy.tensor$basic(
// CHECK-SAME: %[[ARG0:.*]]: !torch.vtensor) -> (!torch.vtensor, !torch.vtensor) {
// CHECK: return %[[ARG0]], %[[ARG0]] : !torch.vtensor, !torch.vtensor
func @torch.copy.tensor$basic(%arg0: !torch.vtensor) -> (!torch.vtensor, !torch.vtensor) {
%0 = torch.copy.to_tensor %arg0 : !torch.tensor
%1 = torch.copy.to_vtensor %0 : !torch.vtensor
%2 = torch.copy.to_vtensor %0 : !torch.vtensor
return %1, %2 : !torch.vtensor, !torch.vtensor
}
// CHECK-LABEL: func @one_mutation_in_a_block(
// CHECK-SAME: %[[ARG0:.*]]: !torch.vtensor,
// CHECK-SAME: %[[ARG1:.*]]: !torch.vtensor) -> (!torch.vtensor, !torch.vtensor) {
// CHECK: return %[[ARG0]], %[[ARG1]] : !torch.vtensor, !torch.vtensor
func @one_mutation_in_a_block(%arg0: !torch.vtensor, %arg1: !torch.vtensor) -> (!torch.vtensor, !torch.vtensor) {
%0 = torch.copy.to_tensor %arg0 : !torch.tensor
%equal_to_arg0 = torch.copy.to_vtensor %0 : !torch.vtensor
Add operand type invariant to `torch.overwrite.tensor.contents` (#606) This commit adds the invariant to the op `torch.overwrite.tensor.contents` that both of its operands have the same shape and size. In order to maintain the invariant, special handling of this op is added to the `RefineTypes` pass.
2022-02-23 03:41:46 +08:00
torch.overwrite.tensor.contents %arg1 overwrites %0 : !torch.vtensor, !torch.tensor
Make MaximizeValueSemantics a bit smarter. This adds a pattern to MaximizeValueSemantics which does a simple abstract interpretation within a block, which handles simple cases of `torch.overwrite_tensor`, enough to remove all the unnecessary uses of non-value tensors in ResNet right now. Before/after IR: [gist](https://gist.github.com/silvasean/a3e1ef625b19dfc63579f73cd3b543b6) Also, - Split `torch.copy.tensor` into `torch.copy.to_tensor` and `torch.copy.to_vtensor` which convert between value and non-value semantic tensors. This is a much cleaner factorization as they have very separate use cases and properties (e.g. different side effects) - Remove the various canonicalization patterns they had, which were confusing because they resulted in limited forms of maximizing value semantics throughout the pipeline. We should structure our compilation pipeline such that only MaximizeValueSemantics should be maximizing value semantics. - Adjust pass pipeline to only run MaximizeValueSemantics once. - Make OverwriteTensorOp `$value` always be a value tensor and `$overwritten` be a non-value tensor.
2021-06-19 04:47:47 +08:00
%equal_to_arg1 = torch.copy.to_vtensor %0 : !torch.vtensor
return %equal_to_arg0, %equal_to_arg1 : !torch.vtensor, !torch.vtensor
}
// CHECK-LABEL: func @multiple_mutations_in_a_block(
// CHECK-SAME: %[[ARG0:.*]]: !torch.vtensor, %[[ARG1:.*]]: !torch.vtensor,
// CHECK-SAME: %[[ARG2:.*]]: !torch.vtensor) -> (!torch.vtensor, !torch.vtensor, !torch.vtensor, !torch.vtensor) {
// CHECK: return %[[ARG0]], %[[ARG1]], %[[ARG1]], %[[ARG2]] : !torch.vtensor, !torch.vtensor, !torch.vtensor, !torch.vtensor
func @multiple_mutations_in_a_block(%arg0: !torch.vtensor, %arg1: !torch.vtensor, %arg2: !torch.vtensor) -> (!torch.vtensor, !torch.vtensor, !torch.vtensor, !torch.vtensor) {
// The mutable tensor we are overwriting.
%tensor = torch.copy.to_tensor %arg0 : !torch.tensor
// The original value.
%equal_to_arg0 = torch.copy.to_vtensor %tensor : !torch.vtensor
// Overwrite with %arg1
Add operand type invariant to `torch.overwrite.tensor.contents` (#606) This commit adds the invariant to the op `torch.overwrite.tensor.contents` that both of its operands have the same shape and size. In order to maintain the invariant, special handling of this op is added to the `RefineTypes` pass.
2022-02-23 03:41:46 +08:00
torch.overwrite.tensor.contents %arg1 overwrites %tensor : !torch.vtensor, !torch.tensor
Make MaximizeValueSemantics a bit smarter. This adds a pattern to MaximizeValueSemantics which does a simple abstract interpretation within a block, which handles simple cases of `torch.overwrite_tensor`, enough to remove all the unnecessary uses of non-value tensors in ResNet right now. Before/after IR: [gist](https://gist.github.com/silvasean/a3e1ef625b19dfc63579f73cd3b543b6) Also, - Split `torch.copy.tensor` into `torch.copy.to_tensor` and `torch.copy.to_vtensor` which convert between value and non-value semantic tensors. This is a much cleaner factorization as they have very separate use cases and properties (e.g. different side effects) - Remove the various canonicalization patterns they had, which were confusing because they resulted in limited forms of maximizing value semantics throughout the pipeline. We should structure our compilation pipeline such that only MaximizeValueSemantics should be maximizing value semantics. - Adjust pass pipeline to only run MaximizeValueSemantics once. - Make OverwriteTensorOp `$value` always be a value tensor and `$overwritten` be a non-value tensor.
2021-06-19 04:47:47 +08:00
%equal_to_arg1 = torch.copy.to_vtensor %tensor : !torch.vtensor
%equal_to_arg1_again = torch.copy.to_vtensor %tensor : !torch.vtensor
// Overwrite with %arg2
Add operand type invariant to `torch.overwrite.tensor.contents` (#606) This commit adds the invariant to the op `torch.overwrite.tensor.contents` that both of its operands have the same shape and size. In order to maintain the invariant, special handling of this op is added to the `RefineTypes` pass.
2022-02-23 03:41:46 +08:00
torch.overwrite.tensor.contents %arg2 overwrites %tensor : !torch.vtensor, !torch.tensor
Make MaximizeValueSemantics a bit smarter. This adds a pattern to MaximizeValueSemantics which does a simple abstract interpretation within a block, which handles simple cases of `torch.overwrite_tensor`, enough to remove all the unnecessary uses of non-value tensors in ResNet right now. Before/after IR: [gist](https://gist.github.com/silvasean/a3e1ef625b19dfc63579f73cd3b543b6) Also, - Split `torch.copy.tensor` into `torch.copy.to_tensor` and `torch.copy.to_vtensor` which convert between value and non-value semantic tensors. This is a much cleaner factorization as they have very separate use cases and properties (e.g. different side effects) - Remove the various canonicalization patterns they had, which were confusing because they resulted in limited forms of maximizing value semantics throughout the pipeline. We should structure our compilation pipeline such that only MaximizeValueSemantics should be maximizing value semantics. - Adjust pass pipeline to only run MaximizeValueSemantics once. - Make OverwriteTensorOp `$value` always be a value tensor and `$overwritten` be a non-value tensor.
2021-06-19 04:47:47 +08:00
%equal_to_arg2 = torch.copy.to_vtensor %tensor : !torch.vtensor
return %equal_to_arg0, %equal_to_arg1, %equal_to_arg1_again, %equal_to_arg2 : !torch.vtensor, !torch.vtensor, !torch.vtensor, !torch.vtensor
}
Fix handling of view-like ops in `maximize-value-semantics` (#611) This commit adds handling to the `maximize-value-semantics` pass for the case where a view-like op depends on a tensor that has been overwritten by a value tensor. The approach for removing the dependency is to change the input to the view-like op to be a copy of the value tensor that is being used to overwrite. This commit also removes `AtenFill_ScalarOp` and `AtenBernoulli_FloatOp` from the list of view-like ops, since these ops now have a corresponding op with value semantics into which they get converted in the `reduce-op-variants` pass.
2022-02-19 02:19:07 +08:00
// CHECK-LABEL: func @mutation_followed_by_view_like_ops(
Elide `!torch.` prefix in nested dialect types. This leads to much more succinct types in many cases: ``` !torch.list<!torch.int> !torch.list<int> !torch.tuple<!torch.list<!torch.int>, !torch.list<!torch.int>> !torch.tuple<list<int>, list<int>> !torch.optional<!torch.list<!torch.int>> !torch.optional<list<int>> !torch.list<list<list<tensor>>> !torch.list<!torch.list<!torch.list<!torch.tensor>>> ``` I would like to take this further and allow omitting the `!torch.` prefix in all cases, but that's harder -- for example, we currently use `FuncOp` for functions, and so I don't think we can customize the printing there. It seems like it will be a longer road to getting that level of customization.
2022-03-16 07:22:56 +08:00
// CHECK-SAME: %[[VALUE_T:.*]]: !torch.vtensor, %[[OVERWRITER:.*]]: !torch.vtensor, %[[INT_LIST:.*]]: !torch.list<int>) -> !torch.vtensor {
// CHECK: %[[VIEW:.*]] = torch.aten.view %[[OVERWRITER]], %[[INT_LIST]] : !torch.vtensor, !torch.list<int> -> !torch.vtensor
// CHECK: %[[RESULT:.*]] = torch.aten.permute %[[VIEW]], %[[INT_LIST]] : !torch.vtensor, !torch.list<int> -> !torch.vtensor
Fix handling of view-like ops in `maximize-value-semantics` (#611) This commit adds handling to the `maximize-value-semantics` pass for the case where a view-like op depends on a tensor that has been overwritten by a value tensor. The approach for removing the dependency is to change the input to the view-like op to be a copy of the value tensor that is being used to overwrite. This commit also removes `AtenFill_ScalarOp` and `AtenBernoulli_FloatOp` from the list of view-like ops, since these ops now have a corresponding op with value semantics into which they get converted in the `reduce-op-variants` pass.
2022-02-19 02:19:07 +08:00
// CHECK: return %[[RESULT]] : !torch.vtensor
Elide `!torch.` prefix in nested dialect types. This leads to much more succinct types in many cases: ``` !torch.list<!torch.int> !torch.list<int> !torch.tuple<!torch.list<!torch.int>, !torch.list<!torch.int>> !torch.tuple<list<int>, list<int>> !torch.optional<!torch.list<!torch.int>> !torch.optional<list<int>> !torch.list<list<list<tensor>>> !torch.list<!torch.list<!torch.list<!torch.tensor>>> ``` I would like to take this further and allow omitting the `!torch.` prefix in all cases, but that's harder -- for example, we currently use `FuncOp` for functions, and so I don't think we can customize the printing there. It seems like it will be a longer road to getting that level of customization.
2022-03-16 07:22:56 +08:00
func @mutation_followed_by_view_like_ops(%value_t: !torch.vtensor, %overwriter: !torch.vtensor, %int_list: !torch.list<int>) -> !torch.vtensor {
Fix handling of view-like ops in `maximize-value-semantics` (#611) This commit adds handling to the `maximize-value-semantics` pass for the case where a view-like op depends on a tensor that has been overwritten by a value tensor. The approach for removing the dependency is to change the input to the view-like op to be a copy of the value tensor that is being used to overwrite. This commit also removes `AtenFill_ScalarOp` and `AtenBernoulli_FloatOp` from the list of view-like ops, since these ops now have a corresponding op with value semantics into which they get converted in the `reduce-op-variants` pass.
2022-02-19 02:19:07 +08:00
%t = torch.copy.to_tensor %value_t : !torch.tensor
Add operand type invariant to `torch.overwrite.tensor.contents` (#606) This commit adds the invariant to the op `torch.overwrite.tensor.contents` that both of its operands have the same shape and size. In order to maintain the invariant, special handling of this op is added to the `RefineTypes` pass.
2022-02-23 03:41:46 +08:00
torch.overwrite.tensor.contents %overwriter overwrites %t : !torch.vtensor, !torch.tensor
Elide `!torch.` prefix in nested dialect types. This leads to much more succinct types in many cases: ``` !torch.list<!torch.int> !torch.list<int> !torch.tuple<!torch.list<!torch.int>, !torch.list<!torch.int>> !torch.tuple<list<int>, list<int>> !torch.optional<!torch.list<!torch.int>> !torch.optional<list<int>> !torch.list<list<list<tensor>>> !torch.list<!torch.list<!torch.list<!torch.tensor>>> ``` I would like to take this further and allow omitting the `!torch.` prefix in all cases, but that's harder -- for example, we currently use `FuncOp` for functions, and so I don't think we can customize the printing there. It seems like it will be a longer road to getting that level of customization.
2022-03-16 07:22:56 +08:00
%view = torch.aten.view %t, %int_list : !torch.tensor, !torch.list<int> -> !torch.tensor
%result = torch.aten.permute %view, %int_list : !torch.tensor, !torch.list<int> -> !torch.tensor
Fix handling of view-like ops in `maximize-value-semantics` (#611) This commit adds handling to the `maximize-value-semantics` pass for the case where a view-like op depends on a tensor that has been overwritten by a value tensor. The approach for removing the dependency is to change the input to the view-like op to be a copy of the value tensor that is being used to overwrite. This commit also removes `AtenFill_ScalarOp` and `AtenBernoulli_FloatOp` from the list of view-like ops, since these ops now have a corresponding op with value semantics into which they get converted in the `reduce-op-variants` pass.
2022-02-19 02:19:07 +08:00
%value_result = torch.copy.to_vtensor %result : !torch.vtensor
return %value_result : !torch.vtensor
}
Combine maximize-value-semantics rewrite patterns into one pattern (#642) This commit replaces the two rewrite patterns of maximize-value-semantics with a single pattern that captures the behavior of both as well as other edge cases previously not supported. The new pattern works by first performing alias analysis on a subgraph to see if pattern is applicable, then rewriting all non-value tensors to value tensors in a single go.
2022-03-11 01:36:52 +08:00
// CHECK-LABEL: func @mutation_of_view_like_op_result(
Elide `!torch.` prefix in nested dialect types. This leads to much more succinct types in many cases: ``` !torch.list<!torch.int> !torch.list<int> !torch.tuple<!torch.list<!torch.int>, !torch.list<!torch.int>> !torch.tuple<list<int>, list<int>> !torch.optional<!torch.list<!torch.int>> !torch.optional<list<int>> !torch.list<list<list<tensor>>> !torch.list<!torch.list<!torch.list<!torch.tensor>>> ``` I would like to take this further and allow omitting the `!torch.` prefix in all cases, but that's harder -- for example, we currently use `FuncOp` for functions, and so I don't think we can customize the printing there. It seems like it will be a longer road to getting that level of customization.
2022-03-16 07:22:56 +08:00
// CHECK-SAME: %[[VALUE_T:.*]]: !torch.vtensor, %[[OVERWRITER:.*]]: !torch.vtensor, %[[INT_LIST:.*]]: !torch.list<int>) -> !torch.vtensor {
Combine maximize-value-semantics rewrite patterns into one pattern (#642) This commit replaces the two rewrite patterns of maximize-value-semantics with a single pattern that captures the behavior of both as well as other edge cases previously not supported. The new pattern works by first performing alias analysis on a subgraph to see if pattern is applicable, then rewriting all non-value tensors to value tensors in a single go.
2022-03-11 01:36:52 +08:00
// CHECK: return %[[OVERWRITER]] : !torch.vtensor
Elide `!torch.` prefix in nested dialect types. This leads to much more succinct types in many cases: ``` !torch.list<!torch.int> !torch.list<int> !torch.tuple<!torch.list<!torch.int>, !torch.list<!torch.int>> !torch.tuple<list<int>, list<int>> !torch.optional<!torch.list<!torch.int>> !torch.optional<list<int>> !torch.list<list<list<tensor>>> !torch.list<!torch.list<!torch.list<!torch.tensor>>> ``` I would like to take this further and allow omitting the `!torch.` prefix in all cases, but that's harder -- for example, we currently use `FuncOp` for functions, and so I don't think we can customize the printing there. It seems like it will be a longer road to getting that level of customization.
2022-03-16 07:22:56 +08:00
func @mutation_of_view_like_op_result(%value_t: !torch.vtensor, %overwriter: !torch.vtensor, %int_list: !torch.list<int>) -> !torch.vtensor {
Combine maximize-value-semantics rewrite patterns into one pattern (#642) This commit replaces the two rewrite patterns of maximize-value-semantics with a single pattern that captures the behavior of both as well as other edge cases previously not supported. The new pattern works by first performing alias analysis on a subgraph to see if pattern is applicable, then rewriting all non-value tensors to value tensors in a single go.
2022-03-11 01:36:52 +08:00
%t = torch.copy.to_tensor %value_t : !torch.tensor
Elide `!torch.` prefix in nested dialect types. This leads to much more succinct types in many cases: ``` !torch.list<!torch.int> !torch.list<int> !torch.tuple<!torch.list<!torch.int>, !torch.list<!torch.int>> !torch.tuple<list<int>, list<int>> !torch.optional<!torch.list<!torch.int>> !torch.optional<list<int>> !torch.list<list<list<tensor>>> !torch.list<!torch.list<!torch.list<!torch.tensor>>> ``` I would like to take this further and allow omitting the `!torch.` prefix in all cases, but that's harder -- for example, we currently use `FuncOp` for functions, and so I don't think we can customize the printing there. It seems like it will be a longer road to getting that level of customization.
2022-03-16 07:22:56 +08:00
%view = torch.aten.view %t, %int_list : !torch.tensor, !torch.list<int> -> !torch.tensor
Combine maximize-value-semantics rewrite patterns into one pattern (#642) This commit replaces the two rewrite patterns of maximize-value-semantics with a single pattern that captures the behavior of both as well as other edge cases previously not supported. The new pattern works by first performing alias analysis on a subgraph to see if pattern is applicable, then rewriting all non-value tensors to value tensors in a single go.
2022-03-11 01:36:52 +08:00
torch.overwrite.tensor.contents %overwriter overwrites %view : !torch.vtensor, !torch.tensor
%result = torch.copy.to_vtensor %view : !torch.vtensor
return %result : !torch.vtensor
}
// CHECK-LABEL: func @value_tensor_used_after_copy_was_mutated(
// CHECK-SAME: %[[VALUE_T:.*]]: !torch.vtensor,
// CHECK-SAME: %[[OVERWRITER:.*]]: !torch.vtensor) -> (!torch.vtensor, !torch.vtensor) {
// CHECK: return %[[VALUE_T]], %[[OVERWRITER]] : !torch.vtensor, !torch.vtensor
func @value_tensor_used_after_copy_was_mutated(%value_t: !torch.vtensor, %overwriter: !torch.vtensor) -> (!torch.vtensor, !torch.vtensor) {
%t = torch.copy.to_tensor %value_t : !torch.tensor
torch.overwrite.tensor.contents %overwriter overwrites %t : !torch.vtensor, !torch.tensor
%value_mutated_t = torch.copy.to_vtensor %t : !torch.vtensor
return %value_t, %value_mutated_t : !torch.vtensor, !torch.vtensor
}
Make MaximizeValueSemantics a bit smarter. This adds a pattern to MaximizeValueSemantics which does a simple abstract interpretation within a block, which handles simple cases of `torch.overwrite_tensor`, enough to remove all the unnecessary uses of non-value tensors in ResNet right now. Before/after IR: [gist](https://gist.github.com/silvasean/a3e1ef625b19dfc63579f73cd3b543b6) Also, - Split `torch.copy.tensor` into `torch.copy.to_tensor` and `torch.copy.to_vtensor` which convert between value and non-value semantic tensors. This is a much cleaner factorization as they have very separate use cases and properties (e.g. different side effects) - Remove the various canonicalization patterns they had, which were confusing because they resulted in limited forms of maximizing value semantics throughout the pipeline. We should structure our compilation pipeline such that only MaximizeValueSemantics should be maximizing value semantics. - Adjust pass pipeline to only run MaximizeValueSemantics once. - Make OverwriteTensorOp `$value` always be a value tensor and `$overwritten` be a non-value tensor.
2021-06-19 04:47:47 +08:00
// CHECK-LABEL: func @unmodeled_mutation(
Add operand type invariant to `torch.overwrite.tensor.contents` (#606) This commit adds the invariant to the op `torch.overwrite.tensor.contents` that both of its operands have the same shape and size. In order to maintain the invariant, special handling of this op is added to the `RefineTypes` pass.
2022-02-23 03:41:46 +08:00
// CHECK: torch.overwrite.tensor.contents
Make MaximizeValueSemantics a bit smarter. This adds a pattern to MaximizeValueSemantics which does a simple abstract interpretation within a block, which handles simple cases of `torch.overwrite_tensor`, enough to remove all the unnecessary uses of non-value tensors in ResNet right now. Before/after IR: [gist](https://gist.github.com/silvasean/a3e1ef625b19dfc63579f73cd3b543b6) Also, - Split `torch.copy.tensor` into `torch.copy.to_tensor` and `torch.copy.to_vtensor` which convert between value and non-value semantic tensors. This is a much cleaner factorization as they have very separate use cases and properties (e.g. different side effects) - Remove the various canonicalization patterns they had, which were confusing because they resulted in limited forms of maximizing value semantics throughout the pipeline. We should structure our compilation pipeline such that only MaximizeValueSemantics should be maximizing value semantics. - Adjust pass pipeline to only run MaximizeValueSemantics once. - Make OverwriteTensorOp `$value` always be a value tensor and `$overwritten` be a non-value tensor.
2021-06-19 04:47:47 +08:00
func @unmodeled_mutation(%arg0: !torch.vtensor, %arg1: !torch.vtensor) -> !torch.vtensor {
%0 = torch.copy.to_tensor %arg0 : !torch.tensor
Add operand type invariant to `torch.overwrite.tensor.contents` (#606) This commit adds the invariant to the op `torch.overwrite.tensor.contents` that both of its operands have the same shape and size. In order to maintain the invariant, special handling of this op is added to the `RefineTypes` pass.
2022-02-23 03:41:46 +08:00
torch.overwrite.tensor.contents %arg1 overwrites %0 : !torch.vtensor, !torch.tensor
Make MaximizeValueSemantics a bit smarter. This adds a pattern to MaximizeValueSemantics which does a simple abstract interpretation within a block, which handles simple cases of `torch.overwrite_tensor`, enough to remove all the unnecessary uses of non-value tensors in ResNet right now. Before/after IR: [gist](https://gist.github.com/silvasean/a3e1ef625b19dfc63579f73cd3b543b6) Also, - Split `torch.copy.tensor` into `torch.copy.to_tensor` and `torch.copy.to_vtensor` which convert between value and non-value semantic tensors. This is a much cleaner factorization as they have very separate use cases and properties (e.g. different side effects) - Remove the various canonicalization patterns they had, which were confusing because they resulted in limited forms of maximizing value semantics throughout the pipeline. We should structure our compilation pipeline such that only MaximizeValueSemantics should be maximizing value semantics. - Adjust pass pipeline to only run MaximizeValueSemantics once. - Make OverwriteTensorOp `$value` always be a value tensor and `$overwritten` be a non-value tensor.
2021-06-19 04:47:47 +08:00
"some.op"(%0) : (!torch.tensor) -> ()
%result = torch.copy.to_vtensor %0 : !torch.vtensor
return %result : !torch.vtensor
}
// We don't yet handle nontrivial cases involving control flow.
// CHECK-LABEL: func @unimplemented_control_flow(
// CHECK: torch.copy.to_vtensor
func @unimplemented_control_flow(%arg0: !torch.vtensor, %arg1: !torch.vtensor, %cond: !torch.bool) -> (!torch.vtensor, !torch.vtensor) {
%tensor = torch.copy.to_tensor %arg0 : !torch.tensor
%equal_to_arg0 = torch.copy.to_vtensor %tensor : !torch.vtensor
torch.prim.If %cond -> () {
Add operand type invariant to `torch.overwrite.tensor.contents` (#606) This commit adds the invariant to the op `torch.overwrite.tensor.contents` that both of its operands have the same shape and size. In order to maintain the invariant, special handling of this op is added to the `RefineTypes` pass.
2022-02-23 03:41:46 +08:00
torch.overwrite.tensor.contents %arg1 overwrites %tensor : !torch.vtensor, !torch.tensor
Make MaximizeValueSemantics a bit smarter. This adds a pattern to MaximizeValueSemantics which does a simple abstract interpretation within a block, which handles simple cases of `torch.overwrite_tensor`, enough to remove all the unnecessary uses of non-value tensors in ResNet right now. Before/after IR: [gist](https://gist.github.com/silvasean/a3e1ef625b19dfc63579f73cd3b543b6) Also, - Split `torch.copy.tensor` into `torch.copy.to_tensor` and `torch.copy.to_vtensor` which convert between value and non-value semantic tensors. This is a much cleaner factorization as they have very separate use cases and properties (e.g. different side effects) - Remove the various canonicalization patterns they had, which were confusing because they resulted in limited forms of maximizing value semantics throughout the pipeline. We should structure our compilation pipeline such that only MaximizeValueSemantics should be maximizing value semantics. - Adjust pass pipeline to only run MaximizeValueSemantics once. - Make OverwriteTensorOp `$value` always be a value tensor and `$overwritten` be a non-value tensor.
2021-06-19 04:47:47 +08:00
torch.prim.If.yield
} else {
torch.prim.If.yield
}
%equal_to_arg1 = torch.copy.to_vtensor %tensor : !torch.vtensor
return %equal_to_arg0, %equal_to_arg1 : !torch.vtensor, !torch.vtensor
Introduce `!torch.tensor` / `!torch.vtensor` types. This removes our reliance on the numpy dialect and avoids our off-label use of the builtin tnesor type for modeling unknown dtypes. The `!torch.vtensor` (`ValueTensorType`) type is a value-semantic tensor. The `!torch.tensor` (`NonValueTensorType`) type is a non-value-semantic tensor. The new types look as follows syntactically: ``` // Least-static-information, non-value-semantic tensor. !torch.tensor // Explicit form of least-static-information variant. !torch.tensor<*,unk> // Least-static-information, value-semantic tensor. !torch.vtensor // Explicit form of least-static-information variant. !torch.vtensor<*,unk> // Fixed-set of allowable element types, with first-class support for // Torch's frontend signedness semantics. !torch.tensor<*,si32> // First-class support for unknown dtypes. !torch.tensor<[?,?,?],unk> // Standard MLIR representation of `?` for unknown dimensions. !torch.tensor<[?,2,?,4],unk> // Statically shaped / dtyped example. !torch.vtensor<[1,2,3,4],f32> ``` This required fairly significant changes throughout the compiler, but overall it is a big cleanup. We now have a much clearer layering of "the Torch frontend lowering" vs "lowering to std + linalg + etc.". At the C++ level, there is `ValueTensorType`, `NonValueTensorType`. We also have a helper `BaseTensorType` (kind of like ShapedType) which interoperates with those two. Included changes: - New `torch.tensor(dense<0.0> : tensor<5xf32>) : !torch.tensor` op for creating torch tensor literals in the frontend. - Consistently use signedness for the types (except i1 which I didn't touch -- we need to sort out the situation with !basicpy.BoolType there anyway so will be attending to that soon) - Frontend can annotate whether an argument to the function has value semantics. We currently require this, as our backend contract does not currently allow us to even model the non-value-semantic case. Before, the value-semantic assumption was randomly injected in the middle of the pass pipeline. - Move ArrayToTensor (now called MaximizeValueSemantics) and RefinePublicReturn passes to torch dialect. - The TorchToStd and TorchToLinalg passes are now type conversions from `!torch.vtensor` to `tensor` and use the dialect conversion infra. The overall conversion pipeline is set up following the best practices of the "Type Conversions the Not-So-Hard Way" talk. This required introducing `torch-func-builtin-tensorize` and `torch-finalizing-builtin-tensorize` passes analogous to the upstream bufferization passes with the corresponding names (mostly just copypasta from there). - Misc Torch-level canonicalizations -- we now cleanly layer the lowering to std later in the pipeline, so we are gradually lessening our reliance on random std constant folding before we get to that point. Recommended review order: - New types in TorchTypes.td/TorchTypes.h/TorchDialect.cpp - New ops in TorchOps.td / TorchOps.cpp - Less important / more mechanical stuff - Frontend changes. - Pass changes/additions in `Torch/Transforms` and `Conversion/`
2021-05-21 08:07:18 +08:00
}
Generalize support for elementwise ops. We plumb through e2e a fair number of interesting cases: - unary, binary, ternary elementwise ops - ops like `torch.aten.add.Tensor` that also take a scalar parameter - static size-1 broadcasting We allow the static size-1 broadcasting case, but emit a runtime error in the case of dynamic size-1 broadcasting. This seems like a sweet spot subset of things that can be lowered directly to linalg, while not being overly constraining to users. This is consistent with what IREE is doing for CHLO->Linalg lowering as well ([code](https://github.com/google/iree/blob/50bf7a87e465d2048c527bc27724edde40519b7e/iree/compiler/InputConversion/MHLO/BroadcastingToLinalgPatterns.cpp#L1)). To test the static size-1 case, we added support for the `torch.aten.unsqueeze` op and lowering for it through `linalg.tensor_expand_shape`. This involved a generalization of `MaximizeValueSemantics` able to handle it (the solution there also works for `torch.aten.flatten.using_ints` which we need for ResNet anyway) Also, a few minor additional changes: - Add `VerifyInvariantsBeforeBackendLowering` pass, which catches a large class of errors before we get to backend lowering (now that we are doing dialect conversion, the errors are way nicer if we just emit them up front rather than in the guts of a random pattern). - Minor change to RefBackend to allow `linalg.tensor_expand_shape`. Recommended review order: - e2e tests in elementwise.py - `ConvertElementwiseOp` in TorchToLinalg.cpp + elementwise.mlir test - `ConvertAtenUnsqueezeOp` in TorchToLinalg.cpp + unsqueeze.mlir test - RefineTypes.cpp + tests - MaximizeValueSemantics changes + test - VerifyInvariantsBeforeBackendLowering pass + test
2021-06-26 08:25:09 +08:00
Combine maximize-value-semantics rewrite patterns into one pattern (#642) This commit replaces the two rewrite patterns of maximize-value-semantics with a single pattern that captures the behavior of both as well as other edge cases previously not supported. The new pattern works by first performing alias analysis on a subgraph to see if pattern is applicable, then rewriting all non-value tensors to value tensors in a single go.
2022-03-11 01:36:52 +08:00
// CHECK-LABEL: func @non_value_tensor_returned(
// CHECK-SAME: %[[VALUE_T:.*]]: !torch.vtensor) -> !torch.tensor {
// CHECK: %[[T:.*]] = torch.copy.to_tensor %[[VALUE_T]] : !torch.tensor
// CHECK: return %[[T]] : !torch.tensor
func @non_value_tensor_returned(%value_t: !torch.vtensor) -> !torch.tensor {
%t = torch.copy.to_tensor %value_t : !torch.tensor
return %t : !torch.tensor
}
Introduce new shape library design. See the documentation in `docs/shape_lib.md` and `docs/adding_a_shape_function.md` for an overview of the system. This completely overhauls how we represent shape functions. In particular, RefineTypes does not infer shapes anymore (only dtypes). Shape functions are now written in (TorchScript'able) Python. Recommended review order: 1. Read `docs/shape_lib.md` and `docs/adding_a_shape_function.md`. 1. Code and tests for ReifyShapeCalculations, DropShapeCalculations. 1. Code and tests for SimplifyShapeCalculations. 1. shape_lib_gen.py 1. Code and tests for new RefineTypes pass. 1. Random folders/canonicalizers in TorchOps.cpp and associated test in `canonicalize.mlir`. 1. New ReadOnly trait inferred from the registry. 1. Any miscellaneous remaining stuff. Example `-print-ir-after-all` for ElementwiseUnaryModule: [IR lowering dump](https://gist.github.com/silvasean/e4dc8cbc8d00aac7819602e3cbd8e212). Example `-print-ir-after-all` for ElementwiseBinaryModule: [IR lowering dump](https://gist.github.com/silvasean/daf6860ecced732af3568af6b1899113).
2022-03-10 08:44:22 +08:00
// CHECK-LABEL: func @non_value_tensor_returned$with_overwrite(
// CHECK-SAME: %[[ARG0:.*]]: !torch.vtensor,
// CHECK-SAME: %{{.*}}: !torch.vtensor) -> !torch.tensor {
// CHECK: %[[RESULT:.*]] = torch.copy.to_tensor %[[ARG0]] : !torch.tensor
// CHECK: return %[[RESULT]] : !torch.tensor
func @non_value_tensor_returned$with_overwrite(%arg0: !torch.vtensor, %arg1: !torch.vtensor) -> !torch.tensor {
%2 = torch.copy.to_tensor %arg1 : !torch.tensor
torch.overwrite.tensor.contents %arg0 overwrites %2 : !torch.vtensor, !torch.tensor
return %2 : !torch.tensor
}
// CHECK-LABEL: func @non_value_tensor_returned$return_from_multiple_slices(
// CHECK-SAME: %[[ARG0:.*]]: !torch.vtensor,
// CHECK-SAME: %[[ARG1:.*]]: !torch.vtensor) -> (!torch.tensor, !torch.vtensor, !torch.tensor) {
// CHECK: %[[NON_VALUE_TENSOR0:.*]] = torch.copy.to_tensor %[[ARG0]] : !torch.tensor
// CHECK: %[[NON_VALUE_TENSOR1:.*]] = torch.copy.to_tensor %[[ARG1]] : !torch.tensor
// CHECK: return %[[NON_VALUE_TENSOR0]], %[[ARG0]], %[[NON_VALUE_TENSOR1]] : !torch.tensor, !torch.vtensor, !torch.tensor
func @non_value_tensor_returned$return_from_multiple_slices(%arg0: !torch.vtensor, %arg1: !torch.vtensor) -> (!torch.tensor, !torch.vtensor, !torch.tensor) {
%0 = torch.copy.to_tensor %arg0 : !torch.tensor
// Make a vtensor copy and return that, just to have a load-bearing use.
// This test mainly checks the rewriting of the non-value tensor returns
// though.
%1 = torch.copy.to_vtensor %0 : !torch.vtensor
%2 = torch.copy.to_tensor %arg1 : !torch.tensor
return %0, %1, %2 : !torch.tensor, !torch.vtensor, !torch.tensor
}
Generalize support for elementwise ops. We plumb through e2e a fair number of interesting cases: - unary, binary, ternary elementwise ops - ops like `torch.aten.add.Tensor` that also take a scalar parameter - static size-1 broadcasting We allow the static size-1 broadcasting case, but emit a runtime error in the case of dynamic size-1 broadcasting. This seems like a sweet spot subset of things that can be lowered directly to linalg, while not being overly constraining to users. This is consistent with what IREE is doing for CHLO->Linalg lowering as well ([code](https://github.com/google/iree/blob/50bf7a87e465d2048c527bc27724edde40519b7e/iree/compiler/InputConversion/MHLO/BroadcastingToLinalgPatterns.cpp#L1)). To test the static size-1 case, we added support for the `torch.aten.unsqueeze` op and lowering for it through `linalg.tensor_expand_shape`. This involved a generalization of `MaximizeValueSemantics` able to handle it (the solution there also works for `torch.aten.flatten.using_ints` which we need for ResNet anyway) Also, a few minor additional changes: - Add `VerifyInvariantsBeforeBackendLowering` pass, which catches a large class of errors before we get to backend lowering (now that we are doing dialect conversion, the errors are way nicer if we just emit them up front rather than in the guts of a random pattern). - Minor change to RefBackend to allow `linalg.tensor_expand_shape`. Recommended review order: - e2e tests in elementwise.py - `ConvertElementwiseOp` in TorchToLinalg.cpp + elementwise.mlir test - `ConvertAtenUnsqueezeOp` in TorchToLinalg.cpp + unsqueeze.mlir test - RefineTypes.cpp + tests - MaximizeValueSemantics changes + test - VerifyInvariantsBeforeBackendLowering pass + test
2021-06-26 08:25:09 +08:00
// CHECK-LABEL: func @viewlike$basic_unsqueeze(
// CHECK-SAME: %[[ARG:.*]]: !torch.vtensor) -> !torch.vtensor {
// CHECK: %[[INT0:.*]] = torch.constant.int 0
// CHECK: %[[UNSQUEEZE:.*]] = torch.aten.unsqueeze %[[ARG]], %[[INT0]] : !torch.vtensor, !torch.int -> !torch.vtensor
// CHECK: return %[[UNSQUEEZE]] : !torch.vtensor
func @viewlike$basic_unsqueeze(%arg0: !torch.vtensor) -> !torch.vtensor {
%int0 = torch.constant.int 0
%0 = torch.copy.to_tensor %arg0 : !torch.tensor
%1 = torch.aten.unsqueeze %0, %int0 : !torch.tensor, !torch.int -> !torch.tensor
%2 = torch.copy.to_vtensor %1 : !torch.vtensor
return %2 : !torch.vtensor
}
// CHECK-LABEL: func @viewlike$basic_flatten(
// CHECK-SAME: %[[ARG:.*]]: !torch.vtensor) -> !torch.vtensor {
// CHECK: %[[INT0:.*]] = torch.constant.int 0
// CHECK: %[[INTM1:.*]] = torch.constant.int -1
// CHECK: %[[FLATTEN:.*]] = torch.aten.flatten.using_ints %[[ARG]], %[[INT0]], %[[INTM1]] : !torch.vtensor, !torch.int, !torch.int -> !torch.vtensor
// CHECK: return %[[FLATTEN]] : !torch.vtensor
func @viewlike$basic_flatten(%arg0: !torch.vtensor) -> !torch.vtensor {
%start = torch.constant.int 0
%end = torch.constant.int -1
%0 = torch.copy.to_tensor %arg0 : !torch.tensor
%1 = torch.aten.flatten.using_ints %0, %start, %end : !torch.tensor, !torch.int, !torch.int -> !torch.tensor
%2 = torch.copy.to_vtensor %1 : !torch.vtensor
return %2 : !torch.vtensor
}
// CHECK-LABEL: func @viewlike$transitive(
// CHECK-SAME: %[[ARG:.*]]: !torch.vtensor) -> !torch.vtensor {
// CHECK: %[[INT0:.*]] = torch.constant.int 0
// CHECK: %[[UNSQUEEZE0:.*]] = torch.aten.unsqueeze %[[ARG]], %[[INT0]] : !torch.vtensor, !torch.int -> !torch.vtensor
// CHECK: %[[UNSQUEEZE1:.*]] = torch.aten.unsqueeze %[[UNSQUEEZE0]], %[[INT0]] : !torch.vtensor, !torch.int -> !torch.vtensor
// CHECK: return %[[UNSQUEEZE1]] : !torch.vtensor
func @viewlike$transitive(%arg0: !torch.vtensor) -> !torch.vtensor {
%int0 = torch.constant.int 0
%0 = torch.copy.to_tensor %arg0 : !torch.tensor
%1 = torch.aten.unsqueeze %0, %int0 : !torch.tensor, !torch.int -> !torch.tensor
%2 = torch.aten.unsqueeze %1, %int0 : !torch.tensor, !torch.int -> !torch.tensor
%3 = torch.copy.to_vtensor %2 : !torch.vtensor
return %3 : !torch.vtensor
}
// CHECK-LABEL: func @viewlike$transitive_tree(
// CHECK-SAME: %[[ARG:.*]]: !torch.vtensor) -> (!torch.vtensor, !torch.vtensor) {
// CHECK: %[[INT0:.*]] = torch.constant.int 0
// CHECK: %[[UNSQUEEZE0:.*]] = torch.aten.unsqueeze %[[ARG]], %[[INT0]] : !torch.vtensor, !torch.int -> !torch.vtensor
// CHECK: %[[RET0:.*]] = torch.aten.unsqueeze %[[UNSQUEEZE0]], %[[INT0]] : !torch.vtensor, !torch.int -> !torch.vtensor
// CHECK: %[[RET1:.*]] = torch.aten.unsqueeze %[[UNSQUEEZE0]], %[[INT0]] : !torch.vtensor, !torch.int -> !torch.vtensor
// CHECK: return %[[RET0]], %[[RET1]] : !torch.vtensor, !torch.vtensor
func @viewlike$transitive_tree(%arg0: !torch.vtensor) -> (!torch.vtensor, !torch.vtensor) {
%int0 = torch.constant.int 0
%0 = torch.copy.to_tensor %arg0 : !torch.tensor
// %1 has two users.
%1 = torch.aten.unsqueeze %0, %int0 : !torch.tensor, !torch.int -> !torch.tensor
%2 = torch.aten.unsqueeze %1, %int0 : !torch.tensor, !torch.int -> !torch.tensor
%3 = torch.copy.to_vtensor %2 : !torch.vtensor
%4 = torch.aten.unsqueeze %1, %int0 : !torch.tensor, !torch.int -> !torch.tensor
%5 = torch.copy.to_vtensor %4 : !torch.vtensor
return %3, %5 : !torch.vtensor, !torch.vtensor
}
// CHECK-LABEL: func @viewlike$unmodeled_op(
// CHECK-SAME: %[[ARG:.*]]: !torch.vtensor) -> !torch.vtensor {
// CHECK: %[[UNSQUEEZE:.*]] = torch.aten.unsqueeze {{.*}} : !torch.tensor, !torch.int -> !torch.tensor
// CHECK: "some.op"(%[[UNSQUEEZE]]) : (!torch.tensor) -> ()
func @viewlike$unmodeled_op(%arg0: !torch.vtensor) -> !torch.vtensor {
%int0 = torch.constant.int 0
%0 = torch.copy.to_tensor %arg0 : !torch.tensor
%1 = torch.aten.unsqueeze %0, %int0 : !torch.tensor, !torch.int -> !torch.tensor
"some.op"(%1) : (!torch.tensor) -> ()
%2 = torch.copy.to_vtensor %1 : !torch.vtensor
return %2 : !torch.vtensor
}
Add maximize-value-semantics support for multiple non-value tensor inputs (#659) This commit adds value semantics support for ops such as `aten.view_as` and `aten.expand_as` that take two non-value tensors as input.
2022-03-16 09:13:45 +08:00
// CHECK-LABEL: func @viewlike$two_inputs_one_copy(
// CHECK-SAME: %[[ARG:.*]]: !torch.vtensor) -> !torch.vtensor {
// CHECK: %[[EXPAND_AS:.*]] = torch.aten.expand_as %[[ARG]], %[[ARG]] : !torch.vtensor, !torch.vtensor -> !torch.vtensor
// CHECK: return %[[EXPAND_AS]] : !torch.vtensor
func @viewlike$two_inputs_one_copy(%arg0: !torch.vtensor) -> !torch.vtensor {
%0 = torch.copy.to_tensor %arg0 : !torch.tensor
%1 = torch.aten.expand_as %0, %0 : !torch.tensor, !torch.tensor -> !torch.tensor
%2 = torch.copy.to_vtensor %1 : !torch.vtensor
return %2 : !torch.vtensor
}
// CHECK-LABEL: func @viewlike$two_inputs_two_copies(
// CHECK-SAME: %[[ARG0:.*]]: !torch.vtensor,
// CHECK-SAME: %[[ARG1:.*]]: !torch.vtensor) -> !torch.vtensor {
// CHECK: %[[EXPAND_AS:.*]] = torch.aten.expand_as %[[ARG0]], %[[ARG1]] : !torch.vtensor, !torch.vtensor -> !torch.vtensor
// CHECK: return %[[EXPAND_AS]] : !torch.vtensor
func @viewlike$two_inputs_two_copies(%arg0: !torch.vtensor, %arg1: !torch.vtensor) -> !torch.vtensor {
%0 = torch.copy.to_tensor %arg0 : !torch.tensor
%1 = torch.copy.to_tensor %arg1 : !torch.tensor
%2 = torch.aten.expand_as %0, %1 : !torch.tensor, !torch.tensor -> !torch.tensor
%3 = torch.copy.to_vtensor %2 : !torch.vtensor
return %3 : !torch.vtensor
}