Implement display and other format conversions (#26)

* Implement display and other format conversions

* Fix benchmarks

* Fix benchmarks typo

* Comment out benchmark conversions for JLArray

* Add JLArrays support and enable conversion benchmarks

- Introduced JLArrays as a weak dependency in Project.toml.
- Enabled benchmark conversions for JLArray in runbenchmarks.jl.
- Added helper functions for sorting and cumulative sum using AcceleratedKernels in helpers.jl.
- Updated conversion tests to include JLArray support in shared test conversions.
- Created DeviceSparseArraysJLArraysExt.jl for JLArray specific operations.

* Fix Metal errors

Author: albertomercurio

Project.toml

@@ -12,15 +12,18 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [weakdeps]
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 
 [extensions]
+DeviceSparseArraysJLArraysExt = "JLArrays"
 DeviceSparseArraysReactantExt = "Reactant"
 
 [compat]
 AcceleratedKernels = "0.4"
 Adapt = "4"
 ArrayInterface = "7"
+JLArrays = "0.3"
 KernelAbstractions = "0.9"
 LinearAlgebra = "1"
 Reactant = "0.2.164"

README.md

@@ -56,6 +56,8 @@ mul!(c_result, A_device, b)
 ### GPU Backend Usage
 
 ```julia
+using Adapt
+
 # For CUDA backend
 using CUDA
 A_cuda = adapt(CuArray, A_device)

benchmarks/benchmark_utils.jl

@@ -32,8 +32,11 @@ _synchronize_backend(arr) = nothing  # Fallback: no-op for arrays without Kernel
 
 Synchronize KernelAbstractions backend for DeviceSparseArray types.
 """
-function _synchronize_backend(arr::AbstractDeviceSparseArray)
-    backend = KernelAbstractions.get_backend(arr)
+_synchronize_backend(arr::AbstractDeviceSparseArray) = _synchronize_backend(nonzeros(arr))
+
+function _synchronize_backend(x::AbstractArray)
+    backend = KernelAbstractions.get_backend(x)
     KernelAbstractions.synchronize(backend)
     return nothing
 end
+_synchronize_backend(x::JLArray) = nothing  # No-op for Julia Arrays

benchmarks/matrix_benchmarks.jl

@@ -242,34 +242,28 @@ Benchmark Kronecker product (kron) for CSC, CSR, and COO formats.
 - `N`: Size of the matrices (default: 100)
 - `T`: Element type (default: Float64)
 """
-function benchmark_kron!(
-    SUITE,
-    array_constructor,
-    array_type_name;
-    N = 100,
-    T = Float64,
-)
+function benchmark_kron!(SUITE, array_constructor, array_type_name; N = 100, T = Float64)
     # Create sparse matrices with 1% density (smaller matrices since kron grows quadratically)
     sm_a_std = sprand(T, N, N, 0.01)
     sm_b_std = sprand(T, N, N, 0.01)
 
     # Convert to different formats
     sm_a_csc = DeviceSparseMatrixCSC(sm_a_std)
     sm_b_csc = DeviceSparseMatrixCSC(sm_b_std)
-    
+
     sm_a_csr = DeviceSparseMatrixCSR(sm_a_std)
     sm_b_csr = DeviceSparseMatrixCSR(sm_b_std)
-    
+
     sm_a_coo = DeviceSparseMatrixCOO(sm_a_std)
     sm_b_coo = DeviceSparseMatrixCOO(sm_b_std)
 
     # Adapt to device
     dsm_a_csc = adapt(array_constructor, sm_a_csc)
     dsm_b_csc = adapt(array_constructor, sm_b_csc)
-    
+
     dsm_a_csr = adapt(array_constructor, sm_a_csr)
     dsm_b_csr = adapt(array_constructor, sm_b_csr)
-    
+
     dsm_a_coo = adapt(array_constructor, sm_a_coo)
     dsm_b_coo = adapt(array_constructor, sm_b_coo)
 
@@ -291,4 +285,3 @@ function benchmark_kron!(
 
     return nothing
 end
-

ext/DeviceSparseArraysJLArraysExt.jl

@@ -0,0 +1,9 @@
+module DeviceSparseArraysJLArraysExt
+
+using JLArrays: JLArray
+import DeviceSparseArrays
+
+DeviceSparseArrays._sortperm_AK(x::JLArray) = JLArray(sortperm(collect(x)))
+DeviceSparseArrays._cumsum_AK(x::JLArray) = JLArray(cumsum(collect(x)))
+
+end

src/DeviceSparseArrays.jl

@@ -5,6 +5,7 @@ import LinearAlgebra: wrap, copymutable_oftype, __normalize!, kron
 using SparseArrays
 import SparseArrays: SparseVector, SparseMatrixCSC
 import SparseArrays: getcolptr, getrowval, getnzval, nonzeroinds
+import SparseArrays: _show_with_braille_patterns
 
 import ArrayInterface: allowed_getindex, allowed_setindex!
 

src/conversions/conversions.jl

@@ -1,5 +1,126 @@
 # Conversions between CSC, CSR, and COO sparse matrix formats
-# All conversions operate on-device
+
+# ============================================================================
+# SparseMatrixCSC ↔ DeviceSparseMatrix (CSC-CSR-COO) Conversions
+# ============================================================================
+
+DeviceSparseMatrixCSC(A::SparseMatrixCSC) =
+    DeviceSparseMatrixCSC(size(A, 1), size(A, 2), A.colptr, A.rowval, A.nzval)
+
+SparseMatrixCSC(A::DeviceSparseMatrixCSC) = SparseMatrixCSC(
+    size(A, 1),
+    size(A, 2),
+    collect(A.colptr),
+    collect(A.rowval),
+    collect(A.nzval),
+)
+function SparseMatrixCSC(A::Transpose{Tv,<:DeviceSparseMatrixCSC}) where {Tv}
+    SparseMatrixCSC(DeviceSparseMatrixCSR(A))
+end
+function SparseMatrixCSC(A::Adjoint{Tv,<:DeviceSparseMatrixCSC}) where {Tv}
+    SparseMatrixCSC(DeviceSparseMatrixCSR(A))
+end
+
+function DeviceSparseMatrixCSR(A::SparseMatrixCSC)
+    # TODO: Implement a direct CSC to CSR conversion without going through transposition
+    At = transpose(A)
+    At_sparse = transpose(SparseMatrixCSC(At))
+    return DeviceSparseMatrixCSR(At_sparse)
+end
+
+function SparseMatrixCSC(A::DeviceSparseMatrixCSR)
+    # Convert CSR to CSC by creating transposed CSC and then transposing back
+    At_csc =
+        SparseMatrixCSC(A.n, A.m, collect(A.rowptr), collect(A.colval), collect(A.nzval))
+    return SparseMatrixCSC(transpose(At_csc))
+end
+function SparseMatrixCSC(A::Transpose{Tv,<:DeviceSparseMatrixCSR}) where {Tv}
+    At = A.parent
+    SparseMatrixCSC(At.n, At.m, collect(At.rowptr), collect(At.colval), collect(At.nzval))
+end
+function SparseMatrixCSC(A::Adjoint{Tv,<:DeviceSparseMatrixCSR}) where {Tv}
+    At = A.parent
+    SparseMatrixCSC(
+        size(A, 1),
+        size(A, 2),
+        collect(At.rowptr),
+        collect(At.colval),
+        collect(conj.(At.nzval)),
+    )
+end
+
+function DeviceSparseMatrixCOO(A::SparseMatrixCSC)
+    m, n = size(A)
+    rows, cols, vals = findnz(A)
+    return DeviceSparseMatrixCOO(m, n, rows, cols, vals)
+end
+
+function SparseMatrixCSC(A::DeviceSparseMatrixCOO)
+    m, n = size(A)
+    rowind = collect(A.rowind)
+    colind = collect(A.colind)
+    nzval = collect(A.nzval)
+
+    return sparse(rowind, colind, nzval, m, n)
+end
+SparseMatrixCSC(A::Transpose{Tv,<:DeviceSparseMatrixCOO}) where {Tv} = SparseMatrixCSC(
+    size(A, 1),
+    size(A, 2),
+    collect(A.parent.colind),
+    collect(A.parent.rowind),
+    collect(A.parent.nzval),
+)
+SparseMatrixCSC(A::Adjoint{Tv,<:DeviceSparseMatrixCOO}) where {Tv} = SparseMatrixCSC(
+    size(A, 1),
+    size(A, 2),
+    collect(A.parent.colind),
+    collect(A.parent.rowind),
+    collect(conj.(A.parent.nzval)),
+)
+
+# ============================================================================
+# CSC ↔ CSR Conversions
+# ============================================================================
+
+DeviceSparseMatrixCSC(A::DeviceSparseMatrixCSR) =
+    DeviceSparseMatrixCSC(DeviceSparseMatrixCOO(A))
+DeviceSparseMatrixCSC(A::Transpose{Tv,<:DeviceSparseMatrixCSR}) where {Tv} =
+    DeviceSparseMatrixCSC(
+        size(A, 1),
+        size(A, 2),
+        A.parent.rowptr,
+        A.parent.colval,
+        A.parent.nzval,
+    )
+DeviceSparseMatrixCSC(A::Adjoint{Tv,<:DeviceSparseMatrixCSR}) where {Tv} =
+    DeviceSparseMatrixCSC(
+        size(A, 1),
+        size(A, 2),
+        A.parent.rowptr,
+        A.parent.colval,
+        conj.(A.parent.nzval),
+    )
+
+DeviceSparseMatrixCSR(A::DeviceSparseMatrixCSC) =
+    DeviceSparseMatrixCSR(DeviceSparseMatrixCOO(A))
+function DeviceSparseMatrixCSR(
+    A::Transpose{Tv,<:Union{<:SparseMatrixCSC,<:DeviceSparseMatrixCSC}},
+) where {Tv}
+    At = A.parent
+    DeviceSparseMatrixCSR(size(A, 1), size(A, 2), At.colptr, rowvals(At), nonzeros(At))
+end
+function DeviceSparseMatrixCSR(
+    A::Adjoint{Tv,<:Union{<:SparseMatrixCSC,<:DeviceSparseMatrixCSC}},
+) where {Tv}
+    At = A.parent
+    DeviceSparseMatrixCSR(
+        size(A, 1),
+        size(A, 2),
+        At.colptr,
+        rowvals(At),
+        conj.(nonzeros(At)),
+    )
+end
 
 # ============================================================================
 # CSC ↔ COO Conversions
@@ -36,7 +157,8 @@ function DeviceSparseMatrixCSC(A::DeviceSparseMatrixCOO{Tv,Ti}) where {Tv,Ti}
     kernel! = kernel_make_csc_keys!(backend)
     kernel!(keys, A.rowind, A.colind, m; ndrange = (nnz_count,))
 
-    perm = AcceleratedKernels.sortperm(keys)
+    # Sort - use AcceleratedKernels
+    perm = _sortperm_AK(keys)
 
     # Apply permutation to get sorted arrays
     rowind_sorted = A.rowind[perm]
@@ -53,7 +175,7 @@ function DeviceSparseMatrixCSC(A::DeviceSparseMatrixCOO{Tv,Ti}) where {Tv,Ti}
 
     # Compute cumulative sum
     allowed_setindex!(colptr, 1, 1) # TODO: Is there a better way to do this?
-    colptr[2:end] .= AcceleratedKernels.cumsum(colptr[2:end]) .+ 1
+    colptr[2:end] .= _cumsum_AK(colptr[2:end]) .+ 1
 
     return DeviceSparseMatrixCSC(m, n, colptr, rowind_sorted, nzval_sorted)
 end
@@ -94,7 +216,7 @@ function DeviceSparseMatrixCSR(A::DeviceSparseMatrixCOO{Tv,Ti}) where {Tv,Ti}
     kernel!(keys, A.rowind, A.colind, n; ndrange = (nnz_count,))
 
     # Sort - use AcceleratedKernels
-    perm = AcceleratedKernels.sortperm(keys)
+    perm = _sortperm_AK(keys)
 
     # Apply permutation to get sorted arrays
     rowind_sorted = A.rowind[perm]
@@ -111,7 +233,7 @@ function DeviceSparseMatrixCSR(A::DeviceSparseMatrixCOO{Tv,Ti}) where {Tv,Ti}
 
     # Compute cumulative sum
     allowed_setindex!(rowptr, 1, 1) # TODO: Is there a better way to do this?
-    rowptr[2:end] .= AcceleratedKernels.cumsum(rowptr[2:end]) .+ 1
+    rowptr[2:end] .= _cumsum_AK(rowptr[2:end]) .+ 1
 
     return DeviceSparseMatrixCSR(m, n, rowptr, colind_sorted, nzval_sorted)
 end

src/core.jl

@@ -15,6 +15,12 @@ const AbstractDeviceSparseMatrix{Tv,Ti} = AbstractDeviceSparseArray{Tv,Ti,2}
 const AbstractDeviceSparseVecOrMat{Tv,Ti} =
     Union{AbstractDeviceSparseVector{Tv,Ti},AbstractDeviceSparseMatrix{Tv,Ti}}
 
+const AbstractDeviceSparseMatrixInclAdjointAndTranspose = Union{
+    AbstractDeviceSparseMatrix,
+    Adjoint{<:Any,<:AbstractDeviceSparseMatrix},
+    Transpose{<:Any,<:AbstractDeviceSparseMatrix},
+}
+
 Base.sum(A::AbstractDeviceSparseArray) = sum(nonzeros(A))
 
 function LinearAlgebra.rmul!(A::AbstractDeviceSparseArray, x::Number)
@@ -43,43 +49,17 @@ end
 
 KernelAbstractions.get_backend(A::AbstractDeviceSparseArray) = get_backend(nonzeros(A))
 
-trans_adj_wrappers(fmt) = (
-    (T -> :($fmt{$T}), false, false, identity, T -> :($T)),
-    (T -> :(Transpose{$T,<:$fmt{$T}}), true, false, A -> :(parent($A)), T -> :($T<:Real)),
-    (
-        T -> :(Transpose{$T,<:$fmt{$T}}),
-        true,
-        false,
-        A -> :(parent($A)),
-        T -> :($T<:Complex),
-    ),
-    (T -> :(Adjoint{$T,<:$fmt{$T}}), true, true, A -> :(parent($A)), T -> :($T)),
-)
+# called by `show(io, MIME("text/plain"), ::AbstractDeviceSparseMatrixInclAdjointAndTranspose)`
+function Base.print_array(io::IO, A::AbstractDeviceSparseMatrixInclAdjointAndTranspose)
+    S = SparseMatrixCSC(A)
+    if max(size(S)...) < 16
+        Base.print_matrix(io, S)
+    else
+        _show_with_braille_patterns(io, S)
+    end
+end
 
 # Generic addition between AbstractDeviceSparseMatrix and DenseMatrix
-"""
-    +(A::AbstractDeviceSparseMatrix, B::DenseMatrix)
-
-Add a sparse matrix `A` to a dense matrix `B`, returning a dense matrix.
-All backends must be compatible.
-
-# Examples
-```jldoctest
-julia> using DeviceSparseArrays, SparseArrays
-
-julia> A = DeviceSparseMatrixCSC(sparse([1, 2], [1, 2], [1.0, 2.0], 3, 3));
-
-julia> B = ones(3, 3);
-
-julia> C = A + B;
-
-julia> collect(C)
-3×3 Matrix{Float64}:
- 2.0  1.0  1.0
- 1.0  3.0  1.0
- 1.0  1.0  1.0
-```
-"""
 function Base.:+(A::AbstractDeviceSparseMatrix, B::DenseMatrix)
     size(A) == size(B) || throw(
         DimensionMismatch(
@@ -101,10 +81,18 @@ function Base.:+(A::AbstractDeviceSparseMatrix, B::DenseMatrix)
     return C
 end
 
-"""
-    +(B::DenseMatrix, A::AbstractDeviceSparseMatrix)
-
-Add a dense matrix `B` to a sparse matrix `A`, returning a dense matrix.
-This is the commutative version of `A + B`.
-"""
 Base.:+(B::DenseMatrix, A::AbstractDeviceSparseMatrix) = A + B
+
+# Keep this at the end of the file
+trans_adj_wrappers(fmt) = (
+    (T -> :($fmt{$T}), false, false, identity, T -> :($T)),
+    (T -> :(Transpose{$T,<:$fmt{$T}}), true, false, A -> :(parent($A)), T -> :($T<:Real)),
+    (
+        T -> :(Transpose{$T,<:$fmt{$T}}),
+        true,
+        false,
+        A -> :(parent($A)),
+        T -> :($T<:Complex),
+    ),
+    (T -> :(Adjoint{$T,<:$fmt{$T}}), true, true, A -> :(parent($A)), T -> :($T)),
+)

src/helpers.jl

@@ -6,3 +6,6 @@ _check_type(::Type{T}, v::AbstractArray{T}) where {T} = true
 _check_type(::Type{T}, v::AbstractArray) where {T} = false
 
 _get_eltype(::AbstractArray{T}) where {T} = T
+
+_sortperm_AK(x) = AcceleratedKernels.sortperm(x)
+_cumsum_AK(x) = AcceleratedKernels.cumsum(x)

src/matrix_coo/matrix_coo.jl

@@ -70,7 +70,13 @@ function DeviceSparseMatrixCOO(
 } where {Ti<:Integer,Tv}
     Ti2 = _get_eltype(rowind)
     Tv2 = _get_eltype(nzval)
-    DeviceSparseMatrixCOO{Tv2,Ti2,RowIndT,ColIndT,NzValT}(m, n, rowind, colind, nzval)
+    DeviceSparseMatrixCOO{Tv2,Ti2,RowIndT,ColIndT,NzValT}(
+        m,
+        n,
+        copy(rowind),
+        copy(colind),
+        copy(nzval),
+    )
 end
 
 # Conversion from SparseMatrixCSC to COO
@@ -95,16 +101,6 @@ function DeviceSparseMatrixCOO(A::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
     return DeviceSparseMatrixCOO(m, n, rowind, colind, nzval)
 end
 
-# Conversion from COO to SparseMatrixCSC
-function SparseMatrixCSC(A::DeviceSparseMatrixCOO)
-    m, n = size(A)
-    rowind = collect(A.rowind)
-    colind = collect(A.colind)
-    nzval = collect(A.nzval)
-
-    return sparse(rowind, colind, nzval, m, n)
-end
-
 Adapt.adapt_structure(to, A::DeviceSparseMatrixCOO) = DeviceSparseMatrixCOO(
     A.m,
     A.n,