Performance issue

[ghost]

Hi, i am using the very good pkg Yao.jl and find something weird.

i try to implement the shift rule method to get the gradient of a circuit. I find the function is unexpected slow.

here are my code

using Yao

function gradient_rot(ham::AbstractBlock, state::AbstractRegister, cir::AbstractBlock, k::Int)
    num_qubits = nqubits(cir) 
    
    dispatch!(+, cir.blocks[k], pi/2)
    pos = expect(ham, copy(state)=>cir)
    
    dispatch!(-, cir.blocks[k], pi)
    neg = expect(ham, copy(state)=>cir)
    
    dispatch!(+, cir.blocks[k], pi/2)
    
    return (pos - neg) / 2 |> real
end

function gradient(ham::AbstractBlock, state::AbstractRegister, cir::AbstractBlock)
    grads = zeros(Float64, length(cir.blocks))
    for k = 1:length(cir.blocks)
        grads[k] = gradient_rot(ham, state, cir, k)
    end
    
    return grads
end

c = chain(1, put(1, 1=>Rx(0.23)))

ham = kron(Z)
state = rand_state(nqubits(c))
dispatch!(c, :random)

@time grads1 = expect'(ham, copy(state)=>c)[2];
@time grads2 = gradient(ham, copy(state), c);

@time begin
    grads3 = zeros(Float64, length(c.blocks))
    for k = 1:length(c.blocks)
        grads3[k] = gradient_rot(ham, state, c, k)
    end
end

in my devices, the results are

0.000037 seconds (50 allocations: 1.953 KiB)
0.037294 seconds (41.22 k allocations: 2.300 MiB, 99.82% compilation time)
0.000029 seconds (59 allocations: 2.203 KiB)

i noticed that the gradient function is unexpected slow. however, when implemented in the main code (like grads3), the speed is much faster.

I do not know How to optimize this function, could someone help me?

[Roger-luo]

Hi I think from the timing it actually prints that 99% of the time is actually compilation time, in Julia if you want to measure the actual timing in your loop you should either use the BenchmarkTools package or measure the timing multiple times.

[ghost]

Hi I think from the timing it actually prints that 99% of the time is actually compilation time, in Julia if you want to measure the actual timing in your loop you should either use the BenchmarkTools package or measure the timing multiple times.

thanks for your reply.

yeah, i also use the BenchmarkTools package, the outputs are more reasonable (shift rule method is several times slower than AD).

in my perspective, the compilation only happens at the first time i run this function. but, i already run the gradient function several times, and use @time to measure the time cost. i do not know why there is still 99.82% compilation time.

[Roger-luo]

yeah I think somehow the second call is recompiling the method again inside the script, so @time does not give you the actual timing. It is always preferred to use @btime or @benchmark, on my laptop, shift rule has similar timing with reverse mode, which is expected since there is only one parameter with one gate in your example.

I'm not sure why the compile result is not cached here, if you wish to investigate deeper https://github.com/JuliaDebug/Cthulhu.jl is the tool you want to use to find out method invalidations. But it probably not worth the effort since after compiled this function is quite fast.