Widen ndarray linear index to int64

python/quadrants/lang/_func_base.py

@@ -83,6 +83,20 @@ 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 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,
+        _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
@@ -727,14 +741,28 @@ 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.")
+            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. 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; 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
                 if v.flags.c_contiguous:

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<int64>(), 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<int64>(),
+          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<TensorType>()) {
     // 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<TensorType>()) {

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<ArgLoadStmt>();
   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);

quadrants/program/ndarray.cpp

@@ -1,8 +1,11 @@
+#include <limits>
 #include <numeric>
 
+#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,7 +35,7 @@ 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,11 +47,19 @@ 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<int>::max()) {
-    ErrorEmitter(QuadrantsIndexWarning(), &dbg_info,
-                 "Ndarray index might be out of int32 boundary but int64 indexing is "
-                 "not supported yet.");
+  // 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<int>::max()) {
+      ErrorEmitter(QuadrantsIndexWarning(), &dbg_info,
+                   "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);
 }
@@ -63,7 +74,7 @@ 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 +89,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<int>::max()) {
-    ErrorEmitter(QuadrantsIndexWarning(), &dbg_info,
-                 "Ndarray index might be out of int32 boundary but int64 indexing is "
-                 "not supported yet.");
-  }
 }
 
 Ndarray::Ndarray(DeviceAllocation &devalloc,

tests/python/test_ndarray_indexing_i64.py

@@ -0,0 +1,90 @@
+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():
+    @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)