Skip to content

fix(rocsparse): write the base and use n for csxsldu_preprocess column walks (AISPARSE-906) #14480

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kliegeois merged 1 commit into from
Oct 28, 2026
Merged

fix(rocsparse): write the base and use n for csxsldu_preprocess column walks (AISPARSE-906) #14480

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
203 changes: 201 additions & 2 deletions projects/rocsparse/clients/unittests/unit_test_host_internal_branches.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -930,8 +930,8 @@ TEST_F(HostInternalBranches, argdescr_getters_and_verbose)
EXPECT_EQ(rocsparse_create_handle(nullptr), rocsparse_status_invalid_pointer);
}

// Buffer-size has three direction combinations. Preprocess's two mixed
// diag-type launches are not on the itilu0 path.
// Buffer-size has three direction combinations. Preprocess's column walk and
// the two mixed diag-type launches are not on the itilu0 path.
TEST_F(HostInternalBranches, csxsldu_directions)
{
constexpr auto row = rocsparse_direction_row;
Expand Down Expand Up @@ -997,6 +997,205 @@ TEST_F(HostInternalBranches, csxsldu_directions)
EXPECT_EQ(unnz, 2);
EXPECT_EQ(lptr, (std::vector<rocsparse_int>{0, 0, 0}));
EXPECT_EQ(uptr, (std::vector<rocsparse_int>{0, 1, 2}));

// Column direction with column factors: both factors in the input direction, so the
// pointers go to lptr / uptr directly. Every element of lptr / uptr starts at the
// sentinel -7, including element 0. A zero handle buffer makes preprocess allocate
// its scratch.
auto* raw = reinterpret_cast<_rocsparse_handle*>(handle);
size_t saved_buffer = raw->buffer_size;
raw->buffer_size = 0;
lnnz = unnz = -1;
EXPECT_EQ(sldu_preprocess(handle,
column,
rocsparse_diag_type_unit,
column,
rocsparse_diag_type_unit,
column,
scratch.ptr,
&lnnz,
&unnz,
&lptr,
&uptr),
rocsparse_status_success);
raw->buffer_size = saved_buffer;
EXPECT_EQ(lnnz, 0);
EXPECT_EQ(unnz, 0);
EXPECT_EQ(lptr, (std::vector<rocsparse_int>{0, 0, 0}));
EXPECT_EQ(uptr, (std::vector<rocsparse_int>{0, 0, 0}));

// Same walk with a non-unit lower factor: the diagonal goes to L.
lnnz = unnz = -1;
EXPECT_EQ(sldu_preprocess(handle,
column,
rocsparse_diag_type_non_unit,
column,
rocsparse_diag_type_unit,
column,
scratch.ptr,
&lnnz,
&unnz,
&lptr,
&uptr),
rocsparse_status_success);
EXPECT_EQ(lnnz, 2);
EXPECT_EQ(unnz, 0);
EXPECT_EQ(lptr, (std::vector<rocsparse_int>{0, 1, 2}));
EXPECT_EQ(uptr, (std::vector<rocsparse_int>{0, 0, 0}));
ASSERT_EQ(hipDeviceSynchronize(), hipSuccess);
}

namespace
{
struct sldu_case
{
rocsparse_direction dir;
rocsparse_int m, n;
std::vector<rocsparse_int> ptr, ind;
rocsparse_index_base base;
rocsparse_diag_type ldiag;
rocsparse_direction ldir;
rocsparse_index_base lbase;
rocsparse_diag_type udiag;
rocsparse_direction udir;
rocsparse_index_base ubase;
};

// Host reference of the L / U pointers that preprocess computes in the input
// direction: one entry per row (row direction) or column (column direction) plus
// one, starting at the factor's index base. At most one diagonal type is non-unit,
// and the diagonal goes to that factor.
void sldu_reference(const sldu_case& c,
std::vector<rocsparse_int>& lptr,
std::vector<rocsparse_int>& uptr)
{
const rocsparse_int nseq = (c.dir == rocsparse_direction_row) ? c.m : c.n;
lptr.assign(1, c.lbase);
uptr.assign(1, c.ubase);
for(rocsparse_int s = 0; s < nseq; ++s)
{
rocsparse_int l = 0, u = 0;
for(rocsparse_int k = c.ptr[s] - c.base; k < c.ptr[s + 1] - c.base; ++k)
{
const rocsparse_int other = c.ind[k] - c.base;
const rocsparse_int row = (c.dir == rocsparse_direction_row) ? s : other;
const rocsparse_int col = (c.dir == rocsparse_direction_row) ? other : s;
if(row > col)
++l;
else if(row < col)
++u;
else if(c.ldiag == rocsparse_diag_type_non_unit)
++l;
else if(c.udiag == rocsparse_diag_type_non_unit)
++u;
}
lptr.push_back(lptr.back() + l);
uptr.push_back(uptr.back() + u);
}
}
} // namespace

