Merge 0d5b5d9 into 77fea11

benchmark/benchmarks.jl

@@ -10,6 +10,9 @@ using LinearAlgebra, SparseArrays
 
 bench(n, U, loc::Tuple) = @benchmarkable instruct!(st, $U, $loc) setup=(st=statevec(rand_state($n)))
 bench(n, U, loc::Tuple, control_locs::Tuple, control_bits::Tuple) = @benchmarkable instruct!(st, $U, $loc, $control_locs, $control_bits) setup=(st=statevec(rand_state($n)))
+bench(n, B::Int, U, loc::Tuple) = @benchmarkable instruct!(st, $U, $loc) setup=(st=statevec(rand_state($n, nbatch=$B)))
+bench(n, B::Int, U, loc::Tuple, control_locs::Tuple, control_bits::Tuple) = @benchmarkable instruct!(st, $U, $loc, $control_locs, $control_bits) setup=(st=statevec(rand_state($n, nbatch=$B)))
+
 
 const SUITE = BenchmarkGroup()
 SUITE["specialized"] = BenchmarkGroup()
@@ -135,3 +138,131 @@ for T in [ComplexF64], n in 1:2:10, m in 3:2:6, U in matrices(T, 1<<n)
     perms = randperm(n+m)
     SUITE["matrices"]["controlled"]["random"][n, m, string(T), string(typeof(U))] = bench(n+m, U, Tuple(perms[1:n]), Tuple(perms[n+1:n+m]), ntuple(x->1, m))
 end
