fix(reduce): support 64-bit indexing for reductions on views past 2^31 and large arrays

[nicoloangileri]

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure:
• I understand it is strictly prohibited to use AI to write PR description

AI usage disclosure

Code, tests, and verification were developed with Claude Code; the fix and underlying logic were fully reviewed, verified, and understood by me.

Fixes #4539

Problem Description

1. Reductions on views extending past 2^31: Slicing or copying works normally, but running reductions (sum, max, etc.) on a view like x[:, 7:] (which has fewer than 2^31 elements, but whose memory offset reaches beyond 2^31) triggers an indexing error: reading the wrong row on GPU, and producing garbage values or crashing on CPU.
2. CPU full-array reduction over 2^31 elements: Reducing an entire array containing more than 2^31 elements on CPU returns the neutral value (0 for sum and max).

What was changed

• CPU: Switched positions and sizes from int to int64_t in cpu/reduce.cpp.
• GPU: Enabled 64-bit indexing when data_size() > INT32_MAX, matching the logic in conv.cpp and rope.cpp.

Testing

• Added a regression test for reductions on views past 2^31 and large arrays (> 2^31). The test fails before the patch and passes after, on both CPU and GPU.
• The test allocates ~2 GB and is skipped when LOW_MEMORY is set (similar to test_large_binary).
• Full test suite passed on GPU (906 tests).
• Full test suite passed on CPU in a CPU-only build (906 tests).

Out of Scope

• The binary operations bug (mx.dequantize returns garbage on CPU past 2^31 unpacked bits (int size in cpu/binary.h) #4479) is handled separately in PR Fix cpu binary ops on large data #4517.

[zcbenz]

Closing as there was already a PR fixing the issue.

[nicoloangileri]

Hi, thanks so much for the update!

Could you point me to the exact PR you're referring to? I can't seem to find one linked to issue #4539, which is actually still open.

Just to clarify: I noticed that PR #4529 fixes #4528 (regarding negative-stride views, Metal only), while #4539 specifically deals with offsets greater than 2^31 on both CPU and GPU.

If that PR covers this case too, that's completely fine—just let me know which one it is so I can take a look!

[zcbenz]

Ah sorry I was confused with the linked issue.