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/Conversion/TorchToLinalg/resize.mlir

155 lines
9.9 KiB
MLIR
Raw Normal View History

OnnxToTorch lowering resize op (#3013) https://github.com/nod-ai/SHARK-Turbine/issues/358 adds a lowering from onnx to linalg for bilinear and nearest resize with support for using scales or sizes to get resize shape. uses coordinate transform half pixel for bilinear mode and asymmetrical for nearest mode. See https://github.com/onnx/onnx/blob/main/docs/Operators.md#Resize. Added two passes -- one for bilinear and the other for nearest.
2024-05-09 05:35:03 +08:00
// RUN: torch-mlir-opt <%s -convert-torch-to-linalg -split-input-file -verify-diagnostics | FileCheck %s
// CHECK-LABEL: func.func @test_resize_sizes_linear
func.func @test_resize_sizes_linear(%arg0: !torch.vtensor<[1,1,2,4],f32>, %arg1: !torch.vtensor<[4]
,si64>) -> !torch.vtensor<[?,?,?,?],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 19 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
// CHECK: %[[generic:.*]] = linalg.generic
Modifies onnx resize lowering to fix numerical issues (#3381) Updates: - some unsupported modes are now going to report a match failure for unsupported coordinate transformation modes. - fixes a bug that was introduced in the last patch for resize (my bad...) - uses actual x and y coordinates for computing weights in bilinear interpolation (rather than eps modified values) - slightly simplifies the bilinear interpolation payload for readability and performance - passes coordinate transformation mode information from an onnx.Resize op to the mode string for the aten._interpolate op. This allows us to perform custom logic in the torch->linalg lowering to support onnx.Resize options without losing the default behaviors of the interpolate op.
2024-05-31 08:34:37 +08:00
// CHECK: %[[extracted:.*]] = tensor.extract %[[x0:.*]][%[[x1:.*]], %[[x2:.*]], %[[x3:.*]], %[[x4:.*]]] : tensor<1x1x2x4xf32>
// CHECK: %[[extracted_7:.*]] = tensor.extract %[[x0]][%[[x1]], %[[x2]]
// CHECK: %[[extracted_8:.*]] = tensor.extract %[[x0]][%[[x1]], %[[x2]]
// CHECK: %[[extracted_9:.*]] = tensor.extract %[[x0]][%[[x1]], %[[x2]]
// CHECK: %[[dx0p00:.*]] = arith.mulf %[[dx0:.*]], %[[extracted]]
// CHECK: %[[dx1p01:.*]] = arith.mulf %[[dx1:.*]], %[[extracted_7]]
// CHECK: %[[sum:.*]] = arith.addf %[[dx0p00]], %[[dx1p01]]
// CHECK: %[[left:.*]] = arith.mulf %[[dy0:.*]], %[[sum]]
// CHECK: %[[dx0p10:.*]] = arith.mulf %[[dx0]], %[[extracted_8]]
// CHECK: %[[dx1p11:.*]] = arith.mulf %[[dx1]], %[[extracted_9]]
// CHECK: %[[sum2:.*]] = arith.addf %[[dx0p10]], %[[dx1p11]]
// CHECK: %[[right:.*]] = arith.mulf %[[dy1:.*]], %[[sum2]]
// CHECK: %[[retval:.*]] = arith.addf %[[left]], %[[right]]
OnnxToTorch lowering resize op (#3013) https://github.com/nod-ai/SHARK-Turbine/issues/358 adds a lowering from onnx to linalg for bilinear and nearest resize with support for using scales or sizes to get resize shape. uses coordinate transform half pixel for bilinear mode and asymmetrical for nearest mode. See https://github.com/onnx/onnx/blob/main/docs/Operators.md#Resize. Added two passes -- one for bilinear and the other for nearest.
2024-05-09 05:35:03 +08:00
%none = torch.constant.none
%none_0 = torch.constant.none
%int0 = torch.constant.int 0
%false = torch.constant.bool false
%true = torch.constant.bool true
%str = torch.constant.str "bilinear"
%int2 = torch.constant.int 2
%0 = torch.aten.select.int %arg1, %int0, %int2 : !torch.vtensor<[4],si64>, !torch.int, !torch.int -> !torch.vtensor<[1],si64>
%1 = torch.aten.item %0 : !torch.vtensor<[1],si64> -> !torch.int
%int3 = torch.constant.int 3
%2 = torch.aten.select.int %arg1, %int0, %int3 : !torch.vtensor<[4],si64>, !torch.int, !torch.int -> !torch.vtensor<[1],si64>
%3 = torch.aten.item %2 : !torch.vtensor<[1],si64> -> !torch.int
%4 = torch.prim.ListConstruct %1, %3 : (!torch.int, !torch.int) -> !torch.list<int>
%5 = torch.aten.__interpolate.size_list_scale_list %arg0, %4, %none_0, %str, %false, %none_0, %false : !torch.vtensor<[1,1,2,4],f32>, !torch.list<int>, !torch.none, !torch.str, !torch.bool, !torch.none, !torch.bool -> !torch.vtensor<[?,?,?,?],f32>
return %5 : !torch.vtensor<[?,?,?,?],f32>
}
// -----
func.func @test_resize_sizes_nearest(%arg0: !torch.vtensor<[1,1,2,4],f32>, %arg1: !torch.vtensor<[4],si64>) -> !torch.vtensor<[?,?,?,?],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 19 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
// CHECK: %[[GENERIC:.*]] = linalg.generic
onnx.Resize and aten._interpolate : allow n spatial dims. (#3368) The old lowering only had logic for 2d (i.e. images). this patch allows interpolation for n spatial dims, which is required for some 3d vision models such as - onnx/models/pytorch-3dunet_vaiq_int8 which successfully compiles and runs with this patch.
2024-05-21 04:35:27 +08:00
// CHECK: %[[x11:.*]] = linalg.index 0 : index
// CHECK: %[[x12:.*]] = linalg.index 1 : index
OnnxToTorch lowering resize op (#3013) https://github.com/nod-ai/SHARK-Turbine/issues/358 adds a lowering from onnx to linalg for bilinear and nearest resize with support for using scales or sizes to get resize shape. uses coordinate transform half pixel for bilinear mode and asymmetrical for nearest mode. See https://github.com/onnx/onnx/blob/main/docs/Operators.md#Resize. Added two passes -- one for bilinear and the other for nearest.
2024-05-09 05:35:03 +08:00
// CHECK: %[[x13:.*]] = linalg.index 2 : index
// CHECK: %[[x14:.*]] = linalg.index 3 : index
[ONNX][TorchToLinalg] Add support for dynamic dims in Interpolate lowering (#3351) Addresses [Shark-Turbine #196](https://github.com/nod-ai/SHARK-TestSuite/issues/196) Related tracker [Shark-Turbine #566](https://github.com/nod-ai/SHARK-Turbine/issues/566) Related onnx.Resize issues [Shark-Turbine #616](https://github.com/nod-ai/SHARK-Turbine/issues/616)
2024-05-18 03:18:57 +08:00
// CHECK: %[[x15:.*]] = arith.sitofp %[[c2_i64:.*]] : i64 to f32
OnnxToTorch lowering resize op (#3013) https://github.com/nod-ai/SHARK-Turbine/issues/358 adds a lowering from onnx to linalg for bilinear and nearest resize with support for using scales or sizes to get resize shape. uses coordinate transform half pixel for bilinear mode and asymmetrical for nearest mode. See https://github.com/onnx/onnx/blob/main/docs/Operators.md#Resize. Added two passes -- one for bilinear and the other for nearest.
2024-05-09 05:35:03 +08:00
// CHECK: %[[x19:.*]] = arith.sitofp %[[x6:.*]] : i64 to f32
// CHECK: %[[x21:.*]] = arith.divf %[[x19]], %[[x15]] : f32
// CHECK: %[[x23:.*]] = arith.index_cast %[[x13]] : index to i64
// CHECK: %[[x24:.*]] = arith.sitofp %[[x23]] : i64 to f32
// CHECK: %[[x25:.*]] = arith.divf %[[x24]], %[[x21]] : f32
onnx.Resize and aten._interpolate : allow n spatial dims. (#3368) The old lowering only had logic for 2d (i.e. images). this patch allows interpolation for n spatial dims, which is required for some 3d vision models such as - onnx/models/pytorch-3dunet_vaiq_int8 which successfully compiles and runs with this patch.
2024-05-21 04:35:27 +08:00
// CHECK: %[[x29:.*]] = math.floor %[[x25]] : f32
// CHECK: %[[x31:.*]] = arith.fptosi %[[x29]] : f32 to i64
// CHECK: %[[x32:.*]] = arith.index_cast %[[x31]] : i64 to index
// CHECK: %[[x16:.*]] = arith.sitofp %[[c4_i64:.*]] : i64 to f32
// CHECK: %[[x20:.*]] = arith.sitofp %[[x7:.*]] : i64 to f32
// CHECK: %[[x22:.*]] = arith.divf %[[x20]], %[[x16]] : f32
OnnxToTorch lowering resize op (#3013) https://github.com/nod-ai/SHARK-Turbine/issues/358 adds a lowering from onnx to linalg for bilinear and nearest resize with support for using scales or sizes to get resize shape. uses coordinate transform half pixel for bilinear mode and asymmetrical for nearest mode. See https://github.com/onnx/onnx/blob/main/docs/Operators.md#Resize. Added two passes -- one for bilinear and the other for nearest.
2024-05-09 05:35:03 +08:00
// CHECK: %[[x26:.*]] = arith.index_cast %[[x14]] : index to i64
// CHECK: %[[x27:.*]] = arith.sitofp %[[x26]] : i64 to f32
// CHECK: %[[x28:.*]] = arith.divf %[[x27]], %[[x22]] : f32
// CHECK: %[[x30:.*]] = math.floor %[[x28]] : f32
// CHECK: %[[x33:.*]] = arith.fptosi %[[x30]] : f32 to i64
// CHECK: %[[x34:.*]] = arith.index_cast %[[x33]] : i64 to index
onnx.Resize and aten._interpolate : allow n spatial dims. (#3368) The old lowering only had logic for 2d (i.e. images). this patch allows interpolation for n spatial dims, which is required for some 3d vision models such as - onnx/models/pytorch-3dunet_vaiq_int8 which successfully compiles and runs with this patch.
2024-05-21 04:35:27 +08:00
// CHECK: %[[extracted:.*]] = tensor.extract %[[x0:.*]][%[[x11]], %[[x12]], %[[x32]], %[[x34]]] : tensor<1x1x2x4xf32>
OnnxToTorch lowering resize op (#3013) https://github.com/nod-ai/SHARK-Turbine/issues/358 adds a lowering from onnx to linalg for bilinear and nearest resize with support for using scales or sizes to get resize shape. uses coordinate transform half pixel for bilinear mode and asymmetrical for nearest mode. See https://github.com/onnx/onnx/blob/main/docs/Operators.md#Resize. Added two passes -- one for bilinear and the other for nearest.
2024-05-09 05:35:03 +08:00
// CHECK: linalg.yield %[[extracted]] : f32
%none = torch.constant.none
%none_0 = torch.constant.none
%int0 = torch.constant.int 0
%false = torch.constant.bool false
%true = torch.constant.bool true
%str = torch.constant.str "nearest"
%int2 = torch.constant.int 2
%0 = torch.aten.select.int %arg1, %int0, %int2 : !torch.vtensor<[4],si64>, !torch.int, !torch.int -> !torch.vtensor<[1],si64>
%1 = torch.aten.item %0 : !torch.vtensor<[1],si64> -> !torch.int
%int3 = torch.constant.int 3
%2 = torch.aten.select.int %arg1, %int0, %int3 : !torch.vtensor<[4],si64>, !torch.int, !torch.int -> !torch.vtensor<[1],si64>
%3 = torch.aten.item %2 : !torch.vtensor<[1],si64> -> !torch.int
%4 = torch.prim.ListConstruct %1, %3 : (!torch.int, !torch.int) -> !torch.list<int>
%5 = torch.aten.__interpolate.size_list_scale_list %arg0, %4, %none_0, %str, %false, %none_0, %false : !torch.vtensor<[1,1,2,4],f32>, !torch.list<int>, !torch.none, !torch.str, !torch.bool, !torch.none, !torch.bool -> !torch.vtensor<[?,?,?,?],f32>
return %5 : !torch.vtensor<[?,?,?,?],f32>
}
onnx.Resize and aten._interpolate : allow n spatial dims. (#3368) The old lowering only had logic for 2d (i.e. images). this patch allows interpolation for n spatial dims, which is required for some 3d vision models such as - onnx/models/pytorch-3dunet_vaiq_int8 which successfully compiles and runs with this patch.
2024-05-21 04:35:27 +08:00
// -----
func.func @test_resize_nearest_1d(%arg0: !torch.vtensor<[?,?,?],f32>, %arg1: !torch.vtensor<[3],si64>) -> !torch.vtensor<[?,?,?],f32> {
// CHECK: %[[GENERIC:.*]] = linalg.generic
// CHECK: %[[x11:.*]] = linalg.index 0 : index
// CHECK: %[[x12:.*]] = linalg.index 1 : index
// CHECK: %[[x13:.*]] = linalg.index 2 : index
// CHECK: %[[x15:.*]] = arith.sitofp %[[c2_i64:.*]] : i64 to f32
// CHECK: %[[x19:.*]] = arith.sitofp %[[x6:.*]] : i64 to f32
// CHECK: %[[x21:.*]] = arith.divf %[[x19]], %[[x15]] : f32
// CHECK: %[[x23:.*]] = arith.index_cast %[[x13]] : index to i64
// CHECK: %[[x24:.*]] = arith.sitofp %[[x23]] : i64 to f32
// CHECK: %[[x25:.*]] = arith.divf %[[x24]], %[[x21]] : f32
// CHECK: %[[x29:.*]] = math.floor %[[x25]] : f32
// CHECK: %[[x31:.*]] = arith.fptosi %[[x29]] : f32 to i64
// CHECK: %[[x32:.*]] = arith.index_cast %[[x31]] : i64 to index
// CHECK: %[[extracted:.*]] = tensor.extract %[[x0:.*]][%[[x11]], %[[x12]], %[[x32]]] : tensor<?x?x?xf32>
// CHECK: linalg.yield %[[extracted]] : f32
%none = torch.constant.none
%none_0 = torch.constant.none
%int0 = torch.constant.int 0
%false = torch.constant.bool false
%true = torch.constant.bool true
%str = torch.constant.str "nearest"
%int2 = torch.constant.int 2
%0 = torch.aten.select.int %arg1, %int0, %int2 : !torch.vtensor<[3],si64>, !torch.int, !torch.int -> !torch.vtensor<[1],si64>
%1 = torch.aten.item %0 : !torch.vtensor<[1],si64> -> !torch.int
%4 = torch.prim.ListConstruct %1 : (!torch.int) -> !torch.list<int>
%5 = torch.aten.__interpolate.size_list_scale_list %arg0, %4, %none_0, %str, %false, %none_0, %false : !torch.vtensor<[?,?,?],f32>, !torch.list<int>, !torch.none, !torch.str, !torch.bool, !torch.none, !torch.bool -> !torch.vtensor<[?,?,?],f32>
return %5 : !torch.vtensor<[?,?,?],f32>
}
// -----
func.func @test_resize_nearest_3d(%arg0: !torch.vtensor<[?,?,?,?,?],f32>, %arg1: !torch.vtensor<[5],si64>) -> !torch.vtensor<[?,?,?,?,?],f32> {
// CHECK: %[[GENERIC:.*]] = linalg.generic
// CHECK: %[[x11:.*]] = linalg.index 0 : index
// CHECK: %[[x12:.*]] = linalg.index 1 : index
// CHECK: %[[x13:.*]] = linalg.index 2 : index
// CHECK: %[[x14:.*]] = linalg.index 3 : index
// CHECK: %[[index4:.*]] = linalg.index 4 : index
// CHECK: %[[x15:.*]] = arith.sitofp %[[c2_i64:.*]] : i64 to f32
// CHECK: %[[x19:.*]] = arith.sitofp %[[x6:.*]] : i64 to f32
// CHECK: %[[x21:.*]] = arith.divf %[[x19]], %[[x15]] : f32
// CHECK: %[[x23:.*]] = arith.index_cast %[[x13]] : index to i64
// CHECK: %[[x24:.*]] = arith.sitofp %[[x23]] : i64 to f32
// CHECK: %[[x25:.*]] = arith.divf %[[x24]], %[[x21]] : f32
// CHECK: %[[x29:.*]] = math.floor %[[x25]] : f32
// CHECK: %[[x31:.*]] = arith.fptosi %[[x29]] : f32 to i64
// CHECK: %[[x32:.*]] = arith.index_cast %[[x31]] : i64 to index
// CHECK: %[[x34:.*]] = arith.index_cast %[[Wfptosi:.*]] : i64 to index
// CHECK: %[[x35:.*]] = arith.index_cast %[[Dfptosi:.*]] : i64 to index
// CHECK: %[[extracted:.*]] = tensor.extract %[[x0:.*]][%[[x11]], %[[x12]], %[[x32]], %[[x34]], %[[x35]]] : tensor<?x?x?x?x?xf32>
// CHECK: linalg.yield %[[extracted]] : f32
%none = torch.constant.none
%none_0 = torch.constant.none
%int0 = torch.constant.int 0
%false = torch.constant.bool false
%true = torch.constant.bool true
%str = torch.constant.str "nearest"
%int2 = torch.constant.int 2
%0 = torch.aten.select.int %arg1, %int0, %int2 : !torch.vtensor<[5],si64>, !torch.int, !torch.int -> !torch.vtensor<[1],si64>
%1 = torch.aten.item %0 : !torch.vtensor<[1],si64> -> !torch.int
%int3 = torch.constant.int 3
%2 = torch.aten.select.int %arg1, %int0, %int3 : !torch.vtensor<[5],si64>, !torch.int, !torch.int -> !torch.vtensor<[1],si64>
%3 = torch.aten.item %2 : !torch.vtensor<[1],si64> -> !torch.int
%int4 = torch.constant.int 4
%4 = torch.aten.select.int %arg1, %int0, %int4 : !torch.vtensor<[5],si64>, !torch.int, !torch.int -> !torch.vtensor<[1],si64>
%5 = torch.aten.item %0 : !torch.vtensor<[1],si64> -> !torch.int
%6 = torch.prim.ListConstruct %1, %3, %5: (!torch.int, !torch.int, !torch.int) -> !torch.list<int>
%7 = torch.aten.__interpolate.size_list_scale_list %arg0, %6, %none_0, %str, %false, %none_0, %false : !torch.vtensor<[?,?,?,?,?],f32>, !torch.list<int>, !torch.none, !torch.str, !torch.bool, !torch.none, !torch.bool -> !torch.vtensor<[?,?,?,?,?],f32>
return %7 : !torch.vtensor<[?,?,?,?,?],f32>
}