mask Class — pytorch Architecture

Source Code

aten/src/ATen/cpu/vec/vec256/vec256_complex_float.h lines 71–126

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