Skip to content

Commit d6feb21

Browse files
sriram-siloaisriram-siloai
authored
Added support for Maxpool (Autopad) (#3774)
Added autopad. and passed 3 tests test_maxpool_2d_precomputed_same_upper test_maxpool_2d_same_lower' test_maxpool_2d_same_upper Address : nod-ai/SHARK-ModelDev#843 2 attributes yet to complete : storage_order, indices output
1 parent 2f9a68c commit d6feb21

File tree

2 files changed

+109
-3
lines changed

2 files changed

+109
-3
lines changed

‎lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1087,9 +1087,6 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
10871087
if (binder.customOpNameStringAttr(autoPad, "auto_pad", "NOTSET"))
10881088
return rewriter.notifyMatchFailure(binder.op,
10891089
"auto_pad bind failure");
1090-
if (autoPad != "NOTSET")
1091-
return rewriter.notifyMatchFailure(
1092-
binder.op, "unsupported conversion: auto_pad != NOTSET");
10931090

10941091
Torch::ValueTensorType resultTypeOut;
10951092
Value operand;
@@ -1136,13 +1133,42 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
11361133
return rewriter.notifyMatchFailure(binder.op,
11371134
"dilations bind failure");
11381135

1136+
// set default padding
11391137
if (padding.empty())
11401138
padding.resize(spatial, 0);
11411139
if (strides.empty())
11421140
strides.resize(spatial, 1);
11431141
if (dilations.empty())
11441142
dilations.resize(spatial, 1);
11451143

1144+
auto inputTensorType = cast<Torch::ValueTensorType>(operand.getType());
1145+
1146+
// Padding for the beginning and ending along each spatial axis, it can
1147+
// take any value greater than or equal to 0. The value represent the
1148+
// number of pixels added to the beginning and end part of the
1149+
// corresponding axis. pads format should be as follow [x1_begin,
1150+
// x2_begin…x1_end, x2_end,…], where xi_begin the number of pixels added
1151+
// at the beginning of axis i and xi_end, the number of pixels added at
1152+
// the end of axis i.
1153+
if (autoPad != "NOTSET" && autoPad != "VALID") {
1154+
const bool isSameLower = autoPad == "SAME_LOWER";
1155+
ArrayRef<int64_t> inputShape = inputTensorType.getSizes();
1156+
padding.resize_for_overwrite(2 * spatial);
1157+
for (unsigned dimIdx = 0; dimIdx < spatial; dimIdx++) {
1158+
const int64_t dilatedKernelSize =
1159+
dilations[dimIdx] * (kernel[dimIdx] - 1) + 1;
1160+
int64_t totalPad = ((inputShape[dimIdx + 2] + strides[dimIdx] - 1) /
1161+
strides[dimIdx] -
1162+
1) *
1163+
strides[dimIdx] +
1164+
dilatedKernelSize - inputShape[dimIdx + 2];
1165+
totalPad = totalPad >= 0 ? totalPad : 0;
1166+
padding[dimIdx] =
1167+
isSameLower ? ((totalPad + 1) / 2) : (totalPad / 2);
1168+
padding[spatial + dimIdx] = totalPad - padding[dimIdx];
1169+
}
1170+
}
1171+
11461172
// If the padding is symmetric we can push the padding operation to the
11471173
// torch operator.
11481174
if (padding.size() == static_cast<size_t>(2 * spatial)) {

‎test/Conversion/TorchOnnxToTorch/simple_ops_g_to_p.mlir

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -730,6 +730,86 @@ func.func @test_maxpool_pad(%arg0: !torch.vtensor<[1,64,111,111],f32>) -> !torch
730730
return %0 : !torch.vtensor<[1,64,56,56],f32>
731731
}
732732

733+
// -----
734+
735+
// CHECK-LABEL: func.func @test_maxpool_2d_same_lower
736+
func.func @test_maxpool_2d_same_lower(%arg0: !torch.vtensor<[1,3,32,32],f32>) -> !torch.vtensor<[1,3,32,32],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 12 : si64} {
737+
// CHECK: %[[int1:.*]] = torch.constant.int 1
738+
// CHECK: %[[int0:.*]] = torch.constant.int 0
739+
// CHECK: %[[int1_0:.*]] = torch.constant.int 1
740+
// CHECK: %[[int0_1:.*]] = torch.constant.int 0
741+
// CHECK: %[[list0:.*]] = torch.prim.ListConstruct %[[int1]], %[[int0]], %[[int1_0]], %[[int0_1]] : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
742+
// CHECK: %[[FLOAT0:.*]] = torch.constant.float -1.7976931348623157E+308
743+
// CHECK: %[[FUNC1:.*]] = torch.aten.constant_pad_nd %arg0, %[[list0]], %[[FLOAT0]] : !torch.vtensor<[1,3,32,32],f32>, !torch.list<int>, !torch.float -> !torch.vtensor<[1,3,33,33],f32>
744+
// CHECK: %[[int2:.*]] = torch.constant.int 2
745+
// CHECK: %[[int2_2:.*]] = torch.constant.int 2
746+
// CHECK: %[[list1:.*]] = torch.prim.ListConstruct %[[int2]], %[[int2_2]] : (!torch.int, !torch.int) -> !torch.list<int>
747+
// CHECK: %[[int0_3:.*]] = torch.constant.int 0
748+
// CHECK: %[[int0_4:.*]] = torch.constant.int 0
749+
// CHECK: %[[list2:.*]] = torch.prim.ListConstruct %[[int0_3]], %[[int0_4]] : (!torch.int, !torch.int) -> !torch.list<int>
750+
// CHECK: %[[int1_5:.*]] = torch.constant.int 1
751+
// CHECK: %[[int1_6:.*]] = torch.constant.int 1
752+
// CHECK: %[[list3:.*]] = torch.prim.ListConstruct %[[int1_5]], %[[int1_6]] : (!torch.int, !torch.int) -> !torch.list<int>
753+
// CHECK: %[[int1_7:.*]] = torch.constant.int 1
754+
// CHECK: %[[int1_8:.*]] = torch.constant.int 1
755+
// CHECK: %[[list4:.*]] = torch.prim.ListConstruct %[[int1_7]], %[[int1_8]] : (!torch.int, !torch.int) -> !torch.list<int>
756+
// CHECK: %[[FALSE:.*]] = torch.constant.bool false
757+
// CHECK: %[[FUNC6:.*]] = torch.aten.max_pool2d %[[FUNC1]], %[[list1]], %[[list3]], %[[list2]], %[[list4]], %[[FALSE]] : !torch.vtensor<[1,3,33,33],f32>, !torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.bool -> !torch.vtensor<[1,3,32,32],f32>
758+
%0 = torch.operator "onnx.MaxPool"(%arg0) {torch.onnx.auto_pad = "SAME_LOWER", torch.onnx.kernel_shape = [2 : si64, 2 : si64]} : (!torch.vtensor<[1,3,32,32],f32>) -> !torch.vtensor<[1,3,32,32],f32>
759+
return %0 : !torch.vtensor<[1,3,32,32],f32>
760+
}
761+
762+
// -----
763+
764+
// CHECK-LABEL: func.func @test_maxpool_2d_same_upper
765+
func.func @test_maxpool_2d_same_upper(%arg0: !torch.vtensor<[1,3,32,32],f32>) -> !torch.vtensor<[1,3,32,32],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 12 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
766+
// CHECK: %[[int0:.*]] = torch.constant.int 0
767+
// CHECK: %[[int1:.*]] = torch.constant.int 1
768+
// CHECK: %[[int0_0:.*]] = torch.constant.int 0
769+
// CHECK: %[[int1_1:.*]] = torch.constant.int 1
770+
// CHECK: %[[list0:.*]] = torch.prim.ListConstruct %[[int0]], %[[int1]], %[[int0_0]], %[[int1_1]] : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
771+
// CHECK: %[[FLOAT0:.*]] = torch.constant.float -1.7976931348623157E+308
772+
// CHECK: %[[FUNC1:.*]] = torch.aten.constant_pad_nd %arg0, %[[list0]], %[[FLOAT0]] : !torch.vtensor<[1,3,32,32],f32>, !torch.list<int>, !torch.float -> !torch.vtensor<[1,3,33,33],f32>
773+
// CHECK: %[[int2:.*]] = torch.constant.int 2
774+
// CHECK: %[[int2_2:.*]] = torch.constant.int 2
775+
// CHECK: %[[list1:.*]] = torch.prim.ListConstruct %[[int2]], %[[int2_2]] : (!torch.int, !torch.int) -> !torch.list<int>
776+
// CHECK: %[[int0_3:.*]] = torch.constant.int 0
777+
// CHECK: %[[int0_4:.*]] = torch.constant.int 0
778+
// CHECK: %[[list2:.*]] = torch.prim.ListConstruct %[[int0_3]], %[[int0_4]] : (!torch.int, !torch.int) -> !torch.list<int>
779+
// CHECK: %[[int1_5:.*]] = torch.constant.int 1
780+
// CHECK: %[[int1_6:.*]] = torch.constant.int 1
781+
// CHECK: %[[list3:.*]] = torch.prim.ListConstruct %[[int1_5]], %[[int1_6]] : (!torch.int, !torch.int) -> !torch.list<int>
782+
// CHECK: %[[int1_7:.*]] = torch.constant.int 1
783+
// CHECK: %[[int1_8:.*]] = torch.constant.int 1
784+
// CHECK: %[[list4:.*]] = torch.prim.ListConstruct %[[int1_7]], %[[int1_8]] : (!torch.int, !torch.int) -> !torch.list<int>
785+
// CHECK: %[[FALSE:.*]] = torch.constant.bool false
786+
// CHECK: %[[FUNC6:.*]] = torch.aten.max_pool2d %[[FUNC1]], %[[list1]], %[[list3]], %[[list2]], %[[list4]], %[[FALSE]] : !torch.vtensor<[1,3,33,33],f32>, !torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.bool -> !torch.vtensor<[1,3,32,32],f32>
787+
%0 = torch.operator "onnx.MaxPool"(%arg0) {torch.onnx.auto_pad = "SAME_UPPER", torch.onnx.kernel_shape = [2 : si64, 2 : si64]} : (!torch.vtensor<[1,3,32,32],f32>) -> !torch.vtensor<[1,3,32,32],f32>
788+
return %0 : !torch.vtensor<[1,3,32,32],f32>
789+
}
790+
791+
// -----
792+
793+
// CHECK-LABEL: func.func @test_maxpool_2d_precomputed_same_upper
794+
func.func @test_maxpool_2d_precomputed_same_upper(%arg0: !torch.vtensor<[1,1,5,5],f32>) -> !torch.vtensor<[1,1,3,3],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 12 : si64}{
795+
// CHECK: %[[int3:.*]] = torch.constant.int 3
796+
// CHECK: %[[int3_0:.*]] = torch.constant.int 3
797+
// CHECK: %[[list0:.*]] = torch.prim.ListConstruct %[[int3]], %[[int3_0]] : (!torch.int, !torch.int) -> !torch.list<int>
798+
// CHECK: %[[int1:.*]] = torch.constant.int 1
799+
// CHECK: %[[int1_1:.*]] = torch.constant.int 1
800+
// CHECK: %[[list1:.*]] = torch.prim.ListConstruct %[[int1]], %[[int1_1]] : (!torch.int, !torch.int) -> !torch.list<int>
801+
// CHECK: %[[int2:.*]] = torch.constant.int 2
802+
// CHECK: %[[int2_2:.*]] = torch.constant.int 2
803+
// CHECK: %[[list2:.*]] = torch.prim.ListConstruct %[[int2]], %[[int2_2]] : (!torch.int, !torch.int) -> !torch.list<int>
804+
// CHECK: %[[int1_3:.*]] = torch.constant.int 1
805+
// CHECK: %[[int1_4:.*]] = torch.constant.int 1
806+
// CHECK: %[[list3:.*]] = torch.prim.ListConstruct %[[int1_3]], %[[int1_4]] : (!torch.int, !torch.int) -> !torch.list<int>
807+
// CHECK: %[[FALSE:.*]] = torch.constant.bool false
808+
// CHECK: %[[FUNC4:.*]] = torch.aten.max_pool2d %arg0, %[[list0]], %[[list2]], %[[list1]], %[[list3]], %[[FALSE]] : !torch.vtensor<[1,1,5,5],f32>, !torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.bool -> !torch.vtensor<[1,1,3,3],f32>
809+
%0 = torch.operator "onnx.MaxPool"(%arg0) {torch.onnx.auto_pad = "SAME_UPPER", torch.onnx.kernel_shape = [3 : si64, 3 : si64], torch.onnx.strides = [2 : si64, 2 : si64]} : (!torch.vtensor<[1,1,5,5],f32>) -> !torch.vtensor<[1,1,3,3],f32>
810+
return %0 : !torch.vtensor<[1,1,3,3],f32>
811+
}
812+
733813

734814
// -----
735815

0 commit comments

Comments
 (0)