No option to pass the tensor when assembling a zero form

[miguelcoolchips]

What is the reason for not accepting the keyword tensor when assembling zero-arity forms? The following example (arguably a bit contrived) shows how the memory grows when we do not pass the tensor keyword and how we can alleviate the issue by passing it.

import firedrake as fd
from firedrake.petsc import PETSc
import time

mesh = fd.UnitSquareMesh(100, 100)
V = fd.FunctionSpace(mesh, "CG", 1)
u = fd.Function(V)
v = fd.TestFunction(V)

from pyop2 import op2
from firedrake import utils

form = u * fd.dx
tensor = op2.Global(1, [0.0], dtype=utils.ScalarType, comm=form.ufl_domains()[0]._comm)

for i in range(100000):
    PETSc.Sys.Print(f"i = {i}")
    t = time.process_time()
    fd.assemble(form)  # Case 1 black
    # fd.assemble(form, tensor=tensor)  # Case 2 blue
    elapsed_time = time.process_time() - t
    PETSc.Sys.Print(f"Elapsed time: {elapsed_time} seconds")

This code results in the following memory plots

The effect is not as noticeable as the problem size grows, but it is still there (of course, 100,000 iterations are a lot of iterations). I modified firedrake to let it accept tensor as keyword. In general, for reusing the assembly routine many times, is it better to build my own assembler (ZeroFormAssembler)? The thing is that firedrake-adjoint builds a new form for every evaluation of the tape. I would prefer not having to cache the form because I would have to update its coefficients whenever I run the tape.

My firedrake-status is (a bit outdated, but I see there's still no support for tensor in zero-form assembly)

Firedrake Configuration:
    package_manager: True
    minimal_petsc: False
    mpicc: mpicc
    mpicxx: mpicxx
    mpif90: mpif90
    mpiexec: mpirun
    disable_ssh: False
    honour_petsc_dir: True
    with_parmetis: False
    slepc: False
    packages: []
    honour_pythonpath: False
    opencascade: False
    tinyasm: False
    petsc_int_type: int32
    cache_dir: /Users/miguel/repos/firedrake-latest/.cache
    complex: False
    remove_build_files: False
    with_blas: None
Additions:
    None
Environment:
    PYTHONPATH: None
    PETSC_ARCH: default
    PETSC_DIR: /Users/miguel/repos/petsc-latest
Status of components:
---------------------------------------------------------------------------
|Package             |Branch                        |Revision  |Modified  |
---------------------------------------------------------------------------
|COFFEE              |master                        |70c1e66   |False     |
|FInAT               |master                        |cd1d528   |False     |
|PyOP2               |master                        |233f6dc3  |False     |
|fiat                |master                        |9fca261   |False     |
|firedrake           |master                        |18c3c2828 |True      |
|h5py                |firedrake                     |6cc4c912  |False     |
|libspatialindex     |master                        |4768bf3   |True      |
|libsupermesh        |master                        |69012e5   |False     |
|loopy               |main                          |3988272b  |False     |
|pyadjoint           |master                        |4742834   |True      |
|tsfc                |master                        |b6da9e0   |False     |
|ufl                 |master                        |13435fbf  |False     |
---------------------------------------------------------------------------

[dham]

There's no strong reason not to allow the tensor to be passed to zero forms. I'm a bit mystified by this though, the variables in question should be cleaned up by Python once they go out of scope.

[wence-]

So that probably means that assemble is a relevant garbage event. Given that solve calls assemble, maybe we should garbage collect on assemble instead of solve.

I think for mass-lumped explicit schemes that might end up making each step quite a bit slower. Perhaps there needs to be a way of manually explicitly destroying the (nominally collectively) destroyed petsc objects individually that one can call from logically collective routines in firedrake (does PetscObject.destroy() actually destroy, or just shove things on the garbage list?). Firedrake knows that any temporarily scoped objects it makes inside assemble (say) are safe for collective destruction when they go out of scope because you know you are already collective here.

Or is the problem that these objects are not always dropped by their refcount and instead are only collected by the cyclic collector?

[JDBetteridge]

PetscObject.destroy() should work as it always has, the delayed destroy only kicks in if things are cleaned up by the cyclic garbage collector.

[dham]

It's the last of these. However (Jack correct me if I'm wrong), we only actually garbage collect every 1/n collectable events in order to not trigger expensive garbage collection too often.

[wence-]

PetscObject.destroy() should work as it always has, the delayed destroy only kicks in if things are cleaned up by the cyclic garbage collector.

Confused: .destroy just decrefs the petsc refcount right, if things are in cycles (so being cleaned up by the cyclic collector) then .destroy on such an object won't (correctly) actually destroy the C handle.

[JDBetteridge]

@dham We (I) didn't implement the every 1/n cleanup. If it became a performance issue we were going to revisit the idea.

@wence- Yes, the destroy routine is just calling PetscObjectDestroy and subsequently the PETSc decref routine, the PETSc C handle is only destroyed if the PETSc refcount goes to zero.