[Mlir-commits] [mlir] [mlir][tosa] Canonicalize avg_pool2d/max_pool2d no-ops (PR #203571)
Luke Hutton
llvmlistbot at llvm.org
Thu Jun 18 05:31:21 PDT 2026
https://github.com/lhutton1 updated https://github.com/llvm/llvm-project/pull/203571
>From 0c2c18c4045f572b0209697677b8c4735a983c86 Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Tue, 9 Jun 2026 22:25:30 +0100
Subject: [PATCH 1/2] [mlir][tosa] Canonicalize avg_pool2d/max_pool2d no-ops
Removes avg_pool2d/max_pool2d ops with unit kernel, unit stride, and
zero padding, as these are no-ops.
Change-Id: Icd1ed62420c2269ad7e8d5d44bb7dd2c3b7e16f9
---
mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td | 1 +
.../Dialect/Tosa/IR/TosaCanonicalizations.cpp | 42 ++++++
mlir/test/Dialect/Tosa/canonicalize.mlir | 134 +++++++++++++++++-
3 files changed, 172 insertions(+), 5 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index a99fb2fcae547..e7e964ddda40f 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -116,6 +116,7 @@ def Tosa_AvgPool2dOp : Tosa_InferShapedTypeOp<"avg_pool2d", [NoMemoryEffect]> {
}];
let hasVerifier = 1;
+ let hasCanonicalizer = 1;
let assemblyFormat =
"operands attr-dict `:` functional-type(operands, results)";
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 63d40ed4a95a5..47b703ce5bdfa 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -271,6 +271,39 @@ struct AvgPool2dAdaptiveToAvgPool2d
}
};
+struct AvgPool2dIsNoOp : public OpRewritePattern<tosa::AvgPool2dOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tosa::AvgPool2dOp op,
+ PatternRewriter &rewriter) const override {
+ if (op.getInput().getType() != op.getOutput().getType())
+ return rewriter.notifyMatchFailure(
+ op, "expected input and output types to match");
+
+ const auto inputType = llvm::cast<ShapedType>(op.getInput().getType());
+ if (!llvm::isa<FloatType>(inputType.getElementType()))
+ return rewriter.notifyMatchFailure(op,
+ "expected floating-point input type");
+
+ if (!llvm::all_of(op.getKernel(), [](int64_t val) { return val == 1; }))
+ return rewriter.notifyMatchFailure(op, "expected unit kernel");
+
+ if (!llvm::all_of(op.getStride(), [](int64_t val) { return val == 1; }))
+ return rewriter.notifyMatchFailure(op, "expected unit stride");
+
+ if (!llvm::all_of(op.getPad(), [](int64_t val) { return val == 0; }))
+ return rewriter.notifyMatchFailure(op, "expected zero padding");
+
+ rewriter.replaceOp(op, op.getInput());
+ return success();
+ }
+};
+
+void AvgPool2dOp::getCanonicalizationPatterns(RewritePatternSet &results,
+ MLIRContext *context) {
+ results.add<AvgPool2dIsNoOp>(context);
+}
+
void AvgPool2dAdaptiveOp::getCanonicalizationPatterns(
RewritePatternSet &results, MLIRContext *context) {
results.add<AvgPool2dAdaptiveToAvgPool2d>(context);
@@ -286,6 +319,15 @@ struct MaxPool2dIsNoOp : public OpRewritePattern<tosa::MaxPool2dOp> {
ShapedType inputType = llvm::cast<ShapedType>(input.getType());
ShapedType outputType = llvm::cast<ShapedType>(output.getType());
+ if (input.getType() == output.getType() &&
+ llvm::all_of(op.getKernel(), [](int64_t val) { return val == 1; }) &&
+ llvm::all_of(op.getStride(), [](int64_t val) { return val == 1; }) &&
+ llvm::all_of(op.getPad(), [](int64_t val) { return val == 0; }) &&
+ op.getNanMode() == tosa::NanPropagationMode::PROPAGATE) {
+ rewriter.replaceOp(op, input);
+ return success();
+ }
+
if (!inputType.hasStaticShape() || !outputType.hasStaticShape()) {
return failure();
}
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index 2cd040f056db8..bae933d9290e6 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -1707,16 +1707,140 @@ func.func @test_do_not_canonicalize_cast_from_cast_to_block_scaled_unranked(%arg
// -----
+// CHECK-LABEL: @canonicalize_unit_avg_pool2d
+// CHECK-NOT: tosa.avg_pool2d
+// CHECK: return %arg0 : tensor<1x32x32x8xf32>
+func.func @canonicalize_unit_avg_pool2d(%arg0: tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32> {
+ %input_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %output_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.avg_pool2d %arg0, %input_zp, %output_zp {acc_type = f32, kernel = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} :
+ (tensor<1x32x32x8xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x32x32x8xf32>
+ return %0 : tensor<1x32x32x8xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @dont_canonicalize_unit_avg_pool2d_integer
+// CHECK: tosa.avg_pool2d
+func.func @dont_canonicalize_unit_avg_pool2d_integer(%arg0: tensor<1x32x32x8xi8>) -> tensor<1x32x32x8xi8> {
+ %input_zp = "tosa.const"() <{values = dense<1> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %output_zp = "tosa.const"() <{values = dense<1> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %0 = tosa.avg_pool2d %arg0, %input_zp, %output_zp {acc_type = i32, kernel = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} :
+ (tensor<1x32x32x8xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<1x32x32x8xi8>
+ return %0 : tensor<1x32x32x8xi8>
+}
+
+// -----
+
+// CHECK-LABEL: @dont_canonicalize_unit_avg_pool2d_non_zero_padding
+// CHECK: tosa.avg_pool2d
+func.func @dont_canonicalize_unit_avg_pool2d_non_zero_padding(%arg0: tensor<1x32x32x8xf32>) -> tensor<1x32x31x8xf32> {
+ %input_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %output_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.avg_pool2d %arg0, %input_zp, %output_zp {acc_type = f32, kernel = array<i64: 2, 2>, pad = array<i64: 1, 0, 0, 0>, stride = array<i64: 1, 1>} :
+ (tensor<1x32x32x8xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x32x31x8xf32>
+ return %0 : tensor<1x32x31x8xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @dont_canonicalize_unit_avg_pool2d_non_unit_stride
+// CHECK: tosa.avg_pool2d
+func.func @dont_canonicalize_unit_avg_pool2d_non_unit_stride(%arg0: tensor<1x33x32x8xf32>) -> tensor<1x17x32x8xf32> {
+ %input_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %output_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.avg_pool2d %arg0, %input_zp, %output_zp {acc_type = f32, kernel = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 2, 1>} :
+ (tensor<1x33x32x8xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x17x32x8xf32>
+ return %0 : tensor<1x17x32x8xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @dont_canonicalize_unit_avg_pool2d_non_unit_kernel
+// CHECK: tosa.avg_pool2d
+func.func @dont_canonicalize_unit_avg_pool2d_non_unit_kernel(%arg0: tensor<1x32x32x8xf32>) -> tensor<1x31x32x8xf32> {
+ %input_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %output_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.avg_pool2d %arg0, %input_zp, %output_zp {acc_type = f32, kernel = array<i64: 2, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} :
+ (tensor<1x32x32x8xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x31x32x8xf32>
+ return %0 : tensor<1x31x32x8xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @dont_canonicalize_unit_avg_pool2d_input_output_mismatch
+// CHECK: tosa.avg_pool2d
+func.func @dont_canonicalize_unit_avg_pool2d_input_output_mismatch(%arg0: tensor<1x32x32x8xf32>) -> tensor<1x?x32x8xf32> {
+ %input_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %output_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.avg_pool2d %arg0, %input_zp, %output_zp {acc_type = f32, kernel = array<i64: 2, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} :
+ (tensor<1x32x32x8xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x?x32x8xf32>
+ return %0 : tensor<1x?x32x8xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @canonicalize_unit_max_pool2d
+// CHECK-NOT: tosa.max_pool2d
+// CHECK: return %arg0 : tensor<1x32x32x8xf32>
+func.func @canonicalize_unit_max_pool2d(%arg0: tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32> {
+ %0 = tosa.max_pool2d %arg0 {kernel = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>, nan_mode = PROPAGATE} :
+ (tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32>
+ return %0 : tensor<1x32x32x8xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @dont_canonicalize_unit_max_pool2d_ignore_nan
+// CHECK: tosa.max_pool2d
+func.func @dont_canonicalize_unit_max_pool2d_ignore_nan(%arg0: tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32> {
+ %0 = tosa.max_pool2d %arg0 {kernel = array<i64: 1, 1>, nan_mode = IGNORE, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} :
+ (tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32>
+ return %0 : tensor<1x32x32x8xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @dont_canonicalize_unit_max_pool2d_non_zero_padding
+// CHECK: tosa.max_pool2d
+func.func @dont_canonicalize_unit_max_pool2d_non_zero_padding(%arg0: tensor<1x32x32x8xf32>) -> tensor<1x32x31x8xf32> {
+ %0 = tosa.max_pool2d %arg0 {kernel = array<i64: 2, 2>, pad = array<i64: 1, 0, 0, 0>, stride = array<i64: 1, 1>} :
+ (tensor<1x32x32x8xf32>) -> tensor<1x32x31x8xf32>
+ return %0 : tensor<1x32x31x8xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @dont_canonicalize_unit_max_pool2d_non_unit_stride
+// CHECK: tosa.max_pool2d
+func.func @dont_canonicalize_unit_max_pool2d_non_unit_stride(%arg0: tensor<1x33x32x8xf32>) -> tensor<1x17x32x8xf32> {
+ %0 = tosa.max_pool2d %arg0 {kernel = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 2, 1>} :
+ (tensor<1x33x32x8xf32>) -> tensor<1x17x32x8xf32>
+ return %0 : tensor<1x17x32x8xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @dont_canonicalize_unit_max_pool2d_non_unit_kernel
+// CHECK: tosa.max_pool2d
+func.func @dont_canonicalize_unit_max_pool2d_non_unit_kernel(%arg0: tensor<1x32x32x8xf32>) -> tensor<1x31x32x8xf32> {
+ %0 = tosa.max_pool2d %arg0 {kernel = array<i64: 2, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} :
+ (tensor<1x32x32x8xf32>) -> tensor<1x31x32x8xf32>
+ return %0 : tensor<1x31x32x8xf32>
+}
+
+// -----
+
// CHECK-LABEL: @canonicalize_max_pool2d_adaptive
-// CHECK: %[[POOL:.+]] = tosa.max_pool2d %arg0 {kernel = array<i64: 1, 1>, nan_mode = IGNORE, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} : (tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32>
+// CHECK: %[[POOL:.+]] = tosa.max_pool2d %arg0 {kernel = array<i64: 2, 1>, nan_mode = IGNORE, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} : (tensor<1x32x32x8xf32>) -> tensor<1x31x32x8xf32>
// CHECK: return %[[POOL]]
-func.func @canonicalize_max_pool2d_adaptive(%arg0: tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32> {
- %kernel = tosa.const_shape {values = dense<[1, 1]> : tensor<2xindex>} : () -> !tosa.shape<2>
+func.func @canonicalize_max_pool2d_adaptive(%arg0: tensor<1x32x32x8xf32>) -> tensor<1x31x32x8xf32> {
+ %kernel = tosa.const_shape {values = dense<[2, 1]> : tensor<2xindex>} : () -> !tosa.shape<2>
%stride = tosa.const_shape {values = dense<[1, 1]> : 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<1x32x32x8xf32>, !tosa.shape<2>, !tosa.shape<2>, !tosa.shape<4>) -> tensor<1x32x32x8xf32>
- return %0 : tensor<1x32x32x8xf32>
+ (tensor<1x32x32x8xf32>, !tosa.shape<2>, !tosa.shape<2>, !tosa.shape<4>) -> tensor<1x31x32x8xf32>
+ return %0 : tensor<1x31x32x8xf32>
}
// -----
>From d6bb566b0b41cc565ef8f10da385fcc2ee720ee8 Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Thu, 18 Jun 2026 13:28:43 +0100
Subject: [PATCH 2/2] address comments
Change-Id: I3139f84cee5606608e725f3b1e82dd6a2455726e
---
mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp | 10 ++++++----
mlir/test/Dialect/Tosa/canonicalize.mlir | 8 +++-----
2 files changed, 9 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 47b703ce5bdfa..ea569d75b6cc0 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -276,15 +276,17 @@ struct AvgPool2dIsNoOp : public OpRewritePattern<tosa::AvgPool2dOp> {
LogicalResult matchAndRewrite(tosa::AvgPool2dOp op,
PatternRewriter &rewriter) const override {
- if (op.getInput().getType() != op.getOutput().getType())
- return rewriter.notifyMatchFailure(
- op, "expected input and output types to match");
-
const auto inputType = llvm::cast<ShapedType>(op.getInput().getType());
if (!llvm::isa<FloatType>(inputType.getElementType()))
return rewriter.notifyMatchFailure(op,
"expected floating-point input type");
+ if (!matchPattern(op.getInputZp(), m_Constant()) ||
+ !matchPattern(op.getOutputZp(), m_Constant()))
+ return rewriter.notifyMatchFailure(
+ op,
+ "expected input and output zero points to be statically verifiable");
+
if (!llvm::all_of(op.getKernel(), [](int64_t val) { return val == 1; }))
return rewriter.notifyMatchFailure(op, "expected unit kernel");
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index bae933d9290e6..f37070f72840e 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -1768,12 +1768,10 @@ func.func @dont_canonicalize_unit_avg_pool2d_non_unit_kernel(%arg0: tensor<1x32x
// -----
-// CHECK-LABEL: @dont_canonicalize_unit_avg_pool2d_input_output_mismatch
+// CHECK-LABEL: @dont_canonicalize_unit_avg_pool2d_dynamic_zp
// CHECK: tosa.avg_pool2d
-func.func @dont_canonicalize_unit_avg_pool2d_input_output_mismatch(%arg0: tensor<1x32x32x8xf32>) -> tensor<1x?x32x8xf32> {
- %input_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
- %output_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
- %0 = tosa.avg_pool2d %arg0, %input_zp, %output_zp {acc_type = f32, kernel = array<i64: 2, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} :
+func.func @dont_canonicalize_unit_avg_pool2d_dynamic_zp(%arg0: tensor<1x32x32x8xf32>, %zp: tensor<1xf32>) -> tensor<1x?x32x8xf32> {
+ %0 = tosa.avg_pool2d %arg0, %zp, %zp {acc_type = f32, kernel = array<i64: 1, 1>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} :
(tensor<1x32x32x8xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x?x32x8xf32>
return %0 : tensor<1x?x32x8xf32>
}
More information about the Mlir-commits
mailing list