reduce Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/SpmmReduceKernel.cpp lines 26–64

template <typename scalar_t, typename index_t, ReductionType reduce>
inline void _update(at::opmath_type<scalar_t>* out_ptr, int64_t e, int64_t c, const scalar_t val, const scalar_t* other_data, int64_t K) {
  using opmath_t = at::opmath_type<scalar_t>;
  using Vec = vec::Vectorized<scalar_t>;
  using aVec = VecType<scalar_t>;
  constexpr int64_t kVecSize = Vec::size();
  constexpr int64_t kVLEN = kVecSize * 4;

  int64_t k = 0;
  aVec val_vec = aVec((opmath_t)val);
  const scalar_t* other_ptr = other_data + c * K;

  for (; k < K - (K % kVLEN); k += kVLEN) {
    aVec out_vec0 = aVec::loadu(out_ptr + k);
    aVec out_vec1 = aVec::loadu(out_ptr + k + kVecSize);
    aVec out_vec2 = aVec::loadu(out_ptr + k + kVecSize * 2);
    aVec out_vec3 = aVec::loadu(out_ptr + k + kVecSize * 3);

    out_vec0 = update<aVec, reduce>(out_vec0, aVec::loadu(other_ptr + k) * val_vec);
    out_vec1 = update<aVec, reduce>(out_vec1, aVec::loadu(other_ptr + k + kVecSize) * val_vec);
    out_vec2 = update<aVec, reduce>(out_vec2, aVec::loadu(other_ptr + k + kVecSize * 2) * val_vec);
    out_vec3 = update<aVec, reduce>(out_vec3, aVec::loadu(other_ptr + k + kVecSize * 3) * val_vec);

    out_vec0.store(out_ptr + k);
    out_vec1.store(out_ptr + k + kVecSize);
    out_vec2.store(out_ptr + k + kVecSize * 2);
    out_vec3.store(out_ptr + k + kVecSize * 3);
  }
  for (; k < K - (K % kVecSize); k += kVecSize) {
    aVec out_vec = aVec::loadu(out_ptr + k);
    out_vec = update<aVec, reduce>(out_vec, aVec::loadu(other_ptr + k) * val_vec);
    out_vec.store(out_ptr + k);
  }
  for (; k < K; k++) {
    opmath_t out_val = opmath_t(out_ptr[k]);
    out_val = update<opmath_t, reduce>(out_val, opmath_t(other_ptr[k]) * opmath_t(val));
    out_ptr[k] = out_val;
  }
}