Remove reference cycle in VecAccessMixin

[angus-g]

With an associated PETSc Vec, VecAccessMixin deferred its version property to a lambda to avoid allocating the storage until necessary. Unfortunately, this lambda creates a reference cycle to self for all users of the VecAccessMixin. Given that counter accesses should be relatively infrequent, it seems fine to look up the counter type within the method itself.

Description

Related to #4014. To benchmark, I'm using the following script (very similar to the one in the linked issue, but uses 500 timesteps, a timestepper object, and removes explicit GC calls):

from firedrake import *
from firedrake.adjoint import *
# from memory_profiler import profile

def test():
    T_c, rf = rf_generator()
    rf.fwd_call = profile(rf.__call__)
    rf.derivative = profile(rf.derivative)

    for i in range(5):
        rf.fwd_call(T_c)
        rf.derivative()

@profile
def rf_generator(checkpoint_to_disk=True):
    tape = get_working_tape()
    tape.clear_tape()
    continue_annotation()

    mesh = RectangleMesh(100, 100, 1.0, 1.0)

    if checkpoint_to_disk:
        enable_disk_checkpointing()
        mesh = checkpointable_mesh(mesh)

    V = VectorFunctionSpace(mesh, "CG", 2)
    Q = FunctionSpace(mesh, "CG", 1)

    # Define the rotation vector field
    X = SpatialCoordinate(mesh)

    w = Function(V, name="rotation").interpolate(as_vector([-X[1] - 0.5, X[0] - 0.5]))
    T_c = Function(Q, name="control")
    T = Function(Q, name="Temperature")
    T_c.interpolate(0.1 * exp(-0.5 * ((X - as_vector((0.75, 0.5))) / Constant(0.1)) ** 2))
    control = Control(T_c)
    T.assign(T_c)

    # for i in ts:
    for i in tape.timestepper(iter(range(500))):
        T.interpolate(T + inner(grad(T), w) * Constant(0.0001))

    objective = assemble(T**2 * dx)

    pause_annotation()
    return T_c, ReducedFunctional(objective, control)


if __name__ == "__main__":
    test()

I'm also running this on the #4020 branch to automatically enable the SingleDiskStorageSchedule and handle the leak of function within CheckpointFunction. On the pyadjoint side, I am using dolfin-adjoint/pyadjoint#194.

Here's a pretty simple mprof comparison:

In black is the base, without this branch. In blue is the base, but with gc.collect() within Block.recompute (very eager, and expensive, also doesn't apply to the derivative). In red is the result with this branch, without any explicit gc. Individual plots follow, but the rescaling means you have to look a bit more closely.

Base plot

GC plot

This PR

I think there is a still a bit left out there in terms of making expensive allocations delete through refcounting, and perhaps there's a more elegant way of implementing the change proposed here.

[github-actions]

	Tests	Passed ✅	Skipped ⏭️	Failed ❌
Firedrake real	8238 ran	7524 passed	714 skipped	0 failed

[github-actions]

	Tests	Passed ✅	Skipped ⏭️	Failed ❌
Firedrake complex	8317 ran	6614 passed	1703 skipped	0 failed

[connorjward]

This is an excellent spot!

I have absolutely no idea why this is failing tests though... AFAICT the changes you have made shouldn't impact the rest of the code.

[angus-g]

I had to do a few wider modifications around the inheritance of AbstractDat and VecAccessMixin. Hopefully it passes the tests now.

[github-actions]

	Tests	Passed ✅	Skipped ⏭️	Failed ❌
Firedrake default	8201 ran	7476 passed	725 skipped	0 failed

[connorjward]

Sorry for taking so long to review this. I think that your changes demonstrate that VecAccessMixin is a poor class for this as some things have a dat_version but no Vec.