Add ptrace support for any GPUArray (#350)

* Add `ptrace` support for any `GPUArray`

* Add changelog

---------

Co-authored-by: Yi-Te Huang <44385685+ytdHuang@users.noreply.github.com>

Author: albertomercurio

CHANGELOG.md

@@ -8,7 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased](https://github.com/qutip/QuantumToolbox.jl/tree/main)
 
 - Change the structure of block diagonalization functions, using `BlockDiagonalForm` struct and changing the function name from `bdf` to `block_diagonal_form`. ([#349])
-
+- Add **GPUArrays** compatibility for `ptrace` function, by using **KernelAbstractions.jl**. ([#350])
 
 ## [v0.24.0]
 Release date: 2024-12-13
@@ -70,3 +70,4 @@ Release date: 2024-11-13
 [#346]: https://github.com/qutip/QuantumToolbox.jl/issues/346
 [#347]: https://github.com/qutip/QuantumToolbox.jl/issues/347
 [#349]: https://github.com/qutip/QuantumToolbox.jl/issues/349
+[#350]: https://github.com/qutip/QuantumToolbox.jl/issues/350

Project.toml

@@ -29,10 +29,13 @@ StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 
 [extensions]
 QuantumToolboxCUDAExt = "CUDA"
 QuantumToolboxCairoMakieExt = "CairoMakie"
+QuantumToolboxGPUArraysExt = ["GPUArrays", "KernelAbstractions"]
 
 [compat]
 Aqua = "0.8"
@@ -44,9 +47,11 @@ DiffEqCallbacks = "4.2.1 - 4"
 DiffEqNoiseProcess = "5"
 Distributed = "1"
 FFTW = "1.5"
+GPUArrays = "10"
 Graphs = "1.7"
 IncompleteLU = "0.2"
 JET = "0.9"
+KernelAbstractions = "0.9.2"
 LinearAlgebra = "1"
 LinearSolve = "2"
 OrdinaryDiffEqCore = "1"

ext/QuantumToolboxGPUArraysExt.jl

@@ -0,0 +1,33 @@
+module QuantumToolboxGPUArraysExt
+
+using QuantumToolbox
+
+import GPUArrays: AbstractGPUArray
+import KernelAbstractions
+import KernelAbstractions: @kernel, @Const, @index, get_backend, synchronize
+
+@kernel function tr_kernel!(B, @Const(A))
+    # i, j, k = @index(Global, NTuple)
+    # Atomix.@atomic B[i, j] += A[i, j, k, k] # TODO: use Atomix when it will support Complex types
+
+    i, j = @index(Global, NTuple)
+    @inbounds B[i, j] = 0
+    @inbounds for k in 1:size(A, 3)
+        B[i, j] += A[i, j, k, k]
+    end
+end
+
+function QuantumToolbox._map_trace(A::AbstractGPUArray{T,4}) where {T}
+    B = similar(A, size(A, 1), size(A, 2))
+    fill!(B, 0)
+
+    backend = get_backend(A)
+    kernel! = tr_kernel!(backend)
+
+    kernel!(B, A, ndrange = size(A)[1:2])
+    KernelAbstractions.synchronize(backend)
+
+    return B
+end
+
+end

src/qobj/arithmetic_and_attributes.jl

@@ -631,11 +631,13 @@ function _ptrace_oper(QO::AbstractArray, dims::Union{SVector,MVector}, sel)
     topermute = reverse(2 * n_d + 1 .- qtrace_sel)
     ρmat = permutedims(ρmat, topermute) # TODO: use PermutedDimsArray when Julia v1.11.0 is released
     ρmat = reshape(ρmat, prod(dkeep), prod(dkeep), prod(dtrace), prod(dtrace))
-    res = map(tr, eachslice(ρmat, dims = (1, 2)))
+    res = _map_trace(ρmat)
 
     return res, dkeep
 end
 
+_map_trace(A::AbstractArray{T,4}) where {T} = map(tr, eachslice(A, dims = (1, 2)))
+
 @doc raw"""
     purity(ρ::QuantumObject)
 

test/ext-test/gpu/cuda_ext.jl

@@ -106,3 +106,91 @@
     @test all([isapprox(sol_cpu.expect[i], sol_gpu64.expect[i]) for i in 1:length(tlist)])
     @test all([isapprox(sol_cpu.expect[i], sol_gpu32.expect[i]; atol = 1e-6) for i in 1:length(tlist)])
 end
+
+@testset "CUDA ptrace" begin
+    g = fock(2, 1)
+    e = fock(2, 0)
+    α = sqrt(0.7)
+    β = sqrt(0.3) * 1im
+    ψ = α * kron(g, e) + β * kron(e, g) |> cu
+
+    ρ1 = ptrace(ψ, 1)
+    ρ2 = ptrace(ψ, 2)
+    @test ρ1.data isa CuArray
+    @test ρ2.data isa CuArray
+    @test Array(ρ1.data) ≈ [0.3 0.0; 0.0 0.7] atol = 1e-10
+    @test Array(ρ2.data) ≈ [0.7 0.0; 0.0 0.3] atol = 1e-10
+
+    ψ_d = ψ'
+
+    ρ1 = ptrace(ψ_d, 1)
+    ρ2 = ptrace(ψ_d, 2)
+    @test ρ1.data isa CuArray
+    @test ρ2.data isa CuArray
+    @test Array(ρ1.data) ≈ [0.3 0.0; 0.0 0.7] atol = 1e-10
+    @test Array(ρ2.data) ≈ [0.7 0.0; 0.0 0.3] atol = 1e-10
+
+    ρ = ket2dm(ψ)
+    ρ1 = ptrace(ρ, 1)
+    ρ2 = ptrace(ρ, 2)
+    @test ρ.data isa CuArray
+    @test ρ1.data isa CuArray
+    @test ρ2.data isa CuArray
+    @test Array(ρ1.data) ≈ [0.3 0.0; 0.0 0.7] atol = 1e-10
+    @test Array(ρ2.data) ≈ [0.7 0.0; 0.0 0.3] atol = 1e-10
+
+    ψ1 = normalize(g + 1im * e)
+    ψ2 = normalize(g + e)
+    ρ1 = ket2dm(ψ1)
+    ρ2 = ket2dm(ψ2)
+    ρ = kron(ρ1, ρ2) |> cu
+    ρ1_ptr = ptrace(ρ, 1)
+    ρ2_ptr = ptrace(ρ, 2)
+    @test ρ1_ptr.data isa CuArray
+    @test ρ2_ptr.data isa CuArray
+    @test ρ1.data ≈ Array(ρ1_ptr.data) atol = 1e-10
+    @test ρ2.data ≈ Array(ρ2_ptr.data) atol = 1e-10
+
+    ψlist = [rand_ket(2), rand_ket(3), rand_ket(4), rand_ket(5)]
+    ρlist = [rand_dm(2), rand_dm(3), rand_dm(4), rand_dm(5)]
+    ψtotal = tensor(ψlist...) |> cu
+    ρtotal = tensor(ρlist...) |> cu
+    sel_tests = [
+        SVector{0,Int}(),
+        1,
+        2,
+        3,
+        4,
+        (1, 2),
+        (1, 3),
+        (1, 4),
+        (2, 3),
+        (2, 4),
+        (3, 4),
+        (1, 2, 3),
+        (1, 2, 4),
+        (1, 3, 4),
+        (2, 3, 4),
+        (1, 2, 3, 4),
+    ]
+    for sel in sel_tests
+        if length(sel) == 0
+            @test ptrace(ψtotal, sel) ≈ 1.0
+            @test ptrace(ρtotal, sel) ≈ 1.0
+        else
+            @test ptrace(ψtotal, sel) ≈ cu(tensor([ket2dm(ψlist[i]) for i in sel]...))
+            @test ptrace(ρtotal, sel) ≈ cu(tensor([ρlist[i] for i in sel]...))
+        end
+    end
+    @test ptrace(ψtotal, (1, 3, 4)) ≈ ptrace(ψtotal, (4, 3, 1)) # check sort of sel
+    @test ptrace(ρtotal, (1, 3, 4)) ≈ ptrace(ρtotal, (3, 1, 4)) # check sort of sel
+
+    @testset "Type Inference (ptrace)" begin
+        @inferred ptrace(ρ, 1)
+        @inferred ptrace(ρ, 2)
+        @inferred ptrace(ψ_d, 1)
+        @inferred ptrace(ψ_d, 2)
+        @inferred ptrace(ψ, 1)
+        @inferred ptrace(ψ, 2)
+    end
+end