add back sparse matrix utils (#4815)

Author: simone-silvestri

ext/OceananigansAMDGPUExt.jl

@@ -6,6 +6,7 @@ using AMDGPU, AMDGPU.rocSPARSE, AMDGPU.rocFFT
 using Oceananigans.Utils: linear_expand, __linear_ndrange, MappedCompilerMetadata
 using KernelAbstractions: __dynamic_checkbounds, __iterspace
 using KernelAbstractions
+using SparseArrays
 
 import Oceananigans.Architectures as AC
 import Oceananigans.BoundaryConditions as BC
@@ -52,6 +53,16 @@ AC.on_architecture(::ROCGPU, a::StepRangeLen) = a
 AC.on_architecture(arch::Distributed, a::ROCArray) = AC.on_architecture(AC.child_architecture(arch), a)
 AC.on_architecture(arch::Distributed, a::SubArray{<:Any, <:Any, <:ROCArray}) = AC.on_architecture(child_architecture(arch), a)
 
+@inline AC.sparse_matrix_constructors(::AC.GPU{ROCBackend}, A::SparseMatrixCSC) = (ROCArray(A.colptr), ROCArray(A.rowval), ROCArray(A.nzval),  (A.m, A.n))
+@inline AC.sparse_matrix_constructors(::AC.CPU, A::ROCSparseMatrixCSC) = (A.dims[1], A.dims[2], Int64.(Array(A.colPtr)), Int64.(Array(A.rowVal)), Array(A.nzVal))
+@inline AC.sparse_matrix_constructors(::AC.GPU{ROCBackend}, A::ROCSparseMatrixCSC) = (A.colPtr, A.rowVal, A.nzVal,  A.dims)
+
+@inline AC.sparse_matrix(::AC.GPU{ROCBackend}, constr::Tuple) = ROCSparseMatrixCSC(constr...)
+
+@inline AC.on_architecture(::AC.CPU, A::ROCSparseMatrixCSC)     = SparseMatrixCSC(AC.sparse_matrix_constructors(AC.CPU(), A)...)
+@inline AC.on_architecture(::AC.GPU{ROCBackend}, A::SparseMatrixCSC) = ROCSparseMatrixCSC(AC.sparse_matrix_constructors(AC.GPU(), A)...)
+@inline AC.on_architecture(::AC.GPU{ROCBackend}, A::ROCSparseMatrixCSC) = A
+
 function AC.unified_array(::ROCGPU, a::AbstractArray)
     error("unified_array is not implemented for ROCGPU.")
 end

ext/OceananigansCUDAExt.jl

@@ -6,6 +6,7 @@ using CUDA, CUDA.CUSPARSE, CUDA.CUFFT
 using Oceananigans.Utils: linear_expand, __linear_ndrange, MappedCompilerMetadata
 using KernelAbstractions: __dynamic_checkbounds, __iterspace
 using KernelAbstractions
+using SparseArrays
 
 import Oceananigans.Architectures as AC
 import Oceananigans.BoundaryConditions as BC
@@ -73,6 +74,16 @@ AC.on_architecture(::CUDAGPU, a::StepRangeLen) = a
 AC.on_architecture(arch::Distributed, a::CuArray) = AC.on_architecture(AC.child_architecture(arch), a)
 AC.on_architecture(arch::Distributed, a::SubArray{<:Any, <:Any, <:CuArray}) = AC.on_architecture(child_architecture(arch), a)
 
+@inline AC.sparse_matrix_constructors(::AC.GPU{CUDABackend}, A::SparseMatrixCSC) = (CuArray(A.colptr), CuArray(A.rowval), CuArray(A.nzval),  (A.m, A.n))
+@inline AC.sparse_matrix_constructors(::AC.CPU, A::CuSparseMatrixCSC) = (A.dims[1], A.dims[2], Int64.(Array(A.colPtr)), Int64.(Array(A.rowVal)), Array(A.nzVal))
+@inline AC.sparse_matrix_constructors(::AC.GPU{CUDABackend}, A::CuSparseMatrixCSC) = (A.colPtr, A.rowVal, A.nzVal,  A.dims)
+
+@inline AC.sparse_matrix(::AC.GPU{CUDABackend}, constr::Tuple) = CuSparseMatrixCSC(constr...)
+
+@inline AC.on_architecture(::AC.CPU, A::CuSparseMatrixCSC)              = SparseMatrixCSC(AC.sparse_matrix_constructors(AC.CPU(), A)...)
+@inline AC.on_architecture(::AC.GPU{CUDABackend}, A::SparseMatrixCSC)   = CuSparseMatrixCSC(AC.sparse_matrix_constructors(AC.GPU(), A)...)
+@inline AC.on_architecture(::AC.GPU{CUDABackend}, A::CuSparseMatrixCSC) = A
+
 # cu alters the type of `a`, so we convert it back to the correct type
 AC.unified_array(::CUDAGPU, a::AbstractArray) = map(eltype(a), cu(a; unified = true))
 

src/Architectures.jl

@@ -75,6 +75,12 @@ architecture(::SparseMatrixCSC) = CPU()
 architecture(::Type{T}) where {T<:AbstractArray} = architecture(Base.typename(T).wrapper)
 architecture(::Type{Array}) = CPU()
 
+# Utils for sparse matrix manipulation
+@inline sparse_matrix_constructors(::CPU, A::SparseMatrixCSC) = (A.m, A.n, A.colptr, A.rowval, A.nzval)
+@inline sparse_matrix_constructors(::CPU, m::Number, n::Number, constr::Tuple) = (m, n, constr...)
+@inline sparse_matrix_constructors(::GPU, m::Number, n::Number, constr::Tuple) = (constr..., (m, n))
+@inline sparse_matrix(::CPU, constr::Tuple) = SparseMatrixCSC(constr...)
+
 """
     child_architecture(arch)
 
@@ -97,6 +103,7 @@ on_architecture(::CPU, a::Array) = a
 on_architecture(::CPU, a::BitArray) = a
 on_architecture(::CPU, a::SubArray{<:Any, <:Any, <:Array}) = a
 on_architecture(::CPU, a::StepRangeLen) = a
+on_architecture(::CPU, A::SparseMatrixCSC) = A
 
 on_architecture(arch::AbstractSerialArchitecture, a::OffsetArray) =
     OffsetArray(on_architecture(arch, a.parent), a.offsets...)