feat(rocsparse): add the rocsparse_spsort public API

projects/rocsparse/CHANGELOG.md

@@ -5,6 +5,10 @@ Documentation for rocSPARSE is available at
 
 ## (Unreleased) rocSPARSE 5.2.0
 
+### Added
+* Added the `rocsparse_spsort` and `rocsparse_spsort_buffer_size` generic routines, which sort a sparse matrix in place or into a separate target matrix. COO matrices are sorted by row or by column, and CSR and CSC matrices have the indices within each row or column sorted. The values are reordered with the indices. Uniform strided batched matrices are supported, including broadcasting a single source matrix to every sample of the target matrix.
+* Added the `rocsparse_spsort_descr_create`, `rocsparse_spsort_descr_destroy`, and `rocsparse_spsort_set_input` auxiliary routines, together with the `rocsparse_spsort_descr` type and the `rocsparse_spsort_alg`, `rocsparse_spsort_input`, and `rocsparse_spsort_stage` enums.
+
 ### Resolved issues
 * Fixed an overflow issue in `rocsparse_roti` and the generic `rocsparse_rot` routine. When using 64-bit indices and `nnz` >= `2^32`, a 32-bit element-index calculation could overflow, leaving some elements unrotated and causing low-index elements to be processed with incorrect data. The kernel now computes element indices in 64-bit arithmetic and uses a grid-stride loop with the launch grid clamped to the device limit.
 * Fixed an overflow issue in `rocsparse_Xsctr` and the generic `rocsparse_scatter` routine. When using 64-bit indices and `nnz` >= `2^32`, a 32-bit element-index calculation could overflow and prevent some elements from being scattered. The kernel now computes element indices in 64-bit arithmetic and uses a grid-stride loop with the launch grid clamped to the device limit.

projects/rocsparse/docs/how-to/using-rocsparse.rst

@@ -349,6 +349,8 @@ Function name                                        yes no
 :cpp:func:`rocsparse_v2_spmv_buffer_size()`              x
 :cpp:func:`rocsparse_spgeam_buffer_size()`               x
 :cpp:func:`rocsparse_spgeam()`                           x
+:cpp:func:`rocsparse_spsort_buffer_size()`               x
+:cpp:func:`rocsparse_spsort()`                           x
 :cpp:func:`rocsparse_sptrsv_buffer_size()`               x
 :cpp:func:`rocsparse_sptrsm_buffer_size()`               x
 :cpp:func:`rocsparse_spilu0_buffer_size()`               x

projects/rocsparse/docs/reference/api.rst

@@ -210,6 +210,12 @@ Auxiliary functions
 +-----------------------------------------------------+
 |:cpp:func:`rocsparse_spgeam_get_output`              |
 +-----------------------------------------------------+
+|:cpp:func:`rocsparse_spsort_descr_create`            |
++-----------------------------------------------------+
+|:cpp:func:`rocsparse_spsort_descr_destroy`           |
++-----------------------------------------------------+
+|:cpp:func:`rocsparse_spsort_set_input`               |
++-----------------------------------------------------+
 |:cpp:func:`rocsparse_create_spmv_descr`              |
 +-----------------------------------------------------+
 |:cpp:func:`rocsparse_destroy_spmv_descr`             |
@@ -664,6 +670,10 @@ documentation for details about the supported data types and compute types.
 +-----------------------------------------------------+
 |:cpp:func:`rocsparse_spgeam()`                       |
 +-----------------------------------------------------+
+|:cpp:func:`rocsparse_spsort_buffer_size()`           |
++-----------------------------------------------------+
+|:cpp:func:`rocsparse_spsort()`                       |
++-----------------------------------------------------+
 |:cpp:func:`rocsparse_spmat_scale()`                  |
 +-----------------------------------------------------+
 |:cpp:func:`rocsparse_sddmm_buffer_size()`            |

projects/rocsparse/docs/reference/auxiliary.rst

@@ -357,6 +357,21 @@ rocsparse_spgeam_get_output
 
 .. doxygenfunction:: rocsparse_spgeam_get_output
 
+rocsparse_spsort_descr_create
+-----------------------------
+
+.. doxygenfunction:: rocsparse_spsort_descr_create
+
+rocsparse_spsort_descr_destroy
+------------------------------
+
+.. doxygenfunction:: rocsparse_spsort_descr_destroy
+
+rocsparse_spsort_set_input
+--------------------------
+
+.. doxygenfunction:: rocsparse_spsort_set_input
+
 rocsparse_create_spmv_descr
 ---------------------------
 

projects/rocsparse/docs/reference/enumerations.rst

