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/decompose-complex-ops.mlir

106 lines
7.6 KiB
MLIR
Raw Normal View History

Add lowering of aten.matmul op. Lowering of `aten.matmul` op is added from torch to linalg dialect. The different cases correspond to https://pytorch.org/docs/stable/generated/torch.matmul.html. TODO: Broadcasting in case of batch-matmul is yet to be taken care of. Signed-off-by: Prashant Kumar <prashant@nod-labs.com>
2021-10-21 13:15:10 +08:00
// RUN: torch-mlir-opt -torch-decompose-complex-ops -split-input-file %s | FileCheck %s
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
// CHECK-LABEL: func.func @matmul_no_decompose
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK: torch.aten.matmul %arg0, %arg1 : !torch.vtensor<[?,?,?,?,?],f32>, !torch.vtensor<[?,?,?],f32> -> !torch.tensor
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
func.func @matmul_no_decompose(%arg0: !torch.vtensor<[?,?,?,?,?],f32>, %arg1: !torch.vtensor<[?,?,?],f32>) -> !torch.tensor {
Add lowering of aten.matmul op. Lowering of `aten.matmul` op is added from torch to linalg dialect. The different cases correspond to https://pytorch.org/docs/stable/generated/torch.matmul.html. TODO: Broadcasting in case of batch-matmul is yet to be taken care of. Signed-off-by: Prashant Kumar <prashant@nod-labs.com>
2021-10-21 13:15:10 +08:00
%0 = torch.aten.matmul %arg0, %arg1 : !torch.vtensor<[?,?,?,?,?],f32>, !torch.vtensor<[?,?,?],f32> -> !torch.tensor
return %0 : !torch.tensor
}
// -----
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
// CHECK-LABEL: func.func @matmul_decompose_2d
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK: torch.aten.mm %arg0, %arg1 : !torch.vtensor<[?,?],f32>, !torch.vtensor<[?,?],f32> -> !torch.tensor
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
func.func @matmul_decompose_2d(%arg0: !torch.vtensor<[?,?],f32>, %arg1: !torch.vtensor<[?,?],f32>) -> !torch.tensor {
Add lowering of aten.matmul op. Lowering of `aten.matmul` op is added from torch to linalg dialect. The different cases correspond to https://pytorch.org/docs/stable/generated/torch.matmul.html. TODO: Broadcasting in case of batch-matmul is yet to be taken care of. Signed-off-by: Prashant Kumar <prashant@nod-labs.com>
2021-10-21 13:15:10 +08:00
%0 = torch.aten.matmul %arg0, %arg1 : !torch.vtensor<[?,?],f32>, !torch.vtensor<[?,?],f32> -> !torch.tensor
return %0 : !torch.tensor
}
// -----
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
// CHECK-LABEL: func.func @matmul_decompose_3d(
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK: torch.aten.bmm %arg0, %arg1 : !torch.vtensor<[?,?,?],f32>, !torch.vtensor<[?,?,?],f32> -> !torch.tensor
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
func.func @matmul_decompose_3d(%arg0: !torch.vtensor<[?,?,?],f32>, %arg1: !torch.vtensor<[?,?,?],f32>) -> !torch.tensor {
Add lowering of aten.matmul op. Lowering of `aten.matmul` op is added from torch to linalg dialect. The different cases correspond to https://pytorch.org/docs/stable/generated/torch.matmul.html. TODO: Broadcasting in case of batch-matmul is yet to be taken care of. Signed-off-by: Prashant Kumar <prashant@nod-labs.com>
2021-10-21 13:15:10 +08:00
%0 = torch.aten.matmul %arg0, %arg1 : !torch.vtensor<[?,?,?],f32>, !torch.vtensor<[?,?,?],f32> -> !torch.tensor
return %0 : !torch.tensor
}
Add type promotion code to refine types. The types have different levels of categories: where complex > floating > integral > boolean (> means left hand side has higher category). The operands have different levels of priorities where: dimensioned tensor > 0-dim tensor > scalar == wrapped 0-dim tensor. This is represented by the `ResultTypeState.dimResult`, `ResultTypeState.zeroResult` and `ResultTypeState..wrappedResult` in the source code. For operands of the same priorities, the result type should be the highest categories with sufficient width to hold all operands. By default, only the highest priority operands participate in the type promotion logic. Lower priority operands participate if they are in a higher category than any higher priority operands. For example, <[],f32> (lower priority) and <[1], si64> tensor would result in <[?],f32> tensor because floating > integeral. Another example <[],f64> (lower priority) and <[1], f32> tensor would result in <[?], f32> tensor because f32 and f64 are the same category. The ScalarType enum definition, type promotion table, ResultTypeState struct definition and some helpers are copied from aten/src/ATen/native/TypeProperties.* Other references: - https://pytorch.org/docs/stable/tensor_attributes.html#type-promotion-doc - https://github.com/pytorch/pytorch/issues/9515 Other minor changes: 1. Fix `visitExpandLikeOp` to consider cases where the given sizes list size is larger than the input rank. 2. Add back the somehow deleted `torch.aten.softmax.int` tests in decompose-complex-ops.mlir.
2021-10-21 03:31:28 +08:00
[TORCH][MLIR] Add E2E support for `aten._unsafe_view` op. This commit adds decomposition of `aten._unsafe_view` op into `aten.view` op. Signed-Off-By: Prateek Gupta<prateek@nod-labs.com>
2022-02-10 16:11:05 +08:00
// -----
Remove unnecessary decompose-complex-ops tests (#1693) This commit removes lit tests from the `decompose-complex-ops` that are essentially testing a macro expansion, in accordance with https://github.com/llvm/torch-mlir/blob/main/docs/architecture.md#dos-and-donts-for-unit-vs-end-to-end-testing .
2022-12-09 00:22:08 +08:00
// CHECK-LABEL: func @torch.aten.adaptive_avg_pool2d$non_unit_output_size(
[MLIR][TORCH] Add decomposition of aten.adaptive_avg_pool2d op This commit adds the decomposition of `aten.adaptive_avg_pool2d` op into `aten.avg_pool2d` op. The current decomposition only supports cases where input size is equal to the output size. Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-05-13 20:06:24 +08:00
// CHECK-SAME: %[[SELF:.*]]: !torch.vtensor<[?,?,?,?],f32>) -> !torch.vtensor<[?,?,?,?],f32> {
Allow running DecomposeComplexOps more than once (#1671) The current implementation of `DecomposeComplexOps` fails if an op expected to be decomposed does not get decomposed in the first iteration of the `createTorchSimplificationPipeline` in `LowerToBackendContractPass`. However, some graphs require multiple iterations of `createTorchSimplificationPipeline` to fully propagate all statically knowable information, such as dtypes and shapes, to the entire graph, sometimes resulting in the need to run `DecomposeComplexOps` more than once. This commit changes `DecomposeComplexOps` to use a greedy algorithm for pattern application and moves the legalization check of ops to the `LowerToBackendContractPass` to allow for the `DecomposeComplexOps` to run more than once.
2022-12-09 01:26:38 +08:00
// CHECK-DAG: %[[CST0:.*]] = torch.constant.int 0
// CHECK-DAG: %[[CST1:.*]] = torch.constant.int 1
// CHECK-DAG: %[[CST2:.*]] = torch.constant.int 2
// CHECK-DAG: %[[CST3:.*]] = torch.constant.int 3
// CHECK-DAG: %[[CST6:.*]] = torch.constant.int 6
// CHECK-DAG: %[[CST7:.*]] = torch.constant.int 7
// CHECK-DAG: %[[FALSE:.*]] = torch.constant.bool false
// CHECK-DAG: %[[TRUE:.*]] = torch.constant.bool true
// CHECK-DAG: %[[NONE:.*]] = torch.constant.none
[MLIR][TORCH] Add decomposition of aten.adaptive_avg_pool2d op This commit adds the decomposition of `aten.adaptive_avg_pool2d` op into `aten.avg_pool2d` op. The current decomposition only supports cases where input size is equal to the output size. Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-05-13 20:06:24 +08:00
// CHECK: %[[DIM2:.*]] = torch.aten.size.int %[[SELF]], %[[CST2]] : !torch.vtensor<[?,?,?,?],f32>, !torch.int -> !torch.int
// CHECK: %[[DIM3:.*]] = torch.aten.size.int %[[SELF]], %[[CST3]] : !torch.vtensor<[?,?,?,?],f32>, !torch.int -> !torch.int
// CHECK: %[[COND1:.*]] = torch.aten.eq.int %[[DIM2]], %[[CST7]] : !torch.int, !torch.int -> !torch.bool
// CHECK: torch.runtime.assert %[[COND1]], "unimplemented: only support cases where input and output size are equal for non-unit output size"
Allow running DecomposeComplexOps more than once (#1671) The current implementation of `DecomposeComplexOps` fails if an op expected to be decomposed does not get decomposed in the first iteration of the `createTorchSimplificationPipeline` in `LowerToBackendContractPass`. However, some graphs require multiple iterations of `createTorchSimplificationPipeline` to fully propagate all statically knowable information, such as dtypes and shapes, to the entire graph, sometimes resulting in the need to run `DecomposeComplexOps` more than once. This commit changes `DecomposeComplexOps` to use a greedy algorithm for pattern application and moves the legalization check of ops to the `LowerToBackendContractPass` to allow for the `DecomposeComplexOps` to run more than once.
2022-12-09 01:26:38 +08:00
// CHECK: %[[T1:.*]] = torch.aten.sub.int %[[DIM2]], %[[CST6]] : !torch.int, !torch.int -> !torch.int
[MLIR][TORCH] Add decomposition of aten.adaptive_avg_pool2d op This commit adds the decomposition of `aten.adaptive_avg_pool2d` op into `aten.avg_pool2d` op. The current decomposition only supports cases where input size is equal to the output size. Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-05-13 20:06:24 +08:00
// CHECK: %[[COND2:.*]] = torch.aten.eq.int %[[DIM3]], %[[CST7]] : !torch.int, !torch.int -> !torch.bool
// CHECK: torch.runtime.assert %[[COND2]], "unimplemented: only support cases where input and output size are equal for non-unit output size"
Allow running DecomposeComplexOps more than once (#1671) The current implementation of `DecomposeComplexOps` fails if an op expected to be decomposed does not get decomposed in the first iteration of the `createTorchSimplificationPipeline` in `LowerToBackendContractPass`. However, some graphs require multiple iterations of `createTorchSimplificationPipeline` to fully propagate all statically knowable information, such as dtypes and shapes, to the entire graph, sometimes resulting in the need to run `DecomposeComplexOps` more than once. This commit changes `DecomposeComplexOps` to use a greedy algorithm for pattern application and moves the legalization check of ops to the `LowerToBackendContractPass` to allow for the `DecomposeComplexOps` to run more than once.
2022-12-09 01:26:38 +08:00
// CHECK: %[[T2:.*]] = torch.aten.sub.int %[[DIM3]], %[[CST6]] : !torch.int, !torch.int -> !torch.int
// CHECK: %[[KERNEL_SIZE:.*]] = torch.prim.ListConstruct %[[T1]], %[[T2]] : (!torch.int, !torch.int) -> !torch.list<int>
// CHECK: %[[STRIDE:.*]] = torch.prim.ListConstruct %[[CST1]], %[[CST1]] : (!torch.int, !torch.int) -> !torch.list<int>
// CHECK: %[[PADDING:.*]] = torch.prim.ListConstruct %[[CST0]], %[[CST0]] : (!torch.int, !torch.int) -> !torch.list<int>
// CHECK: %[[AVG_POOL:.*]] = torch.aten.avg_pool2d %[[SELF]], %[[KERNEL_SIZE]], %[[STRIDE]], %[[PADDING]], %[[FALSE]], %[[TRUE]], %[[NONE]] : !torch.vtensor<[?,?,?,?],f32>, !torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.bool, !torch.bool, !torch.none -> !torch.vtensor<[?,?,?,?],f32>
Remove unnecessary decompose-complex-ops tests (#1693) This commit removes lit tests from the `decompose-complex-ops` that are essentially testing a macro expansion, in accordance with https://github.com/llvm/torch-mlir/blob/main/docs/architecture.md#dos-and-donts-for-unit-vs-end-to-end-testing .
2022-12-09 00:22:08 +08:00
func.func @torch.aten.adaptive_avg_pool2d$non_unit_output_size(%arg0: !torch.vtensor<[?,?,?,?],f32>) -> !torch.vtensor<[?,?,?,?],f32> {
[MLIR][TORCH] Add decomposition of aten.adaptive_avg_pool2d op This commit adds the decomposition of `aten.adaptive_avg_pool2d` op into `aten.avg_pool2d` op. The current decomposition only supports cases where input size is equal to the output size. Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-05-13 20:06:24 +08:00
%int7 = torch.constant.int 7
%output_size = torch.prim.ListConstruct %int7, %int7 : (!torch.int, !torch.int) -> !torch.list<int>
%0 = torch.aten.adaptive_avg_pool2d %arg0, %output_size : !torch.vtensor<[?,?,?,?],f32>, !torch.list<int> -> !torch.vtensor<[?,?,?,?],f32>
return %0 : !torch.vtensor<[?,?,?,?],f32>
}
[MLIR][TORCH] Add E2E support for aten.clamp_[min|max] op This commit decomposes `aten.clamp_min` and `aten.clamp_max` op into `aten.clamp` op. Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-06-03 15:41:13 +08:00
// -----
Add lowering to linalg for softplus and log1p Follows existing conventions for unary operators.
2022-07-18 03:00:29 +08:00
Remove unnecessary decompose-complex-ops tests (#1693) This commit removes lit tests from the `decompose-complex-ops` that are essentially testing a macro expansion, in accordance with https://github.com/llvm/torch-mlir/blob/main/docs/architecture.md#dos-and-donts-for-unit-vs-end-to-end-testing .
2022-12-09 00:22:08 +08:00
// CHECK-LABEL: func.func @torch.aten.adaptive_avg_pool2d$unit_output_size(
// CHECK-SAME: %[[SELF:.*]]: !torch.vtensor<[?,?,?,?],f32>) -> !torch.vtensor<[?,?,?,?],f32> {
Allow running DecomposeComplexOps more than once (#1671) The current implementation of `DecomposeComplexOps` fails if an op expected to be decomposed does not get decomposed in the first iteration of the `createTorchSimplificationPipeline` in `LowerToBackendContractPass`. However, some graphs require multiple iterations of `createTorchSimplificationPipeline` to fully propagate all statically knowable information, such as dtypes and shapes, to the entire graph, sometimes resulting in the need to run `DecomposeComplexOps` more than once. This commit changes `DecomposeComplexOps` to use a greedy algorithm for pattern application and moves the legalization check of ops to the `LowerToBackendContractPass` to allow for the `DecomposeComplexOps` to run more than once.
2022-12-09 01:26:38 +08:00
// CHECK-DAG: %[[CST0:.*]] = torch.constant.int 0
// CHECK-DAG: %[[CST1:.*]] = torch.constant.int 1
// CHECK-DAG: %[[CST2:.*]] = torch.constant.int 2
// CHECK-DAG: %[[CST3:.*]] = torch.constant.int 3
// CHECK-DAG: %[[FALSE:.*]] = torch.constant.bool false
// CHECK-DAG: %[[TRUE:.*]] = torch.constant.bool true
// CHECK-DAG: %[[NONE:.*]] = torch.constant.none
// CHECK: %[[DIM2:.*]] = torch.aten.size.int %[[SELF]], %[[CST2]] : !torch.vtensor<[?,?,?,?],f32>, !torch.int -> !torch.int
// CHECK: %[[DIM3:.*]] = torch.aten.size.int %[[SELF]], %[[CST3]] : !torch.vtensor<[?,?,?,?],f32>, !torch.int -> !torch.int
// CHECK: %[[KERNEL_SIZE:.*]] = torch.prim.ListConstruct %[[DIM2]], %[[DIM3]] : (!torch.int, !torch.int) -> !torch.list<int>
// CHECK: %[[STRIDE:.*]] = torch.prim.ListConstruct %[[CST1]], %[[CST1]] : (!torch.int, !torch.int) -> !torch.list<int>
// CHECK: %[[PADDING:.*]] = torch.prim.ListConstruct %[[CST0]], %[[CST0]] : (!torch.int, !torch.int) -> !torch.list<int>
// CHECK: %[[AVG_POOL:.*]] = torch.aten.avg_pool2d %[[SELF]], %[[KERNEL_SIZE]], %[[STRIDE]], %[[PADDING]], %[[FALSE]], %[[TRUE]], %[[NONE]] : !torch.vtensor<[?,?,?,?],f32>, !torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.bool, !torch.bool, !torch.none -> !torch.vtensor<[?,?,?,?],f32>
Remove unnecessary decompose-complex-ops tests (#1693) This commit removes lit tests from the `decompose-complex-ops` that are essentially testing a macro expansion, in accordance with https://github.com/llvm/torch-mlir/blob/main/docs/architecture.md#dos-and-donts-for-unit-vs-end-to-end-testing .
2022-12-09 00:22:08 +08:00
func.func @torch.aten.adaptive_avg_pool2d$unit_output_size(%arg0: !torch.vtensor<[?,?,?,?],f32>) -> !torch.vtensor<[?,?,?,?],f32> {
[MLIR][TORCH] Add E2E support for aten.mse_loss op This commit adds decomposition for the `aten.mse_loss` op. Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-10-20 19:02:09 +08:00
%int1 = torch.constant.int 1
Remove unnecessary decompose-complex-ops tests (#1693) This commit removes lit tests from the `decompose-complex-ops` that are essentially testing a macro expansion, in accordance with https://github.com/llvm/torch-mlir/blob/main/docs/architecture.md#dos-and-donts-for-unit-vs-end-to-end-testing .
2022-12-09 00:22:08 +08:00
%output_size = torch.prim.ListConstruct %int1, %int1 : (!torch.int, !torch.int) -> !torch.list<int>
%0 = torch.aten.adaptive_avg_pool2d %arg0, %output_size : !torch.vtensor<[?,?,?,?],f32>, !torch.list<int> -> !torch.vtensor<[?,?,?,?],f32>
return %0 : !torch.vtensor<[?,?,?,?],f32>
[MLIR][TORCH] Add E2E support for aten.mse_loss op This commit adds decomposition for the `aten.mse_loss` op. Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-10-20 19:02:09 +08:00
}
[torch-dialect] fix aten.type_as op's folder (#2283) [torch-dialect] fix torch.type_as op's folder by decomposing it to prim.dtype + aten.to_dtype
2023-07-20 09:51:58 +08:00
// -----
// CHECK-LABEL: func.func @torch.aten.type_as$basic(
// CHECK-SAME: %[[ARG_0:.*]]: !torch.tensor, %[[ARG_1:.*]]: !torch.tensor) -> !torch.tensor {
// CHECK: %[[FALSE:.*]] = torch.constant.bool false
// CHECK: %[[NONE:.*]] = torch.constant.none
// CHECK: %[[DTYPE:.*]] = torch.prim.dtype %[[ARG_1]] : !torch.tensor -> !torch.int
// CHECK: %[[VAR:.*]] = torch.aten.to.dtype %[[ARG_0]], %[[DTYPE]], %[[FALSE]], %[[FALSE]], %[[NONE]] : !torch.tensor, !torch.int, !torch.bool, !torch.bool, !torch.none -> !torch.tensor
// CHECK: return %[[VAR]] : !torch.tensor
func.func @torch.aten.type_as$basic(%arg0: !torch.tensor, %arg1: !torch.tensor) -> !torch.tensor {
%0 = torch.aten.type_as %arg0, %arg1 : !torch.tensor, !torch.tensor -> !torch.tensor
return %0 : !torch.tensor
}
// -----
// CHECK-LABEL: func.func @torch.aten.type_as$fold(
// CHECK-SAME: %[[ARG_0:.*]]: !torch.tensor<[?],f16>, %[[ARG_1:.*]]: !torch.tensor<[?,?],f16>) -> !torch.tensor<[?],f16> {
// CHECK: return %[[ARG_0]] : !torch.tensor<[?],f16>
func.func @torch.aten.type_as$fold(%arg0: !torch.tensor<[?], f16>, %arg1: !torch.tensor<[?,?],f16>) -> !torch.tensor<[?],f16> {
%0 = torch.aten.type_as %arg0, %arg1 : !torch.tensor<[?], f16>, !torch.tensor<[?,?],f16> -> !torch.tensor<[?], f16>
return %0 : !torch.tensor<[?], f16>
}