[Mlir-commits] [mlir] [mlir][tosa] Canonicalize avg_pool2d/max_pool2d no-ops (PR #203571)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jun 12 08:47:50 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-tosa
Author: Luke Hutton (lhutton1)
<details>
<summary>Changes</summary>
Removes avg_pool2d/max_pool2d ops with unit kernel, unit stride, and zero padding, as these are no-ops.
---
Full diff: https://github.com/llvm/llvm-project/pull/203571.diff
3 Files Affected:
- (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td (+1)
- (modified) mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp (+42)
- (modified) mlir/test/Dialect/Tosa/canonicalize.mlir (+129-5)
``````````diff
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>
}
// -----
``````````
</details>
https://github.com/llvm/llvm-project/pull/203571
More information about the Mlir-commits
mailing list