Interpolator: bugfix for reusable matfree adjoint Interpolator in parallel

Author: pbrubeck

firedrake/interpolation.py

@@ -931,9 +931,6 @@ def callable():
     else:
         loops = []
 
-        if access == op2.INC:
-            loops.append(tensor.zero)
-
         # Arguments in the operand are allowed to be from a MixedFunctionSpace
         # We need to split the target space V and generate separate kernels
         if len(arguments) == 2:
@@ -961,12 +958,23 @@ def callable():
         if bcs and rank == 1:
             loops.extend(partial(bc.apply, f) for bc in bcs)
 
-        def callable(loops, f):
-            for l in loops:
-                l()
+        def callable(loops, f, access):
+            if access is op2.WRITE:
+                for l in loops:
+                    l()
+                return f
+            # We are repeatedly incrementing into the same Dat so intermediate halo exchanges
+            # can be skipped.
+            f.dat.local_to_global_begin(access)
+            with f.dat.frozen_halo(access):
+                if access is op2.INC:
+                    f.dat.zero()
+                for l in loops:
+                    l()
+            f.dat.local_to_global_end(access)
             return f
 
-        return partial(callable, loops, f)
+        return partial(callable, loops, f, access)
 
 
 @utils.known_pyop2_safe
@@ -1076,7 +1084,7 @@ def _interpolator(tensor, expr, subset, access, bcs=None):
     coefficient_numbers = kernel.coefficient_numbers
     needs_external_coords = kernel.needs_external_coords
     name = kernel.name
-    kernel = op2.Kernel(ast, name, requires_zeroed_output_arguments=True,
+    kernel = op2.Kernel(ast, name, requires_zeroed_output_arguments=(access is op2.WRITE),
                         flop_count=kernel.flop_count, events=(kernel.event,))
 
     parloop_args = [kernel, cell_set]
@@ -1099,7 +1107,7 @@ def _interpolator(tensor, expr, subset, access, bcs=None):
     if isinstance(tensor, op2.Global):
         parloop_args.append(tensor(access))
     elif isinstance(tensor, op2.Dat):
-        V_dest = arguments[-1].function_space() if isinstance(dual_arg, ufl.Cofunction) else V
+        V_dest = arguments[0].function_space()
         m_ = get_interp_node_map(source_mesh, target_mesh, V_dest)
         parloop_args.append(tensor(access, m_))
     else:
@@ -1159,11 +1167,10 @@ def _interpolator(tensor, expr, subset, access, bcs=None):
                 parloop_args.append(target_ref_coords.dat(op2.READ, m_))
 
     parloop = op2.ParLoop(*parloop_args)
-    parloop_compute_callable = parloop.compute
     if isinstance(tensor, op2.Mat):
-        return parloop_compute_callable, tensor.assemble
+        return parloop, tensor.assemble
     else:
-        return copyin + callables + (parloop_compute_callable, ) + copyout
+        return copyin + callables + (parloop, ) + copyout
 
 
 def get_interp_node_map(source_mesh, target_mesh, fs):

tests/firedrake/regression/test_interpolate.py

@@ -327,6 +327,7 @@ def test_trace():
     assert np.allclose(x_tr_cg.dat.data, x_tr_dir.dat.data)
 
 
+@pytest.mark.parallel(nprocs=[1, 3])
 @pytest.mark.parametrize("rank", (0, 1))
 @pytest.mark.parametrize("mat_type", ("matfree", "aij"))
 @pytest.mark.parametrize("degree", (1, 3))
@@ -566,3 +567,35 @@ def test_mixed_matrix(mode):
     result_explicit = assemble(action(a, u))
     for x, y in zip(result_explicit.subfunctions, result_matfree.subfunctions):
         assert np.allclose(x.dat.data, y.dat.data)
+
+
+@pytest.mark.parallel(nprocs=2)
+@pytest.mark.parametrize("mode", ["forward", "adjoint"])
+@pytest.mark.parametrize("family,degree", [("CG", 1)])
+def test_reuse_interpolate(family, degree, mode):
+    mesh = UnitSquareMesh(1, 1)
+    V = FunctionSpace(mesh, family, degree)
+    rg = RandomGenerator(PCG64(seed=123456789))
+    if mode == "forward":
+        u = Function(V)
+        expr = interpolate(u, V)
+
+    elif mode == "adjoint":
+        u = Function(V.dual())
+        expr = interpolate(TestFunction(V), u)
+
+    I = Interpolator(expr, V)
+
+    for k in range(2):
+        u.assign(k+1)
+        expected = u.dat.data.copy()
+        result = I.assemble()
+
+        # Test that the input was not modified
+        x = u.dat.data
+        assert np.allclose(x, expected)
+
+        # Test for correctness
+        y = result.dat.data
+        assert np.allclose(y, expected)
+        print("pass", k)