TSFC: support same-mesh ufl.Interpolate in Form kernels by pbrubeck · Pull Request #5258 · firedrakeproject/firedrake

pbrubeck · 2026-07-16T08:51:16Z

Summary

TSFC now compiles same-mesh ufl.Interpolate nodes in normal Form kernels, and this PR ports the feature from PyOP2 to pyop3.

Standalone Interpolate and Form(Interpolate) share one code generation path.
Same-mesh interpolation across a MixedFunctionSpace supports mat_type="aij" natively.
Interpolation into a VertexOnlyMesh uses the same TSFC code generation.
Cross-mesh interpolation stays on BaseFormAssembler.

SameMeshInterpolator no longer builds its own op3.loops. It sends the interpolation through the form assembler:

 SameMeshInterpolator._get_callable
-  _build_interpolation_callables      # hand-built op3.loop per sub-expression
+  get_form_assembler(Interpolate, access=...)
+    tsfc_interface.compile_form(..., access)
+      as_pyop3_local_kernel(ast, nargs, access)
+        op3.Function(..., requires_zeroed_output_arguments=True)

MIN and MAX are reduced into the output as MIN_WRITE and MAX_WRITE. TSFC kernels add their result to the output argument, so pyop3 starts these temporaries from zero (see #5543).

TwoFormAssembler._apply_bc now takes the field label of a BC from the BC's own space. Before, a rectangular form such as an interpolation from W = V * V into V failed.

quadrature_space is now a cached property. The adjoint weight, the reciprocal of the number of cells in the iteration set whose closure contains each node, is cached on the mesh, keyed by the dual space and the subset.

Stacked on #5543, which adds MIN/MAX reductions to pyop3 kernel calls.
Depends on firedrakeproject/fiat#294.

Fixes #5468

Evidence

Before (merge of main without the fixes): 7 failures in tests/firedrake/regression/test_interpolate.py

test_subset_interpolation_reduction[min|max]     -> all entries stay at ±finfo.max
test_interpolate_periodic_coords_reduction[...]  -> [inf]
test_mixed_space_bcs                             -> MultipleUniqueEntriesException

After:

serial: interpolation, cross-mesh, vertex-only, MG grid transfer, pyop3 loop and TSFC tests  1006 passed
serial: BC, mixed, real-space, matrix and assembly regression tests                           740 passed
mpiexec -n 2  -m "parallel[2]" interpolation and vertex-only tests                            32 passed
mpiexec -n 3  -m "parallel[3]" interpolation and vertex-only tests                           271 passed

test_subset_adjoint_interpolation (new, 1 and 2 ranks) compares the matrix-free adjoint with the transposed aij matrix, with and without a subset. It fails if the subset is left out of the cache key.

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

	before	after
`Function(V).interpolate(q)`	4.01 ms	3.87 ms
`Cofunction(Vq.dual()).interpolate(r)`	15.1 ms	7.8 ms

Merge Danger

Door: two-way

Blast Radius: interpolation

All same-mesh interpolation, and the Form kernels that contain it, now go through TSFC and the form assembler. op3.Function gets a new attribute, so kernel caches from before this change must be cleared with firedrake-clean. V.quadrature_space() is now the property V.quadrature_space.

The CI commits marked DROP BEFORE MERGE install the FIAT branch and must be dropped before merging.

Contribution assisted by Claude Code (Claude Opus 5.5).

🤖 Generated with Claude Code

pbrubeck · 2026-07-16T11:21:11Z

This PR should reimplement SameMeshInterpolator as a dispatcher to {Zero|One|Two}FormAssembler. We want to unify the TSFC codepaths for Form assembly and same-mesh interpolation.

The DirichletBC setter called the Interpolator directly. It now keeps the Interpolate expression and assembles it with firedrake.assemble on every update, so the update is annotated for pyadjoint. The Interpolator is cached on the Interpolate expression, so it is built only once. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Kernel calls with MIN_WRITE or MAX_WRITE access now pack a temporary that starts from the identity of the reduction, and unpack it with a "min" or "max" assignment. Assignments of type "min" and "max" now get their intents from assignment_type_as_intent, so they can be put in a loop. Fixes #5524 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

TSFC kernels add their result to the output argument. A MIN_WRITE or MAX_WRITE temporary that starts from the identity of the reduction therefore stays at the identity. op3.Function now takes requires_zeroed_output_arguments, like the PyOP2 Kernel did. When it is set, these temporaries start from zero. as_pyop3_local_kernel sets it for TSFC kernels. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…ct/firedrake into pbrubeck/fix-minmax-write

…tsfc # Conflicts: # firedrake/interpolation.py # firedrake/tsfc_interface.py

quadrature_space is now a cached property. The adjoint weight, the reciprocal of the number of cells in the iteration set whose closure contains each node of the dual argument, is cached on the mesh, keyed by the dual space and the subset. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UGJuwu7dPzKbSmATxWo9Qc

@pbrubeck

#5127) * Symbolic coarsening of DirichletBC, and bc data interpolation from other spaces * BCBase.coefficients * Apply suggestion from @pbrubeck * Tape DirichletBC's internal interpolation for pyadjoint DirichletBC data on a different FunctionSpace was interpolated via a raw Interpolator.assemble() call, bypassing pyadjoint annotation. Route it through the top-level assemble() instead, stashing the Interpolate object so BaseFormAssembler reuses its cached Interpolator. DirichletBCBlock's adjoint also assumed the BC data lived on the same space as the BC itself; add the cross-space case, transforming the adjoint value through the adjoint of the (linear) interpolation operator. Add a regression test: a time-dependent heat equation with BC data interpolated from a different space at each step, verified with a Taylor test. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * edit AGENTS.md * Tape DirichletBC's internal interpolation for pyadjoint DirichletBC data on a different FunctionSpace was interpolated via a raw Interpolator.assemble() call, bypassing pyadjoint annotation. Route it through the top-level assemble() instead, stashing the Interpolate object so BaseFormAssembler reuses its cached Interpolator. DirichletBCBlock's adjoint also assumed the BC data lived on the same space as the BC itself; add the cross-space case, transforming the adjoint value through the adjoint of the (linear) interpolation operator. Fix two further gaps found while making a reused-DirichletBC variant of the new test actually differentiable, rather than merely documenting it as broken: - DirichletBC.function_arg's setter now keeps the FloatingType's _ad_args in sync with the assigned value, so reusing a BC via set_value()/mutating its data re-tapes against the current value instead of the one captured at __init__. - NonlinearVariationalSolveBlock._forward_solve now applies its own recovered, per-step bcs to the cached forward_nlvs solver before solving, instead of solving with whatever the shared DirichletBC objects' live state currently happens to hold. Add a regression test, parametrized over constructing a fresh DirichletBC each timestep vs. reusing one via NonlinearVariationalSolver, both verified with a Taylor test. Document the underlying patterns in AGENTS.md in generalized form, so they transfer to future pyadjoint (or similar record/replay) bugs beyond this specific one. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Apply suggestion from @pbrubeck * AGENTS.md: document the _ad_annotate_* mixin pattern Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * AGENTS.md: trim and reorder pyadjoint debugging lessons Sort from cheapest diagnostic wins to deepest architectural gotchas, and cut each lesson down to headline + one anchoring example. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Compose DirichletBC adjoint with its taped boundary-value interpolation Make DirichletBCBlock depend on the Function the BC actually applies (bc._function_arg) rather than the raw user-supplied value. The interpolation (or, for elements without interpolation support, the projection) of the boundary value is re-taped through the annotated assemble/project when the BC is added as a dependency of another block, so its own AssembleBlock/ProjectBlock supplies the adjoint, TLM, and recompute. This deletes the hand-rolled interpolation pullback and the dead constant-handling machinery from DirichletBCBlock. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * AGENTS.md: rewrite pyadjoint notes as composition-first design principles Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * AGENTS.md: make pyadjoint principles generic Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Annotate Projector.project instead of rebuilding a Projector per update DirichletBC's projection fallback goes back to one cached Projector per boundary value: rebuilding a solver stack on every function_arg access crashes when the accessor fires inside PCSetUp_MG's coarsening callback (gmg on extruded hierarchies, where NCE/NCF cannot interpolate). The projection stays on the tape through a new ProjectorMixin that annotates ProjectorBase.project, taping one ProjectBlock per call against the current source value. Also locate the DirichletBCBlocks in test_bdy_control by type rather than by tape position, since the taped boundary-value interpolation now precedes them. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * AGENTS.md: branch pairing, CI triage, and PETSc-callback lessons Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * AGENTS.md: add design and debugging method Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Add CLAUDE.md importing AGENTS.md Claude Code auto-loads CLAUDE.md from the repository root; the @AGENTS.md import inlines the agent instructions so they reach every session without relying on convention or memory. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * adjoint AGENTS.md * Apply suggestion from @pbrubeck * Apply suggestion from @pbrubeck * Drop unrelated AGENTS.md files * Remove unused firedrake.matrix import from bcs.py The import was left over from a merge resolution. firedrake.matrix imports DirichletBC from firedrake.bcs, so the import was circular and made `import firedrake` fail. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Clone DirichletBCs in the replay problem The problem that NonlinearVariationalSolver clones for replay copied the coefficients of F and J but kept the user's DirichletBCs, so a replay applied whatever boundary value the BC held at that time. The clone now reconstructs each DirichletBC with a copy of its boundary value, shared with F when the value is also a coefficient of F, and the checkpointed value is assigned before each replay. This replaces the reassignment of the cloned problem's bcs on every solve. Also shorten the comments on how DirichletBC is taped, and rename a local variable in EquationBCSplit.coefficients to match NonlinearVariationalSolver. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Probe the dual basis to choose interpolation or projection in DirichletBC The DirichletBC setter tried to interpolate the boundary value and fell back to projection on any NotImplementedError raised inside assemble. It now checks V.finat_element.dual_basis first, as in #5258, and builds either the interpolation or the Projector. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Add second-order Taylor test for time-dependent Dirichlet BC --------- Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

leo-collins · 2026-10-29T14:51:47Z

+
+    @cached_on(
+        get_obj=lambda self: get_mesh_topologies(self.ufl_interpolate),
+        get_key=lambda self: (self.dual_arg.function_space(), self.subset),


Should this be keyed on the source mesh too? In the submesh case

It is cached on the source mesh. What does get_mesh_topologies return?

Is it going to cache the weights on all meshes, including submeshes? If so, then if we have submesh S of a parent mesh M, an adjoint interpolation M -> M will store weights for the whole mesh on both M and S. Then if we do an adjoint interpolation (on the same dual space) S -> M we would use the weights for M.

It's now cached on the source mesh

leo-collins · 2026-10-29T14:54:55Z

+    @cached_property
    def quadrature_space(self):


I suppose this is public API so should be deprecated correctly.

It's a very young API, I don't think anyone has been messing with this.

I can change it to a cached method if that makes people happier.

get_form_assembler already passes allocation_integral_types only to matrix assemblers, so the BaseForm visitor does not need to filter it. Nothing imports runtime_quadrature_element from pointeval_utils. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

An output that aliases an input is assembled into a temporary. The temporary must start as a copy of the output when the loop misses some target cells, otherwise their nodes are zeroed on copy-out. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

A fused interpolation must reach every target cell. Instead of counting the cells of the interpolation loop on every rank, check the subset and the submesh hierarchy of the source and target meshes. This avoids an MPI reduction while the BaseForm DAG is traversed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-authored-by: Leo Collins <leocollins511@gmail.com>