map_acc Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/ReduceUtils.h lines 145–180

template <typename scalar_t, typename accumut, typename Op,
          typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 0>
inline void map_acc(
    const Op& vec_fun,
    accumut* output_data,
    const accumut* input_data,
    const scalar_t* input_data2,
    int64_t size) {
  using Vec = vec::Vectorized<scalar_t>;
  using aVec = vec::Vectorized<accumut>;
  int64_t d = 0;
  constexpr int64_t kVecSize = Vec::size();
  constexpr int64_t kaVecSize = aVec::size();
  for (d = 0; d < size - (size % kVecSize); d += kVecSize) {
    Vec data2_vec = Vec::loadu(input_data2 + d);
    auto [data2_avec0, data2_avec1] = convert_to_float<scalar_t>(data2_vec);
    aVec input_vec0 = aVec::loadu(input_data + d);
    aVec input_vec1 = aVec::loadu(input_data + d + kaVecSize);
    vec_fun(input_vec0, data2_avec0).store(output_data + d);
    vec_fun(input_vec1, data2_avec1).store(output_data + d + kaVecSize);
  }
  if (size - d > 0) {
    int64_t tail_size = size - d;
    Vec data2_vec = Vec::loadu(input_data2 + d, tail_size);
    auto [data2_avec0, data2_avec1] = convert_to_float<scalar_t>(data2_vec);
    if (tail_size > kaVecSize) {
      aVec input_vec0 = aVec::loadu(input_data + d);
      aVec input_vec1 = aVec::loadu(input_data + d + kaVecSize, tail_size - kaVecSize);
      vec_fun(input_vec0, data2_avec0).store(output_data + d);
      vec_fun(input_vec1, data2_avec1).store(output_data + d + kaVecSize, tail_size - kaVecSize);
    } else {
      aVec input_vec0 = aVec::loadu(input_data + d, tail_size);
      vec_fun(input_vec0, data2_avec0).store(output_data + d, tail_size);
    }
  }
}