Merge f151952 into 7cdb707

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/register.jl

@@ -1,5 +1,6 @@
 using YaoBase, BitBasis
 import BitBasis: BitStr, BitStr64
+using LinearAlgebra: Transpose
 
 export ArrayReg,
     AdjointArrayReg,
@@ -143,6 +144,9 @@ function Base.copyto!(dst::AdjointArrayReg, src::AdjointArrayReg)
     return dst
 end
 
+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)
+
 # register interface
 YaoBase.nqubits(r::ArrayReg{B}) where B = log2i(length(r.state) ÷ B)
 YaoBase.nactive(r::ArrayReg) = log2dim1(r.state)
@@ -333,7 +337,14 @@ product_state(total::Int, bit_config::Integer; nbatch::Int=1) = product_state(Co
 product_state(::Type{T}, bit_str::BitStr; nbatch::Int=1) where T = ArrayReg{nbatch}(T, bit_str)
 
 function product_state(::Type{T}, total::Int, bit_config::Integer; nbatch::Int=1) where T
-    return ArrayReg{nbatch}(onehot(T, total, bit_config, nbatch))
+    if nbatch == 1
+        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
 
 """
@@ -386,7 +397,7 @@ ArrayReg{2, Complex{Float64}, Array...}
 rand_state(n::Int; nbatch::Int=1) = rand_state(ComplexF64, n; nbatch=nbatch)
 
 function rand_state(::Type{T}, n::Int; nbatch::Int=1) where T
-    raw = randn(T, 1<<n, nbatch)
+    raw = nbatch == 1 ? randn(T, 1<<n, nbatch) : transpose(randn(T, nbatch, 1<<n))
     return normalize!(ArrayReg{nbatch}(raw))
 end
 

src/utils.jl

@@ -221,3 +221,9 @@ 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