Improve GPU support for steadystate

Author: albertomercurio

ext/QuantumToolboxCUDAExt.jl

@@ -2,6 +2,7 @@ module QuantumToolboxCUDAExt
 
 using QuantumToolbox
 using QuantumToolbox: makeVal, getVal
+import QuantumToolbox: _sparse_similar
 import CUDA: cu, CuArray, allowscalar
 import CUDA.CUSPARSE: CuSparseVector, CuSparseMatrixCSC, CuSparseMatrixCSR, AbstractCuSparseArray
 import SparseArrays: SparseVector, SparseMatrixCSC
@@ -106,4 +107,7 @@ QuantumToolbox.to_dense(A::MT) where {MT<:AbstractCuSparseArray} = CuArray(A)
 QuantumToolbox.to_dense(::Type{T1}, A::CuArray{T2}) where {T1<:Number,T2<:Number} = CuArray{T1}(A)
 QuantumToolbox.to_dense(::Type{T}, A::AbstractCuSparseArray) where {T<:Number} = CuArray{T}(A)
 
+QuantumToolbox._sparse_similar(A::CuSparseMatrixCSC, args...) = sparse(args..., fmt = :csc)
+QuantumToolbox._sparse_similar(A::CuSparseMatrixCSR, args...) = sparse(args..., fmt = :csr)
+
 end

src/steadystate.jl

@@ -108,18 +108,18 @@ function _steadystate(L::QuantumObject{SuperOperatorQuantumObject}, solver::Stea
     N = prod(L.dimensions)
     weight = norm(L_tmp, 1) / length(L_tmp)
 
-    v0 = _get_dense_similar(L_tmp, N^2)
+    v0 = _dense_similar(L_tmp, N^2)
     fill!(v0, 0)
     allowed_setindex!(v0, weight, 1) # Because scalar indexing is not allowed on GPU arrays
 
     idx_range = collect(1:N)
-    rows = _get_dense_similar(L_tmp, N)
-    cols = _get_dense_similar(L_tmp, N)
-    vals = _get_dense_similar(L_tmp, N)
+    rows = _dense_similar(L_tmp, N)
+    cols = _dense_similar(L_tmp, N)
+    vals = _dense_similar(L_tmp, N)
     fill!(rows, 1)
     copyto!(cols, N .* (idx_range .- 1) .+ idx_range)
     fill!(vals, weight)
-    Tn = sparse(rows, cols, vals, N^2, N^2)
+    Tn = _sparse_similar(L_tmp, rows, cols, vals, N^2, N^2)
     L_tmp = L_tmp + Tn
 
     (haskey(kwargs, :Pl) || haskey(kwargs, :Pr)) && error("The use of preconditioners must be defined in the solver.")
@@ -155,14 +155,14 @@ function _steadystate(L::QuantumObject{SuperOperatorQuantumObject}, solver::Stea
     N = prod(L.dimensions)
     weight = norm(L_tmp, 1) / length(L_tmp)
 
-    v0 = _get_dense_similar(L_tmp, N^2)
+    v0 = _dense_similar(L_tmp, N^2)
     fill!(v0, 0)
     allowed_setindex!(v0, weight, 1) # Because scalar indexing is not allowed on GPU arrays
 
     idx_range = collect(1:N)
-    rows = _get_dense_similar(L_tmp, N)
-    cols = _get_dense_similar(L_tmp, N)
-    vals = _get_dense_similar(L_tmp, N)
+    rows = _dense_similar(L_tmp, N)
+    cols = _dense_similar(L_tmp, N)
+    vals = _dense_similar(L_tmp, N)
     fill!(rows, 1)
     copyto!(cols, N .* (idx_range .- 1) .+ idx_range)
     fill!(vals, weight)

src/utilities.jl

@@ -130,8 +130,10 @@ end
 
 get_typename_wrapper(A) = Base.typename(typeof(A)).wrapper
 
-_get_dense_similar(A::AbstractArray, args...) = similar(A, args...)
-_get_dense_similar(A::AbstractSparseMatrix, args...) = similar(nonzeros(A), args...)
+_dense_similar(A::AbstractArray, args...) = similar(A, args...)
+_dense_similar(A::AbstractSparseMatrix, args...) = similar(nonzeros(A), args...)
+
+_sparse_similar(A::AbstractArray, args...) = sparse(args...)
 
 _Ginibre_ensemble(n::Int, rank::Int = n) = randn(ComplexF64, n, rank) / sqrt(n)
 

test/ext-test/gpu/cuda_ext.jl

@@ -125,6 +125,31 @@
     @test all([isapprox(sol_cpu.expect[i], sol_gpu32.expect[i]; atol = 1e-6) for i in 1:length(tlist)])
 end
 
+@testset "CUDA steadystate" begin
+    N = 50
+    Δ = 0.01
+    F = 0.1
+    γ = 0.1
+    nth = 2
+
+    a = destroy(N)
+    H = Δ * a' * a + F * (a + a')
+    c_ops = [sqrt(γ * (nth + 1)) * a, sqrt(γ * nth) * a']
+
+    ρ_ss_cpu = steadystate(H, c_ops)
+
+    H_gpu_csc = cu(H)
+    c_ops_gpu_csc = [cu(c_op) for c_op in c_ops]
+    ρ_ss_gpu_csc = steadystate(H_gpu_csc, c_ops_gpu_csc, solver = SteadyStateLinearSolver())
+
+    H_gpu_csr = CuSparseMatrixCSR(H_gpu_csc)
+    c_ops_gpu_csr = [CuSparseMatrixCSR(c_op) for c_op in c_ops_gpu_csc]
+    ρ_ss_gpu_csr = steadystate(H_gpu_csr, c_ops_gpu_csr, solver = SteadyStateLinearSolver())
+
+    @test ρ_ss_cpu.data ≈ Array(ρ_ss_gpu_csc.data) atol = 1e-8*length(ρ_ss_cpu)
+    @test ρ_ss_cpu.data ≈ Array(ρ_ss_gpu_csr.data) atol = 1e-8*length(ρ_ss_cpu)
+end
+
 @testset "CUDA ptrace" begin
     g = fock(2, 1)
     e = fock(2, 0)