[Mlir-commits] [mlir] 126115f - [mlir][tosa] Canonicalize avg_pool2d/max_pool2d no-ops (#203571)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jun 18 08:07:39 PDT 2026
Author: Luke Hutton
Date: 2026-06-18T16:07:33+01:00
New Revision: 126115f92994900cfa7c1e221317922d02278f83
URL: https://github.com/llvm/llvm-project/commit/126115f92994900cfa7c1e221317922d02278f83
DIFF: https://github.com/llvm/llvm-project/commit/126115f92994900cfa7c1e221317922d02278f83.diff
LOG: [mlir][tosa] Canonicalize avg_pool2d/max_pool2d no-ops (#203571)
Removes avg_pool2d/max_pool2d ops with unit kernel, unit stride, and
zero padding, as these are no-ops.
Added:
Modified:
mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
mlir/test/Dialect/Tosa/canonicalize.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index ac98354be9f5e..a333505082b7e 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 20f9d002f5482..b2c6540aa74e2 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -271,6 +271,48 @@ struct AvgPool2dAdaptiveToAvgPool2d
}
};
+struct AvgPool2dIsNoOp : public OpRewritePattern<tosa::AvgPool2dOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tosa::AvgPool2dOp op,
+ PatternRewriter &rewriter) const override {
+ // Prevent canonicalization if input/output shapes don't align
+ 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");
+
+ // For statically known zero points, the verifier ensures zero points are
+ // zero for floating-point types
+ 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");
+
+ 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 +328,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 3333e562e313b..992059e06e541 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -1737,16 +1737,150 @@ 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_non_unit_kernel
+// CHECK: tosa.avg_pool2d
+func.func @dont_canonicalize_unit_avg_pool2d_non_unit_kernel(%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: 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>
+}
+
+// -----
+
+// CHECK-LABEL: @dont_canonicalize_unit_avg_pool2d_dynamic_zp
+// CHECK: tosa.avg_pool2d
+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>
+}
+
+// -----
+
+// 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>
}
// -----
More information about the Mlir-commits
mailing list