Fix the FormSum memory leak

firedrake/assemble.py

@@ -469,8 +469,10 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
                 return sum(weight * arg for weight, arg in zip(expr.weights(), args))
             elif all(isinstance(op, firedrake.Cofunction) for op in args):
                 V, = set(a.function_space() for a in args)
-                res = sum([w*op.dat for (op, w) in zip(args, expr.weights())])
-                return firedrake.Cofunction(V, res)
+                result = firedrake.Cofunction(V)
+                for op, w in zip(args, expr.weights()):
+                    result.dat.axpy(w, op.dat)
+                return result
             elif all(isinstance(op, ufl.Matrix) for op in args):
                 res = tensor.petscmat if tensor else PETSc.Mat()
                 is_set = False

pyop2/types/dat.py

@@ -492,6 +492,20 @@ def norm(self):
         from math import sqrt
         return sqrt(self.inner(self).real)
 
+    def axpy(self, alpha: float, other: 'Dat') -> None:
+        """Compute the operation :math:`y = \\alpha x + y`.
+
+        :arg alpha: a scalar
+        :arg other: the :class:`Dat` to add to this one
+
+        """
+        if isinstance(other._data, np.ndarray):
+            np.add(
+                alpha * other.data_ro, self.data_ro,
+                out=self.data_wo)
+        else:
+            raise NotImplementedError("Not implemented for GPU")
+
     def __pos__(self):
         pos = Dat(self)
         return pos
@@ -1022,6 +1036,21 @@ def inner(self, other):
             ret += s.inner(o)
         return ret
 
+    def axpy(self, alpha: float, other: 'MixedDat') -> None:
+        """Compute the operation :math:`y = \\alpha x + y`.
+
+        :arg alpha: a scalar
+        :arg other: the :class:`Dat` to add to this one
+
+        """
+        for dat_result, dat_other in zip(self, other):
+            if isinstance(dat_result._data, np.ndarray):
+                np.add(
+                    alpha * dat_other.data_ro, dat_result.data_ro,
+                    out=dat_result.data_wo)
+            else:
+                raise NotImplementedError("Not implemented for GPU")
+
     def _op(self, other, op):
         ret = []
         if np.isscalar(other):

pyop2/types/glob.py

@@ -203,6 +203,14 @@ def inner(self, other):
         assert issubclass(type(other), type(self))
         return np.dot(self.data_ro, np.conj(other.data_ro))
 
+    def axpy(self, alpha, other):
+        """Compute the operation :math:`y = \\alpha x + y`.
+        """
+        if isinstance(self._data, np.ndarray):
+            np.add(alpha * other.data_ro, self.data_ro, out=self.data_wo)
+        else:
+            raise NotImplementedError("Not implemented for GPU")
+
 
 # must have comm, can be modified in parloop (implies a reduction)
 class Global(SetFreeDataCarrier, VecAccessMixin):