size Class — pytorch Architecture

Source Code

aten/src/ATen/cpu/vec/vec_base.h lines 1406–1428

template <typename T>
inline std::enable_if_t<
    Vectorized<T>::size() % 2 == 0,
    std::pair<Vectorized<T>, Vectorized<T>>>
deinterleave2(const Vectorized<T>& a, const Vectorized<T>& b) {
  static constexpr int size = Vectorized<T>::size();
  static constexpr int half_size = size / 2;
  T a_arr[size];
  T b_arr[size];
  T buffer1[size];
  T buffer2[size];
  a.store(static_cast<void*>(a_arr));
  b.store(static_cast<void*>(b_arr));
  for (const auto i : c10::irange(half_size)) {
    buffer1[i] = a_arr[i * 2];
    buffer1[half_size + i] = b_arr[i * 2];
    buffer2[i] = a_arr[i * 2 + 1];
    buffer2[half_size + i] = b_arr[i * 2 + 1];
  }
  return std::make_pair(
      Vectorized<T>::loadu(static_cast<void*>(buffer1)),
      Vectorized<T>::loadu(static_cast<void*>(buffer2)));
}