Skip to content

Commit

Permalink
[BUG] Fix SPMM strided view (#2124)
Browse files Browse the repository at this point in the history
With the bug fix #2117 there can be an issue with `z_tmp` memory being uninitialized.
SPMM formula is `Z = alpha . X * Y + beta . Z` so when `beta` is not zero, Z is being read.
The proposed solution in this PR remove the need for an extra allocation and a copy from/to an external buffer, by creating a strided view of the original Z.

Authors:
  - Micka (https://github.com/lowener)

Approvers:
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #2124
  • Loading branch information
lowener authored Jan 24, 2024
1 parent dbb5e66 commit 8743ac9
Showing 1 changed file with 2 additions and 6 deletions.
8 changes: 2 additions & 6 deletions cpp/include/raft/sparse/linalg/spmm.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ namespace linalg {
* @param[in] x input raft::device_csr_matrix_view
* @param[in] y input raft::device_matrix_view
* @param[in] beta scalar
* @param[out] z output raft::device_matrix_view
* @param[inout] z input-output raft::device_matrix_view
*/
template <typename ValueType,
typename IndexType,
Expand All @@ -60,19 +60,15 @@ void spmm(raft::resources const& handle,
{
bool is_row_major = detail::is_row_major(y, z);

auto z_tmp = raft::make_device_matrix<ValueType, IndexType>(handle, z.extent(0), z.extent(1));
auto z_tmp_view = raft::make_device_strided_matrix_view<ValueType, IndexType, LayoutPolicyZ>(
z_tmp.data_handle(), z.extent(0), z.extent(1), is_row_major ? z.stride(0) : z.stride(1));
z.data_handle(), z.extent(0), z.extent(1), is_row_major ? z.stride(0) : z.stride(1));

auto descr_x = detail::create_descriptor(x);
auto descr_y = detail::create_descriptor(y);
auto descr_z = detail::create_descriptor(z_tmp_view);

detail::spmm(handle, trans_x, trans_y, is_row_major, alpha, descr_x, descr_y, beta, descr_z);

raft::copy(
z.data_handle(), z_tmp.data_handle(), z_tmp.size(), raft::resource::get_cuda_stream(handle));

RAFT_CUSPARSE_TRY_NO_THROW(cusparseDestroySpMat(descr_x));
RAFT_CUSPARSE_TRY_NO_THROW(cusparseDestroyDnMat(descr_y));
RAFT_CUSPARSE_TRY_NO_THROW(cusparseDestroyDnMat(descr_z));
Expand Down

0 comments on commit 8743ac9

Please sign in to comment.