kMaxDepth Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/moments_utils.h lines 116–188

template <typename T, int64_t kMaxDepth>
std::pair<opmath_t<T>, opmath_t<T>> RowwiseMomentsImpl(const T* X, int64_t N, int64_t ddof = 0) {
  using math_t = opmath_t<T>;

  constexpr int64_t kVecSize = vec::Vectorized<T>::size();
  constexpr int64_t kAccVecSize = vec::Vectorized<math_t>::size();
  const int64_t n = N / kVecSize;
  const int64_t m = divup(n, kChunkSize);
  const int64_t depth = utils::CeilLog2(m);

  using Vec = vec::Vectorized<math_t>;
  const Vec kZeroVec(math_t(0));
  std::array<int64_t, kMaxDepth> m0_stk = {{0}};
  std::array<Vec, kMaxDepth> m1_stk;
  m1_stk.fill(kZeroVec);
  std::array<Vec, kMaxDepth> m2_stk;
  m2_stk.fill(kZeroVec);

  for (const auto i : c10::irange(m)) {
    const T* X_ptr = X + i * kChunkSize * kVecSize;
    const int64_t m0 = std::min(kChunkSize, n - i * kChunkSize);
    static std::array<Vec, kChunkSize> c_vecs = ([]() {
      std::array<Vec, kChunkSize> result;
      for (const auto i : c10::irange(kChunkSize)) {
        result[i] = Vec(math_t(1) / static_cast<math_t>(i + 1));
      }
      return result;
    })();
    UpdateMomentsVec(m0, X_ptr, c_vecs, m0_stk[0], m1_stk[0], m2_stk[0]);

    int64_t mask = i + 1;
    for (int64_t j = 1; j < depth && (mask & 1) == 0; ++j) {
      AddMomentsVec(
          m0_stk[j - 1],
          m1_stk[j - 1],
          m2_stk[j - 1],
          m0_stk[j],
          m1_stk[j],
          m2_stk[j]);
      m0_stk[j - 1] = 0;
      m1_stk[j - 1] = kZeroVec;
      m2_stk[j - 1] = kZeroVec;
      mask >>= 1;
    }
  }
  for (const auto i : c10::irange(1, depth)) {
    AddMomentsVec(
        m0_stk[i], m1_stk[i], m2_stk[i], m0_stk[0], m1_stk[0], m2_stk[0]);
  }

  std::array<math_t, kAccVecSize> m1_arr{};
  std::array<math_t, kAccVecSize> m2_arr{};
  m1_stk[0].store(m1_arr.data());
  m2_stk[0].store(m2_arr.data());

  int64_t m0 = 0;
  math_t m1 = 0;
  math_t m2 = 0;
  for (int64_t i = n * kVecSize; i < N; ++i) {
    math_t x = static_cast<math_t>(X[i]);
    const math_t delta = x - m1;
    ++m0;
    m1 += delta / static_cast<math_t>(m0);
    m2 += delta * (x - m1);
  }
  // for BFloat16, each vector in m1_arr/m2_arr holds 2*n accumulated result
  int64_t m0_add = n * kVecSize / kAccVecSize;
  for (const auto i : c10::irange(kAccVecSize)) {
    AddMoments(m0_add, m1_arr[i], m2_arr[i], m0, m1, m2);
  }

  return std::make_pair(m1, m2 / static_cast<math_t>(N - ddof));
}