+
+
+# batched
+SUITE["batched specialized"] = BenchmarkGroup()
+
+@info "generating benchmark for batched specialized operators for single qubits"
+# Specialized Gate Instruction
+SUITE["batched specialized"]["single qubit"] = BenchmarkGroup()
+## single qubit benchmark
+for U in YaoArrayRegister.SPECIALIZATION_LIST, n in 5:4:15, B in 10:20:60
+    SUITE["batched specialized"]["single qubit"][string(U), n, B] = bench(n, B, Val(U), (1, ))
+end
+
+SUITE["batched specialized"]["single control"] = BenchmarkGroup()
+for U in YaoArrayRegister.SPECIALIZATION_LIST, n in 5:4:15, B in 10:20:60
+    SUITE["batched specialized"]["single control"][string(U), n, B, (2, ), (1, )] =
+        bench(n, B, Val(U), (1, ), (2, ), (1, ))
+end
+
+SUITE["batched specialized"]["multi control"] = BenchmarkGroup()
+
+for U in YaoArrayRegister.SPECIALIZATION_LIST, n in 5:4:15, B in 10:20:60
+    control_locs = Tuple(2:n); control_bits = ntuple(x->1, n-1)
+    SUITE["batched specialized"]["multi control"][string(U), n, B, 2:n, control_locs] =
+        bench(n, B, Val(U), (1, ), control_locs, control_bits)
+end
+
+SUITE["batched specialized"]["multi qubit"] = BenchmarkGroup()
+const location_sparsity = 0.4
+@info "generating benchmark for specialized operators for multi qubits"
+## multi qubit benchmark
+for U in YaoArrayRegister.SPECIALIZATION_LIST, n in 5:4:15, B in 10:20:60
+    perms = randperm(n)[1:ceil(Int, location_sparsity * n)]
+    SUITE["batched specialized"]["multi qubit"][string(U), n, B] = bench(n, B, Val(U), Tuple(perms))
+end
+
+SUITE["batched specialized"]["multi qubit multi control"] = BenchmarkGroup()
+SUITE["batched specialized"]["single qubit multi control"] = BenchmarkGroup()
+
+const control_rate = 0.3
+for U in YaoArrayRegister.SPECIALIZATION_LIST, n in 5:4:15, B in 10:20:60
+    num_controls = ceil(Int, n * control_rate)
+    perms = randperm(n)
+    control_locs = Tuple(perms[1:num_controls]); control_bits = ntuple(x->rand(0:1), num_controls)
+    perms = perms[num_controls+1:num_controls+round(Int, location_sparsity * n)]
+
+    SUITE["batched specialized"]["multi qubit multi control"][string(U), n, B, num_controls] = bench(n, B, Val(U), Tuple(perms), control_locs, control_bits)
+    SUITE["batched specialized"]["single qubit multi control"][string(U), n, B, num_controls] = bench(n, B, Val(U), (perms[1], ), control_locs, control_bits)
+end
+
+for n in 5:4:15, B in 10:20:60
+    SUITE["batched specialized"]["multi qubit"]["SWAP", n, B] = bench(n, B, Val(:SWAP), (1, 2))
+    SUITE["batched specialized"]["multi qubit"]["SWAP", "random", n, B] = bench(n, B, Val(:SWAP), Tuple(randperm(n)[1:2]))
+end
+
+# General Instructions (matrices based)
+SUITE["batched matrices"] = BenchmarkGroup()
+SUITE["batched matrices"]["contiguous"] = BenchmarkGroup()
+SUITE["batched matrices"]["contiguous"]["ordered"] = BenchmarkGroup()
+SUITE["batched matrices"]["contiguous"]["random"] = BenchmarkGroup()
+
+SUITE["batched matrices"]["in-contiguous"] = BenchmarkGroup()
+SUITE["batched matrices"]["in-contiguous"]["ordered"] = BenchmarkGroup()
+SUITE["batched matrices"]["in-contiguous"]["random"] = BenchmarkGroup()
+
+SUITE["batched matrices"]["single qubit"] = BenchmarkGroup()
+SUITE["batched matrices"]["single qubit"]["ordered"] = BenchmarkGroup()
+SUITE["batched matrices"]["single qubit"]["random"] = BenchmarkGroup()
+
+
+## General Matrix Instruction
+function matrices(::Type{T}, N) where T
+    list = Any[
+        rand_unitary(T, N), # dense matrices
+        # SparseMatrixCSC
+        sprand_hermitian(T, N, 0.1),
+        # PermMatrix
+        pmrand(T, N),
+        Diagonal(rand(T, N))
+    ]
+    if N < 100
+        # StaticArrays
+        push!(list, @SArray(rand(T, N, N)))
+        push!(list, @MArray(rand(T, N, N)))
+    end
+    return list
+end
+
+# default test type is ComplexF64
+matrices(N) = matrices(ComplexF64, N)
+
+@info "generating benchmark for batched contiguous matrices locs"
+### contiguous
+for n in 1:2:5, T in [ComplexF64], U in matrices(T, 1<<n), B in 10:20:60
+    # contiguous ordered address
+    SUITE["batched matrices"]["contiguous"]["ordered"][n, B, string(T), string(typeof(U))] = bench(n, B, U, Tuple(1:n))
+    # contiguous random address
+    SUITE["batched matrices"]["contiguous"]["random"][n, B, string(T), string(typeof(U))] = bench(n, B, U, Tuple(randperm(n)))
+end
+
+@info "generating benchmark for in-contiguous matrices locs"
+### in-contiguous
+for m in 1:3, T in [ComplexF64], U in matrices(T, 1 << m), B in 10:20:60
+    n = 10; N = 1 << n
+    # in-contiguous ordered address
+    SUITE["batched matrices"]["in-contiguous"]["ordered"][m, B, string(T), string(typeof(U))] = bench(n, B, U, Tuple(sort(randperm(n)[1:m])))
+    # in-contiguous random address
+    SUITE["batched matrices"]["in-contiguous"]["random"][m, B, string(T), string(typeof(U))] = bench(n, B, U, Tuple(randperm(n)[1:m]))
+end
+
+@info "generating benchmark for single qubit matrices"
+### single qubit
+for T in [ComplexF64], U in matrices(T, 2), n in 1:4:25, B in 10:20:60
+    SUITE["batched matrices"]["single qubit"]["ordered"][string(T), B, string(typeof(U)), n] = bench(n, B, U, (rand(1:n), ))
+    SUITE["batched matrices"]["single qubit"]["random"][string(T), B, string(typeof(U)), n] = bench(n, B, U, (rand(1:n), ))
+end
+
+SUITE["batched matrices"]["controlled"] = BenchmarkGroup()
+SUITE["batched matrices"]["controlled"]["ordered"] = BenchmarkGroup()
+SUITE["batched matrices"]["controlled"]["random"] = BenchmarkGroup()
+
+test_bench(n, U, loc, control_locs, control_bits) = instruct!(statevec(rand_state(n)), U, loc, control_locs, control_bits)
+for T in [ComplexF64], n in 1:2:5, m in 3:2:6, U in matrices(T, 1<<n), B in 10:20:60
+    SUITE["batched matrices"]["controlled"]["ordered"][n, B, m, string(T), string(typeof(U))] = bench(n+m, B, U, Tuple(1:n), Tuple(n+1:n+m), ntuple(x->1, m))
+
+    perms = randperm(n+m)
+    SUITE["batched matrices"]["controlled"]["random"][n, B, m, string(T), string(typeof(U))] = bench(n+m, B, U, Tuple(perms[1:n]), Tuple(perms[n+1:n+m]), ntuple(x->1, m))
+end

