Allocations and time on compile are ever-increasing

[mofeing]

Recently I found this strange behavior

julia> @btime Reactant.with_debug() do
           @jit Ops.abs(x)
       end
  37.419 ms (147556 allocations: 6.07 MiB)
2-element ConcreteRArray{Int64, 1}:
 1
 1

julia> @btime @jit Ops.abs(x)
  58.990 ms (297529 allocations: 12.08 MiB)
2-element ConcreteRArray{Int64, 1}:
 1
 1

julia> @btime @jit Ops.abs(x)
  74.636 ms (411819 allocations: 16.63 MiB)
2-element ConcreteRArray{Int64, 1}:
 1
 1

julia> @btime Reactant.with_debug() do
           @jit Ops.abs(x)
       end
  86.295 ms (505732 allocations: 20.37 MiB)
2-element ConcreteRArray{Int64, 1}:
 1
 1

julia> @btime Reactant.with_debug() do
           @jit Ops.abs(x)
       end
  97.753 ms (586540 allocations: 23.59 MiB)
2-element ConcreteRArray{Int64, 1}:
 1
 1

julia> @btime Reactant.with_debug() do
           @compile Ops.abs(x)
       end
  100.258 ms (652550 allocations: 26.15 MiB)

ulia> @btime @compile Ops.abs(x)
  111.832 ms (721709 allocations: 28.93 MiB)
Reactant.Compiler.Thunk{Symbol("##abs_reactant#83526")}()

Why are allocations (and time, but this looks like a consequence) ever increasing on each call to compile? For sure we are not freeing the previously compiled functions who live in the Thunks but that shouldn't affect?

[jumerckx]

I can't reproduce allocations, even though compile time does seem to increase slowly.

using Reactant, BenchmarkTools

x = ConcreteRArray(rand(1024))

for _ in 1:100
     @btime Reactant.Ops.with_debug() do
          @jit Reactant.Ops.abs(x)
     end
end

  86.708 ms (19863 allocations: 996.98 KiB)
  90.069 ms (19863 allocations: 996.98 KiB)
  93.267 ms (19863 allocations: 996.98 KiB)
  91.644 ms (19863 allocations: 996.98 KiB)
  95.635 ms (19863 allocations: 996.98 KiB)
  103.573 ms (19863 allocations: 996.98 KiB)
  100.355 ms (19863 allocations: 996.98 KiB)
  106.399 ms (19863 allocations: 996.98 KiB)
  103.719 ms (19863 allocations: 996.98 KiB)
  111.025 ms (19863 allocations: 996.98 KiB)
  94.440 ms (19863 allocations: 996.98 KiB)
  97.324 ms (19863 allocations: 996.98 KiB)
  104.927 ms (19863 allocations: 996.98 KiB)
  112.291 ms (19863 allocations: 996.98 KiB)

Julia Version 1.11.1
Commit 8f5b7ca12ad (2024-10-16 10:53 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 12 × Intel(R) Core(TM) i7-9750H CPU @ 2.60GHz
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, skylake)
Threads: 1 default, 0 interactive, 1 GC (on 12 virtual cores)

[glou-nes]

Got similar result as @mofeing using @jumerckx code.

  91.916 ms (69670 allocations: 2.96 MiB)
  66.066 ms (148296 allocations: 6.10 MiB)
  77.496 ms (231286 allocations: 9.41 MiB)
  87.358 ms (299718 allocations: 12.15 MiB)
  97.592 ms (357958 allocations: 14.47 MiB)
  108.565 ms (416932 allocations: 16.84 MiB)
  116.130 ms (466430 allocations: 18.80 MiB)
  125.490 ms (512294 allocations: 20.63 MiB)
  133.661 ms (556702 allocations: 22.40 MiB)
  140.108 ms (603294 allocations: 24.26 MiB)
  146.170 ms (641878 allocations: 25.80 MiB)
  248.894 ms (679006 allocations: 27.28 MiB)
  166.485 ms (714678 allocations: 28.71 MiB)

Julia Version 1.11.1
Build Info:

    Note: This is an unofficial build, please report bugs to the project
    responsible for this build and not to the Julia project unless you can
    reproduce the issue using official builds available at https://julialang.org/downloads

Platform Info:
  OS: Linux (x86_64-pc-linux-gnu)
  CPU: 24 × 12th Gen Intel(R) Core(TM) i7-12850HX
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, alderlake)
Threads: 1 default, 0 interactive, 1 GC (on 24 virtual cores)
Environment:
  JULIA_DEPOT_PATH = ~/depot

[mofeing]

Allocations might be platform dependent?

I'm using macOS aarch64 (M1 pro)

[jumerckx]

Allocations might be platform dependent?

Gaëtan and I are both on Linux (x86_64-pc-linux-gnu) it seems. (For me this is actually WSL, not sure about Gaëtan)

[wsmoses]

x/ref EnzymeAD/Enzyme.jl#2152

[glou-nes]

I'm using EndeavourOS. I tracked allocations using julia --track-allocations. After several hundred calls to @jit. I got this mem file:

./EnzymeCore/SypLX/src/rules.jl.115739.mem
     2112         return TypeVar(gensym(), Annotation{T})
     5312     tt = map(T->_annotate(Base.rewrap_unionall(T, TT0)), TT.parameters[2:end])
     1088     return ft, tt
     7808     TT = Tuple{<:FwdConfig, <:Annotation{ft}, Type{<:Annotation}, tt...}
     7808     TT = Tuple{<:RevConfig, <:Annotation{ft}, Type{<:Annotation}, tt...}
     9104     sig = Base.signature_type(f, tt)
     4896     result = Core.Compiler.findall(sig, method_table; limit=-1)
  3142144                 add_mt_backedge!(cspec, mt, sig)

[wsmoses]

cc @vchuravy

[wsmoses]

can you check if this persists with the recently released Enzyme.jl 0.13.19?

[avik-pal]

Seems to persist with 0.13.19

  100.544 ms (56143 allocations: 2.43 MiB)
  109.824 ms (114179 allocations: 4.78 MiB)
  124.340 ms (167839 allocations: 6.89 MiB)
  132.645 ms (213391 allocations: 8.71 MiB)
  141.765 ms (257701 allocations: 10.49 MiB)
  155.150 ms (300755 allocations: 12.23 MiB)
  163.322 ms (342559 allocations: 13.86 MiB)

[wsmoses]

this is fixed locally with #326 and latest Enzyme. Note that this still indicates Enzyme.jl's interpreter rule stuff is being sad, but we can disable rules

julia> for _ in 1:100
            @btime Reactant.with_debug() do
                 @jit Reactant.Ops.abs(x)
            end
       end
  227.179 ms (14440 allocations: 1012.28 KiB)
  227.738 ms (14440 allocations: 1012.28 KiB)
  227.988 ms (14440 allocations: 1012.28 KiB)
  229.066 ms (14440 allocations: 1012.28 KiB)
  229.683 ms (14440 allocations: 1012.28 KiB)
  231.294 ms (14440 allocations: 1012.28 KiB)
  232.302 ms (14440 allocations: 1012.28 KiB)
  233.247 ms (14440 allocations: 1012.28 KiB)
  233.694 ms (14440 allocations: 1012.28 KiB)
  234.768 ms (14440 allocations: 1012.28 KiB)
  235.695 ms (14440 allocations: 1012.28 KiB)
  237.790 ms (14440 allocations: 1012.28 KiB)
  238.352 ms (14440 allocations: 1012.28 KiB)
  239.087 ms (14440 allocations: 1012.28 KiB)
  240.611 ms (14440 allocations: 1012.28 KiB)
  241.681 ms (14440 allocations: 1012.28 KiB)
  241.488 ms (14440 allocations: 1012.28 KiB)
  242.259 ms (14440 allocations: 1012.28 KiB)
  242.878 ms (14440 allocations: 1012.28 KiB)
  243.833 ms (14440 allocations: 1012.28 KiB)

[wsmoses]

time does climb a bit tho weirdly

[jumerckx]

With the latest interpreter stuff, increasing times (not allocations) are still a thing.

julia> for _ in 1:100
            @btime Reactant.Ops.with_debug() do
                 @jit Reactant.Ops.abs(x)
            end
       end
  62.994 ms (10989 allocations: 535.89 KiB)
  64.359 ms (10989 allocations: 535.89 KiB)
  66.648 ms (10989 allocations: 535.89 KiB)
  69.255 ms (10989 allocations: 535.89 KiB)
  70.272 ms (10989 allocations: 535.89 KiB)
  70.745 ms (10989 allocations: 535.89 KiB)
  71.283 ms (10989 allocations: 535.89 KiB)
  72.623 ms (10989 allocations: 535.89 KiB)
  73.641 ms (10989 allocations: 535.89 KiB)
  76.192 ms (10989 allocations: 535.89 KiB)
  77.077 ms (10989 allocations: 535.89 KiB)
  78.479 ms (10989 allocations: 535.89 KiB)
  80.126 ms (10989 allocations: 535.89 KiB)

[mofeing]

yeah, but the scaling is much lower

julia> for _ in 1:10
           @btime @jit Ops.abs(x)
       end
  15.026 ms (10978 allocations: 535.72 KiB)
  15.111 ms (10978 allocations: 535.72 KiB)
  16.605 ms (10978 allocations: 535.72 KiB)
  16.844 ms (10978 allocations: 535.72 KiB)
  18.649 ms (10978 allocations: 535.72 KiB)
  19.280 ms (10978 allocations: 535.72 KiB)
  19.619 ms (10978 allocations: 535.72 KiB)
  19.835 ms (10978 allocations: 535.72 KiB)
  20.569 ms (10978 allocations: 535.72 KiB)
  22.509 ms (10978 allocations: 535.72 KiB)

julia> for _ in 1:10
           @btime Ops.with_debug() do
               @jit Ops.abs(x)
           end
       end
  25.965 ms (11814 allocations: 583.87 KiB)
  26.740 ms (11814 allocations: 583.87 KiB)
  27.133 ms (11814 allocations: 583.87 KiB)
  27.839 ms (11814 allocations: 583.87 KiB)
  28.385 ms (11814 allocations: 583.87 KiB)
  29.134 ms (11814 allocations: 583.87 KiB)
  29.415 ms (11814 allocations: 583.87 KiB)
  30.080 ms (11814 allocations: 583.87 KiB)
  30.586 ms (11814 allocations: 583.87 KiB)
  31.217 ms (11814 allocations: 583.87 KiB)