Use cudaMallocAsync in GB_cuda_cumsum

Author: VidithM

CUDA/template/GB_cuda_cumsum.cuh

@@ -1,5 +1,5 @@
 //------------------------------------------------------------------------------
-// GraphBLAS/CUDA/template/GB_cuda_cumsum: cumlative sum of array on the GPU(s)
+// GraphBLAS/CUDA/template/GB_cuda_cumsum: cumlative sum of an array on the GPU(s)
 //------------------------------------------------------------------------------
 
 // SuiteSparse:GraphBLAS, Timothy A. Davis, (c) 2017-2025, All Rights Reserved.
@@ -12,9 +12,9 @@
 
 #include <cub/cub.cuh>
 
-#define GB_FREE_ALL             \
-{                               \
-    cudaFree (d_temp_storage) ; \
+#define GB_FREE_ALL                             \
+{                                               \
+    cudaFreeAsync (d_temp_storage, stream) ;    \
 }
 
 typedef enum GB_cuda_cumsum_type
@@ -23,7 +23,7 @@ typedef enum GB_cuda_cumsum_type
     GB_CUDA_CUMSUM_INCLUSIVE
 } GB_cuda_cumsum_type ;
 
-__host__ GrB_Info GB_cuda_cumsum    // compute the cumulative sum of an array
+__host__ GrB_Info GB_cuda_cumsum // compute the cumsum of an array
 (
     int64_t *__restrict__ out,   // size n or n+1, output.
     int64_t *__restrict__ in,    // size n or n+1, input
@@ -43,7 +43,7 @@ __host__ GrB_Info GB_cuda_cumsum    // compute the cumulative sum of an array
     ASSERT (out != NULL) ;
     ASSERT (n >= 0) ;
 
-    void *d_temp_storage = NULL;
+    void *d_temp_storage = NULL ;
     size_t temp_storage_bytes = 0 ;
 
     switch (type)
@@ -57,20 +57,23 @@ __host__ GrB_Info GB_cuda_cumsum    // compute the cumulative sum of an array
                 in, out, n, stream) ;
     }
 
-    CUDA_OK (cudaMalloc (&d_temp_storage, temp_storage_bytes)) ;
+    CUDA_OK (cudaMallocAsync (&d_temp_storage, temp_storage_bytes, stream)) ;
+    CUDA_OK (cudaStreamSynchronize (stream)) ;
 
     // Run
     switch (type)
     {
         case GB_CUDA_CUMSUM_INCLUSIVE:
-            cub::DeviceScan::InclusiveSum (d_temp_storage, temp_storage_bytes,
-                in, out, n, stream) ;
+            cub::DeviceScan::InclusiveSum (d_temp_storage,  
+                temp_storage_bytes, in, out, n, stream) ;
             break;
         default:
-            cub::DeviceScan::ExclusiveSum (d_temp_storage, temp_storage_bytes,
-                in, out, n, stream) ;
+            cub::DeviceScan::ExclusiveSum (d_temp_storage,
+                temp_storage_bytes, in, out, n, stream) ;
     }
-    cudaFree (d_temp_storage) ;
+
+    cudaFreeAsync (d_temp_storage, stream) ;
+    cudaStreamSynchronize (stream) ;
 
     return GrB_SUCCESS;
 }