Add SparseMatricesCSR.jl Ext

Project.toml

@@ -37,10 +37,12 @@ UnsafeAtomics = "013be700-e6cd-48c3-b4a1-df204f14c38f"
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"
 
 [extensions]
 AMDGPUChainRulesCoreExt = "ChainRulesCore"
 AMDGPUEnzymeCoreExt = "EnzymeCore"
+AMDGPUSparseMatricesCSRExt = "SparseMatricesCSR"
 
 [compat]
 AbstractFFTs = "1.0"
@@ -64,6 +66,7 @@ PrettyTables = "3"
 ROCmDeviceLibs_jll = "=5.6.1, =6.2.1"
 Random123 = "1.6"
 RandomNumbers = "1.5"
+SparseMatricesCSR = "0.6.9"
 SpecialFunctions = "2"
 StaticArraysCore = "1"
 UnsafeAtomics = "0.3"

ext/AMDGPUSparseMatricesCSRExt.jl

@@ -0,0 +1,39 @@
+module AMDGPUSparseMatricesCSRExt
+
+using AMDGPU
+import AMDGPU.rocSPARSE:
+    ROCSparseMatrixCSR, ROCSparseMatrixCSC, ROCSparseMatrixCOO, ROCSparseMatrixBSR,
+    SparseMatrixCSC
+using SparseMatricesCSR
+import SparseMatricesCSR: SparseMatrixCSR
+import Adapt
+
+# CPU → GPU
+AMDGPU.rocSPARSE.ROCSparseMatrixCSR{T}(Mat::SparseMatrixCSR) where {T} =
+    AMDGPU.rocSPARSE.ROCSparseMatrixCSR{T}(
+        AMDGPU.ROCVector{Cint}(Mat.rowptr),
+        AMDGPU.ROCVector{Cint}(Mat.colval),
+        AMDGPU.ROCVector{T}(Mat.nzval),
+        size(Mat),
+    )
+AMDGPU.rocSPARSE.ROCSparseMatrixCSR(Mat::SparseMatrixCSR{<:Any, T}) where {T} =
+    AMDGPU.rocSPARSE.ROCSparseMatrixCSR{T}(Mat)
+
+AMDGPU.rocSPARSE.ROCSparseMatrixCSC{T}(Mat::SparseMatrixCSR) where {T} = ROCSparseMatrixCSC(ROCSparseMatrixCSR{T}(Mat))
+AMDGPU.rocSPARSE.ROCSparseMatrixCOO{T}(Mat::SparseMatrixCSR) where {T} = ROCSparseMatrixCOO(ROCSparseMatrixCSR{T}(Mat))
+AMDGPU.rocSPARSE.ROCSparseMatrixBSR{T}(Mat::SparseMatrixCSR, blockdim) where {T} = ROCSparseMatrixBSR(ROCSparseMatrixCSR{T}(Mat), blockdim)
+
+# GPU → CPU
+SparseMatricesCSR.SparseMatrixCSR(A::ROCSparseMatrixCSR) = SparseMatrixCSR{1}(size(A)..., Array(A.rowPtr), Array(A.colVal), Array(A.nzVal))
+SparseMatricesCSR.SparseMatrixCSR(A::ROCSparseMatrixCOO) = SparseMatrixCSR(ROCSparseMatrixCSR(A))
+SparseMatricesCSR.SparseMatrixCSR(A::ROCSparseMatrixCSC) = SparseMatrixCSR(ROCSparseMatrixCSR(A))
+SparseMatricesCSR.SparseMatrixCSR(A::ROCSparseMatrixBSR) = SparseMatrixCSR(ROCSparseMatrixCSR(A))
+
+# Adapt
+Adapt.adapt_storage(::Type{AMDGPU.ROCArray}, xs::SparseMatrixCSR) =
+    AMDGPU.rocSPARSE.ROCSparseMatrixCSR(xs)
+Adapt.adapt_storage(::Type{AMDGPU.ROCArray{T}}, xs::SparseMatrixCSR) where {T} =
+    AMDGPU.rocSPARSE.ROCSparseMatrixCSR{T}(xs)
+Adapt.adapt_storage(::Type{Array}, mat::ROCSparseMatrixCSR) = SparseMatrixCSR(mat)
+
+end

test/Project.toml

@@ -12,11 +12,12 @@ JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

test/hip_rocsparse/sparse_matrices_csr.jl

@@ -0,0 +1,31 @@
+using SparseMatricesCSR
+using SparseArrays
+using AMDGPU
+using AMDGPU.rocSPARSE
+using Test
+
+@assert AMDGPU.functional(:rocsparse)
+
+@testset "SparseMatricesCSRExt" begin
+
+    for (n, bd, p) in [(100, 5, 0.02), (5, 1, 0.8), (4, 2, 0.5)]
+        @testset "conversions between ROCSparseMatrices (n, bd, p) = ($n, $bd, $p)" begin
+            _A = sprand(n, n, p)
+            A = SparseMatrixCSR(_A)
+            blockdim = bd
+            for ROCSparseMatrixType1 in (ROCSparseMatrixCSC, ROCSparseMatrixCSR, ROCSparseMatrixCOO, ROCSparseMatrixBSR)
+                dA1 = ROCSparseMatrixType1 == ROCSparseMatrixBSR ? ROCSparseMatrixType1(A, blockdim) : ROCSparseMatrixType1(A)
+                @testset "conversion $ROCSparseMatrixType1 --> SparseMatrixCSR" begin
+                    @test SparseMatrixCSR(dA1) ≈ A
+                end
+                for ROCSparseMatrixType2 in (ROCSparseMatrixCSC, ROCSparseMatrixCSR, ROCSparseMatrixCOO, ROCSparseMatrixBSR)
+                    ROCSparseMatrixType1 == ROCSparseMatrixType2 && continue
+                    dA2 = ROCSparseMatrixType2 == ROCSparseMatrixBSR ? ROCSparseMatrixType2(dA1, blockdim) : ROCSparseMatrixType2(dA1)
+                    @testset "conversion $ROCSparseMatrixType1 --> $ROCSparseMatrixType2" begin
+                        @test collect(dA1) ≈ collect(dA2)
+                    end
+                end
+            end
+        end
+    end
+end