Performance issue with complicated loops in function

[zhangcx93]

Describe the bug
When I'm running my script, I found that the first time run of the main function containing for loop is unreasonably slow.
By putting the code outside of the main function to the global makes the first step in loop much faster and the following step the same as warmed up code.

To reproduce

The Minimal Working Example (MWE) for this bug:

using CUDA
using Statistics

function some_func(x, y)
    a = CUDA.rand(1000)
    b = CUDA.rand(1000)
    c = a .* b
    return c .* x .* y
end

function main()
    result = []
    count = 0
    for i in 1:3
        t = time()
        for j in 0:20
            for k in 1:20
                count += 1
                for m in 1:20
                    a = CUDA.rand(1000)
                    b = CUDA.rand(1000)
                    c = some_func(a, b)
                    push!(result, mean(c))
                end
            end
        end
        println(time() - t)
    end
end

main()

The printed timing is
First time:

42.087000131607056
33.35199999809265
34.365999937057495

Second time:

1.1009998321533203
0.8480000495910645
1.4190001487731934

Without main, just global loop

10.562000036239624
0.935999870300293
1.316999912261963

Expected behavior

The top loop in main should be around 10s for the first step and around 1s for the following, but not 42s for the first step and around 33s for the following step. 42s and 33s are too much for compilation here.

Version info

Details on Julia:

Julia Version 1.6.1
Commit 6aaedecc44 (2021-04-23 05:59 UTC)
Platform Info:
  OS: Windows (x86_64-w64-mingw32)
  CPU: Intel(R) Core(TM) i9-9900K CPU @ 3.60GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-11.0.1 (ORCJIT, skylake)
Environment:
  JULIA_GPU_ALLOWSCALAR = false

Details on CUDA:

CUDA toolkit 11.3.1, artifact installation
CUDA driver 11.3.0
NVIDIA driver 466.77.0

Libraries:
- CUBLAS: 11.5.1
- CURAND: 10.2.4
- CUFFT: 10.4.2
- CUSOLVER: 11.1.2
- CUSPARSE: 11.6.0
- CUPTI: 14.0.0
- NVML: 11.0.0+466.77
- CUDNN: 8.20.0 (for CUDA 11.3.0)
- CUTENSOR: 1.3.0 (for CUDA 11.2.0)

Toolchain:
- Julia: 1.6.1
- LLVM: 11.0.1
- PTX ISA support: 3.2, 4.0, 4.1, 4.2, 4.3, 5.0, 6.0, 6.1, 6.3, 6.4, 6.5, 7.0
- Device capability support: sm_35, sm_37, sm_50, sm_52, sm_53, sm_60, sm_61, sm_62, sm_70, sm_72, sm_75, sm_80

1 device:
  0: NVIDIA GeForce RTX 2080 Ti (sm_75, 8.582 GiB / 11.000 GiB available)

[zhangcx93]

The CPU version of the same code's timing is:
First time

0.2480001449584961
0.21000003814697266
0.20399999618530273

Second time

0.20300006866455078
0.20399999618530273
0.21299982070922852

This indicates there might be something wrong with loop with CUDA... I think these code are GPU friendly and the size is not too small...

[maleadt]

That's because Julia interprets code when it's at the top level, whereas if you put it in a function you 'force' compilation.

[zhangcx93]

Well, if you try to rewrite the code in CPU version, you'll find that even been put into a function, as "force compilation mode", the second and following step is much faster than the first time(same as what you'll get from @Btime), while the first step is doing heavy compilation.
And if the top level code is pure interpreted, how the second and following step the same speed as compiled code and only the first time in loop is slow?

[zhangcx93]

Also the compilation time should not scale with the length of the loop right?

[stillyslalom]

There does appear to be a regression here (xref from Discourse).

On CUDA.jl v3.1.0,

julia> main() # first run
29.039000034332275
1.7300000190734863
1.61899995803833

On v3.3.0,

julia> main() # first run
80.84599995613098
61.09299993515015
63.625

julia> main() # second run
1.5520000457763672
1.7239999771118164
1.375999927520752

[maleadt]

Yeah, something seems off, even though this also seems compilation heuristic-related.

[zhangcx93]

using CUDA

function some_func(x, y)
    a = CUDA.rand(1000)
    b = CUDA.rand(1000)
    c = a .* b
    return c .* x .* y
end

function main()
    N = 20
    a = CUDA.rand(1000)
    b = CUDA.rand(1000)
    for i in 1:N
        for j in 0:800
            c = some_func(a, b)
        end
    end
end

@time main()
@time main()

The compilation time should approximately be the diff between first and the second @time result, by changing the N value, you can see it's scaling with N... This should not happen.

[maleadt]

Profiling this function reveals that all the time is spent in eval as called from Memoization.jl, and switching to that package was a recent change indeed. cc @marius311

./boot.jl:360, eval [inlined]
/home/tim/Julia/depot/packages/Memoization/6Lcw7/src/Memoization.jl:46, macro expansion [inlined]
/home/tim/Julia/depot/packages/Memoization/6Lcw7/src/Memoization.jl:42, MethodInstance for get_cache(::Type{IdDict}, ::typeof(CUDA.release))
/home/tim/Julia/depot/packages/Memoization/6Lcw7/src/Memoization.jl:125, soft_limit [inlined]
/home/tim/Julia/pkg/CUDA/src/pool.jl:111, MethodInstance for var"#actual_alloc#160"(::Bool, ::CuStream, ::typeof(CUDA.actual_alloc), ::Int64, ::Bool)

[marius311]

Will take a look, that eval is only supposed to be called during warmup.

[marius311]

@zhangcx93 There was indeed a bug in Memoization affecting your case, which I just fixed in the linked PR. See if ] add Memoization#f704c1f fixes your issues (I checked your MWE and it does fix that). Will tag this version shortly, barring anything else popping up.