Skip to content

portable: add fast path for contiguous innermost-dim reduction in amax.out and amin.out#21142

Open
vacu9708 wants to merge 1 commit into
pytorch:mainfrom
vacu9708:amax_fast_path
Open

portable: add fast path for contiguous innermost-dim reduction in amax.out and amin.out#21142
vacu9708 wants to merge 1 commit into
pytorch:mainfrom
vacu9708:amax_fast_path

Conversation

@vacu9708

Copy link
Copy Markdown
Contributor

Motivation

sum.IntList_out, mean.out, and var_mean.correction_out special-case the contiguous, single innermost-dim reduction shape with a direct pointer loop instead of the generic ReduceOverDimListPlan, which recomputes a stride-based starting index per output element via get_init_index (#18789, #18790, #18791). amax.out and amin.out do not have this fast path.

Tests

No new tests were added. The existing tests already cover both fast and generic paths.
sum's (#18789), mean's (#18790), and var_mean's (#18791) fast-path PRs had no test changes either.

Test Shape dim Fast path? Dtypes covered
AllRealInputOutputPasses {2,3,4} {2} (innermost) Yes ET_FORALL_REALHBBF16_TYPES (incl. Bool, Half, BFloat16)
AllRealInputOutputPasses {2,3,4} {-2} (not innermost) No, fallback same
AllRealInputOutputPasses {2,3,4} {0,1} (multi-dim) No, fallback same
AllRealInputOutputPasses {2,2,4} null / empty No, fallback same
InfinityAndNANTest {2,3,4} {-1} (innermost) Yes ET_FORALL_FLOATHBF16_TYPES

Same coverage in both op_amax_test.cpp and op_amin_test.cpp.

$ ninja -C cmake-out-bench portable_kernels_test -j16
[6/6] Linking CXX executable kernels/test/portable_kernels_test

$ ./cmake-out-bench/kernels/test/portable_kernels_test
[==========] Running 1625 tests from 190 test suites.
[  PASSED  ] 1521 tests.
# 0 FAILED, 209 SKIPPED (pre-existing DISABLED/Aten-only tests, unrelated to this change)

$ ./cmake-out-bench/kernels/test/portable_kernels_test --gtest_filter="OpAmaxOutTest.*:OpAminOutTest.*"
[==========] 10 tests from 2 test suites ran.
[  PASSED  ] 10 tests.

$ lintrunner -a kernels/portable/cpu/op_amax.cpp kernels/portable/cpu/op_amin.cpp
ok No lint issues.

….out and amin.out

Adds a fast path to amax.out and amin.out: when the input tensor is
contiguous and reduced over just its innermost dimension, compute the
result with a direct loop instead of the generic path's per-row index
recomputation. Output is unchanged, and existing tests already cover
both the fast path and the fallback.

Continues the reduction fast-path work in pytorch#18789, pytorch#18790, pytorch#18791.

Signed-off-by: Youngsik Yang <vacu9708@gmail.com>
@vacu9708
vacu9708 requested a review from manuelcandales as a code owner July 22, 2026 14:11
@pytorch-bot

pytorch-bot Bot commented Jul 22, 2026

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

⚠️ 12 Awaiting Approval

As of commit 618eaf7 with merge base 666ba0b (image):

AWAITING APPROVAL - The following workflows need approval before CI can run:

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 Jul 22, 2026
@github-actions

Copy link
Copy Markdown

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.

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

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants