patch for cusparseLt 0.6.1

cupy_backends/cuda/libs/cusparselt.pxd

@@ -13,18 +13,20 @@ cpdef enum:
     CUSPARSELT_MAT_BATCH_STRIDE = 1  # READ/WRITE
 
     # cusparseComputeType
-    CUSPARSE_COMPUTE_16F = 0
-    CUSPARSE_COMPUTE_32I = 1
-    CUSPARSE_COMPUTE_TF32 = 2
-    CUSPARSE_COMPUTE_TF32_FAST = 3
+    CUSPARSE_COMPUTE_32I = 0
+    CUSPARSE_COMPUTE_16F = 1
+    CUSPARSE_COMPUTE_32F = 2
 
     # cusparseLtMatmulDescAttribute_t
     CUSPARSELT_MATMUL_ACTIVATION_RELU = 0             # READ/WRITE
     CUSPARSELT_MATMUL_ACTIVATION_RELU_UPPERBOUND = 1  # READ/WRITE
     CUSPARSELT_MATMUL_ACTIVATION_RELU_THRESHOLD = 2   # READ/WRITE
     CUSPARSELT_MATMUL_ACTIVATION_GELU = 3             # READ/WRITE
-    CUSPARSELT_MATMUL_BIAS_STRIDE = 4                 # READ/WRITE
-    CUSPARSELT_MATMUL_BIAS_POINTER = 5                # READ/WRITE
+    CUSPARSELT_MATMUL_ACTIVATION_GELU_SCALING = 4     # READ/WRITE
+    CUSPARSELT_MATMUL_ALPHA_VECTOR_SCALING = 5        # READ/WRITE
+    CUSPARSELT_MATMUL_BETA_VECTOR_SCALING = 6         # READ/WRITE
+    CUSPARSELT_MATMUL_BIAS_POINTER = 7                # READ/WRITE
+    CUSPARSELT_MATMUL_BIAS_STRIDE = 8                 # READ/WRITE
 
     # cusparseLtMatmulAlg_t
     CUSPARSELT_MATMUL_ALG_DEFAULT = 0
@@ -33,6 +35,14 @@ cpdef enum:
     CUSPARSELT_MATMUL_ALG_CONFIG_ID = 0      # NOQA, READ/WRITE
     CUSPARSELT_MATMUL_ALG_CONFIG_MAX_ID = 1  # NOQA, READ-ONLY
     CUSPARSELT_MATMUL_SEARCH_ITERATIONS = 2  # NOQA, READ/WRITE
+    CUSPARSELT_MATMUL_SPLIT_K = 3            # NOQA, READ/WRITE
+    CUSPARSELT_MATMUL_SPLIT_K_MODE =4        # NOQA, READ/WRITE
+    CUSPARSELT_MATMUL_SPLIT_K_BUFFERS=5      # NOQA, READ/WRITE
+
+    # cusparseLtSplitKMode_t
+    CUSPARSELT_INVALID_MODE = 0
+    CUSPARSELT_SPLIT_K_MODE_ONE_KERNEL  = 1
+    CUSPARSELT_SPLIT_K_MODE_TWO_KERNELS = 2
 
     # cusparseLtPruneAlg_t
     CUSPARSELT_PRUNE_SPMMA_TILE = 0

cupy_backends/cuda/libs/cusparselt.pyx

@@ -11,6 +11,13 @@ from cupy_backends.cuda.api cimport runtime
 from cupy_backends.cuda.libs import cusparse as _cusparse
 
 
+###############################################################################
+# Types
+###############################################################################
+cdef extern from *:
+    ctypedef void* LibraryPropertyType 'libraryPropertyType_t'
+
+
 cdef extern from '../../cupy_cusparselt.h' nogil:
     ctypedef int cusparseStatus_t 'cusparseStatus_t'
     ctypedef int cusparseOrder_t 'cusparseOrder_t'
@@ -36,11 +43,14 @@ cdef extern from '../../cupy_cusparselt.h' nogil:
     ctypedef int cusparseOperation_t 'cusparseOperation_t'
     ctypedef int cusparseLtMatmulAlg_t 'cusparseLtMatmulAlg_t'
     ctypedef int cusparseLtMatmulAlgAttribute_t 'cusparseLtMatmulAlgAttribute_t'  # NOQA
+    ctypedef int cusparseLtSplitKMode_t 'cusparseLtSplitKMode_t'
     ctypedef int cusparseLtPruneAlg_t 'cusparseLtPruneAlg_t'
 
     # Management Functions
     cusparseStatus_t cusparseLtInit(cusparseLtHandle_t* handle)
     cusparseStatus_t cusparseLtDestroy(const cusparseLtHandle_t* handle)
+    cusparseStatus_t cusparseLtGetVersion(const cusparseLtHandle_t* handle, int* version)
+    cusparseStatus_t cusparseLtGetProperty(LibraryPropertyType propertyType, int* value)
 
     # Matmul Functions
     cusparseStatus_t cusparseLtDenseDescriptorInit(
@@ -66,6 +76,11 @@ cdef extern from '../../cupy_cusparselt.h' nogil:
         const cusparseLtMatDescriptor_t* matDescr,
         cusparseLtMatDescAttribute_t matAttribute,
         void* data, size_t dataSize)
+    cusparseStatus_t cusparseLtMatDescSetAttribute(
+        const cusparseLtHandle_t* handle,
+        const cusparseLtMatDescriptor_t* matDescr,
+        cusparseLtMatDescAttribute_t matAttribute,
+        void* data, size_t dataSize)
     cusparseStatus_t cusparseLtMatmulDescriptorInit(
         const cusparseLtHandle_t* handle,
         cusparseLtMatmulDescriptor_t* matMulDescr,
@@ -95,24 +110,33 @@ cdef extern from '../../cupy_cusparselt.h' nogil:
         const cusparseLtHandle_t* handle,
         cusparseLtMatmulAlgSelection_t* algSelection,
         cusparseLtMatmulAlgAttribute_t attribute,
-        const void* data, size_t ataSize)
-    cusparseStatus_t cusparseLtMatmulGetWorkspace(
+        const void* data, size_t dataSize)
+    cusparseStatus_t cusparseLtMatmulAlgGetAttribute(
         const cusparseLtHandle_t* handle,
         const cusparseLtMatmulAlgSelection_t* algSelection,
+        cusparseLtMatmulAlgAttribute_t attribute,
+        void* data, size_t dataSize)
+    cusparseStatus_t cusparseLtMatmulGetWorkspace(
+        const cusparseLtHandle_t* handle,
+        const cusparseLtMatmulPlan_t* plan,
         size_t* workspaceSize)
     cusparseStatus_t cusparseLtMatmulPlanInit(
         const cusparseLtHandle_t* handle,
         cusparseLtMatmulPlan_t* plan,
         const cusparseLtMatmulDescriptor_t* matmulDescr,
-        const cusparseLtMatmulAlgSelection_t* algSelection,
-        size_t workspaceSize)
+        const cusparseLtMatmulAlgSelection_t* algSelection)
     cusparseStatus_t cusparseLtMatmulPlanDestroy(
         const cusparseLtMatmulPlan_t* plan)
     cusparseStatus_t cusparseLtMatmul(
         const cusparseLtHandle_t* handle, const cusparseLtMatmulPlan_t* plan,
         const void* alpha, const void* d_A, const void* d_B,
         const void* beta, const void* d_C, void* d_D,
         void* workspace, runtime.Stream* streams, int32_t numStreams)
+    cusparseStatus_t cusparseLtMatmulSearch(
+        const cusparseLtHandle_t* handle, cusparseLtMatmulPlan_t* plan,
+        const void* alpha, const void* d_A, const void* d_B,
+        const void* beta, const void* d_C, void* d_D,
+        void* workspace, runtime.Stream* streams, int32_t numStreams)
 
     # Helper Functions
     cusparseStatus_t cusparseLtSpMMAPrune(
@@ -123,7 +147,7 @@ cdef extern from '../../cupy_cusparselt.h' nogil:
     cusparseStatus_t cusparseLtSpMMAPruneCheck(
         const cusparseLtHandle_t* handle,
         const cusparseLtMatmulDescriptor_t* matmulDescr,
-        const void* d_in, int* valid, runtime.Stream stream)
+        const void* d_in, int* d_valid, runtime.Stream stream)
     cusparseStatus_t cusparseLtSpMMAPrune2(
         const cusparseLtHandle_t* handle,
         const cusparseLtMatDescriptor_t* sparseMatDescr,
@@ -136,19 +160,22 @@ cdef extern from '../../cupy_cusparselt.h' nogil:
         runtime.Stream stream)
     cusparseStatus_t cusparseLtSpMMACompressedSize(
         const cusparseLtHandle_t* handle, const cusparseLtMatmulPlan_t* plan,
-        size_t* compressedSize)
+        size_t* compressedSize,
+        size_t* compressBufferSize)
     cusparseStatus_t cusparseLtSpMMACompress(
         const cusparseLtHandle_t* handle, const cusparseLtMatmulPlan_t* plan,
-        const void* d_dense, void* d_compressed, runtime.Stream stream)
+        const void* d_dense, void* d_compressed, void* d_compressed_buffer,
+        runtime.Stream stream)
     cusparseStatus_t cusparseLtSpMMACompressedSize2(
         const cusparseLtHandle_t* handle,
         const cusparseLtMatDescriptor_t* sparseMatDescr,
-        size_t* compressedSize)
+        size_t* compressedSize,
+        size_t* compressBufferSize)
     cusparseStatus_t cusparseLtSpMMACompress2(
         const cusparseLtHandle_t* handle,
         const cusparseLtMatDescriptor_t* sparseMatDescr,
         int isSparseA, cusparseOperation_t op, const void* d_dense,
-        void* d_compressed, runtime.Stream stream)
+        void* d_compressed, void* d_compressed_buffer, runtime.Stream stream)
 
     # Build-time version
     int CUSPARSELT_VERSION
@@ -233,6 +260,7 @@ cdef class MatmulPlan:
         return <intptr_t>self._ptr
 
 
+
 ###############################################################################
 # Error handling
 ###############################################################################
@@ -290,23 +318,23 @@ cpdef matDescriptorDestroy(MatDescriptor matDescr):
     check_status(status)
 
 cpdef matDescSetAttribute(Handle handle, MatDescriptor matDescr,
-                          matAttribute, size_t data, size_t dataSize):
+                          matAttribute, data, size_t dataSize):
     """Sets the attribute related to matrix descriptor."""
     status = cusparseLtMatDescSetAttribute(
         <const cusparseLtHandle_t*> handle._ptr,
         <cusparseLtMatDescriptor_t*> matDescr._ptr,
         <cusparseLtMatDescAttribute_t> matAttribute,
-        <const void*> data, dataSize)
+        <const void*> data, <size_t> dataSize)
     check_status(status)
 
 cpdef matDescGetAttribute(Handle handle, MatDescriptor matDescr,
-                          matAttribute, size_t data, size_t dataSize):
+                          matAttribute, data, size_t dataSize):
     """Gets the attribute related to matrix descriptor."""
     status = cusparseLtMatDescGetAttribute(
         <const cusparseLtHandle_t*> handle._ptr,
         <const cusparseLtMatDescriptor_t*> matDescr._ptr,
         <cusparseLtMatDescAttribute_t> matAttribute,
-        <void*> data, dataSize)
+        <void*> data, <size_t> dataSize)
     check_status(status)
 
 cpdef matmulDescriptorInit(Handle handle,
@@ -331,23 +359,23 @@ cpdef matmulDescriptorInit(Handle handle,
     check_status(status)
 
 cpdef matmulDescSetAttribute(Handle handle, MatmulDescriptor matmulDescr,
-                             matmulAttribute, size_t data, size_t dataSize):
+                             matmulAttribute, data, size_t dataSize):
     """Sets the attribute related to matmul descriptor."""
     status = cusparseLtMatmulDescSetAttribute(
         <const cusparseLtHandle_t*> handle._ptr,
         <cusparseLtMatmulDescriptor_t*> matmulDescr._ptr,
         <cusparseLtMatmulDescAttribute_t> matmulAttribute,
-        <const void*> data, dataSize)
+        <const void*> data, <size_t> dataSize)
     check_status(status)
 
 cpdef matmulDescGetAttribute(Handle handle, MatmulDescriptor matmulDescr,
-                             matmulAttribute, size_t data, size_t dataSize):
+                             matmulAttribute, data, size_t dataSize):
     """Gets the attribute related to matmul descriptor."""
     status = cusparseLtMatmulDescGetAttribute(
         <const cusparseLtHandle_t*> handle._ptr,
         <const cusparseLtMatmulDescriptor_t*> matmulDescr._ptr,
         <cusparseLtMatmulDescAttribute_t> matmulAttribute,
-        <void*> data, dataSize)
+        <void*> data, <size_t> dataSize)
     check_status(status)
 
 cpdef matmulAlgSelectionInit(Handle handle, MatmulAlgSelection algSelection,
@@ -361,7 +389,8 @@ cpdef matmulAlgSelectionInit(Handle handle, MatmulAlgSelection algSelection,
     check_status(status)
 
 cpdef matmulAlgSetAttribute(Handle handle, MatmulAlgSelection algSelection,
-                            attribute, size_t data, size_t dataSize):
+                            attribute,
+                            data, size_t dataSize):
     """Sets the attribute related to algorithm selection descriptor."""
     status = cusparseLtMatmulAlgSetAttribute(
         <const cusparseLtHandle_t*> handle._ptr,
@@ -370,28 +399,37 @@ cpdef matmulAlgSetAttribute(Handle handle, MatmulAlgSelection algSelection,
         <const void*> data, <size_t> dataSize)
     check_status(status)
 
+cpdef matmulAlgGetAttribute(Handle handle, MatmulAlgSelection algSelection,
+                            attribute,
+                            data, size_t dataSize):
+    """Gets the attribute related to algorithm selection descriptor."""
+    status = cusparseLtMatmulAlgGetAttribute(
+        <const cusparseLtHandle_t*> handle._ptr,
+        <const cusparseLtMatmulAlgSelection_t*> algSelection._ptr,
+        <cusparseLtMatmulAlgAttribute_t> attribute,
+        <void*> data, <size_t> dataSize)
+    check_status(status)
+
 cpdef size_t matmulGetWorkspace(Handle handle,
-                                MatmulAlgSelection algSelection):
+                                MatmulPlan plan):
     """Determines the required workspace size"""
     cdef size_t workspaceSize
     status = cusparseLtMatmulGetWorkspace(
         <const cusparseLtHandle_t*> handle._ptr,
-        <const cusparseLtMatmulAlgSelection_t*> algSelection._ptr,
-        &workspaceSize)
+        <const cusparseLtMatmulPlan_t*> plan._ptr,
+        <size_t*> &workspaceSize)
     check_status(status)
     return workspaceSize
 
 cpdef matmulPlanInit(Handle handle, MatmulPlan plan,
                      MatmulDescriptor matmulDescr,
-                     MatmulAlgSelection algSelection,
-                     size_t workspaceSize):
+                     MatmulAlgSelection algSelection):
     """Initializes the plan."""
     status = cusparseLtMatmulPlanInit(
         <const cusparseLtHandle_t*> handle._ptr,
         <cusparseLtMatmulPlan_t*> plan._ptr,
         <const cusparseLtMatmulDescriptor_t*> matmulDescr._ptr,
-        <const cusparseLtMatmulAlgSelection_t*> algSelection._ptr,
-        <size_t> workspaceSize)
+        <const cusparseLtMatmulAlgSelection_t*> algSelection._ptr)
     check_status(status)
 
 cpdef matmulPlanDestroy(MatmulPlan plan):
@@ -412,7 +450,19 @@ cpdef matmul(Handle handle, MatmulPlan plan,
         <void*> workspace, <runtime.Stream*> NULL, <int32_t> 0)
     check_status(status)
 
-###############################################################################
+cpdef matmulSearch(Handle handle, MatmulPlan plan,
+             size_t alpha, size_t d_A, size_t d_B,
+             size_t beta, size_t d_C, size_t d_D, size_t workspace):
+    """Evaluates all available algorithms for the matrix multiplication"""
+    status = cusparseLtMatmulSearch(
+        <const cusparseLtHandle_t*> handle._ptr,
+        <cusparseLtMatmulPlan_t*> plan._ptr,
+        <const void*> alpha, <const void*> d_A, <const void*> d_B,
+        <const void*> beta, <const void*> d_C, <void*> d_D,
+        <void*> workspace, <runtime.Stream*> NULL, <int32_t> 0)
+    check_status(status)
+
+############<void*> ###################################################################
 # cuSPARSELt: Helper Functions
 ###############################################################################
 
