Skip to content

Fix signed-integer overflow (UB) in roll and tile shape arithmetic - #3604

Merged
zcbenz merged 1 commit into
ml-explore:mainfrom
devYRPauli:fix/shape-arithmetic-overflow
Jun 4, 2026
Merged

zcbenz merged 1 commit into
ml-explore:mainfrom
devYRPauli:fix/shape-arithmetic-overflow

Conversation

@devYRPauli

Copy link
Copy Markdown
Contributor

Summary

Fixes #3601. Several shape ops computed sizes/shifts in 32-bit int without overflow checks, so large-or-degenerate but otherwise valid (int32-range) inputs triggered C++ signed-integer overflow (UB), reachable from the safe public API. UBSan-confirmed sites on current main:

  • tile: reps[i] * shape[i] overflows int (ops.cpp:1352).
  • roll: negating a shift of INT_MIN in (-sh) % size (ops.cpp:6353).
  • roll: summing shifts into an int total_shift accumulator (both the Shape-shift and Shape-shift + axis overloads).

The flatten / Flatten::output_shape product path the issue also mentions is already 64-bit-safe on main, so it is not touched here.

Fix

Compute these in int64_t and narrow with the existing safe_cast helper, which raises a graceful std::overflow_error for out-of-int32 results instead of overflowing. roll's INT_MIN negation is performed in 64-bit so it no longer overflows.

Verification

Built CPU-only with -DUSE_UBSAN=ON:

  • Before: UBSan reports signed integer overflow / negation of -2147483648 cannot be represented at the sites above.
  • After: UBSan is silent on all three; the overflowing cases raise std::overflow_error, and roll by INT_MIN returns the correct (identity) result for a size-4 axis.
  • Full C++ test suite passes (237 cases / 3190 assertions), no regressions.

Added TEST_CASE("roll and tile shape overflow").

roll and tile computed shifts/sizes in 32-bit int without overflow checks, so
large-but-valid int32 inputs hit signed-integer overflow (UB) reachable from the
public API: tile's reps[i]*shape[i], roll's INT_MIN negation, and roll's int
shift-sum accumulator. Compute these in int64 and use safe_cast to reject
out-of-range results with a graceful error. Add a regression test.

@zcbenz zcbenz left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me, thanks!

@zcbenz
zcbenz merged commit 47e27c6 into ml-explore:main Jun 4, 2026
16 checks passed
@devYRPauli

Copy link
Copy Markdown
Contributor Author

Thanks for the review and merge @zcbenz, appreciate the quick turnaround. Glad to get the overflow fix in. Happy to follow up if anything else around the shape arithmetic comes up later.

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.

Signed int-overflow (UB) in shape-op size/rank arithmetic (roll, tile, flatten) for large/degenerate shapes

2 participants