Skip to content

[CPU] Make ReduceSum exact for integer types #33599

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
shibaacharya wants to merge 3 commits into
base: main
Choose a base branch
from
+127 −71
Open

[CPU] Make ReduceSum exact for integer types #33599

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7ab95b5
Make CPU ReduceSum exact for integer types
shibaacharya Oct 26, 2026
d3d2bb0
Address review: CPU-only ReduceSum int64 tests covering all fast paths
shibaacharya Oct 26, 2026
4959c5d
Add scalar-path ReduceSum int64 test
shibaacharya Oct 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
146 changes: 75 additions & 71 deletions onnxruntime/core/providers/cpu/reduction/reduction_ops.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,43 @@ inline bool reduce_isnan<int64_t>(int64_t) { return false; }
template <typename T>
inline constexpr bool kReduceUseKahan = std::is_same_v<T, int64_t> || std::is_same_v<T, uint64_t>;

// Exact 128-bit integer accumulator (two's complement hi:lo) used by ReduceSum for integral types.
// Unlike double it keeps every bit of int64/uint64 inputs, and intermediate partial sums outside
// the range of T cannot corrupt the result. Only the final value is saturated to T.
template <typename T>
struct ExactIntSum {
static_assert(std::is_integral_v<T> && sizeof(T) <= sizeof(uint64_t));
uint64_t lo = 0;
uint64_t hi = 0; // unsigned so that wrap-around is well defined

inline void add(T v) {
const uint64_t u = static_cast<uint64_t>(v);
lo += u;
hi += (lo < u ? 1 : 0);
if constexpr (std::is_signed_v<T>) {
if (v < 0) hi -= 1;
}
}

inline T get() const {
constexpr T t_max = std::numeric_limits<T>::max();
constexpr T t_min = std::numeric_limits<T>::min();
if (static_cast<int64_t>(hi) < 0) {
if constexpr (std::is_signed_v<T>) {
const int64_t low = static_cast<int64_t>(lo);
if (hi == ~uint64_t{0} && low < 0 && low >= static_cast<int64_t>(t_min)) {
return static_cast<T>(low);
}
}
return t_min;
}
if (hi == 0 && lo <= static_cast<uint64_t>(t_max)) {
return static_cast<T>(lo);
}
return t_max;
}
};