benchmark/runbench.jl

@@ -1,6 +1,6 @@
 using PkgBenchmark
 
-current = BenchmarkConfig(id="multithreading", env = Dict("JULIA_NUM_THREADS"=>4), juliacmd=`julia -O3`)
-baseline = BenchmarkConfig(id="master", env = Dict("JULIA_NUM_THREADS"=>1), juliacmd=`julia -O3`)
+current = BenchmarkConfig(id="transpose_storage", juliacmd=`julia -O3`)
+baseline = BenchmarkConfig(id="master", juliacmd=`julia -O3`)
 results = judge("YaoArrayRegister", current, baseline)
 export_markdown("report.md", results)

src/instruct.jl

@@ -438,7 +438,7 @@ function YaoBase.instruct!(
     a = T(cos(theta/2))
     c = T(-im * sin(theta/2))
     e = T(exp(-im/2*theta))
-    for b in itercontrol(log2i(size(state, 1)), [control_locs...], [control_bits...])
+    for b in itercontrol(log2i(size(state, 1)), Int[control_locs...], Int[control_bits...])
         if b&mask1==0
             i = b+1
             i_ = b ⊻ mask12 + 1

src/register.jl

@@ -4,6 +4,7 @@ import BitBasis: BitStr, BitStr64
 export ArrayReg,
     AdjointArrayReg,
     ArrayRegOrAdjointArrayReg,
+    transpose_storage,
     # YaoBase
     nqubits,
     nactive,
@@ -129,6 +130,9 @@ to `copy`.
 """
 ArrayReg(r::ArrayReg{B}) where B = ArrayReg{B}(copy(r.state))
 
+transpose_storage(reg::ArrayReg{B,T,<:Transpose}) where {B,T} = ArrayReg{B}(copy(reg.state))
+transpose_storage(reg::ArrayReg{B,T}) where {B,T} = ArrayReg{B}(transpose(copy(transpose(reg.state))))
+
 Base.copy(r::ArrayReg) = ArrayReg(r)
 Base.similar(r::ArrayRegOrAdjointArrayReg{B}) where B = ArrayReg{B}(similar(state(r)))
 
@@ -302,7 +306,7 @@ ArrayReg{2, Complex{Float32}, Array...}
 product_state(bit_str::BitStr; nbatch::Int=1) = product_state(ComplexF64, bit_str; nbatch=nbatch)
 
 """
-    product_state([T=ComplexF64], total::Int, bit_config::Integer; nbatch=1)
+    product_state([T=ComplexF64], total::Int, bit_config::Integer; nbatch=1, no_transpose_storage=false)
 
 Create an [`ArrayReg`](@ref) with bit configuration `bit_config`, total number of bits `total`.
 See also [`zero_state`](@ref), [`rand_state`](@ref), [`uniform_state`](@ref).
@@ -328,12 +332,19 @@ ArrayReg{1, Complex{Float32}, Array...}
     This interface will not check whether the number of required digits
     for the bit configuration matches the total number of bits.
 """
-product_state(total::Int, bit_config::Integer; nbatch::Int=1) = product_state(ComplexF64, total, bit_config; nbatch=nbatch)
+product_state(total::Int, bit_config::Integer; kwargs...) = product_state(ComplexF64, total, bit_config; kwargs...)
 
-product_state(::Type{T}, bit_str::BitStr; nbatch::Int=1) where T = ArrayReg{nbatch}(T, bit_str)
+product_state(::Type{T}, bit_str::BitStr{N}; kwargs...) where {T,N} = product_state(T, N, buffer(bit_str); kwargs...)
 
-function product_state(::Type{T}, total::Int, bit_config::Integer; nbatch::Int=1) where T
-    return ArrayReg{nbatch}(onehot(T, total, bit_config, nbatch))
+function product_state(::Type{T}, total::Int, bit_config::Integer; nbatch::Int=1, no_transpose_storage::Bool=false) where T
+    if nbatch == 1 || no_transpose_storage
+        raw = onehot(T, total, bit_config, nbatch)
+    else
+        raw = zeros(T, nbatch, 1<<total)
+        raw[:,Int(bit_config)+1] .= 1
+        raw = transpose(raw)
+    end
+    return ArrayReg{nbatch}(raw)
 end
 
 """
@@ -358,12 +369,12 @@ ArrayReg{3, Complex{Float32}, Array...}
     active qubits: 4/4
 ```
 """
-zero_state(n::Int; nbatch::Int=1) = zero_state(ComplexF64, n; nbatch=nbatch)
-zero_state(::Type{T}, n::Int; nbatch::Int=1) where T = product_state(T, n, 0; nbatch=nbatch)
+zero_state(n::Int; kwargs...) = zero_state(ComplexF64, n; kwargs...)
+zero_state(::Type{T}, n::Int; kwargs...) where T = product_state(T, n, 0; kwargs...)
 
 
 """
-    rand_state([T=ComplexF64], n::Int; nbatch::Int=1)
+    rand_state([T=ComplexF64], n::Int; nbatch=1, no_transpose_storage=false)
 
 Create a random [`ArrayReg`](@ref) with total number of qubits `n`.
 
@@ -383,15 +394,15 @@ ArrayReg{2, Complex{Float64}, Array...}
     active qubits: 4/4
 ```
 """
-rand_state(n::Int; nbatch::Int=1) = rand_state(ComplexF64, n; nbatch=nbatch)
+rand_state(n::Int; kwargs...) = rand_state(ComplexF64, n; kwargs...)
 
-function rand_state(::Type{T}, n::Int; nbatch::Int=1) where T
-    raw = randn(T, 1<<n, nbatch)
+function rand_state(::Type{T}, n::Int; nbatch::Int=1, no_transpose_storage::Bool=false) where T
+    raw = nbatch == 1 || no_transpose_storage ? randn(T, 1<<n, nbatch) : transpose(randn(T, nbatch, 1<<n))
     return normalize!(ArrayReg{nbatch}(raw))
 end
 
 """
-    uniform_state([T=ComplexF64], n; nbatch=1)
+    uniform_state([T=ComplexF64], n; nbatch=1, no_transpose_storage=false)
 
 Create a uniform state: ``\\frac{1}{2^n} \\sum_k |k⟩``. This state
 can also be created by applying [`H`](@ref) (Hadmard gate) on ``|00⋯00⟩`` state.
@@ -408,8 +419,11 @@ ArrayReg{2, Complex{Float64}, Array...}
     active qubits: 4/4
 ```
 """
-uniform_state(n::Int; nbatch::Int=1) = uniform_state(ComplexF64, n; nbatch=nbatch)
-uniform_state(::Type{T}, n::Int; nbatch::Int=1) where T = ArrayReg{nbatch}(ones(T, 1<<n, nbatch) ./ sqrt(1<<n))
+uniform_state(n::Int; kwargs...) = uniform_state(ComplexF64, n; kwargs...)
+function uniform_state(::Type{T}, n::Int; nbatch::Int=1, no_transpose_storage::Bool=false) where T
+    raw = nbatch == 1 || no_transpose_storage ? ones(T, 1<<n, nbatch) : transpose(ones(T, nbatch, 1<<n))
+    normalize!(ArrayReg{nbatch}(raw))
+end
 
 """
     oneto(r::ArrayReg, n::Int=nqubits(r))

src/utils.jl

@@ -221,3 +221,14 @@ end
     end
     state
 end
+
+#### Yao Base patch ####
+using YaoBase
+function YaoBase.batched_kron!(C::Array{T, 3}, A::AbstractArray{T1, 3}, B::AbstractArray{T2, 3}) where {T, T1, T2}
+    YaoBase.batched_kron!(C, convert(Array{T,3}, A), convert(Array{T,3}, B))
+end
+
+using LinearAlgebra: Transpose
+Base.convert(::Type{Transpose{T, Matrix{T}}}, arr::AbstractMatrix{T}) where T = transpose(Matrix(transpose(arr)))
+Base.convert(t::Type{Transpose{T, Matrix{T}}}, arr::Transpose{T}) where T = invoke(convert, Tuple{Type{Transpose{T, Matrix{T}}}, Transpose}, t, arr)
+

test/instruct.jl

@@ -22,6 +22,10 @@ using YaoBase.Const
         instruct!(instruct!(copy(ST), U1, 3), U1, 1)
     @test instruct!(copy(REG), kron(U1, U1), (3, 1)) ≈
         instruct!(instruct!(copy(REG), U1, 3), U1, 1)
+    @test instruct!(transpose_storage(REG), kron(U1, U1), (3, 1)) ≈
+        instruct!(instruct!(copy(REG), U1, 3), U1, 1)
+    @test instruct!(transpose_storage(REG), kron(U1, U1), (3, 1)) ≈
+        instruct!(instruct!(transpose_storage(REG), U1, 3), U1, 1)
 
     @test instruct!(reshape(copy(ST), :, 1), kron(U1, U1), (3, 1)) ≈
         instruct!(instruct!(reshape(copy(ST), :, 1), U1, 3), U1, 1)

test/register.jl

@@ -15,11 +15,23 @@ using Test, YaoArrayRegister, BitBasis, LinearAlgebra
 
     st = rand(ComplexF64, 4, 6)
     @test state(adjoint(ArrayReg(st))) == adjoint(st)
+
+    reg = ArrayReg{6}(st)
+    regt = transpose_storage(reg)
+    @test regt.state isa Transpose
+    @test regt == reg
+    @test transpose_storage(regt).state isa Matrix
+    @test transpose_storage(regt) == reg
 end
 
 @testset "test $T initialization methods" for T in [ComplexF64, ComplexF32, ComplexF16]
     @testset "test product state" begin
+        st = state(product_state(T, bit"100"; nbatch=1))
+        @test !(st isa Transpose)
+        st = state(product_state(T, bit"100"; nbatch=2, no_transpose_storage=true))
+        @test !(st isa Transpose)
         st = state(product_state(T, bit"100"; nbatch=2))
+        @test st isa Transpose
         for k in 1:2
             @test st[:, k] ≈ onehot(T, bit"100")
         end
@@ -31,22 +43,37 @@ end
         @test eltype(product_state(Float64, 4, 0).state) == Float64
     end
     @testset "test zero state" begin
+        st = state(zero_state(T, 3; nbatch=1))
+        @test !(st isa Transpose)
+        st = state(zero_state(T, 3; nbatch=2, no_transpose_storage=true))
+        @test !(st isa Transpose)
         st = state(zero_state(T, 4; nbatch=4))
+        @test st isa Transpose
         for k in 1:4
             @test st[:, k] ≈ onehot(T, 4, 0)
         end
         @test eltype(zero_state(Float64, 4).state) == Float64
     end
     @testset "test rand state" begin
+        st = state(rand_state(T, 3; nbatch=1))
+        @test !(st isa Transpose)
+        st = state(rand_state(T, 3; nbatch=2, no_transpose_storage=true))
+        @test !(st isa Transpose)
         # NOTE: we only check if the state is normalized
         st = state(rand_state(T, 4, nbatch=2))
+        @test st isa Transpose
         for k in 1:2
             @test norm(st[:, k]) ≈ 1.0
         end
         @test eltype(rand_state(Float64, 4).state) == Float64
     end
     @testset "test uniform state" begin
+        st = state(uniform_state(T, 3; nbatch=1))
+        @test !(st isa Transpose)
+        st = state(uniform_state(T, 3; nbatch=2, no_transpose_storage=true))
+        @test !(st isa Transpose)
         st = state(uniform_state(T, 4; nbatch=2))
+        @test st isa Transpose
         for k in 1:2
             for each in st[:, k]
                 @test each ≈ 1/sqrt(16)