Fix full reduction reading outside non-contiguous input views - #5439
Open
zheng1 wants to merge 2 commits into
Open
Fix full reduction reading outside non-contiguous input views#5439zheng1 wants to merge 2 commits into
zheng1 wants to merge 2 commits into
Conversation
The flattened reduction kernels address the input linearly, but the Python side never made the input contiguous. For a view whose storage still holds elements outside it, such as base[:, :3] or base[::2], the kernel read those discarded elements and returned a wrong result. Reproduced on Ascend 910B for min, amax, amin, prod, mean, argmax and argmin. prod and amin run the shared implementation, so this is not backend specific. The same one line fix is already used for sum and mean in flagos-ai#1786 and for max in flagos-ai#273. amin_ writes in place, so it reads through a contiguous copy and keeps writing back to the original tensor. Adds a full reduction regression test per operator. The values outside the view are chosen so that reading them changes the result, rather than relying on random data to expose it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
zheng1
requested review from
0x45f,
103yiran,
Caeruleann,
bin913,
douxetpur,
hellojack163,
huangyiqun,
tengqm and
w1120029931-bit
as code owners
August 13, 2026 04:45
3 tasks
Two issues on the dim=None path: 1. keepdim. The output was allocated with torch.empty([]) unconditionally, so keepdim=True was dropped. PyTorch returns shape [1] * inp.dim(). 2. Block size. block_size was derived only from sqrt(numel) with no upper bound. For large inputs it went past what the device handles on this path and the returned index pointed at an element that was not the maximum. It is now capped at 1024, the same bound argmin uses. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
zheng1
force-pushed
the
fix/ascend-reduction-noncontiguous
branch
from
August 13, 2026 09:49
959bc89 to
f33d8b0
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR Category
Operator
Type of Change
Bug Fix
Description
The flattened reduction kernels address the input linearly (
inp + offset), butthe Python side never made the input contiguous. For a view whose storage still
holds elements outside it, such as
base[:, :3]orbase[::2], the kernel readsthose discarded elements and returns a wrong result.
Reproduced on Ascend 910B for
min,amax,amin,prod,mean,argmaxandargmin.prodandaminrun the shared implementation undersrc/flag_gems/ops/, so this is not backend specific. The bug is a mismatchbetween linear addressing of storage and the view's strides, which is
independent of the hardware.
The same one line fix already exists elsewhere in the tree for the same reason:
sumandmeanin #1786, andmaxin #273. Some operators in the same familygot it and some did not, so this fills in the rest.
amin_writes in place, so it reads through a contiguous copy and keeps writingback to the original tensor rather than rebinding the name.
Each operator gets a full reduction regression test. The values outside the view
are chosen so that reading them changes the result, instead of relying on random
data to expose the bug.
Issue
None.
Progress
Performance
No performance impact.
contiguous()returns the same tensor when the input isalready contiguous, which is the case on every existing code path, so no copy is
introduced. A copy happens only for inputs that previously produced a wrong
result.
Verified on Ascend 910B against a pure master checkout on the same device: the
20 new regression cases go from failing to passing, and the overall failure count
for the touched suites drops from 144 to 65. The remaining difference is
pre-existing flakiness in
argmaxandargmintie breaking, which reproduces atthe same rate on master.