Skip to content

Add SparseMatricesCSR.jl Ext #906

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
luraess merged 4 commits into from
May 30, 2026
+103 −1
Merged

Add SparseMatricesCSR.jl Ext #906

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d0ebffe
init
Abdelrahman912 May 9, 2026
86a0763
Merge remote-tracking branch 'upstream/main' into csr-ext
Abdelrahman912 May 28, 2026
422c90f
restrict dispatch
Abdelrahman912 May 28, 2026
cd19e25
add conversions
Abdelrahman912 May 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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"
Expand Down
52 changes: 52 additions & 0 deletions ext/AMDGPUSparseMatricesCSRExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
module AMDGPUSparseMatricesCSRExt

using AMDGPU
import AMDGPU.rocSPARSE:
ROCSparseMatrixCSR, ROCSparseMatrixCSC, ROCSparseMatrixCOO, ROCSparseMatrixBSR,
SparseMatrixCSC
using SparseMatricesCSR
import SparseMatricesCSR: SparseMatrixCSR
import Adapt

# CPU → GPU
# 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{1, T}) where {T} =
AMDGPU.rocSPARSE.ROCSparseMatrixCSR{T}(Mat)

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))
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
5 changes: 5 additions & 0 deletions src/sparse/conversions.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down
3 changes: 2 additions & 1 deletion test/Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
41 changes: 41 additions & 0 deletions test/hip_rocsparse/sparse_matrices_csr.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
using SparseMatricesCSR
using SparseArrays
using AMDGPU
using AMDGPU.rocSPARSE
using Adapt
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
Comment on lines +12 to +30
Copy link
Copy Markdown
Contributor Author

@Abdelrahman912 Abdelrahman912 May 29, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more. 

Error During Test at /var/lib/buildkite-agent/builds/amdgpu1-luraess-com/julialang/amdgpu-dot-jl/test/hip_rocsparse/sparse_matrices_csr.jl:12
  Got exception outside of a @test
  MethodError: no method matching (AMDGPU.rocSPARSE.ROCSparseMatrixBSR{Float64})(::AMDGPU.rocSPARSE.ROCSparseMatrixCSC{Float64, Int32}, ::Int64)

  Closest candidates are:
    (AMDGPU.rocSPARSE.ROCSparseMatrixBSR{Float64})(::AMDGPU.rocSPARSE.ROCSparseMatrixCSR{Float64}, ::Integer; dir, inda, indc)
     @ AMDGPU /var/lib/buildkite-agent/builds/amdgpu1-luraess-com/julialang/amdgpu-dot-jl/src/sparse/conversions.jl:224
    (AMDGPU.rocSPARSE.ROCSparseMatrixBSR{T})(::SparseArrays.SparseMatrixCSC, ::Any) where T
     @ AMDGPU /var/lib/buildkite-agent/builds/amdgpu1-luraess-com/julialang/amdgpu-dot-jl/src/sparse/array.jl:410
    (AMDGPU.rocSPARSE.ROCSparseMatrixBSR{T})(::SparseMatricesCSR.SparseMatrixCSR{1}, ::Any) where T
     @ AMDGPUSparseMatricesCSRExt /var/lib/buildkite-agent/builds/amdgpu1-luraess-com/julialang/amdgpu-dot-jl/ext/AMDGPUSparseMatricesCSRExt.jl:26
    ...

  Stacktrace:
    [1] AMDGPU.rocSPARSE.ROCSparseMatrixBSR(x::AMDGPU.rocSPARSE.ROCSparseMatrixCSC{Float64, Int32}, blockdim::Int64)
      @ AMDGPU.rocSPARSE /var/lib/buildkite-agent/builds/amdgpu1-luraess-com/julialang/amdgpu-dot-jl/src/sparse/array.jl:417

Probably this is a side tangent to this PR, but it's just so happen that it's the main culprit of the CI fail, due to this test nested loop which essentially converts every sparse matrix to every other sparse matrix. That said,
@luraess is there a reason why there is no conversion between CSC & BSR ? Like this one:
https://github.com/JuliaGPU/CUDA.jl/blob/54c758682d909f32d656fcbe8c43c20062850927/lib/cusparse/src/conversions.jl#L710-L711

If no need for them, I can just update the test, otherwise I can add those conversions, which hopefully will make CI happy.

Copy link
Copy Markdown
Member

@luraess luraess May 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

otherwise I can add those conversions

That would be the best way if you're willing to do so. Thanks!

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