is_same_v Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/group_norm_kernel.cpp lines 120–160

template <typename T>
std::enable_if_t<!std::is_same_v<T, at::opmath_type<T>>,
  std::tuple<at::opmath_type<T>, at::opmath_type<T>>>
ColumnwiseMoments(
    const T* X_data,
    int64_t HxW,
    int64_t C,
    int64_t D) {
  using opmath_t = at::opmath_type<T>;
  using Vec = vec::Vectorized<T>;
  using fVec = vec::Vectorized<opmath_t>;
  constexpr int64_t K = Vec::size();
  const int64_t inner_size = D / K * K;
  fVec acc0_fvec{0}, acc1_fvec{0}, zero{0};
  for (const auto m : c10::irange(HxW)) {
    const T* X_ptr = X_data + m * C;
    int64_t d = 0;
    for (; d < inner_size; d += K) {
      Vec x_bvec = Vec::loadu(X_ptr + d);
      auto [x_fvec0, x_fvec1] = convert_to_float<T>(x_bvec);
      acc0_fvec += x_fvec0 + x_fvec1;
      acc1_fvec += x_fvec0 * x_fvec0 + x_fvec1 * x_fvec1;
    }
    if (D - d > 0) {
      Vec x_bvec = Vec::loadu(X_ptr + d, D - d);
      auto [x_fvec0, x_fvec1] = convert_to_float<T>(x_bvec);
      if (D - d > fVec::size()) {
        x_fvec1 = fVec::set(zero, x_fvec1, D - d - fVec::size());
        acc0_fvec += x_fvec0 + x_fvec1;
        acc1_fvec += x_fvec0 * x_fvec0 + x_fvec1 * x_fvec1;
      } else {
        x_fvec0 = fVec::set(zero, x_fvec0, D - d);
        acc0_fvec += x_fvec0;
        acc1_fvec += x_fvec0 * x_fvec0;
      }
    }
  }
  opmath_t mean_val = vec::vec_reduce_all([](fVec& x, fVec& y) { return x + y; }, acc0_fvec);
  opmath_t rstd_val = vec::vec_reduce_all([](fVec& x, fVec& y) { return x + y; }, acc1_fvec);
  return std::tuple<opmath_t, opmath_t>(mean_val, rstd_val);
}