Upstream approximate division to Oceananigans

[Mikolaj-A-Kowalski]

Starting January we should probably move to start up-streaming the performance-beneficial changes. Division is a good place to start.

• Use the rcp.approx.ftz.f64 (or f32) as the first guess in the newton_division
• Try to replace some reciprocals in operators (e.g. \delata) as well.
  • study performance impact
  • see if there are problems with regression tests as a result

[Mikolaj-A-Kowalski]

I repeated the study of using rcp.approx.ftz.f64 in the newton_div. Now with a 'PR' ready implementation (although based on the refactor-benchmarks branch, I did not rebase yet on main due to some version conflicts with ClimaOcean)

The reports are here: reports.tar.gz

On the final version of the test benchmark with the tripolar -grid and realistic immersed boundary from ClimaOcean we are looking at few per-cent speedup. Not as much as before but not bad.

The next step I think is to try to replace all divisions on the GPU with the approximate version. In my opinion it will not have drastic effect on the numerical precision for the following reason:

• Cases where denominator can get close to 0 already tend to include 'regularisation constant` which makes certain that there will be no subnormal numbers in the denominator (hence flush-to-zero will not affect it)
• Optimisation of replacing a/b => a * (1/b) is already applied quite often. Hence we are more often dealing with reciprocal than division.
• The approximate version of the reciprocal on the GPU should not be bad... I will need to set up more through testing (on random input...?) but rcp.approx.ftz.f64 should give us roughly $\epsilon \approx 2^{-20}$ which after two Newton steps should bring us to $\epsilon \approx 2^{-80}$, so we are looking at the error in the few ULP I would recon, with maybe different behaviour for large numbers close to infinity. My hope it will not end up being that different in practice...

[Mikolaj-A-Kowalski]

Note, it is considerably faster with a single newton step instead of two (~6% improvement becomes ~9%). Maybe I should also try a case of Halley iteration as proposed here.

It will be 3 fmas instead of 2 fmas and 2 muls... and theoretically should still deliver good $\mathcal{O} (2^{-60}$ error so it shouldn’t be worse than w Netwons performance wise.

The problem is that I am not entirely sure where does this formulation of Halley method for reciprocal comes from. I was not able to recover it myself just yet (however I am certain it is correct)

[Mikolaj-A-Kowalski]

So the 'all approximate' division seems to accelerate the kernels by 50%!

all_approx_div.tar.gz

So the conclusion is that it may be something worth fighting for. But there are also some red flags sadly...

• NCU tells me that the Avg. Active Threads Per Warp jumped to 32 (from ~23)! This is bad because it should mean either of two things: 1) either I broke something in the kernel 2) branching was due to the slow-path of the division, so it being gone means that there was a lot of divisions needing slow path 😬 , which may spell trouble related to precision.

There are also some interesting observations to do:

• 'Stall No Instruction' stall reason was a bit of mystery to us I remember. With no division slow-paths it basically disappears. I think we can now safely have a mental model that it is related to either function calls or branching to distant instruction address

[Mikolaj-A-Kowalski]

On one hand I fear it may be a bit orthogonal... but on the other I think it will be crucial. I will aim to make some small code for evaluation of the precision of a reciprocal vs properly rounded implementation and will spend some time trying to reverse engineer the CUDA's f64 division to try to understand when exactly slow path is needed.

[Mikolaj-A-Kowalski]

So the conclusion is that it may be something worth fighting for. But there are also some red flags sadly...

• NCU tells me that the Avg. Active Threads Per Warp jumped to 32 (from ~23)! This is bad because it should mean either of two things: 1) either I broke something in the kernel 2) branching was due to the slow-path of the division, so it being gone means that there was a lot of divisions needing slow path 😬 , which may spell trouble related to precision.

Turns out I had a stupid typo in the Newton iteration so the kernel was broken!
The good thing is that with the fixed version we are seeing good ~20% improvement, but not much change in the number of active wraps! Hence the fear about slow path being common can be reduced.

Fixed set of reports is here: reports.tar.gz

[Mikolaj-A-Kowalski]

Looking into the rcp implementation now can confirm 100% that the slow path is taken only by the subnormal numbers: NaNs and Infs (EDIT: and above 4.492034101493843e307). For all other values the fast path should sufficient to get correctly rounded result.

EDIT: Overflow... false confidence is the worse confidence :-/

[Mikolaj-A-Kowalski]

To the long story short is that we will have no problem to have an accurate f64 reciprocal in the range [2.2250738585072014e-308: 4.494237123190219e307 ). [Aka first 'normal' f64 number to first f64 number for which the 'chopped' rcp.approx.ftz reciprocal underflows]

The best way to do it is to steal the fast path of the Nvidia's reciprocal. Then the question is basically how 'safe' we want to be with regards to things like rounding and under/overflow behaviour. I would argue we can be quite adventurous if we manage to include alternative division as a user-option, opt-in feature.

Using the pure fast-path gives roughly ~15% runtime improvement in the advection kernel. It has however bad underflow behaviour (i.e. returns NaN for subnormal or 0 input). Naive fix (via predicate and conditional move) to return Inf in that cases drops the improvement to ~9%. I am not quite sure why it is... it introduces extra dependencies between instructions and requires couple more registers, but the effect is more than I would have expected.

I haven't run any regression tests or longer simulations to verify the effect on numerical accuracy (should be well below $ \sqrt{\epsilon}$ though, that or simulation crashes because we hit under/overflow.)

[Mikolaj-A-Kowalski]

Opened CliMA/Oceananigans.jl#5140 to deal with the newton_div case

[Mikolaj-A-Kowalski]

Greg's suggestions:

• Define new Number type to allow changing the division.
• Change dispatch of newton_div to use 'option type' and disable when reactant is alive