class ReduceAggregatorBase {
public:
// Fast reduction: see OptimizeShapeForFastReduce's comment.
Expand Down Expand Up @@ -221,54 +258,30 @@ class ReduceAggregator : public ReduceAggregatorBase {
template <typename T>
class ReduceAggregatorSum : public ReduceAggregator<T, T> {
protected:
// For integer types, accumulate in double to avoid signed overflow UB.
// When overflow occurs, results saturate to T::max()/T::lowest() rather than wrapping.
// Integer sums are accumulated exactly (see ExactIntSum) to avoid signed overflow UB and the
// precision loss of double for int64 values > 2^53. When the final result is outside the range of T
// it saturates to T::max()/T::lowest() rather than wrapping.
// Note: NumPy/PyTorch use modular (wrapping) arithmetic for integers, but wrapping
// is non-deterministic in ORT due to parallelized reduction order. Saturation provides
// well-defined behavior. Other EPs (CUDA/DML) may still wrap.
// Double has range ~1.8e308, sufficient for any practical sum of int32/int64 values.
// For int64 values > 2^53, double cannot represent every integer exactly;
// Kahan summation minimizes but does not eliminate rounding in that regime.
double double_accumulator_ = 0.0;
double kahan_compensation_ = 0.0;
ExactIntSum<std::conditional_t<std::is_integral_v<T>, T, int32_t>> int_accumulator_;

public:
inline ReduceAggregatorSum(int64_t N, const T&) : ReduceAggregator<T, T>(N, 0) {}
inline void update(const T& v) {
if constexpr (std::is_integral_v<T>) {
if constexpr (kReduceUseKahan<T>) {
double y = static_cast<double>(v) - kahan_compensation_;
double t = double_accumulator_ + y;
kahan_compensation_ = (t - double_accumulator_) - y;
double_accumulator_ = t;
} else {
double_accumulator_ += static_cast<double>(v);
}
int_accumulator_.add(v);
} else {
this->accumulator_ += v;
}
}
static T aggall(const T* from_data, int64_t size) {
if constexpr (std::is_integral_v<T>) {
double sum = 0.0;
if constexpr (kReduceUseKahan<T>) {
double comp = 0.0;
for (size_t i = 0, n = onnxruntime::narrow<size_t>(size); i < n; ++i) {
double y = static_cast<double>(from_data[i]) - comp;
double t = sum + y;
comp = (t - sum) - y;
sum = t;
}
} else {
for (size_t i = 0, n = onnxruntime::narrow<size_t>(size); i < n; ++i) {
sum += static_cast<double>(from_data[i]);
}
ExactIntSum<T> sum;
for (size_t i = 0, n = onnxruntime::narrow<size_t>(size); i < n; ++i) {
sum.add(from_data[i]);
}
constexpr double t_max = static_cast<double>(std::numeric_limits<T>::max());
constexpr double t_min = static_cast<double>(std::numeric_limits<T>::min());
if (sum >= t_max) return std::numeric_limits<T>::max();
if (sum <= t_min) return std::numeric_limits<T>::min();
return static_cast<T>(sum);
return sum.get();
} else {
return Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, 1>>(
from_data, onnxruntime::narrow<size_t>(size))
Expand All @@ -280,11 +293,7 @@ class ReduceAggregatorSum : public ReduceAggregator<T, T> {
}
inline T get_value() {
if constexpr (std::is_integral_v<T>) {
constexpr double t_max = static_cast<double>(std::numeric_limits<T>::max());
constexpr double t_min = static_cast<double>(std::numeric_limits<T>::min());
if (double_accumulator_ >= t_max) return std::numeric_limits<T>::max();
if (double_accumulator_ <= t_min) return std::numeric_limits<T>::min();
return static_cast<T>(double_accumulator_);
return int_accumulator_.get();
} else {
return this->accumulator_;
}
Expand Down Expand Up @@ -328,19 +337,11 @@ class ReduceAggregatorSum : public ReduceAggregator<T, T> {
tp, onnxruntime::narrow<std::ptrdiff_t>(N), ParallelReduceFastCost(1, n_rows, sizeof(T), 6),
[data, out, N, n_rows](ptrdiff_t begin, ptrdiff_t end) {
for (ptrdiff_t col = begin; col < end; ++col) {
double sum = 0.0;
ExactIntSum<T> sum;
for (int64_t row = 0; row < n_rows; ++row) {
sum += static_cast<double>(data[row * N + col]);
}
constexpr double t_max = static_cast<double>(std::numeric_limits<T>::max());
constexpr double t_min = static_cast<double>(std::numeric_limits<T>::lowest());
if (sum >= t_max) {
out[col] = std::numeric_limits<T>::max();
} else if (sum <= t_min) {
out[col] = std::numeric_limits<T>::lowest();
} else {
out[col] = static_cast<T>(sum);
sum.add(data[row * N + col]);
}
out[col] = sum.get();
}
});
} else {
Expand Down Expand Up @@ -372,19 +373,11 @@ class ReduceAggregatorSum : public ReduceAggregator<T, T> {
[data, fast_shape, stridei, strideo, out, N](ptrdiff_t begin, ptrdiff_t last) {
for (ptrdiff_t d = begin; d < last; ++d) {
for (int64_t col = 0; col < N; ++col) {
double sum = 0.0;
ExactIntSum<T> sum;
for (int64_t row = 0; row < fast_shape[1]; ++row) {
sum += static_cast<double>(data[stridei * d + row * N + col]);
}
constexpr double t_max = static_cast<double>(std::numeric_limits<T>::max());
constexpr double t_min = static_cast<double>(std::numeric_limits<T>::lowest());
if (sum >= t_max) {
out[strideo * d + col] = std::numeric_limits<T>::max();
} else if (sum <= t_min) {
out[strideo * d + col] = std::numeric_limits<T>::lowest();
} else {
out[strideo * d + col] = static_cast<T>(sum);
sum.add(data[stridei * d + row * N + col]);
}
out[strideo * d + col] = sum.get();
}
}
});
Expand Down Expand Up @@ -420,21 +413,13 @@ class ReduceAggregatorSum : public ReduceAggregator<T, T> {
[data, out, d0, d2, inc](ptrdiff_t begin, ptrdiff_t last) {
for (ptrdiff_t d = begin; d < last; ++d) {
const T* p = data + d * d2;
double sum = 0.0;
ExactIntSum<T> sum;
for (int64_t i = 0; i < d0; ++i, p += inc) {
for (int64_t j = 0; j < d2; ++j) {
sum += static_cast<double>(p[j]);
sum.add(p[j]);
}
}
constexpr double t_max = static_cast<double>(std::numeric_limits<T>::max());
constexpr double t_min = static_cast<double>(std::numeric_limits<T>::lowest());
if (sum >= t_max) {
out[d] = std::numeric_limits<T>::max();
} else if (sum <= t_min) {
out[d] = std::numeric_limits<T>::lowest();
} else {
out[d] = static_cast<T>(sum);
}
out[d] = sum.get();
}
});
} else {
Expand Down Expand Up @@ -516,8 +501,27 @@ class ReduceAggregatorSumSquare : public ReduceAggregator<T, TVAL> {

template <typename T>
class ReduceAggregatorMean : public ReduceAggregatorSum<T> {
protected:
// Mean keeps accumulating integers in double (with Kahan compensation for 64-bit types).
double double_accumulator_ = 0.0;
double kahan_compensation_ = 0.0;

public:
inline ReduceAggregatorMean(int64_t N, const T&) : ReduceAggregatorSum<T>(N, 0) {}
inline void update(const T& v) {
if constexpr (std::is_integral_v<T>) {
if constexpr (kReduceUseKahan<T>) {
double y = static_cast<double>(v) - kahan_compensation_;
double t = double_accumulator_ + y;
kahan_compensation_ = (t - double_accumulator_) - y;
double_accumulator_ = t;
} else {
double_accumulator_ += static_cast<double>(v);
}
} else {
this->accumulator_ += v;
}
}
static T aggall(const T* from_data, int64_t size) {
if constexpr (std::is_integral_v<T>) {
// Accumulate in double, divide, then saturate back to T.
Expand Down
52 changes: 52 additions & 0 deletions onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
#include "test/common/tensor_op_test_utils.h"
#include "test/providers/provider_test_utils.h"
#include "test/providers/cpu/reduction/reduction_test_cases.h"
#include "core/framework/session_options.h"
#include "core/providers/cpu/reduction/reduction_ops.h"
#include "test/util/include/default_providers.h"

Expand Down Expand Up @@ -7028,6 +7029,57 @@ TEST(ReductionOpTest, ReduceSum_int64_Overflow_Saturates) {
{kTensorrtExecutionProvider, kOpenVINOExecutionProvider, kDmlExecutionProvider, kWebGpuExecutionProvider});
}

// CPU-only: other EPs (e.g. CUDA) still accumulate integers through double.
// One intra-op thread keeps the fast-reduce dispatch thresholds deterministic.
static void RunReduceSumInt64Cpu(const std::vector<int64_t>& shape, const std::vector<int64_t>& axes,
const std::vector<int64_t>& data, const std::vector<int64_t>& out_shape,
const std::vector<int64_t>& expected) {
OpTester test("ReduceSum", 13);
test.AddAttribute("keepdims", static_cast<int64_t>(0));
test.AddInput<int64_t>("data", shape, data);
test.AddInput<int64_t>("axes", {static_cast<int64_t>(axes.size())}, axes);
test.AddOutput<int64_t>("reduced", out_shape, expected);
SessionOptions so;
so.intra_op_param.thread_pool_size = 1;
test.Config(so).ConfigEp(DefaultCpuExecutionProvider()).RunWithConfig();
}

TEST(ReductionOpTest, ReduceSum_int64_ExactAbove2Pow53) {
const int64_t p53 = int64_t{1} << 53;
const int64_t big = std::numeric_limits<int64_t>::max();

// Small dimensions bypass FastReduceRK and exercise scalar update()/get_value().
RunReduceSumInt64Cpu({2, 2}, {0}, {p53 + 1, p53, -p53, 1}, {2}, {1, p53 + 1});

// Full reductions exercise aggall().
RunReduceSumInt64Cpu({2}, {0}, {p53 + 1, -p53}, {}, {1});
RunReduceSumInt64Cpu({2}, {0}, {p53, 1}, {}, {p53 + 1});
// Partial sum exceeds the int64 range, then cancels.
RunReduceSumInt64Cpu({4}, {0}, {big, big, -big, -big}, {}, {0});

// FastReduceKR: reduce the last axis.
RunReduceSumInt64Cpu({2, 2}, {1}, {p53 + 1, -p53, p53, 1}, {2}, {1, p53 + 1});

// FastReduceRK: reduce the first axis; needs rows > 16 and max(dim) > 256 with one thread.
{
std::vector<int64_t> data(512 * 2, 0);
data[0] = p53 + 1; // row 0, col 0
data[1] = p53; // row 0, col 1
data[2] = -p53; // row 1, col 0
data[3] = 1; // row 1, col 1
RunReduceSumInt64Cpu({512, 2}, {0}, data, {2}, {1, p53 + 1});
}

// FastReduceKRK: reduce the middle axis; needs dim 0 >= 2 with one thread.
RunReduceSumInt64Cpu({2, 3, 2}, {1},
{p53 + 1, p53, -p53, 1, 0, 0,
big, -big, big, -big, -big, big},
{2, 2}, {1, p53 + 1, big, -big});

// FastReduceRKR: reduce the first and last axes; needs dim 1 >= 2 with one thread.
RunReduceSumInt64Cpu({2, 2, 2}, {0, 2}, {p53 + 1, -p53, p53, 1, 0, 0, 0, 0}, {2}, {1, p53 + 1});
}

TEST(ReductionOpTest, ReduceMean_int32_LargeValues_NoOverflow) {
OpTester test("ReduceMean", 18);
// Values that would overflow if summed naively in int32, but mean fits in range.
Expand Down
Loading