binary_op_as_fp32 Class — pytorch Architecture

Source Code

aten/src/ATen/cpu/vec/vec512/vec512_bfloat16.h lines 813–825

template <typename T, typename Op>
static inline Vectorized<T> binary_op_as_fp32(
    const Vectorized<T>& a,
    const Vectorized<T>& b,
    Op op) {
  __m512 a_lo, a_hi;
  __m512 b_lo, b_hi;
  cvt_to_fp32<T>(__m512i(a), a_lo, a_hi);
  cvt_to_fp32<T>(__m512i(b), b_lo, b_hi);
  auto o1 = op(a_lo, b_lo);
  auto o2 = op(a_hi, b_hi);
  return cvt_from_fp32<T>(o1, o2);
}