@@ -428,16 +478,16 @@ cpdef spMMAPrune(Handle handle, MatmulDescriptor matmulDescr,
     check_status(status)
 
 cpdef spMMAPruneCheck(Handle handle, MatmulDescriptor matmulDescr,
-                      size_t d_in, size_t valid):
+                      size_t d_in, size_t d_valid):
     """Checks the correctness of the pruning structure"""
     cdef intptr_t stream = stream_module.get_current_stream_ptr()
     status = cusparseLtSpMMAPruneCheck(
         <const cusparseLtHandle_t*> handle._ptr,
         <const cusparseLtMatmulDescriptor_t*> matmulDescr._ptr,
-        <const void*> d_in, <int*> valid, <runtime.Stream> stream)
+        <const void*> d_in, <int*> d_valid, <runtime.Stream> stream)
     check_status(status)
 
-cpdef spMMAPrune2(Handle handle, MatDescriptor sparseMatDescr, isSparseA,
+cpdef spMMAPrune2(Handle handle, MatDescriptor sparseMatDescr, int isSparseA,
                   op, size_t d_in, size_t d_out, pruneAlg):
     """Prunes a dense matrix d_in"""
     cdef intptr_t stream = stream_module.get_current_stream_ptr()
@@ -460,47 +510,49 @@ cpdef spMMAPruneCheck2(Handle handle, MatDescriptor sparseMatDescr, isSparseA,
         <int*> d_valid, <runtime.Stream> stream)
     check_status(status)
 
-cpdef size_t spMMACompressedSize(Handle handle, MatmulPlan plan):
+cpdef spMMACompressedSize(Handle handle, MatmulPlan plan):
     """Provides the size of the compressed matrix"""
-    cdef size_t compressedSize
+    cdef size_t compressedSize, compressBufferSize
     status = cusparseLtSpMMACompressedSize(
         <const cusparseLtHandle_t*> handle._ptr,
         <const cusparseLtMatmulPlan_t*> plan._ptr,
-        &compressedSize)
+        <size_t*> &compressedSize, <size_t*> &compressBufferSize)
     check_status(status)
-    return compressedSize
+    return compressedSize, compressBufferSize
 
-cpdef spMMACompress(Handle handle, MatmulPlan plan,
-                    size_t d_dense, size_t d_compressed):
+cpdef spMMACompress(Handle handle, MatmulPlan plan, size_t d_dense,
+                    size_t d_compressed, size_t d_compressed_buffer):
     """Compresses a dense matrix d_dense."""
     cdef intptr_t stream = stream_module.get_current_stream_ptr()
     status = cusparseLtSpMMACompress(
         <const cusparseLtHandle_t*> handle._ptr,
-        <const cusparseLtMatmulPlan_t*> plan._ptr,
-        <const void*> d_dense, <void*> d_compressed, <runtime.Stream> stream)
+        <const cusparseLtMatmulPlan_t*> plan._ptr, <const void*> d_dense,
+        <void*> d_compressed, <void*> d_compressed_buffer,
+        <runtime.Stream> stream)
     check_status(status)
 
-cpdef size_t spMMACompressedSize2(Handle handle, MatDescriptor sparseMatDescr):
+cpdef spMMACompressedSize2(Handle handle, MatDescriptor sparseMatDescr):
     """Provides the size of the compressed matrix"""
-    cdef size_t compressedSize
+    cdef size_t compressedSize, compressBufferSize
     status = cusparseLtSpMMACompressedSize2(
         <const cusparseLtHandle_t*> handle._ptr,
         <const cusparseLtMatDescriptor_t*> sparseMatDescr._ptr,
-        &compressedSize)
+        <size_t*> &compressedSize, <size_t*> &compressBufferSize)
     check_status(status)
-    return compressedSize
+    return compressedSize, compressBufferSize
 
 cpdef spMMACompress2(Handle handle, MatDescriptor sparseMatDescr,
-                     isSparseA, op, size_t d_dense, size_t d_compressed):
+                     isSparseA, op, size_t d_dense,
+                     size_t d_compressed, size_t d_compressed_buffer):
     """Compresses a dense matrix d_dense."""
     cdef intptr_t stream = stream_module.get_current_stream_ptr()
     status = cusparseLtSpMMACompress2(
         <const cusparseLtHandle_t*> handle._ptr,
         <const cusparseLtMatDescriptor_t*> sparseMatDescr._ptr,
         <int> isSparseA, <cusparseOperation_t> op, <const void*> d_dense,
-        <void*> d_compressed, <runtime.Stream> stream)
+        <void*> d_compressed, <void*> d_compressed_buffer,
+        <runtime.Stream> stream)
     check_status(status)
 
-
 def get_build_version():
     return CUSPARSELT_VERSION