// Non-square inputs, one-based indices, distinct L / U bases, strictly lower and upper
// entries, and factors in the other direction (their pointers go to the scratch buffer,
// L first, U after it). Every pointer and both counts are compared with the host
// reference; pointer arrays preprocess must not write keep their -7 sentinel.
TEST_F(HostInternalBranches, csxsldu_preprocess_reference)
{
constexpr auto row = rocsparse_direction_row;
constexpr auto column = rocsparse_direction_column;
constexpr auto unit = rocsparse_diag_type_unit;
constexpr auto nonu = rocsparse_diag_type_non_unit;
constexpr auto zero = rocsparse_index_base_zero;
constexpr auto one = rocsparse_index_base_one;

// 2x3: (0,0) (0,1) (0,2) / (1,0) (1,1) (1,2), stored by column.
const std::vector<rocsparse_int> p23{0, 2, 4, 6}, i23{0, 1, 0, 1, 0, 1};
// The same matrix one-based.
const std::vector<rocsparse_int> p23b{1, 3, 5, 7}, i23b{1, 2, 1, 2, 1, 2};
// 3x2: (0,0) / (1,0) (1,1) / (2,0), stored by column.
const std::vector<rocsparse_int> p32{0, 3, 4}, i32{0, 1, 2, 1};
// 3x2: (0,0) / (1,0) (1,1) / (2,0) (2,1), stored by row, one-based.
const std::vector<rocsparse_int> r32b{1, 2, 4, 6}, j32b{1, 1, 2, 1, 2};

const sldu_case cases[] = {
{column, 2, 3, p23, i23, zero, unit, column, zero, nonu, column, zero},
{column, 3, 2, p32, i32, zero, nonu, column, zero, unit, column, zero},
{column, 2, 3, p23b, i23b, one, unit, column, one, nonu, column, zero},
{column, 2, 3, p23, i23, zero, unit, row, one, nonu, row, one},
{column, 2, 3, p23, i23, zero, nonu, column, zero, unit, row, one},
{column, 2, 3, p23b, i23b, one, nonu, row, one, unit, column, zero},
{row, 3, 2, r32b, j32b, one, unit, row, one, nonu, column, one},
// 3x0 by column: one pointer and no entries; ind holds a dummy so that the
// device allocation is not empty.
{column, 3, 0, {1}, {1}, one, unit, column, one, nonu, column, zero},
};

int index = 0;
for(const sldu_case& c : cases)
{
SCOPED_TRACE("case " + std::to_string(index++));
const rocsparse_int nseq = (c.dir == row) ? c.m : c.n;
const rocsparse_int nnz = static_cast<rocsparse_int>(c.ind.size());
const std::vector<rocsparse_int> sentinel(nseq + 1, -7);
device_vector<rocsparse_int> ptr{c.ptr};
device_vector<rocsparse_int> ind{c.ind};
device_vector<float> val{std::vector<float>(nnz, 1.0f)};
device_vector<rocsparse_int> lptr{sentinel};
device_vector<rocsparse_int> uptr{sentinel};
device_vector<rocsparse_int> scratch{std::vector<rocsparse_int>(2 * (nseq + 1), -7)};
ASSERT_TRUE(ptr.ptr && ind.ptr && val.ptr && lptr.ptr && uptr.ptr && scratch.ptr);

rocsparse_int lnnz = -1, unnz = -1;
const rocsparse_status status
= rocsparse::csxsldu_preprocess_template<float, rocsparse_int, rocsparse_int>(
handle,
c.dir,
c.m,
c.n,
nnz,
ptr.ptr,
ind.ptr,
val.ptr,
c.base,
c.ldiag,
c.ldir,
&lnnz,
lptr.ptr,
c.lbase,
c.udiag,
c.udir,
&unnz,
uptr.ptr,
c.ubase,
scratch.ptr);
EXPECT_EQ(status, rocsparse_status_success);
if(status != rocsparse_status_success)
continue;

std::vector<rocsparse_int> want_l, want_u;
sldu_reference(c, want_l, want_u);
const bool l_scratch = (c.ldir != c.dir);
const bool u_scratch = (c.udir != c.dir);
const size_t u_offset = l_scratch ? nseq + 1 : 0;
const auto got_l = l_scratch ? to_host(scratch.ptr, nseq + 1) : to_host(lptr);
const auto got_u = u_scratch ? to_host(scratch.ptr + u_offset, nseq + 1) : to_host(uptr);
EXPECT_EQ(got_l, want_l);
EXPECT_EQ(got_u, want_u);
EXPECT_EQ(lnnz, want_l.back() - c.lbase);
EXPECT_EQ(unnz, want_u.back() - c.ubase);
if(l_scratch)
EXPECT_EQ(to_host(lptr), sentinel);
if(u_scratch)
EXPECT_EQ(to_host(uptr), sentinel);
}

// Both diagonals non-unit is rejected by the count kernel dispatch, in both directions.
for(const rocsparse_direction d : {row, column})
{
rocsparse_int lnnz = -1, unnz = -1;
EXPECT_EQ(sldu_preprocess(handle, d, nonu, d, nonu, d, nullptr, &lnnz, &unnz),
rocsparse_status_invalid_value);
}
ASSERT_EQ(hipDeviceSynchronize(), hipSuccess);
}

