[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 f40ee53d5a991df1f00bf3747e1dd092240777ac 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] Lower MAX_POOL2D_ADAPTIVE to Linalg

Legalize max_pool2d_adaptive with constant shape operands through the
existing max_pool2d conversion.

Assisted-by: Codex
---
 .../TosaToLinalg/TosaToLinalgNamed.cpp        |  3 +++
 .../TosaToLinalg/TosaToLinalgNamedPass.cpp    |  1 +
 .../TosaToLinalg/tosa-to-linalg-named.mlir    | 24 +++++++++++++++++++
 3 files changed, 28 insertions(+)

diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index 4a11a927090bf..c787cc5885e8f 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -1120,6 +1120,9 @@ class TransposeConverter : public OpRewritePattern<tosa::TransposeOp> {
 void mlir::tosa::populateTosaToLinalgNamedConversionPatterns(
     const TypeConverter &converter, RewritePatternSet *patterns,
     const TosaToLinalgNamedOptions &options) {
+  MaxPool2dAdaptiveOp::getCanonicalizationPatterns(*patterns,
+                                                   patterns->getContext());
+
   if (options.preferConv2DKernelLayoutHWCF) {
     patterns->add<ConvConverter<tosa::Conv2DOp, linalg::Conv2DNhwcHwcfOp,
                                 linalg::Conv2DNhwcHwcfQOp>>(
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.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
index a4abc687bfe38..39bec5484b894 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
@@ -106,6 +106,30 @@ func.func @max_pool(%arg0: tensor<1x6x34x62xf32>) -> () {
   return
 }
 
+// -----
+
+// Lower constant adaptive pooling parameters, including nonzero padding,
+// through max_pool2d in PROPAGATE mode.
+// CHECK-LABEL: @max_pool_adaptive_propagate
+func.func @max_pool_adaptive_propagate(%arg0: tensor<1x5x6x2xf32>) -> tensor<1x3x3x2xf32> {
+  // CHECK-NOT: tosa.max_pool2d_adaptive
+  // CHECK-DAG: [[PAD_VALUE:%.+]] = arith.constant -3.40282347E+38 : f32
+  // CHECK-DAG: [[PADDED:%.+]] = tensor.pad %arg0 low[0, 1, 1, 0] high[0, 0, 1, 0]
+  // CHECK-DAG:   tensor.yield [[PAD_VALUE]]
+  // CHECK-DAG: [[INIT_VALUE:%.+]] = arith.constant -3.40282347E+38 : f32
+  // CHECK-DAG: [[INIT:%.+]] = tensor.empty()
+  // CHECK-DAG: [[FILL:%.+]] = linalg.fill ins([[INIT_VALUE]] : f32) outs([[INIT]] : tensor<1x3x3x2xf32>)
+  // CHECK-DAG: [[KERNEL:%.+]] = tensor.empty() : tensor<2x4xf32>
+  // CHECK: [[POOL:%.+]] = linalg.pooling_nhwc_max {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>} ins([[PADDED]], [[KERNEL]] : tensor<1x6x8x2xf32>, tensor<2x4xf32>) outs([[FILL]] : tensor<1x3x3x2xf32>)
+  // CHECK: return [[POOL]] : tensor<1x3x3x2xf32>
+  %kernel = tosa.const_shape values(dense<[2, 4]> : tensor<2xindex>) : () -> !tosa.shape<2>
+  %stride = tosa.const_shape values(dense<2> : tensor<2xindex>) : () -> !tosa.shape<2>
+  %pad = tosa.const_shape values(dense<[1, 0, 1, 1]> : tensor<4xindex>) : () -> !tosa.shape<4>
+  %0 = tosa.max_pool2d_adaptive %arg0, %kernel, %stride, %pad nan_mode<PROPAGATE> :
+    (tensor<1x5x6x2xf32>, !tosa.shape<2>, !tosa.shape<2>, !tosa.shape<4>) -> tensor<1x3x3x2xf32>
+  return %0 : tensor<1x3x3x2xf32>
+}
+
 // CHECK-LABEL: @max_pool_padded
 func.func @max_pool_padded(%arg0: tensor<1x6x34x62xf32>) -> () {
   // CHECK-DAG: [[CONST:%.+]] = arith.constant -3.40282347E+38 : f32