Force fields to have same eltype (#28)

* Force fields to have same eltype

* Relax types

* Update src/matrix_csc/matrix_csc.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Update src/matrix_csr/matrix_csr.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Copy arrays in COO format

* Fix errors in docstring

* Fix error on `tr` function

* Refactor to replace `allowed_getindex` with `@allowscalar` for improved scalar handling in matrix operations

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: albertomercurio

Project.toml

@@ -1,31 +1,28 @@
 name = "DeviceSparseArrays"
 uuid = "da3fe0eb-88a8-4d14-ae1a-857c283e9c70"
-version = "0.1.0"
 authors = ["Alberto Mercurio <alberto.mercurio96@gmail.com> and contributors"]
+version = "0.1.0"
 
 [deps]
 AcceleratedKernels = "6a4ca0a5-0e36-4168-a932-d9be78d558f1"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
-ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 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"
+GPUArraysCore = "0.2.0"
 JLArrays = "0.3"
 KernelAbstractions = "0.9"
 LinearAlgebra = "1"
-Reactant = "0.2.164"
 SparseArrays = "1"
 julia = "1.10"

ext/DeviceSparseArraysReactantExt.jl

@@ -7,7 +7,7 @@ import SparseArrays: SparseVector, SparseMatrixCSC
 import SparseArrays: getcolptr, getrowval, getnzval, nonzeroinds
 import SparseArrays: _show_with_braille_patterns
 
-import ArrayInterface: allowed_getindex, allowed_setindex!
+import GPUArraysCore: @allowscalar
 
 import KernelAbstractions
 import KernelAbstractions:

src/DeviceSparseArrays.jl

@@ -174,7 +174,7 @@ function DeviceSparseMatrixCSC(A::DeviceSparseMatrixCOO{Tv,Ti}) where {Tv,Ti}
     kernel!(colptr, colind_sorted; ndrange = (nnz_count,))
 
     # Compute cumulative sum
-    allowed_setindex!(colptr, 1, 1) # TODO: Is there a better way to do this?
+    @allowscalar colptr[1] = 1 # TODO: Is there a better way to do this?
     colptr[2:end] .= _cumsum_AK(colptr[2:end]) .+ 1
 
     return DeviceSparseMatrixCSC(m, n, colptr, rowind_sorted, nzval_sorted)
@@ -232,7 +232,7 @@ function DeviceSparseMatrixCSR(A::DeviceSparseMatrixCOO{Tv,Ti}) where {Tv,Ti}
     kernel!(rowptr, rowind_sorted; ndrange = (nnz_count,))
 
     # Compute cumulative sum
-    allowed_setindex!(rowptr, 1, 1) # TODO: Is there a better way to do this?
+    @allowscalar rowptr[1] = 1 # TODO: Is there a better way to do this?
     rowptr[2:end] .= _cumsum_AK(rowptr[2:end]) .+ 1
 
     return DeviceSparseMatrixCSR(m, n, rowptr, colind_sorted, nzval_sorted)

src/conversions/conversions.jl

@@ -1,11 +1,3 @@
-#=
-A method to check that an AbstractArray is of a given element type.
-This is needed because we can implement new methods for different arrays (e.g., Reactant.jl)
-=#
-_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
-
+# Helper functions to call AcceleratedKernels methods
 _sortperm_AK(x) = AcceleratedKernels.sortperm(x)
 _cumsum_AK(x) = AcceleratedKernels.cumsum(x)

src/helpers.jl

@@ -1,7 +1,7 @@
 # DeviceSparseMatrixCOO implementation
 
 """
-    DeviceSparseMatrixCOO{Tv,Ti,RowIndT<:AbstractVector,ColIndT<:AbstractVector,NzValT<:AbstractVector} <: AbstractDeviceSparseMatrix{Tv,Ti}
+    DeviceSparseMatrixCOO{Tv,Ti,RowIndT<:AbstractVector{Ti},ColIndT<:AbstractVector{Ti},NzValT<:AbstractVector{Tv}} <: AbstractDeviceSparseMatrix{Tv,Ti}
 
 Coordinate (COO) sparse matrix with generic storage vectors for row indices,
 column indices, and nonzero values. Buffer types (e.g. `Vector`, GPU array
@@ -16,28 +16,29 @@ types) enable dispatch on device characteristics.
 """
 struct DeviceSparseMatrixCOO{
     Tv,
-    Ti<:Integer,
-    RowIndT<:AbstractVector,
-    ColIndT<:AbstractVector,
-    NzValT<:AbstractVector,
+    Ti,
+    RowIndT<:AbstractVector{Ti},
+    ColIndT<:AbstractVector{Ti},
+    NzValT<:AbstractVector{Tv},
 } <: AbstractDeviceSparseMatrix{Tv,Ti}
     m::Int
     n::Int
     rowind::RowIndT
     colind::ColIndT
     nzval::NzValT
-    function DeviceSparseMatrixCOO{Tv,Ti,RowIndT,ColIndT,NzValT}(
+
+    function DeviceSparseMatrixCOO(
         m::Integer,
         n::Integer,
         rowind::RowIndT,
         colind::ColIndT,
         nzval::NzValT,
     ) where {
         Tv,
-        Ti<:Integer,
-        RowIndT<:AbstractVector,
-        ColIndT<:AbstractVector,
-        NzValT<:AbstractVector,
+        Ti,
+        RowIndT<:AbstractVector{Ti},
+        ColIndT<:AbstractVector{Ti},
+        NzValT<:AbstractVector{Tv},
     }
         get_backend(rowind) == get_backend(colind) == get_backend(nzval) ||
             throw(ArgumentError("All storage vectors must be on the same device/backend."))
@@ -46,39 +47,19 @@ struct DeviceSparseMatrixCOO{
         n >= 0 || throw(ArgumentError("n must be non-negative"))
         SparseArrays.sparse_check_Ti(m, n, Ti)
 
-        _check_type(Ti, rowind) || throw(ArgumentError("rowind must be of type $Ti"))
-        _check_type(Ti, colind) || throw(ArgumentError("colind must be of type $Ti"))
-        _check_type(Tv, nzval) || throw(ArgumentError("nzval must be of type $Tv"))
-
         length(rowind) == length(colind) == length(nzval) ||
             throw(ArgumentError("rowind, colind, and nzval must have same length"))
 
-        return new(Int(m), Int(n), rowind, colind, nzval)
+        return new{Tv,Ti,RowIndT,ColIndT,NzValT}(
+            Int(m),
+            Int(n),
+            copy(rowind),
+            copy(colind),
+            copy(nzval),
+        )
     end
 end
 
-function DeviceSparseMatrixCOO(
-    m::Integer,
-    n::Integer,
-    rowind::RowIndT,
-    colind::ColIndT,
-    nzval::NzValT,
-) where {
-    RowIndT<:AbstractVector{Ti},
-    ColIndT<:AbstractVector{Ti},
-    NzValT<:AbstractVector{Tv},
-} where {Ti<:Integer,Tv}
-    Ti2 = _get_eltype(rowind)
-    Tv2 = _get_eltype(nzval)
-    DeviceSparseMatrixCOO{Tv2,Ti2,RowIndT,ColIndT,NzValT}(
-        m,
-        n,
-        copy(rowind),
-        copy(colind),
-        copy(nzval),
-    )
-end
-
 # Conversion from SparseMatrixCSC to COO
 function DeviceSparseMatrixCOO(A::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
     m, n = size(A)
@@ -179,7 +160,7 @@ function LinearAlgebra.tr(A::DeviceSparseMatrixCOO)
     kernel = kernel_tr(backend)
     kernel(res, getrowind(A), getcolind(A), nonzeros(A); ndrange = (length(nonzeros(A)),))
 
-    return allowed_getindex(res, 1)
+    return @allowscalar res[1]
 end
 
 # Matrix-Vector and Matrix-Matrix multiplication

src/matrix_coo/matrix_coo.jl

@@ -1,7 +1,7 @@
 # DeviceSparseMatrixCSC implementation
 
 """
-    DeviceSparseMatrixCSC{Tv,Ti,ColPtrT<RowValT,NzValT} <: AbstractDeviceSparseMatrix{Tv,Ti}
+    DeviceSparseMatrixCSC{Tv,Ti,ColPtrT,RowValT,NzValT} <: AbstractDeviceSparseMatrix{Tv,Ti}
 
 Compressed Sparse Column (CSC) matrix with generic storage vectors for column
 pointer, row indices, and nonzero values. Buffer types (e.g. `Vector`, GPU array
@@ -16,28 +16,29 @@ types) enable dispatch on device characteristics.
 """
 struct DeviceSparseMatrixCSC{
     Tv,
-    Ti<:Integer,
-    ColPtrT<:AbstractVector,
-    RowValT<:AbstractVector,
-    NzValT<:AbstractVector,
+    Ti,
+    ColPtrT<:AbstractVector{Ti},
+    RowValT<:AbstractVector{Ti},
+    NzValT<:AbstractVector{Tv},
 } <: AbstractDeviceSparseMatrix{Tv,Ti}
     m::Int
     n::Int
     colptr::ColPtrT
     rowval::RowValT
     nzval::NzValT
-    function DeviceSparseMatrixCSC{Tv,Ti,ColPtrT,RowValT,NzValT}(
+
+    function DeviceSparseMatrixCSC(
         m::Integer,
         n::Integer,
         colptr::ColPtrT,
         rowval::RowValT,
         nzval::NzValT,
     ) where {
         Tv,
-        Ti<:Integer,
-        ColPtrT<:AbstractVector,
-        RowValT<:AbstractVector,
-        NzValT<:AbstractVector,
+        Ti,
+        ColPtrT<:AbstractVector{Ti},
+        RowValT<:AbstractVector{Ti},
+        NzValT<:AbstractVector{Tv},
     }
         get_backend(colptr) == get_backend(rowval) == get_backend(nzval) ||
             throw(ArgumentError("All storage vectors must be on the same device/backend."))
@@ -47,35 +48,20 @@ struct DeviceSparseMatrixCSC{
         SparseArrays.sparse_check_Ti(m, n, Ti)
         # SparseArrays.sparse_check(n, colptr, rowval, nzval) # TODO: this uses scalar indexing
 
-        _check_type(Ti, colptr) || throw(ArgumentError("colptr must be of type $Ti"))
-        _check_type(Ti, rowval) || throw(ArgumentError("rowval must be of type $Ti"))
-        _check_type(Tv, nzval) || throw(ArgumentError("nzval must be of type $Tv"))
-
         length(colptr) == n + 1 || throw(ArgumentError("colptr length must be n+1"))
         length(rowval) == length(nzval) ||
             throw(ArgumentError("rowval and nzval must have same length"))
 
-        return new(Int(m), Int(n), copy(colptr), copy(rowval), copy(nzval))
+        return new{Tv,Ti,ColPtrT,RowValT,NzValT}(
+            Int(m),
+            Int(n),
+            copy(colptr),
+            copy(rowval),
+            copy(nzval),
+        )
     end
 end
 
-function DeviceSparseMatrixCSC(
-    m::Integer,
-    n::Integer,
-    colptr::ColPtrT,
-    rowval::RowValT,
-    nzval::NzValT,
-) where {
-    ColPtrT<:AbstractVector{Ti},
-    RowValT<:AbstractVector{Ti},
-    NzValT<:AbstractVector{Tv},
-} where {Ti<:Integer,Tv}
-    Ti2 = _get_eltype(colptr)
-    Tv2 = _get_eltype(nzval)
-    DeviceSparseMatrixCSC{Tv2,Ti2,ColPtrT,RowValT,NzValT}(m, n, colptr, rowval, nzval)
-end
-
-
 Adapt.adapt_structure(to, A::DeviceSparseMatrixCSC) = DeviceSparseMatrixCSC(
     A.m,
     A.n,
@@ -163,7 +149,7 @@ function LinearAlgebra.tr(A::DeviceSparseMatrixCSC)
     kernel = kernel_tr(backend)
     kernel(res, getcolptr(A), getrowval(A), getnzval(A); ndrange = (n,))
 
-    return allowed_getindex(res, 1)
+    return @allowscalar res[1]
 end
 
 # Matrix-Vector and Matrix-Matrix multiplication

src/matrix_csc/matrix_csc.jl

@@ -1,7 +1,7 @@
 # DeviceSparseMatrixCSR implementation
 
 """
-    DeviceSparseMatrixCSR{Tv,Ti,RowPtrT<:ColValT,NzValT} <: AbstractDeviceSparseMatrix{Tv,Ti}
+    DeviceSparseMatrixCSR{Tv,Ti,RowPtrT,ColValT,NzValT} <: AbstractDeviceSparseMatrix{Tv,Ti}
 
 Compressed Sparse Row (CSR) matrix with generic storage vectors for row
 pointer, column indices, and nonzero values. Buffer types (e.g. `Vector`, GPU array
@@ -16,28 +16,29 @@ types) enable dispatch on device characteristics.
 """
 struct DeviceSparseMatrixCSR{
     Tv,
-    Ti<:Integer,
-    RowPtrT<:AbstractVector,
-    ColValT<:AbstractVector,
-    NzValT<:AbstractVector,
+    Ti,
+    RowPtrT<:AbstractVector{Ti},
+    ColValT<:AbstractVector{Ti},
+    NzValT<:AbstractVector{Tv},
 } <: AbstractDeviceSparseMatrix{Tv,Ti}
     m::Int
     n::Int
     rowptr::RowPtrT
     colval::ColValT
     nzval::NzValT
-    function DeviceSparseMatrixCSR{Tv,Ti,RowPtrT,ColValT,NzValT}(
+
+    function DeviceSparseMatrixCSR(
         m::Integer,
         n::Integer,
         rowptr::RowPtrT,
         colval::ColValT,
         nzval::NzValT,
     ) where {
         Tv,
-        Ti<:Integer,
-        RowPtrT<:AbstractVector,
-        ColValT<:AbstractVector,
-        NzValT<:AbstractVector,
+        Ti,
+        RowPtrT<:AbstractVector{Ti},
+        ColValT<:AbstractVector{Ti},
+        NzValT<:AbstractVector{Tv},
     }
         get_backend(rowptr) == get_backend(colval) == get_backend(nzval) ||
             throw(ArgumentError("All storage vectors must be on the same device/backend."))
@@ -47,34 +48,20 @@ struct DeviceSparseMatrixCSR{
         SparseArrays.sparse_check_Ti(m, n, Ti)
         # SparseArrays.sparse_check(m, rowptr, colval, nzval) # TODO: this uses scalar indexing
 
-        _check_type(Ti, rowptr) || throw(ArgumentError("rowptr must be of type $Ti"))
-        _check_type(Ti, colval) || throw(ArgumentError("colval must be of type $Ti"))
-        _check_type(Tv, nzval) || throw(ArgumentError("nzval must be of type $Tv"))
-
         length(rowptr) == m + 1 || throw(ArgumentError("rowptr length must be m+1"))
         length(colval) == length(nzval) ||
             throw(ArgumentError("colval and nzval must have same length"))
 
-        return new(Int(m), Int(n), copy(rowptr), copy(colval), copy(nzval))
+        return new{Tv,Ti,RowPtrT,ColValT,NzValT}(
+            Int(m),
+            Int(n),
+            copy(rowptr),
+            copy(colval),
+            copy(nzval),
+        )
     end
 end
 
-function DeviceSparseMatrixCSR(
-    m::Integer,
-    n::Integer,
-    rowptr::RowPtrT,
-    colval::ColValT,
-    nzval::NzValT,
-) where {
-    RowPtrT<:AbstractVector{Ti},
-    ColValT<:AbstractVector{Ti},
-    NzValT<:AbstractVector{Tv},
-} where {Ti<:Integer,Tv}
-    Ti2 = _get_eltype(rowptr)
-    Tv2 = _get_eltype(nzval)
-    DeviceSparseMatrixCSR{Tv2,Ti2,RowPtrT,ColValT,NzValT}(m, n, rowptr, colval, nzval)
-end
-
 Adapt.adapt_structure(to, A::DeviceSparseMatrixCSR) = DeviceSparseMatrixCSR(
     A.m,
     A.n,
@@ -163,7 +150,7 @@ function LinearAlgebra.tr(A::DeviceSparseMatrixCSR)
     kernel = kernel_tr(backend)
     kernel(res, getrowptr(A), getcolval(A), nonzeros(A); ndrange = (m,))
 
-    return allowed_getindex(res, 1)
+    return @allowscalar res[1]
 end
 
 # Matrix-Vector and Matrix-Matrix multiplication