diff --git a/lib/BloqadeCUDA/Project.toml b/lib/BloqadeCUDA/Project.toml index 615dc94b2..c8d155421 100644 --- a/lib/BloqadeCUDA/Project.toml +++ b/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" diff --git a/lib/BloqadeCUDA/src/BloqadeCUDA.jl b/lib/BloqadeCUDA/src/BloqadeCUDA.jl index e481cf0b4..4e5a4d05b 100644 --- a/lib/BloqadeCUDA/src/BloqadeCUDA.jl +++ b/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 diff --git a/lib/BloqadeCUDA/src/patch.jl b/lib/BloqadeCUDA/src/patch.jl index 413328919..c14c31fcc 100644 --- a/lib/BloqadeCUDA/src/patch.jl +++ b/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 +=# \ No newline at end of file diff --git a/lib/BloqadeCUDA/src/sample.jl b/lib/BloqadeCUDA/src/sample.jl new file mode 100644 index 000000000..b764c8108 --- /dev/null +++ b/lib/BloqadeCUDA/src/sample.jl @@ -0,0 +1,2 @@ +function sample(weights::CuVector) +end diff --git a/lib/BloqadeCUDA/test/subspace.jl b/lib/BloqadeCUDA/test/subspace.jl new file mode 100644 index 000000000..42afd0d8b --- /dev/null +++ b/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) + diff --git a/lib/BloqadeExpr/src/space.jl b/lib/BloqadeExpr/src/space.jl index 6e2812477..d3859810f 100644 --- a/lib/BloqadeExpr/src/space.jl +++ b/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