int-reduce-int64-narrowing-overflow: fix sum/prod int64 accumulator narrowing - #45
Closed
arames wants to merge 4 commits into
Closed
int-reduce-int64-narrowing-overflow: fix sum/prod int64 accumulator narrowing#45arames wants to merge 4 commits into
arames wants to merge 4 commits into
Conversation
torch.sum/torch.prod promote narrower integer inputs to an int64 accumulator to avoid overflow, but the reduction lowering cast the accumulator down to int32 first via the shared _NARROW_TORCH_DTYPE table, making the reduce itself wrap at int32 instead of int64. Add get_unnarrowed_output_element_type_from_node, used only by the sum/prod lowerings, to target the true promoted dtype instead.
Cover the dim_IntList+keepdim branch and a negative-value prod case, per review feedback on the int64 accumulator narrowing fix.
Internal issue numbering isn't meaningful outside this repo's own history; keep the docstrings self-contained.
Revert to the pre-existing docstring style for these functions; rationale belongs in the PR description, not inline comments.
arames
marked this pull request as ready for review
July 21, 2026 17:32
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
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.
Summary
torch.sum/torch.prodon a narrower integer input (e.g.int32) promotethe accumulator to
int64per PyTorch's own type-promotion contract,specifically so accumulation doesn't overflow (confirmed:
torch.sum(torch.tensor([2**31-1, 2**31-1], dtype=torch.int32))returns4294967294,dtype=torch.int64, no overflow).replace_sum_dim_intlist/replace_prod_default/replace_prod_dim_int(
coreai_torch/_aten_to_core.py) cast the input to a "target type" beforecalling
coreai.reduce_sum/coreai.reduce_product— but that target typecame from
get_output_element_type_from_node(coreai_torch/_utils.py),which unconditionally narrows
int64toint32via_NARROW_TORCH_DTYPE.That made the reduction itself run (and overflow) in
int32instead ofint64, silently wrapping identically on every backend (interpreter, cpu,gpu, ane) — this corrupts the lowered IR itself, upstream of any backend.
Root cause
_NARROW_TORCH_DTYPE's int64-to-int32 narrowing (_utils.py) is correctfor most call sites (plain value casts where the model owner accepts
truncation), but wrong for an accumulator that PyTorch specifically widens
to avoid overflowing during the operation — narrowing it before the
reduction runs changes the reduction's numeric result, not just its output
representation.
Fix
Added
get_unnarrowed_output_element_type_from_node(coreai_torch/_utils.py),identical to
get_output_element_type_from_nodebut without the_NARROW_TORCH_DTYPEstep, and switched the three sum/prod reductionlowerings to use it.
_NARROW_TORCH_DTYPEandget_output_element_type_from_nodeitself are unchanged — they're relied onby ~15 other lowering call sites (elementwise casts,
cumsum,where,nonzero, etc.) where narrowing is intentional and out of scope here.CoreAI's IR already supports
int64(si64) as a first-class type, so thisisn't a runtime limitation — just a lowering policy fix.
Known related gap, intentionally not fixed here:
replace_cumsumhasthe same accumulator-narrowing shape (
torch.cumsumalso promotes toint64) and is not touched by this PR — keeping this fix scoped tosum/prodas originally triaged. Flagging it here so it isn't lost:cumsum's narrowing appears to additionally crash rather than silentlywrap when pressed with an overflowing input, so it likely needs separate,
dedicated attention.
Test plan
tests/ops/test_ops.pycoveringtorch.sum/torch.prod(default, explicitdim_IntList+keepdim,and
dim_intvariants) onint32inputs that overflowint32but notint64when reduced, assertingdtype == torch.int64and fullnumerical correctness via
validate_numerical_output.ruff check/ruff format/pre-commit runclean.uv run pytest tests/ -n auto -q -p no:rerunfailures -p no:randomlypasses except pre-existing flaky failures in
tests/composite_ops/test_gated_delta_update.py(
@pytest.mark.flaky(reruns=3)), confirmed pre-existing/flaky byreproducing a failure in that same file on a clean
origin/maincheckout.
see linked test dir) now passes on all four backends
(interpreter/cpu/gpu/ane) instead of silently wrapping.
Regression test dir (internal tracking, not part of this repo):
wiki/work/experiments/coreai/tests/regression/010-int-reduce-int64-narrowing-overflow/