Add missing blocks to the dycore Jacobian

[dennisYatunin]

Purpose

This PR adds a couple of missing blocks to the dycore Jacobian and modifies the linear solver to handle them. This improves the stability of two longrun simulations that were previously failing (Held-Suarez with topography and aquaplanet with topography).

Content

• Account for the ${\displaystyle \frac{\partial u^3}{\partial u_h}}$ derivative, which scales with $g^{3h}$ and is therefore significant in simulations with topography. This adds the tridiagonal ${\displaystyle \frac{\partial \mathcal{E}\rho }{\partial u_h}}$, ${\displaystyle \frac{\partial \mathcal{E}\rho e_{tot}}{\partial u_h}}$, and ${\displaystyle \frac{\partial \mathcal{E}\rho \chi }{\partial u_h}}$ blocks to the Jacobian.
• Account for the ${\displaystyle \frac{\partial \kappa }{\partial u_h}}$ derivative. This scales with $g^{hh}$ and $g^{3h}$, so it is relevant both with and without topography. However, it only adds the bidiagonal ${\displaystyle \frac{\partial \mathcal{E}{u}_3}{\partial u_h}}$ block to the Jacobian, which, along with the ${\displaystyle \frac{\partial \mathcal{E}{u}_3}{\partial u_3}}$ block, is smaller than similar blocks by a factor of roughly $({\displaystyle \frac{u}{c}}{)}^2$. In the future, we might want to drop both of these blocks if they do not significantly improve stability. Since they do not affect the structure of the linear solver, we will keep them in for now.
• Account for the ${\displaystyle \frac{\partial p}{\partial \rho q_{tot}}}$ derivative, which is significant in simulations with moisture. This adds the bidiagonal ${\displaystyle \frac{\partial \mathcal{E}{u}_3}{\partial \rho q_{tot}}}$ block to the Jacobian, along with the tridiagonal ${\displaystyle \frac{\partial \mathcal{E}\rho e_{tot}}{\partial \rho q_{tot}}}$ block in simulations with implicit diffusion.
• Account for the ${\displaystyle \frac{\partial e_{tot}}{\partial \rho }}$ and ${\displaystyle \frac{\partial \chi }{\partial \rho }}$ derivatives. These add the tridiagonal ${\displaystyle \frac{\partial \mathcal{E}\rho e_{tot}}{\partial \rho }}$ and ${\displaystyle \frac{\partial \mathcal{E}\rho \chi }{\partial \rho }}$ blocks to the Jacobian in simulations with implicit diffusion.
• Simplify update_implicit_equation_jacobian! and cache several matrix Fields in order to improve performance. Remove outdated code comments (new documentation will be ported from Overleaf in a future PR).
• Slightly increase the allocation limits of some performance tests and bump the ref counter.

---

• I have read and checked the items on the review checklist.

[szy21]

Is this ready for testing? Maybe @akshaysridhar could try running some topography simulations.

[dennisYatunin]

@szy21 There seem to be a couple of performance+GPU issues, but other than that this is ready for testing. I'm trying out the longruns here.

[szy21]

I think the topography run breaks faster: https://buildkite.com/clima/climaatmos-longruns/builds/1294#018c6032-e583-4164-b3fe-3df08ce1277c :(

[dennisYatunin]

CPU longruns: https://buildkite.com/clima/climaatmos-longruns/builds/1303
GPU longruns: https://buildkite.com/clima/climaatmos-gpulongruns/builds/10

@szy21 Would you mind taking a look at the results before I merge this in?

[szy21]

I looked at some of the dycore simulations that are done and they look good. Feel free to merge this in. Thanks!

[szy21]

The GPU pipeline seems to be 40% slower though. Is this as expected? (The CPU timeline is likely slower too as many simulations are timing out.)

[dennisYatunin]

Yeah this PR replaces the simple diagonal solver for the diffusion matrix with an iterative solver. This needs to run on every iteration of the higher-level solver for the full Jacobian, so it'll have a significant performance effect. This PR also adds a bunch of additional blocks to the Jacobian, and these currently take a while to compute.

[szy21]

Ok, sounds good to me. Just want to let people who work on performance know (@simonbyrne @sriharshakandala).

[dennisYatunin]

From the flame graphs, it looks like the new blocks are the main source of the performance hit. So, we might want to change something in ClimaCore to speed up long chains of matrix-matrix multiplications. I'll look into this when I get back.

[szy21]

I did want to know whether this helps with the topography runs though. I'm ok with merging it in for now as 40% is not too bad, or @akshaysridhar can test with this branch. Either way is fine.

[dennisYatunin]

Hmm, looks like more memory increases are required. Before I merge this in, I'll try dropping some of the new Jacobian blocks in simulations without topography, which will probably improve the performance of those simulations.

[dennisYatunin]

CPU longruns: https://buildkite.com/clima/climaatmos-longruns/builds/1355
GPU longruns: https://buildkite.com/clima/climaatmos-gpulongruns/builds/24

Performance issues appear to be resolved. I'll merge this in when the longruns are done.

[szy21]

The longrun looks fine. Feel free to merge it.