sum_product_per_row Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/WeightNormKernel.cpp lines 247–273

template <typename scalar_t>
inline std::enable_if_t<is_reduced_floating_point_v<scalar_t>, void>
sum_product_per_row(
    float* out_ptr,
    const scalar_t* grad_w_ptr,
    const scalar_t* v_ptr,
    int64_t size) {
  using bVec = vec::Vectorized<scalar_t>;
  using fVec = vec::Vectorized<float>;
  int64_t d = 0;
  for (; d < size - (size % bVec::size()); d += bVec::size()) {
    bVec grad_w_bvec = bVec::loadu(grad_w_ptr + d);
    auto [grad_w_fvec0, grad_w_fvec1] = vec::convert_to_float<scalar_t>(grad_w_bvec);
    bVec v_bvec = bVec::loadu(v_ptr + d);
    auto [v_fvec0, v_fvec1] = vec::convert_to_float<scalar_t>(v_bvec);

    fVec out_fvec0 = fVec::loadu(out_ptr + d) + grad_w_fvec0 * v_fvec0;
    fVec out_fvec1 = fVec::loadu(out_ptr + d + fVec::size()) + grad_w_fvec1 * v_fvec1;
    out_fvec0.store(out_ptr + d);
    out_fvec1.store(out_ptr + d + fVec::size());
  }
  for(; d < size; ++d) {
    float grad_w_val = float(grad_w_ptr[d]);
    float v_val = float(v_ptr[d]);
    out_ptr[d] += grad_w_val * v_val;
  }
}