[Mlir-commits] [mlir] bfd7810 - [mlir][tosa] Lower AVG_POOL2D_ADAPTIVE to Linalg (#225637)

[llvmlistbot at llvm.org]

Author: Jack Frankland
Date: 2026-10-01T10:10:33+01:00
New Revision: bfd7810d5d91f7123eb7428d0db06ebf2bb4bd73

URL: https://github.com/llvm/llvm-project/commit/bfd7810d5d91f7123eb7428d0db06ebf2bb4bd73
DIFF: https://github.com/llvm/llvm-project/commit/bfd7810d5d91f7123eb7428d0db06ebf2bb4bd73.diff

LOG: [mlir][tosa] Lower AVG_POOL2D_ADAPTIVE to Linalg (#225637)

Register the adaptive average-pool canonicalization with the named
TOSA-to-Linalg conversion and mark the operation illegal so constant
shape operands lower through the established avg_pool2d path.

Add coverage for padding-aware floating-point pooling and quantized
integer pooling with nonzero zero points.

Assited-by: Codex

Added: 
    

Modified: 
    mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
    mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named-invalid.mlir
    mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-pipeline.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
index 15e5bdd3826be..780901eab90ee 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
@@ -60,6 +60,7 @@ struct TosaToLinalgNamed
     target.addIllegalOp<tosa::MaxPool2dOp>();
     target.addIllegalOp<tosa::MaxPool2dAdaptiveOp>();
     target.addIllegalOp<tosa::AvgPool2dOp>();
+    target.addIllegalOp<tosa::AvgPool2dAdaptiveOp>();
     target.addIllegalOp<tosa::MatMulOp>();
     target.addIllegalOp<tosa::TransposeOp>();
 

diff  --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named-invalid.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named-invalid.mlir
index 35d8eadd77c71..d9de3f223f5d0 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named-invalid.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named-invalid.mlir
@@ -10,3 +10,18 @@ func.func @max_pool2d_adaptive_not_canonicalized(
     (tensor<1x4x4x1xf32>, !tosa.shape<2>, !tosa.shape<2>, !tosa.shape<4>) -> tensor<1x2x2x1xf32>
   return %0 : tensor<1x2x2x1xf32>
 }
+
+// -----
+
+func.func @avg_pool2d_adaptive_not_canonicalized(
+    %arg0: tensor<1x4x4x1xf32>) -> tensor<1x2x2x1xf32> {
+  %input_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %output_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %kernel = tosa.const_shape values(dense<[2, 2]> : tensor<2xindex>) : () -> !tosa.shape<2>
+  %stride = tosa.const_shape values(dense<[2, 2]> : tensor<2xindex>) : () -> !tosa.shape<2>
+  %pad = tosa.const_shape values(dense<[0, 0, 0, 0]> : tensor<4xindex>) : () -> !tosa.shape<4>
+  // expected-error at +1 {{failed to legalize operation 'tosa.avg_pool2d_adaptive'}}
+  %0 = tosa.avg_pool2d_adaptive %arg0, %input_zp, %output_zp, %kernel, %stride, %pad acc_type(f32) :
+    (tensor<1x4x4x1xf32>, tensor<1xf32>, tensor<1xf32>, !tosa.shape<2>, !tosa.shape<2>, !tosa.shape<4>) -> tensor<1x2x2x1xf32>
+  return %0 : tensor<1x2x2x1xf32>
+}

diff  --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-pipeline.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-pipeline.mlir
index 69d5a02ac38f5..cd34715359666 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-pipeline.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-pipeline.mlir
@@ -49,6 +49,25 @@ func.func @max_pool2d_adaptive(
 
 // -----
 
+// Check that canonicalization rewrites avg_pool2d_adaptive before named
+// TOSA-to-Linalg conversion.
+// CHECK-LABEL: func.func @avg_pool2d_adaptive(
+// CHECK-NOT: tosa.avg_pool2d_adaptive
+// CHECK: linalg.pooling_nhwc_sum
+func.func @avg_pool2d_adaptive(
+    %arg0: tensor<1x4x4x1xf32>) -> tensor<1x2x2x1xf32> {
+  %input_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %output_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %kernel = tosa.const_shape values(dense<[2, 2]> : tensor<2xindex>) : () -> !tosa.shape<2>
+  %stride = tosa.const_shape values(dense<[2, 2]> : tensor<2xindex>) : () -> !tosa.shape<2>
+  %pad = tosa.const_shape values(dense<[0, 0, 0, 0]> : tensor<4xindex>) : () -> !tosa.shape<4>
+  %0 = tosa.avg_pool2d_adaptive %arg0, %input_zp, %output_zp, %kernel, %stride, %pad acc_type(f32) :
+    (tensor<1x4x4x1xf32>, tensor<1xf32>, tensor<1xf32>, !tosa.shape<2>, !tosa.shape<2>, !tosa.shape<4>) -> tensor<1x2x2x1xf32>
+  return %0 : tensor<1x2x2x1xf32>
+}
+
+// -----
+
 // CHECK-LABEL: rescale_doubleround
 func.func @rescale_doubleround(%arg0: tensor<8x9x7x14xi32>) -> tensor<8x9x7x14xi8> {
   %0 = "tosa.const"() <{values = dense<0> : tensor<14xi32>}> : () -> tensor<14xi32>