maximum Class — pytorch Architecture

Source Code

aten/src/ATen/cpu/vec/vec512/vec512_complex_float.h lines 1152–1165

template <>
Vectorized<c10::complex<float>> inline maximum(
    const Vectorized<c10::complex<float>>& a,
    const Vectorized<c10::complex<float>>& b) {
  auto zero_vector = _mm512_set1_epi32(0);
  auto abs_a = a.abs_2_();
  auto abs_b = b.abs_2_();
  auto mask = _mm512_cmp_ps_mask(abs_a, abs_b, _CMP_LT_OQ);
  auto max = _mm512_mask_blend_ps(mask, a, b);
  // Exploit the fact that all-ones is a NaN.
  auto isnan_mask = _mm512_cmp_ps_mask(abs_a, abs_b, _CMP_UNORD_Q);
  auto isnan = _mm512_mask_set1_epi32(zero_vector, isnan_mask, 0xFFFFFFFF);
  return _mm512_or_ps(max, _mm512_castsi512_ps(isnan));
}