NLVS: Cache more adjoint solvers/forms and fix Hessian

[JHopeCollins]

No description provided.

[sghelichkhani]

Hi team! Excellent work and thanks for doing this. I just wanted to flag something before this lands so we don't end up pulling in opposite directions.

I've got a small fix sitting on a branch against main that makes the adjoint replay clone the appctx dictionary the same way it already clones the form coefficients. The problem it solves is that when a preconditioner reads UFL expressions out of appctx, for us in gadopt and for the schur complement is MassInvPC, the replayed forward and adjoint solves end up pointing at the user's original coefficients rather than the deep copies the tape keeps in sync, so they quietly use stale values.

I was just about to do a PR for that on main. But since your work reorganises exactly that area and your refactor is the more structural and better change, I am thinking it's the natural place to do this properly rather than me landing a fix on main that you'd then have to unpick. So I am going to do PR targetting this one if you are OK. If not and you don't like the idea just tell me off and I wait for this to land and do my thing.

[JHopeCollins]

Hi Sia,

Thanks for raising this. It's a good point, I don't know how we haven't run into this before!
I think you are right that it will be easier to integrate it into this PR rather than duplicating the work. I will have a look at your suggested fix and have a think about how it will fit into the rest of the changes.

[sghelichkhani]

don't know how we haven't run into this before!

It only is a problem if the object on appctx changes during taping, so whatever is in there is different at the end of taping and the first fwd block. in our case is a temperature (coming from a different solve) dependent viscosity.

[sghelichkhani]

A quick heads-up since this PR is reshaping the same file: I have just opened #5175, which fixes #5172 by recording the projection's effective solver parameters on the ProjectBlock (and recording same-space projections as assignments). It only touches ProjectBlock construction from the annotation side, so it should be orthogonal to the CachedSolverBlock changes here, but it does add a few lines near solve_init_params territory, so a merge conflict is possible whichever lands second. I have also filed #5176 for the fact that direct Projector use is never annotated; the natural fix there, a block that holds the projector and reuses its cached solver, is really the projector analogue of what this PR does for the NLVS, so I would rather build that on top of your refactor than in parallel with it.

[JHopeCollins]

A quick heads-up since this PR is reshaping the same file: I have just opened #5175, which fixes #5172 by recording the projection's effective solver parameters on the ProjectBlock (and recording same-space projections as assignments). It only touches ProjectBlock construction from the annotation side, so it should be orthogonal to the CachedSolverBlock changes here, but it does add a few lines near solve_init_params territory, so a merge conflict is possible whichever lands second. I have also filed #5176 for the fact that direct Projector use is never annotated; the natural fix there, a block that holds the projector and reuses its cached solver, is really the projector analogue of what this PR does for the NLVS, so I would rather build that on top of your refactor than in parallel with it.

Hi Sia, thanks again for raising this. I think we should follow the same process as the appctx fix. The fix into release will stand, and will then get merged into main. Then when Angus and I get to dealing with the project block in this PR we will look at how best to incorporate your fix.