addmv_out_sparse_csr Class — pytorch Architecture

Source Code

aten/src/ATen/native/sparse/SparseBlasImpl.cpp lines 359–392

template<typename scalar_t, typename idx_t>
void addmv_out_sparse_csr(
    const Tensor& mat,
    const Tensor& vec,
    const Scalar& beta,
    const Scalar& alpha,
    const Tensor& result) {
  auto cont_values = mat.values().contiguous();
  if (mat.layout() == kSparseBsr) {
    addmv_sparse_bsr(cont_values.data_ptr<scalar_t>(),
        mat.crow_indices().data_ptr<idx_t>(),
        mat.col_indices().data_ptr<idx_t>(),
        mat.size(0),
        mat.values().size(1),
        mat.values().size(2),
        vec.data_ptr<scalar_t>(),
        vec.stride(0),
        alpha.to<scalar_t>(),
        beta.to<scalar_t>(),
        result.data_ptr<scalar_t>(),
        result.stride(0));
  } else {
    addmv_sparse_csr(cont_values.data_ptr<scalar_t>(),
        mat.crow_indices().data_ptr<idx_t>(),
        mat.col_indices().data_ptr<idx_t>(),
        mat.size(0),
        vec.data_ptr<scalar_t>(),
        vec.stride(0),
        alpha.to<scalar_t>(),
        beta.to<scalar_t>(),
        result.data_ptr<scalar_t>(),
        result.stride(0));
  }
}