mask Class — pytorch Architecture

Source Code

aten/src/ATen/cpu/vec/vec128/vec128_int_aarch64.h lines 299–310

template <int64_t mask>
Vectorized<int64_t> Vectorized<int64_t>::blend(
    const Vectorized<int64_t>& a,
    const Vectorized<int64_t>& b) {
  // Build an array of flags: each bit of element is 1 if the corresponding bit
  // in 'mask' is set, 0 otherwise.
  uint64x2_t maskArray = {
      (mask & 1LL) ? 0xFFFFFFFFFFFFFFFF : 0,
      (mask & 2LL) ? 0xFFFFFFFFFFFFFFFF : 0};
  // Use BSL to select elements from b where the mask is 1, else from a
  return vbslq_s64(maskArray, b.values, a.values);
}