ONNXToTosa: MaxPoolSingleOut: Support dilations whens its 1

Xilinx · Mar 12, 2024 · 9189bd7 · 9189bd7
1 parent d266832
commit 9189bd7
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 2 deletions.
diff --git a/src/Conversion/ONNXToTOSA/NN/MaxPoolSingleOut.cpp b/src/Conversion/ONNXToTOSA/NN/MaxPoolSingleOut.cpp
@@ -48,9 +48,11 @@ class ONNXMaxPoolSingleOutOpLoweringToTOSA : public ConversionPattern {
       return rewriter.notifyMatchFailure(
           op, "memrefs as inputs are unsupported by TOSA");
     }
-    if (dilations) {
+    auto isOne = [](IntegerAttr attr) { return attr.getValue().isOne(); };
+    if (dilations &&
+        !llvm::all_of(dilations.getAsRange<IntegerAttr>(), isOne)) {
       return rewriter.notifyMatchFailure(
-          maxpoolOp, "dilations attribute is unsupported by TOSA");
+          maxpoolOp, "dilations != 1 is unsupported by TOSA");
     }
     if (storageOrder && storageOrder.getSInt() != 0) {
       return rewriter.notifyMatchFailure(

diff --git a/test/mlir/conversion/onnx_to_tosa/NN/MaxPoolSingleOut.mlir b/test/mlir/conversion/onnx_to_tosa/NN/MaxPoolSingleOut.mlir
@@ -116,3 +116,26 @@ func.func @test_default_maxpoolsingleout_same_upper_ceil_mode(%arg0 : tensor<5x5
 // CHECK-DAG:     "tosa.const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi32>}> : () -> tensor<4xi32>
 // CHECK-DAG:     %[[TRANS_MPOOL_RES:.*]] = tosa.transpose %[[MPOOL_RES]], %3 : (tensor<5x4x4x5xf32>, tensor<4xi32>) -> tensor<5x5x4x4xf32>
 // CHECK-DAG:     return %[[TRANS_MPOOL_RES]] : tensor<5x5x4x4xf32>
+
+// -----
+
+func.func @test_maxpoolsingleout_dilation1(%arg0 : tensor<5x5x32x32xf32>) -> tensor<5x5x30x30xf32> {
+  %0 = "onnx.MaxPoolSingleOut"(%arg0) {kernel_shape = [3,3], dilations = [1,1]} : (tensor<5x5x32x32xf32>) -> tensor<5x5x30x30xf32>
+  return %0 : tensor<5x5x30x30xf32>
+}
+// CHECK-LABEL: func.func @test_maxpoolsingleout_dilation1(%arg0: tensor<5x5x32x32xf32>) -> tensor<5x5x30x30xf32> {
+// CHECK-DAG:   "tosa.const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi32>}> : () -> tensor<4xi32>
+// CHECK-DAG:   %[[TRANS_ARG:.*]] = tosa.transpose %arg0, %0 : (tensor<5x5x32x32xf32>, tensor<4xi32>) -> tensor<5x32x32x5xf32>
+// CHECK-DAG:   %[[MPOOL_RES:.*]] = tosa.max_pool2d %[[TRANS_ARG]] {kernel = array<i64: 3, 3>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} : (tensor<5x32x32x5xf32>) -> tensor<5x30x30x5xf32>
+// CHECK-DAG:   "tosa.const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi32>}> : () -> tensor<4xi32>
+// CHECK-DAG:   %[[TRANS_MPOOL_RES:.*]] = tosa.transpose %[[MPOOL_RES]], %3 : (tensor<5x30x30x5xf32>, tensor<4xi32>) -> tensor<5x5x30x30xf32>
+// CHECK-DAG:   return %[[TRANS_MPOOL_RES]] : tensor<5x5x30x30xf32>
+
+// -----
+
+func.func @test_maxpoolsingleout_dilation2(%arg0 : tensor<5x5x32x32xf32>) -> tensor<5x5x30x30xf32> {
+  %0 = "onnx.MaxPoolSingleOut"(%arg0) {kernel_shape = [3,3], dilations = [2,2]} : (tensor<5x5x32x32xf32>) -> tensor<5x5x30x30xf32>
+  return %0 : tensor<5x5x30x30xf32>
+}
+// CHECK-LABEL: func.func @test_maxpoolsingleout_dilation2(%arg0: tensor<5x5x32x32xf32>) -> tensor<5x5x30x30xf32> {
+// CHECK: onnx.MaxPoolSingleOut