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

[Jack Frankland]

https://github.com/FranklandJack updated https://github.com/llvm/llvm-project/pull/225637

>From 2f285cfcb496f3bb00cc92423b58424aa01ccfdc Mon Sep 17 00:00:00 2001
From: Jack Frankland <jack.frankland at arm.com>
Date: Wed, 16 Sep 2026 10:57:35 +0100
Subject: [PATCH] [mlir][tosa] Lower AVG_POOL2D_ADAPTIVE to Linalg

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
---
 .../TosaToLinalg/TosaToLinalgNamed.cpp        |  3 ++
 .../TosaToLinalg/TosaToLinalgNamedPass.cpp    |  1 +
 .../TosaToLinalg/tosa-to-linalg-named.mlir    | 42 +++++++++++++++++++
 3 files changed, 46 insertions(+)

diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index 4a11a927090bf..e583d5bba16c4 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -1141,5 +1141,8 @@ void mlir::tosa::populateTosaToLinalgNamedConversionPatterns(
   patterns->add<
       MaxPool2dConverter
     >(converter, patterns->getContext());
+
+  tosa::AvgPool2dAdaptiveOp::getCanonicalizationPatterns(
+      *patterns, patterns->getContext());
   // clang-format on
 }
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
index df066a5cac261..11d56e1a22bec 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
@@ -59,6 +59,7 @@ struct TosaToLinalgNamed
     target.addIllegalOp<tosa::DepthwiseConv2DOp>();
     target.addIllegalOp<tosa::MaxPool2dOp>();
     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.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
index a4abc687bfe38..f78765cd3bf20 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
@@ -465,6 +465,48 @@ func.func @avg_pool_dyn(%arg0: tensor<?x6x34x62xf32>) -> (tensor<?x5x33x62xf32>)
 
 // -----
 
+// Verify adaptive average pooling resolves shape operands and reuses the
+// padding-aware floating-point average-pool lowering.
+// CHECK-LABEL: @avg_pool_adaptive_f32
+func.func @avg_pool_adaptive_f32(%arg0: tensor<1x5x6x2xf32>) -> tensor<1x3x3x2xf32> {
+  // CHECK: %[[PADDED:.+]] = tensor.pad %arg0 low[0, 1, 0, 0] high[0, 1, 0, 0]
+  // CHECK: %[[WINDOW:.+]] = tensor.empty() : tensor<3x2xf32>
+  // CHECK: %[[POOL:.+]] = linalg.pooling_nhwc_sum
+  // CHECK-SAME: strides = dense<2> : vector<2xi64>
+  // CHECK-SAME: ins(%[[PADDED]], %[[WINDOW]] : tensor<1x7x6x2xf32>, tensor<3x2xf32>)
+  // CHECK: linalg.generic
+  // CHECK: arith.divf
+  // CHECK-NOT: tosa.avg_pool2d_adaptive
+  %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<[3, 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<[1, 1, 0, 0]> : tensor<4xindex>) : () -> !tosa.shape<4>
+  %0 = tosa.avg_pool2d_adaptive %arg0, %input_zp, %output_zp, %kernel, %stride, %pad acc_type(f32) : (tensor<1x5x6x2xf32>, tensor<1xf32>, tensor<1xf32>, !tosa.shape<2>, !tosa.shape<2>, !tosa.shape<4>) -> tensor<1x3x3x2xf32>
+  return %0 : tensor<1x3x3x2xf32>
+}
+
+// -----
+
+// Verify adaptive average pooling preserves integer zero points and the i32
+// accumulator path when resolving its shape operands.
+// CHECK-LABEL: @avg_pool_adaptive_i8
+func.func @avg_pool_adaptive_i8(%arg0: tensor<1x4x4x1xi8>) -> tensor<1x2x2x1xi8> {
+  // CHECK: linalg.pooling_nhwc_sum
+  // CHECK: tosa.apply_scale
+  // CHECK: arith.addi
+  // CHECK-NOT: tosa.avg_pool2d_adaptive
+  %input_zp = "tosa.const"() <{values = dense<-3> : tensor<1xi8>}> : () -> tensor<1xi8>
+  %output_zp = "tosa.const"() <{values = dense<5> : tensor<1xi8>}> : () -> tensor<1xi8>
+  %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(i32) : (tensor<1x4x4x1xi8>, tensor<1xi8>, tensor<1xi8>, !tosa.shape<2>, !tosa.shape<2>, !tosa.shape<4>) -> tensor<1x2x2x1xi8>
+  return %0 : tensor<1x2x2x1xi8>
+}
+
+// -----
+
 // CHECK: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (0)>
 // CHECK: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>