Add AMDGPU sparse matrix support extension

.buildkite/AMDGPU_Ext.yml

@@ -0,0 +1,25 @@
+steps:
+  - label: "AMDGPU Julia {{matrix.version}}"
+    matrix:
+      setup:
+        version:
+          - "1.10" # oldest
+          - "1"    # latest
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "{{matrix.version}}"
+      - JuliaCI/julia-test#v1:
+          test_args: "--quickfail"
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
+          dirs:
+            - src
+            - ext
+    agents:
+      queue: "juliagpu"
+      rocm: "*"
+      rocmgpu: "*"
+    env:
+      GROUP: "AMDGPU_Ext"
+      SECRET_CODECOV_TOKEN: "ZfhQu/IcRLqNyZ//ZNs5sjBPaV76IHfU5gui52Qn+Rp8tOurukqgScuyDt+3HQ4R0hJYBw1/Nqg6jmBsvWSc9NEUx8kGsUJFHfN3no0+b+PFxA8oJkWc9EpyIsjht5ZIjlsFWR3f0DpPqMEle/QyWOPcal63CChXR8oAoR+Fz1Bh8GkokLlnC8F9Ugp9xBlu401GCbyZhvLTZnNIgK5yy9q8HBJnBg1cPOhI81J6JvYpEmcIofEzFV/qkfpTUPclu43WNoFX2DZPzbxilf3fsAd5/+nRkRfkNML8KiN4mnmjHxPPbuY8F5zC/PS5ybXtDpfvaMQc01WApXCkZk0ZAQ==;U2FsdGVkX1+eDT7dqCME5+Ox5i8GvWRTQbwiP/VYjapThDbxXFDeSSIC6Opmon+M8go22Bun3bat6Fzie65ang=="
+    timeout_in_minutes: 60

.buildkite/pipeline.yml

@@ -20,6 +20,16 @@ steps:
             - "src/**"
             - "ext/QuantumToolboxCUDAExt.jl"
             - "test/runtests.jl"
-            - "test/ext-test/gpu/**"
+            - "test/ext-test/gpu/cuda/**"
             - "Project.toml"
           target: ".buildkite/CUDA_Ext.yml"
+      - staticfloat/forerunner: # AMDGPU.jl tests
+          watch:
+            - ".buildkite/pipeline.yml"
+            - ".buildkite/AMDGPU_Ext.yml"
+            - "src/**"
+            - "ext/QuantumToolboxAMDGPUExt.jl"
+            - "test/runtests.jl"
+            - "test/ext-test/gpu/amdgpu/**"
+            - "Project.toml"
+          target: ".buildkite/AMDGPU_Ext.yml"

Project.toml

@@ -30,19 +30,22 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 
 [weakdeps]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 
 [extensions]
+QuantumToolboxAMDGPUExt = "AMDGPU"
 QuantumToolboxCUDAExt = "CUDA"
 QuantumToolboxChainRulesCoreExt = "ChainRulesCore"
 QuantumToolboxGPUArraysExt = ["GPUArrays", "KernelAbstractions"]
 QuantumToolboxMakieExt = "Makie"
 
 [compat]
+AMDGPU = "2.1"
 ArrayInterface = "6, 7"
 CUDA = "5.0 - 5.8, 5.9.4 - 5"
 ChainRulesCore = "1"

ext/QuantumToolboxAMDGPUExt.jl

