Repository: google/gemma.cpp
File: gemma/ops.h (C/C++)

Code before:
  const hn::ScalableTag<float> df;

  for (size_t idx_row = 0; idx_row < kOuter; ++idx_row) {
    const size_t row_ofs = mat_ofs + idx_row * kInner;
    if constexpr (kAdd) {
      out[idx_row] = hwy::ConvertScalarTo<float>(add[idx_row]) +
                     Dot(df, mat, row_ofs, vec_aligned, kInner);
    } else {
      out[idx_row] = Dot(df, mat, row_ofs, vec_aligned, kInner);
    }
  }
}

template <bool kAdd, size_t kOuter, size_t kInner, typename VecT, typename AddT,
          size_t kCapacity>
HWY_INLINE void MatVecAddLoop(
    const CompressedArray<hwy::bfloat16_t, kCapacity>& mat,
    const size_t mat_ofs,
    const VecT* HWY_RESTRICT vec_aligned,
    const AddT* HWY_RESTRICT add,
    float* HWY_RESTRICT out) {
  PROFILER_ZONE("MatVecAddLoop");

  const hn::ScalableTag<float> df;

  const auto vec_dequant = hwy::AllocateAligned<float>(kInner);
  ToEvenOddF32(vec_aligned, kInner, vec_dequant.get());

  for (size_t idx_row = 0; idx_row < kOuter; ++idx_row) {
    const size_t row_ofs = mat_ofs + idx_row * kInner;
    if constexpr (kAdd) {
      out[idx_row] = hwy::ConvertScalarTo<float>(add[idx_row]) +
                     Dot<true>(df, mat, row_ofs, vec_dequant.get(), kInner);
    } else {
      out[idx_row] = Dot<true>(df, mat, row_ofs, vec_dequant.get(), kInner);
    }
  }
}

template <size_t kOuter, size_t kInner, typename ArrayT, typename VecT>
HWY_INLINE void MatVecLoop(const ArrayT& mat, const size_t mat_ofs,
                           const VecT* HWY_RESTRICT vec_aligned,
                           float* HWY_RESTRICT out) {
  MatVecAddLoop<false, kOuter, kInner>(
      mat, mat_ofs, vec_aligned, /*add=*/(VecT*)nullptr, out);
}

// Simple version without tiling nor threading, but two offsets/outputs.
template <bool kAdd, size_t kOuter, size_t kInner, typename ArrayT,
          typename VecT, typename AddT>
HWY_INLINE void TwoOfsMatVecAddLoop(const ArrayT& mat, const size_t mat_ofs0,

Review comment (performance):
Allocation can be quite slow, let's move this into gemma.cc's Activations. That would require plumbing through an extra tmp arg here, and the std::array storage should probably be the largest per-call size * max number of threads (say 128 or 256). Would you prefer if I made this change?

Code after:
  const hn::ScalableTag<float> df;

  for (size_t idx_row = 0; idx_row < kOuter; ++idx_row) {
    const size_t row_ofs = mat_ofs + idx_row * kInner;
    if constexpr (kAdd) {
      out[idx_row] = hwy::ConvertScalarTo<float>(add[idx_row]) +
                     Dot(df, mat, row_ofs, vec_aligned, kInner);
    } else {
      out[idx_row] = Dot(df, mat, row_ofs, vec_aligned, kInner);
    }
  }
}

#if !defined(HWY_NATIVE_DOT_BF16) || !HWY_NATIVE_DOT_BF16
template <bool kAdd, size_t kOuter, size_t kInner, typename VecT, typename AddT,
          size_t kCapacity>
HWY_INLINE void MatVecAddLoop(
    const CompressedArray<hwy::bfloat16_t, kCapacity>& mat,
    const size_t mat_ofs,
    const VecT* HWY_RESTRICT vec_aligned,
    const AddT* HWY_RESTRICT add, float* HWY_RESTRICT even_odd,
    float* HWY_RESTRICT out) {
  PROFILER_ZONE("MatVecAddLoop");

  // Sanity check: we can write without race conditions.
  if (HWY_IS_TSAN) {
    even_odd[0] = hwy::ConvertScalarTo<float>(vec_aligned[0]);
    even_odd[kInner - 1] = -even_odd[0];
  }

  const hn::ScalableTag<float> df;
  ToEvenOddF32(vec_aligned, kInner, even_odd);
  for (size_t idx_row = 0; idx_row < kOuter; ++idx_row) {
    const size_t row_ofs = mat_ofs + idx_row * kInner;
    if constexpr (kAdd) {
      out[idx_row] = hwy::ConvertScalarTo<float>(add[idx_row]) +
                     Dot<true>(df, mat, row_ofs, even_odd, kInner);
    } else {
      out[idx_row] = Dot<true>(df, mat, row_ofs, even_odd, kInner);
    }
  }
}
#endif

// even_odd is precomputed for the current thread.
template <size_t kOuter, size_t kInner, typename ArrayT, typename VecT>
HWY_INLINE void MatVecLoop(const ArrayT& mat, const size_t mat_ofs,
                           const VecT* HWY_RESTRICT vec_aligned,
                           float* HWY_RESTRICT even_odd,
                           float* HWY_RESTRICT out) {
  MatVecAddLoop</*kAdd=*/false, kOuter, kInner>(