Performance issues with repeated gemv! calls in for loop
- Dominant language
- Julia
- Stars
- 586
- Forks
- 108
- Avg merge
- 1d 5h
- Merged PRs (30d)
- 44
Description
Minified example:
```julia
using Enzyme, LinearAlgebra
LinearAlgebra.BLAS.set_num_threads(1)
Enzyme.Compiler.bitcode_replacement!(false)
@inline function coupled_springs(K, m, x0, v0, T)
Ktmp = zero(K)
xtmp = zero(x0)
vtmp = zero(v0)
N = length(m)
pX = pointer(xtmp)
pY = pointer(vtmp)
sX = 1
sY = 1
pA = pointer(Ktmp)
for j in 1:5000
# @inline mul!(vtmp, Ktmp, xtmp, dt, 1)
GC.@preserve Ktmp xtmp vtmp ccall((:dgemv_64_, LinearAlgebra.BLAS.libblastrampoline), Cvoid,
(Ref{UInt8}, Ref{LinearAlgebra.BLAS.BlasInt}, Ref{LinearAlgebra.BLAS.BlasInt}, Ref{Float64},
Ptr{Float64}, Ref{LinearAlgebra.BLAS.BlasInt}, Ptr{Float64}, Ref{LinearAlgebra.BLAS.BlasInt},
Ref{Float64}, Ptr{Float64}, Ref{LinearAlgebra.BLAS.BlasInt}, Clong),
'N', N, N, 0.001,
pA, N, pX, sX,
1.0, pY, sY, 1)
end
return @inbounds xtmp[1]
end
function make_args(N)
K = rand(N, N)
K[diagind(K)] .= 0
m = 0.5 .+ 0.5 * rand(N)
x0 = randn(N)
v0 = zeros(N)
T = 1.0
return K, m, x0, v0, T
end
@inline function enzyme_inputs(K, m, x0, v0, T)
dK = zero(K)
dm = zero(m)
dx0 = zero(x0)
dv0 = zero(v0)
return Duplicated(K, dK), Duplicated(m, dm), Duplicated(x0, dx0), Duplicated(v0, dv0), Const(T)
end
function enzyme_gradient(args...)
inputs = enzyme_inputs(args...)
dK = inputs[1].dval
Enzyme.autodiff(Reverse, Const(coupled_springs), inputs...)
return dK
end
N = 200
args = make_args(N)
@time coupled_springs(args...) # compilation
@time enzyme_gradient(args...) # compilation
println("\nPrimal:")
@time coupled_springs(args...) # around 0.02 seconds
println("Enzyme:")
@time enzyme_gradient(args...) # around 1.1 seconds
```
A vtunes profiler of the enzyme gradient call reveals that most time is spent in memcpy:

On Julia 1.10-beta1 with Enzyme main.
Contributor guide
Assessment
This issue has not been assessed yet.