Skip to content

Fixes to mean and better precision on sum (cpu) #3687

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
willyborn wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Fixes to mean and better precision on sum (cpu) #3687

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
058db8a
sum on CPU now has the same precision as other platforms
willyborn Aug 2, 2025
4a8dca4
Fixes to mean, on all platforms
willyborn Aug 8, 2025
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
45 changes: 43 additions & 2 deletions src/backend/cpu/kernel/mean.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,45 @@ struct MeanOp {
}
};

template<typename Ti, typename To, typename Tw>
struct MeanOpWithCorrection {
common::Transform<Ti, To, af_add_t> transform;
To runningMean;
Tw runningCount;
To correction;
MeanOpWithCorrection(Ti mean, Tw count)
: transform()
, runningMean(transform(mean))
, runningCount(count)
, correction(scalar<To>(0)) {}

void operator()(Ti newMean, Tw newCount) {
runningCount += newCount;
if ((newCount != 0) || (runningCount != 0)) { //
// Since only 1 pass is used, the rounding errors will become
// important because the longer the serie the larger the
// difference between the 2 numbers to sum See:
// https://en.wikipedia.org/wiki/Kahan_summation_algorithm to
// reduce this error
// correction is zero for the first time around
To y =
(transform(newMean) - runningMean) * (newCount / runningCount) -
correction;
// Alas, runningMean is big, y small, so low-order digits of y
// are lost
To t = runningMean + y;
// (t - runningMean) cancels the high order part of y
// subtracting y recovers negative (low part of y)
correction = (t - runningMean) - y;
// Algebraically, correction should always be zero. Beware
// overly-agressive optimizing compilers!
runningMean = t;
// Next time around, the lost low part will be added to y in a
// fresh attempt
}
}
};

