BetaBackward Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/group_norm_kernel.cpp lines 817–839

template <typename PT, typename opmath_t>
std::enable_if_t<std::is_same_v<PT, opmath_t>, void>
BetaBackward(int64_t N, int64_t C, const opmath_t* db, PT* dbeta) {
  using Vec = at::vec::Vectorized<PT>;
  constexpr int64_t K = Vec::size();
  Vec acc_vec{0}, zero{0};
  const int64_t inner_size = C / K * K;
  int64_t i = 0;
  for (; i < inner_size; i += K) {
    for (const auto n : c10::irange(N)) {
      acc_vec += Vec::loadu(db + n * C + i);
    }
    acc_vec.store(dbeta + i);
    acc_vec = Vec::set(acc_vec, zero);
  }
  if (C - i > 0) {
    for (const auto n : c10::irange(N)) {
      acc_vec += Vec::loadu(db + n * C + i, C - i);
    }
    acc_vec.store(dbeta + i, C - i);
    acc_vec = Vec::set(acc_vec, zero, C - i);
  }
}