From 364b86d25b82ce0c3b059ea0aa7c033fc78c1e7d Mon Sep 17 00:00:00 2001 From: jamesETsmith Date: Fri, 10 Apr 2026 15:49:54 -0400 Subject: [PATCH 1/5] Widening index values to int64 for llvm codegen, ndarray nelement, and removes some warning logic around int32. It does not make quadrants globally indexed by int64. Co-authored-by: Cursor --- python/quadrants/lang/_func_base.py | 4 ++-- quadrants/codegen/llvm/codegen_llvm.cpp | 21 ++++++++++++------ quadrants/program/ndarray.cpp | 25 +++++++++------------- tests/python/test_ndarray_indexing_i64.py | 26 +++++++++++++++++++++++ 4 files changed, 52 insertions(+), 24 deletions(-) create mode 100644 tests/python/test_ndarray_indexing_i64.py diff --git a/python/quadrants/lang/_func_base.py b/python/quadrants/lang/_func_base.py index 6be0c7bb64..7a9fa29ffe 100644 --- a/python/quadrants/lang/_func_base.py +++ b/python/quadrants/lang/_func_base.py @@ -727,14 +727,14 @@ def _recursive_set_args( # array shapes. is_soa = needed_arg_type.layout == Layout.SOA array_shape = v.shape - if math.prod(array_shape) > np.iinfo(np.int32).max: - warnings.warn("Ndarray index might be out of int32 boundary but int64 indexing is not supported yet.") needed_arg_dtype = needed_arg_type.dtype if needed_arg_dtype is None or id(needed_arg_dtype) in primitive_types.type_ids: element_dim = 0 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): + warnings.warn("Ndarray dimensions above int32 are not supported yet.") if isinstance(v, np.ndarray): # Check ndarray flags is expensive (~250ns), so it is important to order branches according to hit stats if v.flags.c_contiguous: diff --git a/quadrants/codegen/llvm/codegen_llvm.cpp b/quadrants/codegen/llvm/codegen_llvm.cpp index d085c5c6b3..ba72b6e782 100644 --- a/quadrants/codegen/llvm/codegen_llvm.cpp +++ b/quadrants/codegen/llvm/codegen_llvm.cpp @@ -1773,26 +1773,33 @@ void TaskCodeGenLLVM::visit(ExternalPtrStmt *stmt) { int num_array_args = num_indices - num_element_indices; const size_t element_shape_index_offset = num_array_args; + auto *i64_ty = llvm::Type::getInt64Ty(*llvm_context); for (int i = 0; i < num_array_args; i++) { auto raw_arg = builder->CreateGEP( struct_type, llvm_val[stmt->base_ptr], - {tlctx->get_constant(0), tlctx->get_constant(TypeFactory::SHAPE_POS_IN_NDARRAY), tlctx->get_constant(i)}); - raw_arg = builder->CreateLoad(tlctx->get_data_type(PrimitiveType::i32), raw_arg); - sizes[i] = raw_arg; + {tlctx->get_constant(0), + tlctx->get_constant(TypeFactory::SHAPE_POS_IN_NDARRAY), + tlctx->get_constant(i)}); + raw_arg = + builder->CreateLoad(tlctx->get_data_type(PrimitiveType::i32), raw_arg); + sizes[i] = builder->CreateSExt(raw_arg, i64_ty); } - auto linear_index = tlctx->get_constant(0); + auto linear_index = tlctx->get_constant(get_data_type(), 0); size_t size_var_index = 0; for (int i = 0; i < num_indices; i++) { if (i >= element_shape_index_offset && i < element_shape_index_offset + num_element_indices) { // Indexing TensorType-elements - llvm::Value *size_var = tlctx->get_constant(stmt->element_shape[i - element_shape_index_offset]); + llvm::Value *size_var = tlctx->get_constant( + get_data_type(), + stmt->element_shape[i - element_shape_index_offset]); linear_index = builder->CreateMul(linear_index, size_var); } else { // Indexing array dimensions linear_index = builder->CreateMul(linear_index, sizes[size_var_index++]); } - linear_index = builder->CreateAdd(linear_index, llvm_val[stmt->indices[i]]); + auto index = builder->CreateSExtOrBitCast(llvm_val[stmt->indices[i]], i64_ty); + linear_index = builder->CreateAdd(linear_index, index); } QD_ASSERT(size_var_index == num_indices - num_element_indices); @@ -1808,7 +1815,7 @@ void TaskCodeGenLLVM::visit(ExternalPtrStmt *stmt) { if (operand_dtype->is()) { // Access PtrOffset via: base_ptr + offset * sizeof(element) - auto address_offset = builder->CreateSExt(linear_index, llvm::Type::getInt64Ty(*llvm_context)); + auto address_offset = linear_index; auto stmt_ret_type = stmt->ret_type.ptr_removed(); if (stmt_ret_type->is()) { diff --git a/quadrants/program/ndarray.cpp b/quadrants/program/ndarray.cpp index 5bc78fd99f..0e0f82d9f6 100644 --- a/quadrants/program/ndarray.cpp +++ b/quadrants/program/ndarray.cpp @@ -32,7 +32,10 @@ Ndarray::Ndarray(Program *prog, shape(shape_), layout(layout_), dbg_info(dbg_info_), - nelement_(std::accumulate(std::begin(shape_), std::end(shape_), 1, std::multiplies<>())), + nelement_(std::accumulate(std::begin(shape_), + std::end(shape_), + (std::size_t)1, + std::multiplies<>())), element_size_(data_type_size(dtype)), prog_(prog) { // Now that we have two shapes which may be concatenated differently @@ -44,13 +47,8 @@ Ndarray::Ndarray(Program *prog, } else if (layout == ExternalArrayLayout::kSOA) { total_shape_.insert(total_shape_.begin(), element_shape.begin(), element_shape.end()); } - auto total_num_scalar = std::accumulate(std::begin(total_shape_), std::end(total_shape_), 1LL, std::multiplies<>()); - if (total_num_scalar > std::numeric_limits::max()) { - ErrorEmitter(QuadrantsIndexWarning(), &dbg_info, - "Ndarray index might be out of int32 boundary but int64 indexing is " - "not supported yet."); - } - ndarray_alloc_ = prog->allocate_memory_on_device(nelement_ * element_size_, prog->result_buffer); + ndarray_alloc_ = prog->allocate_memory_on_device(nelement_ * element_size_, + prog->result_buffer); } Ndarray::Ndarray(DeviceAllocation &devalloc, @@ -63,7 +61,10 @@ Ndarray::Ndarray(DeviceAllocation &devalloc, shape(shape), layout(layout), dbg_info(dbg_info), - nelement_(std::accumulate(std::begin(shape), std::end(shape), 1, std::multiplies<>())), + nelement_(std::accumulate(std::begin(shape), + std::end(shape), + (std::size_t)1, + std::multiplies<>())), element_size_(data_type_size(dtype)) { // When element_shape is specified but layout is not, default layout is AOS. auto element_shape = data_type_shape(dtype); @@ -78,12 +79,6 @@ Ndarray::Ndarray(DeviceAllocation &devalloc, } else if (layout == ExternalArrayLayout::kSOA) { total_shape_.insert(total_shape_.begin(), element_shape.begin(), element_shape.end()); } - auto total_num_scalar = std::accumulate(std::begin(total_shape_), std::end(total_shape_), 1LL, std::multiplies<>()); - if (total_num_scalar > std::numeric_limits::max()) { - ErrorEmitter(QuadrantsIndexWarning(), &dbg_info, - "Ndarray index might be out of int32 boundary but int64 indexing is " - "not supported yet."); - } } Ndarray::Ndarray(DeviceAllocation &devalloc, diff --git a/tests/python/test_ndarray_indexing_i64.py b/tests/python/test_ndarray_indexing_i64.py new file mode 100644 index 0000000000..5b5cfc06a6 --- /dev/null +++ b/tests/python/test_ndarray_indexing_i64.py @@ -0,0 +1,26 @@ +import re + +import numpy as np + +import quadrants as qd + +from tests import test_utils + + +@test_utils.test(arch=[qd.cpu]) +def test_ndarray_external_ptr_uses_i64_linear_index(): + @qd.kernel + def read(arr: qd.types.NDArray[qd.i32, 2], i: qd.i32, j: qd.i32) -> qd.i32: + return arr[i, j] + + np_arr = np.arange(16, dtype=np.int32).reshape(4, 4) + assert read(np_arr, 1, 2) == np_arr[1, 2] + + compiled = read._primal._last_compiled_kernel_data + assert compiled is not None + + llvm_ir = compiled._debug_dump_to_string() + + assert re.search(r"sext i32 .* to i64", llvm_ir) + assert "mul nsw i64" in llvm_ir + assert re.search(r"getelementptr i32, ptr .* i64 ", llvm_ir) From 1de096c3f472bad885307d43d5f11c8f41f2efa5 Mon Sep 17 00:00:00 2001 From: ptcherni Date: Mon, 13 Jul 2026 15:18:35 -0500 Subject: [PATCH 2/5] Add tests demonstrating int32 linear-index overflow 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 --- tests/python/test_ndarray_indexing_i64.py | 64 +++++++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/tests/python/test_ndarray_indexing_i64.py b/tests/python/test_ndarray_indexing_i64.py index 5b5cfc06a6..047d859cf3 100644 --- a/tests/python/test_ndarray_indexing_i64.py +++ b/tests/python/test_ndarray_indexing_i64.py @@ -1,11 +1,75 @@ import re import numpy as np +import psutil +import pytest import quadrants as qd from tests import test_utils +# ExternalPtrStmt flattens an N-D ndarray access into a single linear element +# index: for a 2-D array of shape (D0, D1) the codegen emits +# linear = i * D1 + j +# (see TaskCodeGenLLVM::visit(ExternalPtrStmt) in codegen/llvm/codegen_llvm.cpp). +# Before this fix 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 pushes the last valid linear index past INT32_MAX: +# shape = (2, 2**30 + 1) -> last index [1, 2**30] -> linear = 2**31 + 1 +_D1 = 2**30 + 1 +_I = 1 +_J = 2**30 +_TRUE_LINEAR = _I * _D1 + _J # == 2**31 + 1, exceeds np.iinfo(np.int32).max + + +@test_utils.test(arch=[qd.cpu]) +def test_i32_linear_index_overflows_but_i64_is_correct(): + # Reproduces the exact arithmetic ExternalPtrStmt performs. The i32 kernel + # mirrors the pre-fix codegen and must wrap to a wrong value; the i64 kernel + # mirrors the fix and must stay correct. + @qd.kernel + def linear_i32(d1: qd.i32, i: qd.i32, j: qd.i32) -> qd.i32: + return i * d1 + j + + @qd.kernel + def linear_i64(d1: qd.i64, i: qd.i64, j: qd.i64) -> qd.i64: + return i * d1 + j + + assert _TRUE_LINEAR > np.iinfo(np.int32).max + + # Pre-fix behavior: i32 accumulation overflows and wraps to a negative offset. + overflowed = linear_i32(_D1, _I, _J) + assert overflowed != _TRUE_LINEAR + assert overflowed < 0 + + # Post-fix behavior: i64 accumulation yields the true linear index. + assert linear_i64(_D1, _I, _J) == _TRUE_LINEAR + + +# ~2 GB for the int8 backing array plus headroom for the device-side copy. +_REQUIRED_BYTES = 5 * 1024**3 + + +@pytest.mark.skipif( + psutil.virtual_memory().available < _REQUIRED_BYTES, + reason="needs >5 GB RAM to allocate an ndarray with more than 2**31 elements", +) +@test_utils.test(arch=[qd.cpu]) +def test_ndarray_read_past_int32_index_boundary(): + # End-to-end regression guard: on the pre-fix codegen the i32 linear index + # for [1, 2**30] wraps negative and this read returns garbage / segfaults. + @qd.kernel + def read(arr: qd.types.NDArray[qd.i8, 2], i: qd.i32, j: qd.i32) -> qd.i8: + return arr[i, j] + + np_arr = np.zeros((2, _D1), dtype=np.int8) + sentinel = np.int8(7) + np_arr[_I, _J] = sentinel + + assert read(np_arr, _I, _J) == sentinel + @test_utils.test(arch=[qd.cpu]) def test_ndarray_external_ptr_uses_i64_linear_index(): From 50e90dccd35670c518e8d84068eda6518917162c Mon Sep 17 00:00:00 2001 From: ptcherni Date: Thu, 16 Jul 2026 17:54:13 -0500 Subject: [PATCH 3/5] [codegen] Warn on int32 linear-index overflow for non-LLVM backends 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. --- python/quadrants/lang/_func_base.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/python/quadrants/lang/_func_base.py b/python/quadrants/lang/_func_base.py index 7a9fa29ffe..52dbdb8eac 100644 --- a/python/quadrants/lang/_func_base.py +++ b/python/quadrants/lang/_func_base.py @@ -83,6 +83,19 @@ def _kernel_coverage_enabled() -> bool: _arch_cuda = _qd_core.Arch.cuda _is_cpython = sys.implementation.name == "cpython" +# 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 +# keep a product-of-dimensions overflow warning for those backends below. +_ARCHS_WITH_I64_LINEAR_INDEX = frozenset( + { + _qd_core.Arch.x64, + _qd_core.Arch.arm64, + _qd_core.Arch.cuda, + _qd_core.Arch.amdgpu, + } +) + # PERF: Frozen-dataclass dispatch caching. # # When a frozen dataclass (e.g. Genesis's StructConstraintState with ~43 fields) is passed to a kernel, the per-launch @@ -735,6 +748,18 @@ def _recursive_set_args( 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): warnings.warn("Ndarray dimensions above int32 are not supported yet.") + elif ( + impl.current_cfg().arch not in _ARCHS_WITH_I64_LINEAR_INDEX + and math.prod(v.shape) > np.iinfo(np.int32).max + ): + # No single dimension overflows int32, but the flattened element count does. Only the LLVM + # backends accumulate the linear index in int64; on the SPIR-V backends this offset is still + # computed in int32 and will overflow, so warn instead of silently reading the wrong element. + warnings.warn( + "Ndarray total element count exceeds the int32 boundary; on this backend the linear index " + "is computed in int32 and may overflow. int64 linear indexing is currently supported only " + "on the LLVM backends (CPU/CUDA/AMDGPU)." + ) if isinstance(v, np.ndarray): # Check ndarray flags is expensive (~250ns), so it is important to order branches according to hit stats if v.flags.c_contiguous: From 50c655683141bf0038fb8039b6005e749b28fe43 Mon Sep 17 00:00:00 2001 From: ptcherni Date: Thu, 16 Jul 2026 23:01:09 -0500 Subject: [PATCH 4/5] [codegen] Warn on int32 ndarray overflow for owned arrays on non-LLVM 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. --- quadrants/program/ndarray.cpp | 29 +++++++++++++++++++---------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/quadrants/program/ndarray.cpp b/quadrants/program/ndarray.cpp index 0e0f82d9f6..db3d90b9da 100644 --- a/quadrants/program/ndarray.cpp +++ b/quadrants/program/ndarray.cpp @@ -1,8 +1,11 @@ +#include #include +#include "quadrants/common/exceptions.h" #include "quadrants/program/adstack_size_expr_eval.h" #include "quadrants/program/ndarray.h" #include "quadrants/program/program.h" +#include "quadrants/rhi/arch.h" #include "fp16.h" #ifdef QD_WITH_LLVM @@ -32,10 +35,7 @@ Ndarray::Ndarray(Program *prog, shape(shape_), layout(layout_), dbg_info(dbg_info_), - nelement_(std::accumulate(std::begin(shape_), - std::end(shape_), - (std::size_t)1, - std::multiplies<>())), + nelement_(std::accumulate(std::begin(shape_), std::end(shape_), (std::size_t)1, std::multiplies<>())), element_size_(data_type_size(dtype)), prog_(prog) { // Now that we have two shapes which may be concatenated differently @@ -47,8 +47,20 @@ Ndarray::Ndarray(Program *prog, } else if (layout == ExternalArrayLayout::kSOA) { total_shape_.insert(total_shape_.begin(), element_shape.begin(), element_shape.end()); } - ndarray_alloc_ = prog->allocate_memory_on_device(nelement_ * element_size_, - prog->result_buffer); + // On non-LLVM backends (SPIR-V: Vulkan/Metal/...) the ndarray linear offset in + // TaskCodegen::visit(ExternalPtrStmt) is still flattened in int32, so an owned ndarray whose total + // element count exceeds int32 will overflow. The LLVM backends (CPU/CUDA/AMDGPU) accumulate the offset + // in int64 and are safe, so the warning is scoped to the non-LLVM backends. + if (!arch_uses_llvm(prog->compile_config().arch)) { + auto total_num_scalar = std::accumulate(std::begin(total_shape_), std::end(total_shape_), 1LL, std::multiplies<>()); + if (total_num_scalar > std::numeric_limits::max()) { + ErrorEmitter(QuadrantsIndexWarning(), &dbg_info, + "Ndarray total element count exceeds the int32 boundary; on this backend the linear index " + "is computed in int32 and may overflow. int64 linear indexing is currently supported only " + "on the LLVM backends (CPU/CUDA/AMDGPU)."); + } + } + ndarray_alloc_ = prog->allocate_memory_on_device(nelement_ * element_size_, prog->result_buffer); } Ndarray::Ndarray(DeviceAllocation &devalloc, @@ -61,10 +73,7 @@ Ndarray::Ndarray(DeviceAllocation &devalloc, shape(shape), layout(layout), dbg_info(dbg_info), - nelement_(std::accumulate(std::begin(shape), - std::end(shape), - (std::size_t)1, - std::multiplies<>())), + nelement_(std::accumulate(std::begin(shape), std::end(shape), (std::size_t)1, std::multiplies<>())), element_size_(data_type_size(dtype)) { // When element_shape is specified but layout is not, default layout is AOS. auto element_shape = data_type_shape(dtype); From 55ab53897896971e3fc786e1d3dbccf5319bc237 Mon Sep 17 00:00:00 2001 From: ptcherni Date: Fri, 17 Jul 2026 16:16:47 -0500 Subject: [PATCH 5/5] [codegen] Widen SPIR-V ndarray linear index to int64 when shaderInt64 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 --- python/quadrants/lang/_func_base.py | 23 +++++++++++-------- quadrants/codegen/spirv/spirv_codegen.cpp | 28 +++++++++++++++-------- quadrants/program/ndarray.cpp | 17 +++++++------- 3 files changed, 41 insertions(+), 27 deletions(-) diff --git a/python/quadrants/lang/_func_base.py b/python/quadrants/lang/_func_base.py index 52dbdb8eac..4853579800 100644 --- a/python/quadrants/lang/_func_base.py +++ b/python/quadrants/lang/_func_base.py @@ -83,10 +83,11 @@ def _kernel_coverage_enabled() -> bool: _arch_cuda = _qd_core.Arch.cuda _is_cpython = sys.implementation.name == "cpython" -# 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 -# keep a product-of-dimensions overflow warning for those backends below. +# The ndarray ExternalPtrStmt linear index is accumulated in int64 in the LLVM codegen backends +# (see TaskCodeGenLLVM::visit(ExternalPtrStmt)) and, on the SPIR-V backends (Vulkan/Metal/...), whenever the +# device advertises 64-bit integers (DeviceCapability.spirv_has_int64 -> TaskCodegen::visit widens to i64). +# Only a SPIR-V device without shaderInt64 still flattens the offset in int32 and can overflow, so the +# product-of-dimensions warning below is gated on both the arch and the device capability. _ARCHS_WITH_I64_LINEAR_INDEX = frozenset( { _qd_core.Arch.x64, @@ -751,14 +752,16 @@ def _recursive_set_args( elif ( impl.current_cfg().arch not in _ARCHS_WITH_I64_LINEAR_INDEX and math.prod(v.shape) > np.iinfo(np.int32).max + and not impl.get_runtime().prog.get_device_caps().get(_qd_core.DeviceCapability.spirv_has_int64) ): - # No single dimension overflows int32, but the flattened element count does. Only the LLVM - # backends accumulate the linear index in int64; on the SPIR-V backends this offset is still - # computed in int32 and will overflow, so warn instead of silently reading the wrong element. + # No single dimension overflows int32, but the flattened element count does. The LLVM backends + # accumulate the linear index in int64, and the SPIR-V backends do too when the device advertises + # shaderInt64. Only a SPIR-V device without shaderInt64 still flattens in int32 and would overflow, + # so warn (this is a cold path -- only huge ndarrays on a non-i64 device reach here). warnings.warn( - "Ndarray total element count exceeds the int32 boundary; on this backend the linear index " - "is computed in int32 and may overflow. int64 linear indexing is currently supported only " - "on the LLVM backends (CPU/CUDA/AMDGPU)." + "Ndarray total element count exceeds the int32 boundary; this device lacks 64-bit integer " + "support (shaderInt64), so the linear index is computed in int32 and may overflow. int64 " + "linear indexing requires the LLVM backends (CPU/CUDA/AMDGPU) or a SPIR-V device with shaderInt64." ) if isinstance(v, np.ndarray): # Check ndarray flags is expensive (~250ns), so it is important to order branches according to hit stats diff --git a/quadrants/codegen/spirv/spirv_codegen.cpp b/quadrants/codegen/spirv/spirv_codegen.cpp index eb509cf464..cd9d9a20ee 100644 --- a/quadrants/codegen/spirv/spirv_codegen.cpp +++ b/quadrants/codegen/spirv/spirv_codegen.cpp @@ -861,7 +861,13 @@ void TaskCodegen::visit(ExternalTensorShapeAlongAxisStmt *stmt) { void TaskCodegen::visit(ExternalPtrStmt *stmt) { // Used mostly for transferring data between host (e.g. numpy array) and // device. - spirv::Value linear_offset = ir_->int_immediate_number(ir_->i32_type(), 0); + // Flatten the multi-dimensional index into a linear (byte) offset. When the device advertises 64-bit integers we + // accumulate in i64 so that ndarrays whose element count exceeds INT32_MAX do not wrap and silently address the + // wrong element (matches the int64 fix already applied on the LLVM backends). Devices without shaderInt64 keep the + // historical i32 arithmetic; the Python/C++ launch path emits an overflow warning for those. + const bool index_use_i64 = caps_->get(DeviceCapability::spirv_has_int64); + const spirv::SType index_type = index_use_i64 ? ir_->i64_type() : ir_->i32_type(); + spirv::Value linear_offset = ir_->int_immediate_number(index_type, 0); const auto *argload = stmt->base_ptr->as(); const auto arg_id = argload->arg_id; { @@ -885,17 +891,18 @@ void TaskCodegen::visit(ExternalPtrStmt *stmt) { spirv::Value size_var; // Use immediate numbers to flatten index for element shapes. if (i >= element_shape_index_offset && i < element_shape_index_offset + element_shape.size()) { - size_var = ir_->uint_immediate_number(ir_->i32_type(), element_shape[i - element_shape_index_offset]); + size_var = ir_->uint_immediate_number(index_type, element_shape[i - element_shape_index_offset]); } else { - size_var = ir_->query_value(size_var_names[size_var_names_idx++]); + // Shapes are stored as i32 in the args buffer; widen to the accumulation type when using i64. + size_var = ir_->cast(index_type, ir_->query_value(size_var_names[size_var_names_idx++])); } - spirv::Value indices = ir_->query_value(stmt->indices[i]->raw_name()); + spirv::Value indices = ir_->cast(index_type, ir_->query_value(stmt->indices[i]->raw_name())); linear_offset = ir_->mul(linear_offset, size_var); linear_offset = ir_->add(linear_offset, indices); } size_t type_size = ir_->get_primitive_type_size(stmt->ret_type.ptr_removed()); - linear_offset = ir_->make_value(spv::OpShiftLeftLogical, ir_->i32_type(), linear_offset, - ir_->int_immediate_number(ir_->i32_type(), log2int(type_size))); + linear_offset = ir_->make_value(spv::OpShiftLeftLogical, index_type, linear_offset, + ir_->int_immediate_number(index_type, log2int(type_size))); if (caps_->get(DeviceCapability::spirv_has_no_integer_wrap_decoration)) { ir_->decorate(spv::OpDecorate, linear_offset, spv::DecorationNoSignedWrap); } @@ -908,7 +915,8 @@ void TaskCodegen::visit(ExternalPtrStmt *stmt) { spirv::Value addr_ptr = ir_->make_access_chain(ir_->get_pointer_type(ir_->u64_type(), spv::StorageClassUniform), get_buffer_value(BufferType::Args, PrimitiveType::i32), indices); spirv::Value base_addr = ir_->load_variable(addr_ptr, ir_->u64_type()); - spirv::Value addr = ir_->add(base_addr, ir_->make_value(spv::OpSConvert, ir_->u64_type(), linear_offset)); + // cast() sign-extends an i32 offset to u64, or bitcasts an already-64-bit i64 offset to u64. + spirv::Value addr = ir_->add(base_addr, ir_->cast(ir_->u64_type(), linear_offset)); ir_->register_value(stmt->raw_name(), addr); // Save decomposed base pointer and element index so at_buffer() can @@ -917,8 +925,10 @@ void TaskCodegen::visit(ExternalPtrStmt *stmt) { // per-element reinterpret_cast from ulong arithmetic is miscompiled // when the stored value is loop-invariant. size_t type_size = ir_->get_primitive_type_size(stmt->ret_type.ptr_removed()); - spirv::Value elem_index = ir_->make_value(spv::OpShiftRightLogical, ir_->i32_type(), linear_offset, - ir_->int_immediate_number(ir_->i32_type(), log2int(type_size))); + // Keep the element index in the same width as the offset accumulation so OpPtrAccessChain indexes the correct + // element for >INT32_MAX-element ndarrays on int64-capable devices. + spirv::Value elem_index = ir_->make_value(spv::OpShiftRightLogical, index_type, linear_offset, + ir_->int_immediate_number(index_type, log2int(type_size))); physical_ptr_components_[stmt] = {base_addr, elem_index}; } else { ir_->register_value(stmt->raw_name(), linear_offset); diff --git a/quadrants/program/ndarray.cpp b/quadrants/program/ndarray.cpp index db3d90b9da..0424ff110e 100644 --- a/quadrants/program/ndarray.cpp +++ b/quadrants/program/ndarray.cpp @@ -47,17 +47,18 @@ Ndarray::Ndarray(Program *prog, } else if (layout == ExternalArrayLayout::kSOA) { total_shape_.insert(total_shape_.begin(), element_shape.begin(), element_shape.end()); } - // On non-LLVM backends (SPIR-V: Vulkan/Metal/...) the ndarray linear offset in - // TaskCodegen::visit(ExternalPtrStmt) is still flattened in int32, so an owned ndarray whose total - // element count exceeds int32 will overflow. The LLVM backends (CPU/CUDA/AMDGPU) accumulate the offset - // in int64 and are safe, so the warning is scoped to the non-LLVM backends. - if (!arch_uses_llvm(prog->compile_config().arch)) { + // The ndarray linear offset in TaskCodegen::visit(ExternalPtrStmt) is flattened in int64 on the LLVM backends + // (CPU/CUDA/AMDGPU) and, on the SPIR-V backends (Vulkan/Metal/...), whenever the device advertises 64-bit + // integers (DeviceCapability::spirv_has_int64). Only a SPIR-V device *without* shaderInt64 still flattens in + // int32, so an owned ndarray whose total element count exceeds int32 would overflow there. Scope the warning to + // exactly that case. + if (!arch_uses_llvm(prog->compile_config().arch) && !prog->get_device_caps().get(DeviceCapability::spirv_has_int64)) { auto total_num_scalar = std::accumulate(std::begin(total_shape_), std::end(total_shape_), 1LL, std::multiplies<>()); if (total_num_scalar > std::numeric_limits::max()) { ErrorEmitter(QuadrantsIndexWarning(), &dbg_info, - "Ndarray total element count exceeds the int32 boundary; on this backend the linear index " - "is computed in int32 and may overflow. int64 linear indexing is currently supported only " - "on the LLVM backends (CPU/CUDA/AMDGPU)."); + "Ndarray total element count exceeds the int32 boundary; this device lacks 64-bit integer " + "support (shaderInt64), so the linear index is computed in int32 and may overflow. int64 linear " + "indexing requires the LLVM backends (CPU/CUDA/AMDGPU) or a SPIR-V device with shaderInt64."); } } ndarray_alloc_ = prog->allocate_memory_on_device(nelement_ * element_size_, prog->result_buffer);