Skip to content

int-reduce-int64-narrowing-overflow: fix sum/prod int64 accumulator narrowing - #45

Closed
arames wants to merge 4 commits into
apple:mainfrom
arames:user/arames/fuzz/int-reduce-int64-narrowing-overflow
Closed

int-reduce-int64-narrowing-overflow: fix sum/prod int64 accumulator narrowing#45
arames wants to merge 4 commits into
apple:mainfrom
arames:user/arames/fuzz/int-reduce-int64-narrowing-overflow

Conversation

@arames

@arames arames commented Jul 18, 2026

Copy link
Copy Markdown

Summary

torch.sum/torch.prod on a narrower integer input (e.g. int32) promote
the accumulator to int64 per 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)) returns
4294967294, 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" before
calling coreai.reduce_sum/coreai.reduce_product — but that target type
came from get_output_element_type_from_node (coreai_torch/_utils.py),
which unconditionally narrows int64 to int32 via _NARROW_TORCH_DTYPE.
That made the reduction itself run (and overflow) in int32 instead of
int64, 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 correct
for 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_node but without the
_NARROW_TORCH_DTYPE step, and switched the three sum/prod reduction
lowerings to use it. _NARROW_TORCH_DTYPE and
get_output_element_type_from_node itself are unchanged — they're relied on
by ~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 this
isn't a runtime limitation — just a lowering policy fix.

Known related gap, intentionally not fixed here: replace_cumsum has
the same accumulator-narrowing shape (torch.cumsum also promotes to
int64) and is not touched by this PR — keeping this fix scoped to
sum/prod as originally triaged. Flagging it here so it isn't lost:
cumsum's narrowing appears to additionally crash rather than silently
wrap when pressed with an overflowing input, so it likely needs separate,
dedicated attention.

Test plan

  • Added regression tests in tests/ops/test_ops.py covering
    torch.sum/torch.prod (default, explicit dim_IntList + keepdim,
    and dim_int variants) on int32 inputs that overflow int32 but not
    int64 when reduced, asserting dtype == torch.int64 and full
    numerical correctness via validate_numerical_output.
  • ruff check/ruff format/pre-commit run clean.
  • Full suite: uv run pytest tests/ -n auto -q -p no:rerunfailures -p no:randomly
    passes except pre-existing flaky failures in
    tests/composite_ops/test_gated_delta_update.py
    (@pytest.mark.flaky(reruns=3)), confirmed pre-existing/flaky by
    reproducing a failure in that same file on a clean origin/main
    checkout.
  • Confirms the previously-filed regression repro (internal tracking;
    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/

arames added 4 commits July 17, 2026 22:31
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
arames marked this pull request as ready for review July 21, 2026 17:32
@arames arames closed this Jul 21, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant