is_same_v Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/moments_utils.h lines 82–110

template <typename T>
inline std::enable_if_t<!std::is_same_v<T, at::opmath_type<T>>, void>
UpdateMomentsVec(
    int64_t m0,
    const T* X_ptr,
    const std::array<vec::Vectorized<at::opmath_type<T>>, kChunkSize>& c_vecs,
    int64_t& m0_stk0,
    vec::Vectorized<at::opmath_type<T>>& m1_stk0,
    vec::Vectorized<at::opmath_type<T>>& m2_stk0) {
  using Vec = vec::Vectorized<T>;
  using fVec = vec::Vectorized<at::opmath_type<T>>;
  fVec m1_fvec0(0), m1_fvec1(0);
  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 = 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);
}