sum_norm_per_row Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/WeightNormKernel.cpp lines 65–87

template <typename scalar_t>
inline std::enable_if_t<is_reduced_floating_point_v<scalar_t>, void>
sum_norm_per_row(
    float* out_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 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) + v_fvec0 * v_fvec0;
    fVec out_fvec1 = fVec::loadu(out_ptr + d + fVec::size()) + v_fvec1 * v_fvec1;
    out_fvec0.store(out_ptr + d);
    out_fvec1.store(out_ptr + d + fVec::size());
  }
  for(; d < size; ++d) {
    float v_val = float(v_ptr[d]);
    out_ptr[d] += v_val * v_val;
  }
}