diff --git a/Project.toml b/Project.toml index df5e1c31f..e94765dc7 100644 --- a/Project.toml +++ b/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" diff --git a/ext/AMDGPUSparseMatricesCSRExt.jl b/ext/AMDGPUSparseMatricesCSRExt.jl new file mode 100644 index 000000000..0cd0e895a --- /dev/null +++ b/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 diff --git a/test/Project.toml b/test/Project.toml index 6185fb8b9..7023e261f 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/hip_rocsparse/sparse_matrices_csr.jl b/test/hip_rocsparse/sparse_matrices_csr.jl new file mode 100644 index 000000000..991b9668e --- /dev/null +++ b/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