Oceananigans.jl: How to reduce compile time for GPU code?

[ali-ramadhan]

Currently it takes 6-7 minutes to compile before the model will start running on the GPU. This is expected but is there a way to compile once and run multiple times, or even to reduce the compile time? Maybe this is the price we pay for zero-cost high-level abstractions.

[vchuravy]

Normally Julia will cache compilation, to reduce compilation costs. Sadly for the GPU we currently have to turn off all the caches and each GPU function basically compiles the entire world.

[ali-ramadhan]

Ah that makes sense. Shouldn't matter for the real-world use case (7 minute compile time << 100+ hour wall clock time simulation). I just need to pick up a hobby I can do in 6-7 minute bursts.

[maleadt]

Do you really compile that many kernels? Or, are those 6-7 minutes really spent in CUDAnative compilation? I wonder, because individual kernel compilation tends to be relatively short (100-200 ms). Maybe we could put some timers in the CUDAnative source code a la https://github.com/JuliaGPU/CuArrays.jl/pull/279. Let's track that here: https://github.com/JuliaGPU/CUDAnative.jl/issues/354

Either way, caching code is tricky because only Julia knows when code has to be invalidated. Within a session, we can look up the age of a method, but across sessions there isn't such a mechanism.
Created an issue: https://github.com/JuliaGPU/CUDAnative.jl/issues/353

[ali-ramadhan]

Thank you for your feedback @maleadt! What is CUDAnative compilation? If you mean the precompilation phase when CUDAnative is first loaded, then it's not that as I start timing after all packages are loaded.

I thought 6-7 minutes was normal/expected as @vchuravy et al. reported similar GPU compilation times for their shallow water model: https://github.com/JuliaLabs/ShallowWaterBench

I haven't done any rigorous benchmarking yet but out of those 6 minutes, ~1.5 minutes are spent on compiling code that creates CuFFT plans (the first plan takes 1.5 minutes then the others take <1 second). From watching the log I'm guessing the other 4.5 minutes are evenly split between setting up the model (creating CuArrays, initializing them, etc.) and the first time step (where the kernels are getting compiled presumably).

I don't think we have that many kernels (just 5 bigger ones) but one of them
https://github.com/ali-ramadhan/Oceananigans.jl/blob/2b64d584c79ece0429f2421335ddb6bc0c6c6663/src/time_steppers.jl#L213
has several layers of inlining (it's inlining the majority of the functions in operators/ops_regular_cartesian_grid_elementwise.jl) after which it probably balloons up to be a pretty big kernel. They also have tons of arguments crammed in as the structs I was passing weren't isbitstype (working on this #59).

I should come back and update this issue once we do some proper benchmarking (note to self: nvprof seems like it's being deprecated in favor of Nsight).

Caching kernels between sessions sounds tough but will definitely look into timing compilations in CUDAnative, might provide some insight on how to speed things up.

[ali-ramadhan]

Some idea (not sure if they would help):

• Use https://github.com/JuliaLang/PackageCompiler.jl
• Check to see if Julia can compile using multiple cores. Right now it only uses 1, when there are sometimes up to 28 cores available.

Also: https://github.com/vchuravy/GPUifyLoops.jl/issues/46 is now relevant to this issue.

[maleadt]

Nono, the problematic modules you identified trigger some quadratic behavior that should just be linear. I'll fix it ASAP.

[ali-ramadhan]

Ah interesting, would explain why it's been 46 minutes and it's still compiling haha.

@lcw may have been hitting the same issue.

[lcw]

Awesome! I am guessing that I am hitting the same thing.

[ali-ramadhan]

Thanks to @maleadt for https://github.com/JuliaGPU/CUDAnative.jl/pull/369 and @vchuravy for https://github.com/vchuravy/GPUifyLoops.jl/pull/49, GPU compile time is now down to 17 seconds!

We'll be able to enjoy these compile times once #147 is finished and merged in.

   _       _ _(_)_     |  Documentation: https://docs.julialang.org
  (_)     | (_) (_)    |
   _ _   _| |_  __ _   |  Type "?" for help, "]?" for Pkg help.
  | | | | | | |/ _` |  |
  | | |_| | | | (_| |  |  Version 1.1.0 (2019-01-21)
 _/ |\__'_|_|_|\__'_|  |  Official https://julialang.org/ release
|__/                   |

julia> using CUDAnative
julia> using Oceananigans
julia> model = Model(N=(128, 128, 128), L=(100, 100, 100), arch=GPU(), float_type=Float32);
julia> time_step!(model, 1, 1)
julia> CUDAnative.timings()
 ────────────────────────────────────────────────────────────────────────────────────
                                             Time                   Allocations      
                                     ──────────────────────   ───────────────────────
          Tot / % measured:                164s / 10.7%           6.24GiB / 29.4%    
 Section                     ncalls     time   %tot     avg     alloc   %tot      avg
 ────────────────────────────────────────────────────────────────────────────────────
 LLVM middle-end                 14    9.08s  52.0%   649ms    637MiB  33.9%  45.5MiB
   IR generation                 14    7.00s  40.1%   500ms    567MiB  30.2%  40.5MiB
     linking                     14    2.49s  14.2%   178ms      224B  0.00%    16.0B
     emission                    14    2.28s  13.0%   163ms    406MiB  21.6%  29.0MiB
     rewrite                     14    2.13s  12.2%   152ms    157MiB  8.34%  11.2MiB
       hide unreachable       2.02k    540ms  3.09%   267μs   20.9MiB  1.11%  10.6KiB
         find                 2.02k    295ms  1.69%   146μs    902KiB  0.05%     456B
         predecessors         2.02k    162ms  0.93%  80.1μs   13.2MiB  0.70%  6.68KiB
         replace              2.02k   68.7ms  0.39%  34.0μs    368KiB  0.02%     186B
       lower throw               14    528ms  3.03%  37.7ms   47.8MiB  2.55%  3.42MiB
       hide trap                 14   65.8ms  0.38%  4.70ms   4.23MiB  0.23%   309KiB
     clean-up                    14   98.9ms  0.57%  7.06ms   4.85MiB  0.26%   355KiB
   optimization                  14    1.99s  11.4%   142ms   69.8MiB  3.71%  4.98MiB
   device library                14   83.7ms  0.48%  5.98ms   11.5KiB  0.00%     839B
   runtime library               14   6.44ms  0.04%   460μs   7.11KiB  0.00%     520B
   verification                  14   2.97ms  0.02%   212μs     0.00B  0.00%    0.00B
 Julia front-end                 14    6.70s  38.3%   478ms   1.11GiB  60.6%  81.3MiB
 CUDA object generation          14    919ms  5.26%  65.6ms   31.9MiB  1.70%  2.28MiB
   linking                       14    780ms  4.47%  55.7ms   14.0MiB  0.75%  1.00MiB
   compilation                   14    139ms  0.79%  9.91ms   17.9MiB  0.95%  1.28MiB
 LLVM back-end                   14    769ms  4.40%  54.9ms   71.1MiB  3.78%  5.08MiB
   machine-code generation       14    124ms  0.71%  8.89ms    192KiB  0.01%  13.7KiB
   preparation                   14   38.2ms  0.22%  2.73ms   2.94MiB  0.16%   215KiB
 ────────────────────────────────────────────────────────────────────────────────────