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..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)