[Mlir-commits] [mlir] [mlir][tosa] Lower AVG_POOL2D_ADAPTIVE to Linalg (PR #225637)
Jack Frankland
llvmlistbot at llvm.org
Mon Sep 28 04:00:52 PDT 2026
https://github.com/FranklandJack updated https://github.com/llvm/llvm-project/pull/225637
>From 365f44d986da86b398c7e9795631c32d197475c6 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 | 88 +++++++++++++++++++
3 files changed, 92 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..5bbf81f3854fb 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
@@ -465,6 +465,94 @@ 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: %[[ZERO:.+]] = arith.constant 0 : i32
+ // CHECK: %[[POOL_EMPTY:.+]] = tensor.empty() : tensor<1x2x2x1xi32>
+ // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[ZERO]] : i32) outs(%[[POOL_EMPTY]] : tensor<1x2x2x1xi32>)
+ // CHECK: %[[WINDOW:.+]] = tensor.empty() : tensor<2x2xi32>
+ // CHECK: %[[POOL:.+]] = linalg.pooling_nhwc_sum
+ // CHECK-SAME: dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>
+ // CHECK-SAME: ins(%arg0, %[[WINDOW]] : tensor<1x4x4x1xi8>, tensor<2x2xi32>)
+ // CHECK-SAME: outs(%[[FILL]] : tensor<1x2x2x1xi32>) -> tensor<1x2x2x1xi32>
+ // CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x2x2x1xi8>
+ // CHECK: %[[GENERIC:.+]] = linalg.generic
+ // CHECK-SAME: indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+ // CHECK-SAME: ins(%[[POOL]] : tensor<1x2x2x1xi32>)
+ // CHECK-SAME: outs(%[[EMPTY]] : tensor<1x2x2x1xi8>)
+ // CHECK: ^bb0(%[[IN:.+]]: i32, %{{.+}}: i8)
+
+ // Compute the input zero-point offset and remove it from the pool sum.
+ // CHECK: %[[COUNT:.+]] = arith.muli %{{[0-9]+}}, %{{[0-9]+}}
+ // CHECK: %[[ICAST:.+]] = arith.index_cast %[[COUNT]]
+ // CHECK: %[[INPUT_ZP:.+]] = arith.constant -3
+ // CHECK: %[[INPUT_OFFSET:.+]] = arith.muli %[[ICAST]], %[[INPUT_ZP]]
+ // CHECK: %[[CENTERED:.+]] = arith.subi %[[IN]], %[[INPUT_OFFSET]]
+
+ // Compute the multiplier and shift values for average pooling.
+ // CHECK: %[[C1:.+]] = arith.constant 1
+ // CHECK: %[[C32:.+]] = arith.constant 32
+ // CHECK: %[[ISUB:.+]] = arith.subi %[[ICAST]], %[[C1]]
+ // CHECK: %[[CTLZ:.+]] = math.ctlz %[[ISUB]]
+ // CHECK: %[[SUB:.+]] = arith.subi %[[C32]], %[[CTLZ]]
+ // CHECK: %[[EXT:.+]] = arith.extui %[[SUB]]
+ // CHECK: %[[CBIG:.+]] = arith.constant 1073741825
+ // CHECK: %[[SHL:.+]] = arith.shli %[[CBIG]], %[[EXT]]
+ // CHECK: %[[IEXT:.+]] = arith.extui %[[ICAST]]
+ // CHECK: %[[DIV:.+]] = arith.divui %[[SHL]], %[[IEXT]]
+ // CHECK: %[[TRUNC_MUL:.+]] = arith.trunci %[[DIV]]
+ // CHECK: %[[TRUNC_SHIFT:.+]] = arith.trunci %[[SUB]]
+ // CHECK: %[[C30:.+]] = arith.constant 30
+ // CHECK: %[[SHIFT:.+]] = arith.addi %[[TRUNC_SHIFT]], %[[C30]] : i8
+ // CHECK: %[[SCALED:.+]] = tosa.apply_scale %[[CENTERED]], %[[TRUNC_MUL]], %[[SHIFT]] rounding_mode<SINGLE_ROUND>
+
+ // Add the output zero point and normalize to the result type.
+ // CHECK: %[[OUTPUT_ZP:.+]] = arith.constant 5
+ // CHECK: %[[OFFSET:.+]] = arith.addi %[[SCALED]], %[[OUTPUT_ZP]]
+ // CHECK: %[[CMIN:.+]] = arith.constant -128
+ // CHECK: %[[CMAX:.+]] = arith.constant 127
+ // CHECK: %[[LOW:.+]] = arith.maxsi %[[CMIN]], %[[OFFSET]]
+ // CHECK: %[[CLAMP:.+]] = arith.minsi %[[CMAX]], %[[LOW]]
+ // CHECK: %[[TRUNC:.+]] = arith.trunci %[[CLAMP]]
+ // CHECK: linalg.yield %[[TRUNC]]
+ // CHECK: return %[[GENERIC]] : tensor<1x2x2x1xi8>
+ // 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)>
More information about the Mlir-commits
mailing list