Fix assemble with multiple kernels and subdomains (#3135)

* Fix issue 3125

---------

Co-authored-by: David A. Ham <[email protected]>

Author: connorjward

firedrake/assemble.py

@@ -692,7 +692,7 @@ def replace_tensor(self, tensor):
             return
 
         # TODO We should have some proper checks here
-        for lknl, parloop in zip(self.local_kernels, self.parloops):
+        for (lknl, _), parloop in zip(self.local_kernels, self.parloops):
             data = _FormHandler.index_tensor(tensor, self._form, lknl.indices, self.diagonal)
             parloop.arguments[0].data = data
         self._tensor = tensor
@@ -703,6 +703,14 @@ def execute_parloops(self):
 
     @cached_property
     def local_kernels(self):
+        """Return local kernels and their subdomain IDs.
+
+        Returns
+        -------
+        tuple
+            Collection of ``(local_kernel, subdomain_id)`` 2-tuples, one for
+            each possible combination.
+        """
         try:
             topology, = set(d.topology for d in self._form.ufl_domains())
         except ValueError:
@@ -714,16 +722,26 @@ def local_kernels(self):
                 raise NotImplementedError("Assembly with multiple meshes is not supported")
 
         if isinstance(self._form, ufl.Form):
-            return tsfc_interface.compile_form(self._form, "form", diagonal=self.diagonal,
-                                               parameters=self._form_compiler_params)
+            kernels = tsfc_interface.compile_form(
+                self._form, "form", diagonal=self.diagonal,
+                parameters=self._form_compiler_params
+            )
         elif isinstance(self._form, slate.TensorBase):
-            return slac.compile_expression(self._form, compiler_parameters=self._form_compiler_params)
+            kernels = slac.compile_expression(
+                self._form,
+                compiler_parameters=self._form_compiler_params
+            )
         else:
             raise AssertionError
+        return tuple(
+            (k, subdomain_id) for k in kernels for subdomain_id in k.kinfo.subdomain_id
+        )
 
     @cached_property
     def all_integer_subdomain_ids(self):
-        return tsfc_interface.gather_integer_subdomain_ids(self.local_kernels)
+        return tsfc_interface.gather_integer_subdomain_ids(
+            {k for k, _ in self.local_kernels}
+        )
 
     @cached_property
     def global_kernels(self):
@@ -732,20 +750,28 @@ def global_kernels(self):
                 self._form, tsfc_knl, subdomain_id, self.all_integer_subdomain_ids,
                 diagonal=self.diagonal, unroll=self.needs_unrolling(tsfc_knl, self._bcs)
             )
-            for tsfc_knl in self.local_kernels
-            for subdomain_id in tsfc_knl.kinfo.subdomain_id
+            for tsfc_knl, subdomain_id in self.local_kernels
         )
 
     @cached_property
     def parloops(self):
-        return tuple(
-            ParloopBuilder(
-                self._form, lknl, gknl, self._tensor, subdomain_id,
-                self.all_integer_subdomain_ids, diagonal=self.diagonal,
-                lgmaps=self.collect_lgmaps(lknl, self._bcs)).build()
-            for lknl, gknl in zip(self.local_kernels, self.global_kernels)
-            for subdomain_id in lknl.kinfo.subdomain_id
-        )
+        loops = []
+        for (local_kernel, subdomain_id), global_kernel in zip(
+            self.local_kernels, self.global_kernels
+        ):
+            loops.append(
+                ParloopBuilder(
+                    self._form,
+                    local_kernel,
+                    global_kernel,
+                    self._tensor,
+                    subdomain_id,
+                    self.all_integer_subdomain_ids,
+                    diagonal=self.diagonal,
+                    lgmaps=self.collect_lgmaps(local_kernel, self._bcs)
+                ).build()
+            )
+        return tuple(loops)
 
     def needs_unrolling(self, local_knl, bcs):
         """Do we need to address matrix elements directly rather than in

tests/regression/test_assemble.py

@@ -266,3 +266,19 @@ def test_assemble_mixed_function_sparse():
     f.sub(4).interpolate(Constant(3.0))
     v = assemble((inner(f[1], f[1]) + inner(f[4], f[4])) * dx)
     assert np.allclose(v, 13.0)
+
+
+def test_3125():
+    # see https://github.com/firedrakeproject/firedrake/issues/3125
+    mesh = UnitSquareMesh(3, 3)
+    V = VectorFunctionSpace(mesh, "CG", 2)
+    W = FunctionSpace(mesh, "CG", 1)
+    Z = MixedFunctionSpace([V, W])
+    z = Function(Z)
+    u, p = split(z)
+    tst = TestFunction(Z)
+    v, q = split(tst)
+    d = Function(W)
+    F = inner(z, tst)*dx + inner(u, v)/(d+p)*dx(2, degree=10)
+    # should run without error
+    solve(F == 0, z)