mask Class — pytorch Architecture

Source Code

aten/src/ATen/cpu/vec/vec512/vec512_complex_double.h lines 72–128

  template <int64_t mask>
  static Vectorized<c10::complex<double>> blend(
      const Vectorized<c10::complex<double>>& a,
      const Vectorized<c10::complex<double>>& b) {
    // convert c10::complex<V> index mask to V index mask: xy -> xxyy
    // NOLINTNEXTLINE(clang-diagnostic-warning)
    switch (mask) {
      case 0:
        return a;
      case 1:
        return _mm512_mask_blend_pd(
            0x03, a.values, b.values); // b0000 0001 = b0000 0011
      case 2:
        return _mm512_mask_blend_pd(
            0x0C, a.values, b.values); // b0000 0010 = b0000 1100
      case 3:
        return _mm512_mask_blend_pd(
            0x0F, a.values, b.values); // b0000 0011 = b0000 1111
      case 4:
        return _mm512_mask_blend_pd(
            0x30, a.values, b.values); // b0000 0100 = b0011 0000
      case 5:
        return _mm512_mask_blend_pd(
            0x33, a.values, b.values); // b0000 0101 = b0011 0011
      case 6:
        return _mm512_mask_blend_pd(
            0x3C, a.values, b.values); // b0000 0110 = b0011 1100
      case 7:
        return _mm512_mask_blend_pd(
            0x3F, a.values, b.values); // b0000 0111 = b0011 1111
      case 8:
        return _mm512_mask_blend_pd(
            0xC0, a.values, b.values); // b0000 1000 = b1100 0000
      case 9:
        return _mm512_mask_blend_pd(
            0xC3, a.values, b.values); // b0000 1001 = b1100 0011
      case 10:
        return _mm512_mask_blend_pd(
            0xCC, a.values, b.values); // b0000 1010 = b1100 1100
      case 11:
        return _mm512_mask_blend_pd(
            0xCF, a.values, b.values); // b0000 1011 = b1100 1111
      case 12:
        return _mm512_mask_blend_pd(
            0xF0, a.values, b.values); // b0000 1100 = b1111 0000
      case 13:
        return _mm512_mask_blend_pd(
            0xF3, a.values, b.values); // b0000 1101 = b1111 0011
      case 14:
        return _mm512_mask_blend_pd(
            0xFC, a.values, b.values); // b0000 1110 = b1111 1100
      case 15:
        return _mm512_mask_blend_pd(
            0xFF, a.values, b.values); // b0000 1111 = b1111 1111
    }
    return b;
  }