Update to cuQuantum 23.10 (#2210)

kshyatt · maleadt · web-flow · commit 2b97ab265cc3 · 2024-01-12T13:01:19.000+01:00
Co-authored-by: Tim Besard &lt;tim.besard@gmail.com&gt;
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -106,18 +106,12 @@ steps:
           setup:
             cuda:
               - "11.4"
+              - "12.0"
             package:
               - "cuDNN"
               - "cuTENSOR"
               - "cuStateVec"
               - "cuTensorNet"
-          adjustments:
-            - with:
-                cuda: "12.0"
-                package: "cuDNN"
-            - with:
-                cuda: "12.0"
-                package: "cuTENSOR"
         plugins:
           - JuliaCI/julia#v1:
               version: "1.9"
diff --git a/lib/custatevec/Project.toml b/lib/custatevec/Project.toml
@@ -13,5 +13,5 @@ cuQuantum_jll = "b75408ef-6fdf-5d74-b65a-7df000ad18e6"
 CEnum = "0.2, 0.3, 0.4"
 CUDA = "~5.1"
 CUDA_Runtime_Discovery = "0.2"
-cuQuantum_jll = "~23.6"
+cuQuantum_jll = "~23.10"
 julia = "1.6"
diff --git a/lib/custatevec/src/cuStateVec.jl b/lib/custatevec/src/cuStateVec.jl
@@ -1,7 +1,7 @@
 module cuStateVec
 
 using CUDA
-using CUDA: CUstream, cudaDataType, @checked, HandleCache, with_workspace, libraryPropertyType
+using CUDA: CUstream, cudaDataType, cudaEvent_t, @checked, HandleCache, with_workspace, libraryPropertyType
 using CUDA: unsafe_free!, retry_reclaim, initialize_context, isdebug
 
 using CEnum: @cenum
diff --git a/lib/custatevec/src/libcustatevec.jl b/lib/custatevec/src/libcustatevec.jl
diff --git a/lib/custatevec/src/statevec.jl b/lib/custatevec/src/statevec.jl
@@ -16,6 +16,20 @@ function applyMatrix!(sv::CuStateVec, matrix::Union{Matrix, CuMatrix}, adjoint::
     sv
 end
 
+function applyMatrixBatched!(sv::CuStateVec, n_svs::Int, map_type::custatevecMatrixMapType_t, matrix_inds::Vector{Int}, matrix::Union{Vector, CuVector}, n_matrices::Int, adjoint::Bool, targets::Vector{<:Integer}, controls::Vector{<:Integer}, controlValues::Vector{<:Integer}=fill(one(Int32), length(controls)))
+    sv_stride    = div(length(sv.data), n_svs)
+    n_index_bits = Int(log2(div(length(sv.data), n_svs)))
+    function bufferSize()
+        out = Ref{Csize_t}()
+        custatevecApplyMatrixBatchedGetWorkspaceSize(handle(), eltype(sv), n_index_bits, n_svs, sv_stride, map_type, matrix_inds, matrix, eltype(matrix), CUSTATEVEC_MATRIX_LAYOUT_COL, Int32(adjoint), n_matrices, length(targets), length(controls), compute_type(eltype(sv), eltype(matrix)), out)
+        out[]
+    end
+    with_workspace(bufferSize) do buffer
+        custatevecApplyMatrixBatched(handle(), sv.data, eltype(sv), n_index_bits, n_svs, sv_stride, map_type, matrix_inds, matrix, eltype(matrix), CUSTATEVEC_MATRIX_LAYOUT_COL, Int32(adjoint), n_matrices, convert(Vector{Int32}, targets), length(targets), convert(Vector{Int32}, controls), convert(Vector{Int32}, controlValues), length(controls), compute_type(eltype(sv), eltype(matrix)), buffer, length(buffer))
+    end
+    sv
+end
+
 function applyGeneralizedPermutationMatrix!(sv::CuStateVec, permutation::Union{Vector{<:Integer}, CuVector{<:Integer}}, diagonals::Union{Vector, CuVector}, adjoint::Bool, targets::Vector{<:Integer}, controls::Vector{<:Integer}, controlValues::Vector{<:Integer}=fill(one(Int32), length(controls)))
     function bufferSize()
         out = Ref{Csize_t}()
@@ -29,8 +43,8 @@ function applyGeneralizedPermutationMatrix!(sv::CuStateVec, permutation::Union{V
 end
 
 function abs2SumOnZBasis(sv::CuStateVec, basisInds::Vector{<:Integer})
-    abs2sum0 = Ref{Float64}[]
-    abs2sum1 = Ref{Float64}[]
+    abs2sum0 = Ref{Float64}(0.0)
+    abs2sum1 = Ref{Float64}(0.0)
     custatevecAbs2SumOnZBasis(handle(), sv.data, eltype(sv), sv.nbits, abs2sum0, abs2sum1, basisInds, length(basisInds))
     return abs2sum0[], abs2sum1[]
 end
@@ -52,12 +66,35 @@ function collapseByBitString!(sv::CuStateVec, bitstring::Union{Vector{<:Integer}
     sv
 end
 
+function collapseByBitStringBatched!(sv::CuStateVec, n_svs::Int, bitstrings::Vector{<:Integer}, bitordering::Vector{<:Integer}, norms::Vector{Float64})
+    function bufferSize()
+        out = Ref{Csize_t}()
+        custatevecCollapseByBitStringBatchedGetWorkspaceSize(handle(), n_svs, convert(Vector{custatevecIndex_t}, bitstrings), norms, out)
+        out[]
+    end
+    sv_stride    = div(length(sv.data), n_svs)
+    n_index_bits = Int(log2(div(length(sv.data), n_svs)))
+    with_workspace(bufferSize) do buffer
+        custatevecCollapseByBitStringBatched(handle(), sv.data, eltype(sv), n_index_bits, n_svs, sv_stride, convert(Vector{custatevecIndex_t}, bitstrings), convert(Vector{Int32}, bitordering), n_index_bits, norms, buffer, length(buffer))
+    end
+    sv
+end
+
 function abs2SumArray(sv::CuStateVec, bitordering::Vector{<:Integer}, maskBitString::Vector{<:Integer}, maskOrdering::Vector{<:Integer})
     abs2sum = Vector{Float64}(undef, 2^length(bitordering))
     custatevecAbs2SumArray(handle(), sv.data, eltype(sv), sv.nbits, abs2sum, convert(Vector{Int32}, bitordering), length(bitordering), convert(Vector{Int32}, maskBitString), convert(Vector{Int32}, maskOrdering), length(maskOrdering))
     return abs2sum
 end
 
+function abs2SumArrayBatched(sv::CuStateVec, n_svs::Int, bitordering::Vector{<:Integer}, maskBitStrings::Vector{<:Integer}, maskOrdering::Vector{<:Integer})
+    abs2sum      = zeros(Float64, n_svs * 2^length(bitordering))
+    sv_stride    = div(length(sv.data), n_svs)
+    n_index_bits = Int(log2(div(length(sv.data), n_svs)))
+    sum_stride   = 2^length(bitordering)
+    custatevecAbs2SumArrayBatched(handle(), sv.data, eltype(sv), n_index_bits, n_svs, sv_stride, abs2sum, sum_stride, convert(Vector{Int32}, bitordering), length(bitordering), convert(Vector{Int32}, maskBitStrings), convert(Vector{Int32}, maskOrdering), length(maskOrdering))
+    return abs2sum
+end
+
 function batchMeasure!(sv::CuStateVec, bitordering::Vector{<:Integer}, randnum::Float64, collapse::custatevecCollapseOp_t=CUSTATEVEC_COLLAPSE_NONE)
     0.0 <= randnum < 1.0 || throw(ArgumentError("randnum $randnum must be in the interval [0, 1)."))
     bitstring = zeros(Int32, length(bitordering))
@@ -86,9 +123,9 @@ function expectation(sv::CuStateVec, matrix::Union{Matrix, CuMatrix}, basis_bits
     return expVal[], residualNorm[]
 end
 
-function expectationsOnPauliBasis(sv::CuStateVec, pauliOps::Vector{Pauli}, basisInds::Vector{Vector{<:Integer}})
-    exp_vals = Vector{Float64}(undef, length(pauliOps))
-    cupaulis = CuStateVecPauli.(pauliOps)
+function expectationsOnPauliBasis(sv::CuStateVec, pauliOps::Vector{Vector{Pauli}}, basisInds::Vector{Vector{Int}})
+    exp_vals = zeros(Float64, length(pauliOps))
+    cupaulis = [[CuStateVecPauli(O) for O in op] for op in pauliOps]
     custatevecComputeExpectationsOnPauliBasis(handle(), sv.data, eltype(sv), sv.nbits, exp_vals, cupaulis, length(pauliOps), convert(Vector{Vector{Int32}}, basisInds), length.(basisInds))
     return exp_vals
 end
@@ -141,3 +178,11 @@ function testMatrixType(matrix::Union{Matrix, CuMatrix}, adjoint::Bool, matrix_t
     end
     return residualNorm[]
 end
+
+function accessorSet(a::CuStateVecAccessor, external_buf::Union{Vector, CuVector}, i_begin::Int, i_end::Int)
+    custatevecAccessorSet(handle(), a, external_buf, i_begin, i_end)
+end
+
+function accessorGet(a::CuStateVecAccessor, external_buf::Union{Vector, CuVector}, i_begin::Int, i_end::Int)
+    custatevecAccessorGet(handle(), a, external_buf, i_begin, i_end)
+end
diff --git a/lib/custatevec/src/types.jl b/lib/custatevec/src/types.jl
@@ -79,6 +79,7 @@ CuStateVec(v::CuVector{T}) where {T} = CuStateVec{T}(v, UInt32(log2(length(v))))
 CuStateVec(v::Vector{T}) where {T}   = CuStateVec(CuVector{T}(v))
 
 Base.eltype(sv::CuStateVec{T}) where T = T
+Base.copy(sv::CuStateVec{T}) where {T} = CuStateVec(copy(sv.data))
 
 mutable struct CuStateVecSampler
     handle::custatevecSamplerDescriptor_t
@@ -94,3 +95,16 @@ mutable struct CuStateVecSampler
 end
 
 Base.unsafe_convert(::Type{custatevecSamplerDescriptor_t}, desc::CuStateVecSampler) = desc.handle
+
+mutable struct CuStateVecAccessor
+    handle::custatevecAccessorDescriptor_t
+    ws_size::Csize_t
+    function CuStateVecAccessor(sv::CuStateVec, bit_ordering::Vector{Int}, mask_bit_string::Vector{Int}, mask_ordering::Vector{Int})
+        desc_ref   = Ref{custatevecAccessorDescriptor_t}()
+        extra_size = Ref{Csize_t}(0)
+        custatevecAccessorCreate(handle(), pointer(sv.data), eltype(sv), sv.nbits, desc_ref, bit_ordering, length(bit_ordering), mask_bit_string, mark_ordering, length(mask_bit_string), extra_size)
+        obj = new(desc_ref[], extra_size[])
+        finalizer(custatevecAccessorDestroy, obj)
+        obj
+    end
+end
diff --git a/lib/custatevec/test/runtests.jl b/lib/custatevec/test/runtests.jl
@@ -8,7 +8,7 @@ using cuStateVec
 @info "cuStateVec version: $(cuStateVec.version())"
 
 @testset "cuStateVec" begin
-    import cuStateVec: CuStateVec, applyMatrix!, applyPauliExp!, applyGeneralizedPermutationMatrix!, expectation, expectationsOnPauliBasis, sample, testMatrixType, Pauli, PauliX, PauliY, PauliZ, PauliI, measureOnZBasis!, swapIndexBits!
+    import cuStateVec: CuStateVec, applyMatrix!, applyMatrixBatched!, applyPauliExp!, applyGeneralizedPermutationMatrix!, expectation, expectationsOnPauliBasis, sample, testMatrixType, Pauli, PauliX, PauliY, PauliZ, PauliI, measureOnZBasis!, swapIndexBits!, abs2SumOnZBasis, collapseOnZBasis!, batchMeasure!, abs2SumArray, collapseByBitString!, abs2SumArrayBatched, collapseByBitStringBatched!
 
     @testset "applyMatrix! and expectation" begin
         # build a simple state and compute expectations
@@ -19,7 +19,8 @@ using cuStateVec
             Z = convert(Matrix{elty}, [1 0; 0 -1])
             sv = CuStateVec(elty, n_q)
             sv = applyMatrix!(sv, H, false, Int32[0], Int32[])
-            exp, res = expectation(sv, Z, Int32[0])
+            sv = applyMatrix!(sv, H, false, Int32[1], Int32[])
+            exp, res = expectation(sv, Z, Int32[1])
             @test exp ≈ 0.0 atol=1e-6
             exp, res = expectation(sv, X, Int32[0])
             @test exp ≈ 1.0 atol=1e-6
@@ -38,6 +39,69 @@ using cuStateVec
             exp, res = expectation(sv, X, Int32[0])
             @test exp ≈ 0.0 atol=1e-6
         end
+        # with expectationsOnPauliBasis
+        n_q = 2
+        @testset for elty in [ComplexF32, ComplexF64]
+            H = convert(Matrix{elty}, (1/√2).*[1 1; 1 -1])
+            X = convert(Matrix{elty}, [0 1; 1 0])
+            Z = convert(Matrix{elty}, [1 0; 0 -1])
+            sv = CuStateVec(elty, n_q)
+            sv = applyMatrix!(sv, H, false, Int32[0], Int32[])
+            sv = applyMatrix!(sv, H, false, Int32[1], Int32[])
+            pauli_ops = [cuStateVec.Pauli[cuStateVec.PauliX()], cuStateVec.Pauli[cuStateVec.PauliX()]]
+            exp_vals = expectationsOnPauliBasis(sv, pauli_ops, [[0], [1]])
+            @test exp_vals[1] ≈ 1.0 atol=1e-6
+            @test exp_vals[2] ≈ 1.0 atol=1e-6
+        end
+    end
+    @testset "applyMatrixBatched! and expectation" begin
+        # build a simple state and compute expectations
+        n_q = 2
+        @testset for elty in [ComplexF32, ComplexF64]
+            H = convert(Matrix{elty}, (1/√2).*[1 1; 1 -1])
+            X = convert(Matrix{elty}, [0 1; 1 0])
+            Z = convert(Matrix{elty}, [1 0; 0 -1])
+            @testset for n_svs in (1, 2)
+                @testset for (mapping, mat_inds, n_mats) in (
+                                                             (cuStateVec.CUSTATEVEC_MATRIX_MAP_TYPE_MATRIX_INDEXED, collect(0:n_svs-1), n_svs), 
+                                                             (cuStateVec.CUSTATEVEC_MATRIX_MAP_TYPE_MATRIX_INDEXED, fill(0, n_svs), 1), 
+                                                             (cuStateVec.CUSTATEVEC_MATRIX_MAP_TYPE_BROADCAST, fill(0, n_svs), 1),
+                                                  )
+                    batched_vec = CUDA.zeros(elty, n_svs*2^(n_q))
+                    for sv_ix in 0:n_svs-1
+                        CUDA.@allowscalar batched_vec[sv_ix*(2^n_q) + 1] = one(elty)
+                    end
+                    sv = CuStateVec(batched_vec) # padded state vector
+                    H_batch = CuVector{elty}(repeat(vec(H), n_mats))
+                    sv = applyMatrixBatched!(sv, n_svs, mapping, mat_inds, H_batch, n_mats, false, Int32[0], Int32[])
+                    CUDA.@allowscalar begin
+                        for sv_ix in 0:n_svs-1
+                            ix_begin = sv_ix*2^n_q + 1
+                            ix_end   = (sv_ix+1)*2^n_q
+                            sv_ = CuStateVec(sv.data[ix_begin:ix_end])
+                            exp, res = expectation(sv_, Z, Int32[0])
+                            @test exp ≈ 0.0 atol=1e-6
+                            exp, res = expectation(sv_, X, Int32[0])
+                            @test exp ≈ 1.0 atol=1e-6
+                        end
+                    end
+                end
+            end
+        end
+        # build a simple state with controls and compute expectations
+        n_q = 2
+        @testset for elty in [ComplexF32, ComplexF64]
+            H = convert(Matrix{elty}, (1/√2).*[1 1; 1 -1])
+            X = convert(Matrix{elty}, [0 1; 1 0])
+            Z = convert(Matrix{elty}, [1 0; 0 -1])
+            sv = CuStateVec(elty, n_q)
+            sv = applyMatrix!(sv, H, false, Int32[0], Int32[])
+            sv = applyMatrix!(sv, X, false, Int32[1], Int32[0]) # CNOT
+            exp, res = expectation(sv, Z, Int32[0])
+            @test exp ≈ 0.0 atol=1e-6
+            exp, res = expectation(sv, X, Int32[0])
+            @test exp ≈ 0.0 atol=1e-6
+        end
     end
     @testset "applyMatrix! and sample" begin
         # build a simple state and compute samples
@@ -74,6 +138,22 @@ using cuStateVec
             @test collect(sv_result.data) ≈ h_sv_result
         end
     end
+    @testset "abs2sumOnZBasis and collapseOnZBasis!" begin
+        @testset for elty in [ComplexF32, ComplexF64]
+            h_sv = 1.0/√8 .* elty[0.0, im, 0.0, im, 0.0, im, 0.0, im]
+            h_sv_result_0 = 1.0/√2 * elty[0.0, 0.0, 0.0, im, 0.0, im,  0.0, 0.0]
+            h_sv_result_1 = 1.0/√2 * elty[0.0, im, 0.0, 0.0, 0.0, 0.0, 0.0, im]
+            sv   = CuStateVec(h_sv)
+            abs2sum0, abs2sum1 = abs2SumOnZBasis(sv, [0, 1, 2])
+            abs2sum = abs2sum0 + abs2sum1
+            for (parity, norm, h_sv_result) in ((0, abs2sum0, h_sv_result_0), (1, abs2sum1, h_sv_result_1))
+                d_sv = copy(sv)
+                d_sv = collapseOnZBasis!(d_sv, parity, [0, 1, 2], norm)
+                sv_result  = collect(d_sv.data)
+                @test sv_result ≈ h_sv_result
+            end
+        end
+    end
     @testset "measureOnZBasis" begin
         @testset for elty in [ComplexF32, ComplexF64]
             h_sv = 1.0/√8 .* elty[0.0, im, 0.0, im, 0.0, im, 0.0, im]
@@ -84,6 +164,61 @@ using cuStateVec
             @test sv_result ≈ h_sv_result
         end
     end
+    @testset "abs2SumArray and collapseByBitString!" begin
+        nq = 3
+        bit_ordering = [2, 1, 0]
+        @testset for elty in [ComplexF32, ComplexF64]
+            h_sv = elty[0.0, 0.1*im, 0.1+0.1*im, 0.1+0.2*im, 0.2+0.2*im, 0.3+0.3im, 0.3+0.4*im, 0.4+0.5*im]
+            h_sv_result = elty[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.3+0.4*im, 0.0]
+            sv   = CuStateVec(h_sv)
+            abs2sum = abs2SumArray(sv, bit_ordering, Int[], Int[])
+            bitstr = [1, 1, 0]
+            d_sv = copy(sv)
+            d_sv = collapseByBitString!(d_sv, bitstr, bit_ordering, 1.)
+            sv_result  = collect(d_sv.data)
+            @test sv_result ≈ h_sv_result
+        end
+    end
+    @testset "abs2SumArrayBatched" begin
+        bit_ordering = [1]
+        @testset for elty in [ComplexF32, ComplexF64]
+            @testset for n_svs in (2,)
+                h_sv = elty[0.0, 0.1*im, 0.1 + 0.1*im, 0.1 + 0.2*im, 0.2+0.2*im, 0.3+0.3*im, 0.3+0.4*im, 0.4+0.5*im, 0.25+0.25*im, 0.25+0.25*im, 0.25+0.25*im, 0.25+0.25*im, 0.25+0.25*im, 0.25+0.25*im, 0.25+0.25*im, 0.25+0.25*im]
+                a2s_result = real(elty)[0.27, 0.73, 0.5, 0.5]
+                sv      = CuStateVec(h_sv)
+                abs2sum = abs2SumArrayBatched(sv, n_svs, bit_ordering, Int[], Int[])
+                @test abs2sum ≈ a2s_result
+            end
+        end
+    end
+    @testset "collapseByBitStringBatched!" begin
+        bit_ordering = [0, 1, 2]
+        @testset for elty in [ComplexF32, ComplexF64]
+            @testset for n_svs in (2,)
+                h_sv = elty[0.0, 0.1*im, 0.1 + 0.1*im, 0.1 + 0.2*im, 0.2+0.2*im, 0.3+0.3*im, 0.3+0.4*im, 0.4+0.5*im, 0.0, 0.1*im, 0.1+0.1*im, 0.1+0.2*im, 0.2+0.2*im, 0.3+0.3*im, 0.3+0.4*im, 0.4*0.5*im]
+                h_sv_result = elty[0.0, im, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.6+0.8*im, 0.0]
+                sv      = CuStateVec(h_sv)
+                bitstr = [0b001, 0b110]
+                d_sv = copy(sv)
+                d_sv = collapseByBitStringBatched!(d_sv, n_svs, bitstr, bit_ordering, [0.01, 0.25])
+                sv_result  = collect(d_sv.data)
+                @test sv_result ≈ h_sv_result
+            end
+        end
+    end
+    @testset "batchMeasure!" begin
+        nq = 3
+        bit_ordering = [2, 1, 0]
+        @testset for elty in [ComplexF32, ComplexF64]
+            h_sv = elty[0.0, 0.1*im, 0.1+0.1*im, 0.1+0.2*im, 0.2+0.2*im, 0.3+0.3im, 0.3+0.4*im, 0.4+0.5*im]
+            h_sv_result = elty[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.6+0.8*im, 0.0]
+            sv   = CuStateVec(h_sv)
+            sv, bitstr = batchMeasure!(sv, bit_ordering, 0.5, cuStateVec.CUSTATEVEC_COLLAPSE_NORMALIZE_AND_ZERO)
+            sv_result  = collect(sv.data)
+            @test sv_result ≈ h_sv_result
+            @test bitstr == [1, 1, 0]
+        end
+    end
     @testset "swapIndexBits" begin
         @testset for elty in [ComplexF32, ComplexF64]
             # 0.1|000> + 0.4|011> - 0.4|101> - 0.3im|111>
diff --git a/lib/cutensornet/Project.toml b/lib/cutensornet/Project.toml
@@ -15,7 +15,7 @@ cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 CEnum = "0.2, 0.3, 0.4"
 CUDA = "~5.1"
 CUDA_Runtime_Discovery = "0.2"
-cuQuantum_jll = "~22.11"
+cuQuantum_jll = "~23.10"
 cuTENSOR = "~1.0, ~1.1, ~1.2"
 julia = "1.6"
 LinearAlgebra = "1"
diff --git a/lib/cutensornet/src/cuTensorNet.jl b/lib/cutensornet/src/cuTensorNet.jl
@@ -3,7 +3,7 @@ module cuTensorNet
 using LinearAlgebra
 using CUDA
 using CUDA: CUstream, cudaDataType, @checked, HandleCache, with_workspace
-using CUDA: retry_reclaim, initialize_context, isdebug
+using CUDA: retry_reclaim, initialize_context, isdebug, cuDoubleComplex
 
 using cuTENSOR
 using cuTENSOR: CuTensor
diff --git a/lib/cutensornet/src/libcutensornet.jl b/lib/cutensornet/src/libcutensornet.jl
diff --git a/lib/cutensornet/src/types.jl b/lib/cutensornet/src/types.jl
diff --git a/res/wrap/custatevec.toml b/res/wrap/custatevec.toml
