Move kernels outside functions and use AcceleratedKernels with JLBackend fallback

Co-authored-by: albertomercurio <[email protected]>

Author: Copilot

src/conversions/conversion_kernels.jl

@@ -33,3 +33,39 @@ end
         nzval_out[j] = nzval_in[j]
     end
 end
+
+# Kernel for creating sort keys for COO → CSC conversion
+@kernel inbounds=true function kernel_make_csc_keys!(
+    keys,
+    @Const(rowind),
+    @Const(colind),
+    @Const(n),
+)
+    i = @index(Global)
+    keys[i] = colind[i] * n + rowind[i]
+end
+
+# Kernel for creating sort keys for COO → CSR conversion
+@kernel inbounds=true function kernel_make_csr_keys!(
+    keys,
+    @Const(rowind),
+    @Const(colind),
+    @Const(m),
+)
+    i = @index(Global)
+    keys[i] = rowind[i] * m + colind[i]
+end
+
+# Kernel for counting entries per column (for COO → CSC)
+@kernel inbounds=true function kernel_count_per_col!(colptr, @Const(colind_sorted))
+    i = @index(Global)
+    col = colind_sorted[i]
+    @atomic colptr[col + 1] += 1
+end
+
+# Kernel for counting entries per row (for COO → CSR)
+@kernel inbounds=true function kernel_count_per_row!(rowptr, @Const(rowind_sorted))
+    i = @index(Global)
+    row = rowind_sorted[i]
+    @atomic rowptr[row + 1] += 1
+end

src/conversions/conversions.jl

@@ -1,5 +1,8 @@
 # Conversions between CSC, CSR, and COO sparse matrix formats
-# All conversions operate entirely on-device without CPU transfers
+# All conversions operate on-device, with CPU fallback only for JLBackend
+
+# Helper function to check if backend is JLBackend (which doesn't support AcceleratedKernels)
+_is_jlbackend(backend) = string(typeof(backend)) == "JLBackend"
 
 # ============================================================================
 # CSC ↔ COO Conversions
@@ -71,29 +74,22 @@ function DeviceSparseMatrixCSC(A::DeviceSparseMatrixCOO{Tv,Ti}) where {Tv,Ti}
     backend = get_backend(A.nzval)
 
     # Create keys for sorting: column first, then row
-    # We use n * rowind + colind to create a unique sortable key
     keys = similar(A.rowind, Ti, nnz_count)
 
     # Create keys on device
-    @kernel inbounds=true function make_keys!(
-        keys,
-        @Const(rowind),
-        @Const(colind),
-        @Const(n)
-    )
-        i = @index(Global)
-        keys[i] = colind[i] * n + rowind[i]
-    end
-
-    kernel! = make_keys!(backend)
+    kernel! = kernel_make_csc_keys!(backend)
     kernel!(keys, A.rowind, A.colind, n; ndrange = (nnz_count,))
 
-    # Sort - collect to CPU and use Base.sortperm since AcceleratedKernels
-    # doesn't work reliably on all backends (e.g., JLBackend)
-    keys_cpu = collect(keys)
-    perm_cpu = sortperm(keys_cpu)
-    # Adapt back to the original backend
-    perm = Adapt.adapt_structure(backend, perm_cpu)
+    # Sort - use AcceleratedKernels for GPU, CPU fallback for JLBackend
+    if _is_jlbackend(backend)
+        # JLBackend doesn't support AcceleratedKernels - use CPU fallback
+        keys_cpu = collect(keys)
+        perm_cpu = sortperm(keys_cpu)
+        perm = Adapt.adapt_structure(backend, perm_cpu)
+    else
+        # Use AcceleratedKernels for GPU and standard CPU backends
+        perm = AcceleratedKernels.sortperm(keys)
+    end
 
     # Apply permutation to get sorted arrays
     rowind_sorted = A.rowind[perm]
@@ -105,22 +101,23 @@ function DeviceSparseMatrixCSC(A::DeviceSparseMatrixCOO{Tv,Ti}) where {Tv,Ti}
     fill!(colptr, zero(Ti))
 
     # Count entries per column
-    @kernel inbounds=true function count_per_col!(colptr, @Const(colind_sorted))
-        i = @index(Global)
-        col = colind_sorted[i]
-        @atomic colptr[col+1] += 1
-    end
-
-    kernel! = count_per_col!(backend)
+    kernel! = kernel_count_per_col!(backend)
     kernel!(colptr, colind_sorted; ndrange = (nnz_count,))
 
