Widen ndarray linear index to int64#776
Conversation
…d removes some warning logic around int32. It does not make quadrants globally indexed by int64. Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
💡 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".
| else: | ||
| element_dim = needed_arg_dtype.ndim | ||
| array_shape = v.shape[element_dim:] if is_soa else v.shape[:-element_dim] | ||
| if any(dim > np.iinfo(np.int32).max for dim in array_shape): |
There was a problem hiding this comment.
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
left a comment
There was a problem hiding this comment.
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
Adds a deterministic simulation of the ExternalPtrStmt linear-index arithmetic (i * D1 + j) showing i32 accumulation wraps negative while i64 stays correct, plus a memory-gated end-to-end test that reads an ndarray with more than 2**31 elements at the overflow-triggering index. Co-authored-by: Cursor <cursoragent@cursor.com>
|
@hughperkins |
|
The criteria is that correctness should be ~identical across all GPU platforms. Example of what happens if this is not the case:
To justify changing the i64 for AMDGPU, you would need a platform-agnostic test that, on main, without your changes:
|
Restore a product-of-dimensions overflow warning for the SPIR-V backends (Vulkan/Metal/...), which still flatten the ndarray ExternalPtr linear offset in int32. The int64 accumulation fix only applies to the LLVM codegen backends (CPU/CUDA/AMDGPU), so without this those backends would silently read the wrong element for ndarrays whose element count exceeds int32. Addresses the Codex P2 review comment.
|
Thanks @hughperkins. Digging into where the change actually lives, I think the framing shifts a bit: The fix is in the shared LLVM codegen — TaskCodeGenLLVM::visit(ExternalPtrStmt) in codegen/llvm/codegen_llvm.cpp — which is used identically by CPU, CUDA, and AMDGPU. On main, that path accumulates the linear index in i32 and only SExts to i64 at the end, so CUDA overflows exactly the same way AMDGPU does. Because of that, the test you asked for (on main: passes on CUDA, fails on AMDGPU) can't really exist — CUDA isn't correct here either. This isn't an AMD-vs-CUDA bitwidth divergence; it's a universal overflow in the LLVM path, and the fix makes CPU/CUDA/AMDGPU consistent. (The [AMDGPU] tag is misleading — I'll retitle.) The real cross-platform delta your principle is pointing at is LLVM backends vs SPIR-V backends: Vulkan/Metal (spirv_codegen.cpp) still flatten the offset in i32 and would silently overflow. To keep correctness parity, I've restored a product-of-dimensions overflow warning for the SPIR-V backends (also addressing Codex's P2), so no backend silently returns the wrong element. Widening the SPIR-V path to i64 is a larger follow-up (needs the Int64 device capability and has perf implications), which I'm happy to do separately if you'd prefer full parity now. |
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 50e90dccd3
ℹ️ 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".
| ndarray_alloc_ = prog->allocate_memory_on_device(nelement_ * element_size_, | ||
| prog->result_buffer); |
There was a problem hiding this comment.
Preserve overflow warning for owned ndarrays
On Vulkan/Metal, TaskCodegen::visit(ExternalPtrStmt) still flattens ndarray offsets in i32, but this constructor no longer emits the old total_shape_ overflow warning. The new Python warning only runs when launching numpy/torch external arrays; qd.ndarray(...) objects take the earlier Ndarray branch in _recursive_set_args and are passed through set_arg_ndarray_impl, so a shape like (2, 2**30 + 1) can now be created and indexed on SPIR-V without any warning even though the generated offset wraps. Please keep a backend-specific warning for Quadrants-owned ndarrays as well.
Useful? React with 👍 / 👎.
… backends The previous push restored the overflow warning only for external numpy/torch arrays (the Python launch path). Quadrants-owned qd.ndarray() objects are constructed via Ndarray(Program*, ...) and bypass that check, so on SPIR-V backends (Vulkan/Metal) an owned ndarray whose element count exceeds int32 could be created and indexed with no warning even though the generated linear offset wraps. Restore the total_shape_ product-of-dimensions overflow warning in that constructor, gated on !arch_uses_llvm() so it only fires on the non-LLVM backends (the LLVM backends now accumulate the offset in int64). Addresses the follow-up Codex P2 review comment.
|
@codex review |
|
Codex Review: Didn't find any major issues. Already looking forward to the next diff. Reviewed commit: ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
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". |
|
Ok, so it sounds like this PR is not addressing a difference in correctness between back-ends? Could you give me a little more background/details into the motivations of this PR? Like, sure, it's possible to construct an artificial test that fails without this PR, but what is your actual underlying motivations/intentions? In genesis-world for example, we are not hitting this limitation, because physics simulations get slow well before we hit such tensor sizes. Things that I'm wary about:
|
|
You're right that this isn't a cross-backend correctness delta, and you're right that genesis-world doesn't hit it. Physics sims get slow long before a single ndarray exceeds ~2.1B elements, so for the primary workload this limit is effectively unreachable. I don't want to oversell it as something Genesis needs today. The concrete benefit is a latent memory-safety bug, not "large-tensor support." There are actually two failure modes on main for an ndarray whose element count exceeds INT32_MAX, and both are silent:
Please let me know if you have any more questions. |
|
Yes, the fact that it is silent is a great point. |
|
I'll try running Genesis benchmarks against this branch. |
|
@codex review |
|
|
||
| # The ndarray ExternalPtrStmt linear index is accumulated in int64 only in the LLVM codegen backends | ||
| # (see TaskCodeGenLLVM::visit(ExternalPtrStmt)). The SPIR-V backends (Vulkan/Metal/...) still flatten the | ||
| # offset in int32, so an ndarray whose total element count exceeds int32 can silently overflow there. We |
There was a problem hiding this comment.
this states that spirv is still broken.
strong preference that we fix across all backends please.
There was a problem hiding this comment.
(or alterntaively, put an overflow warnign for all backends)
There was a problem hiding this comment.
Follow-up: widen the SPIR-V linear index to int64 (commit 55ab538)
This extends the fix to the SPIR-V backends (Vulkan/Metal), closing the cross-backend parity gap.
Change. TaskCodegen::visit(ExternalPtrStmt) previously flattened the multi-dimensional ndarray index into a linear offset in int32, so on SPIR-V an ndarray with more than INT32_MAX elements would wrap and silently address the wrong element — the same overflow the earlier commits fixed for the LLVM backends. It now accumulates the offset (and the derived byte offset / element index) in int64 when the device advertises DeviceCapability::spirv_has_int64, feeding the 64-bit offset directly into the physical-storage-buffer address add and OpPtrAccessChain. Devices without shaderInt64 keep the historical int32 arithmetic.
Warnings. The C++ Ndarray(Program*) constructor and the Python external-array launch path now only warn on a SPIR-V device that actually lacks shaderInt64 — every other case (LLVM backends, or SPIR-V with int64) now indexes correctly, so the diagnostic is no longer misleading.
What I validated locally
Compiles in a Vulkan-enabled build (QD_WITH_VULKAN=ON, alongside AMDGPU + CUDA).
Ran on a software Vulkan device (Mesa lavapipe, which reports shaderInt64 + bufferDeviceAddress): multi-dimensional qd.ndarray indexing matches NumPy, and an int64 reduction over the array is correct.
Confirmed via SPIR-V disassembly that the generated offset math is now 64-bit (OpSConvert/OpIMul/OpIAdd/OpShiftLeftLogical on %long, then OpBitcast to %ulong), versus the previous int32 + OpSConvert path.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: af6db1987f
ℹ️ 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".
| auto index = builder->CreateSExtOrBitCast(llvm_val[stmt->indices[i]], i64_ty); | ||
| linear_index = builder->CreateAdd(linear_index, index); |
There was a problem hiding this comment.
Document backend-specific large-array indexing
This change expands the user-visible ndarray indexing range beyond INT32_MAX on LLVM backends while leaving Vulkan/Metal limited to int32 linear offsets, but it makes no corresponding change under docs/. The repository's AGENTS.md explicitly requires documentation for public API or usage changes; add end-user documentation describing the newly supported large-array behavior and the backend-specific limitation.
Useful? React with 👍 / 👎.
| # array shapes. | ||
| is_soa = needed_arg_type.layout == Layout.SOA | ||
| array_shape = v.shape | ||
| if math.prod(array_shape) > np.iinfo(np.int32).max: |
There was a problem hiding this comment.
wait, there was already a warning?
There was a problem hiding this comment.
Yes — main had an unconditional one here: warnings.warn("...int64 indexing is not supported yet."). But it was only a warning (the offset still overflowed and returned the wrong element), and its message is now false for CPU/CUDA/AMDGPU since those accumulate in int64. Keeping it as-is would warn on exactly the backends this PR fixes. So I kept a warning but scoped it to the backends that still flatten in int32 — and with the new SPIR-V int64 commit, only to devices without shaderInt64.
There was a problem hiding this comment.
I would prefer to just keep the warning as-is, without changing anything, without a strong reason.
What is your motivation for creating this PR? I'm sensing the motivation is primarily 'this is broken and it should be fixed'?
There was a problem hiding this comment.
(I'm happy to upgrade from warning to throwing an exception, if that would make you more comfortable?)
There was a problem hiding this comment.
Here is the motivation from a colleague of mine who worked on this:
Things that I'm wary about:
changes without a clear concrete benefit, since any change involves risk, of adding new bugs, and of performance regressions
They benefit was that they enabled workloads with more than 32k envs (which we explicitly needed)performance regressions, since 64-bit types tend to be slower than 32-bit types, e.g. because use more memory bandwidth
Sure 64 bit types are slow, but we explicitly don't switch to using int64 for indexes, we narrowed the scope and used it for things like flattened indices which show up a lot and can easily exceed the bounds of int32correctness regressions, and crash bugs, which changing bit-widths can definitely both create, especially since we'd need to change this across all platforms
The hope was that the limited scope of these changes would also reduce or eliminate these types of problems.
There was a problem hiding this comment.
Quick update, and a correction to my earlier comment. When I said genesis-world doesn't hit this — that physics slows down before an ndarray crosses ~2.1B elements — I was reasoning from my own assumptions. Since then a colleague who worked on the original change filled me in: they did hit it in practice, on a real workload running >32k parallel envs, where the aggregate ndarray crosses INT32_MAX even though the per-env tensors are modest. So there is a concrete, non-artificial workload behind this, not just the memory-safety argument — I was wrong to imply nobody needed the capability.
That also speaks to the warning-vs-fix question. For that workload, neither the existing warning nor upgrading it to an exception helps — it needs to run at that size, and that requires the flattened-offset arithmetic not to overflow. I'm fully on board with your instinct that silent is the worst case: making overflow throw instead of warn is strictly better where we can't support the size (e.g. a device without shaderInt64). But where we can support it, the fix lets the workload actually execute rather than failing loudly.
So the two motivations are complementary:
- Concrete capability — a real >32k-env workload that overflows INT32_MAX on main and works with this PR.
- Latent memory-safety — the silent wrong-address / under-allocation bug in the shared LLVM path (the point you flagged as a good one), which the same change closes for CPU/CUDA/AMDGPU.
And the SPIR-V commit (55ab538) extends the same narrow widening to Vulkan/Metal (capability-gated on shaderInt64, i32 + warning fallback otherwise), which is the cross-backend parity you asked for — so no backend is left silently overflowing.
There was a problem hiding this comment.
hmmm, this is ndarray only. what happens if we change between ndarray and field? (we do this in genesis-world for example: ndarray for development mode, compiles quickly; field for performance mode, runs quickly, but compiles slowly). Will this introduce a correctness disparity between ndarray and field?
|
Wait, looking a the code, it looks like there is already a warning, in main? |
… is available TaskCodegen::visit(ExternalPtrStmt) flattened the multi-dimensional ndarray index into a linear offset in int32, so on the SPIR-V backends (Vulkan/Metal) an ndarray whose element count exceeds int32 would wrap and silently address the wrong element -- the same overflow the earlier commits fixed for the LLVM backends (CPU/CUDA/AMDGPU). Accumulate the linear offset (and the derived byte offset / element index) in int64 when the device advertises DeviceCapability::spirv_has_int64. Sizes and indices are widened with IRBuilder::cast (sign-extend), and the physical-storage -buffer add/OpPtrAccessChain now consume the 64-bit offset directly (cast() bitcasts i64->u64 and sign-extends i32->u64). Devices without shaderInt64 keep the historical int32 arithmetic. Re-scope the product-of-dimensions overflow warnings accordingly: the C++ Ndarray(Program*) constructor and the Python external-array launch path now only warn on a SPIR-V device that actually lacks shaderInt64, since every other case (LLVM backends, or SPIR-V with int64) now indexes correctly. Verified in a Vulkan-enabled build on a software Vulkan device (Mesa lavapipe, shaderInt64+bufferDeviceAddress): multi-dimensional ndarray indexing matches numpy, and the generated SPIR-V computes the linear offset in 64-bit (OpIMul/ OpIAdd/OpShiftLeftLogical on %long, OpBitcast to %ulong) instead of the previous int32 + OpSConvert. Real AMD/NVIDIA-GPU Vulkan perf and Metal are left to CI and hardware owners. Co-authored-by: Cursor <cursoragent@cursor.com>
|
@codex review |
|
Codex Review: Didn't find any major issues. Chef's kiss. Reviewed commit: ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
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". |
|
Oh this is for ndarray, and I ran for field. I'll run for ndarray. |
|
(Note: how to read the benchmark report:
|


Summary
Backport of ROCm/quadrants@e0ec4416a.
Widens index values to int64 in the LLVM codegen path (ndarray
ExternalPtrStmtlinear indexing), makes
Ndarray::nelement_asize_taccumulation, and removesthe 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
known-flaky amdgpu block-scan tests (
test_block_inclusive,test_exclusive_scan_composition) also present in the base branch — no new regressions.test_ndarray_external_ptr_uses_i64_linear_indexpasses, verifyingi64 linear indexing in the generated LLVM IR.
To make the motivation concrete, I added two tests to
tests/python/test_ndarray_indexing_i64.pythat show why the linear index must be accumulated ini64.visit(ExternalPtrStmt)flattens an N-D ndarray access into a single linear element index — for a 2-D array of shape(D0, D1)it emitslinear = i * D1 + j. Before this PR that accumulation was done ini32and only sign-extended toi64for the final GEP, so any array with more than2**31elements 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 is2**31 + 1— just pastINT32_MAX.test_i32_linear_index_overflows_but_i64_is_correct(always runs, CPU-only) — reproduces the exacti * D1 + jarithmetic with two kernels, one accumulating ini32(pre-fix behavior) and one ini64(the fix), and asserts the i32 version wraps to a negative value while the i64 version stays correct:test_ndarray_read_past_int32_index_boundary(end-to-end regression guard) — allocates a real>2**31-elementint8ndarray, writes a sentinel at[1, 2**30], and reads it back through the actualExternalPtrcodegen 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.skipifon available memory (needs ~2 GB backing + headroom) so it self-skips on small runners. This test adds apsutilimport, consistent with the existingtests/python/test_memory.py.How to run
Or directly with pytest:
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.