@@ -322,6 +322,21 @@ rocsparse_spgeam_output
 
 .. doxygenenum:: rocsparse_spgeam_output
 
+rocsparse_spsort_alg
+--------------------
+
+.. doxygenenum:: rocsparse_spsort_alg
+
+rocsparse_spsort_input
+----------------------
+
+.. doxygenenum:: rocsparse_spsort_input
+
+rocsparse_spsort_stage
+----------------------
+
+.. doxygenenum:: rocsparse_spsort_stage
+
 rocsparse_sparse_to_dense_alg
 -----------------------------
 

projects/rocsparse/docs/reference/generic.rst

@@ -147,6 +147,16 @@ rocsparse_spgeam()
 
 .. doxygenfunction:: rocsparse_spgeam
 
+rocsparse_spsort_buffer_size()
+------------------------------
+
+.. doxygenfunction:: rocsparse_spsort_buffer_size
+
+rocsparse_spsort()
+------------------
+
+.. doxygenfunction:: rocsparse_spsort
+
 rocsparse_spmat_scale()
 -----------------------
 

projects/rocsparse/docs/reference/types.rst

@@ -139,4 +139,9 @@ rocsparse_spgeam_descr
 ----------------------
 
 .. doxygentypedef:: rocsparse_spgeam_descr
+
+rocsparse_spsort_descr
+----------------------
+
+.. doxygentypedef:: rocsparse_spsort_descr
 

projects/rocsparse/library/include/internal/generic/rocsparse_spsort.h

