[AMDGPU] Widen ndarray linear index to int64

[paveltc]

Summary

Backport of ROCm/quadrants@e0ec4416a.
Widens index values to int64 in the LLVM codegen path (ndarray ExternalPtrStmt
linear indexing), makes Ndarray::nelement_ a size_t accumulation, and removes
the obsolete int32-boundary warning logic. This does not make Quadrants globally
int64-indexed; it only prevents int32 overflow in the generated linear-index
arithmetic for large ndarrays.
Cross-backend change (CPU + AMDGPU); conflicts on cherry-pick were purely
clang-format cosmetic and resolved to match the original commit.

Tests

• Result: 5098 passed, 142 skipped, 28 xfailed. The 21 failures are pre-existing
  known-flaky amdgpu block-scan tests (test_block_inclusive,
  test_exclusive_scan_composition) also present in the base branch — no new regressions.
• New CPU test test_ndarray_external_ptr_uses_i64_linear_index passes, verifying
  i64 linear indexing in the generated LLVM IR.

To make the motivation concrete, I added two tests to tests/python/test_ndarray_indexing_i64.py that show why the linear index must be accumulated in i64.

visit(ExternalPtrStmt) flattens an N-D ndarray access into a single linear element index — for a 2-D array of shape (D0, D1) it emits linear = i * D1 + j. Before this PR that accumulation was done in i32 and only sign-extended to i64 for the final GEP, so any array with more than 2**31 elements overflowed before the extend and produced a wrong (often negative) address. The smallest 2-D shape that triggers it is (2, 2**30 + 1), indexed at [1, 2**30], where the true linear index is 2**31 + 1 — just past INT32_MAX.

1. test_i32_linear_index_overflows_but_i64_is_correct (always runs, CPU-only) — reproduces the exact i * D1 + j arithmetic with two kernels, one accumulating in i32 (pre-fix behavior) and one in i64 (the fix), and asserts the i32 version wraps to a negative value while the i64 version stays correct:

   TRUE      = 2147483649
   i32 accum = -2147483647   # wraps negative -> wrong address
   i64 accum = 2147483649    # correct
   

2. test_ndarray_read_past_int32_index_boundary (end-to-end regression guard) — allocates a real >2**31-element int8 ndarray, writes a sentinel at [1, 2**30], and reads it back through the actual ExternalPtr codegen path. On the pre-fix code the wrapped negative offset returns garbage; with this PR it correctly returns the sentinel. It's gated with @pytest.mark.skipif on available memory (needs ~2 GB backing + headroom) so it self-skips on small runners. This test adds a psutil import, consistent with the existing tests/python/test_memory.py.

How to run

# both tests
python3 tests/run_tests.py tests/python/test_ndarray_indexing_i64.py -v

# just the always-on simulation test (no large allocation)
python3 tests/run_tests.py -k test_i32_linear_index_overflows_but_i64_is_correct -v

# force the memory-gated end-to-end test to run (needs ~2GB+ free RAM)
python3 tests/run_tests.py -k test_ndarray_read_past_int32_index_boundary -v

Or directly with pytest:

python3 -m pytest tests/python/test_ndarray_indexing_i64.py -v

Verified against this branch's build: the simulation test passes and the end-to-end read returns the correct sentinel; on the pre-fix codegen the i32 accumulation wraps negative, which is exactly the failure this PR resolves.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 364b86d25b

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Keep product warning for SPIR-V backends

When this runs on Vulkan/Metal, ndarray addressing is still flattened in SPIR-V as an i32 value (quadrants/codegen/spirv/spirv_codegen.cpp:864-894), so a shape like (65536, 65536) has no single dimension above int32 but still wraps the linear offset and reads/writes the wrong element. This branch handles external NumPy/Torch arrays for every arch, so replacing the old product check with only a per-dimension check removes the only warning for unsupported SPIR-V launches; please keep the product warning for non-LLVM backends or widen the SPIR-V path too.

Useful? React with 👍 / 👎.

[hughperkins]

Using a different indexing bitwidth on different platforms conflicts with the ideal that Quadrants code should run unchanged across all platforms.

We could consider globally migrate all indices to int64, but:

• currently there is no obvious end-user need, and
• likely will have a performance impact

[paveltc]

@hughperkins
Thanks — I agree with the principle, but I think this change is narrower than "different indexing bitwidth on different platforms."
It's not a language-level indexing change. Kernel index args keep their declared types. The PR only touches the internal address arithmetic in TaskCodeGenLLVM::visit(ExternalPtrStmt) — how the flat offset into the ndarray buffer is computed.
It's an overflow fix, not a new i64 mode. The final offset was always 64-bit; the old code accumulated in i32 and did a single CreateSExt at the end, so arrays with >2³¹ elements overflowed before the extend. This just moves the widening earlier so the accumulation itself can't overflow.
To make the overflow concrete, I also pushed two tests (tests/python/test_ndarray_indexing_i64.py): a deterministic simulation showing the i32 accumulation wraps negative while i64 stays correct, and a memory-gated end-to-end read of a >2³¹-element ndarray.
Please let me know what kind of fix you would like to see if you disagree with this.

[hughperkins]

The criteria is that correctness should be ~identical across all GPU platforms.

Example of what happens if this is not the case:

1. User writes a script that runs just fine on AMD
   • fails to run correctly on other platforms, such as CUDA or Metal
2. User writes a script that runs just fne on CUDA
   • fails to run correctly on other platforms, such as AMDGPU

To justify changing the i64 for AMDGPU, you would need a platform-agnostic test that, on main, without your changes:

• passes on CUDA
• fails on AMDGPU
  ... and then provide a fix which narrows the correctness delta between AMDGPU and other GPU platforms.