[Mlir-commits] [mlir] [mlir][tosa] Lower MAX_POOL2D_ADAPTIVE to Linalg (PR #225675)

[Jack Frankland]

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

>From da86372f001970c24c797a412641df9b917b2387 Mon Sep 17 00:00:00 2001
From: Jack Frankland <jack.frankland at arm.com>
Date: Wed, 23 Sep 2026 09:38:38 +0100
Subject: [PATCH] [mlir][tosa] Reject uncanonicalized MAX_POOL2D_ADAPTIVE

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

Test both sides of the pass-ordering contract: direct named conversion rejects the uncanonicalized operation, while the full TOSA-to-Linalg pipeline canonicalizes and lowers it successfully.

Assisted-by: Codex
---
 .../TosaToLinalg/TosaToLinalgNamedPass.cpp     |  1 +
 .../tosa-to-linalg-named-invalid.mlir          | 12 ++++++++++++
 .../TosaToLinalg/tosa-to-linalg-pipeline.mlir  | 18 ++++++++++++++++++
 3 files changed, 31 insertions(+)
 create mode 100644 mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named-invalid.mlir

diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
index df066a5cac261..15e5bdd3826be 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
@@ -58,6 +58,7 @@ struct TosaToLinalgNamed
     target.addIllegalOp<tosa::Conv3DOp>();
     target.addIllegalOp<tosa::DepthwiseConv2DOp>();
     target.addIllegalOp<tosa::MaxPool2dOp>();
+    target.addIllegalOp<tosa::MaxPool2dAdaptiveOp>();
     target.addIllegalOp<tosa::AvgPool2dOp>();
     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
new file mode 100644
index 0000000000000..35d8eadd77c71
--- /dev/null
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named-invalid.mlir
@@ -0,0 +1,12 @@
+// RUN: mlir-opt --split-input-file -pass-pipeline="builtin.module(func.func(tosa-to-linalg-named))" %s -verify-diagnostics
+
+func.func @max_pool2d_adaptive_not_canonicalized(
+    %arg0: tensor<1x4x4x1xf32>) -> tensor<1x2x2x1xf32> {
+  %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.max_pool2d_adaptive'}}
+  %0 = tosa.max_pool2d_adaptive %arg0, %kernel, %stride, %pad nan_mode<IGNORE> :
+    (tensor<1x4x4x1xf32>, !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 1667735b06ffe..69d5a02ac38f5 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-pipeline.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-pipeline.mlir
@@ -31,6 +31,24 @@ func.func @avg_pool2d_with_unsupported_quant_type(%arg0: tensor<1x7x7x9x!quant.u
 
 // -----
 
+// Check that canonicalization rewrites max_pool2d_adaptive before named
+// TOSA-to-Linalg conversion.
+// CHECK-LABEL: func.func @max_pool2d_adaptive(
+// CHECK-NOT: tosa.max_pool2d_adaptive
+// CHECK: linalg.generic
+// CHECK: arith.maximumf
+func.func @max_pool2d_adaptive(
+    %arg0: tensor<1x4x4x1xf32>) -> tensor<1x2x2x1xf32> {
+  %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.max_pool2d_adaptive %arg0, %kernel, %stride, %pad nan_mode<IGNORE> :
+    (tensor<1x4x4x1xf32>, !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>