maximum Class — pytorch Architecture

Source Code

aten/src/ATen/cpu/vec/vec512/vec512_double.h lines 412–423

template <>
Vectorized<double> inline maximum(
    const Vectorized<double>& a,
    const Vectorized<double>& b) {
  auto zero_vec = _mm512_set1_epi64(0);
  Vectorized<double> max = _mm512_max_pd(a, b);
  auto isnan_mask = _mm512_cmp_pd_mask(a, b, _CMP_UNORD_Q);
  auto isnan = _mm512_castsi512_pd(
      _mm512_mask_set1_epi64(zero_vec, isnan_mask, 0xFFFFFFFFFFFFFFFF));
  // Exploit the fact that all-ones is a NaN.
  return _mm512_or_pd(max, isnan);
}