Evaluate JIT arrays past 2^31 elements and sum with compensation on CPU

The CPU JIT evaluation loop counted elements in int, so an array of more than 2^31 elements was left unevaluated and every reduction over it returned 0; the linear index is now dim_t through Node::calc. The all-elements add reduction also lost the low-order bits in one pass (1e9 float ones summed to 2^24), so it now carries a Kahan correction, as #3687 does for the dimension-wise sum. Both halves of #3571 on the CPU backend.
melonakos · Sep 12, 2026 · Sep 12, 2026 · 1ddd9b12bc2a7f801e33fbea15f3968409bd4901
commit 1ddd9b12bc2a7f801e33fbea15f3968409bd4901
diff --git a/src/backend/common/jit/Node.hpp b/src/backend/common/jit/Node.hpp
@@ -199,7 +199,7 @@ class Node {
         UNUSED(w);
     }
 
-    virtual void calc(int idx, int lim) {
+    virtual void calc(dim_t idx, int lim) {
         UNUSED(idx);
         UNUSED(lim);
     }

diff --git a/src/backend/cpu/jit/BinaryNode.hpp b/src/backend/cpu/jit/BinaryNode.hpp
@@ -50,7 +50,7 @@ class BinaryNode : public TNode<compute_t<To>> {
         m_op.eval(this->m_val, lhs->m_val, rhs->m_val, lim);
     }
 
-    void calc(int idx, int lim) final {
+    void calc(dim_t idx, int lim) final {
         UNUSED(idx);
         auto lhs = static_cast<TNode<compute_t<Ti>> *>(m_children[0].get());
         auto rhs = static_cast<TNode<compute_t<Ti>> *>(m_children[1].get());

diff --git a/src/backend/cpu/jit/BufferNode.hpp b/src/backend/cpu/jit/BufferNode.hpp
@@ -85,7 +85,7 @@ class BufferNode : public TNode<T> {
         }
     }
 
-    void calc(int idx, int lim) final {
+    void calc(dim_t idx, int lim) final {
         using Tc = compute_t<T>;
 
         T *in_ptr   = m_ptr + idx;

diff --git a/src/backend/cpu/jit/UnaryNode.hpp b/src/backend/cpu/jit/UnaryNode.hpp
@@ -52,7 +52,7 @@ class UnaryNode : public TNode<To> {
         m_op.eval(TNode<To>::m_val, child->m_val, lim);
     }
 
-    void calc(int idx, int lim) final {
+    void calc(dim_t idx, int lim) final {
         UNUSED(idx);
         auto child = static_cast<TNode<Ti> *>(m_children[0].get());
         m_op.eval(TNode<To>::m_val, child->m_val, lim);

diff --git a/src/backend/cpu/kernel/Array.hpp b/src/backend/cpu/kernel/Array.hpp
@@ -157,11 +157,12 @@ void evalMultiple(std::vector<Param<T>> arrays,
     int num_nodes        = node_clones.size();
     int num_output_nodes = cloned_output_nodes.size();
     if (is_linear) {
-        int num = arrays[0].dims().elements();
-        int cnum =
-            jit::VECTOR_LENGTH * std::ceil(double(num) / jit::VECTOR_LENGTH);
-        for (int i = 0; i < cnum; i += jit::VECTOR_LENGTH) {
-            int lim = std::min(jit::VECTOR_LENGTH, num - i);
+        // dim_t throughout: arrays past 2^31 elements were silently left
+        // unevaluated when this counted in int (#3571)
+        const dim_t num = arrays[0].dims().elements();
+        for (dim_t i = 0; i < num; i += jit::VECTOR_LENGTH) {
+            int lim = static_cast<int>(
+                std::min<dim_t>(jit::VECTOR_LENGTH, num - i));
             for (int n = 0; n < num_nodes; n++) {
                 node_clones[n]->calc(i, lim);
             }

diff --git a/src/backend/cpu/kernel/reduce.hpp b/src/backend/cpu/kernel/reduce.hpp
@@ -199,6 +199,49 @@ struct reduce_all {
     }
 };
 
+// Summing every element in one pass loses the low-order bits once the
+// running total is large (1e9 ones summed as float stops at 2^24), so the
+// add reduction carries a Kahan correction term. See #3571.
+template<typename Ti, typename To>
+struct reduce_all<af_add_t, Ti, To> {
+    common::Transform<data_t<Ti>, compute_t<To>, af_add_t> transform;
+
+    void operator()(Param<To> out, CParam<Ti> in, bool change_nan,
+                    double nanval) {
+        af::dim4 dims            = in.dims();
+        af::dim4 strides         = in.strides();
+        const data_t<Ti> *inPtr  = in.get();
+        data_t<To> *const outPtr = out.get();
+
+        compute_t<To> out_val = common::Binary<compute_t<To>, af_add_t>::init();
+        compute_t<To> correction = compute_t<To>(0);
+
+        for (dim_t l = 0; l < dims[3]; l++) {
+            dim_t off3 = l * strides[3];
+            for (dim_t k = 0; k < dims[2]; k++) {
+                dim_t off2 = k * strides[2];
+                for (dim_t j = 0; j < dims[1]; j++) {
+                    dim_t off1 = j * strides[1];
+                    for (dim_t i = 0; i < dims[0]; i++) {
+                        dim_t idx = i + off1 + off2 + off3;
+
+                        compute_t<To> in_val = transform(inPtr[idx]);
+                        if (change_nan) {
+                            in_val = IS_NAN(in_val) ? nanval : in_val;
+                        }
+                        compute_t<To> y = in_val - correction;
+                        compute_t<To> t = out_val + y;
+                        correction      = (t - out_val) - y;
+                        out_val         = t;
+                    }
+                }
+            }
+        }
+
+        *outPtr = data_t<To>(out_val);
+    }
+};
+
 }  // namespace kernel
 }  // namespace cpu
 }  // namespace arrayfire
diff --git a/test/reduce.cpp b/test/reduce.cpp
@@ -2535,3 +2535,13 @@ TEST(Reduce, SNIPPET_algorithm_func_sum) {
     TEMP_FORMAT_TESTS_allTestByKey(form, anyTrueByKey);
 
 FOREACH_TEMP_FORMAT(TEMP_FORMATS_TESTS)
+
+// A single-precision sum of 1e8 ones stops at 2^24 when accumulated in one
+// pass without compensation (#3571). The result must be exact.
+TEST(Reduce, SumOfManyOnesIsExact_ISSUE_3571) {
+    const dim_t n = 100000000;
+    array ones    = constant(1.f, n);
+    ASSERT_EQ((float)n, af::sum<float>(ones));
+    ASSERT_EQ((double)n, af::sum<double>(ones));
+    ASSERT_NEAR(std::sqrt((double)n), af::norm(ones), 1e-3);
+}