[Caffe2] Improve AddMomentsVec and UpdateMomentsVec (pytorch#167664)

Summary: RowwiseMomentsImpl accounts for about 0.4% cpu time of AdRanker: https://fburl.com/strobelight/ywf79nw3. It primarily calls AddMomentsVec and UpdateMomentsVec. These two routines are written using Pytorch's VecLib, meaning the utilized operators translate into intrinsics. Unfortunately, the compiler makes less transformations and optimizations when intrinsics are used. Therefore, if we carefully decouple and re-order operations, the emitted instruction sequence improves. Here we can see the dissassembly for the old and new AddMomentsVec: https://godbolt.org/z/83fxYvKfv We can see a much better instruction sequence is achieved in the new implementation. Test Plan: AdRanker ServiceLab Differential Revision: D86805648 Pull Request resolved: pytorch#167664 Approved by: https://github.com/mcfi

Author: Nicoshev

aten/src/ATen/native/cpu/moments_utils.h

@@ -46,8 +46,11 @@ C10_ALWAYS_INLINE void AddMomentsVec(
  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;
- m1 += c_vec * delta;
- m2 += m2_add + delta * delta * c_vec * Vec(static_cast<T>(m0));
+ 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;
 }
 
@@ -65,9 +68,11 @@ UpdateMomentsVec(
  Vec m2_vec(0);
  for (const auto j : c10::irange(m0)) {
  const Vec x_vec = Vec::loadu(X_ptr + j * Vec::size());
+ const Vec tmpVec = c_vecs[j];
  const Vec delta_vec = x_vec - m1_vec;
- m1_vec += delta_vec * c_vecs[j];
- m2_vec += delta_vec * (x_vec - m1_vec);
+ m1_vec = fmadd(tmpVec, delta_vec, m1_vec);
+ const Vec tmpVec2 = x_vec - m1_vec;
+ m2_vec = fmadd(delta_vec, tmpVec2, m2_vec);
  }
  AddMomentsVec(m0, m1_vec, m2_vec, m0_stk0, m1_stk0, m2_stk0);
 }
@@ -89,13 +94,16 @@ UpdateMomentsVec(
  fVec m2_fvec0(0), m2_fvec1(0);
  for (const auto j : c10::irange(m0)) {
  const Vec x_bvec = Vec::loadu(X_ptr + j * Vec::size());
+ const fVec tmpVec = c_vecs[j];
  auto [x_fvec0, x_fvec1] = convert_to_float<T>(x_bvec);
  const fVec delta_fvec0 = x_fvec0 - m1_fvec0;
  const fVec delta_fvec1 = x_fvec1 - m1_fvec1;
- m1_fvec0 += delta_fvec0 * c_vecs[j];
- m1_fvec1 += delta_fvec1 * c_vecs[j];
- m2_fvec0 += delta_fvec0 * (x_fvec0 - m1_fvec0);
- m2_fvec1 += delta_fvec1 * (x_fvec1 - m1_fvec1);
+ m1_fvec0 = fmadd(delta_fvec0, tmpVec, m1_fvec0);
+ m1_fvec1 = fmadd(delta_fvec1, tmpVec, m1_fvec1);
+ const fVec delta_fvec2 = x_fvec0 - m1_fvec0;
+ const fVec delta_fvec3 = x_fvec1 - m1_fvec1;
+ m2_fvec0 = fmadd(delta_fvec0, delta_fvec2, m2_fvec0);
+ m2_fvec1 = fmadd(delta_fvec1, delta_fvec3, m2_fvec1);
  }
  AddMomentsVec(m0, m1_fvec0, m2_fvec0, m0_stk0, m1_stk0, m2_stk0);
  AddMomentsVec(m0, m1_fvec1, m2_fvec1, m0_stk0, m1_stk0, m2_stk0);