addmv_sparse_csr Class — pytorch Architecture

Source Code

aten/src/ATen/native/sparse/SparseBlasImpl.cpp lines 305–326

template<typename scalar_t, typename idx_t>
void addmv_sparse_csr(
    const scalar_t* mat_values,
    const idx_t* crow_index,
    const idx_t* col_index,
    const int64_t mat_rows,
    const scalar_t* vec,
    const size_t vec_stride,
    const scalar_t alpha,
    const scalar_t beta,
    scalar_t* result,
    const size_t result_stride) {
  at::parallel_for(0, mat_rows, 0, [&](int64_t rstart, int64_t rend) {
    for(const auto row: c10::irange(rstart, rend)) {
      scalar_t acc(0);
      for(const auto idx: c10::irange(crow_index[row], crow_index[row + 1])) {
        acc += mat_values[idx] * vec[col_index[idx] * vec_stride];
      }
      result[row * result_stride] = acc * alpha + result[row * result_stride] * beta;
    }
  });
}