Cofunction vesus Function on RHS

[Ig-dolci]

I encountered a result that raised some questions. Let me summarise this:

In the code below, q_s is a Cofunction added on the right-hand side (RHS) of the equation. I executed two tests:
Test 1: The code is implemented exactly as shown, with the Confunction q_s directly included in the RHS.
Test 2: I first transform q_s into its Riesz representation using q_s = q_s.riesz_representation(). I then rewrite F as: F = time_term + a - q_s * v * dx and set up the problem as lin_var = LinearVariationalProblem(lhs(F), rhs(F), u_np1)
Both approaches should solve the same system. However, I noticed differences in runtime and memory usage between the two implementations, as illustrated in the figure below. Is this difference in performance expected? Could this behavior indicate a bug or an inefficiency in how the system handles Cofunction compared to its Riesz representation?

from firedrake import *
import numpy as np

mesh = UnitSquareMesh(100, 100)

source_locations = [0.5, 0.5]
dt = 0.001  # time step in seconds
final_time = 1.  # final time in seconds

V = FunctionSpace(mesh, "CG", 1)
x, z = SpatialCoordinate(mesh)
c = Function(V).interpolate(1.5)


source_mesh = VertexOnlyMesh(mesh, [source_locations])
V_s = FunctionSpace(source_mesh, "DG", 0)
d_s = Function(V_s)
d_s.interpolate(1.0)
source_cofunction = assemble(d_s * TestFunction(V_s) * dx)
q_s = Cofunction(V.dual()).interpolate(source_cofunction)

u = TrialFunction(V)
v = TestFunction(V)
u_np1 = Function(V, name="u_np1") # timestep n+1
u_n = Function(V, name="u_n") # timestep n
u_nm1 = Function(V, name="u_nm1") # timestep n-1
m = (1 / (c * c))
time_term =  m * (u - 2.0 * u_n + u_nm1) / Constant(dt**2) * v * dx
a = dot(grad(u_n), grad(v)) * dx
F = time_term + a
lin_var = LinearVariationalProblem(lhs(F), rhs(F) +  q_s, u_np1)
solver = LinearVariationalSolver(lin_var)
for _ in range(50):
    print("step = ", _)
    for step in range(int(final_time / dt) + 1):
        solver.solve()
        u_nm1.assign(u_n)
        u_n.assign(u_np1)

[connorjward]

I think the difference is in when we assemble the two things. When we have a Cofunction then we are assembling a FormSum instead of a Form. These use different Assembler types (source). Form assembly (with FormAssembler/ParloopFormAssembler) is quite heavily optimised whereas the FormSum (with BaseFormAssembler) does what I expect is quite an expensive DAG traversal each time.

I expect that one could optimise BaseFormAssembler in a clever way to reduce this cost.

[ksagiyam]

It seems the following already exhibits memory leak:

from firedrake import *


mesh = UnitIntervalMesh(1)
V = FunctionSpace(mesh, "CG", 1)
a = Cofunction(V.dual())
for _ in range(1000000):
    assemble(a + a)

[connorjward]

It looks like you are only running a short simulation - pyop2.compilation only takes half a second - so perhaps this is OK.

[Ig-dolci]

Yes. This is a flamegraph only for

for _ in range(1:
    print("step = ", _)
    for step in range(int(final_time / dt) + 1):
        solver.solve()
        u_nm1.assign(u_n)
        u_n.assign(u_np1)

which is 1000 of time-steps.

[dham]

Look at all the unaccounted for time in SNESFunctionEval. I bet that's assemble being super slow for this case.

[wence-]

My crystal ball says the cofunction numbering is increasing and turning up in the compiled form each time

[ksagiyam]

It looks that the kernel is not cached on the output/input Dats appropriately in pyop2:

diff --git a/pyop2/types/dat.py b/pyop2/types/dat.py
index fb877c1a8..6144474a1 100644
--- a/pyop2/types/dat.py
+++ b/pyop2/types/dat.py
@@ -383,8 +383,16 @@ class AbstractDat(DataCarrier, EmptyDataMixin, abc.ABC):
         else:
             self._check_shape(other)
             globalp = False
+            if hasattr(other, "_op_kernel_cache"):
+                if hasattr(self, "_op_kernel_cache"):
+                    self._op_kernel_cache.update(other._op_kernel_cache)
+                else:
+                    self._op_kernel_cache = other._op_kernel_cache
         parloop(self._op_kernel(op, globalp, other.dtype),
                 self.dataset.set, self(Access.READ), other(Access.READ), ret(Access.WRITE))
+        ret._op_kernel_cache = self._op_kernel_cache
+        if not globalp:
+            other._op_kernel_cache = self._op_kernel_cache
         return ret
 
     def _iop_kernel(self, op, globalp, other_is_self, dtype):

If I apply the above patch, I can see that the memory usage is flattened in the following example:

from firedrake import *
import gc


def f():
    mesh = UnitIntervalMesh(1)
    V = FunctionSpace(mesh, "CG", 1)
    a = Cofunction(V.dual())
    for _ in range(20000):
        _ = assemble(a + a)  # or 1 * a.dat + 1 * a.dat
        gc.collect()


f()
for _ in range(100):
    gc.collect()

I still have not investigated the direct cause of memory leak.