From d0ebffec8778531afd338cce805d10438ece2234 Mon Sep 17 00:00:00 2001 From: abdelrahman912 Date: Sat, 9 May 2026 03:04:40 +0200 Subject: [PATCH 1/3] init --- Project.toml | 3 ++ ext/AMDGPUSparseMatricesCSRExt.jl | 39 +++++++++++++++++++++++ test/Project.toml | 3 +- test/hip_rocsparse/sparse_matrices_csr.jl | 31 ++++++++++++++++++ 4 files changed, 75 insertions(+), 1 deletion(-) create mode 100644 ext/AMDGPUSparseMatricesCSRExt.jl create mode 100644 test/hip_rocsparse/sparse_matrices_csr.jl 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 From 422c90f3a115b0bb8d15e867880ec4ddcb787716 Mon Sep 17 00:00:00 2001 From: abdelrahman912 Date: Thu, 28 May 2026 23:16:07 +0200 Subject: [PATCH 2/3] restrict dispatch --- ext/AMDGPUSparseMatricesCSRExt.jl | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/ext/AMDGPUSparseMatricesCSRExt.jl b/ext/AMDGPUSparseMatricesCSRExt.jl index 0cd0e895a..a72e65716 100644 --- a/ext/AMDGPUSparseMatricesCSRExt.jl +++ b/ext/AMDGPUSparseMatricesCSRExt.jl @@ -9,19 +9,32 @@ import SparseMatricesCSR: SparseMatrixCSR import Adapt # CPU → GPU -AMDGPU.rocSPARSE.ROCSparseMatrixCSR{T}(Mat::SparseMatrixCSR) where {T} = +# NOTE: `SparseMatricesCSR.SparseMatrixCSR{Bi,Tv,Ti}` has a type parameter `Bi` for the index base (0 or 1), +# but `ROCSPARSE` expects 1-based pointer. Therefore, we restrict dispatch to `SparseMatrixCSR{1}`. +AMDGPU.rocSPARSE.ROCSparseMatrixCSR{T}(Mat::SparseMatrixCSR{1}) 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(Mat::SparseMatrixCSR{1, 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) +AMDGPU.rocSPARSE.ROCSparseMatrixCSC{T}(Mat::SparseMatrixCSR{1}) where {T} = ROCSparseMatrixCSC(ROCSparseMatrixCSR{T}(Mat)) +AMDGPU.rocSPARSE.ROCSparseMatrixCOO{T}(Mat::SparseMatrixCSR{1}) where {T} = ROCSparseMatrixCOO(ROCSparseMatrixCSR{T}(Mat)) +AMDGPU.rocSPARSE.ROCSparseMatrixBSR{T}(Mat::SparseMatrixCSR{1}, blockdim) where {T} = ROCSparseMatrixBSR(ROCSparseMatrixCSR{T}(Mat), blockdim) + +# Error for unsupported index base. +@noinline _unsupported_index_base(::SparseMatrixCSR{Bi}) where {Bi} = throw(ArgumentError( + "`ROCSPARSE` expects 1-based index `SparseMatrixCSR{1}` to be converted to rocSPARSE types, " * + "but got a $Bi-based `SparseMatrixCSR{$Bi}`.")) + +AMDGPU.rocSPARSE.ROCSparseMatrixCSR{T}(Mat::SparseMatrixCSR) where {T} = _unsupported_index_base(Mat) +AMDGPU.rocSPARSE.ROCSparseMatrixCSR(Mat::SparseMatrixCSR) = _unsupported_index_base(Mat) +AMDGPU.rocSPARSE.ROCSparseMatrixCSC{T}(Mat::SparseMatrixCSR) where {T} = _unsupported_index_base(Mat) +AMDGPU.rocSPARSE.ROCSparseMatrixCOO{T}(Mat::SparseMatrixCSR) where {T} = _unsupported_index_base(Mat) +AMDGPU.rocSPARSE.ROCSparseMatrixBSR{T}(Mat::SparseMatrixCSR, ::Any) where {T} = _unsupported_index_base(Mat) # GPU → CPU SparseMatricesCSR.SparseMatrixCSR(A::ROCSparseMatrixCSR) = SparseMatrixCSR{1}(size(A)..., Array(A.rowPtr), Array(A.colVal), Array(A.nzVal)) From cd19e255507111f227defd9e41cfbb903a1a3291 Mon Sep 17 00:00:00 2001 From: abdelrahman912 Date: Sat, 30 May 2026 15:48:11 +0200 Subject: [PATCH 3/3] add conversions --- src/sparse/conversions.jl | 5 +++++ test/hip_rocsparse/sparse_matrices_csr.jl | 10 ++++++++++ 2 files changed, 15 insertions(+) diff --git a/src/sparse/conversions.jl b/src/sparse/conversions.jl index 6053122ed..569e95255 100644 --- a/src/sparse/conversions.jl +++ b/src/sparse/conversions.jl @@ -338,6 +338,11 @@ ROCSparseMatrixCOO(csc::ROCSparseMatrixCSC) = ROCSparseMatrixCOO(ROCSparseMatrix ROCSparseMatrixBSR(coo::ROCSparseMatrixCOO, blockdim) = ROCSparseMatrixBSR(ROCSparseMatrixCSR(coo), blockdim) # no direct conversion ROCSparseMatrixCOO(bsr::ROCSparseMatrixBSR) = ROCSparseMatrixCOO(ROCSparseMatrixCSR(bsr)) # no direct conversion +### BSR to CSC and vice-versa + +ROCSparseMatrixBSR(csc::ROCSparseMatrixCSC, blockdim) = ROCSparseMatrixBSR(ROCSparseMatrixCSR(csc), blockdim) # no direct conversion +ROCSparseMatrixCSC(bsr::ROCSparseMatrixBSR) = ROCSparseMatrixCSC(ROCSparseMatrixCSR(bsr)) # no direct conversion + ## sparse to dense, and vice-versa for (cname,rname,elty) in ( diff --git a/test/hip_rocsparse/sparse_matrices_csr.jl b/test/hip_rocsparse/sparse_matrices_csr.jl index 991b9668e..9db31ee8d 100644 --- a/test/hip_rocsparse/sparse_matrices_csr.jl +++ b/test/hip_rocsparse/sparse_matrices_csr.jl @@ -2,6 +2,7 @@ using SparseMatricesCSR using SparseArrays using AMDGPU using AMDGPU.rocSPARSE +using Adapt using Test @assert AMDGPU.functional(:rocsparse) @@ -28,4 +29,13 @@ using Test end end end + + @testset "non-1-based SparseMatrixCSR is rejected" begin + A0 = SparseMatrixCSR{0}(3, 3, Cint[0, 1, 2, 3], Cint[0, 1, 2], [1.0, 2.0, 3.0]) + @test_throws ArgumentError ROCSparseMatrixCSR(A0) + @test_throws ArgumentError ROCSparseMatrixCSC(A0) + @test_throws ArgumentError ROCSparseMatrixCOO(A0) + @test_throws ArgumentError ROCSparseMatrixBSR(A0, 1) + @test_throws ArgumentError adapt(ROCArray, A0) + end end