dst_n Class — pytorch Architecture

Source Code

aten/src/ATen/cpu/vec/vec256/vec256_mask.h lines 12–42

template <typename T, int dst_n, typename mask_t, int mask_n>
struct VecMaskLoad<
    T,
    dst_n,
    mask_t,
    mask_n,
    typename std::enable_if_t<
        (mask_n == dst_n * 2 && dst_n >= 1) &&
            (std::is_same_v<T, float> || std::is_same_v<T, int32_t>),
        void>> {
  static inline VectorizedN<T, dst_n> apply(
      const T* ptr,
      const VecMask<mask_t, mask_n>& vec_mask) {
    VectorizedN<mask_t, 2> tmp_vec;
    VectorizedN<T, dst_n> result;
    for (int i = 0; i < dst_n; i++) {
      tmp_vec[0] = vec_mask[2 * i];
      tmp_vec[1] = vec_mask[2 * i + 1];
      auto int64_mask = VecMask<mask_t, 2>(tmp_vec).template cast<int64_t, 2>();
      auto int_mask = int64_mask.template cast<int, 1>()[0];
      if constexpr (std::is_same_v<T, float>) {
        result[i] = Vectorized<T>(
            _mm256_maskload_ps(ptr + i * Vectorized<T>::size(), int_mask));
      } else {
        result[i] = Vectorized<T>(
            _mm256_maskload_epi32(ptr + i * Vectorized<T>::size(), int_mask));
      }
    }
    return result;
  }
};