@@ -0,0 +1,159 @@
+/*! \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.
+ *
+ * ************************************************************************ */
+
+#ifndef ROCSPARSE_SPSORT_H
+#define ROCSPARSE_SPSORT_H
+
+#include "../../rocsparse-types.h"
+#include "rocsparse/rocsparse-export.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/*! \ingroup generic_module
+*  \details
+*  \p rocsparse_spsort_buffer_size returns the size of the required buffer to execute the given stage of the SpSort operation.
+*
+*  \note
+*  This routine is executed on the host. It does not launch any work on the device, and it
+*  supports execution in a hipGraph context.
+*
+*  @param[in]
+*  handle       handle to the rocSPARSE library context queue.
+*  @param[in]
+*  descr        SpSort descriptor.
+*  @param[in]
+*  source       source sparse matrix descriptor, the matrix to sort.
+*  @param[in]
+*  target       target sparse matrix descriptor, the sorted output matrix. \p target can be the same
+*               descriptor as \p source to sort in place.
+*  @param[in]
+*  stage        SpSort stage for the SpSort computation.
+*  @param[out]
+*  buffer_size_in_bytes  size in bytes of the temporary storage buffer.
+*  @param[out]
+*  p_error      error descriptor created if the returned status is not \ref rocsparse_status_success. A null pointer can be passed if an error descriptor is not required.
+*
+*  \retval rocsparse_status_success the operation completed successfully.
+*  \retval rocsparse_status_invalid_handle the library context was not initialized.
+*  \retval rocsparse_status_invalid_pointer \p descr, \p source, \p target, or \p buffer_size_in_bytes pointer is invalid.
+*  \retval rocsparse_status_invalid_size \p source and \p target do not have the same dimensions or number of non-zeros, or the batch strides make the batches overlap.
+*  \retval rocsparse_status_invalid_value \p source and \p target do not have the same format, index types, data type or index base, the data type is \ref rocsparse_datatype_u8_r, \ref rocsparse_datatype_i32_r or \ref rocsparse_datatype_u32_r, which are not supported, the batch count of \p source is neither the batch count of \p target nor one with a batch stride of zero, or \p source is in COO format and the direction is not set.
+*  \retval rocsparse_status_not_implemented \p source is not in COO, CSR or CSC format.
+*/
+ROCSPARSE_EXPORT
+rocsparse_status rocsparse_spsort_buffer_size(rocsparse_handle            handle,
+                                              rocsparse_spsort_descr      descr,
+                                              rocsparse_const_spmat_descr source,
+                                              rocsparse_spmat_descr       target,
+                                              rocsparse_spsort_stage      stage,
+                                              size_t*                     buffer_size_in_bytes,
+                                              rocsparse_error*            p_error);
+
+/*! \ingroup generic_module
+*  \brief Sparse matrix sorting.
+*
+*  \details
+*  \p rocsparse_spsort sorts the source sparse matrix and writes the result to the target sparse
+*  matrix. The target matrix must be created with the same format, dimensions, number of
+*  non-zeros, index types, data type and index base as the source matrix. If the same descriptor
+*  is passed for \p source and \p target, the matrix is sorted in place. Otherwise, the arrays of
+*  the target matrix must not overlap the arrays of the source matrix.
+*
+*  The supported formats are \ref rocsparse_format_coo, \ref rocsparse_format_csr and
+*  \ref rocsparse_format_csc :
+*  - COO: the \ref rocsparse_direction set with \ref rocsparse_spsort_input_direction selects
+*    the sort order. \ref rocsparse_direction_row sorts the entries by row and then by column,
+*    and \ref rocsparse_direction_column sorts them by column and then by row.
+*  - CSR: the column indices within each row are sorted. The row pointer of the target matrix
+*    is a copy of the row pointer of the source matrix.
+*  - CSC: the row indices within each column are sorted. The column pointer of the target
+*    matrix is a copy of the column pointer of the source matrix.
+*
+*  The direction is only used by the COO format. It does not need to be set for the CSR and
+*  CSC formats, and is ignored if it is.
+*
+*  Uniform strided batched matrices are supported. They are set up with
+*  \ref rocsparse_coo_set_strided_batch, \ref rocsparse_csr_set_strided_batch or
+*  \ref rocsparse_csc_set_strided_batch. Every sample of the batch must have the same sparsity
+*  pattern, with the same offsets and indices in the same order, and only the values may differ.
+*  The indices are sorted once, and the resulting permutation is applied to the values of every
+*  sample.
+*
+*  The batch count is given by the target matrix. The source matrix must either have the same
+*  batch count, or a batch count of one and a batch stride of zero, in which case every sample of
+*  the target matrix is the sorted source matrix. The batch stride must be at least the number of
+*  non-zeros. For the CSR and CSC formats, the offsets batch stride must be either zero, when all
+*  samples share one offsets array, or at least the size of the offsets array.
+*
+*  \note
+*  The \ref rocsparse_spsort_stage_compute stage is non-blocking and executed asynchronously with
+*  respect to the host for the CSR and CSC formats. For the COO format, it is blocking with respect
+*  to the host.
+*
+*  \note
+*  This routine does not support execution in a hipGraph context.
+*
+*  @param[in]
+*  handle       handle to the rocSPARSE library context queue.
+*  @param[in]
+*  descr        SpSort descriptor.
+*  @param[in]
+*  source       source sparse matrix descriptor, the matrix to sort.
+*  @param[inout]
+*  target       target sparse matrix descriptor, the sorted output matrix. \p target can be the same
+*               descriptor as \p source to sort in place.
+*  @param[in]
+*  stage        SpSort stage for the SpSort computation.
+*  @param[in]
+*  buffer_size_in_bytes  size in bytes of the temporary storage buffer. \p buffer_size_in_bytes is
+*               determined by calling \ref rocsparse_spsort_buffer_size.
+*  @param[in]
+*  buffer       temporary storage buffer allocated by the user.
+*  @param[out]
+*  p_error      error descriptor created if the returned status is not \ref rocsparse_status_success. A null pointer can be passed if an error descriptor is not required.
+*
+*  \retval rocsparse_status_success the operation completed successfully.
+*  \retval rocsparse_status_invalid_handle the library context was not initialized.
+*  \retval rocsparse_status_invalid_pointer \p descr, \p source, \p target, or \p buffer pointer is invalid.
+*  \retval rocsparse_status_invalid_size \p source and \p target do not have the same dimensions or number of non-zeros, or the batch strides make the batches overlap.
+*  \retval rocsparse_status_invalid_value \p source and \p target do not have the same format, index types, data type or index base, the data type is \ref rocsparse_datatype_u8_r, \ref rocsparse_datatype_i32_r or \ref rocsparse_datatype_u32_r, which are not supported, the batch count of \p source is neither the batch count of \p target nor one with a batch stride of zero, or \p source is in COO format and the direction is not set.
+*  \retval rocsparse_status_not_implemented \p source is not in COO, CSR or CSC format.
+*/
+ROCSPARSE_EXPORT
+rocsparse_status rocsparse_spsort(rocsparse_handle            handle,
+                                  rocsparse_spsort_descr      descr,
+                                  rocsparse_const_spmat_descr source,
+                                  rocsparse_spmat_descr       target,
+                                  rocsparse_spsort_stage      stage,
+                                  size_t                      buffer_size_in_bytes,
+                                  void*                       buffer,
+                                  rocsparse_error*            p_error);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* ROCSPARSE_SPSORT_H */

projects/rocsparse/library/include/internal/rocsparse-generic.h

