mask Class — pytorch Architecture

Source Code

aten/src/ATen/cpu/vec/sve/vec_float.h lines 58–77

  template <uint64_t mask>
  static Vectorized<float> blend(
      const Vectorized<float>& a,
      const Vectorized<float>& b) {
    // Build an array of flags: each element is 1 if the corresponding bit in
    // 'mask' is set, 0 otherwise.
    __at_align__ int32_t flag_arr[size()];
    for (int i = 0; i < size(); i++) {
      flag_arr[i] = (mask & (1ULL << i)) ? 1 : 0;
    }
    // Load the flag array into an SVE int32 vector.
    svint32_t int_mask = svld1_s32(svptrue_b32(), flag_arr);
    // Compare each lane of int_mask to 0; returns an svbool_t predicate where
    // true indicates a nonzero flag.
    svbool_t blend_mask = svcmpne_n_s32(svptrue_b32(), int_mask, 0);
    // Use svsel to select elements from b where the predicate is true, else
    // from a.
    svfloat32_t result = svsel_f32(blend_mask, b.values, a.values);
    return Vectorized<float>(result);
  }