addmv_sparse_bsr Class — pytorch Architecture

Source Code

aten/src/ATen/native/sparse/SparseBlasImpl.cpp lines 328–357

template<typename scalar_t, typename idx_t>
void addmv_sparse_bsr(
    const scalar_t* mat_values,
    const idx_t* crow_index,
    const idx_t* col_index,
    const int64_t mat_rows,
    const int64_t blocksize_rows,
    const int64_t blocksize_cols,
    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)) {
      const auto block_row = row / blocksize_rows;
      const auto block_row_offset = row % blocksize_rows;
      scalar_t acc(0);
      for(const auto block_idx: c10::irange(crow_index[block_row], crow_index[block_row + 1])) {
        const auto block_offs = (block_idx * blocksize_rows + block_row_offset) * blocksize_cols;
        const auto vec_offs = col_index[block_idx]* blocksize_cols;
        for(const auto idx: c10::irange(blocksize_cols)) {
          acc += mat_values[block_offs + idx] * vec[(vec_offs + idx) * vec_stride];
        }
      }
      result[row * result_stride] = acc * alpha + result[row * result_stride] * beta;
    }
  });
}