CUDA patch

lib/BloqadeCUDA/Project.toml

@@ -5,8 +5,13 @@ version = "0.1.1"
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+CuYao = "b48ca7a8-dd42-11e8-2b8e-1b7706800275"
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+YaoArrayRegister = "e600142f-9330-5003-8abb-0ebd767abc51"
+YaoSubspaceArrayReg = "bd27d05e-4ce1-5e79-84dd-c5d7d508ade2"
 
 [compat]
 Adapt = "3"

lib/BloqadeCUDA/src/BloqadeCUDA.jl

@@ -6,7 +6,11 @@ using CUDA
 using LinearAlgebra
 using CUDA.CUSPARSE
 using CUDA.CUSPARSE: CuSparseMatrixCSC, CuSparseMatrixCSR, AbstractCuSparseMatrix
+using CuYao
+using YaoSubspaceArrayReg
 
 include("patch.jl")
 
+export cu
+
 end

lib/BloqadeCUDA/src/patch.jl

@@ -1,3 +1,4 @@
+@static if !hasmethod(+, Tuple{Diagonal{T,<:CuArray}, CuSparseMatrixCSC{T}} where T)
 # https://github.com/JuliaGPU/CUDA.jl/issues/1469
 
 for SparseMatrixType in [:CuSparseMatrixCSC, :CuSparseMatrixCSR]
@@ -10,3 +11,17 @@ for SparseMatrixType in [:CuSparseMatrixCSC, :CuSparseMatrixCSR]
         end
     end
 end
+
+end
+
+function CuYao.cu(reg::SubspaceArrayReg{D}) where {D}
+
+    natoms = reg.natoms
+    new_state = CuArray(reg.state)
+    new_subspace = Subspace(
+        reg.subspace.nqubits,
+        reg.subspace.map,
+        CuArray(reg.subspace.subspace_v)
+    )
+    return SubspaceArrayReg{D}(new_state,new_subspace)
+end

lib/BloqadeCUDA/src/sample.jl

@@ -0,0 +1,2 @@
+function sample(weights::CuVector)
+end

lib/BloqadeCUDA/test/subspace.jl

@@ -0,0 +1,65 @@
+using Test
+using Random
+using CUDA
+using YaoSubspaceArrayReg
+using YaoArrayRegister
+using StatsBase
+using CuYao
+using GPUArrays
+using Adapt
+using BloqadeCUDA
+
+
+struct Sampler{A,B}
+    cum_prob::A
+    values::B
+end
+
+struct CuWeights{T<:Real}
+    values::CuVector{T}
+    sum::T
+end
+
+
+function Adapt.adapt_structure(to, sampler::Sampler)
+    cum_prob = Adapt.adapt_structure(to, sampler.cum_prob)
+    values = Adapt.adapt_structure(to, sampler.values)
+    Sampler(cum_prob,values)
+end
+
+
+function StatsBase.Weights(values::CuVector{T}) where {T<:Real}
+    return CuWeights{eltype(values)}(values,sum(values))
+end
+
+function (sampler::Sampler)(x)
+    i = searchsortedfirst(sampler.cum_prob, x)
+    i = clamp(i, firstindex(sampler.values), lastindex(sampler.values))
+    @inbounds sampler.values[i]
+end
+
+
+function sample(rng::AbstractRNG, subspace_v::CuVector, weights::CuWeights,nshots::Integer)
+    dices = rand(rng, nshots)
+    sampler = Sampler(cumsum(weights.values),subspace_v)
+
+    Array(sampler.(dices))
+end
+
+
+
+space = Subspace(10, sort(randperm(1 << 10)[1:76] .- 1))
+r = SubspaceArrayReg(randn(ComplexF64, 76), space)
+normalize!(r)
+dr = cu(r)
+
+
+weights = Weights(abs2.(relaxedvec(dr)))
+subspace_v = vec(dr.subspace)
+
+samples = sample(CURAND.default_rng(),subspace_v,weights,1000)
+
+# weights = Weights(abs2.(relaxedvec(dr)))
+# @which measure(r)
+
+# measure(dr)

lib/BloqadeExpr/src/space.jl

@@ -112,3 +112,7 @@ Base.to_index(ss::Subspace) = ss.subspace_v .+ 1
 function Base.:(==)(x::Subspace, y::Subspace)
     return (x.nqubits == y.nqubits) && (x.map == y.map) && (x.subspace_v == y.subspace_v)
 end
+
+function Adapt.adapt_structure(to, x::Subspace{T, S <: AbstractVector{T}}) where T
+    return Subspace{T,S}(x.nqubits,x.map, Adapt.adapt_structure(to, x.subspace_v))
+end

lib/YaoSubspaceArrayReg/src/type.jl

@@ -42,7 +42,7 @@ YaoArrayRegister.relaxedvec(reg::SubspaceArrayReg) = reg.state
 YaoArrayRegister.datatype(reg::SubspaceArrayReg) = eltype(reg.state)
 
 function Adapt.adapt_structure(to, x::SubspaceArrayReg{D}) where D
-    return SubspaceArrayReg{D}(Adapt.adapt(to, x.state), x.subspace)
+    return SubspaceArrayReg{D}(Adapt.adapt(to, x.state), Adapt.adapt_structure(to, x.subspace))
 end
 
 """