Skip to content

Replace AVG_POOL2D with REDUCE_SUM in DecomposeMeanDimPass#19242

Open
mcremon-meta wants to merge 1 commit intomainfrom
export-D101418199
Open

Replace AVG_POOL2D with REDUCE_SUM in DecomposeMeanDimPass#19242
mcremon-meta wants to merge 1 commit intomainfrom
export-D101418199

Conversation

@mcremon-meta
Copy link
Copy Markdown
Contributor

@mcremon-meta mcremon-meta commented May 1, 2026

Summary:
Replace the avg_pool2d decomposition path in DecomposeMeanDimPass with
REDUCE_SUM + MUL(1/N) for all mean.dim reductions.

AVG_POOL2D can only pool over spatial (H×W) axes in TOSA/NHWC layout,
which forces the compiler to insert TRANSPOSE ops when the reduction is
over channels (common in LayerNorm). REDUCE_SUM works on any axis without
layout constraints, avoiding those transposes entirely.

Reviewed By: 3l1

Differential Revision: D101418199

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell

Summary:
Replace the avg_pool2d decomposition path in DecomposeMeanDimPass with
REDUCE_SUM + MUL(1/N) for all mean.dim reductions.

AVG_POOL2D can only pool over spatial (H×W) axes in TOSA/NHWC layout,
which forces the compiler to insert TRANSPOSE ops when the reduction is
over channels (common in LayerNorm). REDUCE_SUM works on any axis without
layout constraints, avoiding those transposes entirely.

Reviewed By: 3l1

Differential Revision: D101418199
@mcremon-meta mcremon-meta requested a review from digantdesai as a code owner May 1, 2026 03:51
@pytorch-bot
Copy link
Copy Markdown

pytorch-bot Bot commented May 1, 2026

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/19242

Note: Links to docs will display an error until the docs builds have been completed.

❌ 5 New Failures, 3 Unrelated Failures

As of commit 8dff93f with merge base 94d2881 (image):

NEW FAILURES - The following jobs have failed:

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label May 1, 2026
@meta-codesync
Copy link
Copy Markdown
Contributor

meta-codesync Bot commented May 1, 2026

@mcremon-meta has exported this pull request. If you are a Meta employee, you can view the originating Diff in D101418199.

@github-actions
Copy link
Copy Markdown

github-actions Bot commented May 1, 2026

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

@gggekov
Copy link
Copy Markdown
Collaborator

gggekov commented May 1, 2026

Thanks for the PR, Matthias! I see a few failing tests, could you take a look into them?
You are indeed reducing the number of permutes in a few cases, can you please update the lower number of expected permutes in backends/arm/test/misc/test_transpose_counts.py ?
I am also seeing a failure in backends/arm/test/ops/test_layer_norm.py for the test_native_layer_norm_16a8w_u85_INT suite and for backends/arm/test/ops/test_cond -k test_cond_u85_INT for the one_arg_two_outputs test case, wonder why is that?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/trunk CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. module: arm Issues related to arm backend

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants