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

682 lines
51 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
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK-LABEL: func @matmul_no_decompose
// CHECK: torch.aten.matmul %arg0, %arg1 : !torch.vtensor<[?,?,?,?,?],f32>, !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
func @matmul_no_decompose(%arg0: !torch.vtensor<[?,?,?,?,?],f32>, %arg1: !torch.vtensor<[?,?,?],f32>) -> !torch.tensor {
%0 = torch.aten.matmul %arg0, %arg1 : !torch.vtensor<[?,?,?,?,?],f32>, !torch.vtensor<[?,?,?],f32> -> !torch.tensor
return %0 : !torch.tensor
}
// -----
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK-LABEL: func @matmul_decompose_2d
// CHECK: torch.aten.mm %arg0, %arg1 : !torch.vtensor<[?,?],f32>, !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
func @matmul_decompose_2d(%arg0: !torch.vtensor<[?,?],f32>, %arg1: !torch.vtensor<[?,?],f32>) -> !torch.tensor {
%0 = torch.aten.matmul %arg0, %arg1 : !torch.vtensor<[?,?],f32>, !torch.vtensor<[?,?],f32> -> !torch.tensor
return %0 : !torch.tensor
}
// -----
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK-LABEL: func @matmul_decompose_3d(
// CHECK: torch.aten.bmm %arg0, %arg1 : !torch.vtensor<[?,?,?],f32>, !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
func @matmul_decompose_3d(%arg0: !torch.vtensor<[?,?,?],f32>, %arg1: !torch.vtensor<[?,?,?],f32>) -> !torch.tensor {
%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
// -----
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
// CHECK-LABEL: func @torch.aten.softmax.int(
// CHECK-SAME: %[[T:.*]]: !torch.tensor<[2,3],f32>,
// CHECK-SAME: %[[DIM:.*]]: !torch.int) -> !torch.tensor<[2,3],f32> {
// CHECK: %[[DTYPE:.*]] = torch.constant.none
Modify softmax decomposition to be more numerically stable. The softmax decomposition is modified according to https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.pytorch to account for numerical stability. Also, modified aten.argmax lowering to handle negative dimension.
2022-02-01 03:56:32 +08:00
// CHECK: %[[KEEP_DIM0:.*]] = torch.constant.bool true
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK: %[[VAL:.*]], %[[IND:.*]] = torch.aten.max.dim %[[T]], %[[DIM]], %[[KEEP_DIM0]] :
Modify softmax decomposition to be more numerically stable. The softmax decomposition is modified according to https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.pytorch to account for numerical stability. Also, modified aten.argmax lowering to handle negative dimension.
2022-02-01 03:56:32 +08:00
// CHECK-SAME: !torch.tensor<[2,3],f32>, !torch.int, !torch.bool -> !torch.tensor<[?,?],f32>, !torch.tensor<[?,?],si64>
// CHECK: %[[FLOAT1:.*]] = torch.constant.float 1.000000e+00
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK: %[[SUB:.*]] = torch.aten.sub.Tensor %[[T]], %[[VAL]], %[[FLOAT1]] : !torch.tensor<[2,3],f32>,
Modify softmax decomposition to be more numerically stable. The softmax decomposition is modified according to https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.pytorch to account for numerical stability. Also, modified aten.argmax lowering to handle negative dimension.
2022-02-01 03:56:32 +08:00
// CHECK-SAME: !torch.tensor<[?,?],f32>, !torch.float -> !torch.tensor<[2,3],f32>
// CHECK: %[[EXP:.*]] = torch.aten.exp %[[SUB]] : !torch.tensor<[2,3],f32> -> !torch.tensor<[2,3],f32>
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: %[[DIM_LIST:.*]] = torch.prim.ListConstruct %[[DIM]] : (!torch.int) -> !torch.list<int>
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
// CHECK: %[[KEEP_DIM:.*]] = torch.constant.bool true
// CHECK: %[[SUM_DTYPE:.*]] = torch.constant.none
// CHECK: %[[SUM:.*]] = torch.aten.sum.dim_IntList %[[EXP]], %[[DIM_LIST]], %[[KEEP_DIM]], %[[SUM_DTYPE]] :
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: !torch.tensor<[2,3],f32>, !torch.list<int>, !torch.bool, !torch.none -> !torch.tensor<[?,?],f32>
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
// CHECK: %[[SOFTMAX:.*]] = torch.aten.div.Tensor %[[EXP]], %[[SUM]] : !torch.tensor<[2,3],f32>, !torch.tensor<[?,?],f32> -> !torch.tensor<[2,3],f32>
// CHECK: %[[RET:.*]] = torch.tensor_static_info_cast %[[SOFTMAX]] : !torch.tensor<[2,3],f32> to !torch.tensor<[2,3],f32>
// CHECK: return %[[RET]] : !torch.tensor<[2,3],f32>
func @torch.aten.softmax.int(%t: !torch.tensor<[2,3],f32>, %dim: !torch.int) -> !torch.tensor<[2,3],f32> {
%dtype = torch.constant.none
%ret = torch.aten.softmax.int %t, %dim, %dtype: !torch.tensor<[2,3],f32>, !torch.int, !torch.none -> !torch.tensor<[2,3],f32>
return %ret : !torch.tensor<[2,3],f32>
}
Modify softmax decomposition to be more numerically stable. The softmax decomposition is modified according to https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.pytorch to account for numerical stability. Also, modified aten.argmax lowering to handle negative dimension.
2022-02-01 03:56:32 +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
// -----
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
// CHECK-LABEL: func @torch.aten.softmax.int$cst_dim(
// CHECK-SAME: %[[T:.*]]: !torch.tensor<[2,3],f32>) -> !torch.tensor<[2,3],f32> {
// CHECK: %[[DTYPE:.*]] = torch.constant.none
// CHECK: %[[DIM:.*]] = torch.constant.int 1
Modify softmax decomposition to be more numerically stable. The softmax decomposition is modified according to https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.pytorch to account for numerical stability. Also, modified aten.argmax lowering to handle negative dimension.
2022-02-01 03:56:32 +08:00
// CHECK: %[[TRU:.*]] = torch.constant.bool true
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK: %[[VAL:.*]], %[[IND:.*]] = torch.aten.max.dim %[[T]], %[[DIM]], %[[TRU]] : !torch.tensor<[2,3],f32>, !torch.int, !torch.bool ->
Modify softmax decomposition to be more numerically stable. The softmax decomposition is modified according to https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.pytorch to account for numerical stability. Also, modified aten.argmax lowering to handle negative dimension.
2022-02-01 03:56:32 +08:00
// CHECK-SAME: !torch.tensor<[2,1],f32>, !torch.tensor<[2,1],si64>
// CHECK: %[[FLOAT1:.*]] = torch.constant.float 1.000000e+00
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK: %[[SUB:.*]] = torch.aten.sub.Tensor %[[T]], %[[VAL]], %[[FLOAT1]] : !torch.tensor<[2,3],f32>,
Modify softmax decomposition to be more numerically stable. The softmax decomposition is modified according to https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.pytorch to account for numerical stability. Also, modified aten.argmax lowering to handle negative dimension.
2022-02-01 03:56:32 +08:00
// CHECK-SAME: !torch.tensor<[2,1],f32>, !torch.float -> !torch.tensor<[2,3],f32>
// CHECK: %[[EXP:.*]] = torch.aten.exp %[[SUB]] : !torch.tensor<[2,3],f32> -> !torch.tensor<[2,3],f32>
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: %[[DIM_LIST:.*]] = torch.prim.ListConstruct %[[DIM]] : (!torch.int) -> !torch.list<int>
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
// CHECK: %[[KEEP_DIM:.*]] = torch.constant.bool true
// CHECK: %[[SUM_DTYPE:.*]] = torch.constant.none
// CHECK: %[[SUM:.*]] = torch.aten.sum.dim_IntList %[[EXP]], %[[DIM_LIST]], %[[KEEP_DIM]], %[[SUM_DTYPE]] :
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 !torch.tensor<[2,3],f32>, !torch.list<int>, !torch.bool, !torch.none -> !torch.tensor<[2,1],f32>
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
// CHECK: %[[SOFTMAX:.*]] = torch.aten.div.Tensor %[[EXP]], %[[SUM]] : !torch.tensor<[2,3],f32>, !torch.tensor<[2,1],f32> -> !torch.tensor<[2,3],f32>
// CHECK: %[[RET:.*]] = torch.tensor_static_info_cast %[[SOFTMAX]] : !torch.tensor<[2,3],f32> to !torch.tensor<[2,3],f32>
// CHECK: return %[[RET]] : !torch.tensor<[2,3],f32>
func @torch.aten.softmax.int$cst_dim(%t: !torch.tensor<[2,3],f32>) -> !torch.tensor<[2,3],f32> {
%none = torch.constant.none
%dim = torch.constant.int 1
%ret = torch.aten.softmax.int %t, %dim, %none : !torch.tensor<[2,3],f32>, !torch.int, !torch.none -> !torch.tensor<[2,3],f32>
return %ret : !torch.tensor<[2,3],f32>
}
[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
// -----
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
// CHECK-LABEL: func @torch.aten.softmax.int$dyn_shape(
// CHECK-SAME: %[[T:.*]]: !torch.tensor<[?,?],f32>) -> !torch.tensor<[?,?],f32> {
// CHECK: %[[DTYPE:.*]] = torch.constant.none
// CHECK: %[[DIM:.*]] = torch.constant.int 1
Modify softmax decomposition to be more numerically stable. The softmax decomposition is modified according to https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.pytorch to account for numerical stability. Also, modified aten.argmax lowering to handle negative dimension.
2022-02-01 03:56:32 +08:00
// CHECK: %[[TRU:.*]] = torch.constant.bool true
// CHECK: %[[VAL:.*]], %[[IND:.*]] = torch.aten.max.dim %[[T]], %[[DIM]], %[[TRU]] : !torch.tensor<[?,?],f32>, !torch.int, !torch.bool ->
// CHECK-SAME: !torch.tensor<[?,1],f32>, !torch.tensor<[?,1],si64>
// CHECK: %[[FLOAT1:.*]] = torch.constant.float 1.000000e+00
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK: %[[SUB:.*]] = torch.aten.sub.Tensor %[[T]], %[[VAL]], %[[FLOAT1]] : !torch.tensor<[?,?],f32>,
Modify softmax decomposition to be more numerically stable. The softmax decomposition is modified according to https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.pytorch to account for numerical stability. Also, modified aten.argmax lowering to handle negative dimension.
2022-02-01 03:56:32 +08:00
// CHECK-SAME: !torch.tensor<[?,1],f32>, !torch.float -> !torch.tensor<[?,?],f32>
// CHECK: %[[EXP:.*]] = torch.aten.exp %[[SUB]] : !torch.tensor<[?,?],f32> -> !torch.tensor<[?,?],f32>
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: %[[DIM_LIST:.*]] = torch.prim.ListConstruct %[[DIM]] : (!torch.int) -> !torch.list<int>
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
// CHECK: %[[KEEP_DIM:.*]] = torch.constant.bool true
// CHECK: %[[SUM_DTYPE:.*]] = torch.constant.none
// CHECK: %[[SUM:.*]] = torch.aten.sum.dim_IntList %[[EXP]], %[[DIM_LIST]], %[[KEEP_DIM]], %[[SUM_DTYPE]] :
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: !torch.tensor<[?,?],f32>, !torch.list<int>, !torch.bool, !torch.none -> !torch.tensor<[?,1],f32>
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
// CHECK: %[[SOFTMAX:.*]] = torch.aten.div.Tensor %[[EXP]], %[[SUM]] : !torch.tensor<[?,?],f32>, !torch.tensor<[?,1],f32> -> !torch.tensor<[?,?],f32>
// CHECK: %[[RET:.*]] = torch.tensor_static_info_cast %[[SOFTMAX]] : !torch.tensor<[?,?],f32> to !torch.tensor<[?,?],f32>
// CHECK: return %[[RET]] : !torch.tensor<[?,?],f32>
func @torch.aten.softmax.int$dyn_shape(%t: !torch.tensor<[?,?],f32>) -> !torch.tensor<[?,?],f32> {
%none = torch.constant.none
%dim = torch.constant.int 1
%ret = torch.aten.softmax.int %t, %dim, %none : !torch.tensor<[?,?],f32>, !torch.int, !torch.none -> !torch.tensor<[?,?],f32>
return %ret : !torch.tensor<[?,?],f32>
}
[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
// -----
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
// CHECK-LABEL: func @torch.aten.softmax.int$unknown_shape(
// CHECK-SAME: %[[T:.*]]: !torch.tensor<*,f32>) -> !torch.tensor<*,f32> {
// CHECK: %[[DTYPE:.*]] = torch.constant.none
// CHECK: %[[DIM:.*]] = torch.constant.int 1
Modify softmax decomposition to be more numerically stable. The softmax decomposition is modified according to https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.pytorch to account for numerical stability. Also, modified aten.argmax lowering to handle negative dimension.
2022-02-01 03:56:32 +08:00
// CHECK: %[[TRU:.*]] = torch.constant.bool true
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK: %[[VAL:.*]], %[[IND:.*]] = torch.aten.max.dim %[[T]], %[[DIM]], %[[TRU]] : !torch.tensor<*,f32>, !torch.int, !torch.bool
Modify softmax decomposition to be more numerically stable. The softmax decomposition is modified according to https://github.com/pytorch/functorch/blob/main/functorch/_src/decompositions.pytorch to account for numerical stability. Also, modified aten.argmax lowering to handle negative dimension.
2022-02-01 03:56:32 +08:00
// CHECK-SAME: -> !torch.tensor<*,f32>, !torch.tensor<*,si64>
// CHECK: %[[FLOAT1:.*]] = torch.constant.float 1.000000e+00
// CHECK: %[[SUB:.*]] = torch.aten.sub.Tensor %[[T]], %[[VAL]], %[[FLOAT1]] : !torch.tensor<*,f32>, !torch.tensor<*,f32>,
// CHECK-SAME: !torch.float -> !torch.tensor<*,f32>
// CHECK: %[[EXP:.*]] = torch.aten.exp %[[SUB]] : !torch.tensor<*,f32> -> !torch.tensor<*,f32>
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: %[[DIM_LIST:.*]] = torch.prim.ListConstruct %[[DIM]] : (!torch.int) -> !torch.list<int>
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
// CHECK: %[[KEEP_DIM:.*]] = torch.constant.bool true
// CHECK: %[[SUM_DTYPE:.*]] = torch.constant.none
// CHECK: %[[SUM:.*]] = torch.aten.sum.dim_IntList %[[EXP]], %[[DIM_LIST]], %[[KEEP_DIM]], %[[SUM_DTYPE]] :
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: !torch.tensor<*,f32>, !torch.list<int>, !torch.bool, !torch.none -> !torch.tensor<*,f32>
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
// CHECK: %[[SOFTMAX:.*]] = torch.aten.div.Tensor %[[EXP]], %[[SUM]] : !torch.tensor<*,f32>, !torch.tensor<*,f32> -> !torch.tensor<*,f32>
// CHECK: %[[RET:.*]] = torch.tensor_static_info_cast %[[SOFTMAX]] : !torch.tensor<*,f32> to !torch.tensor<*,f32>
// CHECK: return %[[RET]] : !torch.tensor<*,f32>
func @torch.aten.softmax.int$unknown_shape(%t: !torch.tensor<*,f32>) -> !torch.tensor<*,f32> {
%none = torch.constant.none
%dim = torch.constant.int 1
%ret = torch.aten.softmax.int %t, %dim, %none : !torch.tensor<*,f32>, !torch.int, !torch.none -> !torch.tensor<*,f32>
return %ret : !torch.tensor<*,f32>
}
Add e2e support for aten._softmax_backward_data. Decompose aten._softmax_backward_data into aten math ops. Also decompose `aten.size` to facilitate decomposing _softmax_backward_data.
2021-11-08 23:56:40 +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
// -----
Add e2e support for aten._softmax_backward_data. Decompose aten._softmax_backward_data into aten math ops. Also decompose `aten.size` to facilitate decomposing _softmax_backward_data.
2021-11-08 23:56:40 +08:00
// CHECK-LABEL: func @torch.aten.size(
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: %[[T:.*]]: !torch.vtensor<[?,3],f32>) -> !torch.list<int> {
Add e2e support for aten._softmax_backward_data. Decompose aten._softmax_backward_data into aten math ops. Also decompose `aten.size` to facilitate decomposing _softmax_backward_data.
2021-11-08 23:56:40 +08:00
// CHECK: %[[CST0:.*]] = torch.constant.int 0
// CHECK: %[[DIM0:.*]] = torch.aten.size.int %[[T]], %[[CST0]] : !torch.vtensor<[?,3],f32>, !torch.int -> !torch.int
// CHECK: %[[CST1:.*]] = torch.constant.int 1
// CHECK: %[[DIM1:.*]] = torch.aten.size.int %[[T]], %[[CST1]] : !torch.vtensor<[?,3],f32>, !torch.int -> !torch.int
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: %[[SIZE:.*]] = torch.prim.ListConstruct %[[DIM0]], %[[DIM1]] : (!torch.int, !torch.int) -> !torch.list<int>
// CHECK: return %[[SIZE]] : !torch.list<int>
func @torch.aten.size(%arg0: !torch.vtensor<[?,3],f32>) -> !torch.list<int> {
%0 = torch.aten.size %arg0 : !torch.vtensor<[?,3],f32> -> !torch.list<int>
return %0 : !torch.list<int>
Add e2e support for aten._softmax_backward_data. Decompose aten._softmax_backward_data into aten math ops. Also decompose `aten.size` to facilitate decomposing _softmax_backward_data.
2021-11-08 23:56:40 +08:00
}
[MLIR][TORCH] Add E2E support for torch.arange op This commit adds lowering of `aten.arange.start_step` op. This commit decomposes `aten.arange` and `aten.arange.start` into `aten.arange.start_step` op. Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2021-12-23 21:22:45 +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
// -----
[MLIR][TORCH] Add E2E support for torch.arange op This commit adds lowering of `aten.arange.start_step` op. This commit decomposes `aten.arange` and `aten.arange.start` into `aten.arange.start_step` op. Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2021-12-23 21:22:45 +08:00
// CHECK-LABEL: func @torch.aten.arange() -> !torch.vtensor<[?],si64> {
// CHECK: %[[CST5:.*]] = torch.constant.int 5
// CHECK: %[[CSTN:.*]] = torch.constant.none
// CHECK: %[[CST0:.*]] = torch.constant.int 0
// CHECK: %[[CST1:.*]] = torch.constant.int 1
// CHECK: %[[RESULT:.*]] = torch.aten.arange.start_step %[[CST0]], %[[CST5]], %[[CST1]], %[[CSTN]], %[[CSTN]], %[[CSTN]], %[[CSTN]] :
// CHECK-SAME: !torch.int, !torch.int, !torch.int, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[?],si64>
// CHECK: return %[[RESULT]] : !torch.vtensor<[?],si64>
func @torch.aten.arange() -> !torch.vtensor<[?],si64> {
%int5 = torch.constant.int 5
%none = torch.constant.none
%0 = torch.aten.arange %int5, %none, %none, %none, %none : !torch.int, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[?],si64>
return %0 : !torch.vtensor<[?],si64>
}
[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
// -----
[MLIR][TORCH] Add E2E support for torch.arange op This commit adds lowering of `aten.arange.start_step` op. This commit decomposes `aten.arange` and `aten.arange.start` into `aten.arange.start_step` op. Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2021-12-23 21:22:45 +08:00
// CHECK-LABEL: func @torch.aten.arange.start() -> !torch.vtensor<[?],si64> {
// CHECK: %[[CST10:.*]] = torch.constant.int 10
// CHECK: %[[CST0:.*]] = torch.constant.int 0
// CHECK: %[[CSTN:.*]] = torch.constant.none
// CHECK: %[[CST1:.*]] = torch.constant.int 1
// CHECK: %[[RESULT:.*]] = torch.aten.arange.start_step %[[CST0]], %[[CST10]], %[[CST1]], %[[CSTN]], %[[CSTN]], %[[CSTN]], %[[CSTN]] :
// CHECK-SAME: !torch.int, !torch.int, !torch.int, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[?],si64>
// CHECK: return %[[RESULT]] : !torch.vtensor<[?],si64>
func @torch.aten.arange.start() -> !torch.vtensor<[?],si64> {
%int10 = torch.constant.int 10
%int0 = torch.constant.int 0
%none = torch.constant.none
%0 = torch.aten.arange.start %int0, %int10, %none, %none, %none, %none : !torch.int, !torch.int, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[?],si64>
return %0 : !torch.vtensor<[?],si64>
}
Add lowering of `aten.max.dim` op. Lowering of `aten.max.dim` op has been added.
2022-01-25 16:53:55 +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 hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK-LABEL: func @torch.aten.argmax(
// CHECK-SAME: %[[INP:.*]]: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[1,?],si64> {
// CHECK: %[[CST0:.*]] = torch.constant.int 0
// CHECK: %[[TRUE:.*]] = torch.constant.bool true
// CHECK: %[[VAL:.*]], %[[IND:.*]] = torch.aten.max.dim %[[INP]], %[[CST0]], %[[TRUE]] :
// CHECK-SAME: !torch.vtensor<[?,?],f32>, !torch.int, !torch.bool -> !torch.vtensor<[1,?],f32>, !torch.vtensor<[1,?],si64>
// CHECK: return %[[IND]] : !torch.vtensor<[1,?],si64>
Add lowering of `aten.max.dim` op. Lowering of `aten.max.dim` op has been added.
2022-01-25 16:53:55 +08:00
func @torch.aten.argmax(%arg0: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[1,?],si64> {
%int0 = torch.constant.int 0
%true = torch.constant.bool true
%0 = torch.aten.argmax %arg0, %int0, %true : !torch.vtensor<[?,?],f32>, !torch.int, !torch.bool -> !torch.vtensor<[1,?],si64>
return %0 : !torch.vtensor<[1,?],si64>
}
[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 hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK-LABEL: func @torch.aten.argmax$reduceall(
// CHECK-SAME: %[[INP:.*]]: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[],si64> {
// CHECK: %[[NONE:.*]] = torch.constant.none
// CHECK: %[[FALSE:.*]] = torch.constant.bool false
// CHECK: %[[CST0:.*]] = torch.constant.int 0
// CHECK: %[[CST1:.*]] = torch.constant.int 1
// CHECK: %[[FLATTEN:.*]] = torch.aten.flatten.using_ints %[[INP]], %[[CST0]], %[[CST1]] :
// CHECK-SAME: !torch.vtensor<[?,?],f32>, !torch.int, !torch.int -> !torch.vtensor<[?],f32>
// CHECK: %[[VAL:.*]], %[[IND:.*]] = torch.aten.max.dim %[[FLATTEN]], %[[CST0]], %[[FALSE]] :
// CHECK-SAME: !torch.vtensor<[?],f32>, !torch.int, !torch.bool -> !torch.vtensor<[],f32>, !torch.vtensor<[],si64>
// CHECK: return %[[IND]] : !torch.vtensor<[],si64>
Add lowering of `aten.max.dim` op. Lowering of `aten.max.dim` op has been added.
2022-01-25 16:53:55 +08:00
func @torch.aten.argmax$reduceall(%arg0: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[],si64> {
%none = torch.constant.none
%false = torch.constant.bool false
%0 = torch.aten.argmax %arg0, %none, %false : !torch.vtensor<[?,?],f32>, !torch.none, !torch.bool -> !torch.vtensor<[],si64>
return %0 : !torch.vtensor<[],si64>
}
[TORCH] Add aten.std e2e support
2022-01-30 01:10:50 +08:00
// -----
// CHECK-LABEL: func @torch.aten.square(
// CHECK-SAME: %[[INPUT:.*]]: !torch.vtensor<[?,?,?],f32>) -> !torch.vtensor<[?,?,?],f32> {
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK: %[[SQUARE:.*]] = torch.aten.mul.Tensor %[[INPUT]], %[[INPUT]] :
// CHECK-SAME: !torch.vtensor<[?,?,?],f32>, !torch.vtensor<[?,?,?],f32> -> !torch.vtensor<[?,?,?],f32>
[TORCH] Add aten.std e2e support
2022-01-30 01:10:50 +08:00
// CHECK: return %[[SQUARE]] : !torch.vtensor<[?,?,?],f32>
func @torch.aten.square(%arg0: !torch.vtensor<[?,?,?],f32>) -> !torch.vtensor<[?,?,?],f32> {
%0 = torch.aten.square %arg0 : !torch.vtensor<[?,?,?],f32> -> !torch.vtensor<[?,?,?],f32>
return %0 : !torch.vtensor<[?,?,?],f32>
}
// -----
// CHECK-LABEL: func @torch.aten.var$unbiased(
// CHECK-SAME: %[[INPUT:.*]]: !torch.vtensor<[?,?,?],f32>) -> !torch.vtensor<[],f32> {
// CHECK: %[[UNBIASED:.*]] = torch.constant.bool true
// CHECK: %[[DTYPE:.*]] = torch.constant.none
// CHECK: %[[SUM:.*]] = torch.aten.sum %[[INPUT]], %[[DTYPE]] : !torch.vtensor<[?,?,?],f32>, !torch.none -> !torch.vtensor<[],f32>
// CHECK: %[[NUM_ELEMENTS:.*]] = torch.aten.numel %[[INPUT]] : !torch.vtensor<[?,?,?],f32> -> !torch.int
// CHECK: %[[MEAN:.*]] = torch.aten.div.Scalar %[[SUM]], %[[NUM_ELEMENTS]] : !torch.vtensor<[],f32>, !torch.int -> !torch.vtensor<[],f32>
// CHECK: %[[ALPHA:.*]] = torch.constant.float 1.000000e+00
// CHECK: %[[SUB_MEAN:.*]] = torch.aten.sub.Tensor %[[INPUT]], %[[MEAN]], %[[ALPHA]] : !torch.vtensor<[?,?,?],f32>, !torch.vtensor<[],f32>, !torch.float -> !torch.vtensor<[?,?,?],f32>
// CHECK: %[[SUB_MEAN_SQUARE:.*]] = torch.aten.mul.Tensor %[[SUB_MEAN]], %[[SUB_MEAN]] : !torch.vtensor<[?,?,?],f32>, !torch.vtensor<[?,?,?],f32> -> !torch.vtensor<[?,?,?],f32>
// CHECK: %[[SUB_MEAN_SQUARE_SUM:.*]] = torch.aten.sum %[[SUB_MEAN_SQUARE]], %[[DTYPE]] : !torch.vtensor<[?,?,?],f32>, !torch.none -> !torch.vtensor<[],f32>
// CHECK: %[[SUB_MEAN_SQUARE_NUM_ELEMENTS:.*]] = torch.aten.numel %[[SUB_MEAN_SQUARE]] : !torch.vtensor<[?,?,?],f32> -> !torch.int
// CHECK: %[[CST1:.*]] = torch.constant.int 1
// CHECK: %[[NUM_ELEMENTS_SUB1:.*]] = torch.aten.sub.int %[[SUB_MEAN_SQUARE_NUM_ELEMENTS]], %[[CST1]] : !torch.int, !torch.int -> !torch.int
// CHECK: %[[UNBIASED_VAR:.*]] = torch.aten.div.Scalar %[[SUB_MEAN_SQUARE_SUM]], %[[NUM_ELEMENTS_SUB1]] : !torch.vtensor<[],f32>, !torch.int -> !torch.vtensor<[],f32>
// CHECK: return %[[UNBIASED_VAR]] : !torch.vtensor<[],f32>
func @torch.aten.var$unbiased(%arg0: !torch.vtensor<[?,?,?],f32>) -> !torch.vtensor<[],f32> {
%true = torch.constant.bool true
%0 = torch.aten.var %arg0, %true: !torch.vtensor<[?,?,?],f32>, !torch.bool -> !torch.vtensor<[],f32>
return %0 : !torch.vtensor<[],f32>
}
// -----
// CHECK-LABEL: func @torch.aten.var$biased(
// CHECK-SAME: %[[INPUT:.*]]: !torch.vtensor<[?,?,?],f32>) -> !torch.vtensor<[],f32> {
// CHECK: %[[UNBIASED:.*]] = torch.constant.bool false
// CHECK: %[[DTYPE:.*]] = torch.constant.none
// CHECK: %[[SUM:.*]] = torch.aten.sum %[[INPUT]], %[[DTYPE]] : !torch.vtensor<[?,?,?],f32>, !torch.none -> !torch.vtensor<[],f32>
// CHECK: %[[NUM_ELEMENTS:.*]] = torch.aten.numel %[[INPUT]] : !torch.vtensor<[?,?,?],f32> -> !torch.int
// CHECK: %[[MEAN:.*]] = torch.aten.div.Scalar %[[SUM]], %[[NUM_ELEMENTS]] : !torch.vtensor<[],f32>, !torch.int -> !torch.vtensor<[],f32>
// CHECK: %[[ALPHA:.*]] = torch.constant.float 1.000000e+00
// CHECK: %[[SUB_MEAN:.*]] = torch.aten.sub.Tensor %[[INPUT]], %[[MEAN]], %[[ALPHA]] : !torch.vtensor<[?,?,?],f32>, !torch.vtensor<[],f32>, !torch.float -> !torch.vtensor<[?,?,?],f32>
// CHECK: %[[SUB_MEAN_SQUARE:.*]] = torch.aten.mul.Tensor %[[SUB_MEAN]], %[[SUB_MEAN]] : !torch.vtensor<[?,?,?],f32>, !torch.vtensor<[?,?,?],f32> -> !torch.vtensor<[?,?,?],f32>
// CHECK: %[[SUB_MEAN_SQUARE_SUM:.*]] = torch.aten.sum %[[SUB_MEAN_SQUARE]], %[[DTYPE]] : !torch.vtensor<[?,?,?],f32>, !torch.none -> !torch.vtensor<[],f32>
// CHECK: %[[SUB_MEAN_SQUARE_NUM_ELEMENTS:.*]] = torch.aten.numel %[[SUB_MEAN_SQUARE]] : !torch.vtensor<[?,?,?],f32> -> !torch.int
// CHECK: %[[BIASED_VAR:.*]] = torch.aten.div.Scalar %[[SUB_MEAN_SQUARE_SUM]], %[[SUB_MEAN_SQUARE_NUM_ELEMENTS]] : !torch.vtensor<[],f32>, !torch.int -> !torch.vtensor<[],f32>
// CHECK: return %[[BIASED_VAR]] : !torch.vtensor<[],f32>
func @torch.aten.var$biased(%arg0: !torch.vtensor<[?,?,?],f32>) -> !torch.vtensor<[],f32> {
%false = torch.constant.bool false
%0 = torch.aten.var %arg0, %false: !torch.vtensor<[?,?,?],f32>, !torch.bool -> !torch.vtensor<[],f32>
return %0 : !torch.vtensor<[],f32>
}
// -----
// CHECK-LABEL: func @torch.aten.std$unbiased(
// CHECK-SAME: %[[INPUT:.*]]: !torch.vtensor<[?,?,?],f32>) -> !torch.vtensor<[],f32> {
// CHECK: %[[UNBIASED:.*]] = torch.constant.bool true
// CHECK: %[[DTYPE:.*]] = torch.constant.none
// CHECK: %[[SUM:.*]] = torch.aten.sum %[[INPUT]], %[[DTYPE]] : !torch.vtensor<[?,?,?],f32>, !torch.none -> !torch.vtensor<[],f32>
// CHECK: %[[NUM_ELEMENTS:.*]] = torch.aten.numel %[[INPUT]] : !torch.vtensor<[?,?,?],f32> -> !torch.int
// CHECK: %[[MEAN:.*]] = torch.aten.div.Scalar %[[SUM]], %[[NUM_ELEMENTS]] : !torch.vtensor<[],f32>, !torch.int -> !torch.vtensor<[],f32>
// CHECK: %[[ALPHA:.*]] = torch.constant.float 1.000000e+00
// CHECK: %[[SUB_MEAN:.*]] = torch.aten.sub.Tensor %[[INPUT]], %[[MEAN]], %[[ALPHA]] : !torch.vtensor<[?,?,?],f32>, !torch.vtensor<[],f32>, !torch.float -> !torch.vtensor<[?,?,?],f32>
// CHECK: %[[SUB_MEAN_SQUARE:.*]] = torch.aten.mul.Tensor %[[SUB_MEAN]], %[[SUB_MEAN]] : !torch.vtensor<[?,?,?],f32>, !torch.vtensor<[?,?,?],f32> -> !torch.vtensor<[?,?,?],f32>
// CHECK: %[[SUB_MEAN_SQUARE_SUM:.*]] = torch.aten.sum %[[SUB_MEAN_SQUARE]], %[[DTYPE]] : !torch.vtensor<[?,?,?],f32>, !torch.none -> !torch.vtensor<[],f32>
// CHECK: %[[SUB_MEAN_SQUARE_NUM_ELEMENTS:.*]] = torch.aten.numel %[[SUB_MEAN_SQUARE]] : !torch.vtensor<[?,?,?],f32> -> !torch.int
// CHECK: %[[CST1:.*]] = torch.constant.int 1
// CHECK: %[[NUM_ELEMENTS_SUB1:.*]] = torch.aten.sub.int %[[SUB_MEAN_SQUARE_NUM_ELEMENTS]], %[[CST1]] : !torch.int, !torch.int -> !torch.int
// CHECK: %[[UNBIASED_VAR:.*]] = torch.aten.div.Scalar %[[SUB_MEAN_SQUARE_SUM]], %[[NUM_ELEMENTS_SUB1]] : !torch.vtensor<[],f32>, !torch.int -> !torch.vtensor<[],f32>
// CHECK: %[[UNBIASED_STD:.*]] = torch.aten.sqrt %[[UNBIASED_VAR]] : !torch.vtensor<[],f32> -> !torch.vtensor<[],f32>
// CHECK: return %[[UNBIASED_STD]] : !torch.vtensor<[],f32>
func @torch.aten.std$unbiased(%arg0: !torch.vtensor<[?,?,?],f32>) -> !torch.vtensor<[],f32> {
%true = torch.constant.bool true
%0 = torch.aten.std %arg0, %true: !torch.vtensor<[?,?,?],f32>, !torch.bool -> !torch.vtensor<[],f32>
return %0 : !torch.vtensor<[],f32>
}
// -----
// CHECK-LABEL: func @torch.aten.std$biased(
// CHECK-SAME: %[[INPUT:.*]]: !torch.vtensor<[?,?,?],f32>) -> !torch.vtensor<[],f32> {
// CHECK: %[[UNBIASED:.*]] = torch.constant.bool false
// CHECK: %[[DTYPE:.*]] = torch.constant.none
// CHECK: %[[SUM:.*]] = torch.aten.sum %[[INPUT]], %[[DTYPE]] : !torch.vtensor<[?,?,?],f32>, !torch.none -> !torch.vtensor<[],f32>
// CHECK: %[[NUM_ELEMENTS:.*]] = torch.aten.numel %[[INPUT]] : !torch.vtensor<[?,?,?],f32> -> !torch.int
// CHECK: %[[MEAN:.*]] = torch.aten.div.Scalar %[[SUM]], %[[NUM_ELEMENTS]] : !torch.vtensor<[],f32>, !torch.int -> !torch.vtensor<[],f32>
// CHECK: %[[ALPHA:.*]] = torch.constant.float 1.000000e+00
// CHECK: %[[SUB_MEAN:.*]] = torch.aten.sub.Tensor %[[INPUT]], %[[MEAN]], %[[ALPHA]] : !torch.vtensor<[?,?,?],f32>, !torch.vtensor<[],f32>, !torch.float -> !torch.vtensor<[?,?,?],f32>
// CHECK: %[[SUB_MEAN_SQUARE:.*]] = torch.aten.mul.Tensor %[[SUB_MEAN]], %[[SUB_MEAN]] : !torch.vtensor<[?,?,?],f32>, !torch.vtensor<[?,?,?],f32> -> !torch.vtensor<[?,?,?],f32>
// CHECK: %[[SUB_MEAN_SQUARE_SUM:.*]] = torch.aten.sum %[[SUB_MEAN_SQUARE]], %[[DTYPE]] : !torch.vtensor<[?,?,?],f32>, !torch.none -> !torch.vtensor<[],f32>
// CHECK: %[[SUB_MEAN_SQUARE_NUM_ELEMENTS:.*]] = torch.aten.numel %[[SUB_MEAN_SQUARE]] : !torch.vtensor<[?,?,?],f32> -> !torch.int
// CHECK: %[[BIASED_VAR:.*]] = torch.aten.div.Scalar %[[SUB_MEAN_SQUARE_SUM]], %[[SUB_MEAN_SQUARE_NUM_ELEMENTS]] : !torch.vtensor<[],f32>, !torch.int -> !torch.vtensor<[],f32>
// CHECK: %[[BIASED_STD:.*]] = torch.aten.sqrt %[[BIASED_VAR]] : !torch.vtensor<[],f32> -> !torch.vtensor<[],f32>
// CHECK: return %[[BIASED_STD]] : !torch.vtensor<[],f32>
func @torch.aten.std$biased(%arg0: !torch.vtensor<[?,?,?],f32>) -> !torch.vtensor<[],f32> {
%false = torch.constant.bool false
%0 = torch.aten.std %arg0, %false: !torch.vtensor<[?,?,?],f32>, !torch.bool -> !torch.vtensor<[],f32>
return %0 : !torch.vtensor<[],f32>
}
[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 hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK-LABEL: func @torch.aten._unsafe_view$static
// CHECK-SAME: (%[[ARG0:.*]]: !torch.vtensor<[1,512,32],f32>)
// CHECK: %[[LIST:.*]] = torch.prim.ListConstruct
// CHECK-NOT: torch.aten._unsafe_view
// CHECK-NEXT: %[[RES:.*]] = torch.aten.view %[[ARG0]], %[[LIST]]
// CHECK-NEXT: return
[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
func @torch.aten._unsafe_view$static(%arg0: !torch.vtensor<[1,512,32],f32>) -> !torch.vtensor<[1,2,256,32],f32> {
%c1 = torch.constant.int 1
%c2 = torch.constant.int 2
%c256 = torch.constant.int 256
%c32 = torch.constant.int 32
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
%0 = torch.prim.ListConstruct %c1, %c2, %c256, %c32 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
%1 = torch.aten._unsafe_view %arg0, %0 : !torch.vtensor<[1,512,32],f32>, !torch.list<int> -> !torch.vtensor<[1,2,256,32],f32>
[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
return %1 : !torch.vtensor<[1,2,256,32],f32>
}
// -----
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK-LABEL: func @torch.aten._unsafe_view$dynamic
// CHECK-SAME: (%[[ARG0:.*]]: !torch.vtensor<[?,?,?],f32>)
// CHECK: %[[LIST:.*]] = torch.prim.ListConstruct
// CHECK-NOT: torch.aten._unsafe_view
// CHECK-NEXT: %[[RES:.*]] = torch.aten.view %[[ARG0]], %[[LIST]]
// CHECK-NEXT: return
[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
func @torch.aten._unsafe_view$dynamic(%arg0: !torch.vtensor<[?,?,?],f32>) -> !torch.vtensor<[512,32],f32> {
%c256 = torch.constant.int 512
%c32 = torch.constant.int 32
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
%0 = torch.prim.ListConstruct %c256, %c32 : (!torch.int, !torch.int) -> !torch.list<int>
%1 = torch.aten._unsafe_view %arg0, %0 : !torch.vtensor<[?,?,?],f32>, !torch.list<int> -> !torch.vtensor<[512,32],f32>
[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
return %1 : !torch.vtensor<[512,32],f32>
}
Add decomposition of _log_softmax op. Decompose _log_softmax into log(softmax(x)).
2022-02-10 15:05:23 +08:00
// -----
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK-LABEL: func @torch.aten._log_softmax(
// CHECK-SAME: %[[INP:.*]]: !torch.vtensor<[?,?,?],f32>) -> !torch.vtensor<[?,?,?],f32> {
// CHECK: %[[INT0:.*]] = torch.constant.int 0
// CHECK: %[[FALSE:.*]] = torch.constant.bool false
// CHECK: %[[TRUE:.*]] = torch.constant.bool true
// CHECK: %[[VAL:.*]], %[[IND:.*]] = torch.aten.max.dim %[[INP]], %[[INT0]], %[[TRUE]] :
// CHECK-SAME: !torch.vtensor<[?,?,?],f32>, !torch.int, !torch.bool -> !torch.vtensor<[1,?,?],f32>, !torch.vtensor<[1,?,?],si64>
// CHECK: %[[FLOAT1:.*]] = torch.constant.float 1.000000e+00
// CHECK: %[[SUB:.*]] = torch.aten.sub.Tensor %[[INP]], %[[VAL]], %[[FLOAT1]] : !torch.vtensor<[?,?,?],f32>, !torch.vtensor<[1,?,?],f32>, !torch.float -> !torch.vtensor<[?,?,?],f32>
// CHECK: %[[EXP:.*]] = torch.aten.exp %[[SUB]] : !torch.vtensor<[?,?,?],f32> -> !torch.vtensor<[?,?,?],f32>
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: %[[PRIM:.*]] = torch.prim.ListConstruct %[[INT0]] : (!torch.int) -> !torch.list<int>
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK: %[[TRU:.*]] = torch.constant.bool true
// CHECK: %[[NONE:.*]] = torch.constant.none
// CHECK: %[[SUM_DIM:.*]] = torch.aten.sum.dim_IntList %[[EXP]], %[[PRIM]], %[[TRU]], %[[NONE]] :
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: !torch.vtensor<[?,?,?],f32>, !torch.list<int>, !torch.bool, !torch.none -> !torch.vtensor<[1,?,?],f32>
Modify aten._log_softmax op decomposition for numerical stability. `aten.log_softmax` is decomposed to be more numerically stable.
2022-02-11 16:39:34 +08:00
// CHECK: %[[LOG:.*]] = torch.aten.log %[[SUM_DIM]] : !torch.vtensor<[1,?,?],f32> -> !torch.vtensor<[1,?,?],f32>
// CHECK: %[[FLOAT_1:.*]] = torch.constant.float 1.000000e+00
Add aten.hardtanh e2e support.
2022-02-09 04:57:23 +08:00
// CHECK: %[[SUB1:.*]] = torch.aten.sub.Tensor %[[SUB]], %[[LOG]], %[[FLOAT_1]] : !torch.vtensor<[?,?,?],f32>,
Modify aten._log_softmax op decomposition for numerical stability. `aten.log_softmax` is decomposed to be more numerically stable.
2022-02-11 16:39:34 +08:00
// CHECK-SAME: !torch.vtensor<[1,?,?],f32>, !torch.float -> !torch.vtensor<[?,?,?],f32>
// CHECK: return %[[SUB1]] : !torch.vtensor<[?,?,?],f32>
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
func @torch.aten._log_softmax(%arg0: !torch.vtensor<[?,?,?],f32> loc(unknown)) -> !torch.vtensor<[?,?,?],f32> {
Add decomposition of _log_softmax op. Decompose _log_softmax into log(softmax(x)).
2022-02-10 15:05:23 +08:00
%int0 = torch.constant.int 0
%false = torch.constant.bool false
%0 = torch.aten._log_softmax %arg0, %int0, %false : !torch.vtensor<[?,?,?],f32>, !torch.int, !torch.bool -> !torch.vtensor<[?,?,?],f32>
return %0 : !torch.vtensor<[?,?,?],f32>
}
Add aten.bernoulli decomposition. aten.bernoulli is decomposed to aten.gtTensor(aten.uniform(x), x).
2022-02-04 19:43:25 +08:00
// -----
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK-LABEL: func @torch.aten.bernoulli
[LINALG] Fix `aten.bernoulli` op lowering - This commit adds E2E support for `aten.rand_like` and `aten.bernoulli_.Tensor` ops. - The `aten.bernoulli(x)` was implemented as: `aten.bernoulli(x) = rand_like(x) < 0.5`, assuming 0.5 as default probability, whereas according to the pytorch documentation: https://pytorch.org/docs/stable/generated/torch.bernoulli.html#torch.bernoulli the input x in `aten.bernoulli(x)` is itself a tensor containing probabilities to be used for drawing the binary random number. - So this commit fixes the `aten.bernoulli(x)` implementation as: `aten.bernoulli(x) = rand_like(x) < x`. - It also fixes the case where the input to `aten.bernoulli_.float` is an integer tensor. In this case the input must be casted to float type before passing it as operand to `aten.rand_like` op. `aten.bernoulli_.float(x, p) = rand_like(float(x)) < p`. Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com>
2022-02-26 00:35:04 +08:00
// CHECK-SAME: (%[[INP:.*]]: !torch.vtensor<[?,?,?],f64>) -> !torch.vtensor {
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK: %[[NONE:.*]] = torch.constant.none
[LINALG] Fix `aten.bernoulli` op lowering - This commit adds E2E support for `aten.rand_like` and `aten.bernoulli_.Tensor` ops. - The `aten.bernoulli(x)` was implemented as: `aten.bernoulli(x) = rand_like(x) < 0.5`, assuming 0.5 as default probability, whereas according to the pytorch documentation: https://pytorch.org/docs/stable/generated/torch.bernoulli.html#torch.bernoulli the input x in `aten.bernoulli(x)` is itself a tensor containing probabilities to be used for drawing the binary random number. - So this commit fixes the `aten.bernoulli(x)` implementation as: `aten.bernoulli(x) = rand_like(x) < x`. - It also fixes the case where the input to `aten.bernoulli_.float` is an integer tensor. In this case the input must be casted to float type before passing it as operand to `aten.rand_like` op. `aten.bernoulli_.float(x, p) = rand_like(float(x)) < p`. Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com>
2022-02-26 00:35:04 +08:00
// CHECK: %[[INT7:.*]] = torch.constant.int 7
// CHECK: %[[FALSE:.*]] = torch.constant.bool false
// CHECK: %[[NONE_0:.*]] = torch.constant.none
// CHECK: %[[CON2FLOAT:.*]] = torch.aten.to.dtype %[[INP]], %[[INT7]], %[[FALSE]], %[[FALSE]], %[[NONE_0]] :
// CHECK-SAME: !torch.vtensor<[?,?,?],f64>, !torch.int, !torch.bool, !torch.bool, !torch.none -> !torch.vtensor<[?,?,?],f64>
// CHECK: %[[NONE_1:.*]] = torch.constant.none
// CHECK: %[[NONE_2:.*]] = torch.constant.none
// CHECK: %[[FLOAT0:.*]] = torch.constant.float 0.000000e+00
// CHECK: %[[FLOAT1:.*]] = torch.constant.float 1.000000e+00
Improve "pseudo" op terminology. The term "pseudo" is very vague and was getting confusing (I felt I had to explain it in every comment referencing it). Instead, rework the "pseudo" ops to instead be named: - MLIR Syntax: `torch.valsem.*` - C++ / ODS: `ValsemVariant*Op` This makes it clear what the concept is, and avoids confusion with other things that might be called "pseudo", since these are very specific and should be 100% consistently named w.r.t. the non-valsem-variant ops that they correspond to.
2022-03-16 07:57:33 +08:00
// CHECK: %[[UNF:.*]] = torch.valsem.aten.uniform %[[CON2FLOAT]], %[[FLOAT0]], %[[FLOAT1]], %[[NONE_2]] : !torch.vtensor<[?,?,?],f64>, !torch.float, !torch.float, !torch.none -> !torch.vtensor<[?,?,?],f64>
[LINALG] Fix `aten.bernoulli` op lowering - This commit adds E2E support for `aten.rand_like` and `aten.bernoulli_.Tensor` ops. - The `aten.bernoulli(x)` was implemented as: `aten.bernoulli(x) = rand_like(x) < 0.5`, assuming 0.5 as default probability, whereas according to the pytorch documentation: https://pytorch.org/docs/stable/generated/torch.bernoulli.html#torch.bernoulli the input x in `aten.bernoulli(x)` is itself a tensor containing probabilities to be used for drawing the binary random number. - So this commit fixes the `aten.bernoulli(x)` implementation as: `aten.bernoulli(x) = rand_like(x) < x`. - It also fixes the case where the input to `aten.bernoulli_.float` is an integer tensor. In this case the input must be casted to float type before passing it as operand to `aten.rand_like` op. `aten.bernoulli_.float(x, p) = rand_like(float(x)) < p`. Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com>
2022-02-26 00:35:04 +08:00
// CHECK: %[[CMP:.*]] = torch.aten.lt.Tensor %[[UNF]], %[[INP]] : !torch.vtensor<[?,?,?],f64>, !torch.vtensor<[?,?,?],f64> -> !torch.vtensor<[?,?,?],i1>
// CHECK: %[[INT7_2:.*]] = torch.constant.int 7
// CHECK: %[[FALSE_2:.*]] = torch.constant.bool false
// CHECK: %[[NONE_3:.*]] = torch.constant.none
// CHECK: %[[TODTYPE:.*]] = torch.aten.to.dtype %[[CMP]], %[[INT7_2]], %[[FALSE_2]], %[[FALSE_2]], %[[NONE_3]] :
// CHECK-SAME: !torch.vtensor<[?,?,?],i1>, !torch.int, !torch.bool, !torch.bool, !torch.none -> !torch.vtensor<[?,?,?],f64>
// CHECK: %[[CAST:.*]] = torch.tensor_static_info_cast %[[TODTYPE]] : !torch.vtensor<[?,?,?],f64> to !torch.vtensor
// CHECK: return %[[CAST]] : !torch.vtensor
func @torch.aten.bernoulli(%arg0: !torch.vtensor<[?,?,?],f64>) -> !torch.vtensor {
%none = torch.constant.none
%0 = torch.aten.bernoulli %arg0, %none : !torch.vtensor<[?,?,?],f64>, !torch.none -> !torch.vtensor<[?,?,?],f64>
%1 = torch.tensor_static_info_cast %0 : !torch.vtensor<[?,?,?],f64> to !torch.vtensor
return %1 : !torch.vtensor
}
// -----
Improve "pseudo" op terminology. The term "pseudo" is very vague and was getting confusing (I felt I had to explain it in every comment referencing it). Instead, rework the "pseudo" ops to instead be named: - MLIR Syntax: `torch.valsem.*` - C++ / ODS: `ValsemVariant*Op` This makes it clear what the concept is, and avoids confusion with other things that might be called "pseudo", since these are very specific and should be 100% consistently named w.r.t. the non-valsem-variant ops that they correspond to.
2022-03-16 07:57:33 +08:00
// CHECK-LABEL: func @torch.valsem.aten.bernoulli.float
[LINALG] Fix `aten.bernoulli` op lowering - This commit adds E2E support for `aten.rand_like` and `aten.bernoulli_.Tensor` ops. - The `aten.bernoulli(x)` was implemented as: `aten.bernoulli(x) = rand_like(x) < 0.5`, assuming 0.5 as default probability, whereas according to the pytorch documentation: https://pytorch.org/docs/stable/generated/torch.bernoulli.html#torch.bernoulli the input x in `aten.bernoulli(x)` is itself a tensor containing probabilities to be used for drawing the binary random number. - So this commit fixes the `aten.bernoulli(x)` implementation as: `aten.bernoulli(x) = rand_like(x) < x`. - It also fixes the case where the input to `aten.bernoulli_.float` is an integer tensor. In this case the input must be casted to float type before passing it as operand to `aten.rand_like` op. `aten.bernoulli_.float(x, p) = rand_like(float(x)) < p`. Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com>
2022-02-26 00:35:04 +08:00
// CHECK-SAME: (%[[INP:.*]]: !torch.vtensor<[?,?,?],f64>) -> !torch.vtensor {
// CHECK: %[[NONE:.*]] = torch.constant.none
// CHECK: %[[PROB:.*]] = torch.constant.float 4.000000e-01
// CHECK: %[[PROB_TENSOR:.*]] = torch.prim.NumToTensor.Scalar %[[PROB]] : !torch.float -> !torch.vtensor<[],f64>
// CHECK: %[[INT7:.*]] = torch.constant.int 7
// CHECK: %[[FALSE:.*]] = torch.constant.bool false
// CHECK: %[[NONE_0:.*]] = torch.constant.none
// CHECK: %[[CON2FLOAT:.*]] = torch.aten.to.dtype %[[INP]], %[[INT7]], %[[FALSE]], %[[FALSE]], %[[NONE_0]] :
// CHECK-SAME: !torch.vtensor<[?,?,?],f64>, !torch.int, !torch.bool, !torch.bool, !torch.none -> !torch.vtensor<[?,?,?],f64>
// CHECK: %[[NONE_1:.*]] = torch.constant.none
// CHECK: %[[NONE_2:.*]] = torch.constant.none
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK: %[[FLOAT0:.*]] = torch.constant.float 0.000000e+00
// CHECK: %[[FLOAT1:.*]] = torch.constant.float 1.000000e+00
Improve "pseudo" op terminology. The term "pseudo" is very vague and was getting confusing (I felt I had to explain it in every comment referencing it). Instead, rework the "pseudo" ops to instead be named: - MLIR Syntax: `torch.valsem.*` - C++ / ODS: `ValsemVariant*Op` This makes it clear what the concept is, and avoids confusion with other things that might be called "pseudo", since these are very specific and should be 100% consistently named w.r.t. the non-valsem-variant ops that they correspond to.
2022-03-16 07:57:33 +08:00
// CHECK: %[[UNF:.*]] = torch.valsem.aten.uniform %[[CON2FLOAT]], %[[FLOAT0]], %[[FLOAT1]], %[[NONE_2]] : !torch.vtensor<[?,?,?],f64>, !torch.float, !torch.float, !torch.none -> !torch.vtensor<[?,?,?],f64>
[LINALG] Fix `aten.bernoulli` op lowering - This commit adds E2E support for `aten.rand_like` and `aten.bernoulli_.Tensor` ops. - The `aten.bernoulli(x)` was implemented as: `aten.bernoulli(x) = rand_like(x) < 0.5`, assuming 0.5 as default probability, whereas according to the pytorch documentation: https://pytorch.org/docs/stable/generated/torch.bernoulli.html#torch.bernoulli the input x in `aten.bernoulli(x)` is itself a tensor containing probabilities to be used for drawing the binary random number. - So this commit fixes the `aten.bernoulli(x)` implementation as: `aten.bernoulli(x) = rand_like(x) < x`. - It also fixes the case where the input to `aten.bernoulli_.float` is an integer tensor. In this case the input must be casted to float type before passing it as operand to `aten.rand_like` op. `aten.bernoulli_.float(x, p) = rand_like(float(x)) < p`. Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com>
2022-02-26 00:35:04 +08:00
// CHECK: %[[CMP:.*]] = torch.aten.lt.Tensor %[[UNF]], %[[PROB_TENSOR]] : !torch.vtensor<[?,?,?],f64>, !torch.vtensor<[],f64> -> !torch.vtensor<[?,?,?],i1>
// CHECK: %[[INT7_2:.*]] = torch.constant.int 7
// CHECK: %[[FALSE_2:.*]] = torch.constant.bool false
// CHECK: %[[NONE_3:.*]] = torch.constant.none
// CHECK: %[[TODTYPE:.*]] = torch.aten.to.dtype %[[CMP]], %[[INT7_2]], %[[FALSE_2]], %[[FALSE_2]], %[[NONE_3]] :
// CHECK-SAME: !torch.vtensor<[?,?,?],i1>, !torch.int, !torch.bool, !torch.bool, !torch.none -> !torch.vtensor<[?,?,?],f64>
// CHECK: %[[CAST:.*]] = torch.tensor_static_info_cast %[[TODTYPE]] : !torch.vtensor<[?,?,?],f64> to !torch.vtensor
// CHECK: return %[[CAST]] : !torch.vtensor
Improve "pseudo" op terminology. The term "pseudo" is very vague and was getting confusing (I felt I had to explain it in every comment referencing it). Instead, rework the "pseudo" ops to instead be named: - MLIR Syntax: `torch.valsem.*` - C++ / ODS: `ValsemVariant*Op` This makes it clear what the concept is, and avoids confusion with other things that might be called "pseudo", since these are very specific and should be 100% consistently named w.r.t. the non-valsem-variant ops that they correspond to.
2022-03-16 07:57:33 +08:00
func @torch.valsem.aten.bernoulli.float(%arg0: !torch.vtensor<[?,?,?],f64>) -> !torch.vtensor {
[LINALG] Fix `aten.bernoulli` op lowering - This commit adds E2E support for `aten.rand_like` and `aten.bernoulli_.Tensor` ops. - The `aten.bernoulli(x)` was implemented as: `aten.bernoulli(x) = rand_like(x) < 0.5`, assuming 0.5 as default probability, whereas according to the pytorch documentation: https://pytorch.org/docs/stable/generated/torch.bernoulli.html#torch.bernoulli the input x in `aten.bernoulli(x)` is itself a tensor containing probabilities to be used for drawing the binary random number. - So this commit fixes the `aten.bernoulli(x)` implementation as: `aten.bernoulli(x) = rand_like(x) < x`. - It also fixes the case where the input to `aten.bernoulli_.float` is an integer tensor. In this case the input must be casted to float type before passing it as operand to `aten.rand_like` op. `aten.bernoulli_.float(x, p) = rand_like(float(x)) < p`. Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com>
2022-02-26 00:35:04 +08:00
%none = torch.constant.none
%prob = torch.constant.float 4.000000e-01
Improve "pseudo" op terminology. The term "pseudo" is very vague and was getting confusing (I felt I had to explain it in every comment referencing it). Instead, rework the "pseudo" ops to instead be named: - MLIR Syntax: `torch.valsem.*` - C++ / ODS: `ValsemVariant*Op` This makes it clear what the concept is, and avoids confusion with other things that might be called "pseudo", since these are very specific and should be 100% consistently named w.r.t. the non-valsem-variant ops that they correspond to.
2022-03-16 07:57:33 +08:00
%0 = torch.valsem.aten.bernoulli.float %arg0, %prob, %none : !torch.vtensor<[?,?,?],f64>, !torch.float, !torch.none -> !torch.vtensor<[?,?,?],f64>
[LINALG] Fix `aten.bernoulli` op lowering - This commit adds E2E support for `aten.rand_like` and `aten.bernoulli_.Tensor` ops. - The `aten.bernoulli(x)` was implemented as: `aten.bernoulli(x) = rand_like(x) < 0.5`, assuming 0.5 as default probability, whereas according to the pytorch documentation: https://pytorch.org/docs/stable/generated/torch.bernoulli.html#torch.bernoulli the input x in `aten.bernoulli(x)` is itself a tensor containing probabilities to be used for drawing the binary random number. - So this commit fixes the `aten.bernoulli(x)` implementation as: `aten.bernoulli(x) = rand_like(x) < x`. - It also fixes the case where the input to `aten.bernoulli_.float` is an integer tensor. In this case the input must be casted to float type before passing it as operand to `aten.rand_like` op. `aten.bernoulli_.float(x, p) = rand_like(float(x)) < p`. Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com>
2022-02-26 00:35:04 +08:00
%1 = torch.tensor_static_info_cast %0 : !torch.vtensor<[?,?,?],f64> to !torch.vtensor
return %1 : !torch.vtensor
}
// -----
Improve "pseudo" op terminology. The term "pseudo" is very vague and was getting confusing (I felt I had to explain it in every comment referencing it). Instead, rework the "pseudo" ops to instead be named: - MLIR Syntax: `torch.valsem.*` - C++ / ODS: `ValsemVariant*Op` This makes it clear what the concept is, and avoids confusion with other things that might be called "pseudo", since these are very specific and should be 100% consistently named w.r.t. the non-valsem-variant ops that they correspond to.
2022-03-16 07:57:33 +08:00
// CHECK-LABEL: func @torch.valsem.aten.bernoulli.Tensor(
[LINALG] Fix `aten.bernoulli` op lowering - This commit adds E2E support for `aten.rand_like` and `aten.bernoulli_.Tensor` ops. - The `aten.bernoulli(x)` was implemented as: `aten.bernoulli(x) = rand_like(x) < 0.5`, assuming 0.5 as default probability, whereas according to the pytorch documentation: https://pytorch.org/docs/stable/generated/torch.bernoulli.html#torch.bernoulli the input x in `aten.bernoulli(x)` is itself a tensor containing probabilities to be used for drawing the binary random number. - So this commit fixes the `aten.bernoulli(x)` implementation as: `aten.bernoulli(x) = rand_like(x) < x`. - It also fixes the case where the input to `aten.bernoulli_.float` is an integer tensor. In this case the input must be casted to float type before passing it as operand to `aten.rand_like` op. `aten.bernoulli_.float(x, p) = rand_like(float(x)) < p`. Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com>
2022-02-26 00:35:04 +08:00
// CHECK-SAME: %[[INP:.*]]: !torch.vtensor<[?,?,?],f64>,
// CHECK-SAME: %[[PROB:.*]]: !torch.vtensor<[?,?,?],f64>) -> !torch.vtensor {
// CHECK: %[[NONE:.*]] = torch.constant.none
// CHECK: %[[INT7:.*]] = torch.constant.int 7
Add aten.hardtanh e2e support.
2022-02-09 04:57:23 +08:00
// CHECK: %[[FALSE:.*]] = torch.constant.bool false
[LINALG] Fix `aten.bernoulli` op lowering - This commit adds E2E support for `aten.rand_like` and `aten.bernoulli_.Tensor` ops. - The `aten.bernoulli(x)` was implemented as: `aten.bernoulli(x) = rand_like(x) < 0.5`, assuming 0.5 as default probability, whereas according to the pytorch documentation: https://pytorch.org/docs/stable/generated/torch.bernoulli.html#torch.bernoulli the input x in `aten.bernoulli(x)` is itself a tensor containing probabilities to be used for drawing the binary random number. - So this commit fixes the `aten.bernoulli(x)` implementation as: `aten.bernoulli(x) = rand_like(x) < x`. - It also fixes the case where the input to `aten.bernoulli_.float` is an integer tensor. In this case the input must be casted to float type before passing it as operand to `aten.rand_like` op. `aten.bernoulli_.float(x, p) = rand_like(float(x)) < p`. Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com>
2022-02-26 00:35:04 +08:00
// CHECK: %[[NONE_0:.*]] = torch.constant.none
// CHECK: %[[CON2FLOAT:.*]] = torch.aten.to.dtype %[[INP]], %[[INT7]], %[[FALSE]], %[[FALSE]], %[[NONE_0]] :
// CHECK-SAME: !torch.vtensor<[?,?,?],f64>, !torch.int, !torch.bool, !torch.bool, !torch.none -> !torch.vtensor<[?,?,?],f64>
// CHECK: %[[NONE_1:.*]] = torch.constant.none
// CHECK: %[[NONE_2:.*]] = torch.constant.none
// CHECK: %[[FLOAT0:.*]] = torch.constant.float 0.000000e+00
// CHECK: %[[FLOAT1:.*]] = torch.constant.float 1.000000e+00
Improve "pseudo" op terminology. The term "pseudo" is very vague and was getting confusing (I felt I had to explain it in every comment referencing it). Instead, rework the "pseudo" ops to instead be named: - MLIR Syntax: `torch.valsem.*` - C++ / ODS: `ValsemVariant*Op` This makes it clear what the concept is, and avoids confusion with other things that might be called "pseudo", since these are very specific and should be 100% consistently named w.r.t. the non-valsem-variant ops that they correspond to.
2022-03-16 07:57:33 +08:00
// CHECK: %[[UNF:.*]] = torch.valsem.aten.uniform %[[CON2FLOAT]], %[[FLOAT0]], %[[FLOAT1]], %[[NONE_2]] : !torch.vtensor<[?,?,?],f64>, !torch.float, !torch.float, !torch.none -> !torch.vtensor<[?,?,?],f64>
[LINALG] Fix `aten.bernoulli` op lowering - This commit adds E2E support for `aten.rand_like` and `aten.bernoulli_.Tensor` ops. - The `aten.bernoulli(x)` was implemented as: `aten.bernoulli(x) = rand_like(x) < 0.5`, assuming 0.5 as default probability, whereas according to the pytorch documentation: https://pytorch.org/docs/stable/generated/torch.bernoulli.html#torch.bernoulli the input x in `aten.bernoulli(x)` is itself a tensor containing probabilities to be used for drawing the binary random number. - So this commit fixes the `aten.bernoulli(x)` implementation as: `aten.bernoulli(x) = rand_like(x) < x`. - It also fixes the case where the input to `aten.bernoulli_.float` is an integer tensor. In this case the input must be casted to float type before passing it as operand to `aten.rand_like` op. `aten.bernoulli_.float(x, p) = rand_like(float(x)) < p`. Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com>
2022-02-26 00:35:04 +08:00
// CHECK: %[[CMP:.*]] = torch.aten.lt.Tensor %[[UNF]], %[[PROB]] : !torch.vtensor<[?,?,?],f64>, !torch.vtensor<[?,?,?],f64> -> !torch.vtensor<[?,?,?],i1>
// CHECK: %[[INT7_2:.*]] = torch.constant.int 7
// CHECK: %[[FALSE_2:.*]] = torch.constant.bool false
// CHECK: %[[NONE_3:.*]] = torch.constant.none
// CHECK: %[[TODTYPE:.*]] = torch.aten.to.dtype %[[CMP]], %[[INT7_2]], %[[FALSE_2]], %[[FALSE_2]], %[[NONE_3]] :
// CHECK-SAME: !torch.vtensor<[?,?,?],i1>, !torch.int, !torch.bool, !torch.bool, !torch.none -> !torch.vtensor<[?,?,?],f64>
// CHECK: %[[CAST:.*]] = torch.tensor_static_info_cast %[[TODTYPE]] : !torch.vtensor<[?,?,?],f64> to !torch.vtensor
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// CHECK: return %[[CAST]] : !torch.vtensor
Improve "pseudo" op terminology. The term "pseudo" is very vague and was getting confusing (I felt I had to explain it in every comment referencing it). Instead, rework the "pseudo" ops to instead be named: - MLIR Syntax: `torch.valsem.*` - C++ / ODS: `ValsemVariant*Op` This makes it clear what the concept is, and avoids confusion with other things that might be called "pseudo", since these are very specific and should be 100% consistently named w.r.t. the non-valsem-variant ops that they correspond to.
2022-03-16 07:57:33 +08:00
func @torch.valsem.aten.bernoulli.Tensor(%arg0: !torch.vtensor<[?,?,?],f64>, %arg1: !torch.vtensor<[?,?,?],f64>) -> !torch.vtensor {
Add aten.bernoulli decomposition. aten.bernoulli is decomposed to aten.gtTensor(aten.uniform(x), x).
2022-02-04 19:43:25 +08:00
%none = torch.constant.none
Improve "pseudo" op terminology. The term "pseudo" is very vague and was getting confusing (I felt I had to explain it in every comment referencing it). Instead, rework the "pseudo" ops to instead be named: - MLIR Syntax: `torch.valsem.*` - C++ / ODS: `ValsemVariant*Op` This makes it clear what the concept is, and avoids confusion with other things that might be called "pseudo", since these are very specific and should be 100% consistently named w.r.t. the non-valsem-variant ops that they correspond to.
2022-03-16 07:57:33 +08:00
%0 = torch.valsem.aten.bernoulli.Tensor %arg0, %arg1, %none : !torch.vtensor<[?,?,?],f64>, !torch.vtensor<[?,?,?],f64>, !torch.none -> !torch.vtensor<[?,?,?],f64>
[LINALG] Fix `aten.bernoulli` op lowering - This commit adds E2E support for `aten.rand_like` and `aten.bernoulli_.Tensor` ops. - The `aten.bernoulli(x)` was implemented as: `aten.bernoulli(x) = rand_like(x) < 0.5`, assuming 0.5 as default probability, whereas according to the pytorch documentation: https://pytorch.org/docs/stable/generated/torch.bernoulli.html#torch.bernoulli the input x in `aten.bernoulli(x)` is itself a tensor containing probabilities to be used for drawing the binary random number. - So this commit fixes the `aten.bernoulli(x)` implementation as: `aten.bernoulli(x) = rand_like(x) < x`. - It also fixes the case where the input to `aten.bernoulli_.float` is an integer tensor. In this case the input must be casted to float type before passing it as operand to `aten.rand_like` op. `aten.bernoulli_.float(x, p) = rand_like(float(x)) < p`. Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com>
2022-02-26 00:35:04 +08:00
%1 = torch.tensor_static_info_cast %0 : !torch.vtensor<[?,?,?],f64> to !torch.vtensor
Add aten.bernoulli decomposition. aten.bernoulli is decomposed to aten.gtTensor(aten.uniform(x), x).
2022-02-04 19:43:25 +08:00
return %1 : !torch.vtensor
}
Remove hacky aten.select.int lowering code
2022-02-12 03:34:05 +08:00
// -----
// CHECK-LABEL: func @torch.aten.select.int(
// CHECK-SAME: %[[T:.*]]: !torch.vtensor<[?,?],si64>) -> !torch.vtensor<[?],si64> {
// CHECK: %[[CST0:.*]] = torch.constant.int 0
// CHECK: %[[CST1:.*]] = torch.constant.int 1
// CHECK: %[[END:.*]] = torch.aten.add.int %[[CST0]], %[[CST1]] : !torch.int, !torch.int -> !torch.int
// CHECK: %[[SLICE:.*]] = torch.aten.slice.Tensor %[[T]], %[[CST0]], %[[CST0]], %[[END]], %[[CST1]] :
// CHECK-SAME: !torch.vtensor<[?,?],si64>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.vtensor<[1,?],si64>
// CHECK: %[[SELECT:.*]] = torch.aten.squeeze.dim %[[SLICE]], %[[CST0]] :
// CHECK-SAME: !torch.vtensor<[1,?],si64>, !torch.int -> !torch.vtensor<[?],si64>
// CHECK: return %[[SELECT]] : !torch.vtensor<[?],si64>
func @torch.aten.select.int(%arg0: !torch.vtensor<[?,?],si64>) -> !torch.vtensor<[?],si64> {
%int0 = torch.constant.int 0
%0 = torch.aten.select.int %arg0, %int0, %int0 : !torch.vtensor<[?,?],si64>, !torch.int, !torch.int -> !torch.vtensor<[?],si64>
return %0 : !torch.vtensor<[?],si64>
}
[LINALG] Add E2E support for `aten.Hardsigmoid` op This commit adds lowering of `aten.Hardsigmoid` op. Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com>
2022-02-14 22:46:44 +08:00
// -----
// CHECK-LABEL: func @torch.aten.hardsigmoid(
Add aten.hardtanh e2e support.
2022-02-09 04:57:23 +08:00
// CHECK-SAME: %[[INPUT:.*]]: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[?,?],f32> {
// CHECK: %[[CST1:.*]] = torch.constant.int 1
// CHECK: %[[CST2:.*]] = torch.constant.int 3
// CHECK: %[[CST6:.*]] = torch.constant.int 6
// CHECK: %[[ADD:.*]] = torch.aten.add.Scalar %[[INPUT]], %[[CST2]], %[[CST1]] : !torch.vtensor<[?,?],f32>, !torch.int, !torch.int -> !torch.vtensor<[?,?],f32>
// CHECK: %[[DIV:.*]] = torch.aten.div.Scalar %[[ADD]], %[[CST6]] : !torch.vtensor<[?,?],f32>, !torch.int -> !torch.vtensor<[?,?],f32>
// CHECK: %[[CST0:.*]] = torch.constant.int 0
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: %[[SIZES:.*]] = torch.prim.ListConstruct : () -> !torch.list<int>
[LINALG] Add E2E support for `aten.Hardsigmoid` op This commit adds lowering of `aten.Hardsigmoid` op. Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com>
2022-02-14 22:46:44 +08:00
// CHECK: %[[NONE:.*]] = torch.constant.none
Add aten.hardtanh e2e support.
2022-02-09 04:57:23 +08:00
// CHECK: %[[EMPTY:.*]] = torch.aten.empty.memory_format %[[SIZES]], %[[NONE]], %[[NONE]], %[[NONE]], %[[NONE]], %[[NONE]] :
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: !torch.list<int>, !torch.none, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[],f32>
Improve "pseudo" op terminology. The term "pseudo" is very vague and was getting confusing (I felt I had to explain it in every comment referencing it). Instead, rework the "pseudo" ops to instead be named: - MLIR Syntax: `torch.valsem.*` - C++ / ODS: `ValsemVariant*Op` This makes it clear what the concept is, and avoids confusion with other things that might be called "pseudo", since these are very specific and should be 100% consistently named w.r.t. the non-valsem-variant ops that they correspond to.
2022-03-16 07:57:33 +08:00
// CHECK: %[[CST1_TENSOR:.*]] = torch.valsem.aten.fill.Scalar %[[EMPTY]], %[[CST1]] : !torch.vtensor<[],f32>, !torch.int -> !torch.vtensor<[],f32>
Add aten.hardtanh e2e support.
2022-02-09 04:57:23 +08:00
// CHECK: %[[MIN:.*]] = torch.aten.minimum %[[CST1_TENSOR]], %[[DIV]] : !torch.vtensor<[],f32>, !torch.vtensor<[?,?],f32> -> !torch.vtensor<[?,?],f32>
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: %[[SIZES:.*]] = torch.prim.ListConstruct : () -> !torch.list<int>
Add aten.hardtanh e2e support.
2022-02-09 04:57:23 +08:00
// CHECK: %[[NONE_1:.*]] = torch.constant.none
// CHECK: %[[EMPTY_1:.*]] = torch.aten.empty.memory_format %[[SIZES]], %[[NONE_1]], %[[NONE_1]], %[[NONE_1]], %[[NONE_1]], %[[NONE_1]] :
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: !torch.list<int>, !torch.none, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[],f32>
Improve "pseudo" op terminology. The term "pseudo" is very vague and was getting confusing (I felt I had to explain it in every comment referencing it). Instead, rework the "pseudo" ops to instead be named: - MLIR Syntax: `torch.valsem.*` - C++ / ODS: `ValsemVariant*Op` This makes it clear what the concept is, and avoids confusion with other things that might be called "pseudo", since these are very specific and should be 100% consistently named w.r.t. the non-valsem-variant ops that they correspond to.
2022-03-16 07:57:33 +08:00
// CHECK: %[[CST0_TENSOR:.*]] = torch.valsem.aten.fill.Scalar %[[EMPTY_1]], %[[CST0]] : !torch.vtensor<[],f32>, !torch.int -> !torch.vtensor<[],f32>
Add aten.hardtanh e2e support.
2022-02-09 04:57:23 +08:00
// CHECK: %[[RET:.*]] = torch.aten.maximum %[[CST0_TENSOR]], %[[MIN]] : !torch.vtensor<[],f32>, !torch.vtensor<[?,?],f32> -> !torch.vtensor<[?,?],f32>
// CHECK: return %[[RET]] : !torch.vtensor<[?,?],f32>
// CHECK: }
[LINALG] Add E2E support for `aten.Hardsigmoid` op This commit adds lowering of `aten.Hardsigmoid` op. Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com>
2022-02-14 22:46:44 +08:00
func @torch.aten.hardsigmoid(%arg0: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[?,?],f32> {
%0 = torch.aten.hardsigmoid %arg0 : !torch.vtensor<[?,?],f32> -> !torch.vtensor<[?,?],f32>
return %0 : !torch.vtensor<[?,?],f32>
}
[LINALG] Decompose aten_hardswish op. `aten.hardswish` op is decomposed into (x/6) * Relu6(x+3).
2022-02-15 21:14:32 +08:00
// -----
// CHECK-LABEL: func @torch.aten.hardswish(
// CHECK-SAME: %[[INP:.*]]: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[?,?],f32> {
// CHECK: %[[INT1:.*]] = torch.constant.int 1
// CHECK: %[[INT3:.*]] = torch.constant.int 3
// CHECK: %[[INT6:.*]] = torch.constant.int 6
// CHECK: %[[ADD:.*]] = torch.aten.add.Scalar %[[INP]], %[[INT3]], %[[INT1]] : !torch.vtensor<[?,?],f32>, !torch.int, !torch.int -> !torch.vtensor<[?,?],f32>
// CHECK: %[[RELU:.*]] = torch.aten.relu %[[ADD]] : !torch.vtensor<[?,?],f32> -> !torch.vtensor<[?,?],f32>
// CHECK: %[[INT6_:.*]] = torch.constant.int 6
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: %[[LIST:.*]] = torch.prim.ListConstruct : () -> !torch.list<int>
Add aten.hardtanh e2e support.
2022-02-09 04:57:23 +08:00
// CHECK: %[[NONE:.*]] = torch.constant.none
// CHECK: %[[MEM:.*]] = torch.aten.empty.memory_format %[[LIST]], %[[NONE]], %[[NONE]], %[[NONE]], %[[NONE]], %[[NONE]] :
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: !torch.list<int>, !torch.none, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[],f32>
Improve "pseudo" op terminology. The term "pseudo" is very vague and was getting confusing (I felt I had to explain it in every comment referencing it). Instead, rework the "pseudo" ops to instead be named: - MLIR Syntax: `torch.valsem.*` - C++ / ODS: `ValsemVariant*Op` This makes it clear what the concept is, and avoids confusion with other things that might be called "pseudo", since these are very specific and should be 100% consistently named w.r.t. the non-valsem-variant ops that they correspond to.
2022-03-16 07:57:33 +08:00
// CHECK: %[[FILL:.*]] = torch.valsem.aten.fill.Scalar %[[MEM]], %[[INT6_]] : !torch.vtensor<[],f32>, !torch.int -> !torch.vtensor<[],f32>
Add aten.hardtanh e2e support.
2022-02-09 04:57:23 +08:00
// CHECK: %[[MIN:.*]] = torch.aten.minimum %[[RELU]], %[[FILL]] : !torch.vtensor<[?,?],f32>, !torch.vtensor<[],f32> -> !torch.vtensor<[?,?],f32>
[LINALG] Decompose aten_hardswish op. `aten.hardswish` op is decomposed into (x/6) * Relu6(x+3).
2022-02-15 21:14:32 +08:00
// CHECK: %[[DIV:.*]] = torch.aten.div.Scalar %[[MIN]], %[[INT6]] : !torch.vtensor<[?,?],f32>, !torch.int -> !torch.vtensor<[?,?],f32>
// CHECK: %[[MUL:.*]] = torch.aten.mul.Tensor %[[DIV]], %[[INP]] : !torch.vtensor<[?,?],f32>, !torch.vtensor<[?,?],f32> -> !torch.vtensor<[?,?],f32>
// CHECK: return %[[MUL]] : !torch.vtensor<[?,?],f32>
func @torch.aten.hardswish(%arg0: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[?,?],f32> {
%0 = torch.aten.hardswish %arg0 : !torch.vtensor<[?,?],f32> -> !torch.vtensor<[?,?],f32>
return %0 : !torch.vtensor<[?,?],f32>
}
[MLIR][TORCH] Add E2E support for aten.new_zeros, aten.new_ones op This commit adds lowering of `aten.new_zeros` and `aten.new_ones` op Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-02-28 14:14:40 +08:00
Add aten.hardtanh e2e support.
2022-02-09 04:57:23 +08:00
// -----
// CHECK-LABEL: func @torch.aten.hardtanh(
// CHECK-SAME: %[[INPUT:.*]]: !torch.vtensor<[?],f32>,
// CHECK-SAME: %[[MIN_VAL:.*]]: !torch.float,
// CHECK-SAME: %[[MAX_VAL:.*]]: !torch.float) -> !torch.vtensor<[?],f32> {
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: %[[SIZES:.*]] = torch.prim.ListConstruct : () -> !torch.list<int>
Add aten.hardtanh e2e support.
2022-02-09 04:57:23 +08:00
// CHECK: %[[NONE:.*]] = torch.constant.none
// CHECK: %[[EMPTY:.*]] = torch.aten.empty.memory_format %[[SIZES]], %[[NONE]], %[[NONE]], %[[NONE]], %[[NONE]], %[[NONE]] :
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: !torch.list<int>, !torch.none, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[],f32>
Improve "pseudo" op terminology. The term "pseudo" is very vague and was getting confusing (I felt I had to explain it in every comment referencing it). Instead, rework the "pseudo" ops to instead be named: - MLIR Syntax: `torch.valsem.*` - C++ / ODS: `ValsemVariant*Op` This makes it clear what the concept is, and avoids confusion with other things that might be called "pseudo", since these are very specific and should be 100% consistently named w.r.t. the non-valsem-variant ops that they correspond to.
2022-03-16 07:57:33 +08:00
// CHECK: %[[MIN_TENSOR:.*]] = torch.valsem.aten.fill.Scalar %[[EMPTY]], %[[MIN_VAL]] : !torch.vtensor<[],f32>, !torch.float -> !torch.vtensor<[],f32>
Add aten.hardtanh e2e support.
2022-02-09 04:57:23 +08:00
// CHECK: %[[MIN:.*]] = torch.aten.maximum %[[INPUT]], %[[MIN_TENSOR]] : !torch.vtensor<[?],f32>, !torch.vtensor<[],f32> -> !torch.vtensor<[?],f32>
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: %[[SIZES:.*]] = torch.prim.ListConstruct : () -> !torch.list<int>
Add aten.hardtanh e2e support.
2022-02-09 04:57:23 +08:00
// CHECK: %[[NONE:.*]] = torch.constant.none
// CHECK: %[[VAL_10:.*]] = torch.aten.empty.memory_format %[[SIZES]], %[[NONE]], %[[NONE]], %[[NONE]], %[[NONE]], %[[NONE]] :
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: !torch.list<int>, !torch.none, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[],f32>
Improve "pseudo" op terminology. The term "pseudo" is very vague and was getting confusing (I felt I had to explain it in every comment referencing it). Instead, rework the "pseudo" ops to instead be named: - MLIR Syntax: `torch.valsem.*` - C++ / ODS: `ValsemVariant*Op` This makes it clear what the concept is, and avoids confusion with other things that might be called "pseudo", since these are very specific and should be 100% consistently named w.r.t. the non-valsem-variant ops that they correspond to.
2022-03-16 07:57:33 +08:00
// CHECK: %[[MAX_TENSOR:.*]] = torch.valsem.aten.fill.Scalar %[[VAL_10]], %[[MAX_VAL]] : !torch.vtensor<[],f32>, !torch.float -> !torch.vtensor<[],f32>
Add aten.hardtanh e2e support.
2022-02-09 04:57:23 +08:00
// CHECK: %[[RET:.*]] = torch.aten.minimum %[[MAX_TENSOR]], %[[MIN]] : !torch.vtensor<[],f32>, !torch.vtensor<[?],f32> -> !torch.vtensor<[?],f32>
// CHECK: return %[[RET]] : !torch.vtensor<[?],f32>
func @torch.aten.hardtanh(%arg0: !torch.vtensor<[?],f32>, %min: !torch.float, %max: !torch.float) -> !torch.vtensor<[?],f32> {
%0 = torch.aten.hardtanh %arg0, %min, %max : !torch.vtensor<[?],f32>, !torch.float, !torch.float -> !torch.vtensor<[?],f32>
return %0 : !torch.vtensor<[?],f32>
}
[MLIR][TORCH] Add E2E support for aten.new_zeros, aten.new_ones op This commit adds lowering of `aten.new_zeros` and `aten.new_ones` op Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-02-28 14:14:40 +08:00
// -----
// CHECK-LABEL: func @torch.aten.new_zeros
// CHECK-SAME: %[[INP:.*]]: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[2,3],f32> {
// CHECK: %[[NONE:.*]] = torch.constant.none
// CHECK: %[[INT2:.*]] = torch.constant.int 2
// CHECK: %[[INT3:.*]] = torch.constant.int 3
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: %[[SIZE:.*]] = torch.prim.ListConstruct %[[INT2]], %[[INT3]] : (!torch.int, !torch.int) -> !torch.list<int>
// CHECK: %[[RES:.*]] = torch.aten.zeros %[[SIZE]], %[[NONE]], %[[NONE]], %[[NONE]], %[[NONE]] : !torch.list<int>, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[2,3],f32>
[MLIR][TORCH] Add E2E support for aten.new_zeros, aten.new_ones op This commit adds lowering of `aten.new_zeros` and `aten.new_ones` op Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-02-28 14:14:40 +08:00
// CHECK: return %[[RES]] : !torch.vtensor<[2,3],f32>
// CHECK: }
func @torch.aten.new_zeros(%arg0: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[2,3],f32> {
%none = torch.constant.none
%int2 = torch.constant.int 2
%int3 = torch.constant.int 3
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
%0 = torch.prim.ListConstruct %int2, %int3 : (!torch.int, !torch.int) -> !torch.list<int>
%1 = torch.aten.new_zeros %arg0, %0, %none, %none, %none, %none : !torch.vtensor<[?,?],f32>, !torch.list<int>, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[2,3],f32>
[MLIR][TORCH] Add E2E support for aten.new_zeros, aten.new_ones op This commit adds lowering of `aten.new_zeros` and `aten.new_ones` op Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-02-28 14:14:40 +08:00
return %1 : !torch.vtensor<[2,3],f32>
}
// -----
// CHECK-LABEL: func @torch.aten.new_ones
// CHECK-SAME: %[[INP:.*]]: !torch.vtensor<[?,?],si64>) -> !torch.vtensor<[3,4],si64> {
// CHECK: %[[NONE:.*]] = torch.constant.none
// CHECK: %[[INT3:.*]] = torch.constant.int 3
// CHECK: %[[INT4:.*]] = torch.constant.int 4
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: %[[SIZE:.*]] = torch.prim.ListConstruct %[[INT3]], %[[INT4]] : (!torch.int, !torch.int) -> !torch.list<int>
// CHECK: %[[RES:.*]] = torch.aten.ones %[[SIZE]], %[[NONE]], %[[NONE]], %[[NONE]], %[[NONE]] : !torch.list<int>, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[3,4],si64>
[MLIR][TORCH] Add E2E support for aten.new_zeros, aten.new_ones op This commit adds lowering of `aten.new_zeros` and `aten.new_ones` op Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-02-28 14:14:40 +08:00
// CHECK: return %[[RES]] : !torch.vtensor<[3,4],si64>
// CHECK: }
func @torch.aten.new_ones(%arg0: !torch.vtensor<[?,?],si64>) -> !torch.vtensor<[3,4],si64> {
%none = torch.constant.none
%int3 = torch.constant.int 3
%int4 = torch.constant.int 4
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
%0 = torch.prim.ListConstruct %int3, %int4 : (!torch.int, !torch.int) -> !torch.list<int>
%1 = torch.aten.new_ones %arg0, %0, %none, %none, %none, %none : !torch.vtensor<[?,?],si64>, !torch.list<int>, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[3,4],si64>
[MLIR][TORCH] Add E2E support for aten.new_zeros, aten.new_ones op This commit adds lowering of `aten.new_zeros` and `aten.new_ones` op Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-02-28 14:14:40 +08:00
return %1 : !torch.vtensor<[3,4],si64>
}
Decompose aten.silu op Decomposition of aten.silu.op is added as silu(x) = x * sigmoid(x).
2022-03-02 01:30:58 +08:00
// -----
// CHECK-LABEL: func @torch.aten.silu(
// CHECK-SAME: %[[INP:.*]]: !torch.vtensor<[?,?],f32>) -> !torch.vtensor {
// CHECK: %[[SIGMOID:.*]] = torch.aten.sigmoid %[[INP]] : !torch.vtensor<[?,?],f32> -> !torch.vtensor
// CHECK: %[[MUL:.*]] = torch.aten.mul.Tensor %[[SIGMOID]], %[[INP]] : !torch.vtensor, !torch.vtensor<[?,?],f32> -> !torch.vtensor
// CHECK: return %[[MUL]] : !torch.vtensor
func @torch.aten.silu(%arg0: !torch.vtensor<[?,?],f32> loc(unknown)) -> !torch.vtensor {
%0 = torch.aten.silu %arg0 : !torch.vtensor<[?,?],f32> -> !torch.vtensor
return %0 : !torch.vtensor
}
[MLIR][TORCH] Add E2E support for aten.full op This commit decomposes `aten.full` op into `aten.empty` and `aten.fill` op. Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-03-03 21:41:14 +08:00
// -----
// CHECK-LABEL: func @torch.aten.full
// CHECK-SAME: () -> !torch.vtensor<[2,3],f32> {
// CHECK: %[[FLOAT5:.*]] = torch.constant.float 5.000000e+00
// CHECK: %[[INT3:.*]] = torch.constant.int 3
// CHECK: %[[INT2:.*]] = torch.constant.int 2
// CHECK: %[[NONE:.*]] = torch.constant.none
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: %[[SIZE:.*]] = torch.prim.ListConstruct %[[INT2]], %[[INT3]] : (!torch.int, !torch.int) -> !torch.list<int>
[MLIR][TORCH] Add E2E support for aten.full op This commit decomposes `aten.full` op into `aten.empty` and `aten.fill` op. Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-03-03 21:41:14 +08:00
// CHECK: %[[MEM_FORMAT:.*]] = torch.constant.none
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: %[[EMPTY:.*]] = torch.aten.empty.memory_format %[[SIZE]], %[[NONE]], %[[NONE]], %[[NONE]], %[[NONE]], %[[MEM_FORMAT]] : !torch.list<int>, !torch.none, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[2,3],f32>
Improve "pseudo" op terminology. The term "pseudo" is very vague and was getting confusing (I felt I had to explain it in every comment referencing it). Instead, rework the "pseudo" ops to instead be named: - MLIR Syntax: `torch.valsem.*` - C++ / ODS: `ValsemVariant*Op` This makes it clear what the concept is, and avoids confusion with other things that might be called "pseudo", since these are very specific and should be 100% consistently named w.r.t. the non-valsem-variant ops that they correspond to.
2022-03-16 07:57:33 +08:00
// CHECK: %[[RES:.*]] = torch.valsem.aten.fill.Scalar %[[EMPTY]], %[[FLOAT5]] : !torch.vtensor<[2,3],f32>, !torch.float -> !torch.vtensor<[2,3],f32>
[MLIR][TORCH] Add E2E support for aten.full op This commit decomposes `aten.full` op into `aten.empty` and `aten.fill` op. Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-03-03 21:41:14 +08:00
// CHECK: return %[[RES]] : !torch.vtensor<[2,3],f32>
func @torch.aten.full() -> !torch.vtensor<[2,3],f32> {
%float5.000000e00 = torch.constant.float 5.000000e+00
%int3 = torch.constant.int 3
%int2 = torch.constant.int 2
%none = torch.constant.none
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
%0 = torch.prim.ListConstruct %int2, %int3 : (!torch.int, !torch.int) -> !torch.list<int>
%1 = torch.aten.full %0, %float5.000000e00, %none, %none, %none, %none : !torch.list<int>, !torch.float, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[2,3],f32>
[MLIR][TORCH] Add E2E support for aten.full op This commit decomposes `aten.full` op into `aten.empty` and `aten.fill` op. Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-03-03 21:41:14 +08:00
return %1 : !torch.vtensor<[2,3],f32>
}
[MLIR][TORCH] Add E2E support for aten.full_like op This commit decomposes `aten.full_like` op into `aten.empty_like` and `aten.fill` op. Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-03-03 22:25:22 +08:00
// -----
// CHECK-LABEL: func @torch.aten.full_like(
// CHECK-SAME: %[[INP:.*]]: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[?,?],f32> {
// CHECK: %[[INT5:.*]] = torch.constant.int 5
// CHECK: %[[NONE:.*]] = torch.constant.none
// CHECK: %[[INT0:.*]] = torch.constant.int 0
// CHECK: %[[DIM0:.*]] = torch.aten.size.int %[[INP]], %[[INT0]] : !torch.vtensor<[?,?],f32>, !torch.int -> !torch.int
// CHECK: %[[INT1:.*]] = torch.constant.int 1
// CHECK: %[[DIM1:.*]] = torch.aten.size.int %[[INP]], %[[INT1]] : !torch.vtensor<[?,?],f32>, !torch.int -> !torch.int
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: %[[SIZE:.*]] = torch.prim.ListConstruct %[[DIM0]], %[[DIM1]] : (!torch.int, !torch.int) -> !torch.list<int>
// CHECK: %[[EMPTY:.*]] = torch.aten.empty.memory_format %[[SIZE]], %[[NONE]], %[[NONE]], %[[NONE]], %[[NONE]], %[[NONE]] : !torch.list<int>, !torch.none, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[?,?],f32>
Improve "pseudo" op terminology. The term "pseudo" is very vague and was getting confusing (I felt I had to explain it in every comment referencing it). Instead, rework the "pseudo" ops to instead be named: - MLIR Syntax: `torch.valsem.*` - C++ / ODS: `ValsemVariant*Op` This makes it clear what the concept is, and avoids confusion with other things that might be called "pseudo", since these are very specific and should be 100% consistently named w.r.t. the non-valsem-variant ops that they correspond to.
2022-03-16 07:57:33 +08:00
// CHECK: %[[RES:.*]] = torch.valsem.aten.fill.Scalar %[[EMPTY]], %[[INT5]] : !torch.vtensor<[?,?],f32>, !torch.int -> !torch.vtensor<[?,?],f32>
[MLIR][TORCH] Add E2E support for aten.full_like op This commit decomposes `aten.full_like` op into `aten.empty_like` and `aten.fill` op. Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-03-03 22:25:22 +08:00
// CHECK: return %[[RES]] : !torch.vtensor<[?,?],f32>
func @torch.aten.full_like(%arg0: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[?,?],f32> {
%int5 = torch.constant.int 5
%none = torch.constant.none
%0 = torch.aten.full_like %arg0, %int5, %none, %none, %none, %none, %none : !torch.vtensor<[?,?],f32>, !torch.int, !torch.none, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[?,?],f32>
return %0 : !torch.vtensor<[?,?],f32>
}
[MLIR][TORCH] Add E2E support for aten.index_put op This commit decomposes `aten.index_put` op into `valsem.aten.index_put_impl` op. Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-03-10 23:18:08 +08:00
// -----
// CHECK-LABEL: func @torch.aten.index_put(
// CHECK-SAME: %[[INP:.*]]: !torch.vtensor<[?],f32>, %[[INDEX:.*]]: !torch.vtensor<[?],si64>,
// CHECK-SAME: %[[VALUES:.*]]: !torch.vtensor<[?],f32>,
// CHECK-SAME: %[[ACCUM:.*]]: !torch.bool) -> !torch.vtensor<[?],f32> {
// CHECK: %[[INDICES:.*]] = torch.prim.ListConstruct %[[INDEX]] : (!torch.vtensor<[?],si64>) -> !torch.list<vtensor>
// CHECK: %[[FALSE:.*]] = torch.constant.bool false
// CHECK: %[[RES:.*]] = torch.valsem.aten.index_put_impl %[[INP]], %[[INDICES]], %[[VALUES]], %[[ACCUM]], %[[FALSE]] : !torch.vtensor<[?],f32>, !torch.list<vtensor>, !torch.vtensor<[?],f32>, !torch.bool, !torch.bool -> !torch.vtensor<[?],f32>
// CHECK: return %[[RES]] : !torch.vtensor<[?],f32>
func @torch.aten.index_put(%input: !torch.vtensor<[?],f32>, %index: !torch.vtensor<[?],si64>, %values: !torch.vtensor<[?],f32>, %accumulate : !torch.bool) -> !torch.vtensor<[?],f32> {
%indices = torch.prim.ListConstruct %index : (!torch.vtensor<[?],si64>) -> !torch.list<vtensor>
%0 = torch.aten.index_put %input, %indices, %values, %accumulate : !torch.vtensor<[?],f32>, !torch.list<vtensor>, !torch.vtensor<[?],f32>, !torch.bool -> !torch.vtensor<[?],f32>
return %0 : !torch.vtensor<[?],f32>
}
[MLIR][TORCH] Add E2E support for aten.expand_as op This commit decomposes `aten.expand_as` op into `aten.broadcast_to` op. Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>
2022-03-14 16:12:37 +08:00
// -----
// CHECK-LABEL: func @torch.aten.expand_as(
// CHECK-SAME: %[[INP:.*]]: !torch.vtensor<[?,1,1],f32>, %[[OTHER:.*]]: !torch.vtensor<[?,?,?],f32>) -> !torch.vtensor<[?,?,?],f32> {
// CHECK: %[[INT0:.*]] = torch.constant.int 0
// CHECK: %[[DIM0:.*]] = torch.aten.size.int %[[OTHER]], %[[INT0]] : !torch.vtensor<[?,?,?],f32>, !torch.int -> !torch.int
// CHECK: %[[INT1:.*]] = torch.constant.int 1
// CHECK: %[[DIM1:.*]] = torch.aten.size.int %[[OTHER]], %[[INT1]] : !torch.vtensor<[?,?,?],f32>, !torch.int -> !torch.int
// CHECK: %[[INT2:.*]] = torch.constant.int 2
// CHECK: %[[DIM2:.*]] = torch.aten.size.int %[[OTHER]], %[[INT2]] : !torch.vtensor<[?,?,?],f32>, !torch.int -> !torch.int
// CHECK: %[[SIZE:.*]] = torch.prim.ListConstruct %[[DIM0]], %[[DIM1]], %[[DIM2]] : (!torch.int, !torch.int, !torch.int) -> !torch.list<int>
// CHECK: %[[RES:.*]] = torch.aten.broadcast_to %[[INP]], %[[SIZE]] : !torch.vtensor<[?,1,1],f32>, !torch.list<int> -> !torch.vtensor<[?,?,?],f32>
// CHECK: return %[[RES]] : !torch.vtensor<[?,?,?],f32>
func @torch.aten.expand_as(%arg0: !torch.vtensor<[?,1,1],f32>, %arg1: !torch.vtensor<[?,?,?],f32>) -> !torch.vtensor<[?,?,?],f32> {
%0 = torch.aten.expand_as %arg0, %arg1 : !torch.vtensor<[?,1,1],f32>, !torch.vtensor<[?,?,?],f32> -> !torch.vtensor<[?,?,?],f32>
return %0 : !torch.vtensor<[?,?,?],f32>
}