Skip to content

Handle oversized tuple-valued wp.tid() dimensions #1815

Description

@shi-eric

Follow-up to #1802.

wp.tid() returns signed 32-bit thread coordinates. #1802 checks the retained leading extent for kernels that use scalar wp.tid(), but tuple-valued calls still have unhandled overflow cases.

The multidimensional launch bounds representation is relevant here:

template <int N> struct launch_bounds_t {
    int shape[N];
    size_t size;
    size_t coord_mult;
};

launch_coord() uses the non-leading entries in shape for division and modulo. This leaves a few gaps:

  • A kernel that unpacks wp.tid() into two or more coordinates, such as i, j = wp.tid(), but never calls scalar wp.tid(), does not trigger the scalar validation. A leading extent greater than 2**31 can therefore overflow the first returned coordinate.
  • A non-leading extent of exactly 2**31 wraps when stored in shape. For dim=(2, 2**31), linear index 2**31 should produce (1, 0), but the current reconstruction produces (0, -2147483648).
  • When a kernel uses both scalar and tuple-valued wp.tid(), the existing scalar validation only covers the leading coordinate.

The total thread count is separate from the coordinate limits. A grid such as (65536, 65536) has more than 2**31 threads while each coordinate remains representable, so that case should continue to work.

The follow-up should define and enforce the tuple-valued limits consistently across direct, recorded, and JAX launch paths. Possible approaches include validating dimensions before packing them into the current bounds, changing the bounds representation, or combining the two. A larger representation would use more memory, so the implementation choice should remain open. The immediate goal is to prevent unsupported dimensions from silently producing incorrect coordinates.

Metadata

Metadata

Assignees

Labels

bugSomething isn't working

Type

Projects

No projects

Milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions