Repository navigation
apply foreach_point in CM cache - #4788
Mikolaj-A-Kowalski wants to merge 8 commits into
Conversation
|
📖 Docs preview for this PR: https://clima.github.io/ClimaAtmos.jl/previews/PR4788/ |
|
@imreddyTeja @dennisYatunin @sjavis I have spend some time trying to eliminate the stack in the heaviest CM cache kernel and in the end what did the trick was to parametrise the This basically allows the values to be passed to the function not via variables, but via instructions. This, in turn, allows the compiler to avoid copying the I do assume that the values of Also to shrink the size of the closure I have used Below is the summary of the kernel runtimes from the NCU reports (and original NCU reports as well)
I would like now to clean up this PR and move it out of the draft status. For that I have two questions:
|
a95021e to
c13e8dc
Compare
Co-authored by: Teja Reddy <treddy0205@gmail.com>
Gives a healthy increase of 6% but still uses a lot of stack do to the function call in the kernel. The benefit seems to originate from better cache efficiency (all threads read the same data)
Instead of trying to access parameters as variables, we move them to the types and place in the instructions. Removes most of the local memory traffic by preventing the compiler from creating the copy of the parameter packs on the stack (since now they are embedded in the type). We observe ~30% speedup.
It does not fix the type instability per se (it must happen at some point since the return type depends on the value of arguments with type parametrised with values). We, however can move it higher in the call stack where it may be less problematic.
70db19a to
f5b6b85
Compare
dennisYatunin
left a comment
There was a problem hiding this comment.
Moving the parameters into instruction memory seems like the right call here, since they are guaranteed to be constant for every simulation. I suspect we'll need to generalize this pattern to keep fusing more kernels.
However, this implementation is still limited by inlining---if the CloudMicrophysics.jl functions that take in the parameters aren't inlined, LLVM might not be able to perform SROA (scalar reduction of aggregates) and infer that only a few parameters need to be stored in registers. To guarantee minimal register usage, the Val wrapper (or any singleton equivalent) would need to be propagated all the way through the CloudMicrophysics.jl source code. For sufficiently complex thermodynamics functions, the same would have to be done in Thermodynamics.jl. I'm happy to see that such extensive changes aren't required for this example, though; maybe we need a lot more complexity before they become necessary.
| let thp_val=Val(thp), cmp_val=Val(cmp), dt=dt, nsubs=nsubs, nsubs_quad=nsubs_quad, | ||
| sgs_quad=sgs_quad, corr_Tq=corr_Tq, α=α |
There was a problem hiding this comment.
What's the purpose of this let-block? I've only seen something like let dt = dt used to avoid boxing inside a closure where dt is reassigned, but this isn't inside a closure. The do-block after foreach_point is a closure, but the let-block is outside of that.
There was a problem hiding this comment.
Here it is not necessary per see, but I am wrapping the variables in the let block to ensure they are captured by value. In the other case I had a problem that a variable would get boxed and captured by reference if it was reassigned in the parent scope of the foreach_block (i.e. outside of closure)
| # Inner constructor: `cmp` and `thp` are captured as `Val`s so that the | ||
| # parameter structs live in the type domain (type parameters `MP`/`TPS`), | ||
| # not in the evaluator's fields. They may be passed either as plain structs | ||
| # or already `Val`-wrapped; both are unwrapped before specialization. | ||
| function Microphysics1MEvaluator( | ||
| scheme::S, mp, tps, ρ::FT, | ||
| q_rai::FT, q_sno::FT, | ||
| λ::FT, λ_lagrange::FT, mu_S::FT, α::FT, | ||
| dt::FT, nsubs::Int, args::Args, | ||
| ) where {S, FT, Args <: Tuple} | ||
| return new{S, unwrap_value(mp), unwrap_value(tps), FT, Args}( | ||
| scheme, ρ, q_rai, q_sno, λ, λ_lagrange, mu_S, α, dt, nsubs, args, | ||
| ) | ||
| end |
There was a problem hiding this comment.
If you're always going to store mp and tps as type parameters, there's no benefit to passing them in without Val wrappers. Doing so would probably trigger dynamic dispatch, which can't be compiled on GPUs. So, instead of adding this inner constructor, you can keep the original default constructor and use unwrap_val(eval.mp)/unwrap_val(eval.tps) below.
There was a problem hiding this comment.
I'm afraid I don't quite follow I am not 100% sure what you mean by
So, instead of adding this inner constructor, you can keep the original default constructor and use unwrap_val(eval.mp)/unwrap_val(eval.tps) below.
I think you are right that I didn't do thinks idiomatically, but my impression is that we are kind of protected from dynamic dispatch depending how we call this constructor.
If we wrap the parameters in Val like:
Microphysics1MEvaluator(scheme, Val(mp), Val(thp) ...)
or mp isa Value && thp is a Value the constructor should not have type instability and we would avoid dynamic dispatch.
It will only happen if we call the constructor with the values of the arguments.
I guess the more full-proof approach would be to just replace mp -> ::Val{mp} in the argument list Is this what you meant?
|
@haakon-e - The debate is here |
Purpose
Apply the foreach_point to the longest running kernel in the
set_microphysics_tendency_cache!functionTo-do
Content
At the moment we don't really observe any notable runtime improvement:
Taken from the NCU reports below
reports.tar.gz
Looking at what the profiler is saying it feels like the main bottleneck for this kernel is the traffic to the local memory (~ 100 million memory requests to local memory vs around 800k issued to the global memory).
In my opinion it is caused by the usage of the stack by the kernel. The compiler does not seem to be able assert that the structs we pass as arguments are not mutated, hence it feels compelled to create stack, which uses local memory and creates the heavy memory traffic.
I don't think it is the case of simple register spillover. In that case I would expect to see similar number of stores and loads (at least it feels like it should be the case). Here we have a heavy imbalance (72 million loads vs 31 million stores) .