-    # Build cumulative sum on CPU (collect, compute, adapt back)
-    colptr_cpu = collect(colptr)
-    colptr_cpu[1] = 1
-    for i = 2:(n+1)
-        colptr_cpu[i] += colptr_cpu[i-1]
+    # Compute cumulative sum - use CPU fallback for JLBackend
+    if _is_jlbackend(backend) || backend isa KernelAbstractions.CPU
+        # For CPU-like backends, use CPU cumsum
+        colptr_cpu = collect(colptr)
+        colptr_cpu[1] = 1
+        for i = 2:(n + 1)
+            colptr_cpu[i] += colptr_cpu[i - 1]
+        end
+        colptr = Adapt.adapt_structure(backend, colptr_cpu)
+    else
+        # For GPU backends, use AcceleratedKernels scan
+        colptr[1] = 1
+        colptr[2:end] .= AcceleratedKernels.cumsum(colptr[2:end]) .+ 1
     end
-    colptr = Adapt.adapt_structure(backend, colptr_cpu)
 
     return DeviceSparseMatrixCSC(m, n, colptr, rowind_sorted, nzval_sorted)
 end
@@ -195,29 +192,22 @@ function DeviceSparseMatrixCSR(A::DeviceSparseMatrixCOO{Tv,Ti}) where {Tv,Ti}
     backend = get_backend(A.nzval)
 
     # Create keys for sorting: row first, then column
-    # We use m * colind + rowind to create a unique sortable key
     keys = similar(A.rowind, Ti, nnz_count)
 
     # Create keys on device
-    @kernel inbounds=true function make_keys!(
-        keys,
-        @Const(rowind),
-        @Const(colind),
-        @Const(m)
-    )
-        i = @index(Global)
-        keys[i] = rowind[i] * m + colind[i]
-    end
-
-    kernel! = make_keys!(backend)
+    kernel! = kernel_make_csr_keys!(backend)
     kernel!(keys, A.rowind, A.colind, m; ndrange = (nnz_count,))
 
-    # Sort - collect to CPU and use Base.sortperm since AcceleratedKernels
-    # doesn't work reliably on all backends (e.g., JLBackend)
-    keys_cpu = collect(keys)
-    perm_cpu = sortperm(keys_cpu)
-    # Adapt back to the original backend
-    perm = Adapt.adapt_structure(backend, perm_cpu)
+    # Sort - use AcceleratedKernels for GPU, CPU fallback for JLBackend
+    if _is_jlbackend(backend)
+        # JLBackend doesn't support AcceleratedKernels - use CPU fallback
+        keys_cpu = collect(keys)
+        perm_cpu = sortperm(keys_cpu)
+        perm = Adapt.adapt_structure(backend, perm_cpu)
+    else
+        # Use AcceleratedKernels for GPU and standard CPU backends
+        perm = AcceleratedKernels.sortperm(keys)
+    end
 
     # Apply permutation to get sorted arrays
     rowind_sorted = A.rowind[perm]
@@ -229,22 +219,23 @@ function DeviceSparseMatrixCSR(A::DeviceSparseMatrixCOO{Tv,Ti}) where {Tv,Ti}
     fill!(rowptr, zero(Ti))
 
     # Count entries per row
-    @kernel inbounds=true function count_per_row!(rowptr, @Const(rowind_sorted))
-        i = @index(Global)
-        row = rowind_sorted[i]
-        @atomic rowptr[row+1] += 1
-    end
-
-    kernel! = count_per_row!(backend)
+    kernel! = kernel_count_per_row!(backend)
     kernel!(rowptr, rowind_sorted; ndrange = (nnz_count,))
 
-    # Build cumulative sum on CPU (collect, compute, adapt back)
-    rowptr_cpu = collect(rowptr)
-    rowptr_cpu[1] = 1
-    for i = 2:(m+1)
-        rowptr_cpu[i] += rowptr_cpu[i-1]
+    # Compute cumulative sum - use CPU fallback for JLBackend
+    if _is_jlbackend(backend) || backend isa KernelAbstractions.CPU
+        # For CPU-like backends, use CPU cumsum
+        rowptr_cpu = collect(rowptr)
+        rowptr_cpu[1] = 1
+        for i = 2:(m + 1)
+            rowptr_cpu[i] += rowptr_cpu[i - 1]
+        end
+        rowptr = Adapt.adapt_structure(backend, rowptr_cpu)
+    else
+        # For GPU backends, use AcceleratedKernels scan
+        rowptr[1] = 1
+        rowptr[2:end] .= AcceleratedKernels.cumsum(rowptr[2:end]) .+ 1
     end
-    rowptr = Adapt.adapt_structure(backend, rowptr_cpu)
 
     return DeviceSparseMatrixCSR(m, n, rowptr, colind_sorted, nzval_sorted)
 end