@@ -0,0 +1,108 @@
+module QuantumToolboxAMDGPUExt
+
+using QuantumToolbox
+using QuantumToolbox: makeVal, getVal
+import QuantumToolbox: _sparse_similar, _convert_eltype_wordsize
+import AMDGPU: roc, ROCArray, allowscalar
+import AMDGPU.rocSPARSE: ROCSparseVector, ROCSparseMatrixCSC, ROCSparseMatrixCSR, AbstractROCSparseArray
+import SparseArrays: SparseVector, SparseMatrixCSC, sparse
+import AMDGPU.Adapt: adapt
+
+allowscalar(false)
+
+@doc raw"""
+    ROCArray(A::QuantumObject)
+
+If `A.data` is a dense array, return a new [`QuantumObject`](@ref) where `A.data` is in the type of `AMDGPU.ROCArray` for gpu calculations.
+"""
+ROCArray(A::QuantumObject) = QuantumObject(ROCArray(A.data), A.type, A.dimensions)
+
+@doc raw"""
+    ROCArray{T}(A::QuantumObject)
+
+If `A.data` is a dense array, return a new [`QuantumObject`](@ref) where `A.data` is in the type of `AMDGPU.ROCArray` with element type `T` for gpu calculations.
+"""
+ROCArray{T}(A::QuantumObject) where {T} = QuantumObject(ROCArray{T}(A.data), A.type, A.dimensions)
+
+@doc raw"""
+    ROCSparseVector(A::QuantumObject)
+
+If `A.data` is a sparse vector, return a new [`QuantumObject`](@ref) where `A.data` is in the type of `AMDGPU.rocSPARSE.ROCSparseVector` for gpu calculations.
+"""
+ROCSparseVector(A::QuantumObject) = QuantumObject(ROCSparseVector(A.data), A.type, A.dimensions)
+
+@doc raw"""
+    ROCSparseVector{T}(A::QuantumObject)
+
+If `A.data` is a sparse vector, return a new [`QuantumObject`](@ref) where `A.data` is in the type of `AMDGPU.rocSPARSE.ROCSparseVector` with element type `T` for gpu calculations.
+"""
+ROCSparseVector{T}(A::QuantumObject) where {T} = QuantumObject(ROCSparseVector{T}(A.data), A.type, A.dimensions)
+
+@doc raw"""
+    ROCSparseMatrixCSC(A::QuantumObject)
+
+If `A.data` is in the type of `SparseMatrixCSC`, return a new [`QuantumObject`](@ref) where `A.data` is in the type of `AMDGPU.rocSPARSE.ROCSparseMatrixCSC` for gpu calculations.
+"""
+ROCSparseMatrixCSC(A::QuantumObject) = QuantumObject(ROCSparseMatrixCSC(A.data), A.type, A.dimensions)
+
+@doc raw"""
+    ROCSparseMatrixCSC{T}(A::QuantumObject)
+
+If `A.data` is in the type of `SparseMatrixCSC`, return a new [`QuantumObject`](@ref) where `A.data` is in the type of `AMDGPU.rocSPARSE.ROCSparseMatrixCSC` with element type `T` for gpu calculations.
+"""
+ROCSparseMatrixCSC{T}(A::QuantumObject) where {T} = QuantumObject(ROCSparseMatrixCSC{T}(A.data), A.type, A.dimensions)
+
+@doc raw"""
+    ROCSparseMatrixCSR(A::QuantumObject)
+
+If `A.data` is in the type of `SparseMatrixCSC`, return a new [`QuantumObject`](@ref) where `A.data` is in the type of `AMDGPU.rocSPARSE.ROCSparseMatrixCSR` for gpu calculations.
+"""
+ROCSparseMatrixCSR(A::QuantumObject) = QuantumObject(ROCSparseMatrixCSR(A.data), A.type, A.dimensions)
+
+@doc raw"""
+    ROCSparseMatrixCSR{T}(A::QuantumObject)
+
+If `A.data` is in the type of `SparseMatrixCSC`, return a new [`QuantumObject`](@ref) where `A.data` is in the type of `AMDGPU.rocSPARSE.ROCSparseMatrixCSR` with element type `T` for gpu calculations.
+"""
+ROCSparseMatrixCSR{T}(A::QuantumObject) where {T} = QuantumObject(ROCSparseMatrixCSR{T}(A.data), A.type, A.dimensions)
+
+@doc raw"""
+    roc(A::QuantumObject; word_size::Int=64)
+
+Return a new [`QuantumObject`](@ref) where `A.data` is in the type of `AMDGPU` arrays for gpu calculations.
+
+# Arguments
+- `A::QuantumObject`: The [`QuantumObject`](@ref)
+- `word_size::Int`: The word size of the element type of `A`, can be either `32` or `64`. Default to `64`.
+"""
+function roc(A::QuantumObject; word_size::Union{Val,Int} = Val(64))
+    _word_size = getVal(makeVal(word_size))
+
+    ((_word_size == 64) || (_word_size == 32)) || throw(DomainError(_word_size, "The word size should be 32 or 64."))
+
+    return roc(A, makeVal(word_size))
+end
+roc(A::QuantumObject, word_size::Union{Val{32},Val{64}}) =
+    QuantumObject(adapt(ROCArray{_convert_eltype_wordsize(eltype(A), word_size)}, A.data), A.type, A.dimensions)
+function roc(
+    A::QuantumObject{ObjType,DimsType,<:SparseVector},
+    word_size::Union{Val{32},Val{64}},
+) where {ObjType<:QuantumObjectType,DimsType<:AbstractDimensions}
+    return ROCSparseVector{_convert_eltype_wordsize(eltype(A), word_size)}(A)
+end
+function roc(
+    A::QuantumObject{ObjType,DimsType,<:SparseMatrixCSC},
+    word_size::Union{Val{32},Val{64}},
+) where {ObjType<:QuantumObjectType,DimsType<:AbstractDimensions}
+    return ROCSparseMatrixCSC{_convert_eltype_wordsize(eltype(A), word_size)}(A)
+end
+
+QuantumToolbox.to_dense(A::MT) where {MT<:AbstractROCSparseArray} = ROCArray(A)
+
+QuantumToolbox.to_dense(::Type{T1}, A::ROCArray{T2}) where {T1<:Number,T2<:Number} = ROCArray{T1}(A)
+QuantumToolbox.to_dense(::Type{T}, A::AbstractROCSparseArray) where {T<:Number} = ROCArray{T}(A)
+
+QuantumToolbox._sparse_similar(A::ROCSparseMatrixCSC, args...) = sparse(args..., fmt = :csc)
+QuantumToolbox._sparse_similar(A::ROCSparseMatrixCSR, args...) = sparse(args..., fmt = :csr)
+
+end

test/ext-test/gpu/amdgpu/Project.toml

@@ -0,0 +1,7 @@
+[deps]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
+QuantumToolbox = "6c2fb7c5-b903-41d2-bc5e-5a7c320b9fab"
+
+[compat]
+AMDGPU = "2.1"