is_same_v Class — pytorch Architecture

Source Code

aten/src/ATen/cpu/vec/vec256/vec256_mask.h lines 75–102

template <typename T, typename mask_t>
struct VecMaskLoad<
    T,
    2,
    mask_t,
    1,
    typename std::enable_if_t<
        std::is_same_v<T, int64_t> || std::is_same_v<T, double>>> {
  static inline VectorizedN<T, 2> apply(
      const T* ptr,
      const VecMask<mask_t, 1>& vec_mask) {
    auto int64_mask = vec_mask.template cast<int64_t, 2>();
    auto result = at::vec::VectorizedN<T, 2>();
    if constexpr (std::is_same_v<T, double>) {
      result[0] = _mm256_maskload_pd(ptr, int64_mask[0]);
      result[1] = _mm256_maskload_pd(
          ptr + at::vec::Vectorized<T>::size(), int64_mask[1]);
    } else {
      result[0] = _mm256_maskload_epi64(
          reinterpret_cast<const long long*>(ptr), int64_mask[0]);
      result[1] = _mm256_maskload_epi64(
          reinterpret_cast<const long long*>(
              ptr + at::vec::Vectorized<T>::size()),
          int64_mask[1]);
    }
    return result;
  }
};