Prevent constant Form reassembly

[pbrubeck]

Note

Disclaimer: This implementation plan was formulated by a core Firedrake developer in collaboration with Gemini.

Summary of Implementation Plan: Avoiding Reassembly into the Same Tensor

This plan introduces a change to PyOP2 to track state changes in global/scalar datasets (Global) alongside standard Dat dependencies. The ParloopAssembler runtime is then updated to leverage these versions, allowing it to skip core execution loops when variational forms, coefficients, constants, and boundary conditions match a previous execution on the same target tensor. This approach is completely memory-safe via a weakref.WeakKeyDictionary caching pattern on the target tensor.

---

1. Implement .dat_version Tracking in PyOP2 Dat and Global

• Files to modify: pyop2/types/dat.py (and/or pyop2/types/abstract.py depending on PyOP2 architecture)

# Inside the AbstractDat or Base class (or added specifically to Dat and Global classes):
class AbstractDat(object):
    def __init__(self, *args, **kwargs):
        # Existing initialization logic...
        self._dat_version = 0

    @property
    def dat_version(self):
        return self._dat_version

    @property
    def data(self):
        # Increment version whenever mutable data access is invoked
        self._dat_version += 1
        return self._data_setter_or_buffer  # Existing PyOP2 logic returning the numpy/PETSc array pointer

---

2. Implement the Assembly Signature Helper

• File to modify: firedrake/assemble.py

def _get_assembly_signature(self, tensor, form, bcs):
    from firedrake.utils import extract_firedrake_constants
    
    tensor_version = tensor.dat.dat_version
    coeff_versions = tuple(c.dat.dat_version for c in form.coefficients())
    constant_versions = tuple(c.dat.dat_version for c in extract_firedrake_constants(form))
    
    if not bcs:
        bc_signature = ()
    else:
        bc_signature = tuple(
            (id(bc), getattr(getattr(bc, "function", None), "dat", type('Mock', (object,), {"dat_version": 0})()).dat_version)
            for bc in bcs
        )
    
    return (tensor_version, coeff_versions, constant_versions, bc_signature)

---

3. Integrate WeakKeyDictionary Cache into ParloopAssembler

• File to modify: firedrake/assemble.py

import weakref

class ParloopAssembler(object):
    def __init__(self, form, *args, **kwargs):
        self.form = form
        self._assembled_tensor_versions = weakref.WeakKeyDictionary()

    def assemble(self, tensor, bcs=None):
        form = self.form
        cached_signature = self._assembled_tensor_versions.get(tensor, None)
        current_signature = self._get_assembly_signature(tensor, form, bcs)

        if cached_signature == current_signature:
            return tensor

        self._assemble_core(tensor, bcs=bcs)
        self._assembled_tensor_versions[tensor] = current_signature

        return tensor