[AMDGPU] Widen ndarray linear index to int64

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:

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/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<int>::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<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)