dst_n Class — pytorch Architecture

Source Code

aten/src/ATen/cpu/vec/vec256/vsx/vec256_mask_vsx.h lines 38–64

template <int dst_n, typename mask_t, int mask_n>
struct VecMaskCast<
    int64_t,
    dst_n,
    mask_t,
    mask_n,
    typename std::enable_if_t<
        (dst_n == 2 * mask_n) &&
        (std::is_same_v<mask_t, float> || std::is_same_v<mask_t, int>)>> {
  static inline VecMask<int64_t, dst_n> apply(
      const VecMask<mask_t, mask_n>& vec_mask) {
    VectorizedN<int64_t, dst_n> result;

    auto int_mask = vec_mask.template cast<int, mask_n>();

    for (int i = 0; i < mask_n; ++i) {
      VectorizedN<int, 1> in_int_n;
      in_int_n[0] = int_mask[i];

      auto int64_vecs = convert<int64_t, 2, int, 1>(in_int_n);

      result[2 * i] = int64_vecs[0];
      result[2 * i + 1] = int64_vecs[1];
    }
    return VecMask<int64_t, dst_n>(result);
  }
};