Expand Down
28 changes: 15 additions & 13 deletions projects/rocsparse/library/src/conversion/rocsparse_csxsldu_preprocess.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,9 +159,7 @@ namespace rocsparse
{
case rocsparse_diag_type_non_unit:
{
// LCOV_EXCL_START
RETURN_IF_ROCSPARSE_ERROR(rocsparse_status_invalid_value);
// LCOV_EXCL_STOP
}
case rocsparse_diag_type_unit:
{
Expand Down Expand Up @@ -223,19 +221,21 @@ rocsparse_status rocsparse::csxsldu_preprocess_template(rocsparse_handle han
: buffer)
: uptr_;

RETURN_IF_HIP_ERROR(
rocsparse_hipMemcpyAsync(lptr, &lbase, sizeof(I), hipMemcpyHostToDevice, handle_->stream));
RETURN_IF_HIP_ERROR(
rocsparse_hipMemcpyAsync(uptr, &ubase, sizeof(I), hipMemcpyHostToDevice, handle_->stream));
RETURN_IF_HIP_ERROR(rocsparse_hipStreamSynchronize(handle_->stream));

switch(dir_)
{
case rocsparse_direction_row:
{
RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync(
lptr, &lbase, sizeof(I), hipMemcpyHostToDevice, handle_->stream));
RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync(
uptr, &ubase, sizeof(I), hipMemcpyHostToDevice, handle_->stream));
RETURN_IF_HIP_ERROR(rocsparse_hipStreamSynchronize(handle_->stream));
J nblocks = (m_ - 1) / nthreads_per_block + 1;
dim3 blocks(nblocks);
rocsparse::csxtril_count_kernel_dispatch<nthreads_per_block, I, J>(
handle_, blocks, threads, ldiag_, udiag_, m_, ptr_, ind_, base_, lptr, uptr);
RETURN_IF_ROCSPARSE_ERROR(
(rocsparse::csxtril_count_kernel_dispatch<nthreads_per_block, I, J>(
handle_, blocks, threads, ldiag_, udiag_, m_, ptr_, ind_, base_, lptr, uptr)));
RETURN_IF_ROCSPARSE_ERROR(rocsparse::inclusive_scan(handle_, m_, lptr));
RETURN_IF_ROCSPARSE_ERROR(rocsparse::inclusive_scan(handle_, m_, uptr));
break;
Expand All @@ -245,18 +245,20 @@ rocsparse_status rocsparse::csxsldu_preprocess_template(rocsparse_handle han
{
J nblocks = (n_ - 1) / nthreads_per_block + 1;
dim3 blocks(nblocks);
rocsparse::csxtril_count_kernel_dispatch<nthreads_per_block, I, J>(
handle_, blocks, threads, udiag_, ldiag_, n_, ptr_, ind_, base_, uptr, lptr);
RETURN_IF_ROCSPARSE_ERROR(
(rocsparse::csxtril_count_kernel_dispatch<nthreads_per_block, I, J>(
handle_, blocks, threads, udiag_, ldiag_, n_, ptr_, ind_, base_, uptr, lptr)));
RETURN_IF_ROCSPARSE_ERROR(rocsparse::inclusive_scan(handle_, n_, lptr));
RETURN_IF_ROCSPARSE_ERROR(rocsparse::inclusive_scan(handle_, n_, uptr));
break;
}
}

const J nseq = (dir_ == rocsparse_direction_row) ? m_ : n_;
RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync(
host_lnnz_, &lptr[m_], sizeof(I), hipMemcpyDeviceToHost, handle_->stream));
host_lnnz_, &lptr[nseq], sizeof(I), hipMemcpyDeviceToHost, handle_->stream));
RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync(
host_unnz_, &uptr[m_], sizeof(I), hipMemcpyDeviceToHost, handle_->stream));
host_unnz_, &uptr[nseq], sizeof(I), hipMemcpyDeviceToHost, handle_->stream));
RETURN_IF_HIP_ERROR(rocsparse_hipStreamSynchronize(handle_->stream));

host_lnnz_[0] -= lbase;
Expand Down
Loading