diff --git a/python/quadrants/lang/_func_base.py b/python/quadrants/lang/_func_base.py index 6be0c7bb64..4853579800 100644 --- a/python/quadrants/lang/_func_base.py +++ b/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: 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/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 5bc78fd99f..0424ff110e 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,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::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::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::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..047d859cf3 --- /dev/null +++ b/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)