zero_mask Class — pytorch Architecture

Source Code

aten/src/ATen/cpu/vec/vec256/zarch/vec256_zarch.h lines 857–868

  template <
      typename U = T,
      std::enable_if_t<std::is_floating_point_v<U>, int> = 0>
  int zero_mask() const {
    auto cmp = (*this == Vectorized<U>(0));
    constexpr auto mask_zero_bits = GetBpermZeroMask<U>();
    ZSimdVectBinary<uint64_t> result0 =
        vec_bperm_u128((ZSimdVectBinary<uint8_t>)cmp.vecb0(), mask_zero_bits);
    ZSimdVectBinary<uint64_t> result1 =
        vec_bperm_u128((ZSimdVectBinary<uint8_t>)cmp.vecb1(), mask_zero_bits);
    return (result0[0] | (result1[0] << (size() / 2)));
  }