Write CUDA kernel

Author: gdalle

ext/SparseMatrixColoringsCUDAExt.jl

@@ -2,7 +2,8 @@ module SparseMatrixColoringsCUDAExt
 
 import SparseMatrixColorings as SMC
 using SparseArrays: SparseMatrixCSC, rowvals, nnz, nzrange
-using CUDA: CuVector, CuMatrix
+using CUDA:
+    @cuda, CuVector, CuMatrix, blockIdx, blockDim, gridDim, threadIdx, launch_configuration
 using CUDA.CUSPARSE: AbstractCuSparseMatrix, CuSparseMatrixCSC, CuSparseMatrixCSR
 
 SMC.matrix_versions(A::AbstractCuSparseMatrix) = (A,)
@@ -53,29 +54,32 @@ function SMC.StarSetColoringResult(
     return SMC.StarSetColoringResult(A, ag, color, result_cpu.group, compressed_indices)
 end
 
-# TODO: write a kernel
-
-function SMC.decompress!(
-    A::CuSparseMatrixCSC, B::CuMatrix, result::SMC.ColumnColoringResult{<:CuSparseMatrixCSC}
-)
-    A.nzVal .= getindex.(Ref(B), result.compressed_indices)
-    return A
-end
-
-function SMC.decompress!(
-    A::CuSparseMatrixCSC, B::CuMatrix, result::SMC.RowColoringResult{<:CuSparseMatrixCSC}
+function update_nzval_from_matrix!(
+    nzVal::AbstractVector, B::AbstractMatrix, compressed_indices::AbstractVector{<:Integer}
 )
-    A.nzVal .= getindex.(Ref(B), result.compressed_indices)
-    return A
+    index = (blockIdx().x - 1) * blockDim().x + threadIdx().x
+    stride = gridDim().x * blockDim().x
+    for k in index:stride:length(nzVal)
+        nzVal[k] = B[compressed_indices[k]]
+    end
+    return nothing
 end
 
-function SMC.decompress!(
-    A::CuSparseMatrixCSC,
-    B::CuMatrix,
-    result::SMC.StarSetColoringResult{<:CuSparseMatrixCSC},
-)
-    A.nzVal .= getindex.(Ref(B), result.compressed_indices)
-    return A
+for R in (:ColumnColoringResult, :RowColoringResult, :StarSetColoringResult)
+    # loop to avoid method ambiguity
+    @eval function SMC.decompress!(
+        A::CuSparseMatrixCSC, B::CuMatrix, result::SMC.$R{<:CuSparseMatrixCSC}
+    )
+        A.nnz == 0 && return A
+        kernel = @cuda launch = false update_nzval_from_matrix!(
+            A.nzVal, B, result.compressed_indices
+        )
+        config = launch_configuration(kernel.fun)
+        threads = min(A.nnz, config.threads)
+        blocks = cld(A.nnz, threads)
+        kernel(A.nzVal, B, result.compressed_indices; threads, blocks)
+        return A
+    end
 end
 
 end