Skip to content

[Bugfix] Compute dynamic layout sizes in i64 and name oversized tensor dims - #1236

Open
Phil-amd wants to merge 1 commit into
mainfrom
fix/issue-1176-i32-shape-overflow
Open

Phil-amd wants to merge 1 commit into
mainfrom
fix/issue-1176-i32-shape-overflow

Conversation

@Phil-amd

@Phil-amd Phil-amd commented Oct 8, 2026 •

Copy link
Copy Markdown
Member

What

Fixes a silent wrong-address bug for layout-dynamic tensors whose total size reaches 2^31, and replaces the bare struct.error raised when a single dynamic dimension does not fit in i32. Addresses #1176.

A (4, 2^29) tensor has size 2^31. Each extent fits in i32, but the size was also computed in i32 and wrapped to INT32_MIN. logical_divide uses that size as a divisor when it decomposes a coordinate, so every lane of the kernel below wrote to lane * 16 instead of its layout position. The wrong addresses stay in bounds, so buffer OOB checking never fires and nothing reports an error.

@flyc.kernel(known_block_size=[64, 1, 1])
def scatter(out: fx.Tensor):   # out = storage.as_strided((4, S1), (16, 1))
    lanes = fx.logical_divide(fx.rocdl.make_buffer_tensor(out, max_size=False), fx.make_layout(1, 1))
    ...
    fx.copy(atom, reg, fx.slice(lanes, (None, fx.Int32(fx.thread_idx.x % 64))))
S1 4 * S1 as i32 offsets written on main with this PR
2^29 - 1 2147483644 0..63 0..63
2^29 -2147483648 0, 16, ..., 1008 0..63

pytest tests/unit/test_dynamic_extent_overflow.py -m large_shape reproduces it (about 2 GiB of device memory).

Why not widen all dynamic multiplication to i64

A new intTupleSize computes sizes in i64, and only layoutSize, fly.size and fly.int_tuple_product use it. intTupleProduct stays i32 because it also builds the divisors for crd2idx / idx2crd. Widening it, which was my first attempt, turned per-thread index math over nested dynamic modes into 64-bit divsi/remsi. A FileCheck case now guards against that.

Rejecting launches whose prod(shape) exceeds INT32_MAX would also block valid multi-dimensional tensors. Accepting dimensions of 2^31 or larger would need ABI and compiler changes, as the issue discussion notes. This PR only reports that case clearly.

Key points

  • Composition stays i32. min(extent, size) is bounded by the i32 extent. minExtent truncates it back, so coordinate div/rem are i32. Only arithmetic on the size itself becomes i64: its multiply and, with a non-trivial tiler, one ceil-div by a constant.

  • Visible type change. fx.size of a multi-dimensional dynamic layout is now Int64. Static layouts, 1-D dynamic layouts and cosize are unchanged; with the default 64-bit strides cosize is already i64. test_static_vs_dynamic returned this size from an i32 function and was updated.

  • Launch diagnostic. tensor shape[0] = 2147483648 does not fit in a signed 32-bit integer ... names the field and its index. The check runs only after pack_into has already failed, so a launch that succeeds pays nothing.

  • Not changed. Static IntAttr values and their divisibility are still int32. Two cases can still wrap silently:

    • A nested dynamic mode used as a divisor keeps an i32 product, so a mode whose extents multiply past INT32_MAX decomposes coordinates wrongly. For shape ((4, 2^30+1), 2) with strides ((1, 4), 1), coordinate 63 maps to offset 18 instead of 63.
    • With use_32bit_stride, cosize is computed from i32 terms: cosize((50000, 50000):(50000, 1)) gives -1794967296 instead of 2500000000.

    Both reproduce when the ops are lowered directly. I have not checked whether any kernel reaches them.

@Phil-amd Phil-amd self-assigned this Oct 8, 2026
…r dims

A dynamic layout's size was an i32 product of i32 extents, so (4, 2**29)
wrapped to INT32_MIN. logical_divide used it as a decomposition divisor
and every lane wrote to a wrong but in-bounds address (#1176). Sizes now
go through intTupleSize (i64), so fx.size of a multi-dim dynamic layout is
Int64. intTupleProduct and coordinate div/rem stay i32, so a nested dynamic
mode whose extents multiply past INT32_MAX still wraps. A dynamic dim
>= 2**31 now fails at launch with a message naming the field instead of a
bare struct.error.
@Phil-amd
Phil-amd force-pushed the fix/issue-1176-i32-shape-overflow branch from 80de7eb to 53ac5d7 Compare October 8, 2026 01:01

This branch has not been deployed

No deployments
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