From 737b658c44668d88cc0a12470c3609556e381faf Mon Sep 17 00:00:00 2001 From: jsandham Date: Fri, 25 Sep 2026 19:49:28 +0000 Subject: [PATCH 1/8] Implement rocsparse_spsort for CSR matrices Sorts the column indices, and the corresponding values, within each row of a CSR matrix, out of place or in place, and for strided batched matrices. The direction must be rocsparse_direction_row. Adds rocsparse::csrsort_buffer_size and rocsparse::csrsort, which take the arrays of both the input and output matrices and loop over the batches. The existing csrsort implementation is turned into templates on the row pointer and column index types, which the rocsparse_csrsort C API now calls. The segmented radix sort gains the mixed 32-bit key and 64-bit value instantiations needed for 64-bit row pointers with 32-bit column indices. Co-authored-by: Cursor --- .../generic/example_rocsparse_spsort.cpp | 179 ++++++ .../internal/generic/rocsparse_spsort.h | 2 + .../library/src/conversion/csrsort_device.h | 6 +- .../src/conversion/rocsparse_csrsort.cpp | 602 ++++++++++++++---- .../src/conversion/rocsparse_csrsort.hpp | 69 ++ .../src/conversion/rocsparse_spsort.cpp | 42 +- .../rocsparse_spsort_buffer_size.cpp | 8 +- .../rocsparse_segmented_radix_sort_keys.cpp | 2 + .../rocsparse_segmented_radix_sort_pairs.cpp | 2 + 9 files changed, 782 insertions(+), 130 deletions(-) create mode 100644 projects/rocsparse/clients/samples/documentation_examples/generic/example_rocsparse_spsort.cpp create mode 100644 projects/rocsparse/library/src/conversion/rocsparse_csrsort.hpp diff --git a/projects/rocsparse/clients/samples/documentation_examples/generic/example_rocsparse_spsort.cpp b/projects/rocsparse/clients/samples/documentation_examples/generic/example_rocsparse_spsort.cpp new file mode 100644 index 000000000000..89e662e7c51f --- /dev/null +++ b/projects/rocsparse/clients/samples/documentation_examples/generic/example_rocsparse_spsort.cpp @@ -0,0 +1,179 @@ +/* ************************************************************************ + * Copyright (C) 2026 Advanced Micro Devices, Inc. All rights Reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + * + * ************************************************************************ */ + +#include +#include + +#include + +#define HIP_CHECK(stat) \ + { \ + if(stat != hipSuccess) \ + { \ + std::cerr << "Error: hip error " << stat << " in line " << __LINE__ << std::endl; \ + return -1; \ + } \ + } + +#define ROCSPARSE_CHECK(stat) \ + { \ + if(stat != rocsparse_status_success) \ + { \ + std::cerr << "Error: rocsparse error " << stat << " in line " << __LINE__ \ + << std::endl; \ + return -1; \ + } \ + } + +//! [doc example] +int main() +{ + // 1 2 3 4 0 5 + // A = 0 2 3 0 0 4 + // 5 0 6 7 8 0 + // 1 0 9 0 6 7 + int m = 4; + int n = 6; + + std::vector hcsr_row_ptr = {0, 5, 8, 12, 16}; + // std::vector hcsr_col_ind = {0, 1, 2, 3, 5, + // 1, 2, 5, + // 0, 2, 3, 4, + // 0, 2, 4, 5}; + // std::vector hcsr_val = {1, 2, 3, 4, 5, 2, 3, 4, 5, 6, 7, 8, 1, 9, 6, 7}; + std::vector hcsr_col_ind = {5, 3, 2, 1, 0, 1, 5, 2, 4, 2, 3, 0, 0, 5, 4, 2}; + std::vector hcsr_val = {5, 4, 3, 2, 1, 2, 4, 3, 8, 6, 7, 5, 1, 7, 6, 9}; + + int nnz = hcsr_row_ptr[m] - hcsr_row_ptr[0]; + + // Offload data to device + int* dcsr_row_ptr; + int* dcsr_col_ind; + float* dcsr_val; + HIP_CHECK(hipMalloc(&dcsr_row_ptr, sizeof(int) * (m + 1))); + HIP_CHECK(hipMalloc(&dcsr_col_ind, sizeof(int) * nnz)); + HIP_CHECK(hipMalloc(&dcsr_val, sizeof(float) * nnz)); + + HIP_CHECK( + hipMemcpy(dcsr_row_ptr, hcsr_row_ptr.data(), sizeof(int) * (m + 1), hipMemcpyHostToDevice)); + HIP_CHECK( + hipMemcpy(dcsr_col_ind, hcsr_col_ind.data(), sizeof(int) * nnz, hipMemcpyHostToDevice)); + HIP_CHECK(hipMemcpy(dcsr_val, hcsr_val.data(), sizeof(float) * nnz, hipMemcpyHostToDevice)); + + rocsparse_handle handle; + rocsparse_error p_error[1] = {}; + rocsparse_spmat_descr matA; + + rocsparse_indextype row_idx_type = rocsparse_indextype_i32; + rocsparse_indextype col_idx_type = rocsparse_indextype_i32; + rocsparse_datatype data_type = rocsparse_datatype_f32_r; + rocsparse_index_base idx_base = rocsparse_index_base_zero; + + ROCSPARSE_CHECK(rocsparse_create_handle(&handle)); + + // Create sparse matrix A + ROCSPARSE_CHECK(rocsparse_create_csr_descr(&matA, + m, + n, + nnz, + dcsr_row_ptr, + dcsr_col_ind, + dcsr_val, + row_idx_type, + col_idx_type, + idx_base, + data_type)); + + rocsparse_spsort_descr spsort_descr; + ROCSPARSE_CHECK(rocsparse_spsort_descr_create(handle, &spsort_descr, p_error)); + + const rocsparse_spsort_alg spsort_alg = rocsparse_spsort_alg_default; + ROCSPARSE_CHECK(rocsparse_spsort_set_input(handle, + spsort_descr, + rocsparse_spsort_input_alg, + &spsort_alg, + sizeof(spsort_alg), + p_error)); + + // Sort the column indices within each row + const rocsparse_direction spsort_direction = rocsparse_direction_row; + ROCSPARSE_CHECK(rocsparse_spsort_set_input(handle, + spsort_descr, + rocsparse_spsort_input_direction, + &spsort_direction, + sizeof(spsort_direction), + p_error)); + + // Passing matA as both the input and the output sorts it in place + // Call spsort to get buffer size + size_t buffer_size; + ROCSPARSE_CHECK(rocsparse_spsort_buffer_size( + handle, spsort_descr, matA, matA, rocsparse_spsort_stage_analysis, &buffer_size, p_error)); + + void* buffer; + HIP_CHECK(hipMalloc(&buffer, buffer_size)); + + // Call spsort to perform analysis + ROCSPARSE_CHECK(rocsparse_spsort(handle, + spsort_descr, + matA, + matA, + rocsparse_spsort_stage_analysis, + buffer_size, + buffer, + p_error)); + + HIP_CHECK(hipFree(buffer)); + + ROCSPARSE_CHECK(rocsparse_spsort_buffer_size( + handle, spsort_descr, matA, matA, rocsparse_spsort_stage_compute, &buffer_size, p_error)); + + HIP_CHECK(hipMalloc(&buffer, buffer_size)); + + // Call spsort to perform computation + ROCSPARSE_CHECK(rocsparse_spsort(handle, + spsort_descr, + matA, + matA, + rocsparse_spsort_stage_compute, + buffer_size, + buffer, + p_error)); + + HIP_CHECK(hipFree(buffer)); + + ROCSPARSE_CHECK(rocsparse_destroy_error(p_error[0])); + ROCSPARSE_CHECK(rocsparse_spsort_descr_destroy(handle, spsort_descr, nullptr)); + + // Clear rocSPARSE + ROCSPARSE_CHECK(rocsparse_destroy_spmat_descr(matA)); + ROCSPARSE_CHECK(rocsparse_destroy_handle(handle)); + + // Clear device memory + HIP_CHECK(hipFree(dcsr_row_ptr)); + HIP_CHECK(hipFree(dcsr_col_ind)); + HIP_CHECK(hipFree(dcsr_val)); + + return 0; +} +//! [doc example] diff --git a/projects/rocsparse/library/include/internal/generic/rocsparse_spsort.h b/projects/rocsparse/library/include/internal/generic/rocsparse_spsort.h index 15ec759e614b..531d26175bf2 100644 --- a/projects/rocsparse/library/include/internal/generic/rocsparse_spsort.h +++ b/projects/rocsparse/library/include/internal/generic/rocsparse_spsort.h @@ -143,6 +143,8 @@ rocsparse_status rocsparse_spsort_buffer_size(rocsparse_handle handle * \retval rocsparse_status_not_implemented \p source is not in COO, CSR or CSC format. * * \par Example +* \snippet example_rocsparse_spsort.cpp doc example +* \par Example * \snippet example_rocsparse_spsort_coo.cpp doc example */ ROCSPARSE_EXPORT diff --git a/projects/rocsparse/library/src/conversion/csrsort_device.h b/projects/rocsparse/library/src/conversion/csrsort_device.h index fd5cbe6d26d5..b8424f7caddb 100644 --- a/projects/rocsparse/library/src/conversion/csrsort_device.h +++ b/projects/rocsparse/library/src/conversion/csrsort_device.h @@ -29,11 +29,11 @@ namespace rocsparse { // Shift CSR offsets - template + template ROCSPARSE_KERNEL(BLOCKSIZE) - void csrsort_shift_kernel(rocsparse_int size, const rocsparse_int* in, rocsparse_int* out) + void csrsort_shift_kernel(int64_t size, const I* in, I* out) { - rocsparse_int gid = hipBlockIdx_x * BLOCKSIZE + hipThreadIdx_x; + const int64_t gid = int64_t(hipBlockIdx_x) * BLOCKSIZE + hipThreadIdx_x; if(gid >= size) { diff --git a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp index 64a609e325d4..4b86ee5acbf8 100644 --- a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp +++ b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp @@ -24,11 +24,200 @@ #include "internal/conversion/rocsparse_csrsort.h" #include "rocsparse_utility.hpp" +#include "../level1/rocsparse_gthr.hpp" #include "csrsort_device.h" #include "rocsparse_control.hpp" - +#include "rocsparse_csrsort.hpp" +#include "rocsparse_gcreate_identity_permutation.hpp" #include "rocsparse_primitives.hpp" +namespace rocsparse +{ + // Number of bits needed to represent the column indices, which are at most n. + static uint32_t csrsort_endbit(int64_t n) + { + // __builtin_clzll is undefined for n == 0 + return (n == 0) ? 0 : 64 - __builtin_clzll(static_cast(n)); + } + + template + static rocsparse_status csrsort_buffer_size_template(rocsparse_handle handle, + int64_t m, + int64_t n, + int64_t nnz, + const I* csr_row_ptr, + const J* csr_col_ind, + size_t* buffer_size) + { + ROCSPARSE_ROUTINE_TRACE; + + if(m == 0 || n == 0 || nnz == 0) + { + *buffer_size = 0; + return rocsparse_status_success; + } + + const uint32_t startbit = 0; + const uint32_t endbit = rocsparse::csrsort_endbit(n); + + // We do not know if sort_pairs or sort_keys will be called, so use the largest buffer between the two + size_t size1; + size_t size2; + RETURN_IF_ROCSPARSE_ERROR( + (rocsparse::primitives::segmented_radix_sort_pairs_buffer_size( + handle, nnz, m, startbit, endbit, &size1))); + RETURN_IF_ROCSPARSE_ERROR( + (rocsparse::primitives::segmented_radix_sort_keys_buffer_size( + handle, nnz, m, startbit, endbit, &size2))); + + *buffer_size = rocsparse::align_size(rocsparse::max(size1, size2)); + + // rocPRIM does not support in-place sorting, so we need additional buffer + // for all temporary arrays + + // columns buffer + *buffer_size += rocsparse::align_size(nnz); + // perm buffer + *buffer_size += rocsparse::align_size(nnz); + // segm buffer + *buffer_size += rocsparse::align_size(m + 1); + + return rocsparse_status_success; + } + + // Sorts the column indices within each row. If perm is not null, the sort applies its + // reordering to perm. + template + static rocsparse_status csrsort_template(rocsparse_handle handle, + int64_t m, + int64_t n, + int64_t nnz, + rocsparse_index_base idx_base, + const I* csr_row_ptr, + J* csr_col_ind, + I* perm, + void* temp_buffer) + { + ROCSPARSE_ROUTINE_TRACE; + + // Quick return if possible + if(m == 0 || n == 0 || nnz == 0) + { + return rocsparse_status_success; + } + + // Stream + hipStream_t stream = handle->stream; + + const uint32_t startbit = 0; + const uint32_t endbit = rocsparse::csrsort_endbit(n); + size_t size; + + if(perm != nullptr) + { + // Sort pairs, if permutation vector is present + RETURN_IF_ROCSPARSE_ERROR( + (rocsparse::primitives::segmented_radix_sort_pairs_buffer_size( + handle, nnz, m, startbit, endbit, &size))); + } + else + { + // Sort keys, if no permutation vector is present + RETURN_IF_ROCSPARSE_ERROR( + (rocsparse::primitives::segmented_radix_sort_keys_buffer_size( + handle, nnz, m, startbit, endbit, &size))); + } + + // Temporary buffer entry points + char* ptr = reinterpret_cast(temp_buffer); + + // columns buffer + J* tmp_cols = reinterpret_cast(ptr); + ptr += rocsparse::align_size(nnz); + + // perm buffer + I* tmp_perm = reinterpret_cast(ptr); + ptr += rocsparse::align_size(nnz); + + // segm buffer + I* tmp_segm = reinterpret_cast(ptr); + ptr += rocsparse::align_size(m + 1); + + // Index base one requires shift of offset positions + if(idx_base == rocsparse_index_base_one) + { +#define CSRSORT_DIM 512 + dim3 csrsort_blocks((m + 1 - 1) / CSRSORT_DIM + 1); + dim3 csrsort_threads(CSRSORT_DIM); + + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR((rocsparse::csrsort_shift_kernel), + csrsort_blocks, + csrsort_threads, + 0, + stream, + m + 1, + csr_row_ptr, + tmp_segm); +#undef CSRSORT_DIM + } + + // rocprim buffer + void* tmp_rocprim = reinterpret_cast(ptr); + + // Switch between offsets + const I* offsets = (idx_base == rocsparse_index_base_one) ? tmp_segm : csr_row_ptr; + + // Sort by columns and obtain permutation vector + + if(perm != nullptr) + { + // Sort by pairs, if permutation vector is present + rocsparse::primitives::double_buffer keys(csr_col_ind, tmp_cols); + rocsparse::primitives::double_buffer vals(perm, tmp_perm); + + RETURN_IF_ROCSPARSE_ERROR( + rocsparse::primitives::segmented_radix_sort_pairs(handle, + keys, + vals, + nnz, + m, + offsets, + offsets + 1, + startbit, + endbit, + size, + tmp_rocprim)); + + if(keys.current() != csr_col_ind) + { + RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync( + csr_col_ind, keys.current(), sizeof(J) * nnz, hipMemcpyDeviceToDevice, stream)); + } + if(vals.current() != perm) + { + RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync( + perm, vals.current(), sizeof(I) * nnz, hipMemcpyDeviceToDevice, stream)); + } + } + else + { + // Sort by keys, if no permutation vector is present + rocsparse::primitives::double_buffer keys(csr_col_ind, tmp_cols); + + RETURN_IF_ROCSPARSE_ERROR(rocsparse::primitives::segmented_radix_sort_keys( + handle, keys, nnz, m, offsets, offsets + 1, startbit, endbit, size, tmp_rocprim)); + + if(keys.current() != csr_col_ind) + { + RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync( + csr_col_ind, keys.current(), sizeof(J) * nnz, hipMemcpyDeviceToDevice, stream)); + } + } + return rocsparse_status_success; + } + +} + extern "C" rocsparse_status rocsparse_csrsort_buffer_size(rocsparse_handle handle, rocsparse_int m, rocsparse_int n, @@ -58,39 +247,9 @@ try ROCSPARSE_CHECKARG_ARRAY(5, nnz, csr_col_ind); ROCSPARSE_CHECKARG_POINTER(6, buffer_size); - if(m == 0 || n == 0 || nnz == 0) - { - *buffer_size = 0; - return rocsparse_status_success; - } - - uint32_t startbit = 0; - uint32_t endbit = rocsparse::clz(n); - - // We do not know if sort_pairs or sort_keys will be called, so use the largest buffer between the two - size_t size1; - size_t size2; - RETURN_IF_ROCSPARSE_ERROR( - (rocsparse::primitives:: - segmented_radix_sort_pairs_buffer_size( - handle, nnz, m, startbit, endbit, &size1))); RETURN_IF_ROCSPARSE_ERROR( - (rocsparse::primitives::segmented_radix_sort_keys_buffer_size( - handle, nnz, m, startbit, endbit, &size2))); - - *buffer_size = rocsparse::max(size1, size2); - - *buffer_size = ((*buffer_size - 1) / 256 + 1) * 256; - - // rocPRIM does not support in-place sorting, so we need additional buffer - // for all temporary arrays - - // columns buffer - *buffer_size += ((sizeof(rocsparse_int) * nnz - 1) / 256 + 1) * 256; - // perm buffer - *buffer_size += ((sizeof(rocsparse_int) * nnz - 1) / 256 + 1) * 256; - // segm buffer - *buffer_size += ((sizeof(rocsparse_int) * (m + 1)) / 256 + 1) * 256; + (rocsparse::csrsort_buffer_size_template( + handle, m, n, nnz, csr_row_ptr, csr_col_ind, buffer_size))); return rocsparse_status_success; // LCOV_EXCL_START @@ -135,125 +294,320 @@ try ROCSPARSE_CHECKARG_ARRAY(6, nnz, csr_col_ind); ROCSPARSE_CHECKARG_ARRAY(8, nnz, temp_buffer); - // Quick return if possible - if(m == 0 || n == 0 || nnz == 0) - { - return rocsparse_status_success; - } + RETURN_IF_ROCSPARSE_ERROR((rocsparse::csrsort_template( + handle, m, n, nnz, descr->base, csr_row_ptr, csr_col_ind, perm, temp_buffer))); - // Stream - hipStream_t stream = handle->stream; + return rocsparse_status_success; + // LCOV_EXCL_START +} +catch(...) +{ + RETURN_ROCSPARSE_EXCEPTION(); +} +// LCOV_EXCL_STOP - uint32_t startbit = 0; - uint32_t endbit = rocsparse::clz(n); - size_t size; +namespace rocsparse +{ + static rocsparse_status gcsrsort_buffer_size(rocsparse_handle handle, + int64_t m, + int64_t n, + int64_t nnz, + rocsparse_indextype ptr_type, + rocsparse_indextype ind_type, + const void* ptr, + const void* ind, + size_t* buffer_size) + { + if(ptr_type == rocsparse_indextype_i32 && ind_type == rocsparse_indextype_i32) + { + RETURN_IF_ROCSPARSE_ERROR( + rocsparse::csrsort_buffer_size_template(handle, + m, + n, + nnz, + reinterpret_cast(ptr), + reinterpret_cast(ind), + buffer_size)); + return rocsparse_status_success; + } + if(ptr_type == rocsparse_indextype_i64 && ind_type == rocsparse_indextype_i32) + { + RETURN_IF_ROCSPARSE_ERROR( + rocsparse::csrsort_buffer_size_template(handle, + m, + n, + nnz, + reinterpret_cast(ptr), + reinterpret_cast(ind), + buffer_size)); + return rocsparse_status_success; + } + if(ptr_type == rocsparse_indextype_i64 && ind_type == rocsparse_indextype_i64) + { + RETURN_IF_ROCSPARSE_ERROR( + rocsparse::csrsort_buffer_size_template(handle, + m, + n, + nnz, + reinterpret_cast(ptr), + reinterpret_cast(ind), + buffer_size)); + return rocsparse_status_success; + } + RETURN_WITH_MESSAGE_IF_ROCSPARSE_ERROR( + rocsparse_status_invalid_value, + "the combination of offsets and indices types is not supported"); + } - if(perm != nullptr) + static rocsparse_status gcsrsort(rocsparse_handle handle, + int64_t m, + int64_t n, + int64_t nnz, + rocsparse_index_base idx_base, + rocsparse_indextype ptr_type, + rocsparse_indextype ind_type, + const void* ptr, + void* ind, + void* perm, + void* buffer) { - // Sort pairs, if permutation vector is present - RETURN_IF_ROCSPARSE_ERROR(( - rocsparse::primitives:: - segmented_radix_sort_pairs_buffer_size( - handle, nnz, m, startbit, endbit, &size))); + if(ptr_type == rocsparse_indextype_i32 && ind_type == rocsparse_indextype_i32) + { + RETURN_IF_ROCSPARSE_ERROR( + rocsparse::csrsort_template(handle, + m, + n, + nnz, + idx_base, + reinterpret_cast(ptr), + reinterpret_cast(ind), + reinterpret_cast(perm), + buffer)); + return rocsparse_status_success; + } + if(ptr_type == rocsparse_indextype_i64 && ind_type == rocsparse_indextype_i32) + { + RETURN_IF_ROCSPARSE_ERROR( + rocsparse::csrsort_template(handle, + m, + n, + nnz, + idx_base, + reinterpret_cast(ptr), + reinterpret_cast(ind), + reinterpret_cast(perm), + buffer)); + return rocsparse_status_success; + } + if(ptr_type == rocsparse_indextype_i64 && ind_type == rocsparse_indextype_i64) + { + RETURN_IF_ROCSPARSE_ERROR( + rocsparse::csrsort_template(handle, + m, + n, + nnz, + idx_base, + reinterpret_cast(ptr), + reinterpret_cast(ind), + reinterpret_cast(perm), + buffer)); + return rocsparse_status_success; + } + RETURN_WITH_MESSAGE_IF_ROCSPARSE_ERROR( + rocsparse_status_invalid_value, + "the combination of offsets and indices types is not supported"); } - else + + // The buffer starts with the permutation array, which tracks where each entry moves + // while the indices are sorted, followed by scratch space shared by the index sort + // and the value permutation. + static size_t csxsort_perm_size(int64_t nnz, rocsparse_indextype perm_indextype) { - // Sort keys, if no permutation vector is present - RETURN_IF_ROCSPARSE_ERROR( - (rocsparse::primitives::segmented_radix_sort_keys_buffer_size( - handle, nnz, m, startbit, endbit, &size))); + return rocsparse::align_size(rocsparse::indextype_sizeof(perm_indextype) * nnz); } +} - // Temporary buffer entry points - char* ptr = reinterpret_cast(temp_buffer); +rocsparse_status rocsparse::csxsort_buffer_size(rocsparse_handle handle, + rocsparse_direction dir, + rocsparse_const_spmat_descr source, + rocsparse_const_spmat_descr target, + size_t* buffer_size_in_bytes) +{ + ROCSPARSE_ROUTINE_TRACE; - // columns buffer - rocsparse_int* tmp_cols = reinterpret_cast(ptr); - ptr += ((sizeof(rocsparse_int) * nnz - 1) / 256 + 1) * 256; + // A CSC matrix is sorted as the CSR layout of its transpose. + const bool by_row = (dir == rocsparse_direction_row); + const int64_t m = by_row ? target->rows : target->cols; + const int64_t n = by_row ? target->cols : target->rows; + const int64_t nnz = target->nnz; + const rocsparse_indextype ptr_type = by_row ? target->row_type : target->col_type; + const rocsparse_indextype ind_type = by_row ? target->col_type : target->row_type; + const void* ptr = by_row ? target->const_row_data : target->const_col_data; + const void* ind = by_row ? target->const_col_data : target->const_row_data; + + size_t sort_buffer_size = 0; + RETURN_IF_ROCSPARSE_ERROR(rocsparse::gcsrsort_buffer_size( + handle, m, n, nnz, ptr_type, ind_type, ptr, ind, &sort_buffer_size)); + + // Values sorted in place are gathered into scratch space first, since the gather cannot + // write over its own input. + const size_t gather_buffer_size + = (target->const_val_data == source->const_val_data) + ? rocsparse::align_size(rocsparse::datatype_sizeof(target->data_type) * nnz) + : 0; + + *buffer_size_in_bytes = rocsparse::csxsort_perm_size(nnz, ptr_type) + + rocsparse::max(sort_buffer_size, gather_buffer_size); - // perm buffer - rocsparse_int* tmp_perm = reinterpret_cast(ptr); - ptr += ((sizeof(rocsparse_int) * nnz - 1) / 256 + 1) * 256; + return rocsparse_status_success; +} - // segm buffer - rocsparse_int* tmp_segm = reinterpret_cast(ptr); - ptr += ((sizeof(rocsparse_int) * (m + 1)) / 256 + 1) * 256; +rocsparse_status rocsparse::csxsort(rocsparse_handle handle, + rocsparse_direction dir, + rocsparse_const_spmat_descr source, + rocsparse_spmat_descr target, + size_t buffer_size_in_bytes, + void* buffer) +{ + ROCSPARSE_ROUTINE_TRACE; - // Index base one requires shift of offset positions - if(descr->base == rocsparse_index_base_one) + size_t required_buffer_size; + RETURN_IF_ROCSPARSE_ERROR( + rocsparse::csxsort_buffer_size(handle, dir, source, target, &required_buffer_size)); + if(buffer_size_in_bytes < required_buffer_size) { -#define CSRSORT_DIM 512 - dim3 csrsort_blocks(m / CSRSORT_DIM + 1); - dim3 csrsort_threads(CSRSORT_DIM); - - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR((rocsparse::csrsort_shift_kernel), - csrsort_blocks, - csrsort_threads, - 0, - stream, - m + 1, - csr_row_ptr, - tmp_segm); -#undef CSRSORT_DIM + RETURN_WITH_MESSAGE_IF_ROCSPARSE_ERROR( + rocsparse_status_invalid_size, + "the buffer is smaller than the size returned by the buffer size query"); } - // rocprim buffer - void* tmp_rocprim = reinterpret_cast(ptr); - - // Switch between offsets - const rocsparse_int* offsets = descr->base == rocsparse_index_base_one ? tmp_segm : csr_row_ptr; - - // Sort by columns and obtain permutation vector - - if(perm != nullptr) + // A CSC matrix is sorted as the CSR layout of its transpose. + const bool by_row = (dir == rocsparse_direction_row); + const int64_t m = by_row ? target->rows : target->cols; + const int64_t n = by_row ? target->cols : target->rows; + const int64_t nnz = target->nnz; + const rocsparse_index_base idx_base = target->idx_base; + const rocsparse_indextype ptr_type = by_row ? target->row_type : target->col_type; + const rocsparse_indextype ind_type = by_row ? target->col_type : target->row_type; + const rocsparse_datatype data_type = target->data_type; + + const void* ptr_source = by_row ? source->const_row_data : source->const_col_data; + const void* ind_source = by_row ? source->const_col_data : source->const_row_data; + void* ptr_target = by_row ? target->row_data : target->col_data; + void* ind_target = by_row ? target->col_data : target->row_data; + + const size_t ptr_size = rocsparse::indextype_sizeof(ptr_type); + const size_t ind_size = rocsparse::indextype_sizeof(ind_type); + const size_t val_size = rocsparse::datatype_sizeof(data_type); + + void* perm = buffer; + void* sort_buffer + = reinterpret_cast(buffer) + rocsparse::csxsort_perm_size(nnz, ptr_type); + + // The batches run one after the other on the handle stream, so they share the buffer. + for(int64_t batch = 0; batch < target->batch_count; ++batch) { - // Sort by pairs, if permutation vector is present - rocsparse::primitives::double_buffer keys(csr_col_ind, tmp_cols); - rocsparse::primitives::double_buffer vals(perm, tmp_perm); - - RETURN_IF_ROCSPARSE_ERROR(rocsparse::primitives::segmented_radix_sort_pairs( - handle, keys, vals, nnz, m, offsets, offsets + 1, startbit, endbit, size, tmp_rocprim)); - - if(keys.current() != csr_col_ind) + const int64_t offsets_offset_source = batch * source->offsets_batch_stride; + const int64_t offsets_offset_target = batch * target->offsets_batch_stride; + const int64_t cv_offset_source = batch * source->columns_values_batch_stride; + const int64_t cv_offset_target = batch * target->columns_values_batch_stride; + + const void* batch_ptr_source + = reinterpret_cast(ptr_source) + offsets_offset_source * ptr_size; + const void* batch_ind_source + = reinterpret_cast(ind_source) + cv_offset_source * ind_size; + const void* batch_val_source + = reinterpret_cast(source->const_val_data) + cv_offset_source * val_size; + void* batch_ptr_target + = reinterpret_cast(ptr_target) + offsets_offset_target * ptr_size; + void* batch_ind_target = reinterpret_cast(ind_target) + cv_offset_target * ind_size; + void* batch_val_target + = reinterpret_cast(target->val_data) + cv_offset_target * val_size; + + // The index sort works in place, so the offsets and indices of source are first copied + // into target. A matrix without rows may have a null offsets array. + if(m > 0 && batch_ptr_target != batch_ptr_source) { - RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync(csr_col_ind, - keys.current(), - sizeof(rocsparse_int) * nnz, + RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync(batch_ptr_target, + batch_ptr_source, + ptr_size * (m + 1), hipMemcpyDeviceToDevice, - stream)); + handle->stream)); } - if(vals.current() != perm) + if(batch_ind_target != batch_ind_source) { - RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync(perm, - vals.current(), - sizeof(rocsparse_int) * nnz, + RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync(batch_ind_target, + batch_ind_source, + ind_size * nnz, hipMemcpyDeviceToDevice, - stream)); + handle->stream)); } - } - else - { - // Sort by keys, if no permutation vector is present - rocsparse::primitives::double_buffer keys(csr_col_ind, tmp_cols); - - RETURN_IF_ROCSPARSE_ERROR(rocsparse::primitives::segmented_radix_sort_keys( - handle, keys, nnz, m, offsets, offsets + 1, startbit, endbit, size, tmp_rocprim)); - if(keys.current() != csr_col_ind) + // The index sort applies its reordering to perm, so it must start as the identity. + RETURN_IF_ROCSPARSE_ERROR( + rocsparse::gcreate_identity_permutation(handle, nnz, ptr_type, perm)); + + RETURN_IF_ROCSPARSE_ERROR(rocsparse::gcsrsort(handle, + m, + n, + nnz, + idx_base, + ptr_type, + ind_type, + batch_ptr_target, + batch_ind_target, + perm, + sort_buffer)); + + // The gather cannot write over its own input, so in place values go through scratch. + const bool in_place_val = (batch_val_target == batch_val_source); + void* sorted_val = in_place_val ? sort_buffer : batch_val_target; + RETURN_IF_ROCSPARSE_ERROR(rocsparse::gthr(handle, + nnz, + data_type, + batch_val_source, + data_type, + sorted_val, + ptr_type, + perm, + rocsparse_index_base_zero)); + + if(in_place_val) { - RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync(csr_col_ind, - keys.current(), - sizeof(rocsparse_int) * nnz, + RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync(batch_val_target, + sorted_val, + val_size * nnz, hipMemcpyDeviceToDevice, - stream)); + handle->stream)); } } + return rocsparse_status_success; - // LCOV_EXCL_START } -catch(...) + +rocsparse_status rocsparse::csrsort_buffer_size(rocsparse_handle handle, + rocsparse_csrsort_alg alg, + rocsparse_const_spmat_descr source, + rocsparse_const_spmat_descr target, + size_t* buffer_size_in_bytes) { - RETURN_ROCSPARSE_EXCEPTION(); + ROCSPARSE_ROUTINE_TRACE; + + RETURN_IF_ROCSPARSE_ERROR(rocsparse::csxsort_buffer_size( + handle, rocsparse_direction_row, source, target, buffer_size_in_bytes)); + return rocsparse_status_success; +} + +rocsparse_status rocsparse::csrsort(rocsparse_handle handle, + rocsparse_csrsort_alg alg, + rocsparse_const_spmat_descr source, + rocsparse_spmat_descr target, + size_t buffer_size_in_bytes, + void* buffer) +{ + ROCSPARSE_ROUTINE_TRACE; + + RETURN_IF_ROCSPARSE_ERROR(rocsparse::csxsort( + handle, rocsparse_direction_row, source, target, buffer_size_in_bytes, buffer)); + return rocsparse_status_success; } -// LCOV_EXCL_STOP diff --git a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.hpp b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.hpp new file mode 100644 index 000000000000..32fa6ecd51d1 --- /dev/null +++ b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.hpp @@ -0,0 +1,69 @@ +/*! \file */ +/* ************************************************************************ + * Copyright (C) 2026 Advanced Micro Devices, Inc. All rights Reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + * + * ************************************************************************ */ + +#pragma once + +#include "rocsparse_handle.hpp" + +namespace rocsparse +{ + typedef enum rocsparse_csrsort_alg_ + { + rocsparse_csrsort_alg_default = 0 + } rocsparse_csrsort_alg; + + // Shared by the CSR and CSC sorts: dir == rocsparse_direction_row sorts the column indices + // within each row of a CSR matrix, and dir == rocsparse_direction_column sorts the row + // indices within each column of a CSC matrix. + rocsparse_status csxsort_buffer_size(rocsparse_handle handle, + rocsparse_direction dir, + rocsparse_const_spmat_descr source, + rocsparse_const_spmat_descr target, + size_t* buffer_size_in_bytes); + + rocsparse_status csxsort(rocsparse_handle handle, + rocsparse_direction dir, + rocsparse_const_spmat_descr source, + rocsparse_spmat_descr target, + size_t buffer_size_in_bytes, + void* buffer); + + rocsparse_status csrsort_buffer_size(rocsparse_handle handle, + rocsparse_csrsort_alg alg, + rocsparse_const_spmat_descr source, + rocsparse_const_spmat_descr target, + size_t* buffer_size_in_bytes); + + // Sorts the column indices and values within each row of the CSR matrix source into the + // CSR matrix target, which must have the same sizes, types, index base and batch layout as + // source, and whose row pointer receives a copy of the row pointer of source. Each array of + // target may either alias the matching array of source, in which case it is sorted in + // place, or not overlap it at all. Each batch is sorted independently. + rocsparse_status csrsort(rocsparse_handle handle, + rocsparse_csrsort_alg alg, + rocsparse_const_spmat_descr source, + rocsparse_spmat_descr target, + size_t buffer_size_in_bytes, + void* buffer); +} diff --git a/projects/rocsparse/library/src/conversion/rocsparse_spsort.cpp b/projects/rocsparse/library/src/conversion/rocsparse_spsort.cpp index c5c77dc7c52e..8bafd0867e89 100644 --- a/projects/rocsparse/library/src/conversion/rocsparse_spsort.cpp +++ b/projects/rocsparse/library/src/conversion/rocsparse_spsort.cpp @@ -26,6 +26,7 @@ #include "rocsparse_utility.hpp" #include "rocsparse_coosort.hpp" +#include "rocsparse_csrsort.hpp" #include "rocsparse_spsort.hpp" namespace rocsparse @@ -65,6 +66,9 @@ namespace rocsparse } case rocsparse_format_csr: case rocsparse_format_csc: + { + return mat->offsets_batch_stride == 0 && mat->columns_values_batch_stride == 0; + } case rocsparse_format_coo_aos: case rocsparse_format_bsr: case rocsparse_format_ell: @@ -89,6 +93,18 @@ namespace rocsparse } case rocsparse_format_csr: case rocsparse_format_csc: + { + // The offsets may be shared by all samples with a stride of zero. + const int64_t offsets_size + = ((mat->format == rocsparse_format_csr) ? mat->rows : mat->cols) + 1; + ROCSPARSE_CHECKARG( + arg, + mat, + (mat->columns_values_batch_stride < mat->nnz + || (mat->offsets_batch_stride != 0 && mat->offsets_batch_stride < offsets_size)), + rocsparse_status_invalid_size); + return rocsparse_status_success; + } case rocsparse_format_coo_aos: case rocsparse_format_bsr: case rocsparse_format_ell: @@ -155,6 +171,19 @@ rocsparse_status rocsparse::spsort_check_arguments(rocsparse_spsort_descr d && rocsparse::enum_utils::is_invalid(descr->get_direction())), rocsparse_status_invalid_value); + // A CSR matrix can only have the column indices within each row sorted, and a CSC matrix + // can only have the row indices within each column sorted. + ROCSPARSE_CHECKARG(1, + descr, + (source->format == rocsparse_format_csr + && descr->get_direction() != rocsparse_direction_row), + rocsparse_status_invalid_value); + ROCSPARSE_CHECKARG(1, + descr, + (source->format == rocsparse_format_csc + && descr->get_direction() != rocsparse_direction_column), + rocsparse_status_invalid_value); + return rocsparse_status_success; } @@ -181,12 +210,12 @@ namespace rocsparse switch(format) { case rocsparse_format_coo: + case rocsparse_format_csr: { return rocsparse_status_success; } // LCOV_EXCL_START - case rocsparse_format_csr: case rocsparse_format_csc: case rocsparse_format_coo_aos: case rocsparse_format_bsr: @@ -212,9 +241,18 @@ namespace rocsparse rocsparse::coosort(handle, dir, source, target, buffer_size_in_bytes, buffer)); return rocsparse_status_success; } + case rocsparse_format_csr: + { + RETURN_IF_ROCSPARSE_ERROR(rocsparse::csrsort(handle, + rocsparse_csrsort_alg_default, + source, + target, + buffer_size_in_bytes, + buffer)); + return rocsparse_status_success; + } // LCOV_EXCL_START - case rocsparse_format_csr: case rocsparse_format_csc: case rocsparse_format_coo_aos: case rocsparse_format_bsr: diff --git a/projects/rocsparse/library/src/conversion/rocsparse_spsort_buffer_size.cpp b/projects/rocsparse/library/src/conversion/rocsparse_spsort_buffer_size.cpp index 4c874c0e451e..1dc799c13b2a 100644 --- a/projects/rocsparse/library/src/conversion/rocsparse_spsort_buffer_size.cpp +++ b/projects/rocsparse/library/src/conversion/rocsparse_spsort_buffer_size.cpp @@ -25,6 +25,7 @@ #include "rocsparse_utility.hpp" #include "rocsparse_coosort.hpp" +#include "rocsparse_csrsort.hpp" #include "rocsparse_spsort.hpp" namespace rocsparse @@ -47,11 +48,11 @@ namespace rocsparse switch(format) { case rocsparse_format_coo: + case rocsparse_format_csr: { *buffer_size_in_bytes = 0; return rocsparse_status_success; } - case rocsparse_format_csr: case rocsparse_format_csc: case rocsparse_format_coo_aos: case rocsparse_format_bsr: @@ -74,6 +75,11 @@ namespace rocsparse return rocsparse_status_success; } case rocsparse_format_csr: + { + RETURN_IF_ROCSPARSE_ERROR(rocsparse::csrsort_buffer_size( + handle, rocsparse_csrsort_alg_default, source, target, buffer_size_in_bytes)); + return rocsparse_status_success; + } case rocsparse_format_csc: case rocsparse_format_coo_aos: case rocsparse_format_bsr: diff --git a/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_keys.cpp b/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_keys.cpp index 67f3aa4988a9..a92a465ef106 100644 --- a/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_keys.cpp +++ b/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_keys.cpp @@ -232,6 +232,8 @@ rocsparse_status rocsparse::primitives::sort_csr_column_indices(rocsparse_handle INSTANTIATE(int32_t, int32_t); INSTANTIATE(int64_t, int64_t); +INSTANTIATE(int32_t, int64_t); +INSTANTIATE(int32_t, const int64_t); #undef INSTANTIATE #define INSTANTIATE(ITYPE, JTYPE) \ diff --git a/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_pairs.cpp b/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_pairs.cpp index 70acc483d558..edec179ca8b5 100644 --- a/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_pairs.cpp +++ b/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_pairs.cpp @@ -134,4 +134,6 @@ INSTANTIATE(int32_t, int32_t, const int32_t); INSTANTIATE(int64_t, int64_t, const int64_t); INSTANTIATE(int32_t, int32_t, int32_t); INSTANTIATE(int64_t, int64_t, int64_t); +INSTANTIATE(int32_t, int64_t, const int64_t); +INSTANTIATE(int32_t, int64_t, int64_t); #undef INSTANTIATE From 243a4372c56978248ef2cf8647d315c0c0e4d8e0 Mon Sep 17 00:00:00 2001 From: jsandham Date: Wed, 30 Sep 2026 13:27:38 +0000 Subject: [PATCH 2/8] Sort uniform CSR and CSC batches with a single index sort Only uniform batches are supported, so every batch shares the sparsity pattern of the first one. Its indices are sorted once, copied with the offsets into every batch of the target, and the resulting permutation is applied to the values of each batch. The batch strides are converted to bytes once, up front. Co-authored-by: Cursor --- .../src/conversion/rocsparse_csrsort.cpp | 109 ++++++++++-------- .../src/conversion/rocsparse_csrsort.hpp | 3 +- 2 files changed, 62 insertions(+), 50 deletions(-) diff --git a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp index 4b86ee5acbf8..d329eb237491 100644 --- a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp +++ b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp @@ -491,73 +491,84 @@ rocsparse_status rocsparse::csxsort(rocsparse_handle handle, const rocsparse_indextype ind_type = by_row ? target->col_type : target->row_type; const rocsparse_datatype data_type = target->data_type; - const void* ptr_source = by_row ? source->const_row_data : source->const_col_data; - const void* ind_source = by_row ? source->const_col_data : source->const_row_data; - void* ptr_target = by_row ? target->row_data : target->col_data; - void* ind_target = by_row ? target->col_data : target->row_data; + const char* ptr_source + = reinterpret_cast(by_row ? source->const_row_data : source->const_col_data); + const char* ind_source + = reinterpret_cast(by_row ? source->const_col_data : source->const_row_data); + const char* val_source = reinterpret_cast(source->const_val_data); + char* ptr_target = reinterpret_cast(by_row ? target->row_data : target->col_data); + char* ind_target = reinterpret_cast(by_row ? target->col_data : target->row_data); + char* val_target = reinterpret_cast(target->val_data); const size_t ptr_size = rocsparse::indextype_sizeof(ptr_type); const size_t ind_size = rocsparse::indextype_sizeof(ind_type); const size_t val_size = rocsparse::datatype_sizeof(data_type); + const size_t val_stride_source = source->columns_values_batch_stride * val_size; + const size_t ptr_stride_target = target->offsets_batch_stride * ptr_size; + const size_t ind_stride_target = target->columns_values_batch_stride * ind_size; + const size_t val_stride_target = target->columns_values_batch_stride * val_size; + void* perm = buffer; void* sort_buffer = reinterpret_cast(buffer) + rocsparse::csxsort_perm_size(nnz, ptr_type); + // Only uniform batches are supported, so every batch has the sparsity pattern of the first + // one. Its indices are sorted once, and the resulting permutation is applied to the values + // of every batch. + + // The index sort works in place, so the offsets and indices of source are first copied into + // target. A matrix without rows may have a null offsets array. + if(m > 0 && ptr_target != ptr_source) + { + RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync( + ptr_target, ptr_source, ptr_size * (m + 1), hipMemcpyDeviceToDevice, handle->stream)); + } + if(ind_target != ind_source) + { + RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync( + ind_target, ind_source, ind_size * nnz, hipMemcpyDeviceToDevice, handle->stream)); + } + + // The index sort applies its reordering to perm, so it must start as the identity. + RETURN_IF_ROCSPARSE_ERROR(rocsparse::gcreate_identity_permutation(handle, nnz, ptr_type, perm)); + + RETURN_IF_ROCSPARSE_ERROR(rocsparse::gcsrsort(handle, + m, + n, + nnz, + idx_base, + ptr_type, + ind_type, + ptr_target, + ind_target, + perm, + sort_buffer)); + // The batches run one after the other on the handle stream, so they share the buffer. for(int64_t batch = 0; batch < target->batch_count; ++batch) { - const int64_t offsets_offset_source = batch * source->offsets_batch_stride; - const int64_t offsets_offset_target = batch * target->offsets_batch_stride; - const int64_t cv_offset_source = batch * source->columns_values_batch_stride; - const int64_t cv_offset_target = batch * target->columns_values_batch_stride; - - const void* batch_ptr_source - = reinterpret_cast(ptr_source) + offsets_offset_source * ptr_size; - const void* batch_ind_source - = reinterpret_cast(ind_source) + cv_offset_source * ind_size; - const void* batch_val_source - = reinterpret_cast(source->const_val_data) + cv_offset_source * val_size; - void* batch_ptr_target - = reinterpret_cast(ptr_target) + offsets_offset_target * ptr_size; - void* batch_ind_target = reinterpret_cast(ind_target) + cv_offset_target * ind_size; - void* batch_val_target - = reinterpret_cast(target->val_data) + cv_offset_target * val_size; - - // The index sort works in place, so the offsets and indices of source are first copied - // into target. A matrix without rows may have a null offsets array. - if(m > 0 && batch_ptr_target != batch_ptr_source) + if(batch > 0) { - RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync(batch_ptr_target, - batch_ptr_source, - ptr_size * (m + 1), - hipMemcpyDeviceToDevice, - handle->stream)); - } - if(batch_ind_target != batch_ind_source) - { - RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync(batch_ind_target, - batch_ind_source, + // The offsets of the target are shared by all its batches when their stride is zero. + char* batch_ptr_target = ptr_target + batch * ptr_stride_target; + if(m > 0 && batch_ptr_target != ptr_target) + { + RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync(batch_ptr_target, + ptr_target, + ptr_size * (m + 1), + hipMemcpyDeviceToDevice, + handle->stream)); + } + RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync(ind_target + batch * ind_stride_target, + ind_target, ind_size * nnz, hipMemcpyDeviceToDevice, handle->stream)); } - // The index sort applies its reordering to perm, so it must start as the identity. - RETURN_IF_ROCSPARSE_ERROR( - rocsparse::gcreate_identity_permutation(handle, nnz, ptr_type, perm)); - - RETURN_IF_ROCSPARSE_ERROR(rocsparse::gcsrsort(handle, - m, - n, - nnz, - idx_base, - ptr_type, - ind_type, - batch_ptr_target, - batch_ind_target, - perm, - sort_buffer)); + const char* batch_val_source = val_source + batch * val_stride_source; + char* batch_val_target = val_target + batch * val_stride_target; // The gather cannot write over its own input, so in place values go through scratch. const bool in_place_val = (batch_val_target == batch_val_source); diff --git a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.hpp b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.hpp index 32fa6ecd51d1..e6adc1bef769 100644 --- a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.hpp +++ b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.hpp @@ -59,7 +59,8 @@ namespace rocsparse // CSR matrix target, which must have the same sizes, types, index base and batch layout as // source, and whose row pointer receives a copy of the row pointer of source. Each array of // target may either alias the matching array of source, in which case it is sorted in - // place, or not overlap it at all. Each batch is sorted independently. + // place, or not overlap it at all. The batches must be uniform: they all share the sparsity + // pattern of the first batch, and only their values differ. rocsparse_status csrsort(rocsparse_handle handle, rocsparse_csrsort_alg alg, rocsparse_const_spmat_descr source, From be6684268dcb1245cf0463d2f2c3bfab91ed1dd2 Mon Sep 17 00:00:00 2001 From: jsandham Date: Wed, 30 Sep 2026 16:06:39 +0000 Subject: [PATCH 3/8] Replace the templated CSR sort with non-templated compute routines Replace csrsort_template and the gcsrsort wrappers with non-templated csrsort_buffer_size_compute and csrsort_compute routines that take the offsets and indices types at runtime, keeping the type dispatch only at the rocPRIM and kernel calls. Pass the buffer size to the index sort so it can check the buffer is large enough, and remove an unused segmented_radix_sort_keys instantiation. Co-authored-by: Cursor --- .../src/conversion/rocsparse_csrsort.cpp | 617 +++++++++++------- .../rocsparse_segmented_radix_sort_keys.cpp | 1 - 2 files changed, 369 insertions(+), 249 deletions(-) diff --git a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp index d329eb237491..69ebb084e12b 100644 --- a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp +++ b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp @@ -40,20 +40,268 @@ namespace rocsparse return (n == 0) ? 0 : 64 - __builtin_clzll(static_cast(n)); } + static bool csrsort_is_supported(rocsparse_indextype ptr_type, rocsparse_indextype ind_type) + { + return (ptr_type == rocsparse_indextype_i32 && ind_type == rocsparse_indextype_i32) + || (ptr_type == rocsparse_indextype_i64 && ind_type == rocsparse_indextype_i32) + || (ptr_type == rocsparse_indextype_i64 && ind_type == rocsparse_indextype_i64); + } + template - static rocsparse_status csrsort_buffer_size_template(rocsparse_handle handle, - int64_t m, - int64_t n, - int64_t nnz, - const I* csr_row_ptr, - const J* csr_col_ind, - size_t* buffer_size) + static rocsparse_status csrsort_rocprim_buffer_size_template(rocsparse_handle handle, + int64_t m, + int64_t nnz, + uint32_t startbit, + uint32_t endbit, + bool with_perm, + size_t* buffer_size) + { + if(with_perm) + { + RETURN_IF_ROCSPARSE_ERROR( + (rocsparse::primitives::segmented_radix_sort_pairs_buffer_size( + handle, nnz, m, startbit, endbit, buffer_size))); + } + else + { + RETURN_IF_ROCSPARSE_ERROR( + (rocsparse::primitives::segmented_radix_sort_keys_buffer_size( + handle, nnz, m, startbit, endbit, buffer_size))); + } + return rocsparse_status_success; + } + + static rocsparse_status csrsort_rocprim_buffer_size(rocsparse_handle handle, + rocsparse_indextype ptr_type, + rocsparse_indextype ind_type, + int64_t m, + int64_t nnz, + uint32_t startbit, + uint32_t endbit, + bool with_perm, + size_t* buffer_size) + { + if(ptr_type == rocsparse_indextype_i32 && ind_type == rocsparse_indextype_i32) + { + RETURN_IF_ROCSPARSE_ERROR((csrsort_rocprim_buffer_size_template( + handle, m, nnz, startbit, endbit, with_perm, buffer_size))); + return rocsparse_status_success; + } + if(ptr_type == rocsparse_indextype_i64 && ind_type == rocsparse_indextype_i32) + { + RETURN_IF_ROCSPARSE_ERROR((csrsort_rocprim_buffer_size_template( + handle, m, nnz, startbit, endbit, with_perm, buffer_size))); + return rocsparse_status_success; + } + if(ptr_type == rocsparse_indextype_i64 && ind_type == rocsparse_indextype_i64) + { + RETURN_IF_ROCSPARSE_ERROR((csrsort_rocprim_buffer_size_template( + handle, m, nnz, startbit, endbit, with_perm, buffer_size))); + return rocsparse_status_success; + } + RETURN_WITH_MESSAGE_IF_ROCSPARSE_ERROR( + rocsparse_status_invalid_value, + "the combination of offsets and indices types is not supported"); + } + + // rocPRIM does not sort in place, so the sorted indices and permutation end up in either + // the input arrays or the temporary arrays, and are copied back into the input arrays. + template + static rocsparse_status csrsort_rocprim_sort_template(rocsparse_handle handle, + int64_t m, + int64_t nnz, + const void* offsets, + void* ind, + void* tmp_ind, + void* perm, + void* tmp_perm, + uint32_t startbit, + uint32_t endbit, + size_t buffer_size, + void* buffer) + { + const I* offsets_ = reinterpret_cast(offsets); + J* ind_ = reinterpret_cast(ind); + I* perm_ = reinterpret_cast(perm); + + rocsparse::primitives::double_buffer keys(ind_, reinterpret_cast(tmp_ind)); + + if(perm_ != nullptr) + { + rocsparse::primitives::double_buffer vals(perm_, reinterpret_cast(tmp_perm)); + + RETURN_IF_ROCSPARSE_ERROR( + rocsparse::primitives::segmented_radix_sort_pairs(handle, + keys, + vals, + nnz, + m, + offsets_, + offsets_ + 1, + startbit, + endbit, + buffer_size, + buffer)); + + if(vals.current() != perm_) + { + RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync(perm_, + vals.current(), + sizeof(I) * nnz, + hipMemcpyDeviceToDevice, + handle->stream)); + } + } + else + { + RETURN_IF_ROCSPARSE_ERROR(rocsparse::primitives::segmented_radix_sort_keys(handle, + keys, + nnz, + m, + offsets_, + offsets_ + 1, + startbit, + endbit, + buffer_size, + buffer)); + } + + if(keys.current() != ind_) + { + RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync( + ind_, keys.current(), sizeof(J) * nnz, hipMemcpyDeviceToDevice, handle->stream)); + } + return rocsparse_status_success; + } + + static rocsparse_status csrsort_rocprim_sort(rocsparse_handle handle, + rocsparse_indextype ptr_type, + rocsparse_indextype ind_type, + int64_t m, + int64_t nnz, + const void* offsets, + void* ind, + void* tmp_ind, + void* perm, + void* tmp_perm, + uint32_t startbit, + uint32_t endbit, + size_t buffer_size, + void* buffer) + { + if(ptr_type == rocsparse_indextype_i32 && ind_type == rocsparse_indextype_i32) + { + RETURN_IF_ROCSPARSE_ERROR((csrsort_rocprim_sort_template(handle, + m, + nnz, + offsets, + ind, + tmp_ind, + perm, + tmp_perm, + startbit, + endbit, + buffer_size, + buffer))); + return rocsparse_status_success; + } + if(ptr_type == rocsparse_indextype_i64 && ind_type == rocsparse_indextype_i32) + { + RETURN_IF_ROCSPARSE_ERROR((csrsort_rocprim_sort_template(handle, + m, + nnz, + offsets, + ind, + tmp_ind, + perm, + tmp_perm, + startbit, + endbit, + buffer_size, + buffer))); + return rocsparse_status_success; + } + if(ptr_type == rocsparse_indextype_i64 && ind_type == rocsparse_indextype_i64) + { + RETURN_IF_ROCSPARSE_ERROR((csrsort_rocprim_sort_template(handle, + m, + nnz, + offsets, + ind, + tmp_ind, + perm, + tmp_perm, + startbit, + endbit, + buffer_size, + buffer))); + return rocsparse_status_success; + } + RETURN_WITH_MESSAGE_IF_ROCSPARSE_ERROR( + rocsparse_status_invalid_value, + "the combination of offsets and indices types is not supported"); + } + + // Shifts one based offsets to zero based offsets, which the segmented sort expects. + static rocsparse_status csrsort_shift_offsets(rocsparse_handle handle, + rocsparse_indextype ptr_type, + int64_t size, + const void* in, + void* out) + { +#define CSRSORT_DIM 512 + dim3 csrsort_blocks((size - 1) / CSRSORT_DIM + 1); + dim3 csrsort_threads(CSRSORT_DIM); + + if(ptr_type == rocsparse_indextype_i32) + { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR((rocsparse::csrsort_shift_kernel), + csrsort_blocks, + csrsort_threads, + 0, + handle->stream, + size, + reinterpret_cast(in), + reinterpret_cast(out)); + return rocsparse_status_success; + } + if(ptr_type == rocsparse_indextype_i64) + { + RETURN_IF_HIPLAUNCHKERNELGGL_ERROR((rocsparse::csrsort_shift_kernel), + csrsort_blocks, + csrsort_threads, + 0, + handle->stream, + size, + reinterpret_cast(in), + reinterpret_cast(out)); + return rocsparse_status_success; + } +#undef CSRSORT_DIM + RETURN_WITH_MESSAGE_IF_ROCSPARSE_ERROR(rocsparse_status_invalid_value, + "the offsets type is not supported"); + } + + static rocsparse_status csrsort_buffer_size_compute(rocsparse_handle handle, + int64_t m, + int64_t n, + int64_t nnz, + rocsparse_indextype ptr_type, + rocsparse_indextype ind_type, + size_t* buffer_size_in_bytes) { ROCSPARSE_ROUTINE_TRACE; + if(!rocsparse::csrsort_is_supported(ptr_type, ind_type)) + { + RETURN_WITH_MESSAGE_IF_ROCSPARSE_ERROR( + rocsparse_status_invalid_value, + "the combination of offsets and indices types is not supported"); + } + if(m == 0 || n == 0 || nnz == 0) { - *buffer_size = 0; + *buffer_size_in_bytes = 0; return rocsparse_status_success; } @@ -61,161 +309,117 @@ namespace rocsparse const uint32_t endbit = rocsparse::csrsort_endbit(n); // We do not know if sort_pairs or sort_keys will be called, so use the largest buffer between the two - size_t size1; - size_t size2; - RETURN_IF_ROCSPARSE_ERROR( - (rocsparse::primitives::segmented_radix_sort_pairs_buffer_size( - handle, nnz, m, startbit, endbit, &size1))); - RETURN_IF_ROCSPARSE_ERROR( - (rocsparse::primitives::segmented_radix_sort_keys_buffer_size( - handle, nnz, m, startbit, endbit, &size2))); + size_t size1 = std::numeric_limits::max(); + size_t size2 = std::numeric_limits::max(); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::csrsort_rocprim_buffer_size( + handle, ptr_type, ind_type, m, nnz, startbit, endbit, true, &size1)); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::csrsort_rocprim_buffer_size( + handle, ptr_type, ind_type, m, nnz, startbit, endbit, false, &size2)); - *buffer_size = rocsparse::align_size(rocsparse::max(size1, size2)); + const size_t ptr_size = rocsparse::indextype_sizeof(ptr_type); + const size_t ind_size = rocsparse::indextype_sizeof(ind_type); // rocPRIM does not support in-place sorting, so we need additional buffer // for all temporary arrays + *buffer_size_in_bytes = rocsparse::align_size(rocsparse::max(size1, size2)); // columns buffer - *buffer_size += rocsparse::align_size(nnz); + *buffer_size_in_bytes += rocsparse::align_size(ind_size * nnz); // perm buffer - *buffer_size += rocsparse::align_size(nnz); + *buffer_size_in_bytes += rocsparse::align_size(ptr_size * nnz); // segm buffer - *buffer_size += rocsparse::align_size(m + 1); + *buffer_size_in_bytes += rocsparse::align_size(ptr_size * (m + 1)); return rocsparse_status_success; } // Sorts the column indices within each row. If perm is not null, the sort applies its // reordering to perm. - template - static rocsparse_status csrsort_template(rocsparse_handle handle, - int64_t m, - int64_t n, - int64_t nnz, - rocsparse_index_base idx_base, - const I* csr_row_ptr, - J* csr_col_ind, - I* perm, - void* temp_buffer) + static rocsparse_status csrsort_compute(rocsparse_handle handle, + int64_t m, + int64_t n, + int64_t nnz, + rocsparse_index_base idx_base, + rocsparse_indextype ptr_type, + rocsparse_indextype ind_type, + const void* csr_row_ptr, + void* csr_col_ind, + void* perm, + size_t buffer_size_in_bytes, + void* buffer) { ROCSPARSE_ROUTINE_TRACE; + size_t required_buffer_size = std::numeric_limits::max(); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::csrsort_buffer_size_compute( + handle, m, n, nnz, ptr_type, ind_type, &required_buffer_size)); + if(buffer_size_in_bytes < required_buffer_size) + { + RETURN_WITH_MESSAGE_IF_ROCSPARSE_ERROR( + rocsparse_status_invalid_size, + "the buffer is smaller than the size required by the index sort"); + } + // Quick return if possible if(m == 0 || n == 0 || nnz == 0) { return rocsparse_status_success; } - // Stream - hipStream_t stream = handle->stream; - const uint32_t startbit = 0; const uint32_t endbit = rocsparse::csrsort_endbit(n); - size_t size; - if(perm != nullptr) - { - // Sort pairs, if permutation vector is present - RETURN_IF_ROCSPARSE_ERROR( - (rocsparse::primitives::segmented_radix_sort_pairs_buffer_size( - handle, nnz, m, startbit, endbit, &size))); - } - else - { - // Sort keys, if no permutation vector is present - RETURN_IF_ROCSPARSE_ERROR( - (rocsparse::primitives::segmented_radix_sort_keys_buffer_size( - handle, nnz, m, startbit, endbit, &size))); - } + size_t rocprim_size = std::numeric_limits::max(); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::csrsort_rocprim_buffer_size( + handle, ptr_type, ind_type, m, nnz, startbit, endbit, perm != nullptr, &rocprim_size)); + + const size_t ptr_size = rocsparse::indextype_sizeof(ptr_type); + const size_t ind_size = rocsparse::indextype_sizeof(ind_type); // Temporary buffer entry points - char* ptr = reinterpret_cast(temp_buffer); + char* ptr = reinterpret_cast(buffer); // columns buffer - J* tmp_cols = reinterpret_cast(ptr); - ptr += rocsparse::align_size(nnz); + void* tmp_cols = ptr; + ptr += rocsparse::align_size(ind_size * nnz); // perm buffer - I* tmp_perm = reinterpret_cast(ptr); - ptr += rocsparse::align_size(nnz); + void* tmp_perm = ptr; + ptr += rocsparse::align_size(ptr_size * nnz); // segm buffer - I* tmp_segm = reinterpret_cast(ptr); - ptr += rocsparse::align_size(m + 1); + void* tmp_segm = ptr; + ptr += rocsparse::align_size(ptr_size * (m + 1)); + + // rocprim buffer + void* tmp_rocprim = ptr; // Index base one requires shift of offset positions + const void* offsets = csr_row_ptr; if(idx_base == rocsparse_index_base_one) { -#define CSRSORT_DIM 512 - dim3 csrsort_blocks((m + 1 - 1) / CSRSORT_DIM + 1); - dim3 csrsort_threads(CSRSORT_DIM); - - RETURN_IF_HIPLAUNCHKERNELGGL_ERROR((rocsparse::csrsort_shift_kernel), - csrsort_blocks, - csrsort_threads, - 0, - stream, - m + 1, - csr_row_ptr, - tmp_segm); -#undef CSRSORT_DIM + RETURN_IF_ROCSPARSE_ERROR( + rocsparse::csrsort_shift_offsets(handle, ptr_type, m + 1, csr_row_ptr, tmp_segm)); + offsets = tmp_segm; } - // rocprim buffer - void* tmp_rocprim = reinterpret_cast(ptr); - - // Switch between offsets - const I* offsets = (idx_base == rocsparse_index_base_one) ? tmp_segm : csr_row_ptr; - // Sort by columns and obtain permutation vector - - if(perm != nullptr) - { - // Sort by pairs, if permutation vector is present - rocsparse::primitives::double_buffer keys(csr_col_ind, tmp_cols); - rocsparse::primitives::double_buffer vals(perm, tmp_perm); - - RETURN_IF_ROCSPARSE_ERROR( - rocsparse::primitives::segmented_radix_sort_pairs(handle, - keys, - vals, - nnz, + RETURN_IF_ROCSPARSE_ERROR(rocsparse::csrsort_rocprim_sort(handle, + ptr_type, + ind_type, m, + nnz, offsets, - offsets + 1, + csr_col_ind, + tmp_cols, + perm, + tmp_perm, startbit, endbit, - size, + rocprim_size, tmp_rocprim)); - - if(keys.current() != csr_col_ind) - { - RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync( - csr_col_ind, keys.current(), sizeof(J) * nnz, hipMemcpyDeviceToDevice, stream)); - } - if(vals.current() != perm) - { - RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync( - perm, vals.current(), sizeof(I) * nnz, hipMemcpyDeviceToDevice, stream)); - } - } - else - { - // Sort by keys, if no permutation vector is present - rocsparse::primitives::double_buffer keys(csr_col_ind, tmp_cols); - - RETURN_IF_ROCSPARSE_ERROR(rocsparse::primitives::segmented_radix_sort_keys( - handle, keys, nnz, m, offsets, offsets + 1, startbit, endbit, size, tmp_rocprim)); - - if(keys.current() != csr_col_ind) - { - RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync( - csr_col_ind, keys.current(), sizeof(J) * nnz, hipMemcpyDeviceToDevice, stream)); - } - } return rocsparse_status_success; } - } extern "C" rocsparse_status rocsparse_csrsort_buffer_size(rocsparse_handle handle, @@ -248,8 +452,13 @@ try ROCSPARSE_CHECKARG_POINTER(6, buffer_size); RETURN_IF_ROCSPARSE_ERROR( - (rocsparse::csrsort_buffer_size_template( - handle, m, n, nnz, csr_row_ptr, csr_col_ind, buffer_size))); + rocsparse::csrsort_buffer_size_compute(handle, + m, + n, + nnz, + rocsparse::get_indextype(), + rocsparse::get_indextype(), + buffer_size)); return rocsparse_status_success; // LCOV_EXCL_START @@ -294,8 +503,30 @@ try ROCSPARSE_CHECKARG_ARRAY(6, nnz, csr_col_ind); ROCSPARSE_CHECKARG_ARRAY(8, nnz, temp_buffer); - RETURN_IF_ROCSPARSE_ERROR((rocsparse::csrsort_template( - handle, m, n, nnz, descr->base, csr_row_ptr, csr_col_ind, perm, temp_buffer))); + // The legacy API does not take the size of temp_buffer, which is assumed to be the size + // returned by rocsparse_csrsort_buffer_size. + size_t buffer_size = std::numeric_limits::max(); + RETURN_IF_ROCSPARSE_ERROR( + rocsparse::csrsort_buffer_size_compute(handle, + m, + n, + nnz, + rocsparse::get_indextype(), + rocsparse::get_indextype(), + &buffer_size)); + + RETURN_IF_ROCSPARSE_ERROR(rocsparse::csrsort_compute(handle, + m, + n, + nnz, + descr->base, + rocsparse::get_indextype(), + rocsparse::get_indextype(), + csr_row_ptr, + csr_col_ind, + perm, + buffer_size, + temp_buffer)); return rocsparse_status_success; // LCOV_EXCL_START @@ -308,116 +539,6 @@ catch(...) namespace rocsparse { - static rocsparse_status gcsrsort_buffer_size(rocsparse_handle handle, - int64_t m, - int64_t n, - int64_t nnz, - rocsparse_indextype ptr_type, - rocsparse_indextype ind_type, - const void* ptr, - const void* ind, - size_t* buffer_size) - { - if(ptr_type == rocsparse_indextype_i32 && ind_type == rocsparse_indextype_i32) - { - RETURN_IF_ROCSPARSE_ERROR( - rocsparse::csrsort_buffer_size_template(handle, - m, - n, - nnz, - reinterpret_cast(ptr), - reinterpret_cast(ind), - buffer_size)); - return rocsparse_status_success; - } - if(ptr_type == rocsparse_indextype_i64 && ind_type == rocsparse_indextype_i32) - { - RETURN_IF_ROCSPARSE_ERROR( - rocsparse::csrsort_buffer_size_template(handle, - m, - n, - nnz, - reinterpret_cast(ptr), - reinterpret_cast(ind), - buffer_size)); - return rocsparse_status_success; - } - if(ptr_type == rocsparse_indextype_i64 && ind_type == rocsparse_indextype_i64) - { - RETURN_IF_ROCSPARSE_ERROR( - rocsparse::csrsort_buffer_size_template(handle, - m, - n, - nnz, - reinterpret_cast(ptr), - reinterpret_cast(ind), - buffer_size)); - return rocsparse_status_success; - } - RETURN_WITH_MESSAGE_IF_ROCSPARSE_ERROR( - rocsparse_status_invalid_value, - "the combination of offsets and indices types is not supported"); - } - - static rocsparse_status gcsrsort(rocsparse_handle handle, - int64_t m, - int64_t n, - int64_t nnz, - rocsparse_index_base idx_base, - rocsparse_indextype ptr_type, - rocsparse_indextype ind_type, - const void* ptr, - void* ind, - void* perm, - void* buffer) - { - if(ptr_type == rocsparse_indextype_i32 && ind_type == rocsparse_indextype_i32) - { - RETURN_IF_ROCSPARSE_ERROR( - rocsparse::csrsort_template(handle, - m, - n, - nnz, - idx_base, - reinterpret_cast(ptr), - reinterpret_cast(ind), - reinterpret_cast(perm), - buffer)); - return rocsparse_status_success; - } - if(ptr_type == rocsparse_indextype_i64 && ind_type == rocsparse_indextype_i32) - { - RETURN_IF_ROCSPARSE_ERROR( - rocsparse::csrsort_template(handle, - m, - n, - nnz, - idx_base, - reinterpret_cast(ptr), - reinterpret_cast(ind), - reinterpret_cast(perm), - buffer)); - return rocsparse_status_success; - } - if(ptr_type == rocsparse_indextype_i64 && ind_type == rocsparse_indextype_i64) - { - RETURN_IF_ROCSPARSE_ERROR( - rocsparse::csrsort_template(handle, - m, - n, - nnz, - idx_base, - reinterpret_cast(ptr), - reinterpret_cast(ind), - reinterpret_cast(perm), - buffer)); - return rocsparse_status_success; - } - RETURN_WITH_MESSAGE_IF_ROCSPARSE_ERROR( - rocsparse_status_invalid_value, - "the combination of offsets and indices types is not supported"); - } - // The buffer starts with the permutation array, which tracks where each entry moves // while the indices are sorted, followed by scratch space shared by the index sort // and the value permutation. @@ -442,12 +563,10 @@ rocsparse_status rocsparse::csxsort_buffer_size(rocsparse_handle hand const int64_t nnz = target->nnz; const rocsparse_indextype ptr_type = by_row ? target->row_type : target->col_type; const rocsparse_indextype ind_type = by_row ? target->col_type : target->row_type; - const void* ptr = by_row ? target->const_row_data : target->const_col_data; - const void* ind = by_row ? target->const_col_data : target->const_row_data; - size_t sort_buffer_size = 0; - RETURN_IF_ROCSPARSE_ERROR(rocsparse::gcsrsort_buffer_size( - handle, m, n, nnz, ptr_type, ind_type, ptr, ind, &sort_buffer_size)); + size_t sort_buffer_size = std::numeric_limits::max(); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::csrsort_buffer_size_compute( + handle, m, n, nnz, ptr_type, ind_type, &sort_buffer_size)); // Values sorted in place are gathered into scratch space first, since the gather cannot // write over its own input. @@ -509,9 +628,10 @@ rocsparse_status rocsparse::csxsort(rocsparse_handle handle, const size_t ind_stride_target = target->columns_values_batch_stride * ind_size; const size_t val_stride_target = target->columns_values_batch_stride * val_size; - void* perm = buffer; - void* sort_buffer - = reinterpret_cast(buffer) + rocsparse::csxsort_perm_size(nnz, ptr_type); + const size_t perm_size = rocsparse::csxsort_perm_size(nnz, ptr_type); + const size_t sort_buffer_size = buffer_size_in_bytes - perm_size; + void* perm = buffer; + void* sort_buffer = reinterpret_cast(buffer) + perm_size; // Only uniform batches are supported, so every batch has the sparsity pattern of the first // one. Its indices are sorted once, and the resulting permutation is applied to the values @@ -533,17 +653,18 @@ rocsparse_status rocsparse::csxsort(rocsparse_handle handle, // The index sort applies its reordering to perm, so it must start as the identity. RETURN_IF_ROCSPARSE_ERROR(rocsparse::gcreate_identity_permutation(handle, nnz, ptr_type, perm)); - RETURN_IF_ROCSPARSE_ERROR(rocsparse::gcsrsort(handle, - m, - n, - nnz, - idx_base, - ptr_type, - ind_type, - ptr_target, - ind_target, - perm, - sort_buffer)); + RETURN_IF_ROCSPARSE_ERROR(rocsparse::csrsort_compute(handle, + m, + n, + nnz, + idx_base, + ptr_type, + ind_type, + ptr_target, + ind_target, + perm, + sort_buffer_size, + sort_buffer)); // The batches run one after the other on the handle stream, so they share the buffer. for(int64_t batch = 0; batch < target->batch_count; ++batch) diff --git a/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_keys.cpp b/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_keys.cpp index a92a465ef106..70fa1d4109c2 100644 --- a/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_keys.cpp +++ b/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_keys.cpp @@ -233,7 +233,6 @@ rocsparse_status rocsparse::primitives::sort_csr_column_indices(rocsparse_handle INSTANTIATE(int32_t, int32_t); INSTANTIATE(int64_t, int64_t); INSTANTIATE(int32_t, int64_t); -INSTANTIATE(int32_t, const int64_t); #undef INSTANTIATE #define INSTANTIATE(ITYPE, JTYPE) \ From bd8173723701c66a257c73ccac78011c9dcf8652 Mon Sep 17 00:00:00 2001 From: jsandham Date: Wed, 30 Sep 2026 20:02:00 +0000 Subject: [PATCH 4/8] Do not require a rocsparse_spsort direction for CSR and CSC matrices The CSR and CSC formats always sort the indices within each row or column, so the direction is ignored for them. Co-authored-by: Cursor --- .../generic/example_rocsparse_spsort.cpp | 9 --------- .../library/src/conversion/rocsparse_spsort.cpp | 13 ------------- 2 files changed, 22 deletions(-) diff --git a/projects/rocsparse/clients/samples/documentation_examples/generic/example_rocsparse_spsort.cpp b/projects/rocsparse/clients/samples/documentation_examples/generic/example_rocsparse_spsort.cpp index 89e662e7c51f..0001a144a042 100644 --- a/projects/rocsparse/clients/samples/documentation_examples/generic/example_rocsparse_spsort.cpp +++ b/projects/rocsparse/clients/samples/documentation_examples/generic/example_rocsparse_spsort.cpp @@ -115,15 +115,6 @@ int main() sizeof(spsort_alg), p_error)); - // Sort the column indices within each row - const rocsparse_direction spsort_direction = rocsparse_direction_row; - ROCSPARSE_CHECK(rocsparse_spsort_set_input(handle, - spsort_descr, - rocsparse_spsort_input_direction, - &spsort_direction, - sizeof(spsort_direction), - p_error)); - // Passing matA as both the input and the output sorts it in place // Call spsort to get buffer size size_t buffer_size; diff --git a/projects/rocsparse/library/src/conversion/rocsparse_spsort.cpp b/projects/rocsparse/library/src/conversion/rocsparse_spsort.cpp index 8bafd0867e89..17b0fbad8c59 100644 --- a/projects/rocsparse/library/src/conversion/rocsparse_spsort.cpp +++ b/projects/rocsparse/library/src/conversion/rocsparse_spsort.cpp @@ -171,19 +171,6 @@ rocsparse_status rocsparse::spsort_check_arguments(rocsparse_spsort_descr d && rocsparse::enum_utils::is_invalid(descr->get_direction())), rocsparse_status_invalid_value); - // A CSR matrix can only have the column indices within each row sorted, and a CSC matrix - // can only have the row indices within each column sorted. - ROCSPARSE_CHECKARG(1, - descr, - (source->format == rocsparse_format_csr - && descr->get_direction() != rocsparse_direction_row), - rocsparse_status_invalid_value); - ROCSPARSE_CHECKARG(1, - descr, - (source->format == rocsparse_format_csc - && descr->get_direction() != rocsparse_direction_column), - rocsparse_status_invalid_value); - return rocsparse_status_success; } From 2a87e0eaa71f3c757afb1f0e61ce769ea2078c0b Mon Sep 17 00:00:00 2001 From: jsandham Date: Wed, 30 Sep 2026 20:59:11 +0000 Subject: [PATCH 5/8] Document that a single CSR or CSC source matrix is sorted into every batch Co-authored-by: Cursor --- .../rocsparse/library/src/conversion/rocsparse_csrsort.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp index 69ebb084e12b..92ae148772d2 100644 --- a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp +++ b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp @@ -635,7 +635,8 @@ rocsparse_status rocsparse::csxsort(rocsparse_handle handle, // Only uniform batches are supported, so every batch has the sparsity pattern of the first // one. Its indices are sorted once, and the resulting permutation is applied to the values - // of every batch. + // of every batch. A single source matrix has zero batch strides, so its values are sorted + // into every batch of target. // The index sort works in place, so the offsets and indices of source are first copied into // target. A matrix without rows may have a null offsets array. From 639c02bf008fef4d7120032af577121632dce1c3 Mon Sep 17 00:00:00 2001 From: jsandham Date: Wed, 30 Sep 2026 22:16:52 +0000 Subject: [PATCH 6/8] Address rocsparse_csrsort review comments Give the size of csrsort_shift_kernel the type of its offsets, and describe the csxsort direction as selecting the CSR or CSC format. Co-authored-by: Cursor --- projects/rocsparse/library/src/conversion/csrsort_device.h | 2 +- .../rocsparse/library/src/conversion/rocsparse_csrsort.cpp | 2 +- .../rocsparse/library/src/conversion/rocsparse_csrsort.hpp | 5 ++--- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/projects/rocsparse/library/src/conversion/csrsort_device.h b/projects/rocsparse/library/src/conversion/csrsort_device.h index b8424f7caddb..837bd6ca5650 100644 --- a/projects/rocsparse/library/src/conversion/csrsort_device.h +++ b/projects/rocsparse/library/src/conversion/csrsort_device.h @@ -31,7 +31,7 @@ namespace rocsparse // Shift CSR offsets template ROCSPARSE_KERNEL(BLOCKSIZE) - void csrsort_shift_kernel(int64_t size, const I* in, I* out) + void csrsort_shift_kernel(I size, const I* in, I* out) { const int64_t gid = int64_t(hipBlockIdx_x) * BLOCKSIZE + hipThreadIdx_x; diff --git a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp index 92ae148772d2..09fdc263fb44 100644 --- a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp +++ b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp @@ -260,7 +260,7 @@ namespace rocsparse csrsort_threads, 0, handle->stream, - size, + static_cast(size), reinterpret_cast(in), reinterpret_cast(out)); return rocsparse_status_success; diff --git a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.hpp b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.hpp index e6adc1bef769..7d35b0055df6 100644 --- a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.hpp +++ b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.hpp @@ -33,9 +33,8 @@ namespace rocsparse rocsparse_csrsort_alg_default = 0 } rocsparse_csrsort_alg; - // Shared by the CSR and CSC sorts: dir == rocsparse_direction_row sorts the column indices - // within each row of a CSR matrix, and dir == rocsparse_direction_column sorts the row - // indices within each column of a CSC matrix. + // Shared by the CSR and CSC sorts: dir selects the format, rocsparse_direction_row for CSR + // and rocsparse_direction_column for CSC. rocsparse_status csxsort_buffer_size(rocsparse_handle handle, rocsparse_direction dir, rocsparse_const_spmat_descr source, From 4a2237234de21504ef562cc32bd59ed8716884f4 Mon Sep 17 00:00:00 2001 From: jsandham Date: Fri, 2 Oct 2026 15:16:13 +0000 Subject: [PATCH 7/8] Rename csxsort_perm_size to csxsort_perm_size_in_bytes Co-authored-by: Cursor --- .../rocsparse/library/src/conversion/rocsparse_csrsort.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp index 09fdc263fb44..fc92bd4dc1a4 100644 --- a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp +++ b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp @@ -542,7 +542,7 @@ namespace rocsparse // The buffer starts with the permutation array, which tracks where each entry moves // while the indices are sorted, followed by scratch space shared by the index sort // and the value permutation. - static size_t csxsort_perm_size(int64_t nnz, rocsparse_indextype perm_indextype) + static size_t csxsort_perm_size_in_bytes(int64_t nnz, rocsparse_indextype perm_indextype) { return rocsparse::align_size(rocsparse::indextype_sizeof(perm_indextype) * nnz); } @@ -575,7 +575,7 @@ rocsparse_status rocsparse::csxsort_buffer_size(rocsparse_handle hand ? rocsparse::align_size(rocsparse::datatype_sizeof(target->data_type) * nnz) : 0; - *buffer_size_in_bytes = rocsparse::csxsort_perm_size(nnz, ptr_type) + *buffer_size_in_bytes = rocsparse::csxsort_perm_size_in_bytes(nnz, ptr_type) + rocsparse::max(sort_buffer_size, gather_buffer_size); return rocsparse_status_success; @@ -628,7 +628,7 @@ rocsparse_status rocsparse::csxsort(rocsparse_handle handle, const size_t ind_stride_target = target->columns_values_batch_stride * ind_size; const size_t val_stride_target = target->columns_values_batch_stride * val_size; - const size_t perm_size = rocsparse::csxsort_perm_size(nnz, ptr_type); + const size_t perm_size = rocsparse::csxsort_perm_size_in_bytes(nnz, ptr_type); const size_t sort_buffer_size = buffer_size_in_bytes - perm_size; void* perm = buffer; void* sort_buffer = reinterpret_cast(buffer) + perm_size; From eef150505c68585baba2aa9faa0db2230f7c6a60 Mon Sep 17 00:00:00 2001 From: jsandham Date: Tue, 6 Oct 2026 15:40:10 +0000 Subject: [PATCH 8/8] Address more rocsparse_csrsort review comments Support 32-bit offsets with 64-bit indices, which needs the matching segmented radix sort instantiations. Set the single offset of a target without rows to the index base when the source has no offsets array, instead of leaving it uninitialized, explain why every batch of the target receives the sorted indices, remove the unused rocsparse_csrsort_alg, and use static_cast in csrsort_shift_kernel. Co-authored-by: Cursor --- .../library/src/conversion/csrsort_device.h | 2 +- .../src/conversion/rocsparse_csrsort.cpp | 61 +++++++++++++++---- .../src/conversion/rocsparse_csrsort.hpp | 7 --- .../src/conversion/rocsparse_spsort.cpp | 8 +-- .../rocsparse_spsort_buffer_size.cpp | 4 +- .../rocsparse_segmented_radix_sort_keys.cpp | 2 + .../rocsparse_segmented_radix_sort_pairs.cpp | 2 + 7 files changed, 57 insertions(+), 29 deletions(-) diff --git a/projects/rocsparse/library/src/conversion/csrsort_device.h b/projects/rocsparse/library/src/conversion/csrsort_device.h index 837bd6ca5650..64929d3a0484 100644 --- a/projects/rocsparse/library/src/conversion/csrsort_device.h +++ b/projects/rocsparse/library/src/conversion/csrsort_device.h @@ -33,7 +33,7 @@ namespace rocsparse ROCSPARSE_KERNEL(BLOCKSIZE) void csrsort_shift_kernel(I size, const I* in, I* out) { - const int64_t gid = int64_t(hipBlockIdx_x) * BLOCKSIZE + hipThreadIdx_x; + const int64_t gid = static_cast(hipBlockIdx_x) * BLOCKSIZE + hipThreadIdx_x; if(gid >= size) { diff --git a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp index fc92bd4dc1a4..55a2432f2cd4 100644 --- a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp +++ b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.cpp @@ -30,6 +30,7 @@ #include "rocsparse_csrsort.hpp" #include "rocsparse_gcreate_identity_permutation.hpp" #include "rocsparse_primitives.hpp" +#include "rocsparse_valset.hpp" namespace rocsparse { @@ -42,9 +43,8 @@ namespace rocsparse static bool csrsort_is_supported(rocsparse_indextype ptr_type, rocsparse_indextype ind_type) { - return (ptr_type == rocsparse_indextype_i32 && ind_type == rocsparse_indextype_i32) - || (ptr_type == rocsparse_indextype_i64 && ind_type == rocsparse_indextype_i32) - || (ptr_type == rocsparse_indextype_i64 && ind_type == rocsparse_indextype_i64); + return (ptr_type == rocsparse_indextype_i32 || ptr_type == rocsparse_indextype_i64) + && (ind_type == rocsparse_indextype_i32 || ind_type == rocsparse_indextype_i64); } template @@ -87,6 +87,12 @@ namespace rocsparse handle, m, nnz, startbit, endbit, with_perm, buffer_size))); return rocsparse_status_success; } + if(ptr_type == rocsparse_indextype_i32 && ind_type == rocsparse_indextype_i64) + { + RETURN_IF_ROCSPARSE_ERROR((csrsort_rocprim_buffer_size_template( + handle, m, nnz, startbit, endbit, with_perm, buffer_size))); + return rocsparse_status_success; + } if(ptr_type == rocsparse_indextype_i64 && ind_type == rocsparse_indextype_i32) { RETURN_IF_ROCSPARSE_ERROR((csrsort_rocprim_buffer_size_template( @@ -205,6 +211,22 @@ namespace rocsparse buffer))); return rocsparse_status_success; } + if(ptr_type == rocsparse_indextype_i32 && ind_type == rocsparse_indextype_i64) + { + RETURN_IF_ROCSPARSE_ERROR((csrsort_rocprim_sort_template(handle, + m, + nnz, + offsets, + ind, + tmp_ind, + perm, + tmp_perm, + startbit, + endbit, + buffer_size, + buffer))); + return rocsparse_status_success; + } if(ptr_type == rocsparse_indextype_i64 && ind_type == rocsparse_indextype_i32) { RETURN_IF_ROCSPARSE_ERROR((csrsort_rocprim_sort_template(handle, @@ -639,11 +661,23 @@ rocsparse_status rocsparse::csxsort(rocsparse_handle handle, // into every batch of target. // The index sort works in place, so the offsets and indices of source are first copied into - // target. A matrix without rows may have a null offsets array. - if(m > 0 && ptr_target != ptr_source) + // target. A matrix without rows may have a null offsets array, in which case the single + // offset of target is set to the index base. + if(ptr_target != nullptr && ptr_target != ptr_source) { - RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync( - ptr_target, ptr_source, ptr_size * (m + 1), hipMemcpyDeviceToDevice, handle->stream)); + if(ptr_source != nullptr) + { + RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync(ptr_target, + ptr_source, + ptr_size * (m + 1), + hipMemcpyDeviceToDevice, + handle->stream)); + } + else + { + RETURN_IF_ROCSPARSE_ERROR( + rocsparse::valset(handle, 1, static_cast(idx_base), ptr_type, ptr_target)); + } } if(ind_target != ind_source) { @@ -667,15 +701,18 @@ rocsparse_status rocsparse::csxsort(rocsparse_handle handle, sort_buffer_size, sort_buffer)); - // The batches run one after the other on the handle stream, so they share the buffer. + // A strided batched CSR or CSC matrix stores its indices with the same batch stride as its + // values, so every batch of target has its own indices, which receive the sorted indices of + // the first batch. Its offsets are copied too, unless their batch stride is zero and all + // batches share them. The batches run one after the other on the handle stream, so they + // share the buffer. for(int64_t batch = 0; batch < target->batch_count; ++batch) { if(batch > 0) { - // The offsets of the target are shared by all its batches when their stride is zero. - char* batch_ptr_target = ptr_target + batch * ptr_stride_target; - if(m > 0 && batch_ptr_target != ptr_target) + if(ptr_target != nullptr && ptr_stride_target != 0) { + char* batch_ptr_target = ptr_target + batch * ptr_stride_target; RETURN_IF_HIP_ERROR(rocsparse_hipMemcpyAsync(batch_ptr_target, ptr_target, ptr_size * (m + 1), @@ -719,7 +756,6 @@ rocsparse_status rocsparse::csxsort(rocsparse_handle handle, } rocsparse_status rocsparse::csrsort_buffer_size(rocsparse_handle handle, - rocsparse_csrsort_alg alg, rocsparse_const_spmat_descr source, rocsparse_const_spmat_descr target, size_t* buffer_size_in_bytes) @@ -732,7 +768,6 @@ rocsparse_status rocsparse::csrsort_buffer_size(rocsparse_handle hand } rocsparse_status rocsparse::csrsort(rocsparse_handle handle, - rocsparse_csrsort_alg alg, rocsparse_const_spmat_descr source, rocsparse_spmat_descr target, size_t buffer_size_in_bytes, diff --git a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.hpp b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.hpp index 7d35b0055df6..242c061d16d4 100644 --- a/projects/rocsparse/library/src/conversion/rocsparse_csrsort.hpp +++ b/projects/rocsparse/library/src/conversion/rocsparse_csrsort.hpp @@ -28,11 +28,6 @@ namespace rocsparse { - typedef enum rocsparse_csrsort_alg_ - { - rocsparse_csrsort_alg_default = 0 - } rocsparse_csrsort_alg; - // Shared by the CSR and CSC sorts: dir selects the format, rocsparse_direction_row for CSR // and rocsparse_direction_column for CSC. rocsparse_status csxsort_buffer_size(rocsparse_handle handle, @@ -49,7 +44,6 @@ namespace rocsparse void* buffer); rocsparse_status csrsort_buffer_size(rocsparse_handle handle, - rocsparse_csrsort_alg alg, rocsparse_const_spmat_descr source, rocsparse_const_spmat_descr target, size_t* buffer_size_in_bytes); @@ -61,7 +55,6 @@ namespace rocsparse // place, or not overlap it at all. The batches must be uniform: they all share the sparsity // pattern of the first batch, and only their values differ. rocsparse_status csrsort(rocsparse_handle handle, - rocsparse_csrsort_alg alg, rocsparse_const_spmat_descr source, rocsparse_spmat_descr target, size_t buffer_size_in_bytes, diff --git a/projects/rocsparse/library/src/conversion/rocsparse_spsort.cpp b/projects/rocsparse/library/src/conversion/rocsparse_spsort.cpp index 17b0fbad8c59..410d60179573 100644 --- a/projects/rocsparse/library/src/conversion/rocsparse_spsort.cpp +++ b/projects/rocsparse/library/src/conversion/rocsparse_spsort.cpp @@ -230,12 +230,8 @@ namespace rocsparse } case rocsparse_format_csr: { - RETURN_IF_ROCSPARSE_ERROR(rocsparse::csrsort(handle, - rocsparse_csrsort_alg_default, - source, - target, - buffer_size_in_bytes, - buffer)); + RETURN_IF_ROCSPARSE_ERROR( + rocsparse::csrsort(handle, source, target, buffer_size_in_bytes, buffer)); return rocsparse_status_success; } diff --git a/projects/rocsparse/library/src/conversion/rocsparse_spsort_buffer_size.cpp b/projects/rocsparse/library/src/conversion/rocsparse_spsort_buffer_size.cpp index 1dc799c13b2a..c83238864038 100644 --- a/projects/rocsparse/library/src/conversion/rocsparse_spsort_buffer_size.cpp +++ b/projects/rocsparse/library/src/conversion/rocsparse_spsort_buffer_size.cpp @@ -76,8 +76,8 @@ namespace rocsparse } case rocsparse_format_csr: { - RETURN_IF_ROCSPARSE_ERROR(rocsparse::csrsort_buffer_size( - handle, rocsparse_csrsort_alg_default, source, target, buffer_size_in_bytes)); + RETURN_IF_ROCSPARSE_ERROR( + rocsparse::csrsort_buffer_size(handle, source, target, buffer_size_in_bytes)); return rocsparse_status_success; } case rocsparse_format_csc: diff --git a/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_keys.cpp b/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_keys.cpp index 70fa1d4109c2..76cfbb2efcfd 100644 --- a/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_keys.cpp +++ b/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_keys.cpp @@ -233,6 +233,8 @@ rocsparse_status rocsparse::primitives::sort_csr_column_indices(rocsparse_handle INSTANTIATE(int32_t, int32_t); INSTANTIATE(int64_t, int64_t); INSTANTIATE(int32_t, int64_t); +INSTANTIATE(int64_t, int32_t); +INSTANTIATE(int64_t, const int32_t); #undef INSTANTIATE #define INSTANTIATE(ITYPE, JTYPE) \ diff --git a/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_pairs.cpp b/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_pairs.cpp index edec179ca8b5..27406b847f3c 100644 --- a/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_pairs.cpp +++ b/projects/rocsparse/library/src/primitives/rocsparse_segmented_radix_sort_pairs.cpp @@ -136,4 +136,6 @@ INSTANTIATE(int32_t, int32_t, int32_t); INSTANTIATE(int64_t, int64_t, int64_t); INSTANTIATE(int32_t, int64_t, const int64_t); INSTANTIATE(int32_t, int64_t, int64_t); +INSTANTIATE(int64_t, int32_t, const int32_t); +INSTANTIATE(int64_t, int32_t, int32_t); #undef INSTANTIATE