[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 e55d4da0fb2a4d8e6508e2b75e2158f520ce1d58 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] Reject uncanonicalized AVG_POOL2D_ADAPTIVE

Mark tosa.avg_pool2d_adaptive illegal in the named TOSA-to-Linalg
conversion so the pass fails if canonicalization has not rewritten it to
tosa.avg_pool2d.

Assisted-by: Codex
---
 .../TosaToLinalg/TosaToLinalgNamedPass.cpp    |  1 +
 .../tosa-to-linalg-named-invalid.mlir         | 15 +++++++++++++++
 .../TosaToLinalg/tosa-to-linalg-pipeline.mlir | 19 +++++++++++++++++++
 3 files changed, 35 insertions(+)

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>