@@ -53,6 +53,7 @@
 #include "generic/rocsparse_spmm.h"
 #include "generic/rocsparse_spmv.h"
 #include "generic/rocsparse_spsm.h"
+#include "generic/rocsparse_spsort.h"
 #include "generic/rocsparse_spsv.h"
 #include "generic/rocsparse_sptrsm.h"
 #include "generic/rocsparse_sptrsv.h"

projects/rocsparse/library/include/rocsparse-auxiliary.h

@@ -1485,6 +1485,80 @@ rocsparse_status rocsparse_create_extract_descr(rocsparse_extract_descr*    desc
 ROCSPARSE_EXPORT
 rocsparse_status rocsparse_destroy_extract_descr(rocsparse_extract_descr descr);
 
+/*! \ingroup aux_module
+*  \brief Sparse matrix SpSort routine descriptor creation.
+*
+*  \details
+*  \p rocsparse_spsort_descr_create creates the descriptor of the \ref rocsparse_spsort_buffer_size and
+*  \ref rocsparse_spsort routines.
+*
+*  @param[in]
+*  handle         the handle to the rocSPARSE library context.
+*  @param[out]
+*  p_spsort_descr pointer to the descriptor of the SpSort routine.
+*  @param[out]
+*  p_error        error descriptor created if the returned status is not \ref rocsparse_status_success. A null pointer can be passed if an error descriptor is not required.
+*
+*  \retval      rocsparse_status_success the operation completed successfully.
+*  \retval      rocsparse_status_invalid_handle \p handle pointer is invalid.
+*  \retval      rocsparse_status_invalid_pointer \p p_spsort_descr pointer is invalid.
+*/
+ROCSPARSE_EXPORT
+rocsparse_status rocsparse_spsort_descr_create(rocsparse_handle        handle,
+                                               rocsparse_spsort_descr* p_spsort_descr,
+                                               rocsparse_error*        p_error);
+
+/*! \ingroup aux_module
+*  \brief Destroy a sparse matrix SpSort descriptor.
+*
+*  \details
+*  \p rocsparse_spsort_descr_destroy destroys the descriptor of the \ref rocsparse_spsort_buffer_size and
+*  \ref rocsparse_spsort routines.
+*
+*  @param[in]
+*  handle         the handle to the rocSPARSE library context.
+*  @param[in]
+*  spsort_descr   descriptor of the SpSort routine.
+*  @param[out]
+*  p_error        error descriptor created if the returned status is not \ref rocsparse_status_success. A null pointer can be passed if an error descriptor is not required.
+*
+*  \retval      rocsparse_status_success the operation completed successfully.
+*  \retval      rocsparse_status_invalid_handle \p handle pointer is invalid.
+*/
+ROCSPARSE_EXPORT
+rocsparse_status rocsparse_spsort_descr_destroy(rocsparse_handle       handle,
+                                                rocsparse_spsort_descr spsort_descr,
+                                                rocsparse_error*       p_error);
+
+/*! \ingroup aux_module
+ *  \brief Set the requested \ref rocsparse_spsort_input data in the SpSort descriptor.
+ *
+ *  @param[in]
+ *  handle      the pointer to the handle to the rocSPARSE library context.
+ *  @param[inout]
+ *  descr       the pointer to the SpSort descriptor.
+ *  @param[in]
+ *  input       one of the values from \ref rocsparse_spsort_input.
+ *  @param[in]
+ *  data        pointer to the input data, of the type listed for \p input in \ref rocsparse_spsort_input.
+ *  @param[in]
+ *  data_size_in_bytes   size in bytes of the input data, which must be the size of that type.
+ *  @param[out]
+ *  p_error        error descriptor created if the returned status is not \ref rocsparse_status_success. A null pointer can be passed if an error descriptor is not required.
+ *
+ *  \retval rocsparse_status_success the operation completed successfully.
+ *  \retval rocsparse_status_invalid_pointer if \p descr or \p data is invalid.
+ *  \retval rocsparse_status_invalid_value if \p input is invalid, or if the analysis stage has already been executed.
+ *  \retval rocsparse_status_invalid_size if \p data_size_in_bytes is invalid.
+ */
+ROCSPARSE_EXPORT
+rocsparse_status rocsparse_spsort_set_input(rocsparse_handle       handle,
+                                            rocsparse_spsort_descr descr,
+                                            rocsparse_spsort_input input,
+                                            const void*            data,
+                                            size_t                 data_size_in_bytes,
+                                            rocsparse_error*       p_error);
+
 /*! \ingroup aux_module
 *  \brief Sparse matrix SpGEAM routine descriptor creation.
 *