From 1f857be3d80aff90b5d643820726db1f7790e04e Mon Sep 17 00:00:00 2001
From: Alberto Mercurio <61953577+albertomercurio@users.noreply.github.com>
Date: Sun, 15 Dec 2024 13:53:49 +0100
Subject: [PATCH] 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>
---
 CHANGELOG.md                          |  3 +-
 Project.toml                          |  5 ++
 ext/QuantumToolboxGPUArraysExt.jl     | 33 ++++++++++
 src/qobj/arithmetic_and_attributes.jl |  4 +-
 test/ext-test/gpu/cuda_ext.jl         | 88 +++++++++++++++++++++++++++
 5 files changed, 131 insertions(+), 2 deletions(-)
 create mode 100644 ext/QuantumToolboxGPUArraysExt.jl

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d61a642f..3e3a8c0d 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/Project.toml b/Project.toml
index 8b71ebaa..32d1ad8e 100644
--- a/Project.toml
+++ b/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"
diff --git a/ext/QuantumToolboxGPUArraysExt.jl b/ext/QuantumToolboxGPUArraysExt.jl
new file mode 100644
index 00000000..d9b68681
--- /dev/null
+++ b/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
diff --git a/src/qobj/arithmetic_and_attributes.jl b/src/qobj/arithmetic_and_attributes.jl
index 6a106148..a3baf214 100644
--- a/src/qobj/arithmetic_and_attributes.jl
+++ b/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)
 
diff --git a/test/ext-test/gpu/cuda_ext.jl b/test/ext-test/gpu/cuda_ext.jl
index 20219b42..5ed3ce18 100644
--- a/test/ext-test/gpu/cuda_ext.jl
+++ b/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