sum somtimes allocates during reverse pass
- Dominant language
- Julia
- Stars
- 586
- Forks
- 108
- Avg merge
- 1d 5h
- Merged PRs (30d)
- 44
Description
This may be machine-dependent, but for specific programs, I noticed that Enzyme is creating a lot of small allocations when calling into sum.
A MWE that reproduces this on one of my machines is
```julia
using Enzyme
using BenchmarkTools
function clrinv!(y, x)
y .= exp.(x)
tot = sum(y)
y ./= tot
return nothing
end
function test(x)
y = similar(x)
clrinv!(y, x)
return sum(abs2, y)
end
julia> @benchmark autodiff(Enzyme.Reverse, test, Active, Duplicated(x, dx)) setup=(;x, dx)
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
Range (min … max): 143.788 μs … 79.601 ms ┊ GC (min … max): 0.00% … 99.66%
Time (median): 148.948 μs ┊ GC (median): 0.00%
Time (mean ± σ): 165.528 μs ± 799.178 μs ┊ GC (mean ± σ): 6.02% ± 4.29%
▃▇██▇▆▅▄▃▂▁ ▂▃▃▃▃▂▁ ▂
█████████████▇█▇▆▆▅▅▄▄▅▄▅▄▅▄▅▁▁▅▁▁▁▁▁▃▁▃▃▁▁▁▁▄▇████████▇▇▆▅▅▆ █
144 μs Histogram: log(frequency) by time 228 μs <
Memory estimate: 257.09 KiB, allocs estimate: 36.
```
If I switch from `sum` to a standard for loop the allocations are greatly reduced
```julia
function sumloop(f, x)
tot = zero(eltype(x))
for i in eachindex(x)
tot += f(x[i])
end
return tot
end
sumloop(x) = sumloop(identity, x)
function clrinvloop!(y, x)
y .= exp.(x)
tot = sumloop(y)
y ./= tot
return nothing
end
function testloop(x)
y = similar(x)
clrinv_fast!(y, x)
return sumloop(abs2, y)
end
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
Range (min … max): 150.708 μs … 86.834 ms ┊ GC (min … max): 0.00% … 99.71%
Time (median): 156.758 μs ┊ GC (median): 0.00%
Time (mean ± σ): 172.895 μs ± 871.867 μs ┊ GC (mean ± σ): 6.38% ± 4.63%
▄▇██▇▆▅▄▃▂▁▁ ▂▂▂▂▂▁▁ ▂
▆████████████▇▇▆▅▆▄▅▁▃▁▃▁▁▁▅▄▅▅▃▃▃▃▁▁▁▁▁▁▁▁▃▁▁▁▁▃▆█████████▇▆ █
151 μs Histogram: log(frequency) by time 234 μs <
Memory estimate: 256.12 KiB, allocs estimate: 5.
```
When using `@profview_allocs` on the autodiff `sum` version I get

while for the loop version the only allocations the call to `similar` which I believe is basically optimal.
Here is my system info and manifest
```julia
Julia Version 1.10.5
Commit 6f3fdf7b362 (2024-08-27 14:19 UTC)
Build Info:
Official https://julialang.org/ release
Platform Info:
OS: Linux (x86_64-linux-gnu)
CPU: 32 × AMD Ryzen 9 7950X 16-Core Processor
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-15.0.7 (ORCJIT, znver3)
Threads: 1 default, 0 interactive, 1 GC (on 32 virtual cores)
Environment:
JULIA_EDITOR = code
JULIA_NUM_THREADS = 1
```
```julia
pkg> st -m
Status `~/Research/EnzymeTest/Manifest.toml`
[fa961155] CEnum v0.5.0
[7da242da] Enzyme v0.12.36 `https://github.com/EnzymeAD/Enzyme.jl.git#main`
[f151be2c] EnzymeCore v0.7.8
[e2ba6199] ExprTools v0.1.10
[61eb1bfa] GPUCompiler v0.27.5
[692b3bcd] JLLWrappers v1.6.0
[929cbde3] LLVM v9.0.0
[d8793406] ObjectFile v0.4.2
[aea7be01] PrecompileTools v1.2.1
[21216c6a] Preferences v1.4.3
[189a3867] Reexport v1.2.2
[ae029012] Requires v1.3.0
[6c6a2e73] Scratch v1.2.1
[53d494c1] StructIO v0.3.1
[a759f4b9] TimerOutputs v0.5.24
[7cc45869] Enzyme_jll v0.0.148+0
[dad2f222] LLVMExtra_jll v0.0.33+0
[0dad84c5] ArgTools v1.1.1
[56f22d72] Artifacts
[2a0f44e3] Base64
[ade2ca70] Dates
[f43a241f] Downloads v1.6.0
[7b1f6079] FileWatching
[b77e0a4c] InteractiveUtils
[4af54fe1] LazyArtifacts
[b27032c2] LibCURL v0.6.4
[76f85450] LibGit2
[8f399da3] Libdl
[37e2e46d] LinearAlgebra
[56ddb016] Logging
[d6f4376e] Markdown
[ca575930] NetworkOptions v1.2.0
[44cfe95a] Pkg v1.10.0
[de0858da] Printf
[3fa0cd96] REPL
[9a3f8284] Random
[ea8e919c] SHA v0.7.0
[9e88b42a] Serialization
[6462fe0b] Sockets
[fa267f1f] TOML v1.0.3
[a4e569a6] Tar v1.10.0
[cf7118a7] UUIDs
[4ec0a83e] Unicode
[e66e0078] CompilerSupportLibraries_jll v1.1.1+0
[deac9b47] LibCURL_jll v8.4.0+0
[e37daf67] LibGit2_jll v1.6.4+0
[29816b5a] LibSSH2_jll v1.11.0+1
[c8ffd9c3] MbedTLS_jll v2.28.2+1
[14a3606d] MozillaCACerts_jll v2023.1.10
[4536629a] OpenBLAS_jll v0.3.23+4
[83775a58] Zlib_jll v1.2.13+1
[8e850b90] libblastrampoline_jll v5.11.0+0
[8e850ede] nghttp2_jll v1.52.0+1
[3f19e933] p7zip_jll v17.4.0+2
```
Contributor guide
Assessment
This issue has not been assessed yet.