fmadd Class — pytorch Architecture

Source Code

aten/src/ATen/cpu/vec/vec128/vec128_bfloat16_neon.h lines 627–645

template <>
Vectorized<c10::BFloat16> inline fmadd(
    const Vectorized<c10::BFloat16>& a,
    const Vectorized<c10::BFloat16>& b,
    const Vectorized<c10::BFloat16>& c) {
#if BF16_ARITHMETIC_SUPPORTED()
  bfloat16x8_t x = a;
  bfloat16x8_t y = b;
  bfloat16x8_t z = c;
  return x * y + z;
#else
  // NOTE [BF16 FMA]: There isn't an FMA that accumulates into BF16!  Also,
  // vbfmlalbq_f32 and vbfmlaltq_f32 take the even and odd-numbered
  // elements, not the bottom and top half, so they don't seem
  // particularly useful here. Ideally we would include dot product in
  // the Vectorized interface...
  return a * b + c;
#endif
}