VsxShiftRightArith Class — pytorch Architecture

Source Code

aten/src/ATen/cpu/vec/vec256/vsx/vsx_helpers.h lines 538–548

template <typename T>
Vectorized<T> VsxShiftRightArith(
    const Vectorized<T>& a,
    const Vectorized<T>& b) {
  const Vectorized<T> max_shift(sizeof(T) * CHAR_BIT - std::is_signed_v<T>);
  const auto mask = (b < Vectorized<T>(0)) | (b >= max_shift);
  const auto shift = Vectorized<T>::blendv(b, max_shift, mask);
  return Vectorized<T>{
      vec_sra(a.vec0(), make_vuint(shift.vec0())),
      vec_sra(a.vec1(), make_vuint(shift.vec1()))};
}