[pbrubeck]

Stacked on #5258. Includes #5567.

Summary

Forms that differ only in their coefficients now share their loops. Function.interpolate builds a new form on every call, and so do multigrid transfers through a quadrature space. Before, each call built and hashed new loops.

 ParloopFormAssembler.parloops(tensor)
-  build the loops for self._form
+  _build_parloops(tensor)    # cached on the mesh, keyed by the signature, the spaces,
+                             # the subdomain data and the assembler options

 ParloopFormAssembler.assemble
-  parloop(**{tensor name: tensor})
+  parloop(**{coefficient argument names: dats of self._form, tensor name: tensor})

ParloopBuilder records the names of the loop arguments that take the coefficients and constants. Two commits make those names reliable:

• Slate: Number coefficients by their position in the expression signature (Slate: Number coefficients by their position in the expression signature #5567). Slate expressions now have a signature() that does not depend on the coefficients, so they share loops too. It also fixes a wrong-result bug: two Slate expressions with different coefficients in the same position shared a kernel.
• pyop3: Give each copy of a tensor a name of its own. Two copies of one Dat were both named <name>_copy, and pyop3 drops a loop argument name that two arguments share.

Evidence

test_loops_shared_between_coefficients (new, 2 ranks, forms of rank 0, 1 and 2) and test_interpolate_new_coefficients (new, 2 ranks, forward and adjoint) assemble forms with new coefficients, a mixed coefficient and a new Constant:

• Before: the loops are built again on every call
  After: one set of loops; all 7 cases fail without the coefficient substitution

test_caching_distinguishes_coefficient_positions (new) assembles Tensor(f*m) + Tensor(2*g*m) + Tensor(3*c*m) with c = f, then c = g:

• Before: the diagonal is 24 M for both (the second reuses the first kernel)
  After: 24 M and 51 M

test_hessian.py::test_burgers[nlvs] failed in CI with KeyError: 'function_13_copy' and passes with the unique copy names.

Repeated interpolation on BDM2 16×16 through Vq = V.quadrature_space, per call:

	before	after	pyop2
Function(V).interpolate(q)	3.87 ms	1.35 ms	1.03 ms
Cofunction(Vq.dual()).interpolate(r)	7.76 ms	5.16 ms	3.78 ms
forward, new q on each call	4.92 ms	1.60 ms	1.31 ms
adjoint, new r on each call	8.11 ms	5.42 ms	4.17 ms

Merge Danger

Door: two-way

Blast Radius: form assembly

A form that shares loops reads the wrong data if the key misses something the loops depend on. The key holds the parent and index of each space (subspaces of a mixed space compare equal) and the subdomain data (the subset of an Interpolate is not in its signature). The loops of the first form of each key keep its dats alive with the mesh. Slate kernels are compiled again once, because their cache key changes.

---

AI assistance: this PR was written with Claude Code (Claude Opus 5.5).

🤖 Generated with Claude Code

https://claude.ai/code/session_01UGJuwu7dPzKbSmATxWo9Qc

[sghelichkhani]

Hi Pablo, thanks for this, this seems to help the regressions in our timings the most. I tested the three commits cherry-picked onto main with the G-ADOPT regression suite. Yet there are two problems.

The first gives wrong values without any error. _parloops_key holds the space, parent and index of each coefficient, but not the layout of its data, so a subfunction of a mixed Function (a view into the parent's data) and a standalone Function on the same subspace share a loop:

from firedrake import *

mesh = UnitSquareMesh(10, 10)
x, y = SpatialCoordinate(mesh)
W = FunctionSpace(mesh, "DG", 1) * FunctionSpace(mesh, "CG", 2)

# Standalone Functions on the two subspaces of W
u1 = Function(W.sub(0)).interpolate(x + 2*y)
v1 = Function(W.sub(1)).interpolate(sin(4*pi*x)*cos(4*pi*y))
print(assemble((u1 - v1)**2*dx))

# The same values in the subfunctions of a Function on W
w = Function(W)
u2, v2 = w.subfunctions
u2.assign(u1)
v2.assign(v1)
print(assemble((u2 - v2)**2*dx))  # should print the same number

On d4209ce this prints 2.9112 and then 0.1828 (pyop2 gives 2.9112 twice, and on 3 ranks the second value can be nan). It also makes test_checkpoint_comm_disk_checkpointing_parallel and test_sub_comm_disk_checkpointing_parallel fail on 3 ranks. Adding the axis trees of the coefficient dats to the key fixes both.

The second is speed. assemble rebuilds the map from loop argument names to coefficient data on every call, and iter_constants walks the whole UFL form each time. Solvers assemble the same forms thousands of times, so in our CI 2d_cylindrical went from 2099 s to 3252 s. Building the map once per assembler removes this.

There is a PR fixing each of these against your branch: #5573 for the key and #5574 for the speed.