template<typename T, typename Tw, int D>
struct mean_weighted_dim {
void operator()(Param<T> output, const dim_t outOffset,
Expand Down Expand Up @@ -74,7 +113,8 @@ struct mean_weighted_dim<T, Tw, 0> {

dim_t istride = istrides[dim];
dim_t wstride = wstrides[dim];
MeanOp<compute_t<T>, compute_t<T>, compute_t<Tw>> Op(0, 0);
MeanOpWithCorrection<compute_t<T>, compute_t<T>, compute_t<Tw>> Op(0,
0);
for (dim_t i = 0; i < idims[dim]; i++) {
Op(compute_t<T>(in[inOffset + i * istride]),
compute_t<Tw>(wt[wtOffset + i * wstride]));
Expand Down Expand Up @@ -113,7 +153,8 @@ struct mean_dim<Ti, Tw, To, 0> {

dim_t istride = istrides[dim];
dim_t end = inOffset + idims[dim] * istride;
MeanOp<compute_t<Ti>, compute_t<To>, compute_t<Tw>> Op(0, 0);
MeanOpWithCorrection<compute_t<Ti>, compute_t<To>, compute_t<Tw>> Op(0,
0);
for (dim_t i = inOffset; i < end; i += istride) {
Op(compute_t<Ti>(in[i]), 1);
}
Expand Down
136 changes: 136 additions & 0 deletions src/backend/cpu/kernel/reduce.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,27 @@ namespace arrayfire {
namespace cpu {
namespace kernel {

// Since only 1 pass is used, the rounding errors will become
// important because the longer the serie the larger the difference
// between the 2 numbers to sum See:
// https://en.wikipedia.org/wiki/Kahan_summation_algorithm to reduce
// this error
template<typename Ti, typename To>
To stable_add(To &correction, To out_val, Ti in_val) {
// correction is zero for the first time around
To y = in_val - correction;
// Alas out_val is big, y small, so low-order digits of y are lost
To t = out_val + y;
// (t - out_val) cancels the high order part of y
// subtracting y recovers negative (low part of y)
correction = (t - out_val) - y;
// Algebraically, correction should always be zero. Beware of
// overly-agressive optimizing compilers!
return t;
// Next time around, the lost low part will be added to y in a fresh
// attempt
}

template<af_op_t op, typename Ti, typename To, int D>
struct reduce_dim {
void operator()(Param<To> out, const dim_t outOffset, CParam<Ti> in,
Expand Down Expand Up @@ -62,6 +83,32 @@ struct reduce_dim<op, Ti, To, 0> {
}
};

template<typename Ti, typename To>
struct reduce_dim<af_add_t, Ti, To, 0> {
common::Transform<data_t<Ti>, compute_t<To>, af_add_t> transform;
common::Binary<compute_t<To>, af_add_t> reduce;
void operator()(Param<To> out, const dim_t outOffset, CParam<Ti> in,
const dim_t inOffset, const int dim, bool change_nan,
double nanval) {
const af::dim4 istrides = in.strides();
const af::dim4 idims = in.dims();

data_t<To> *const outPtr = out.get() + outOffset;
data_t<Ti> const *const inPtr = in.get() + inOffset;
dim_t stride = istrides[dim];

compute_t<To> out_val = common::Binary<compute_t<To>, af_add_t>::init();
compute_t<To> correction = scalar<compute_t<To>>(0);
for (dim_t i = 0; i < idims[dim]; i++) {
compute_t<To> in_val = transform(inPtr[i * stride]);
if (change_nan) in_val = IS_NAN(in_val) ? nanval : in_val;
out_val = stable_add(correction, out_val, in_val);
}

*outPtr = data_t<To>(out_val);
}
};

template<typename Tk>
void n_reduced_keys(Param<Tk> okeys, int *n_reduced, CParam<Tk> keys) {
const af::dim4 kdims = keys.dims();
Expand Down Expand Up @@ -159,6 +206,55 @@ struct reduce_dim_by_key<op, Ti, Tk, To, 0> {
}
};

template<typename Ti, typename Tk, typename To>
struct reduce_dim_by_key<af_add_t, Ti, Tk, To, 0> {
common::Transform<data_t<Ti>, compute_t<To>, af_add_t> transform;
void operator()(Param<To> ovals, const dim_t ovOffset, CParam<Tk> keys,
CParam<Ti> vals, const dim_t vOffset, int *n_reduced,
const int dim, bool change_nan, double nanval) {
const af::dim4 vstrides = vals.strides();
const af::dim4 vdims = vals.dims();

const af::dim4 ovstrides = ovals.strides();

data_t<Tk> const *const inKeysPtr = keys.get();
data_t<Ti> const *const inValsPtr = vals.get();
data_t<To> *const outValsPtr = ovals.get();

int keyidx = 0;
compute_t<Tk> current_key = compute_t<Tk>(inKeysPtr[0]);
compute_t<To> out_val = common::Binary<compute_t<To>, af_add_t>::init();
compute_t<To> correction = scalar<compute_t<To>>(0);

dim_t istride = vstrides[dim];
dim_t ostride = ovstrides[dim];

for (dim_t i = 0; i < vdims[dim]; i++) {
compute_t<Tk> keyval = inKeysPtr[i];

if (keyval == current_key) {
compute_t<To> in_val =
transform(inValsPtr[vOffset + (i * istride)]);
if (change_nan) in_val = IS_NAN(in_val) ? nanval : in_val;
out_val = stable_add(correction, out_val, in_val);

} else {
outValsPtr[ovOffset + (keyidx * ostride)] = out_val;

current_key = keyval;
correction = scalar<compute_t<To>>(0);
out_val = transform(inValsPtr[vOffset + (i * istride)]);
if (change_nan) out_val = IS_NAN(out_val) ? nanval : out_val;
++keyidx;
}

if (i == (vdims[dim] - 1)) {
outValsPtr[ovOffset + (keyidx * ostride)] = out_val;
}
}
}
};

template<af_op_t op, typename Ti, typename To>
struct reduce_all {
common::Transform<data_t<Ti>, compute_t<To>, op> transform;
Expand Down Expand Up @@ -199,6 +295,46 @@ struct reduce_all {
}
};

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) {
// Decrement dimension of select dimension
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 = scalar<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;
}
out_val = stable_add(correction, out_val, in_val);
}
}
}
}

*outPtr = data_t<To>(out_val);
}
};

} // namespace kernel
} // namespace cpu
} // namespace arrayfire
43 changes: 23 additions & 20 deletions src/backend/cpu/mean.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,32 +63,35 @@ Array<T> mean(const Array<T> &in, const Array<Tw> &wt, const int dim) {

template<typename T, typename Tw>
T mean(const Array<T> &in, const Array<Tw> &wt) {
using MeanOpT = kernel::MeanOp<compute_t<T>, compute_t<T>, compute_t<Tw>>;
using MeanOpT =
kernel::MeanOpWithCorrection<compute_t<T>, compute_t<T>, compute_t<Tw>>;
in.eval();
wt.eval();
getQueue().sync();

af::dim4 dims = in.dims();
af::dim4 strides = in.strides();
const T *inPtr = in.get();
const Tw *wtPtr = wt.get();

auto input = compute_t<T>(inPtr[0]);
auto weight = compute_t<Tw>(wtPtr[0]);
MeanOpT Op(input, weight);
const af::dim4 &dims = in.dims();
const af::dim4 &istrides = in.strides();
const af::dim4 &wstrides = wt.strides();
const T *inPtr = in.get();
const Tw *wtPtr = wt.get();

// Split workload in equal parts, to improve accuracy or larger arrays
MeanOpT Op(0, 0);
for (dim_t l = 0; l < dims[3]; l++) {
dim_t off3 = l * strides[3];
dim_t ioff3 = l * istrides[3];
dim_t woff3 = l * wstrides[3];

for (dim_t k = 0; k < dims[2]; k++) {
dim_t off2 = k * strides[2];
dim_t ioff2 = k * istrides[2];
dim_t woff2 = k * wstrides[2];

for (dim_t j = 0; j < dims[1]; j++) {
dim_t off1 = j * strides[1];
dim_t ioff1 = j * istrides[1];
dim_t woff1 = j * wstrides[1];

for (dim_t i = 0; i < dims[0]; i++) {
dim_t idx = i + off1 + off2 + off3;
Op(compute_t<T>(inPtr[idx]), compute_t<Tw>(wtPtr[idx]));
Op(compute_t<T>(inPtr[i + ioff1 + ioff2 + ioff3]),
compute_t<Tw>(wtPtr[i + woff1 + woff2 + woff3]));
}
}
}
Expand All @@ -99,16 +102,16 @@ T mean(const Array<T> &in, const Array<Tw> &wt) {

template<typename Ti, typename Tw, typename To>
To mean(const Array<Ti> &in) {
using MeanOpT = kernel::MeanOp<compute_t<Ti>, compute_t<To>, compute_t<Tw>>;
using MeanOpT = kernel::MeanOpWithCorrection<compute_t<Ti>, compute_t<To>,
compute_t<Tw>>;
in.eval();
getQueue().sync();

af::dim4 dims = in.dims();
af::dim4 strides = in.strides();
const Ti *inPtr = in.get();
const af::dim4 &dims = in.dims();
const af::dim4 &strides = in.strides();
const Ti *inPtr = in.get();

MeanOpT Op(0, 0);

for (dim_t l = 0; l < dims[3]; l++) {
dim_t off3 = l * strides[3];

Expand All @@ -120,7 +123,7 @@ To mean(const Array<Ti> &in) {

for (dim_t i = 0; i < dims[0]; i++) {
dim_t idx = i + off1 + off2 + off3;
Op(compute_t<Ti>(inPtr[idx]), 1);
Op(compute_t<Ti>(inPtr[idx]), 1.);
}
}
}
Expand Down
Loading
Loading