Replace AVG_POOL2D with REDUCE_SUM in DecomposeMeanDimPass#19242
Replace AVG_POOL2D with REDUCE_SUM in DecomposeMeanDimPass#19242mcremon-meta wants to merge 1 commit intomainfrom
Conversation
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
🔗 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 FailuresAs of commit 8dff93f with merge base 94d2881 ( 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. |
|
@mcremon-meta has exported this pull request. If you are a Meta employee, you can view the originating Diff in D101418199. |
This PR needs a
|
|
Thanks for the PR, Matthias! I see a few failing tests, could you take a look into them? |
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