GPU support for steadystate_fourier

Author: Fe-r-oz

ext/QuantumToolboxCUDAExt.jl

@@ -3,11 +3,13 @@ module QuantumToolboxCUDAExt
 using QuantumToolbox
 using QuantumToolbox: makeVal, getVal
 import QuantumToolbox: _sparse_similar, _convert_eltype_wordsize
-import CUDA: cu, CuArray, allowscalar
+import CUDA: cu, CuArray, allowscalar, @allowscalar, has_cuda
 import CUDA.CUSPARSE: CuSparseVector, CuSparseMatrixCSC, CuSparseMatrixCSR, AbstractCuSparseArray
-import SparseArrays: SparseVector, SparseMatrixCSC, sparse
+import SparseArrays: SparseVector, SparseMatrixCSC, sparse, spzeros
 import CUDA.Adapt: adapt
 
+export _safe_setindex!
+
 allowscalar(false)
 
 @doc raw"""
@@ -104,5 +106,10 @@ QuantumToolbox.to_dense(::Type{T}, A::AbstractCuSparseArray) where {T<:Number} =
 
 QuantumToolbox._sparse_similar(A::CuSparseMatrixCSC, args...) = sparse(args..., fmt = :csc)
 QuantumToolbox._sparse_similar(A::CuSparseMatrixCSR, args...) = sparse(args..., fmt = :csr)
-
+_sparse_similar(A::CuSparseMatrixCSC, I::AbstractVector, J::AbstractVector, V::AbstractVector, m::Int, n::Int) =
+    CuSparseMatrixCSC(sparse(I, J, V, m, n))
+_sparse_similar(A::CuSparseMatrixCSC, m::Int, n::Int) = CuSparseMatrixCSC(spzeros(eltype(A), m, n))
+_sparse_similar(A::CuSparseMatrixCSR, I::AbstractVector, J::AbstractVector, V::AbstractVector, m::Int, n::Int) =
+    CuSparseMatrixCSR(sparse(I, J, V, m, n))
+_sparse_similar(A::CuSparseMatrixCSR, m::Int, n::Int) = CuSparseMatrixCSR(spzeros(eltype(A), m, n))
 end

src/utilities.jl

@@ -135,7 +135,6 @@ _dense_similar(A::AbstractSparseMatrix, args...) = similar(nonzeros(A), args...)
 
 _sparse_similar(A::AbstractArray, args...) = sparse(args...)
 _sparse_similar(A::AbstractArray, m::Int, n::Int) = spzeros(eltype(A), m, n)
-_sparse_similar(A::AbstractArray, rows, cols, vals, m::Int, n::Int) = sparse(rows, cols, vals, m, n)
 
 _Ginibre_ensemble(n::Int, rank::Int = n) = randn(ComplexF64, n, rank) / sqrt(n)
 

test/ext-test/gpu/cuda_ext.jl

@@ -154,6 +154,27 @@ end
     @test ρ_ss_cpu.data ≈ Array(ρ_ss_gpu_csr.data) atol = 1e-8 * length(ρ_ss_cpu)
 end
 
+@testset "CUDA steadystate_fourier" begin
+    N = 2
+    ωd = 1.0
+    n_max = 2
+    H0 = cu(sigmaz())
+    Hp = cu(sigmax())
+    Hm = cu(sigmax())
+    c_ops = [cu(sqrt(0.1) * sigmam())]
+    ρ_list1 = steadystate_fourier(H0, Hp, Hm, ωd, c_ops; solver = SteadyStateLinearSolver(), n_max = n_max)
+    ρ0 = steadystate_fourier(
+        H0,
+        Hp,
+        Hm,
+        ωd,
+        c_ops;
+        solver = SSFloquetEffectiveLiouvillian(SteadyStateDirectSolver()),
+        n_max = n_max,
+    )
+    @test isapprox(ρ0, ρ_list1[1]; atol = 1e-6)
+end
+
 @testset "CUDA ptrace" begin
     g = fock(2, 1)
     e = fock(2, 0)