CalcDsDb Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/group_norm_kernel.cpp lines 603–624

template <typename PT, typename opmath_t>
inline std::enable_if_t<std::is_same_v<PT, opmath_t>, void>
CalcDsDb(
    const opmath_t* ds_ptr,
    const opmath_t* db_ptr,
    const PT* gamma_ptr,
    const int64_t d,
    const int64_t K,
    void* ds_arr,
    void* db_arr) {
    vec::Vectorized<opmath_t> ds_vec(0);
    vec::Vectorized<opmath_t> db_vec(0);
    for (int64_t j = 0; j < d; j += K) {
      const vec::Vectorized<PT> gamma_vec = (gamma_ptr == nullptr)
          ? vec::Vectorized<PT>(1)
          : vec::Vectorized<PT>::loadu(gamma_ptr + j);
      ds_vec = ds_vec + vec::Vectorized<PT>::loadu(ds_ptr + j) * gamma_vec;
      db_vec = db_vec + vec::Vectorized<PT>::loadu(db_ptr + j) * gamma_vec;
    }
    ds_vec.store(ds_arr);
    db_vec.store(db_arr);
}