EnzymeAD / EnzymeAD/Enzyme.jl

Performance issues with repeated gemv! calls in for loop

Open
#1,004 2 comments 0 reactions 0 assignees View on GitHub
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:
![image](https://github.com/EnzymeAD/Enzyme.jl/assets/43899491/b3dd51d8-6f39-4889-8f7d-20551dc57cf5)
On Julia 1.10-beta1 with Enzyme main.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.