[sghelichkhani]

Fixes wrong values from shared loops in #5566. A subfunction of a mixed Function is a view into the dat of its parent, and a Function on the same subspace has a dat of its own. Both get the same key, so one reuses the loop built for the other layout.

from firedrake import *

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

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))

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

Before this change it prints 2.9112 and then 0.1828, with it 2.9112 twice. It also fixes test_checkpoint_comm_disk_checkpointing_parallel and test_sub_comm_disk_checkpointing_parallel on 3 ranks, which fail on the base branch.

A related case I have not hit: ParloopBuilder builds a different loop when a coefficient shares its buffer with the output tensor, and that is not in the key either.

[pbrubeck]

I know that 5566 was still failing CI, should we add the example from the PR description to the tests? Or is it covered already by the existing tests?

[sghelichkhani]

Done, I added the example as test_loops_shared_between_coefficient_layouts in test_assemble.py. It fails without the change to the key. The existing tests only caught this indirectly, through the disk checkpointing tests on 3 ranks.

[sghelichkhani]

About the CI here: "Basic repository checks" fails because the base of this PR is your branch and not main or release (check-config: error: argument mode: invalid choice: 'INVALID_BASE_REF'), and "Test default" and "Build documentation" fail on #5566 itself as well, so the CI of this PR does not say much until the base passes. test_assemble.py passes locally on this branch, in serial and on 2 ranks.

[pbrubeck]

About the CI here: "Basic repository checks" fails because the base of this PR is your branch and not main or release (check-config: error: argument mode: invalid choice: 'INVALID_BASE_REF'), and "Test default" and "Build documentation" fail on #5566 itself as well, so the CI of this PR does not say much until the base passes. test_assemble.py passes locally on this branch, in serial and on 2 ranks.

I meant CI failing on 5566. You need to add the base:main label

[pbrubeck]

Are these still needed?

[pbrubeck]

The signature should include ufl-specific information, and the loop nuances should be captured by the axis_trees

[pbrubeck]

Maybe they are only needed for the arguments?