vectorized_inner_reduction Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/Reduce.h lines 80–91

template <typename func_t, typename vec_func_t>
inline void vectorized_inner_reduction(char** data, int64_t n, func_t op, vec_func_t vop) {
  VEC_LOOP_HEADER(func_t, data)
  constexpr int64_t vector_stride = 4 * Vec::size() * sizeof(scalar_t);
  int64_t count = n / (4 * Vec::size());
  if (count > 0) {
    vectorized_reduction(data, count, vector_stride, op, vop, /*reduce=*/true);
  }
  char* ptrs[3] = { data[0], data[0], data[1] };
  int64_t strides[] = { 0, 0, sizeof(scalar_t) };
  basic_loop(ptrs, strides, count * 4 * Vec::size(), n, op);
}