CUDA patch

lib/BloqadeCUDA/Project.toml

@@ -4,9 +4,15 @@ version = "0.1.1"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 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,19 @@ 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,91 @@
+using Test
+using Random
+using CUDA
+using YaoSubspaceArrayReg
+using YaoArrayRegister
+using StatsBase
+using CuYao
+using GPUArrays
+using Adapt
+using BloqadeCUDA
+
+# TODO: use Alias sampler
+struct SamplerWithValues{T,A<:AbstractVector{T},B<:AbstractVector}
+    total::T
+    cumulative_weights::A
+    values::B
+end
+
+function SamplerWithValues(weights,values)
+    length(weights) != length(values) && error("Sampler expects values with same length as weights.")
+    total = sum(weights)
+    cumulative_weights = cumsum(weights)
+    return SamplerWithValues(total,cumulative_weights,values)
+end
+
+function Adapt.adapt_structure(to, sampler::SamplerWithValues)
+    cumulative_weights = Adapt.adapt_structure(to, sampler.cumulative_weights)
+    values = Adapt.adapt_structure(to, sampler.values)
+    SamplerWithValues(sampler.total,cumulative_weights,values)
+end
+
+function (sampler::SamplerWithValues)(x)
+    i = searchsortedfirst(sampler.cumulative_weights, x)
+    i = clamp(i, firstindex(sampler.cumulative_weights), lastindex(sampler.cumulative_weights))
+    @inbounds sampler.values[i]
+end
+
+# TODO: use Alias sampler
+struct Sampler{T,A<:AbstractVector{T}}
+    total::T
+    cumulative_weights::A
+end
+
+function Sampler(weights)
+    total = sum(weights)
+    cumulative_weights = cumsum(weights)
+    return Sampler(total,cumulative_weights)
+end
+
+function Adapt.adapt_structure(to, sampler::Sampler)
+    cumulative_weights = Adapt.adapt_structure(to, sampler.cumulative_weights)
+    Sampler(sampler.total,cumulative_weights)
+end
+
+function (sampler::Sampler)(x)
+    i = searchsortedfirst(sampler.cumulative_weights, x)
+    clamp(i, firstindex(sampler.cumulative_weights), lastindex(sampler.cumulative_weights))
+end
+
+
+function sample(rng::AbstractRNG, weights::CuVector,nshots::Integer)
+    sampler = Sampler(weights)
+
+    dices = sampler.total .* rand(rng, nshots)
+
+    sampler.(dices)
+end
+
+function sample(rng::AbstractRNG, weights::CuVector,nshots::Integer, obs::CuVector)
+    sampler = SamplerWithValues(weights,obs)
+
+    dices = sampler.total .* rand(rng, nshots)
+
+    sampler.(dices)
+end
+
+
+N = 20
+ntot = 2^(N-4)
+space = Subspace(N, sort(randperm(1 << N)[1:ntot] .- 1))
+r = SubspaceArrayReg(randn(ComplexF64, ntot), space)
+normalize!(r)
+dr = cu(r)
+
+
+weights = abs2.(relaxedvec(dr))
+
+using BenchmarkTools 
+
+@benchmark CUDA.@sync sample(CURAND.default_rng(),weights,10000)
+

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