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

Author: albertomercurio

Project.toml

@@ -1,12 +1,12 @@
 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"
@@ -20,7 +20,7 @@ DeviceSparseArraysJLArraysExt = "JLArrays"
 [compat]
 AcceleratedKernels = "0.4"
 Adapt = "4"
-ArrayInterface = "7"
+GPUArraysCore = "0.2.0"
 JLArrays = "0.3"
 KernelAbstractions = "0.9"
 LinearAlgebra = "1"

src/DeviceSparseArrays.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/conversions/conversions.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/matrix_coo/matrix_coo.jl

@@ -160,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_csc/matrix_csc.jl

@@ -149,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_csr/matrix_csr.jl

@@ -150,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

src/vector/vector.jl

@@ -135,7 +135,7 @@ function LinearAlgebra.dot(x::DeviceSparseVector, y::DenseVector)
     kernel = kernel_dot(backend)
     kernel(res, nzval, nzind, y; ndrange = (m,))
 
-    return allowed_getindex(res, 1)
+    return @allowscalar res[1]
 end
 LinearAlgebra.dot(x::DenseVector{T1}, y::DeviceSparseVector{Tv}) where {T1<:Real,Tv<:Real} =
     dot(y, x)