void Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/moments_utils.h lines 36–55

template <typename T>
C10_ALWAYS_INLINE void AddMomentsVec(
    int64_t m0_add,
    const vec::Vectorized<T>& m1_add,
    const vec::Vectorized<T>& m2_add,
    int64_t& m0,
    vec::Vectorized<T>& m1,
    vec::Vectorized<T>& m2) {
  using Vec = vec::Vectorized<T>;
  const int64_t n = m0 + m0_add;
  const T c = n == 0 ? static_cast<T>(0) : static_cast<T>(m0_add) / static_cast<T>(n);
  const Vec c_vec(c);
  const Vec delta = m1_add - m1;
  const Vec m2_tmp = m2 + m2_add;
  const Vec c_vec_delta = c_vec * delta;
  const Vec m0_delta = delta * Vec(static_cast<T>(m0));
  m1 = m1 + c_vec_delta;
  m2 = fmadd(m0_delta, c_vec_delta, m2_tmp);
  m0 = n;
}