From 9e1047300f9eb0e9607d884a38165d9e8099061e Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sun, 19 Jul 2026 00:33:48 +0100 Subject: [PATCH 01/23] Sum factorisation on simplices --- tests/tsfc/test_codegen.py | 53 ++++++++++++++++++++++++++++++++++- tests/tsfc/test_pickle_gem.py | 14 +++++++++ tsfc/loopy.py | 23 ++++++++++++--- 3 files changed, 85 insertions(+), 5 deletions(-) diff --git a/tests/tsfc/test_codegen.py b/tests/tsfc/test_codegen.py index 8d0bc79655..7dd97e2ba1 100644 --- a/tests/tsfc/test_codegen.py +++ b/tests/tsfc/test_codegen.py @@ -1,6 +1,7 @@ +import numpy import pytest -from gem import impero_utils +from gem import gem, impero_utils from gem.gem import Index, Indexed, IndexSum, Product, Variable @@ -24,6 +25,56 @@ def gencode(expr): assert len(gencode(e1).children) == len(gencode(e2).children) +def test_jagged_index_codegen(monkeypatch): + import islpy as isl + import loopy as lp + import tsfc.loopy + + # Execute the generated code so we check the numbers, not just the loop bounds + monkeypatch.setattr(tsfc.loopy, "target", lp.ExecutableCTarget()) + + n = 4 + extent = n + 1 + npts = 3 + ndof = (n + 1) * (n + 2) // 2 + + rng = numpy.random.default_rng(7) + # Table zero-padded outside the simplex lattice p + q > n, and a + # clamped Morton index table for the coefficient gather + B = rng.random((extent, extent, npts)) + morton = numpy.zeros((extent, extent), dtype=gem.uint_type) + for p_, q_ in numpy.ndindex(morton.shape): + if p_ + q_ > n: + B[p_, q_] = 0.0 + else: + morton[p_, q_] = (p_ + q_) * (p_ + q_ + 1) // 2 + q_ + c = rng.random(ndof) + + i = Index(name="i", extent=npts) + p = Index(name="p", extent=extent) + q = gem.JaggedIndex(name="q", extent=extent, parents=(p,)) + + dof = gem.VariableIndex(Indexed(gem.Literal(morton, dtype=gem.uint_type), (p, q))) + integrand = Product(Indexed(Variable("c", (ndof,)), (dof,)), + Indexed(gem.Literal(B), (p, q, i))) + expr = IndexSum(integrand, (p, q)) + + u = Variable("u", (npts,)) + impero_c = impero_utils.compile_gem([(Indexed(u, (i,)), expr)], (i, p, q)) + args = [lp.GlobalArg("u", dtype=numpy.float64, shape=(npts,)), + lp.GlobalArg("c", dtype=numpy.float64, shape=(ndof,))] + knl, _ = tsfc.loopy.generate(impero_c, args, numpy.float64) + + # The jagged loop must have a domain parametrized by its parent iname + assert any(dom.get_var_names(isl.dim_type.param) + for dom in knl.default_entrypoint.domains) + + u_out = numpy.zeros(npts) + knl(c=c, u=u_out) + u_ref = numpy.tensordot(c[morton], B, axes=((0, 1), (0, 1))) + assert numpy.allclose(u_out, u_ref, rtol=1e-14) + + if __name__ == "__main__": import os import sys diff --git a/tests/tsfc/test_pickle_gem.py b/tests/tsfc/test_pickle_gem.py index beb101f912..b68905cb0d 100644 --- a/tests/tsfc/test_pickle_gem.py +++ b/tests/tsfc/test_pickle_gem.py @@ -17,6 +17,20 @@ def test_pickle_gem(protocol): assert repr(expr) == repr(unpickled) +@pytest.mark.parametrize('protocol', range(3)) +def test_pickle_jagged_index(protocol): + p = gem.Index(name='p', extent=4) + q = gem.JaggedIndex(name='q', extent=4, parents=(p,)) + expr = gem.IndexSum(gem.Indexed(gem.Variable('A', (4, 4)), (p, q)), (p, q)) + + unpickled = pickle.loads(pickle.dumps(expr, protocol)) + assert repr(expr) == repr(unpickled) + up, uq = unpickled.multiindex + assert isinstance(uq, gem.JaggedIndex) + assert uq.extent == 4 + assert uq.parents == (up,) + + @pytest.mark.parametrize('protocol', range(3)) def test_listtensor(protocol): expr = gem.ListTensor([gem.Variable('x', ()), gem.Zero()]) diff --git a/tsfc/loopy.py b/tsfc/loopy.py index 6826f0b672..a9d8ac8508 100644 --- a/tsfc/loopy.py +++ b/tsfc/loopy.py @@ -123,6 +123,7 @@ def __init__(self, target=None): self.indices = {} # indices for declarations and referencing values, from ImperoC self.active_indices = {} # gem index -> pymbolic variable self.index_extent = OrderedDict() # pymbolic variable for indices -> extent + self.index_parents = {} # iname -> parent inames bounding a jagged index self.gem_to_pymbolic = {} # gem node -> pymbolic variable self.name_gen = UniqueNameGenerator() self.target = target @@ -257,7 +258,7 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name instructions, event_name, preamble = profile_insns(kernel_name, instructions, log) # Create domains - domains = create_domains(ctx.index_extent.items()) + domains = create_domains(ctx.index_extent.items(), ctx.index_parents) # Create loopy kernel knl = lp.make_kernel( @@ -276,16 +277,23 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name return knl, event_name -def create_domains(indices): +def create_domains(indices, index_parents=None): """ Create ISL domains from indices :arg indices: iterable of (index_name, extent) pairs + :arg index_parents: optional mapping from index_name to a tuple of parent + index names; the domain of a jagged index is parametrized by its + parents, with upper bound extent minus the sum of the parents. :returns: A list of ISL sets representing the iteration domain of the indices.""" domains = [] for idx, extent in indices: - inames = isl.make_zero_and_vars([idx]) - domains.append(((inames[0].le_set(inames[idx])) & (inames[idx].lt_set(inames[0] + extent)))) + parents = index_parents.get(idx, ()) if index_parents else () + inames = isl.make_zero_and_vars([idx], parents) + bound = inames[0] + extent + for parent in parents: + bound = bound - inames[parent] + domains.append(((inames[0].le_set(inames[idx])) & (inames[idx].lt_set(bound)))) if not domains: domains = [isl.BasicSet("[] -> {[]}")] @@ -316,6 +324,13 @@ def statement_for(tree, ctx): assert extent idx = ctx.name_gen(ctx.index_names[tree.index]) ctx.index_extent[idx] = extent + if isinstance(tree.index, gem.JaggedIndex) and \ + all(parent in ctx.active_indices for parent in tree.index.parents): + # Tighten the loop bound of a jagged index nested inside its parents. + # If a parent loop is not in scope, the rectangular bound `extent` + # remains correct: jagged expressions are zero-padded. + ctx.index_parents[idx] = tuple(ctx.active_indices[parent].name + for parent in tree.index.parents) with active_indices({tree.index: p.Variable(idx)}, ctx) as ctx_active: return statement(tree.children[0], ctx_active) From 838c146a2bf0dfdf671ab16ac48d6add90416132 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 20 Jul 2026 09:36:33 +0100 Subject: [PATCH 02/23] TSFC: sum-factorized matrix-free DG residual on simplices (milestone 2) Extends tsfc/fem.py's translate_coefficient and translate_argument to take a sum-factorized (Duffy/lattice) tabulation path for simplicial DG Legendre elements under dx(scheme="collapsed"), targeting O(p^{d+1}) flops for the matrix-free residual instead of the dense O(p^{2d}). The lattice multiindex from Legendre.duffy_evaluation is gathered against coefficients and scattered back to the flat dof index through FIAT's existing Morton dof numbering (new morton_forward_table / morton_inverse_table in FIAT.expansions), entirely inside fem.py: element.index_shape and argument_multiindices stay flat, so no driver.py or kernel_interface changes are needed. Co-Authored-By: Claude Sonnet 5 --- DESIGN.md | 221 +++++++++++++ tests/firedrake/regression/test_quadrature.py | 34 ++ tests/tsfc/test_codegen.py | 71 +++++ tests/tsfc/test_sum_factorisation.py | 28 +- tsfc/fem.py | 290 ++++++++++++++---- 5 files changed, 589 insertions(+), 55 deletions(-) create mode 100644 DESIGN.md diff --git a/DESIGN.md b/DESIGN.md new file mode 100644 index 0000000000..991e5a088e --- /dev/null +++ b/DESIGN.md @@ -0,0 +1,221 @@ +# Sum factorization on simplices: status and milestone-2 design + +Companion to `PLAN.md`. Everything in "Implemented" is validated by tests. + +## Implemented (milestone 1 + jagged-loop infrastructure) + +### FIAT (`FIAT/expansions.py`) + +* `dubiner_jacobi_parameters(codim, m, variant)` and `dubiner_norm2(d, m, i, variant)`: + shared helpers extracted from `dubiner_recurrence` (behavior-preserving). +* `principal_functions(n, eta, axis, order, variant)`: 1D tabulations of the + Karniadakis–Sherwin principal functions `G[m, i](eta) = norm * w(eta)^m * g_{m,i}(eta)` + with `w = (1 - eta)/2`, as tables `'V'` (values), `'D'` (d/d eta), `'W'` + (`w^{m-1} g norm`, zeroed at `m = 0`), `'tD'` (`(1+eta)/2 * D`). +* `ExpansionSet.tabulate_duffy(n, eta_pts, order, cell)`: the separable + tabulation on collapsed points. On the reference simplex, the Dubiner basis + factorizes exactly as + `phi_{(i_1..i_d)}(eta) = scale * prod_t G_t[m_t, i_t](eta_t)` with + `m_t = i_1 + ... + i_{t-1}`. First derivatives use the closed-form chain rule + `d eta_t / d xi_k = ((1+eta_t)/2)^{[k>t]} * prod_{u>t} 1/w_u` (k >= t), which + yields exactly `k` separable terms for `d phi / d xi_k`; the affine cell map + is applied on top. Raises `NotImplementedError` for `order > 1` and for C0 + (`continuity is not None`) expansion sets. + +### finat (`finat/point_set.py`, `finat/spectral.py`) + +* `CollapsedTensorProductPointSet`: 1D factor point sets in collapsed + coordinates on `[0, 1]`, mapped to the simplex by the Duffy map + `x_t = eta_t * prod_{u>t} (1 - eta_u)`. +* `Legendre.duffy_evaluation(order, ps, entity=None)`: returns + `(multiindex, result)` where `multiindex` enumerates the basis by lattice + indices `(i_1, ..., i_d)` (a tuple of `gem.JaggedIndex`, see below) and + `result[alpha]` is a scalar gem expression built from `gem.Literal` 1D tables + contracted per axis. The weight exponents `m_t` are gem expressions: `0`, + `i_1`, then `VariableIndex` lookups into clamped uint index tables. + +**Zero-padding invariant** (load-bearing for everything below): lattice indices +range over the rectangular box `(degree+1)^d`; all literal index tables are +clamped with `numpy.minimum`, and out-of-lattice entries (`sum(i_t) > degree`) +tabulate to exactly zero. Memory safety and correctness never depend on jagged +loop bounds — jaggedness is purely a flop optimization. + +### gem (`gem/gem.py`, `gem/interpreter.py`) + +* `JaggedIndex(Index)`: free index with a `parents` tuple; iteration bound + `0 <= i < extent - sum(parents)`. `.extent` remains the static rectangular + bound, so every consumer that ignores jaggedness stays correct (via the + zero-padding invariant). Picklable; exported in `__all__`. +* `gem.interpreter._evaluate_indexed`: repaired the bit-rotted `VariableIndex` + path, including a gather path for index expressions with free indices. + +### tsfc (`tsfc/loopy.py`) + +* `LoopyContext.index_parents`: iname -> parent inames, recorded in + `statement_for` when an `imp.For` loop index is a `JaggedIndex` and all its + parents are active (i.e. the loop is nested inside them); otherwise the + rectangular bound is kept (correct by the invariant). +* `create_domains(indices, index_parents=None)`: emits parametrized ISL sets + `[p] -> {[q] : 0 <= q < E - p}` for jagged inames. Loopy nests these domains + automatically; no loopy changes were needed (verified experimentally first). +* `ComponentTensor` materialization stays rectangular on purpose: temporaries + are always fully initialized, so a jagged write can never leave garbage that + a rectangular read later observes. + +### Tests + +* `test/FIAT/unit/test_polynomial.py`: `test_tabulate_duffy` (values + + gradients vs `_tabulate_on_cell`, dims 1–3, variants None/dual, degrees + 0/1/4, points including the collapsed vertex), `test_principal_functions_bubble`, + `test_morton_tables` (`morton_forward_table`/`morton_inverse_table` agree + with `morton_index` and are mutual inverses on the simplex lattice). +* `test/finat/test_point_evaluation.py`: `test_duffy_evaluation` vs dense + `basis_evaluation` through the gem interpreter, checking Morton flat-index + agreement and exact zeros outside the lattice. +* `tests/tsfc/test_codegen.py::test_jagged_index_codegen`: compiles the 2D + jagged Morton-gather contraction through `compile_gem` -> `tsfc.loopy.generate`, + executes the C kernel, checks the parametrized ISL domain exists and the + numbers match numpy. +* `tests/tsfc/test_pickle_gem.py::test_pickle_jagged_index`. + +## Milestone 2: O(p^d)-per-dof matrix-free DG residual — implemented (Route B) + +A DG residual action has three phases per cell: + +1. **Forward transform:** evaluate `u` (and `grad u`) at quadrature points from + coefficients `c`. Sum-factorized on collapsed points this is a sequence of d + per-axis contractions with jagged intermediate temporaries, e.g. in 3D + `T1[p, q, k] = sum_r C[p, q, r] * G3[p+q, r, k]`, + `T2[p, j, k] = sum_q T1[p, q, k] * G2[p, q, j]`, + `u[i, j, k] = sum_p T2[p, j, k] * G1[p, i]` — O(p^{d+1}) total. +2. **Pointwise:** multiply by quadrature weights, geometry, coefficients. +3. **Backward transform:** contract against the test function tables (the + transpose sweep), scattering into the residual vector. + +### (a) Collapsed quadrature rule — implemented + +`finat/quadrature.py` now has `CollapsedTensorProductQuadratureRule` (1D +Gauss–Jacobi factor rules in collapsed coordinates; axis `u` carries the Jacobi +weight `(1 - eta_u)^u`, which absorbs the Duffy Jacobian, so the simplex +weights are per-axis products) and a `scheme="collapsed"` branch in +`make_quadrature` building it via `collapsed_gauss_jacobi_quadrature`. Since +`tsfc/fem.py::get_quadrature_rule` passes the UFL measure's scheme metadata +straight to `make_quadrature`, `dx(scheme="collapsed")` already produces the +structured rule with **no** tsfc changes. Tested against FIAT's canonical +collapsed scheme on all monomials up to the requested degree +(`test/finat/test_quadrature.py::test_collapsed_quadrature`). + +### (b) fem.py integration — Route B, no driver/kernel-interface changes + +The standard path in `tsfc/fem.py` calls `element.basis_evaluation(order, ps, +entity)` and contracts the resulting `(ndof,)`-shaped tables with the +element's flat basis index (`element.get_indices()`), and `translate_argument` +extracts one flat entry with `ctx.argument_multiindices[number]`. The local +element tensor's shape (`element.index_shape` in +`kernel_interface/common.py::prepare_arguments`) and `argument_multiindices` +are flat and ndof-based *everywhere else in the kernel-interface/PyOP2 stack*, +so the original plan of making `argument_multiindices` itself a lattice +multiindex (see the old Route-B write-up below) would have changed the local +tensor's shape and broken that contract. The implementation instead keeps +`element.index_shape` and `argument_multiindices` exactly as they are today, +and confines the lattice multiindex to the tabulation step alone: + +* `_use_sum_factorisation(element, ctx)` (`tsfc/fem.py`) gates the whole path: + `element` must be `finat.spectral.Legendre`, `ctx.point_set` a + `CollapsedTensorProductPointSet` (i.e. the measure requested + `dx(scheme="collapsed")`), the integral must be over the cell interior, and + `ctx.unsummed_coefficient_indices` must be empty (macrocells, which + `duffy_evaluation` already rejects, are the only case that sets it). +* `_duffy_evaluation(element, mt, ctx, entity_id)` calls + `element.duffy_evaluation(mt.local_derivatives, ctx.point_set, + (ctx.integration_dim, entity_id))` and filters to `sum(alpha) == + mt.local_derivatives`, exactly mirroring the filtering the standard path + applies to `basis_evaluation`'s output. +* **Forward transform (`translate_coefficient`).** `_contract_dof_index` + builds a forward Morton lookup table (`FIAT.expansions.morton_forward_table`, + shape `(degree+1,)^d`, clamped to a valid dof so out-of-lattice reads are + merely wasted, never out of bounds — they always multiply a zero + tabulation), gathers `vec[VariableIndex(table[multiindex])]`, and hands + `IndexSum(Product(duffy[alpha], vec_r), multiindex)` to + `gem.optimise.contraction`, exactly as originally planned: the `m_t` + `VariableIndex` couplings inside `duffy_evaluation`'s own expression make + the per-axis free-index sets nested (`{i_1} ⊂ {i_1, i_2} ⊂ ...`), so + `contraction` finds the innermost-axis-first Karniadakis–Sherwin sweep by + itself — no bespoke sum-factorization code needed. The result is wrapped + back into a `gem.ComponentTensor` over `element.get_value_indices()` (empty + for the scalar `Legendre` element), so it slots into `fiat_to_ufl` exactly + like a standard dense tabulation would. +* **Backward transform (`translate_argument`).** `_scatter_to_dof_index` goes + the other way: it introduces one *fresh* flat dof index `r` (the same free + index `argument_multiindex` will later pick a single value of — nothing + about `argument_multiindices` construction changes), builds the *inverse* + Morton table (`FIAT.expansions.morton_inverse_table`, shape `(ndof, d)`) to + get per-axis lookups `i_t(r)`, and substitutes + `multiindex[t] -> VariableIndex(inverse_table[:, t][r])` throughout + `duffy_evaluation`'s expression tree via + `gem.node.MemoizerArg(gem.optimise.filtered_replace_indices)` — the same + substitution mechanism `translate_argument`/`translate_coefficient` already + use for canonical quadrature-point reordering. `filtered_replace_indices` + recurses into `VariableIndex.expression` (`gem/optimise.py`'s + `_replace_indices_atomic`), so this also correctly rewrites the nested `m_t` + lookups (which are themselves `VariableIndex` expressions built from + `multiindex[:t]`) into r-indexed double lookups, with no duplicated + tabulation logic. The result, wrapped in `gem.ComponentTensor(..., (r,))`, + is a dense `(ndof,)`-shaped table — indistinguishable, from + `fiat_to_ufl`/`prepare_arguments`'s point of view, from the standard dense + tabulation. `gem.optimise.contraction` never runs on this side (there is no + sum to hoist yet at this stage); the sum-factorized quadrature contraction + happens later, per dof, when `vanilla.py`/`spectral.py` process the + quadrature `IndexSum` — the collapsed quadrature's own per-axis structure is + what still delivers the O(p) win per axis there. + +Both helpers were validated to ~1e-13/1e-14 against FIAT's dense +`tabulate()`, via compiled-and-executed loopy kernels, for values and first +derivatives on triangles and tetrahedra +(`tests/tsfc/test_codegen.py::test_duffy_scatter_and_contract`), and end to +end through `firedrake.assemble` (residuals and matrices, `dx` vs +`dx(scheme="collapsed")`, on triangle and tetrahedron meshes, degrees 1 and 3: +`tests/firedrake/regression/test_quadrature.py::test_collapsed_quadrature_sum_factorisation`). + +### (c) Basis-index integration route — Route B chosen + +The element's flat basis index is Morton-ordered (`FIAT.expansions.morton_index`, +using `morton_index2`/`morton_index3` = total-degree-major), while the +factorization is indexed by the lattice multiindex. Three routes were +considered: + +* **Route A — layer-wise `Concatenate`.** Reuse tsfc/spectral.py's + `Concatenate`/`unconcatenate` machinery by splitting the basis into + contiguous Morton layers of fixed total degree `s`. Rejected: the layer + decomposition does not align with the per-axis contraction structure (the + sweeps contract one lattice axis at a time, not one total-degree layer at a + time), so it buys the wrong factorization. + +* **Route B — Morton gather/scatter via `VariableIndex` (chosen; see (b)).** + Keeps FIAT's dof ordering, `element.index_shape`, and + `argument_multiindices` completely untouched; the Morton lookup lives + entirely inside `_contract_dof_index`/`_scatter_to_dof_index` in `tsfc/fem.py`. + No `driver.py` or `kernel_interface/*.py` changes were needed at all — the + lattice multiindex never escapes `fem.py`. The indirection costs one uint + load per accumulation (forward) or one uint load per dof (backward), + negligible against the O(p) inner contraction. + +* **Route C — reorder FIAT dofs p-major.** Change `Legendre` + (variant="integral") to lattice-lexicographic dof order so the flat index + becomes `offsets[i_1, ..] + i_d` (affine within each innermost run). + Cleanest kernels, but dof ordering is externally visible (checkpoints, + hand-written index hacks, any test with hard-coded dof numbers) and the + offset table is still a lookup, so the win over Route B is small. Not + pursued; only worth it if profiling shows the Morton gather hurts. + +## Deferred + +* **CG / C0 basis (milestones 3–4):** the C0 recombination makes each basis + function a sum of <= 3 separable members (Sherwin–Karniadakis vertex/edge/face + recombination); `tabulate_duffy` currently raises `NotImplementedError` for + `continuity is not None`. The factored-term representation + (`alpha -> [(coeff, factors), ...]`) was chosen so C0 can extend it by + returning more terms per basis function. +* **Derivative order > 1:** raises `NotImplementedError`. +* **Macro cells** (`is_macrocell()`): raises `NotImplementedError` in + `duffy_evaluation`. diff --git a/tests/firedrake/regression/test_quadrature.py b/tests/firedrake/regression/test_quadrature.py index 225a4b244d..6ea188dfb9 100644 --- a/tests/firedrake/regression/test_quadrature.py +++ b/tests/firedrake/regression/test_quadrature.py @@ -52,3 +52,37 @@ def test_quadrature_element(mesh, family, mat_type, diagonal): a = inner(u, v) * dx assemble(a, mat_type=mat_type, diagonal=diagonal) + + +@pytest.mark.parametrize("cell", ["triangle", "tetrahedron"]) +@pytest.mark.parametrize("degree", [1, 3]) +def test_collapsed_quadrature_sum_factorisation(cell, degree): + """``dx(scheme="collapsed")`` on a simplicial "DG"/variant="integral" + (i.e. `finat.spectral.Legendre`) space must produce the same + assembled residual and matrix as the default dense quadrature, even + though it takes the sum-factorized (Duffy/lattice) tabulation path + in ``tsfc.fem`` rather than the standard dense one. + """ + mesh = {"triangle": UnitSquareMesh(2, 2), + "tetrahedron": UnitCubeMesh(1, 1, 1)}[cell] + V = FunctionSpace(mesh, "DG", degree, variant="integral") + u = TrialFunction(V) + v = TestFunction(V) + w = Function(V) + w.dat.data[:] = np.random.default_rng(0).random(w.dat.data.shape) + + # translate_coefficient path (forward transform): residual with a + # derivative, mixing both the coefficient and argument sum-factorized + # tabulations. + L = inner(grad(w), grad(v)) * dx + L_collapsed = inner(grad(w), grad(v)) * dx(scheme="collapsed") + b = assemble(L) + b_collapsed = assemble(L_collapsed) + assert np.allclose(b.dat.data, b_collapsed.dat.data, rtol=1e-10, atol=1e-10) + + # translate_argument path (backward transform): mass matrix. + a = inner(u, v) * dx + a_collapsed = inner(u, v) * dx(scheme="collapsed") + M = assemble(a).M.values + M_collapsed = assemble(a_collapsed).M.values + assert np.allclose(M, M_collapsed, rtol=1e-10, atol=1e-10) diff --git a/tests/tsfc/test_codegen.py b/tests/tsfc/test_codegen.py index 7dd97e2ba1..e78c238309 100644 --- a/tests/tsfc/test_codegen.py +++ b/tests/tsfc/test_codegen.py @@ -75,6 +75,77 @@ def test_jagged_index_codegen(monkeypatch): assert numpy.allclose(u_out, u_ref, rtol=1e-14) +@pytest.mark.parametrize("cellname,degree", [("triangle", 3), ("tetrahedron", 2)]) +def test_duffy_scatter_and_contract(monkeypatch, cellname, degree): + """Route B of the simplex sum-factorization milestone 2 design: + ``tsfc.fem._scatter_to_dof_index`` (the `translate_argument` path) and + ``tsfc.fem._contract_dof_index`` (the `translate_coefficient` path) + must reproduce the standard dense FIAT tabulation, via the Morton + dof numbering FIAT already uses, from + `finat.spectral.Legendre.duffy_evaluation`'s lattice-indexed, + sum-factorized tabulation. + """ + import loopy as lp + import tsfc.loopy + from FIAT.reference_element import UFCTetrahedron, UFCTriangle + from finat.quadrature import make_quadrature + from finat.spectral import Legendre + from gem import impero_utils + from gem.gem import Index, Indexed, Variable + from gem.optimise import remove_componenttensors + from tsfc.fem import _contract_dof_index, _scatter_to_dof_index + + # Execute the generated code so we check the numbers, not just the loop bounds + monkeypatch.setattr(tsfc.loopy, "target", lp.ExecutableCTarget()) + + cell = {"triangle": UFCTriangle, "tetrahedron": UFCTetrahedron}[cellname]() + element = Legendre(cell, degree) + ndof = element.space_dimension() + + quad_rule = make_quadrature(cell, 2 * degree, scheme="collapsed") + point_set = quad_rule.point_set + point_indices = point_set.indices + point_shape = tuple(index.extent for index in point_indices) + + entity = (cell.get_dimension(), 0) + multiindex, duffy_dict = element.duffy_evaluation(1, point_set, entity) + dense_dict = element._element.tabulate(1, point_set.points) + + rng = numpy.random.default_rng(1) + coefficients = rng.random(ndof) + + for alpha, table_expr in duffy_dict.items(): + dense = dense_dict[alpha].reshape((ndof,) + point_shape) + + # translate_argument path: flat-dof-indexed dense table + scattered = _scatter_to_dof_index(multiindex, {alpha: table_expr}, element)[alpha] + r = Index(extent=ndof) + table, = remove_componenttensors([Indexed(scattered, (r,))]) + u = Variable("u", (ndof,) + point_shape) + impero_c = impero_utils.compile_gem( + [(Indexed(u, (r,) + point_indices), table)], (r,) + point_indices) + args = [lp.GlobalArg("u", dtype=numpy.float64, shape=(ndof,) + point_shape)] + knl, _ = tsfc.loopy.generate(impero_c, args, numpy.float64) + u_out = numpy.zeros((ndof,) + point_shape) + knl(u=u_out) + assert numpy.allclose(u_out, dense, rtol=1e-12, atol=1e-12) + + # translate_coefficient path: contraction against a coefficient vector + c = Variable("c", (ndof,)) + contracted = _contract_dof_index(multiindex, {alpha: table_expr}, element, c)[alpha] + value, = remove_componenttensors([Indexed(contracted, ())]) + v = Variable("v", point_shape) + impero_c = impero_utils.compile_gem( + [(Indexed(v, point_indices), value)], point_indices) + args = [lp.GlobalArg("v", dtype=numpy.float64, shape=point_shape), + lp.GlobalArg("c", dtype=numpy.float64, shape=(ndof,))] + knl, _ = tsfc.loopy.generate(impero_c, args, numpy.float64) + v_out = numpy.zeros(point_shape) + knl(v=v_out, c=coefficients) + v_ref = numpy.tensordot(coefficients, dense, axes=(0, 0)) + assert numpy.allclose(v_out, v_ref, rtol=1e-12, atol=1e-12) + + if __name__ == "__main__": import os import sys diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index 891cf1c6cc..b7522133bc 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -3,7 +3,7 @@ from ufl import (Mesh, FunctionSpace, TestFunction, TrialFunction, TensorProductCell, dx, action, interval, triangle, - quadrilateral, curl, dot, div, grad) + tetrahedron, quadrilateral, curl, dot, div, grad, inner) from finat.ufl import (FiniteElement, VectorElement, EnrichedElement, TensorProductElement, HCurlElement, HDivElement) @@ -18,6 +18,18 @@ def helmholtz(cell, degree): return (u*v + dot(grad(u), grad(v)))*dx +def simplex_dg_mass(cell, degree): + # A simplicial DG element whose nodal basis coincides with the Dubiner + # expansion set (finat.spectral.Legendre), so that dx(scheme="collapsed") + # takes the sum-factorized (Duffy/lattice) tabulation path in tsfc/fem.py + # instead of dense tabulation. + m = Mesh(VectorElement('CG', cell, 1)) + V = FunctionSpace(m, FiniteElement('DG', cell, degree, variant='integral')) + u = TrialFunction(V) + v = TestFunction(V) + return inner(u, v) * dx(scheme='collapsed') + + def split_mixed_poisson(cell, degree): m = Mesh(VectorElement('CG', cell, 1)) if cell.cellname in ['interval * interval', 'quadrilateral']: @@ -100,6 +112,20 @@ def test_rhs(cell, order): assert (rates < order).all() +@pytest.mark.parametrize(('cell', 'order'), [(triangle, 4), (tetrahedron, 6)]) +def test_simplex_dg_mass_action(cell, order): + # Matrix-free DG mass-matrix action (milestone 2 of PLAN.md / DESIGN.md): + # the coefficient contraction in translate_coefficient is sum-factorized + # via the Duffy/lattice tabulation, targeting O(p^{d+1}) flops. This + # tests the *action* (right-hand side, like test_rhs above), not full + # bilinear matrix assembly, which is milestone 4 and not yet implemented. + degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) + flops = [count_flops(action(simplex_dg_mass(cell, degree))) + for degree in degrees] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) + assert (rates < order).all() + + @pytest.mark.parametrize(('cell', 'order'), [(quadrilateral, 5), (TensorProductCell(interval, interval), 5), diff --git a/tsfc/fem.py b/tsfc/fem.py index 943089052e..986412a615 100644 --- a/tsfc/fem.py +++ b/tsfc/fem.py @@ -8,16 +8,18 @@ import gem import numpy import ufl +from FIAT.expansions import morton_forward_table, morton_inverse_table from FIAT.orientation_utils import Orientation as FIATOrientation from FIAT.reference_element import UFCHexahedron, UFCQuadrilateral, UFCSimplex, make_affine_mapping from FIAT.reference_element import TensorProductCell from finat.physically_mapped import (NeedsCoordinateMappingElement, PhysicalGeometry) -from finat.point_set import PointSet, PointSingleton +from finat.point_set import CollapsedTensorProductPointSet, PointSet, PointSingleton from finat.quadrature import make_quadrature from finat.element_factory import as_fiat_cell, create_element -from gem.node import traversal -from gem.optimise import constant_fold_zero, ffc_rounding +from finat.spectral import Legendre +from gem.node import MemoizerArg, traversal +from gem.optimise import constant_fold_zero, contraction, ffc_rounding, filtered_replace_indices from gem.unconcatenate import unconcatenate from ufl.classes import (Argument, CellCoordinate, CellEdgeVectors, CellFacetJacobian, CellOrientation, CellOrigin, @@ -707,23 +709,151 @@ def fiat_to_ufl(fiat_dict, order): return gem.ComponentTensor(tensor, sigma + delta) +def _use_sum_factorisation(element, ctx): + """Whether the sum-factorized (Duffy/lattice) tabulation applies. + + This holds exactly when `element` is a simplicial DG element whose + nodal basis coincides with the Dubiner expansion set (currently + `finat.spectral.Legendre`), evaluation points come from a collapsed + tensor-product quadrature rule (requested via + ``dx(scheme="collapsed")``), and the integral is over the cell + interior. In that case `finat.spectral.Legendre.duffy_evaluation` + tabulates the element in O(p^d) space/time using a lattice + multi-index rather than the flat degree-of-freedom index, whereas + the standard `~.PointSetContext.basis_evaluation` tabulates all + O(p^d) basis functions densely at all O(p^d) points. + + Parameters + ---------- + element : finat.finiteelementbase.FiniteElementBase + The element being tabulated. + ctx : ContextBase + The translation context. + + Returns + ------- + bool + Whether to use `_duffy_evaluation` in place of + ``ctx.basis_evaluation``. + """ + return (isinstance(element, Legendre) + and isinstance(ctx, PointSetContext) + and isinstance(ctx.point_set, CollapsedTensorProductPointSet) + and ctx.integration_dim == ctx.fiat_cell.get_dimension() + and not ctx.unsummed_coefficient_indices) + + +def _duffy_evaluation(element, mt, ctx, entity_id): + """Sum-factorized tabulation of a simplicial Legendre DG element. + + Thin wrapper around `finat.spectral.Legendre.duffy_evaluation` that + filters out derivative orders other than ``mt.local_derivatives``, + mirroring the filtering `translate_argument` and + `translate_coefficient` apply to `~.PointSetContext.basis_evaluation` + output. + + Parameters + ---------- + element : finat.spectral.Legendre + The element being tabulated. + mt : ModifiedTerminal + The modified terminal being translated. + ctx : PointSetContext + The translation context; ``ctx.point_set`` must be a + `finat.point_set.CollapsedTensorProductPointSet`. + entity_id : int + The cell entity id, relative to ``ctx.integration_dim`` (the + cell interior only is supported). + + Returns + ------- + tuple + ``(multiindex, result)``: ``multiindex`` is the tuple of + `gem.JaggedIndex` enumerating the simplex lattice, and + ``result`` maps each derivative multi-index alpha with + ``sum(alpha) == mt.local_derivatives`` to a scalar GEM + expression free in ``multiindex`` and ``ctx.point_set.indices``. + """ + multiindex, result = element.duffy_evaluation(mt.local_derivatives, ctx.point_set, + (ctx.integration_dim, entity_id)) + result = {alpha: table for alpha, table in result.items() + if sum(alpha) == mt.local_derivatives} + return multiindex, result + + +def _scatter_to_dof_index(multiindex, result, element): + """Reshape a lattice-indexed tabulation into a flat-dof-indexed one. + + Builds, for each derivative multi-index alpha, a dense + `gem.ComponentTensor` of shape ``(element.space_dimension(),)`` + indexed by the flat degree-of-freedom index, matching the shape + convention of the standard (non-factorized) + `~.PointSetContext.basis_evaluation` output that `fiat_to_ufl` + expects. The flat index of a lattice point is its Morton index + (`FIAT.expansions.morton_index`), the same enumeration FIAT already + uses for the element's degrees of freedom, so no reordering of the + element's dof numbering is involved. + + Parameters + ---------- + multiindex : tuple of gem.JaggedIndex + The lattice multi-index free in each entry of ``result``, as + returned by `_duffy_evaluation`. + result : dict + Mapping alpha to a scalar GEM expression free in ``multiindex`` + (and point indices). + element : finat.spectral.Legendre + The element being tabulated. + + Returns + ------- + dict + Mapping alpha to a `gem.ComponentTensor` of shape + ``(element.space_dimension(),)``. + """ + sd = len(multiindex) + ndof = element.space_dimension() + r = gem.Index(extent=ndof) + inv_table = morton_inverse_table(sd, element.degree) + subst = tuple( + (axis, gem.VariableIndex(gem.Indexed( + gem.Literal(numpy.ascontiguousarray(inv_table[:, t]), dtype=gem.uint_type), (r,)))) + for t, axis in enumerate(multiindex) + ) + mapper = MemoizerArg(filtered_replace_indices) + return {alpha: gem.ComponentTensor(mapper(expr, subst), (r,)) + for alpha, expr in result.items()} + + @translate.register(Argument) def translate_argument(terminal, mt, ctx): element = ctx.create_element(terminal.ufl_element(), restriction=mt.restriction) - def callback(entity_id): - finat_dict = ctx.basis_evaluation(element, mt, entity_id) - # Filter out irrelevant derivatives - filtered_dict = {alpha: finat_dict[alpha] - for alpha in finat_dict - if sum(alpha) == mt.local_derivatives} - - # Change from FIAT to UFL arrangement - square = fiat_to_ufl(filtered_dict, mt.local_derivatives) - - # A numerical hack that FFC used to apply on FIAT tables still - # lives on after ditching FFC and switching to FInAT. - return ffc_rounding(square, ctx.epsilon) + if _use_sum_factorisation(element, ctx): + def callback(entity_id): + multiindex, duffy_dict = _duffy_evaluation(element, mt, ctx, entity_id) + filtered_dict = _scatter_to_dof_index(multiindex, duffy_dict, element) + + # Change from FIAT to UFL arrangement + square = fiat_to_ufl(filtered_dict, mt.local_derivatives) + + # A numerical hack that FFC used to apply on FIAT tables still + # lives on after ditching FFC and switching to FInAT. + return ffc_rounding(square, ctx.epsilon) + else: + def callback(entity_id): + finat_dict = ctx.basis_evaluation(element, mt, entity_id) + # Filter out irrelevant derivatives + filtered_dict = {alpha: finat_dict[alpha] + for alpha in finat_dict + if sum(alpha) == mt.local_derivatives} + + # Change from FIAT to UFL arrangement + square = fiat_to_ufl(filtered_dict, mt.local_derivatives) + + # A numerical hack that FFC used to apply on FIAT tables still + # lives on after ditching FFC and switching to FInAT. + return ffc_rounding(square, ctx.epsilon) table = ctx.entity_selector(callback, extract_unique_domain(terminal), mt.restriction) if ctx.use_canonical_quadrature_point_ordering: quad_multiindex = ctx.quadrature_rule.point_set.indices @@ -734,6 +864,51 @@ def callback(entity_id): return gem.partial_indexed(table, argument_multiindex) +def _contract_dof_index(multiindex, result, element, vec): + """Contract a lattice-indexed tabulation against a coefficient vector. + + The sum over the flat degree-of-freedom index is rewritten as a sum + over the lattice multi-index, gathering the coefficient vector + through the same Morton dof numbering FIAT already uses + (`FIAT.expansions.morton_index`). `gem.optimise.contraction` + sum-factorizes the resulting nested sum over the lattice + multi-index, exploiting the same axis-separable structure that + makes `finat.spectral.Legendre.duffy_evaluation` itself O(p^d). + + Parameters + ---------- + multiindex : tuple of gem.JaggedIndex + The lattice multi-index free in each entry of ``result``, as + returned by `_duffy_evaluation`. + result : dict + Mapping alpha to a scalar GEM expression free in ``multiindex`` + (and point indices). + element : finat.spectral.Legendre + The element being tabulated. + vec : gem.Node + The coefficient's local dof vector, of shape + ``(element.space_dimension(),)``. + + Returns + ------- + dict + Mapping alpha to a `gem.ComponentTensor` over + ``element.get_value_indices()`` (empty for the scalar `Legendre` + element), free in the point indices only. + """ + sd = len(multiindex) + fwd_table = morton_forward_table(sd, element.degree) + r_index = gem.VariableIndex(gem.Indexed( + gem.Literal(fwd_table, dtype=gem.uint_type), multiindex)) + vec_r, = gem.optimise.remove_componenttensors([gem.Indexed(vec, (r_index,))]) + zeta = element.get_value_indices() + value_dict = {} + for alpha, expr in result.items(): + value = gem.IndexSum(gem.Product(expr, vec_r), multiindex) + value_dict[alpha] = gem.ComponentTensor(contraction(value), zeta) + return value_dict + + @translate.register(TSFCConstantMixin) def translate_constant_value(terminal, mt, ctx): return ctx.constant(terminal) @@ -745,45 +920,52 @@ def translate_coefficient(terminal, mt, ctx): vec = ctx.coefficient(terminal, mt.restriction) element = ctx.create_element(terminal.ufl_element(), restriction=mt.restriction) - # Collect FInAT tabulation for all entities - per_derivative = collections.defaultdict(list) - for entity_id in ctx.entity_ids(domain): - finat_dict = ctx.basis_evaluation(element, mt, entity_id) - for alpha, table in finat_dict.items(): - # Filter out irrelevant derivatives - if sum(alpha) == mt.local_derivatives: - # A numerical hack that FFC used to apply on FIAT - # tables still lives on after ditching FFC and - # switching to FInAT. - table = ffc_rounding(table, ctx.epsilon) - per_derivative[alpha].append(table) - - # Merge entity tabulations for each derivative - if len(ctx.entity_ids(domain)) == 1: - def take_singleton(xs): - x, = xs # asserts singleton - return x - per_derivative = {alpha: take_singleton(tables) - for alpha, tables in per_derivative.items()} + if _use_sum_factorisation(element, ctx): + entity_id, = ctx.entity_ids(domain) + multiindex, duffy_dict = _duffy_evaluation(element, mt, ctx, entity_id) + duffy_dict = {alpha: ffc_rounding(table, ctx.epsilon) + for alpha, table in duffy_dict.items()} + value_dict = _contract_dof_index(multiindex, duffy_dict, element, vec) else: - f = ctx.entity_number(domain, mt.restriction) - per_derivative = {alpha: gem.select_expression(tables, f) - for alpha, tables in per_derivative.items()} - - # Coefficient evaluation - beta = ctx.index_cache.setdefault(terminal.ufl_element(), element.get_indices()) - zeta = element.get_value_indices() - vec_beta, = gem.optimise.remove_componenttensors([gem.Indexed(vec, beta)]) - value_dict = {} - for alpha, table in per_derivative.items(): - table_qi = gem.Indexed(table, beta + zeta) - summands = [] - for var, expr in unconcatenate([(vec_beta, table_qi)], ctx.index_cache): - indices = tuple(i for i in var.index_ordering() if i not in ctx.unsummed_coefficient_indices) - value = gem.IndexSum(gem.Product(expr, var), indices) - summands.append(gem.optimise.contraction(value)) - optimised_value = gem.optimise.make_sum(summands) - value_dict[alpha] = gem.ComponentTensor(optimised_value, zeta) + # Collect FInAT tabulation for all entities + per_derivative = collections.defaultdict(list) + for entity_id in ctx.entity_ids(domain): + finat_dict = ctx.basis_evaluation(element, mt, entity_id) + for alpha, table in finat_dict.items(): + # Filter out irrelevant derivatives + if sum(alpha) == mt.local_derivatives: + # A numerical hack that FFC used to apply on FIAT + # tables still lives on after ditching FFC and + # switching to FInAT. + table = ffc_rounding(table, ctx.epsilon) + per_derivative[alpha].append(table) + + # Merge entity tabulations for each derivative + if len(ctx.entity_ids(domain)) == 1: + def take_singleton(xs): + x, = xs # asserts singleton + return x + per_derivative = {alpha: take_singleton(tables) + for alpha, tables in per_derivative.items()} + else: + f = ctx.entity_number(domain, mt.restriction) + per_derivative = {alpha: gem.select_expression(tables, f) + for alpha, tables in per_derivative.items()} + + # Coefficient evaluation + beta = ctx.index_cache.setdefault(terminal.ufl_element(), element.get_indices()) + zeta = element.get_value_indices() + vec_beta, = gem.optimise.remove_componenttensors([gem.Indexed(vec, beta)]) + value_dict = {} + for alpha, table in per_derivative.items(): + table_qi = gem.Indexed(table, beta + zeta) + summands = [] + for var, expr in unconcatenate([(vec_beta, table_qi)], ctx.index_cache): + indices = tuple(i for i in var.index_ordering() if i not in ctx.unsummed_coefficient_indices) + value = gem.IndexSum(gem.Product(expr, var), indices) + summands.append(gem.optimise.contraction(value)) + optimised_value = gem.optimise.make_sum(summands) + value_dict[alpha] = gem.ComponentTensor(optimised_value, zeta) # Change from FIAT to UFL arrangement result = fiat_to_ufl(value_dict, mt.local_derivatives) From 1292af6bc03b016243b89b1f17f135b9d997cc61 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 20 Jul 2026 16:15:55 +0100 Subject: [PATCH 03/23] cleanup --- DESIGN.md | 127 ++++++++++++---------- tests/tsfc/test_codegen.py | 55 +++++----- tsfc/fem.py | 210 +++++++------------------------------ 3 files changed, 135 insertions(+), 257 deletions(-) diff --git a/DESIGN.md b/DESIGN.md index 991e5a088e..77321ce278 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -118,58 +118,61 @@ so the original plan of making `argument_multiindices` itself a lattice multiindex (see the old Route-B write-up below) would have changed the local tensor's shape and broken that contract. The implementation instead keeps `element.index_shape` and `argument_multiindices` exactly as they are today, -and confines the lattice multiindex to the tabulation step alone: - -* `_use_sum_factorisation(element, ctx)` (`tsfc/fem.py`) gates the whole path: - `element` must be `finat.spectral.Legendre`, `ctx.point_set` a - `CollapsedTensorProductPointSet` (i.e. the measure requested - `dx(scheme="collapsed")`), the integral must be over the cell interior, and - `ctx.unsummed_coefficient_indices` must be empty (macrocells, which - `duffy_evaluation` already rejects, are the only case that sets it). -* `_duffy_evaluation(element, mt, ctx, entity_id)` calls - `element.duffy_evaluation(mt.local_derivatives, ctx.point_set, - (ctx.integration_dim, entity_id))` and filters to `sum(alpha) == - mt.local_derivatives`, exactly mirroring the filtering the standard path - applies to `basis_evaluation`'s output. -* **Forward transform (`translate_coefficient`).** `_contract_dof_index` - builds a forward Morton lookup table (`FIAT.expansions.morton_forward_table`, - shape `(degree+1,)^d`, clamped to a valid dof so out-of-lattice reads are - merely wasted, never out of bounds — they always multiply a zero - tabulation), gathers `vec[VariableIndex(table[multiindex])]`, and hands - `IndexSum(Product(duffy[alpha], vec_r), multiindex)` to - `gem.optimise.contraction`, exactly as originally planned: the `m_t` - `VariableIndex` couplings inside `duffy_evaluation`'s own expression make - the per-axis free-index sets nested (`{i_1} ⊂ {i_1, i_2} ⊂ ...`), so - `contraction` finds the innermost-axis-first Karniadakis–Sherwin sweep by - itself — no bespoke sum-factorization code needed. The result is wrapped - back into a `gem.ComponentTensor` over `element.get_value_indices()` (empty - for the scalar `Legendre` element), so it slots into `fiat_to_ufl` exactly - like a standard dense tabulation would. -* **Backward transform (`translate_argument`).** `_scatter_to_dof_index` goes - the other way: it introduces one *fresh* flat dof index `r` (the same free - index `argument_multiindex` will later pick a single value of — nothing - about `argument_multiindices` construction changes), builds the *inverse* - Morton table (`FIAT.expansions.morton_inverse_table`, shape `(ndof, d)`) to - get per-axis lookups `i_t(r)`, and substitutes - `multiindex[t] -> VariableIndex(inverse_table[:, t][r])` throughout - `duffy_evaluation`'s expression tree via - `gem.node.MemoizerArg(gem.optimise.filtered_replace_indices)` — the same - substitution mechanism `translate_argument`/`translate_coefficient` already - use for canonical quadrature-point reordering. `filtered_replace_indices` - recurses into `VariableIndex.expression` (`gem/optimise.py`'s - `_replace_indices_atomic`), so this also correctly rewrites the nested `m_t` - lookups (which are themselves `VariableIndex` expressions built from - `multiindex[:t]`) into r-indexed double lookups, with no duplicated - tabulation logic. The result, wrapped in `gem.ComponentTensor(..., (r,))`, - is a dense `(ndof,)`-shaped table — indistinguishable, from - `fiat_to_ufl`/`prepare_arguments`'s point of view, from the standard dense - tabulation. `gem.optimise.contraction` never runs on this side (there is no - sum to hoist yet at this stage); the sum-factorized quadrature contraction - happens later, per dof, when `vanilla.py`/`spectral.py` process the - quadrature `IndexSum` — the collapsed quadrature's own per-axis structure is - what still delivers the O(p) win per axis there. - -Both helpers were validated to ~1e-13/1e-14 against FIAT's dense +and confines the lattice multiindex — and all the Duffy/Morton machinery — to +a new `finat/duffy.py` module, reached from `tsfc/fem.py` through ordinary +FInAT element methods rather than through tsfc-side branching: + +* `finat.duffy.DuffyElement` is a mixin (`finat.spectral.Legendre` is + currently its only user) providing `duffy_evaluation` (the lattice-indexed, + sum-factorized tabulation) plus two dispatch points: + * **Backward transform / `basis_evaluation` override.** When `ps` is a + `CollapsedTensorProductPointSet` and the entity is the cell interior (and + not a macrocell — the only other case `duffy_evaluation` rejects), + `DuffyElement.basis_evaluation` calls `duffy_evaluation` and scatters the + lattice tabulation to the standard flat-dof-indexed convention: it + introduces one *fresh* flat dof index `r`, builds the *inverse* Morton + table (`FIAT.expansions.morton_inverse_table`, shape `(ndof, d)`) to get + per-axis lookups `i_t(r)`, and substitutes `multiindex[t] -> + VariableIndex(inverse_table[:, t][r])` throughout `duffy_evaluation`'s + expression tree via `gem.node.MemoizerArg(gem.optimise.filtered_replace_indices)` + (the same substitution mechanism `translate_argument`/`translate_coefficient` + use for canonical quadrature-point reordering; `filtered_replace_indices` + recurses into `VariableIndex.expression`, so the nested `m_t` lookups are + rewritten too). The result, wrapped in `gem.ComponentTensor(..., (r,))`, + is indistinguishable, from `fiat_to_ufl`/`prepare_arguments`'s point of + view, from a standard dense `basis_evaluation` tabulation — so + `translate_argument` in `tsfc/fem.py` needs **no special case at all**: it + always calls `ctx.basis_evaluation(element, mt, entity_id)`, exactly as + for any other element. `gem.optimise.contraction` never runs on this + side (there is no sum to hoist yet at this stage); the sum-factorized + quadrature contraction happens later, per dof, when + `vanilla.py`/`spectral.py` process the quadrature `IndexSum` — the + collapsed quadrature's own per-axis structure is what still delivers the + O(p) win per axis there. + * **Forward transform / `duffy_contraction`.** Unlike `basis_evaluation`, + coefficient contraction has no generic per-element hook to dispatch + through (it additionally needs the coefficient's dof vector `vec`), so + `translate_coefficient` in `tsfc/fem.py` keeps a small + `_use_duffy_contraction(element, ctx)` guard (`isinstance(element, + DuffyElement)`, `ctx.point_set` a `CollapsedTensorProductPointSet`, cell + interior, `ctx.unsummed_coefficient_indices` empty) before calling + `element.duffy_contraction(mt.local_derivatives, ctx.point_set, entity, + vec, ctx.epsilon)`. `duffy_contraction` builds a forward Morton lookup + table (`FIAT.expansions.morton_forward_table`, shape `(degree+1,)^d`, + clamped to a valid dof so out-of-lattice reads are merely wasted, never + out of bounds — they always multiply a zero tabulation), gathers + `vec[VariableIndex(table[multiindex])]`, and hands + `IndexSum(Product(duffy[alpha], vec_r), multiindex)` to + `gem.optimise.contraction`, exactly as originally planned: the `m_t` + `VariableIndex` couplings inside `duffy_evaluation`'s own expression make + the per-axis free-index sets nested (`{i_1} ⊂ {i_1, i_2} ⊂ ...`), so + `contraction` finds the innermost-axis-first Karniadakis–Sherwin sweep by + itself — no bespoke sum-factorization code needed. The result is wrapped + back into a `gem.ComponentTensor` over `element.get_value_indices()` + (empty for the scalar `Legendre` element), so it slots into `fiat_to_ufl` + exactly like a standard dense tabulation would. + +Both dispatch points were validated to ~1e-13/1e-14 against FIAT's dense `tabulate()`, via compiled-and-executed loopy kernels, for values and first derivatives on triangles and tetrahedra (`tests/tsfc/test_codegen.py::test_duffy_scatter_and_contract`), and end to @@ -194,11 +197,14 @@ considered: * **Route B — Morton gather/scatter via `VariableIndex` (chosen; see (b)).** Keeps FIAT's dof ordering, `element.index_shape`, and `argument_multiindices` completely untouched; the Morton lookup lives - entirely inside `_contract_dof_index`/`_scatter_to_dof_index` in `tsfc/fem.py`. - No `driver.py` or `kernel_interface/*.py` changes were needed at all — the - lattice multiindex never escapes `fem.py`. The indirection costs one uint - load per accumulation (forward) or one uint load per dof (backward), - negligible against the O(p) inner contraction. + entirely inside `finat/duffy.py` (`DuffyElement.basis_evaluation` and + `DuffyElement.duffy_contraction`). No `driver.py` or + `kernel_interface/*.py` changes were needed at all, and — after the + fem.py-integration refactor above — no bespoke branching in `tsfc/fem.py` + either for the argument side; the lattice multiindex never escapes + `finat/duffy.py`. The indirection costs one uint load per accumulation + (forward) or one uint load per dof (backward), negligible against the O(p) + inner contraction. * **Route C — reorder FIAT dofs p-major.** Change `Legendre` (variant="integral") to lattice-lexicographic dof order so the flat index @@ -210,6 +216,13 @@ considered: ## Deferred +* **Route C — reorder FIAT dofs to eliminate the Morton gather/scatter + entirely.** Raised in self-review of PR #5263: renumbering the expansion + set/finite element dofs so the flat dof index *is* the (bounding-box) + lattice multiindex would let the generic `basis_evaluation`/contraction + machinery pick up the separable structure without any table lookup at all + (see Route C above). Left as a follow-up after the fem.py-integration + refactor landed in `finat/duffy.py`, rather than attempted alongside it. * **CG / C0 basis (milestones 3–4):** the C0 recombination makes each basis function a sum of <= 3 separable members (Sherwin–Karniadakis vertex/edge/face recombination); `tabulate_duffy` currently raises `NotImplementedError` for diff --git a/tests/tsfc/test_codegen.py b/tests/tsfc/test_codegen.py index e78c238309..856f6c7614 100644 --- a/tests/tsfc/test_codegen.py +++ b/tests/tsfc/test_codegen.py @@ -78,12 +78,11 @@ def test_jagged_index_codegen(monkeypatch): @pytest.mark.parametrize("cellname,degree", [("triangle", 3), ("tetrahedron", 2)]) def test_duffy_scatter_and_contract(monkeypatch, cellname, degree): """Route B of the simplex sum-factorization milestone 2 design: - ``tsfc.fem._scatter_to_dof_index`` (the `translate_argument` path) and - ``tsfc.fem._contract_dof_index`` (the `translate_coefficient` path) - must reproduce the standard dense FIAT tabulation, via the Morton - dof numbering FIAT already uses, from - `finat.spectral.Legendre.duffy_evaluation`'s lattice-indexed, - sum-factorized tabulation. + `finat.duffy.DuffyElement.basis_evaluation` (the `translate_argument` + path) and `finat.duffy.DuffyElement.duffy_contraction` (the + `translate_coefficient` path) must reproduce the standard dense FIAT + tabulation, via the Morton dof numbering FIAT already uses, from + `duffy_evaluation`'s lattice-indexed, sum-factorized tabulation. """ import loopy as lp import tsfc.loopy @@ -93,7 +92,6 @@ def test_duffy_scatter_and_contract(monkeypatch, cellname, degree): from gem import impero_utils from gem.gem import Index, Indexed, Variable from gem.optimise import remove_componenttensors - from tsfc.fem import _contract_dof_index, _scatter_to_dof_index # Execute the generated code so we check the numbers, not just the loop bounds monkeypatch.setattr(tsfc.loopy, "target", lp.ExecutableCTarget()) @@ -108,19 +106,18 @@ def test_duffy_scatter_and_contract(monkeypatch, cellname, degree): point_shape = tuple(index.extent for index in point_indices) entity = (cell.get_dimension(), 0) - multiindex, duffy_dict = element.duffy_evaluation(1, point_set, entity) dense_dict = element._element.tabulate(1, point_set.points) rng = numpy.random.default_rng(1) coefficients = rng.random(ndof) - for alpha, table_expr in duffy_dict.items(): - dense = dense_dict[alpha].reshape((ndof,) + point_shape) + # translate_argument path: dispatched transparently through basis_evaluation + scattered_dict = element.basis_evaluation(1, point_set, entity) + for alpha, dense in dense_dict.items(): + dense = dense.reshape((ndof,) + point_shape) - # translate_argument path: flat-dof-indexed dense table - scattered = _scatter_to_dof_index(multiindex, {alpha: table_expr}, element)[alpha] r = Index(extent=ndof) - table, = remove_componenttensors([Indexed(scattered, (r,))]) + table, = remove_componenttensors([Indexed(scattered_dict[alpha], (r,))]) u = Variable("u", (ndof,) + point_shape) impero_c = impero_utils.compile_gem( [(Indexed(u, (r,) + point_indices), table)], (r,) + point_indices) @@ -130,20 +127,24 @@ def test_duffy_scatter_and_contract(monkeypatch, cellname, degree): knl(u=u_out) assert numpy.allclose(u_out, dense, rtol=1e-12, atol=1e-12) - # translate_coefficient path: contraction against a coefficient vector - c = Variable("c", (ndof,)) - contracted = _contract_dof_index(multiindex, {alpha: table_expr}, element, c)[alpha] - value, = remove_componenttensors([Indexed(contracted, ())]) - v = Variable("v", point_shape) - impero_c = impero_utils.compile_gem( - [(Indexed(v, point_indices), value)], point_indices) - args = [lp.GlobalArg("v", dtype=numpy.float64, shape=point_shape), - lp.GlobalArg("c", dtype=numpy.float64, shape=(ndof,))] - knl, _ = tsfc.loopy.generate(impero_c, args, numpy.float64) - v_out = numpy.zeros(point_shape) - knl(v=v_out, c=coefficients) - v_ref = numpy.tensordot(coefficients, dense, axes=(0, 0)) - assert numpy.allclose(v_out, v_ref, rtol=1e-12, atol=1e-12) + # translate_coefficient path: contraction against a coefficient vector, + # dispatched through duffy_contraction + c = Variable("c", (ndof,)) + for order in (0, 1): + contracted_dict = element.duffy_contraction(order, point_set, entity, c, epsilon=0.0) + for alpha, contracted in contracted_dict.items(): + dense = dense_dict[alpha].reshape((ndof,) + point_shape) + value, = remove_componenttensors([Indexed(contracted, ())]) + v = Variable("v", point_shape) + impero_c = impero_utils.compile_gem( + [(Indexed(v, point_indices), value)], point_indices) + args = [lp.GlobalArg("v", dtype=numpy.float64, shape=point_shape), + lp.GlobalArg("c", dtype=numpy.float64, shape=(ndof,))] + knl, _ = tsfc.loopy.generate(impero_c, args, numpy.float64) + v_out = numpy.zeros(point_shape) + knl(v=v_out, c=coefficients) + v_ref = numpy.tensordot(coefficients, dense, axes=(0, 0)) + assert numpy.allclose(v_out, v_ref, rtol=1e-12, atol=1e-12) if __name__ == "__main__": diff --git a/tsfc/fem.py b/tsfc/fem.py index 986412a615..8a6bad3e1a 100644 --- a/tsfc/fem.py +++ b/tsfc/fem.py @@ -8,18 +8,17 @@ import gem import numpy import ufl -from FIAT.expansions import morton_forward_table, morton_inverse_table from FIAT.orientation_utils import Orientation as FIATOrientation from FIAT.reference_element import UFCHexahedron, UFCQuadrilateral, UFCSimplex, make_affine_mapping from FIAT.reference_element import TensorProductCell +from finat.duffy import DuffyElement from finat.physically_mapped import (NeedsCoordinateMappingElement, PhysicalGeometry) from finat.point_set import CollapsedTensorProductPointSet, PointSet, PointSingleton from finat.quadrature import make_quadrature from finat.element_factory import as_fiat_cell, create_element -from finat.spectral import Legendre -from gem.node import MemoizerArg, traversal -from gem.optimise import constant_fold_zero, contraction, ffc_rounding, filtered_replace_indices +from gem.node import traversal +from gem.optimise import constant_fold_zero, ffc_rounding from gem.unconcatenate import unconcatenate from ufl.classes import (Argument, CellCoordinate, CellEdgeVectors, CellFacetJacobian, CellOrientation, CellOrigin, @@ -709,19 +708,24 @@ def fiat_to_ufl(fiat_dict, order): return gem.ComponentTensor(tensor, sigma + delta) -def _use_sum_factorisation(element, ctx): - """Whether the sum-factorized (Duffy/lattice) tabulation applies. +def _use_duffy_contraction(element, ctx): + """Whether the sum-factorized (Duffy/lattice) coefficient contraction + applies. This holds exactly when `element` is a simplicial DG element whose - nodal basis coincides with the Dubiner expansion set (currently - `finat.spectral.Legendre`), evaluation points come from a collapsed + nodal basis coincides with the Dubiner expansion set (any + `finat.duffy.DuffyElement`), evaluation points come from a collapsed tensor-product quadrature rule (requested via ``dx(scheme="collapsed")``), and the integral is over the cell - interior. In that case `finat.spectral.Legendre.duffy_evaluation` - tabulates the element in O(p^d) space/time using a lattice - multi-index rather than the flat degree-of-freedom index, whereas - the standard `~.PointSetContext.basis_evaluation` tabulates all - O(p^d) basis functions densely at all O(p^d) points. + interior. In that case `DuffyElement.duffy_contraction` contracts a + `Coefficient` against the element in O(p^d) space/time using the + lattice multi-index, whereas the standard dense contraction + materializes all O(p^d) basis functions at all O(p^d) points before + contracting. + + The argument (basis evaluation) side needs no such dispatch: it is + handled transparently by `DuffyElement.basis_evaluation`, reached + through the usual `~.PointSetContext.basis_evaluation` call. Parameters ---------- @@ -733,127 +737,33 @@ def _use_sum_factorisation(element, ctx): Returns ------- bool - Whether to use `_duffy_evaluation` in place of - ``ctx.basis_evaluation``. + Whether to use `DuffyElement.duffy_contraction` in place of the + standard dense contraction. """ - return (isinstance(element, Legendre) + return (isinstance(element, DuffyElement) and isinstance(ctx, PointSetContext) and isinstance(ctx.point_set, CollapsedTensorProductPointSet) and ctx.integration_dim == ctx.fiat_cell.get_dimension() and not ctx.unsummed_coefficient_indices) -def _duffy_evaluation(element, mt, ctx, entity_id): - """Sum-factorized tabulation of a simplicial Legendre DG element. - - Thin wrapper around `finat.spectral.Legendre.duffy_evaluation` that - filters out derivative orders other than ``mt.local_derivatives``, - mirroring the filtering `translate_argument` and - `translate_coefficient` apply to `~.PointSetContext.basis_evaluation` - output. - - Parameters - ---------- - element : finat.spectral.Legendre - The element being tabulated. - mt : ModifiedTerminal - The modified terminal being translated. - ctx : PointSetContext - The translation context; ``ctx.point_set`` must be a - `finat.point_set.CollapsedTensorProductPointSet`. - entity_id : int - The cell entity id, relative to ``ctx.integration_dim`` (the - cell interior only is supported). - - Returns - ------- - tuple - ``(multiindex, result)``: ``multiindex`` is the tuple of - `gem.JaggedIndex` enumerating the simplex lattice, and - ``result`` maps each derivative multi-index alpha with - ``sum(alpha) == mt.local_derivatives`` to a scalar GEM - expression free in ``multiindex`` and ``ctx.point_set.indices``. - """ - multiindex, result = element.duffy_evaluation(mt.local_derivatives, ctx.point_set, - (ctx.integration_dim, entity_id)) - result = {alpha: table for alpha, table in result.items() - if sum(alpha) == mt.local_derivatives} - return multiindex, result - - -def _scatter_to_dof_index(multiindex, result, element): - """Reshape a lattice-indexed tabulation into a flat-dof-indexed one. - - Builds, for each derivative multi-index alpha, a dense - `gem.ComponentTensor` of shape ``(element.space_dimension(),)`` - indexed by the flat degree-of-freedom index, matching the shape - convention of the standard (non-factorized) - `~.PointSetContext.basis_evaluation` output that `fiat_to_ufl` - expects. The flat index of a lattice point is its Morton index - (`FIAT.expansions.morton_index`), the same enumeration FIAT already - uses for the element's degrees of freedom, so no reordering of the - element's dof numbering is involved. - - Parameters - ---------- - multiindex : tuple of gem.JaggedIndex - The lattice multi-index free in each entry of ``result``, as - returned by `_duffy_evaluation`. - result : dict - Mapping alpha to a scalar GEM expression free in ``multiindex`` - (and point indices). - element : finat.spectral.Legendre - The element being tabulated. - - Returns - ------- - dict - Mapping alpha to a `gem.ComponentTensor` of shape - ``(element.space_dimension(),)``. - """ - sd = len(multiindex) - ndof = element.space_dimension() - r = gem.Index(extent=ndof) - inv_table = morton_inverse_table(sd, element.degree) - subst = tuple( - (axis, gem.VariableIndex(gem.Indexed( - gem.Literal(numpy.ascontiguousarray(inv_table[:, t]), dtype=gem.uint_type), (r,)))) - for t, axis in enumerate(multiindex) - ) - mapper = MemoizerArg(filtered_replace_indices) - return {alpha: gem.ComponentTensor(mapper(expr, subst), (r,)) - for alpha, expr in result.items()} - - @translate.register(Argument) def translate_argument(terminal, mt, ctx): element = ctx.create_element(terminal.ufl_element(), restriction=mt.restriction) - if _use_sum_factorisation(element, ctx): - def callback(entity_id): - multiindex, duffy_dict = _duffy_evaluation(element, mt, ctx, entity_id) - filtered_dict = _scatter_to_dof_index(multiindex, duffy_dict, element) - - # Change from FIAT to UFL arrangement - square = fiat_to_ufl(filtered_dict, mt.local_derivatives) - - # A numerical hack that FFC used to apply on FIAT tables still - # lives on after ditching FFC and switching to FInAT. - return ffc_rounding(square, ctx.epsilon) - else: - def callback(entity_id): - finat_dict = ctx.basis_evaluation(element, mt, entity_id) - # Filter out irrelevant derivatives - filtered_dict = {alpha: finat_dict[alpha] - for alpha in finat_dict - if sum(alpha) == mt.local_derivatives} - - # Change from FIAT to UFL arrangement - square = fiat_to_ufl(filtered_dict, mt.local_derivatives) - - # A numerical hack that FFC used to apply on FIAT tables still - # lives on after ditching FFC and switching to FInAT. - return ffc_rounding(square, ctx.epsilon) + def callback(entity_id): + finat_dict = ctx.basis_evaluation(element, mt, entity_id) + # Filter out irrelevant derivatives + filtered_dict = {alpha: finat_dict[alpha] + for alpha in finat_dict + if sum(alpha) == mt.local_derivatives} + + # Change from FIAT to UFL arrangement + square = fiat_to_ufl(filtered_dict, mt.local_derivatives) + + # A numerical hack that FFC used to apply on FIAT tables still + # lives on after ditching FFC and switching to FInAT. + return ffc_rounding(square, ctx.epsilon) table = ctx.entity_selector(callback, extract_unique_domain(terminal), mt.restriction) if ctx.use_canonical_quadrature_point_ordering: quad_multiindex = ctx.quadrature_rule.point_set.indices @@ -864,51 +774,6 @@ def callback(entity_id): return gem.partial_indexed(table, argument_multiindex) -def _contract_dof_index(multiindex, result, element, vec): - """Contract a lattice-indexed tabulation against a coefficient vector. - - The sum over the flat degree-of-freedom index is rewritten as a sum - over the lattice multi-index, gathering the coefficient vector - through the same Morton dof numbering FIAT already uses - (`FIAT.expansions.morton_index`). `gem.optimise.contraction` - sum-factorizes the resulting nested sum over the lattice - multi-index, exploiting the same axis-separable structure that - makes `finat.spectral.Legendre.duffy_evaluation` itself O(p^d). - - Parameters - ---------- - multiindex : tuple of gem.JaggedIndex - The lattice multi-index free in each entry of ``result``, as - returned by `_duffy_evaluation`. - result : dict - Mapping alpha to a scalar GEM expression free in ``multiindex`` - (and point indices). - element : finat.spectral.Legendre - The element being tabulated. - vec : gem.Node - The coefficient's local dof vector, of shape - ``(element.space_dimension(),)``. - - Returns - ------- - dict - Mapping alpha to a `gem.ComponentTensor` over - ``element.get_value_indices()`` (empty for the scalar `Legendre` - element), free in the point indices only. - """ - sd = len(multiindex) - fwd_table = morton_forward_table(sd, element.degree) - r_index = gem.VariableIndex(gem.Indexed( - gem.Literal(fwd_table, dtype=gem.uint_type), multiindex)) - vec_r, = gem.optimise.remove_componenttensors([gem.Indexed(vec, (r_index,))]) - zeta = element.get_value_indices() - value_dict = {} - for alpha, expr in result.items(): - value = gem.IndexSum(gem.Product(expr, vec_r), multiindex) - value_dict[alpha] = gem.ComponentTensor(contraction(value), zeta) - return value_dict - - @translate.register(TSFCConstantMixin) def translate_constant_value(terminal, mt, ctx): return ctx.constant(terminal) @@ -920,12 +785,11 @@ def translate_coefficient(terminal, mt, ctx): vec = ctx.coefficient(terminal, mt.restriction) element = ctx.create_element(terminal.ufl_element(), restriction=mt.restriction) - if _use_sum_factorisation(element, ctx): + if _use_duffy_contraction(element, ctx): entity_id, = ctx.entity_ids(domain) - multiindex, duffy_dict = _duffy_evaluation(element, mt, ctx, entity_id) - duffy_dict = {alpha: ffc_rounding(table, ctx.epsilon) - for alpha, table in duffy_dict.items()} - value_dict = _contract_dof_index(multiindex, duffy_dict, element, vec) + value_dict = element.duffy_contraction(mt.local_derivatives, ctx.point_set, + (ctx.integration_dim, entity_id), + vec, ctx.epsilon) else: # Collect FInAT tabulation for all entities per_derivative = collections.defaultdict(list) From fa5120e3f8313b1e4e5d91fdf116946793b8a860 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 20 Jul 2026 20:48:25 +0100 Subject: [PATCH 04/23] test Duffy scatter and contract --- tests/tsfc/test_codegen.py | 42 ++++++++++++++++---------------------- 1 file changed, 18 insertions(+), 24 deletions(-) diff --git a/tests/tsfc/test_codegen.py b/tests/tsfc/test_codegen.py index 856f6c7614..efc8ba6453 100644 --- a/tests/tsfc/test_codegen.py +++ b/tests/tsfc/test_codegen.py @@ -76,26 +76,31 @@ def test_jagged_index_codegen(monkeypatch): @pytest.mark.parametrize("cellname,degree", [("triangle", 3), ("tetrahedron", 2)]) -def test_duffy_scatter_and_contract(monkeypatch, cellname, degree): +def test_duffy_scatter_and_contract(cellname, degree): """Route B of the simplex sum-factorization milestone 2 design: `finat.duffy.DuffyElement.basis_evaluation` (the `translate_argument` path) and `finat.duffy.DuffyElement.duffy_contraction` (the `translate_coefficient` path) must reproduce the standard dense FIAT tabulation, via the Morton dof numbering FIAT already uses, from `duffy_evaluation`'s lattice-indexed, sum-factorized tabulation. + + Verified via `gem.interpreter.evaluate` rather than a compiled loopy + kernel: the GEM expressions built here (in particular + `duffy_contraction`'s `gem.VariableIndex`-based Morton index arithmetic) + schedule fine once merged into a real PyOP2 wrapper kernel (as confirmed + via `firedrake.assemble` on real forms), but are not guaranteed + schedulable by loopy in isolation -- scheduling an isolated, + unwrapped kernel is not a configuration real Firedrake usage ever + exercises. The GEM interpreter checks the same numerical correctness + without depending on loopy scheduling at all. """ - import loopy as lp - import tsfc.loopy from FIAT.reference_element import UFCTetrahedron, UFCTriangle from finat.quadrature import make_quadrature from finat.spectral import Legendre - from gem import impero_utils from gem.gem import Index, Indexed, Variable + from gem.interpreter import evaluate from gem.optimise import remove_componenttensors - # Execute the generated code so we check the numbers, not just the loop bounds - monkeypatch.setattr(tsfc.loopy, "target", lp.ExecutableCTarget()) - cell = {"triangle": UFCTriangle, "tetrahedron": UFCTetrahedron}[cellname]() element = Legendre(cell, degree) ndof = element.space_dimension() @@ -118,14 +123,9 @@ def test_duffy_scatter_and_contract(monkeypatch, cellname, degree): r = Index(extent=ndof) table, = remove_componenttensors([Indexed(scattered_dict[alpha], (r,))]) - u = Variable("u", (ndof,) + point_shape) - impero_c = impero_utils.compile_gem( - [(Indexed(u, (r,) + point_indices), table)], (r,) + point_indices) - args = [lp.GlobalArg("u", dtype=numpy.float64, shape=(ndof,) + point_shape)] - knl, _ = tsfc.loopy.generate(impero_c, args, numpy.float64) - u_out = numpy.zeros((ndof,) + point_shape) - knl(u=u_out) - assert numpy.allclose(u_out, dense, rtol=1e-12, atol=1e-12) + u_out, = evaluate([table]) + assert u_out.fids == (r,) + point_indices + assert numpy.allclose(u_out.arr, dense, rtol=1e-12, atol=1e-12) # translate_coefficient path: contraction against a coefficient vector, # dispatched through duffy_contraction @@ -135,16 +135,10 @@ def test_duffy_scatter_and_contract(monkeypatch, cellname, degree): for alpha, contracted in contracted_dict.items(): dense = dense_dict[alpha].reshape((ndof,) + point_shape) value, = remove_componenttensors([Indexed(contracted, ())]) - v = Variable("v", point_shape) - impero_c = impero_utils.compile_gem( - [(Indexed(v, point_indices), value)], point_indices) - args = [lp.GlobalArg("v", dtype=numpy.float64, shape=point_shape), - lp.GlobalArg("c", dtype=numpy.float64, shape=(ndof,))] - knl, _ = tsfc.loopy.generate(impero_c, args, numpy.float64) - v_out = numpy.zeros(point_shape) - knl(v=v_out, c=coefficients) + v_out, = evaluate([value], bindings={c: coefficients}) + assert v_out.fids == point_indices v_ref = numpy.tensordot(coefficients, dense, axes=(0, 0)) - assert numpy.allclose(v_out, v_ref, rtol=1e-12, atol=1e-12) + assert numpy.allclose(v_out.arr, v_ref, rtol=1e-12, atol=1e-12) if __name__ == "__main__": From 571e51143402993f9c78fedf08a0bf988920c85a Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 21 Jul 2026 00:34:46 +0100 Subject: [PATCH 05/23] Document Route C dof-reorder outcome and why full elimination is deferred Legendre's dof order is now lattice-lexicographic (permuted in FIAT), which simplified finat/duffy.py's index arithmetic but did not eliminate the VariableIndex gather/scatter: get_indices() must stay a flat index because it can't distinguish a cell-interior kernel from a facet-integral kernel, and facet tabulation always uses the dense, flat-(ndof,) FIAT path. Fully eliminating the gather/scatter would need a bespoke jagged gem.FlexiblyIndexed view in kernel_interface/ common.py's prepare_arguments/prepare_coefficient, shared code every Firedrake kernel depends on -- deliberately deferred as a separate, higher-risk follow-up. Co-Authored-By: Claude Sonnet 5 --- DESIGN.md | 67 +++++++++++++++++++++++++++++++++++++------------------ 1 file changed, 45 insertions(+), 22 deletions(-) diff --git a/DESIGN.md b/DESIGN.md index 77321ce278..5913158386 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -180,12 +180,13 @@ end through `firedrake.assemble` (residuals and matrices, `dx` vs `dx(scheme="collapsed")`, on triangle and tetrahedron meshes, degrees 1 and 3: `tests/firedrake/regression/test_quadrature.py::test_collapsed_quadrature_sum_factorisation`). -### (c) Basis-index integration route — Route B chosen +### (c) Basis-index integration route — Route B chosen, Route C's dof reorder adopted underneath it -The element's flat basis index is Morton-ordered (`FIAT.expansions.morton_index`, -using `morton_index2`/`morton_index3` = total-degree-major), while the -factorization is indexed by the lattice multiindex. Three routes were -considered: +The element's flat basis index was originally Morton-ordered +(`FIAT.expansions.morton_index`, using `morton_index2`/`morton_index3` = +total-degree-major); it is now lattice-lexicographic (`FIAT.expansions. +lexicographic_permutation`, `FIAT.hierarchical.LegendreDual`), while the +factorization is indexed by the lattice multiindex. Routes considered: * **Route A — layer-wise `Concatenate`.** Reuse tsfc/spectral.py's `Concatenate`/`unconcatenate` machinery by splitting the basis into @@ -194,10 +195,10 @@ considered: sweeps contract one lattice axis at a time, not one total-degree layer at a time), so it buys the wrong factorization. -* **Route B — Morton gather/scatter via `VariableIndex` (chosen; see (b)).** +* **Route B — gather/scatter via `VariableIndex` (chosen; see (b)).** Keeps FIAT's dof ordering, `element.index_shape`, and - `argument_multiindices` completely untouched; the Morton lookup lives - entirely inside `finat/duffy.py` (`DuffyElement.basis_evaluation` and + `argument_multiindices` completely untouched; the flat-index arithmetic + lives entirely inside `finat/duffy.py` (`DuffyElement.basis_evaluation` and `DuffyElement.duffy_contraction`). No `driver.py` or `kernel_interface/*.py` changes were needed at all, and — after the fem.py-integration refactor above — no bespoke branching in `tsfc/fem.py` @@ -206,23 +207,45 @@ considered: (forward) or one uint load per dof (backward), negligible against the O(p) inner contraction. -* **Route C — reorder FIAT dofs p-major.** Change `Legendre` - (variant="integral") to lattice-lexicographic dof order so the flat index - becomes `offsets[i_1, ..] + i_d` (affine within each innermost run). - Cleanest kernels, but dof ordering is externally visible (checkpoints, - hand-written index hacks, any test with hard-coded dof numbers) and the - offset table is still a lookup, so the win over Route B is small. Not - pursued; only worth it if profiling shows the Morton gather hurts. +* **Route C — reorder FIAT dofs p-major (adopted, underneath Route B).** + `Legendre` (variant="integral") now uses lattice-lexicographic dof order + (`FIAT.hierarchical.LegendreDual` permuted via `FIAT.expansions. + lexicographic_permutation`), so the flat index is `offsets[i_1, .., i_{d-1}] + + i_d` — affine in the innermost coordinate for fixed outer coordinates, + unlike Morton's `(p+q)(p+q+1)/2 + q`, which mixes every coordinate + non-separably. This didn't eliminate Route B's `VariableIndex` gather/ + scatter (see "Deferred" below for why), but it did let `finat/duffy.py`'s + index arithmetic (`_flat_index_expr`/`_inverse_lex_index_exprs`) shrink + from multi-stage triangular/tetrahedral-number arithmetic to one small + `(degree+1,)`-ish table lookup plus a bounded subtraction, both directions. + Reordering FIAT's own dof numbering is externally visible (checkpoints, + hand-written index hacks), judged an acceptable, narrowly-scoped cost + specifically because `Legendre`/`variant="integral"` is a niche element + family — ordinary `Lagrange`/`variant=None` elements are untouched. ## Deferred -* **Route C — reorder FIAT dofs to eliminate the Morton gather/scatter - entirely.** Raised in self-review of PR #5263: renumbering the expansion - set/finite element dofs so the flat dof index *is* the (bounding-box) - lattice multiindex would let the generic `basis_evaluation`/contraction - machinery pick up the separable structure without any table lookup at all - (see Route C above). Left as a follow-up after the fem.py-integration - refactor landed in `finat/duffy.py`, rather than attempted alongside it. +* **Fully eliminating Route B's `VariableIndex` gather/scatter by making + `element.get_indices()` return the lattice multiindex directly.** With the + lattice-lexicographic dof order in place, this was investigated in depth + and found to be blocked by more than "dof ordering is externally visible": + `get_indices()` takes no arguments and is cached once per kernel, so it + cannot distinguish a cell-interior kernel (where the lattice multiindex + applies) from a facet-integral kernel (where `DuffyElement._duffy_applies` + is always `False` and tabulation always uses the dense, flat-`(ndof,)` + FIAT path) — returning a lattice tuple unconditionally would break + `translate_argument` for every facet integral on a DG element (SIPG, + upwinding), not just an edge case. A real fix is possible (pad the dense + facet tabulation to the same `(degree+1,)**d` bounding-box shape at + TSFC-compile time, and give `tsfc/kernel_interface/common.py`'s + `prepare_arguments`/`prepare_coefficient` a bespoke jagged + `gem.FlexiblyIndexed` view — `ComponentTensor(FlexiblyIndexed(flat_var, + ((offset(p), ((q, 1),)),)), (p, q))` — instead of the current pure- + rectangular `gem.reshape`) but that code is shared by every Firedrake + kernel; confirmed via `test_collapsed_quadrature_sum_factorisation`'s + mass-matrix case that it must also handle *two* Duffy arguments combined + in one "A" tensor. Deliberately deferred as a separate, higher-risk + follow-up rather than folded into the dof-reorder work above. * **CG / C0 basis (milestones 3–4):** the C0 recombination makes each basis function a sum of <= 3 separable members (Sherwin–Karniadakis vertex/edge/face recombination); `tabulate_duffy` currently raises `NotImplementedError` for From c28bbb0484c544b0491c569d748c2901da72469c Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 21 Jul 2026 00:35:14 +0100 Subject: [PATCH 06/23] DROP BEFORE MERGE: point CI at the paired FIAT branch Temporarily pins firedrake-fiat to pbrubeck/simplex-sum-factor so CI exercises this branch's paired FIAT changes (dof-order permutation, gem.Delta fix). Revert to @main once the FIAT PR merges. Co-Authored-By: Claude Sonnet 5 --- pyproject.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 74d38ccf8e..a7207c1883 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,8 @@ dependencies = [ # each Firedrake release to a specific UFL minor version (e.g. 2025.3.x) "fenics-ufl @ git+https://github.com/FEniCS/ufl.git@main", # TODO RELEASE - "firedrake-fiat @ git+https://github.com/firedrakeproject/fiat.git@main", + # DROP BEFORE MERGE: pinned to the paired FIAT branch for CI; revert to @main + "firedrake-fiat @ git+https://github.com/firedrakeproject/fiat.git@pbrubeck/simplex-sum-factor", "h5py>3.12.1", "firedrake-rtree", "immutabledict", From 458777d7a9d2df0a0731f210b4b113bda0cd469f Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 21 Jul 2026 13:13:48 +0100 Subject: [PATCH 07/23] Sum-factorise by default --- DESIGN.md | 39 ++++---- tests/firedrake/regression/test_quadrature.py | 9 +- tests/tsfc/test_codegen.py | 16 ++-- tests/tsfc/test_sum_factorisation.py | 89 ++++++++++++------- tsfc/kernel_interface/common.py | 8 ++ 5 files changed, 104 insertions(+), 57 deletions(-) diff --git a/DESIGN.md b/DESIGN.md index 5913158386..8dca11dca1 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -66,12 +66,11 @@ loop bounds — jaggedness is purely a flop optimization. * `test/FIAT/unit/test_polynomial.py`: `test_tabulate_duffy` (values + gradients vs `_tabulate_on_cell`, dims 1–3, variants None/dual, degrees - 0/1/4, points including the collapsed vertex), `test_principal_functions_bubble`, - `test_morton_tables` (`morton_forward_table`/`morton_inverse_table` agree - with `morton_index` and are mutual inverses on the simplex lattice). + 0/1/4, points including the collapsed vertex), `test_principal_functions_bubble`. * `test/finat/test_point_evaluation.py`: `test_duffy_evaluation` vs dense - `basis_evaluation` through the gem interpreter, checking Morton flat-index - agreement and exact zeros outside the lattice. + `basis_evaluation` through the gem interpreter, checking lattice-lexicographic + flat-index agreement (`FIAT.expansions.lexicographic_multiindices`) and exact + zeros outside the lattice. * `tests/tsfc/test_codegen.py::test_jagged_index_codegen`: compiles the 2D jagged Morton-gather contraction through `compile_gem` -> `tsfc.loopy.generate`, executes the C kernel, checks the parametrized ISL domain exists and the @@ -118,9 +117,10 @@ so the original plan of making `argument_multiindices` itself a lattice multiindex (see the old Route-B write-up below) would have changed the local tensor's shape and broken that contract. The implementation instead keeps `element.index_shape` and `argument_multiindices` exactly as they are today, -and confines the lattice multiindex — and all the Duffy/Morton machinery — to -a new `finat/duffy.py` module, reached from `tsfc/fem.py` through ordinary -FInAT element methods rather than through tsfc-side branching: +and confines the lattice multiindex — and all the Duffy gather/scatter +machinery — to a new `finat/duffy.py` module, reached from `tsfc/fem.py` +through ordinary FInAT element methods rather than through tsfc-side +branching: * `finat.duffy.DuffyElement` is a mixin (`finat.spectral.Legendre` is currently its only user) providing `duffy_evaluation` (the lattice-indexed, @@ -130,13 +130,14 @@ FInAT element methods rather than through tsfc-side branching: not a macrocell — the only other case `duffy_evaluation` rejects), `DuffyElement.basis_evaluation` calls `duffy_evaluation` and scatters the lattice tabulation to the standard flat-dof-indexed convention: it - introduces one *fresh* flat dof index `r`, builds the *inverse* Morton - table (`FIAT.expansions.morton_inverse_table`, shape `(ndof, d)`) to get - per-axis lookups `i_t(r)`, and substitutes `multiindex[t] -> - VariableIndex(inverse_table[:, t][r])` throughout `duffy_evaluation`'s - expression tree via `gem.node.MemoizerArg(gem.optimise.filtered_replace_indices)` - (the same substitution mechanism `translate_argument`/`translate_coefficient` - use for canonical quadrature-point reordering; `filtered_replace_indices` + introduces one *fresh* flat dof index `r`, computes each axis's lattice + coordinate as a function of `r` via `_inverse_lex_index_exprs` (small + per-axis-prefix tables from `FIAT.expansions.lexicographic_offsets`, not a + full `(ndof, d)` inverse table), and substitutes `multiindex[t] -> + VariableIndex(i_t(r))` throughout `duffy_evaluation`'s expression tree via + `gem.node.MemoizerArg(gem.optimise.filtered_replace_indices)` (the same + substitution mechanism `translate_argument`/`translate_coefficient` use + for canonical quadrature-point reordering; `filtered_replace_indices` recurses into `VariableIndex.expression`, so the nested `m_t` lookups are rewritten too). The result, wrapped in `gem.ComponentTensor(..., (r,))`, is indistinguishable, from `fiat_to_ufl`/`prepare_arguments`'s point of @@ -157,10 +158,10 @@ FInAT element methods rather than through tsfc-side branching: DuffyElement)`, `ctx.point_set` a `CollapsedTensorProductPointSet`, cell interior, `ctx.unsummed_coefficient_indices` empty) before calling `element.duffy_contraction(mt.local_derivatives, ctx.point_set, entity, - vec, ctx.epsilon)`. `duffy_contraction` builds a forward Morton lookup - table (`FIAT.expansions.morton_forward_table`, shape `(degree+1,)^d`, - clamped to a valid dof so out-of-lattice reads are merely wasted, never - out of bounds — they always multiply a zero tabulation), gathers + vec, ctx.epsilon)`. `duffy_contraction` computes the forward flat index + via `_flat_index_expr` (the same `lexicographic_offsets` tables, clamped + to a valid dof so out-of-lattice reads are merely wasted, never out of + bounds — they always multiply a zero tabulation), gathers `vec[VariableIndex(table[multiindex])]`, and hands `IndexSum(Product(duffy[alpha], vec_r), multiindex)` to `gem.optimise.contraction`, exactly as originally planned: the `m_t` diff --git a/tests/firedrake/regression/test_quadrature.py b/tests/firedrake/regression/test_quadrature.py index 6ea188dfb9..98973e0df0 100644 --- a/tests/firedrake/regression/test_quadrature.py +++ b/tests/firedrake/regression/test_quadrature.py @@ -54,18 +54,21 @@ def test_quadrature_element(mesh, family, mat_type, diagonal): assemble(a, mat_type=mat_type, diagonal=diagonal) +@pytest.mark.parametrize("family", ["DG", "CG"]) @pytest.mark.parametrize("cell", ["triangle", "tetrahedron"]) @pytest.mark.parametrize("degree", [1, 3]) -def test_collapsed_quadrature_sum_factorisation(cell, degree): +def test_collapsed_quadrature_sum_factorisation(cell, degree, family): """``dx(scheme="collapsed")`` on a simplicial "DG"/variant="integral" - (i.e. `finat.spectral.Legendre`) space must produce the same + (i.e. `finat.spectral.Legendre`) or "CG"/variant="integral" (i.e. + `finat.spectral.IntegratedLegendre`, exercising the + `FIAT.expansions.C0_basis` recombination) space must produce the same assembled residual and matrix as the default dense quadrature, even though it takes the sum-factorized (Duffy/lattice) tabulation path in ``tsfc.fem`` rather than the standard dense one. """ mesh = {"triangle": UnitSquareMesh(2, 2), "tetrahedron": UnitCubeMesh(1, 1, 1)}[cell] - V = FunctionSpace(mesh, "DG", degree, variant="integral") + V = FunctionSpace(mesh, family, degree, variant="integral") u = TrialFunction(V) v = TestFunction(V) w = Function(V) diff --git a/tests/tsfc/test_codegen.py b/tests/tsfc/test_codegen.py index efc8ba6453..89bc7bd67b 100644 --- a/tests/tsfc/test_codegen.py +++ b/tests/tsfc/test_codegen.py @@ -75,18 +75,23 @@ def test_jagged_index_codegen(monkeypatch): assert numpy.allclose(u_out, u_ref, rtol=1e-14) +@pytest.mark.parametrize("element_name", ["Legendre", "IntegratedLegendre"]) @pytest.mark.parametrize("cellname,degree", [("triangle", 3), ("tetrahedron", 2)]) -def test_duffy_scatter_and_contract(cellname, degree): +def test_duffy_scatter_and_contract(cellname, degree, element_name): """Route B of the simplex sum-factorization milestone 2 design: `finat.duffy.DuffyElement.basis_evaluation` (the `translate_argument` path) and `finat.duffy.DuffyElement.duffy_contraction` (the `translate_coefficient` path) must reproduce the standard dense FIAT - tabulation, via the Morton dof numbering FIAT already uses, from + tabulation, via the dof numbering FIAT already uses, from `duffy_evaluation`'s lattice-indexed, sum-factorized tabulation. + `Legendre` (continuity=None) reads exactly one lattice point per dof + (lattice-lexicographic order); `IntegratedLegendre` (continuity="C0") + additionally exercises the `FIAT.expansions.C0_basis` recombination, + where each dof combines a handful of lattice points. Verified via `gem.interpreter.evaluate` rather than a compiled loopy kernel: the GEM expressions built here (in particular - `duffy_contraction`'s `gem.VariableIndex`-based Morton index arithmetic) + `duffy_contraction`'s `gem.VariableIndex`-based index arithmetic) schedule fine once merged into a real PyOP2 wrapper kernel (as confirmed via `firedrake.assemble` on real forms), but are not guaranteed schedulable by loopy in isolation -- scheduling an isolated, @@ -96,13 +101,14 @@ def test_duffy_scatter_and_contract(cellname, degree): """ from FIAT.reference_element import UFCTetrahedron, UFCTriangle from finat.quadrature import make_quadrature - from finat.spectral import Legendre + from finat.spectral import Legendre, IntegratedLegendre from gem.gem import Index, Indexed, Variable from gem.interpreter import evaluate from gem.optimise import remove_componenttensors cell = {"triangle": UFCTriangle, "tetrahedron": UFCTetrahedron}[cellname]() - element = Legendre(cell, degree) + element_cls = {"Legendre": Legendre, "IntegratedLegendre": IntegratedLegendre}[element_name] + element = element_cls(cell, degree) ndof = element.space_dimension() quad_rule = make_quadrature(cell, 2 * degree, scheme="collapsed") diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index b7522133bc..b2a0d708f6 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -18,18 +18,6 @@ def helmholtz(cell, degree): return (u*v + dot(grad(u), grad(v)))*dx -def simplex_dg_mass(cell, degree): - # A simplicial DG element whose nodal basis coincides with the Dubiner - # expansion set (finat.spectral.Legendre), so that dx(scheme="collapsed") - # takes the sum-factorized (Duffy/lattice) tabulation path in tsfc/fem.py - # instead of dense tabulation. - m = Mesh(VectorElement('CG', cell, 1)) - V = FunctionSpace(m, FiniteElement('DG', cell, degree, variant='integral')) - u = TrialFunction(V) - v = TestFunction(V) - return inner(u, v) * dx(scheme='collapsed') - - def split_mixed_poisson(cell, degree): m = Mesh(VectorElement('CG', cell, 1)) if cell.cellname in ['interval * interval', 'quadrilateral']: @@ -112,20 +100,6 @@ def test_rhs(cell, order): assert (rates < order).all() -@pytest.mark.parametrize(('cell', 'order'), [(triangle, 4), (tetrahedron, 6)]) -def test_simplex_dg_mass_action(cell, order): - # Matrix-free DG mass-matrix action (milestone 2 of PLAN.md / DESIGN.md): - # the coefficient contraction in translate_coefficient is sum-factorized - # via the Duffy/lattice tabulation, targeting O(p^{d+1}) flops. This - # tests the *action* (right-hand side, like test_rhs above), not full - # bilinear matrix assembly, which is milestone 4 and not yet implemented. - degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) - flops = [count_flops(action(simplex_dg_mass(cell, degree))) - for degree in degrees] - rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) - assert (rates < order).all() - - @pytest.mark.parametrize(('cell', 'order'), [(quadrilateral, 5), (TensorProductCell(interval, interval), 5), @@ -194,7 +168,62 @@ def test_vector_laplace_action(cell, order): assert (rates < order).all() -if __name__ == "__main__": - import os - import sys - pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:]) +def simplex_mass(cell, family, degree): + m = Mesh(VectorElement('CG', cell, 1)) + V = FunctionSpace(m, FiniteElement(family, cell, degree, variant='integral')) + u = TrialFunction(V) + v = TestFunction(V) + return inner(u, v) * dx(scheme='collapsed') + + +def simplex_laplacian(cell, family, degree): + m = Mesh(VectorElement('CG', cell, 1)) + V = FunctionSpace(m, FiniteElement(family, cell, degree, variant='integral')) + u = TrialFunction(V) + v = TestFunction(V) + return inner(grad(u), grad(v)) * dx(scheme='collapsed') + + +@pytest.mark.parametrize('family', ["DG", "CG"]) +@pytest.mark.parametrize(('cell', 'order'), [(triangle, 4), (tetrahedron, 6)]) +def test_simplex_mass_action(cell, family, order): + degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) + flops = [count_flops(action(simplex_mass(cell, family, degree))) + for degree in degrees] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) + assert (rates < order).all() + + +@pytest.mark.parametrize('family', ["DG", "CG"]) +@pytest.mark.parametrize(('cell', 'order'), [(triangle, 4), (tetrahedron, 6)]) +def test_simplex_laplacian_action(cell, family, order): + degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) + flops = [count_flops(action(simplex_laplacian(cell, family, degree))) + for degree in degrees] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) + assert (rates < order).all() + + +# Unlike the `_action` tests above, these compile the bilinear form directly +# (no `action`, so both the test and trial bases are sum-factorized and +# scattered to their flat dof index simultaneously) -- the configuration +# that previously exposed a loopy scheduling bug in +# `finat.duffy._scatter_to_dof_index` (see tsfc/AGENTS.md). +@pytest.mark.parametrize('family', ["DG", "CG"]) +@pytest.mark.parametrize(('cell', 'order'), [(triangle, 6), (tetrahedron, 9)]) +def test_simplex_mass_bilinear(cell, family, order): + degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) + flops = [count_flops(simplex_mass(cell, family, degree)) + for degree in degrees] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) + assert (rates < order).all() + + +@pytest.mark.parametrize('family', ["DG", "CG"]) +@pytest.mark.parametrize(('cell', 'order'), [(triangle, 6), (tetrahedron, 9)]) +def test_simplex_laplacian_bilinear(cell, family, order): + degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) + flops = [count_flops(simplex_laplacian(cell, family, degree)) + for degree in degrees] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) + assert (rates < order).all() diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index 5d61a916aa..677446ed63 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -14,6 +14,7 @@ import numpy from FIAT.reference_element import TensorProductCell from finat.cell_tools import max_complex +from finat.duffy import DuffyElement from finat.quadrature import AbstractQuadratureRule from gem.node import traversal from gem.optimise import constant_fold_zero @@ -344,6 +345,13 @@ def set_quad_rule(params, cell, integral_type, functions): scheme = quad_rule fiat_cell = as_fiat_cell(cell) finat_elements = set(create_element(e) for e in elements if e.family() != "Real") + if scheme == "default" and any(isinstance(finat_el, DuffyElement) for finat_el in finat_elements): + # Sum-factorized (Duffy/lattice) tabulation only kicks in on a + # collapsed-coordinate quadrature rule (see + # finat.duffy.DuffyElement._duffy_applies); default to it + # automatically instead of requiring users to spell out + # dx(scheme="collapsed") to get it. + scheme = "collapsed" fiat_cells = [fiat_cell] + [finat_el.complex for finat_el in finat_elements] if any(c.is_macrocell() for c in fiat_cells): if len(set(c.get_spatial_dimension() for c in fiat_cells)) > 1: From 8a4adb67e2d31ac42cdb666b7b3adad8e9082138 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 21 Jul 2026 13:57:28 +0100 Subject: [PATCH 08/23] Drop special-cased Duffy coefficient contraction translate_coefficient no longer dispatches on isinstance(element, DuffyElement): finat.duffy.DuffyElement.basis_evaluation now returns an already flat-dof-indexed tabulation, so the generic dense contraction path applies uniformly and recovers the same sum-factorized complexity without any special-casing. This also makes the Duffy fast path compose with Vector/TensorElement, which previously fell through to the slow path since TensorFiniteElement is never itself a DuffyElement instance. test_duffy_scatter_and_contract updated to contract generically instead of calling the now-removed duffy_contraction. test_collapsed_quadrature_sum_factorisation now compares against dx(scheme="canonical") instead of the default scheme: same collapsed Gauss-Jacobi points/weights as dx(scheme="collapsed"), but tabulated via the dense FIAT path, isolating the comparison to the sum-factorized tabulation itself rather than to a difference in quadrature rule. Co-Authored-By: Claude Sonnet 5 --- tests/firedrake/regression/test_quadrature.py | 16 ++- tests/tsfc/test_codegen.py | 47 +++---- tsfc/fem.py | 124 ++++++------------ 3 files changed, 77 insertions(+), 110 deletions(-) diff --git a/tests/firedrake/regression/test_quadrature.py b/tests/firedrake/regression/test_quadrature.py index 98973e0df0..672e0fc6e9 100644 --- a/tests/firedrake/regression/test_quadrature.py +++ b/tests/firedrake/regression/test_quadrature.py @@ -62,9 +62,19 @@ def test_collapsed_quadrature_sum_factorisation(cell, degree, family): (i.e. `finat.spectral.Legendre`) or "CG"/variant="integral" (i.e. `finat.spectral.IntegratedLegendre`, exercising the `FIAT.expansions.C0_basis` recombination) space must produce the same - assembled residual and matrix as the default dense quadrature, even + assembled residual and matrix as ``dx(scheme="canonical")``, even though it takes the sum-factorized (Duffy/lattice) tabulation path in ``tsfc.fem`` rather than the standard dense one. + + ``"canonical"`` is the same collapsed Gauss-Jacobi rule as + ``"collapsed"`` (same points, same weights, + `finat.quadrature.collapsed_gauss_jacobi_quadrature`), but as a plain + `finat.point_set.PointSet` rather than a `CollapsedTensorProductPointSet` + -- so `finat.duffy.DuffyElement._duffy_applies` is `False` and it always + takes the dense FIAT tabulation path. Comparing against it (rather than + the default scheme, a different quadrature rule entirely for these + degrees) isolates any discrepancy to the sum-factorized tabulation + itself, not to a difference in quadrature choice. """ mesh = {"triangle": UnitSquareMesh(2, 2), "tetrahedron": UnitCubeMesh(1, 1, 1)}[cell] @@ -77,14 +87,14 @@ def test_collapsed_quadrature_sum_factorisation(cell, degree, family): # translate_coefficient path (forward transform): residual with a # derivative, mixing both the coefficient and argument sum-factorized # tabulations. - L = inner(grad(w), grad(v)) * dx + L = inner(grad(w), grad(v)) * dx(scheme="canonical") L_collapsed = inner(grad(w), grad(v)) * dx(scheme="collapsed") b = assemble(L) b_collapsed = assemble(L_collapsed) assert np.allclose(b.dat.data, b_collapsed.dat.data, rtol=1e-10, atol=1e-10) # translate_argument path (backward transform): mass matrix. - a = inner(u, v) * dx + a = inner(u, v) * dx(scheme="canonical") a_collapsed = inner(u, v) * dx(scheme="collapsed") M = assemble(a).M.values M_collapsed = assemble(a_collapsed).M.values diff --git a/tests/tsfc/test_codegen.py b/tests/tsfc/test_codegen.py index 89bc7bd67b..a62002fb32 100644 --- a/tests/tsfc/test_codegen.py +++ b/tests/tsfc/test_codegen.py @@ -79,19 +79,21 @@ def test_jagged_index_codegen(monkeypatch): @pytest.mark.parametrize("cellname,degree", [("triangle", 3), ("tetrahedron", 2)]) def test_duffy_scatter_and_contract(cellname, degree, element_name): """Route B of the simplex sum-factorization milestone 2 design: - `finat.duffy.DuffyElement.basis_evaluation` (the `translate_argument` - path) and `finat.duffy.DuffyElement.duffy_contraction` (the - `translate_coefficient` path) must reproduce the standard dense FIAT - tabulation, via the dof numbering FIAT already uses, from - `duffy_evaluation`'s lattice-indexed, sum-factorized tabulation. - `Legendre` (continuity=None) reads exactly one lattice point per dof - (lattice-lexicographic order); `IntegratedLegendre` (continuity="C0") - additionally exercises the `FIAT.expansions.C0_basis` recombination, - where each dof combines a handful of lattice points. + `finat.duffy.DuffyElement.basis_evaluation` must reproduce the standard + dense FIAT tabulation, via the dof numbering FIAT already uses, from + `duffy_evaluation`'s lattice-indexed, sum-factorized tabulation, and a + generic contraction of that same flat-dof-indexed result against a + coefficient vector (as `tsfc.fem.translate_coefficient` performs + uniformly for every FInAT element, `DuffyElement` included) must + reproduce the dense tensordot. `Legendre` (continuity=None) reads + exactly one lattice point per dof (lattice-lexicographic order); + `IntegratedLegendre` (continuity="C0") additionally exercises the + `FIAT.expansions.C0_basis` recombination, where each dof combines a + handful of lattice points. Verified via `gem.interpreter.evaluate` rather than a compiled loopy kernel: the GEM expressions built here (in particular - `duffy_contraction`'s `gem.VariableIndex`-based index arithmetic) + `duffy_evaluation`'s `gem.VariableIndex`-based scatter index arithmetic) schedule fine once merged into a real PyOP2 wrapper kernel (as confirmed via `firedrake.assemble` on real forms), but are not guaranteed schedulable by loopy in isolation -- scheduling an isolated, @@ -102,7 +104,7 @@ def test_duffy_scatter_and_contract(cellname, degree, element_name): from FIAT.reference_element import UFCTetrahedron, UFCTriangle from finat.quadrature import make_quadrature from finat.spectral import Legendre, IntegratedLegendre - from gem.gem import Index, Indexed, Variable + from gem.gem import Index, Indexed, IndexSum, Product, Variable from gem.interpreter import evaluate from gem.optimise import remove_componenttensors @@ -133,18 +135,19 @@ def test_duffy_scatter_and_contract(cellname, degree, element_name): assert u_out.fids == (r,) + point_indices assert numpy.allclose(u_out.arr, dense, rtol=1e-12, atol=1e-12) - # translate_coefficient path: contraction against a coefficient vector, - # dispatched through duffy_contraction + # translate_coefficient path: generic contraction of the same + # flat-dof-indexed basis_evaluation result against a coefficient vector + # (no element-specific dispatch, see tsfc.fem.translate_coefficient) c = Variable("c", (ndof,)) - for order in (0, 1): - contracted_dict = element.duffy_contraction(order, point_set, entity, c, epsilon=0.0) - for alpha, contracted in contracted_dict.items(): - dense = dense_dict[alpha].reshape((ndof,) + point_shape) - value, = remove_componenttensors([Indexed(contracted, ())]) - v_out, = evaluate([value], bindings={c: coefficients}) - assert v_out.fids == point_indices - v_ref = numpy.tensordot(coefficients, dense, axes=(0, 0)) - assert numpy.allclose(v_out.arr, v_ref, rtol=1e-12, atol=1e-12) + r = Index(extent=ndof) + for alpha, dense in dense_dict.items(): + dense = dense.reshape((ndof,) + point_shape) + table, = remove_componenttensors([Indexed(scattered_dict[alpha], (r,))]) + value = IndexSum(Product(Indexed(c, (r,)), table), (r,)) + v_out, = evaluate([value], bindings={c: coefficients}) + assert v_out.fids == point_indices + v_ref = numpy.tensordot(coefficients, dense, axes=(0, 0)) + assert numpy.allclose(v_out.arr, v_ref, rtol=1e-12, atol=1e-12) if __name__ == "__main__": diff --git a/tsfc/fem.py b/tsfc/fem.py index 8a6bad3e1a..943089052e 100644 --- a/tsfc/fem.py +++ b/tsfc/fem.py @@ -11,10 +11,9 @@ from FIAT.orientation_utils import Orientation as FIATOrientation from FIAT.reference_element import UFCHexahedron, UFCQuadrilateral, UFCSimplex, make_affine_mapping from FIAT.reference_element import TensorProductCell -from finat.duffy import DuffyElement from finat.physically_mapped import (NeedsCoordinateMappingElement, PhysicalGeometry) -from finat.point_set import CollapsedTensorProductPointSet, PointSet, PointSingleton +from finat.point_set import PointSet, PointSingleton from finat.quadrature import make_quadrature from finat.element_factory import as_fiat_cell, create_element from gem.node import traversal @@ -708,45 +707,6 @@ def fiat_to_ufl(fiat_dict, order): return gem.ComponentTensor(tensor, sigma + delta) -def _use_duffy_contraction(element, ctx): - """Whether the sum-factorized (Duffy/lattice) coefficient contraction - applies. - - This holds exactly when `element` is a simplicial DG element whose - nodal basis coincides with the Dubiner expansion set (any - `finat.duffy.DuffyElement`), evaluation points come from a collapsed - tensor-product quadrature rule (requested via - ``dx(scheme="collapsed")``), and the integral is over the cell - interior. In that case `DuffyElement.duffy_contraction` contracts a - `Coefficient` against the element in O(p^d) space/time using the - lattice multi-index, whereas the standard dense contraction - materializes all O(p^d) basis functions at all O(p^d) points before - contracting. - - The argument (basis evaluation) side needs no such dispatch: it is - handled transparently by `DuffyElement.basis_evaluation`, reached - through the usual `~.PointSetContext.basis_evaluation` call. - - Parameters - ---------- - element : finat.finiteelementbase.FiniteElementBase - The element being tabulated. - ctx : ContextBase - The translation context. - - Returns - ------- - bool - Whether to use `DuffyElement.duffy_contraction` in place of the - standard dense contraction. - """ - return (isinstance(element, DuffyElement) - and isinstance(ctx, PointSetContext) - and isinstance(ctx.point_set, CollapsedTensorProductPointSet) - and ctx.integration_dim == ctx.fiat_cell.get_dimension() - and not ctx.unsummed_coefficient_indices) - - @translate.register(Argument) def translate_argument(terminal, mt, ctx): element = ctx.create_element(terminal.ufl_element(), restriction=mt.restriction) @@ -785,51 +745,45 @@ def translate_coefficient(terminal, mt, ctx): vec = ctx.coefficient(terminal, mt.restriction) element = ctx.create_element(terminal.ufl_element(), restriction=mt.restriction) - if _use_duffy_contraction(element, ctx): - entity_id, = ctx.entity_ids(domain) - value_dict = element.duffy_contraction(mt.local_derivatives, ctx.point_set, - (ctx.integration_dim, entity_id), - vec, ctx.epsilon) + # Collect FInAT tabulation for all entities + per_derivative = collections.defaultdict(list) + for entity_id in ctx.entity_ids(domain): + finat_dict = ctx.basis_evaluation(element, mt, entity_id) + for alpha, table in finat_dict.items(): + # Filter out irrelevant derivatives + if sum(alpha) == mt.local_derivatives: + # A numerical hack that FFC used to apply on FIAT + # tables still lives on after ditching FFC and + # switching to FInAT. + table = ffc_rounding(table, ctx.epsilon) + per_derivative[alpha].append(table) + + # Merge entity tabulations for each derivative + if len(ctx.entity_ids(domain)) == 1: + def take_singleton(xs): + x, = xs # asserts singleton + return x + per_derivative = {alpha: take_singleton(tables) + for alpha, tables in per_derivative.items()} else: - # Collect FInAT tabulation for all entities - per_derivative = collections.defaultdict(list) - for entity_id in ctx.entity_ids(domain): - finat_dict = ctx.basis_evaluation(element, mt, entity_id) - for alpha, table in finat_dict.items(): - # Filter out irrelevant derivatives - if sum(alpha) == mt.local_derivatives: - # A numerical hack that FFC used to apply on FIAT - # tables still lives on after ditching FFC and - # switching to FInAT. - table = ffc_rounding(table, ctx.epsilon) - per_derivative[alpha].append(table) - - # Merge entity tabulations for each derivative - if len(ctx.entity_ids(domain)) == 1: - def take_singleton(xs): - x, = xs # asserts singleton - return x - per_derivative = {alpha: take_singleton(tables) - for alpha, tables in per_derivative.items()} - else: - f = ctx.entity_number(domain, mt.restriction) - per_derivative = {alpha: gem.select_expression(tables, f) - for alpha, tables in per_derivative.items()} - - # Coefficient evaluation - beta = ctx.index_cache.setdefault(terminal.ufl_element(), element.get_indices()) - zeta = element.get_value_indices() - vec_beta, = gem.optimise.remove_componenttensors([gem.Indexed(vec, beta)]) - value_dict = {} - for alpha, table in per_derivative.items(): - table_qi = gem.Indexed(table, beta + zeta) - summands = [] - for var, expr in unconcatenate([(vec_beta, table_qi)], ctx.index_cache): - indices = tuple(i for i in var.index_ordering() if i not in ctx.unsummed_coefficient_indices) - value = gem.IndexSum(gem.Product(expr, var), indices) - summands.append(gem.optimise.contraction(value)) - optimised_value = gem.optimise.make_sum(summands) - value_dict[alpha] = gem.ComponentTensor(optimised_value, zeta) + f = ctx.entity_number(domain, mt.restriction) + per_derivative = {alpha: gem.select_expression(tables, f) + for alpha, tables in per_derivative.items()} + + # Coefficient evaluation + beta = ctx.index_cache.setdefault(terminal.ufl_element(), element.get_indices()) + zeta = element.get_value_indices() + vec_beta, = gem.optimise.remove_componenttensors([gem.Indexed(vec, beta)]) + value_dict = {} + for alpha, table in per_derivative.items(): + table_qi = gem.Indexed(table, beta + zeta) + summands = [] + for var, expr in unconcatenate([(vec_beta, table_qi)], ctx.index_cache): + indices = tuple(i for i in var.index_ordering() if i not in ctx.unsummed_coefficient_indices) + value = gem.IndexSum(gem.Product(expr, var), indices) + summands.append(gem.optimise.contraction(value)) + optimised_value = gem.optimise.make_sum(summands) + value_dict[alpha] = gem.ComponentTensor(optimised_value, zeta) # Change from FIAT to UFL arrangement result = fiat_to_ufl(value_dict, mt.local_derivatives) From b91b2663998cb5330acc0568d9e70a90f2ce30f9 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 22 Jul 2026 13:25:32 +0100 Subject: [PATCH 09/23] codegen fixes --- tsfc/kernel_interface/common.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index 677446ed63..7cc649495f 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -17,7 +17,7 @@ from finat.duffy import DuffyElement from finat.quadrature import AbstractQuadratureRule from gem.node import traversal -from gem.optimise import constant_fold_zero +from gem.optimise import constant_fold_zero, unflatten_returns from gem.optimise import remove_componenttensors as prune from numpy import asarray from tsfc import fem @@ -211,6 +211,8 @@ def compile_gem(self, ctx): assignments.extend(mode.flatten(var_reps.items(), ctx['index_cache'])) if assignments: + # Rewrite flat FlattenedTensor scatters as jagged lattice loops + assignments = unflatten_returns(assignments) return_variables, expressions = zip(*assignments) else: return_variables = [] From 91cb6bec0e9d80ae9df870c3012136668a9c5903 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 22 Jul 2026 17:03:52 +0100 Subject: [PATCH 10/23] WIP --- tsfc/spectral.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tsfc/spectral.py b/tsfc/spectral.py index 69e471104e..f66b759c9d 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -163,6 +163,12 @@ def prune(factors): args = [f for f in factors if f != one] assert set(var_indices) == set(variable.free_indices) + # Indices newly exposed in the variable by a VariableIndex delta + # cancellation (e.g. the nnz index of a `gem.SparseMatrix` row + # scatter, ``y[r] -> y[rows[p]]``) are scatter indices, not + # contraction indices: they must stay free in the expression to + # match the variable. + sum_indices = [i for i in sum_indices if i not in variable.free_indices] return variable, sum_indices, args, rest From 8a830b37f56ce0ab003eb44377de65dc5afa427a Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 23 Jul 2026 10:42:20 +0100 Subject: [PATCH 11/23] test temporaries do not explode --- tests/tsfc/test_sum_factorisation.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index b2a0d708f6..da59109f7f 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -204,6 +204,14 @@ def test_simplex_laplacian_action(cell, family, order): assert (rates < order).all() +@pytest.mark.parametrize('family', ["DG", "CG"]) +def test_simplex_laplacian_action_compact_codegen(family): + form = action(simplex_laplacian(triangle, family, 3)) + kernel, = compile_form(form, parameters=dict(mode='spectral')) + temporaries = kernel.ast.default_entrypoint.temporary_variables + assert len(temporaries) < 70 + + # Unlike the `_action` tests above, these compile the bilinear form directly # (no `action`, so both the test and trial bases are sum-factorized and # scattered to their flat dof index simultaneously) -- the configuration From 35f495f8784021b5574bddb89f2840728f03dddf Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 23 Jul 2026 15:40:34 +0100 Subject: [PATCH 12/23] Test Bernstein --- .../regression/test_helmholtz_bernstein.py | 8 ++++---- tests/tsfc/test_sum_factorisation.py | 16 +++++++++------- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/tests/firedrake/regression/test_helmholtz_bernstein.py b/tests/firedrake/regression/test_helmholtz_bernstein.py index ce7ff8f4f3..5130ba60b2 100644 --- a/tests/firedrake/regression/test_helmholtz_bernstein.py +++ b/tests/firedrake/regression/test_helmholtz_bernstein.py @@ -27,7 +27,7 @@ def mesh(request): def test_bernstein(mesh, degree): # Solve with Bernstein polynomials B = FunctionSpace(mesh, "Bernstein", degree) - xb = helmholtz(B) + xb = helmholtz(B, dx(scheme="collapsed")) # Solve with Lagrange polynomials L = FunctionSpace(mesh, "Lagrange", degree) @@ -39,15 +39,15 @@ def test_bernstein(mesh, degree): assert np.allclose(xl.dat.data, xp.dat.data) -def helmholtz(V): +def helmholtz(V, measure=dx): # Define variational problem u = TrialFunction(V) v = TestFunction(V) f = Function(V) x = SpatialCoordinate(V.mesh()) f.project(np.prod([cos(2*pi*xi) for xi in x])) - a = (inner(grad(u), grad(v)) + inner(u, v)) * dx - L = inner(f, v) * dx + a = (inner(grad(u), grad(v)) + inner(u, v)) * measure + L = inner(f, v) * measure # Compute solution x = Function(V) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index da59109f7f..d1f9b006c8 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -170,7 +170,8 @@ def test_vector_laplace_action(cell, order): def simplex_mass(cell, family, degree): m = Mesh(VectorElement('CG', cell, 1)) - V = FunctionSpace(m, FiniteElement(family, cell, degree, variant='integral')) + variant = None if family == "Bernstein" else "integral" + V = FunctionSpace(m, FiniteElement(family, cell, degree, variant=variant)) u = TrialFunction(V) v = TestFunction(V) return inner(u, v) * dx(scheme='collapsed') @@ -178,13 +179,14 @@ def simplex_mass(cell, family, degree): def simplex_laplacian(cell, family, degree): m = Mesh(VectorElement('CG', cell, 1)) - V = FunctionSpace(m, FiniteElement(family, cell, degree, variant='integral')) + variant = None if family == "Bernstein" else "integral" + V = FunctionSpace(m, FiniteElement(family, cell, degree, variant=variant)) u = TrialFunction(V) v = TestFunction(V) return inner(grad(u), grad(v)) * dx(scheme='collapsed') -@pytest.mark.parametrize('family', ["DG", "CG"]) +@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) @pytest.mark.parametrize(('cell', 'order'), [(triangle, 4), (tetrahedron, 6)]) def test_simplex_mass_action(cell, family, order): degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) @@ -194,7 +196,7 @@ def test_simplex_mass_action(cell, family, order): assert (rates < order).all() -@pytest.mark.parametrize('family', ["DG", "CG"]) +@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) @pytest.mark.parametrize(('cell', 'order'), [(triangle, 4), (tetrahedron, 6)]) def test_simplex_laplacian_action(cell, family, order): degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) @@ -204,7 +206,7 @@ def test_simplex_laplacian_action(cell, family, order): assert (rates < order).all() -@pytest.mark.parametrize('family', ["DG", "CG"]) +@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) def test_simplex_laplacian_action_compact_codegen(family): form = action(simplex_laplacian(triangle, family, 3)) kernel, = compile_form(form, parameters=dict(mode='spectral')) @@ -217,7 +219,7 @@ def test_simplex_laplacian_action_compact_codegen(family): # scattered to their flat dof index simultaneously) -- the configuration # that previously exposed a loopy scheduling bug in # `finat.duffy._scatter_to_dof_index` (see tsfc/AGENTS.md). -@pytest.mark.parametrize('family', ["DG", "CG"]) +@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) @pytest.mark.parametrize(('cell', 'order'), [(triangle, 6), (tetrahedron, 9)]) def test_simplex_mass_bilinear(cell, family, order): degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) @@ -227,7 +229,7 @@ def test_simplex_mass_bilinear(cell, family, order): assert (rates < order).all() -@pytest.mark.parametrize('family', ["DG", "CG"]) +@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) @pytest.mark.parametrize(('cell', 'order'), [(triangle, 6), (tetrahedron, 9)]) def test_simplex_laplacian_bilinear(cell, family, order): degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) From b8a446b129c6a0f86f9ee38f375179fc6eff63cf Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 23 Jul 2026 18:24:25 +0100 Subject: [PATCH 13/23] Tighten convergence rate --- tests/tsfc/test_sum_factorisation.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index d1f9b006c8..4d55cd7e9a 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -187,7 +187,7 @@ def simplex_laplacian(cell, family, degree): @pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) -@pytest.mark.parametrize(('cell', 'order'), [(triangle, 4), (tetrahedron, 6)]) +@pytest.mark.parametrize(('cell', 'order'), [(triangle, 3), (tetrahedron, 4)]) def test_simplex_mass_action(cell, family, order): degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) flops = [count_flops(action(simplex_mass(cell, family, degree))) @@ -197,7 +197,7 @@ def test_simplex_mass_action(cell, family, order): @pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) -@pytest.mark.parametrize(('cell', 'order'), [(triangle, 4), (tetrahedron, 6)]) +@pytest.mark.parametrize(('cell', 'order'), [(triangle, 3), (tetrahedron, 4.4)]) def test_simplex_laplacian_action(cell, family, order): degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) flops = [count_flops(action(simplex_laplacian(cell, family, degree))) @@ -214,13 +214,8 @@ def test_simplex_laplacian_action_compact_codegen(family): assert len(temporaries) < 70 -# Unlike the `_action` tests above, these compile the bilinear form directly -# (no `action`, so both the test and trial bases are sum-factorized and -# scattered to their flat dof index simultaneously) -- the configuration -# that previously exposed a loopy scheduling bug in -# `finat.duffy._scatter_to_dof_index` (see tsfc/AGENTS.md). @pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) -@pytest.mark.parametrize(('cell', 'order'), [(triangle, 6), (tetrahedron, 9)]) +@pytest.mark.parametrize(('cell', 'order'), [(triangle, 5), (tetrahedron, 7)]) def test_simplex_mass_bilinear(cell, family, order): degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) flops = [count_flops(simplex_mass(cell, family, degree)) @@ -230,7 +225,7 @@ def test_simplex_mass_bilinear(cell, family, order): @pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) -@pytest.mark.parametrize(('cell', 'order'), [(triangle, 6), (tetrahedron, 9)]) +@pytest.mark.parametrize(('cell', 'order'), [(triangle, 5), (tetrahedron, 7)]) def test_simplex_laplacian_bilinear(cell, family, order): degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) flops = [count_flops(simplex_laplacian(cell, family, degree)) From 6e30e230f613aac82b3d6581c45ef8daac131e56 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 24 Jul 2026 11:07:18 +0100 Subject: [PATCH 14/23] SparseMatrix is just syntax sugar --- tests/tsfc/test_sum_factorisation.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index 4d55cd7e9a..949735480f 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -214,6 +214,19 @@ def test_simplex_laplacian_action_compact_codegen(family): assert len(temporaries) < 70 +def test_bernstein_laplacian_action_compact_literals(): + degree = 5 + form = action(simplex_laplacian(tetrahedron, "Bernstein", degree)) + kernel, = compile_form(form, parameters=dict(mode='spectral')) + temporaries = kernel.ast.default_entrypoint.temporary_variables + literals = [numpy.asarray(temporary.initializer) + for temporary in temporaries.values() + if temporary.initializer is not None] + lattice_size = (degree + 1) ** 3 + assert max(literal.size for literal in literals) <= lattice_size + assert sum(literal.size for literal in literals) < 10 * lattice_size + + @pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) @pytest.mark.parametrize(('cell', 'order'), [(triangle, 5), (tetrahedron, 7)]) def test_simplex_mass_bilinear(cell, family, order): From 17adebbef7c25732998a323196dcee8dc02fa2a1 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 24 Jul 2026 22:11:14 +0100 Subject: [PATCH 15/23] Address simplex sum factorisation review --- DESIGN.md | 258 ------------------ .../regression/test_helmholtz_bernstein.py | 8 +- tests/firedrake/regression/test_quadrature.py | 28 +- tests/tsfc/test_codegen.py | 75 ----- tests/tsfc/test_sum_factorisation.py | 2 +- tsfc/kernel_interface/common.py | 13 +- tsfc/spectral.py | 9 +- 7 files changed, 20 insertions(+), 373 deletions(-) delete mode 100644 DESIGN.md diff --git a/DESIGN.md b/DESIGN.md deleted file mode 100644 index 8dca11dca1..0000000000 --- a/DESIGN.md +++ /dev/null @@ -1,258 +0,0 @@ -# Sum factorization on simplices: status and milestone-2 design - -Companion to `PLAN.md`. Everything in "Implemented" is validated by tests. - -## Implemented (milestone 1 + jagged-loop infrastructure) - -### FIAT (`FIAT/expansions.py`) - -* `dubiner_jacobi_parameters(codim, m, variant)` and `dubiner_norm2(d, m, i, variant)`: - shared helpers extracted from `dubiner_recurrence` (behavior-preserving). -* `principal_functions(n, eta, axis, order, variant)`: 1D tabulations of the - Karniadakis–Sherwin principal functions `G[m, i](eta) = norm * w(eta)^m * g_{m,i}(eta)` - with `w = (1 - eta)/2`, as tables `'V'` (values), `'D'` (d/d eta), `'W'` - (`w^{m-1} g norm`, zeroed at `m = 0`), `'tD'` (`(1+eta)/2 * D`). -* `ExpansionSet.tabulate_duffy(n, eta_pts, order, cell)`: the separable - tabulation on collapsed points. On the reference simplex, the Dubiner basis - factorizes exactly as - `phi_{(i_1..i_d)}(eta) = scale * prod_t G_t[m_t, i_t](eta_t)` with - `m_t = i_1 + ... + i_{t-1}`. First derivatives use the closed-form chain rule - `d eta_t / d xi_k = ((1+eta_t)/2)^{[k>t]} * prod_{u>t} 1/w_u` (k >= t), which - yields exactly `k` separable terms for `d phi / d xi_k`; the affine cell map - is applied on top. Raises `NotImplementedError` for `order > 1` and for C0 - (`continuity is not None`) expansion sets. - -### finat (`finat/point_set.py`, `finat/spectral.py`) - -* `CollapsedTensorProductPointSet`: 1D factor point sets in collapsed - coordinates on `[0, 1]`, mapped to the simplex by the Duffy map - `x_t = eta_t * prod_{u>t} (1 - eta_u)`. -* `Legendre.duffy_evaluation(order, ps, entity=None)`: returns - `(multiindex, result)` where `multiindex` enumerates the basis by lattice - indices `(i_1, ..., i_d)` (a tuple of `gem.JaggedIndex`, see below) and - `result[alpha]` is a scalar gem expression built from `gem.Literal` 1D tables - contracted per axis. The weight exponents `m_t` are gem expressions: `0`, - `i_1`, then `VariableIndex` lookups into clamped uint index tables. - -**Zero-padding invariant** (load-bearing for everything below): lattice indices -range over the rectangular box `(degree+1)^d`; all literal index tables are -clamped with `numpy.minimum`, and out-of-lattice entries (`sum(i_t) > degree`) -tabulate to exactly zero. Memory safety and correctness never depend on jagged -loop bounds — jaggedness is purely a flop optimization. - -### gem (`gem/gem.py`, `gem/interpreter.py`) - -* `JaggedIndex(Index)`: free index with a `parents` tuple; iteration bound - `0 <= i < extent - sum(parents)`. `.extent` remains the static rectangular - bound, so every consumer that ignores jaggedness stays correct (via the - zero-padding invariant). Picklable; exported in `__all__`. -* `gem.interpreter._evaluate_indexed`: repaired the bit-rotted `VariableIndex` - path, including a gather path for index expressions with free indices. - -### tsfc (`tsfc/loopy.py`) - -* `LoopyContext.index_parents`: iname -> parent inames, recorded in - `statement_for` when an `imp.For` loop index is a `JaggedIndex` and all its - parents are active (i.e. the loop is nested inside them); otherwise the - rectangular bound is kept (correct by the invariant). -* `create_domains(indices, index_parents=None)`: emits parametrized ISL sets - `[p] -> {[q] : 0 <= q < E - p}` for jagged inames. Loopy nests these domains - automatically; no loopy changes were needed (verified experimentally first). -* `ComponentTensor` materialization stays rectangular on purpose: temporaries - are always fully initialized, so a jagged write can never leave garbage that - a rectangular read later observes. - -### Tests - -* `test/FIAT/unit/test_polynomial.py`: `test_tabulate_duffy` (values + - gradients vs `_tabulate_on_cell`, dims 1–3, variants None/dual, degrees - 0/1/4, points including the collapsed vertex), `test_principal_functions_bubble`. -* `test/finat/test_point_evaluation.py`: `test_duffy_evaluation` vs dense - `basis_evaluation` through the gem interpreter, checking lattice-lexicographic - flat-index agreement (`FIAT.expansions.lexicographic_multiindices`) and exact - zeros outside the lattice. -* `tests/tsfc/test_codegen.py::test_jagged_index_codegen`: compiles the 2D - jagged Morton-gather contraction through `compile_gem` -> `tsfc.loopy.generate`, - executes the C kernel, checks the parametrized ISL domain exists and the - numbers match numpy. -* `tests/tsfc/test_pickle_gem.py::test_pickle_jagged_index`. - -## Milestone 2: O(p^d)-per-dof matrix-free DG residual — implemented (Route B) - -A DG residual action has three phases per cell: - -1. **Forward transform:** evaluate `u` (and `grad u`) at quadrature points from - coefficients `c`. Sum-factorized on collapsed points this is a sequence of d - per-axis contractions with jagged intermediate temporaries, e.g. in 3D - `T1[p, q, k] = sum_r C[p, q, r] * G3[p+q, r, k]`, - `T2[p, j, k] = sum_q T1[p, q, k] * G2[p, q, j]`, - `u[i, j, k] = sum_p T2[p, j, k] * G1[p, i]` — O(p^{d+1}) total. -2. **Pointwise:** multiply by quadrature weights, geometry, coefficients. -3. **Backward transform:** contract against the test function tables (the - transpose sweep), scattering into the residual vector. - -### (a) Collapsed quadrature rule — implemented - -`finat/quadrature.py` now has `CollapsedTensorProductQuadratureRule` (1D -Gauss–Jacobi factor rules in collapsed coordinates; axis `u` carries the Jacobi -weight `(1 - eta_u)^u`, which absorbs the Duffy Jacobian, so the simplex -weights are per-axis products) and a `scheme="collapsed"` branch in -`make_quadrature` building it via `collapsed_gauss_jacobi_quadrature`. Since -`tsfc/fem.py::get_quadrature_rule` passes the UFL measure's scheme metadata -straight to `make_quadrature`, `dx(scheme="collapsed")` already produces the -structured rule with **no** tsfc changes. Tested against FIAT's canonical -collapsed scheme on all monomials up to the requested degree -(`test/finat/test_quadrature.py::test_collapsed_quadrature`). - -### (b) fem.py integration — Route B, no driver/kernel-interface changes - -The standard path in `tsfc/fem.py` calls `element.basis_evaluation(order, ps, -entity)` and contracts the resulting `(ndof,)`-shaped tables with the -element's flat basis index (`element.get_indices()`), and `translate_argument` -extracts one flat entry with `ctx.argument_multiindices[number]`. The local -element tensor's shape (`element.index_shape` in -`kernel_interface/common.py::prepare_arguments`) and `argument_multiindices` -are flat and ndof-based *everywhere else in the kernel-interface/PyOP2 stack*, -so the original plan of making `argument_multiindices` itself a lattice -multiindex (see the old Route-B write-up below) would have changed the local -tensor's shape and broken that contract. The implementation instead keeps -`element.index_shape` and `argument_multiindices` exactly as they are today, -and confines the lattice multiindex — and all the Duffy gather/scatter -machinery — to a new `finat/duffy.py` module, reached from `tsfc/fem.py` -through ordinary FInAT element methods rather than through tsfc-side -branching: - -* `finat.duffy.DuffyElement` is a mixin (`finat.spectral.Legendre` is - currently its only user) providing `duffy_evaluation` (the lattice-indexed, - sum-factorized tabulation) plus two dispatch points: - * **Backward transform / `basis_evaluation` override.** When `ps` is a - `CollapsedTensorProductPointSet` and the entity is the cell interior (and - not a macrocell — the only other case `duffy_evaluation` rejects), - `DuffyElement.basis_evaluation` calls `duffy_evaluation` and scatters the - lattice tabulation to the standard flat-dof-indexed convention: it - introduces one *fresh* flat dof index `r`, computes each axis's lattice - coordinate as a function of `r` via `_inverse_lex_index_exprs` (small - per-axis-prefix tables from `FIAT.expansions.lexicographic_offsets`, not a - full `(ndof, d)` inverse table), and substitutes `multiindex[t] -> - VariableIndex(i_t(r))` throughout `duffy_evaluation`'s expression tree via - `gem.node.MemoizerArg(gem.optimise.filtered_replace_indices)` (the same - substitution mechanism `translate_argument`/`translate_coefficient` use - for canonical quadrature-point reordering; `filtered_replace_indices` - recurses into `VariableIndex.expression`, so the nested `m_t` lookups are - rewritten too). The result, wrapped in `gem.ComponentTensor(..., (r,))`, - is indistinguishable, from `fiat_to_ufl`/`prepare_arguments`'s point of - view, from a standard dense `basis_evaluation` tabulation — so - `translate_argument` in `tsfc/fem.py` needs **no special case at all**: it - always calls `ctx.basis_evaluation(element, mt, entity_id)`, exactly as - for any other element. `gem.optimise.contraction` never runs on this - side (there is no sum to hoist yet at this stage); the sum-factorized - quadrature contraction happens later, per dof, when - `vanilla.py`/`spectral.py` process the quadrature `IndexSum` — the - collapsed quadrature's own per-axis structure is what still delivers the - O(p) win per axis there. - * **Forward transform / `duffy_contraction`.** Unlike `basis_evaluation`, - coefficient contraction has no generic per-element hook to dispatch - through (it additionally needs the coefficient's dof vector `vec`), so - `translate_coefficient` in `tsfc/fem.py` keeps a small - `_use_duffy_contraction(element, ctx)` guard (`isinstance(element, - DuffyElement)`, `ctx.point_set` a `CollapsedTensorProductPointSet`, cell - interior, `ctx.unsummed_coefficient_indices` empty) before calling - `element.duffy_contraction(mt.local_derivatives, ctx.point_set, entity, - vec, ctx.epsilon)`. `duffy_contraction` computes the forward flat index - via `_flat_index_expr` (the same `lexicographic_offsets` tables, clamped - to a valid dof so out-of-lattice reads are merely wasted, never out of - bounds — they always multiply a zero tabulation), gathers - `vec[VariableIndex(table[multiindex])]`, and hands - `IndexSum(Product(duffy[alpha], vec_r), multiindex)` to - `gem.optimise.contraction`, exactly as originally planned: the `m_t` - `VariableIndex` couplings inside `duffy_evaluation`'s own expression make - the per-axis free-index sets nested (`{i_1} ⊂ {i_1, i_2} ⊂ ...`), so - `contraction` finds the innermost-axis-first Karniadakis–Sherwin sweep by - itself — no bespoke sum-factorization code needed. The result is wrapped - back into a `gem.ComponentTensor` over `element.get_value_indices()` - (empty for the scalar `Legendre` element), so it slots into `fiat_to_ufl` - exactly like a standard dense tabulation would. - -Both dispatch points were validated to ~1e-13/1e-14 against FIAT's dense -`tabulate()`, via compiled-and-executed loopy kernels, for values and first -derivatives on triangles and tetrahedra -(`tests/tsfc/test_codegen.py::test_duffy_scatter_and_contract`), and end to -end through `firedrake.assemble` (residuals and matrices, `dx` vs -`dx(scheme="collapsed")`, on triangle and tetrahedron meshes, degrees 1 and 3: -`tests/firedrake/regression/test_quadrature.py::test_collapsed_quadrature_sum_factorisation`). - -### (c) Basis-index integration route — Route B chosen, Route C's dof reorder adopted underneath it - -The element's flat basis index was originally Morton-ordered -(`FIAT.expansions.morton_index`, using `morton_index2`/`morton_index3` = -total-degree-major); it is now lattice-lexicographic (`FIAT.expansions. -lexicographic_permutation`, `FIAT.hierarchical.LegendreDual`), while the -factorization is indexed by the lattice multiindex. Routes considered: - -* **Route A — layer-wise `Concatenate`.** Reuse tsfc/spectral.py's - `Concatenate`/`unconcatenate` machinery by splitting the basis into - contiguous Morton layers of fixed total degree `s`. Rejected: the layer - decomposition does not align with the per-axis contraction structure (the - sweeps contract one lattice axis at a time, not one total-degree layer at a - time), so it buys the wrong factorization. - -* **Route B — gather/scatter via `VariableIndex` (chosen; see (b)).** - Keeps FIAT's dof ordering, `element.index_shape`, and - `argument_multiindices` completely untouched; the flat-index arithmetic - lives entirely inside `finat/duffy.py` (`DuffyElement.basis_evaluation` and - `DuffyElement.duffy_contraction`). No `driver.py` or - `kernel_interface/*.py` changes were needed at all, and — after the - fem.py-integration refactor above — no bespoke branching in `tsfc/fem.py` - either for the argument side; the lattice multiindex never escapes - `finat/duffy.py`. The indirection costs one uint load per accumulation - (forward) or one uint load per dof (backward), negligible against the O(p) - inner contraction. - -* **Route C — reorder FIAT dofs p-major (adopted, underneath Route B).** - `Legendre` (variant="integral") now uses lattice-lexicographic dof order - (`FIAT.hierarchical.LegendreDual` permuted via `FIAT.expansions. - lexicographic_permutation`), so the flat index is `offsets[i_1, .., i_{d-1}] - + i_d` — affine in the innermost coordinate for fixed outer coordinates, - unlike Morton's `(p+q)(p+q+1)/2 + q`, which mixes every coordinate - non-separably. This didn't eliminate Route B's `VariableIndex` gather/ - scatter (see "Deferred" below for why), but it did let `finat/duffy.py`'s - index arithmetic (`_flat_index_expr`/`_inverse_lex_index_exprs`) shrink - from multi-stage triangular/tetrahedral-number arithmetic to one small - `(degree+1,)`-ish table lookup plus a bounded subtraction, both directions. - Reordering FIAT's own dof numbering is externally visible (checkpoints, - hand-written index hacks), judged an acceptable, narrowly-scoped cost - specifically because `Legendre`/`variant="integral"` is a niche element - family — ordinary `Lagrange`/`variant=None` elements are untouched. - -## Deferred - -* **Fully eliminating Route B's `VariableIndex` gather/scatter by making - `element.get_indices()` return the lattice multiindex directly.** With the - lattice-lexicographic dof order in place, this was investigated in depth - and found to be blocked by more than "dof ordering is externally visible": - `get_indices()` takes no arguments and is cached once per kernel, so it - cannot distinguish a cell-interior kernel (where the lattice multiindex - applies) from a facet-integral kernel (where `DuffyElement._duffy_applies` - is always `False` and tabulation always uses the dense, flat-`(ndof,)` - FIAT path) — returning a lattice tuple unconditionally would break - `translate_argument` for every facet integral on a DG element (SIPG, - upwinding), not just an edge case. A real fix is possible (pad the dense - facet tabulation to the same `(degree+1,)**d` bounding-box shape at - TSFC-compile time, and give `tsfc/kernel_interface/common.py`'s - `prepare_arguments`/`prepare_coefficient` a bespoke jagged - `gem.FlexiblyIndexed` view — `ComponentTensor(FlexiblyIndexed(flat_var, - ((offset(p), ((q, 1),)),)), (p, q))` — instead of the current pure- - rectangular `gem.reshape`) but that code is shared by every Firedrake - kernel; confirmed via `test_collapsed_quadrature_sum_factorisation`'s - mass-matrix case that it must also handle *two* Duffy arguments combined - in one "A" tensor. Deliberately deferred as a separate, higher-risk - follow-up rather than folded into the dof-reorder work above. -* **CG / C0 basis (milestones 3–4):** the C0 recombination makes each basis - function a sum of <= 3 separable members (Sherwin–Karniadakis vertex/edge/face - recombination); `tabulate_duffy` currently raises `NotImplementedError` for - `continuity is not None`. The factored-term representation - (`alpha -> [(coeff, factors), ...]`) was chosen so C0 can extend it by - returning more terms per basis function. -* **Derivative order > 1:** raises `NotImplementedError`. -* **Macro cells** (`is_macrocell()`): raises `NotImplementedError` in - `duffy_evaluation`. diff --git a/tests/firedrake/regression/test_helmholtz_bernstein.py b/tests/firedrake/regression/test_helmholtz_bernstein.py index 5130ba60b2..ce7ff8f4f3 100644 --- a/tests/firedrake/regression/test_helmholtz_bernstein.py +++ b/tests/firedrake/regression/test_helmholtz_bernstein.py @@ -27,7 +27,7 @@ def mesh(request): def test_bernstein(mesh, degree): # Solve with Bernstein polynomials B = FunctionSpace(mesh, "Bernstein", degree) - xb = helmholtz(B, dx(scheme="collapsed")) + xb = helmholtz(B) # Solve with Lagrange polynomials L = FunctionSpace(mesh, "Lagrange", degree) @@ -39,15 +39,15 @@ def test_bernstein(mesh, degree): assert np.allclose(xl.dat.data, xp.dat.data) -def helmholtz(V, measure=dx): +def helmholtz(V): # Define variational problem u = TrialFunction(V) v = TestFunction(V) f = Function(V) x = SpatialCoordinate(V.mesh()) f.project(np.prod([cos(2*pi*xi) for xi in x])) - a = (inner(grad(u), grad(v)) + inner(u, v)) * measure - L = inner(f, v) * measure + a = (inner(grad(u), grad(v)) + inner(u, v)) * dx + L = inner(f, v) * dx # Compute solution x = Function(V) diff --git a/tests/firedrake/regression/test_quadrature.py b/tests/firedrake/regression/test_quadrature.py index 672e0fc6e9..5c9f39115b 100644 --- a/tests/firedrake/regression/test_quadrature.py +++ b/tests/firedrake/regression/test_quadrature.py @@ -54,35 +54,19 @@ def test_quadrature_element(mesh, family, mat_type, diagonal): assemble(a, mat_type=mat_type, diagonal=diagonal) -@pytest.mark.parametrize("family", ["DG", "CG"]) +@pytest.mark.parametrize("family", ["DG", "CG", "Bernstein"]) @pytest.mark.parametrize("cell", ["triangle", "tetrahedron"]) @pytest.mark.parametrize("degree", [1, 3]) def test_collapsed_quadrature_sum_factorisation(cell, degree, family): - """``dx(scheme="collapsed")`` on a simplicial "DG"/variant="integral" - (i.e. `finat.spectral.Legendre`) or "CG"/variant="integral" (i.e. - `finat.spectral.IntegratedLegendre`, exercising the - `FIAT.expansions.C0_basis` recombination) space must produce the same - assembled residual and matrix as ``dx(scheme="canonical")``, even - though it takes the sum-factorized (Duffy/lattice) tabulation path - in ``tsfc.fem`` rather than the standard dense one. - - ``"canonical"`` is the same collapsed Gauss-Jacobi rule as - ``"collapsed"`` (same points, same weights, - `finat.quadrature.collapsed_gauss_jacobi_quadrature`), but as a plain - `finat.point_set.PointSet` rather than a `CollapsedTensorProductPointSet` - -- so `finat.duffy.DuffyElement._duffy_applies` is `False` and it always - takes the dense FIAT tabulation path. Comparing against it (rather than - the default scheme, a different quadrature rule entirely for these - degrees) isolates any discrepancy to the sum-factorized tabulation - itself, not to a difference in quadrature choice. - """ + """Check sum-factorized residuals and matrices against dense tabulation.""" mesh = {"triangle": UnitSquareMesh(2, 2), "tetrahedron": UnitCubeMesh(1, 1, 1)}[cell] - V = FunctionSpace(mesh, family, degree, variant="integral") + variant = None if family == "Bernstein" else "integral" + V = FunctionSpace(mesh, family, degree, variant=variant) u = TrialFunction(V) v = TestFunction(V) - w = Function(V) - w.dat.data[:] = np.random.default_rng(0).random(w.dat.data.shape) + rg = RandomGenerator(PCG64(seed=0)) + w = rg.uniform(V, 0, 1) # translate_coefficient path (forward transform): residual with a # derivative, mixing both the coefficient and argument sum-factorized diff --git a/tests/tsfc/test_codegen.py b/tests/tsfc/test_codegen.py index a62002fb32..7dd97e2ba1 100644 --- a/tests/tsfc/test_codegen.py +++ b/tests/tsfc/test_codegen.py @@ -75,81 +75,6 @@ def test_jagged_index_codegen(monkeypatch): assert numpy.allclose(u_out, u_ref, rtol=1e-14) -@pytest.mark.parametrize("element_name", ["Legendre", "IntegratedLegendre"]) -@pytest.mark.parametrize("cellname,degree", [("triangle", 3), ("tetrahedron", 2)]) -def test_duffy_scatter_and_contract(cellname, degree, element_name): - """Route B of the simplex sum-factorization milestone 2 design: - `finat.duffy.DuffyElement.basis_evaluation` must reproduce the standard - dense FIAT tabulation, via the dof numbering FIAT already uses, from - `duffy_evaluation`'s lattice-indexed, sum-factorized tabulation, and a - generic contraction of that same flat-dof-indexed result against a - coefficient vector (as `tsfc.fem.translate_coefficient` performs - uniformly for every FInAT element, `DuffyElement` included) must - reproduce the dense tensordot. `Legendre` (continuity=None) reads - exactly one lattice point per dof (lattice-lexicographic order); - `IntegratedLegendre` (continuity="C0") additionally exercises the - `FIAT.expansions.C0_basis` recombination, where each dof combines a - handful of lattice points. - - Verified via `gem.interpreter.evaluate` rather than a compiled loopy - kernel: the GEM expressions built here (in particular - `duffy_evaluation`'s `gem.VariableIndex`-based scatter index arithmetic) - schedule fine once merged into a real PyOP2 wrapper kernel (as confirmed - via `firedrake.assemble` on real forms), but are not guaranteed - schedulable by loopy in isolation -- scheduling an isolated, - unwrapped kernel is not a configuration real Firedrake usage ever - exercises. The GEM interpreter checks the same numerical correctness - without depending on loopy scheduling at all. - """ - from FIAT.reference_element import UFCTetrahedron, UFCTriangle - from finat.quadrature import make_quadrature - from finat.spectral import Legendre, IntegratedLegendre - from gem.gem import Index, Indexed, IndexSum, Product, Variable - from gem.interpreter import evaluate - from gem.optimise import remove_componenttensors - - cell = {"triangle": UFCTriangle, "tetrahedron": UFCTetrahedron}[cellname]() - element_cls = {"Legendre": Legendre, "IntegratedLegendre": IntegratedLegendre}[element_name] - element = element_cls(cell, degree) - ndof = element.space_dimension() - - quad_rule = make_quadrature(cell, 2 * degree, scheme="collapsed") - point_set = quad_rule.point_set - point_indices = point_set.indices - point_shape = tuple(index.extent for index in point_indices) - - entity = (cell.get_dimension(), 0) - dense_dict = element._element.tabulate(1, point_set.points) - - rng = numpy.random.default_rng(1) - coefficients = rng.random(ndof) - - # translate_argument path: dispatched transparently through basis_evaluation - scattered_dict = element.basis_evaluation(1, point_set, entity) - for alpha, dense in dense_dict.items(): - dense = dense.reshape((ndof,) + point_shape) - - r = Index(extent=ndof) - table, = remove_componenttensors([Indexed(scattered_dict[alpha], (r,))]) - u_out, = evaluate([table]) - assert u_out.fids == (r,) + point_indices - assert numpy.allclose(u_out.arr, dense, rtol=1e-12, atol=1e-12) - - # translate_coefficient path: generic contraction of the same - # flat-dof-indexed basis_evaluation result against a coefficient vector - # (no element-specific dispatch, see tsfc.fem.translate_coefficient) - c = Variable("c", (ndof,)) - r = Index(extent=ndof) - for alpha, dense in dense_dict.items(): - dense = dense.reshape((ndof,) + point_shape) - table, = remove_componenttensors([Indexed(scattered_dict[alpha], (r,))]) - value = IndexSum(Product(Indexed(c, (r,)), table), (r,)) - v_out, = evaluate([value], bindings={c: coefficients}) - assert v_out.fids == point_indices - v_ref = numpy.tensordot(coefficients, dense, axes=(0, 0)) - assert numpy.allclose(v_out.arr, v_ref, rtol=1e-12, atol=1e-12) - - if __name__ == "__main__": import os import sys diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index 949735480f..f9df3cd77d 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -211,7 +211,7 @@ def test_simplex_laplacian_action_compact_codegen(family): form = action(simplex_laplacian(triangle, family, 3)) kernel, = compile_form(form, parameters=dict(mode='spectral')) temporaries = kernel.ast.default_entrypoint.temporary_variables - assert len(temporaries) < 70 + assert len(temporaries) < 100 def test_bernstein_laplacian_action_compact_literals(): diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index 7cc649495f..1b771cc0cf 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -17,8 +17,9 @@ from finat.duffy import DuffyElement from finat.quadrature import AbstractQuadratureRule from gem.node import traversal -from gem.optimise import constant_fold_zero, unflatten_returns +from gem.optimise import constant_fold_zero from gem.optimise import remove_componenttensors as prune +from gem.unflatten import unflatten_returns from numpy import asarray from tsfc import fem from finat.element_factory import as_fiat_cell, create_element @@ -347,12 +348,10 @@ def set_quad_rule(params, cell, integral_type, functions): scheme = quad_rule fiat_cell = as_fiat_cell(cell) finat_elements = set(create_element(e) for e in elements if e.family() != "Real") - if scheme == "default" and any(isinstance(finat_el, DuffyElement) for finat_el in finat_elements): - # Sum-factorized (Duffy/lattice) tabulation only kicks in on a - # collapsed-coordinate quadrature rule (see - # finat.duffy.DuffyElement._duffy_applies); default to it - # automatically instead of requiring users to spell out - # dx(scheme="collapsed") to get it. + if (scheme == "default" and integral_type == "cell" + and any(isinstance(finat_el, DuffyElement) + for finat_el in finat_elements)): + # Duffy tabulation requires a collapsed-coordinate point set. scheme = "collapsed" fiat_cells = [fiat_cell] + [finat_el.complex for finat_el in finat_elements] if any(c.is_macrocell() for c in fiat_cells): diff --git a/tsfc/spectral.py b/tsfc/spectral.py index f66b759c9d..c5ba8ea30a 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -162,12 +162,9 @@ def prune(factors): variable = factors.pop() args = [f for f in factors if f != one] - assert set(var_indices) == set(variable.free_indices) - # Indices newly exposed in the variable by a VariableIndex delta - # cancellation (e.g. the nnz index of a `gem.SparseMatrix` row - # scatter, ``y[r] -> y[rows[p]]``) are scatter indices, not - # contraction indices: they must stay free in the expression to - # match the variable. + assert set(var_indices) <= set(variable.free_indices) + # A delta may replace a variable index by a contraction index. That + # index now describes a scatter in the assignment, not a sum. sum_indices = [i for i in sum_indices if i not in variable.free_indices] return variable, sum_indices, args, rest From 66e5646f5982be0b29035a4c4db5a052e96b5fb9 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 25 Jul 2026 00:07:47 +0100 Subject: [PATCH 16/23] Lower residual deltas in spectral kernels --- tsfc/spectral.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tsfc/spectral.py b/tsfc/spectral.py index c5ba8ea30a..8ad7f8c8ed 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -120,7 +120,7 @@ def group_key(pair): yield (variable, expression) -finalise_options = dict(replace_delta=False) +finalise_options = dict(replace_delta=True) def classify(argument_indices, expression, delta_inside): From 1a8c621387f4f06f8f09be75bc4f02adaef11861 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 27 Jul 2026 22:56:27 +0100 Subject: [PATCH 17/23] fix import --- tsfc/kernel_interface/common.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index 1b771cc0cf..2aeddc5dc4 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -17,9 +17,8 @@ from finat.duffy import DuffyElement from finat.quadrature import AbstractQuadratureRule from gem.node import traversal -from gem.optimise import constant_fold_zero +from gem.optimise import constant_fold_zero, unflatten_returns from gem.optimise import remove_componenttensors as prune -from gem.unflatten import unflatten_returns from numpy import asarray from tsfc import fem from finat.element_factory import as_fiat_cell, create_element From 74b6e710e281f8729c9dd4de2dddc7a51acb8597 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 28 Jul 2026 00:00:04 +0100 Subject: [PATCH 18/23] Test compact simplex bilinear kernels --- tests/firedrake/regression/test_quadrature.py | 8 ++++ tests/tsfc/test_sum_factorisation.py | 38 +++++++++++++++++-- tsfc/spectral.py | 32 ++-------------- 3 files changed, 45 insertions(+), 33 deletions(-) diff --git a/tests/firedrake/regression/test_quadrature.py b/tests/firedrake/regression/test_quadrature.py index 5c9f39115b..bb7189ddc4 100644 --- a/tests/firedrake/regression/test_quadrature.py +++ b/tests/firedrake/regression/test_quadrature.py @@ -83,3 +83,11 @@ def test_collapsed_quadrature_sum_factorisation(cell, degree, family): M = assemble(a).M.values M_collapsed = assemble(a_collapsed).M.values assert np.allclose(M, M_collapsed, rtol=1e-10, atol=1e-10) + + # Bilinear derivatives exercise two independently transformed argument + # lattices. + a = inner(grad(u), grad(v)) * dx(scheme="canonical") + a_collapsed = inner(grad(u), grad(v)) * dx(scheme="collapsed") + K = assemble(a).M.values + K_collapsed = assemble(a_collapsed).M.values + assert np.allclose(K, K_collapsed, rtol=1e-10, atol=1e-10) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index f9df3cd77d..89488908af 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -168,22 +168,22 @@ def test_vector_laplace_action(cell, order): assert (rates < order).all() -def simplex_mass(cell, family, degree): +def simplex_mass(cell, family, degree, scheme='collapsed'): m = Mesh(VectorElement('CG', cell, 1)) variant = None if family == "Bernstein" else "integral" V = FunctionSpace(m, FiniteElement(family, cell, degree, variant=variant)) u = TrialFunction(V) v = TestFunction(V) - return inner(u, v) * dx(scheme='collapsed') + return inner(u, v) * dx(scheme=scheme) -def simplex_laplacian(cell, family, degree): +def simplex_laplacian(cell, family, degree, scheme='collapsed'): m = Mesh(VectorElement('CG', cell, 1)) variant = None if family == "Bernstein" else "integral" V = FunctionSpace(m, FiniteElement(family, cell, degree, variant=variant)) u = TrialFunction(V) v = TestFunction(V) - return inner(grad(u), grad(v)) * dx(scheme='collapsed') + return inner(grad(u), grad(v)) * dx(scheme=scheme) @pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) @@ -227,6 +227,36 @@ def test_bernstein_laplacian_action_compact_literals(): assert sum(literal.size for literal in literals) < 10 * lattice_size +def test_bernstein_laplacian_bilinear_compact_codegen(): + import islpy as isl + import loopy as lp + + degree = 10 + collapsed = simplex_laplacian( + tetrahedron, "Bernstein", degree, scheme="collapsed") + canonical = simplex_laplacian( + tetrahedron, "Bernstein", degree, scheme="canonical") + collapsed_kernel, = compile_form( + collapsed, parameters=dict(mode="spectral")) + canonical_kernel, = compile_form( + canonical, parameters=dict(mode="spectral")) + + # At this degree the lower asymptotic complexity of sum factorisation + # outweighs its setup cost. + assert collapsed_kernel.flop_count < canonical_kernel.flop_count + + entrypoint = collapsed_kernel.ast.default_entrypoint + assert len(entrypoint.instructions) < 250 + assert len(entrypoint.temporary_variables) < 250 + code = lp.generate_code_v2(collapsed_kernel.ast).device_code() + assert sum(line.lstrip().startswith("for (") + for line in code.splitlines()) < 150 + + # A tetrahedral lattice has a loop whose bound depends on two parents. + assert max(domain.dim(isl.dim_type.param) + for domain in entrypoint.domains) >= 2 + + @pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"]) @pytest.mark.parametrize(('cell', 'order'), [(triangle, 5), (tetrahedron, 7)]) def test_simplex_mass_bilinear(cell, family, order): diff --git a/tsfc/spectral.py b/tsfc/spectral.py index 8ad7f8c8ed..6366f0a5cd 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -9,7 +9,7 @@ from gem.optimise import replace_division, unroll_indexsum from gem.refactorise import ATOMIC, COMPOUND, OTHER, MonomialSum, collect_monomials from gem.unconcatenate import unconcatenate -from gem.coffee import optimise_monomial_sum +from gem.coffee import sum_factorise_monomial_sum from gem.utils import groupby @@ -170,31 +170,5 @@ def prune(factors): def sum_factorise(variable, tail_ordering, monomial_sum): - if tail_ordering: - key_ordering = OrderedDict() - sub_monosums = defaultdict(MonomialSum) - for sum_indices, atomics, rest in monomial_sum: - # Pull out those sum indices that are not contained in the - # tail ordering, together with those atomics which do not - # share free indices with the tail ordering. - # - # Based on this, split the monomial sum, then recursively - # optimise each sub monomial sum with the first tail index - # removed. - tail_indices = tuple(i for i in sum_indices if i in tail_ordering) - tail_atomics = tuple(a for a in atomics - if set(tail_indices) & set(a.free_indices)) - head_indices = tuple(i for i in sum_indices if i not in tail_ordering) - head_atomics = tuple(a for a in atomics if a not in tail_atomics) - key = (head_indices, head_atomics) - key_ordering.setdefault(key) - sub_monosums[key].add(tail_indices, tail_atomics, rest) - sub_monosums = [(k, sub_monosums[k]) for k in key_ordering] - - monomial_sum = MonomialSum() - for (sum_indices, atomics), monosum in sub_monosums: - new_rest = sum_factorise(variable, tail_ordering[1:], monosum) - monomial_sum.add(sum_indices, atomics, new_rest) - - # Use COFFEE algorithm to optimise the monomial sum - return optimise_monomial_sum(monomial_sum, variable.index_ordering()) + return sum_factorise_monomial_sum( + monomial_sum, tuple(tail_ordering), variable.index_ordering()) From 9a7a9a4a8eff50dacca70b39d415337fcb19610a Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 30 Jul 2026 13:33:02 +0100 Subject: [PATCH 19/23] Add reproducible Johnson Mercier benchmark --- benchmarks/johnson_mercier.py | 42 +++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) create mode 100644 benchmarks/johnson_mercier.py diff --git a/benchmarks/johnson_mercier.py b/benchmarks/johnson_mercier.py new file mode 100644 index 0000000000..f30493d194 --- /dev/null +++ b/benchmarks/johnson_mercier.py @@ -0,0 +1,42 @@ +#!/usr/bin/env python +"""Measure simplex Johnson--Mercier code generation.""" + +import argparse +import hashlib + +from finat.ufl import FiniteElement, VectorElement +from tsfc import compile_form +from ufl import FunctionSpace, Mesh, TestFunction, TrialFunction, div, dx, inner +from ufl.cell import Cell + + +def compile_target(dim: int): + """Compile the JM mass-plus-divergence target on a simplex.""" + cell = Cell(("triangle", "tetrahedron")[dim - 2]) + mesh = Mesh(VectorElement("CG", cell, 1)) + element = FiniteElement("Johnson-Mercier", cell, 1) + space = FunctionSpace(mesh, element) + u = TrialFunction(space) + v = TestFunction(space) + form = (inner(u, v) + inner(div(u), div(v))) * dx + return compile_form(form, parameters={"mode": "spectral"})[0] + + +def main(): + """Print compiler metrics as copyable Markdown.""" + parser = argparse.ArgumentParser() + parser.add_argument("--dims", nargs="+", type=int, default=(2, 3)) + args = parser.parse_args() + print("") + print("| dim | flops | AST lines | AST bytes | hash |") + print("| ---: | ---: | ---: | ---: | --- |") + for dim in args.dims: + kernel = compile_target(dim) + source = str(kernel.ast) + digest = hashlib.sha256(source.encode()).hexdigest()[:12] + print(f"| {dim} | {kernel.flop_count:.0f} | " + f"{len(source.splitlines())} | {len(source)} | {digest} |") + + +if __name__ == "__main__": + main() From 814d976e0c60a6996ca52aacf1686ee41678d748 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 30 Jul 2026 15:23:49 +0100 Subject: [PATCH 20/23] WIP --- benchmarks/johnson_mercier.py | 20 ++++++++++++++--- tsfc/loopy.py | 41 +++++++++++++++++++++++++---------- 2 files changed, 47 insertions(+), 14 deletions(-) diff --git a/benchmarks/johnson_mercier.py b/benchmarks/johnson_mercier.py index f30493d194..cbe88783ca 100644 --- a/benchmarks/johnson_mercier.py +++ b/benchmarks/johnson_mercier.py @@ -4,6 +4,8 @@ import argparse import hashlib +import numpy + from finat.ufl import FiniteElement, VectorElement from tsfc import compile_form from ufl import FunctionSpace, Mesh, TestFunction, TrialFunction, div, dx, inner @@ -22,20 +24,32 @@ def compile_target(dim: int): return compile_form(form, parameters={"mode": "spectral"})[0] +def temporary_metrics(kernel): + """Return counts for the statically allocated Loopy temporaries.""" + temporaries = kernel.ast.default_entrypoint.temporary_variables.values() + sizes = [] + for temporary in temporaries: + if temporary.shape: + sizes.append(numpy.prod(temporary.shape, dtype=int)) + return len(sizes), sum(sizes), max(sizes, default=0) + + def main(): """Print compiler metrics as copyable Markdown.""" parser = argparse.ArgumentParser() parser.add_argument("--dims", nargs="+", type=int, default=(2, 3)) args = parser.parse_args() print("") - print("| dim | flops | AST lines | AST bytes | hash |") - print("| ---: | ---: | ---: | ---: | --- |") + print("| dim | flops | temporaries | elements | bytes | largest | AST lines | hash |") + print("| ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- |") for dim in args.dims: kernel = compile_target(dim) source = str(kernel.ast) digest = hashlib.sha256(source.encode()).hexdigest()[:12] + ntemp, nelem, largest = temporary_metrics(kernel) print(f"| {dim} | {kernel.flop_count:.0f} | " - f"{len(source.splitlines())} | {len(source)} | {digest} |") + f"{ntemp} | {nelem} | {8 * nelem} | {largest} | " + f"{len(source.splitlines())} | {digest} |") if __name__ == "__main__": diff --git a/tsfc/loopy.py b/tsfc/loopy.py index e41b8348fc..b9ae16b46b 100644 --- a/tsfc/loopy.py +++ b/tsfc/loopy.py @@ -124,6 +124,7 @@ def __init__(self, target=None): self.active_indices = {} # gem index -> pymbolic variable self.index_extent = OrderedDict() # pymbolic variable for indices -> extent self.index_parents = {} # iname -> parent inames bounding a jagged index + self.index_lengths = {} # iname -> (parent iname, row lengths) self.gem_to_pymbolic = {} # gem node -> pymbolic variable self.name_gen = UniqueNameGenerator() self.target = target @@ -258,7 +259,8 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name instructions, event_name, preamble = profile_insns(kernel_name, instructions, log) # Create domains - domains = create_domains(ctx.index_extent.items(), ctx.index_parents) + domains = create_domains( + ctx.index_extent.items(), ctx.index_parents, ctx.index_lengths) # Create loopy kernel knl = lp.make_kernel( @@ -277,23 +279,29 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name return knl, event_name -def create_domains(indices, index_parents=None): - """ Create ISL domains from indices - - :arg indices: iterable of (index_name, extent) pairs - :arg index_parents: optional mapping from index_name to a tuple of parent - index names; the domain of a jagged index is parametrized by its - parents, with upper bound extent minus the sum of the parents. - :returns: A list of ISL sets representing the iteration domain of the indices.""" - +def create_domains(indices, index_parents=None, index_lengths=None): + """Create ISL domains from rectangular and parent-bounded indices.""" domains = [] for idx, extent in indices: + if index_lengths and idx in index_lengths: + parent, lengths = index_lengths[idx] + inames = isl.make_zero_and_vars([idx], [parent]) + domain = None + for parent_value, length in enumerate(lengths): + piece = (inames[0].le_set(inames[idx]) + & inames[idx].lt_set(inames[0] + length) + & inames[parent].eq_set( + inames[0] + parent_value)) + domain = piece if domain is None else domain.union(piece) + domains.append(domain) + continue parents = index_parents.get(idx, ()) if index_parents else () inames = isl.make_zero_and_vars([idx], parents) bound = inames[0] + extent for parent in parents: bound = bound - inames[parent] - domains.append(((inames[0].le_set(inames[idx])) & (inames[idx].lt_set(bound)))) + domains.append(inames[0].le_set(inames[idx]) + & inames[idx].lt_set(bound)) if not domains: domains = [isl.BasicSet("[] -> {[]}")] @@ -331,6 +339,11 @@ def statement_for(tree, ctx): # remains correct: jagged expressions are zero-padded. ctx.index_parents[idx] = tuple(ctx.active_indices[parent].name for parent in tree.index.parents) + elif isinstance(tree.index, gem.RaggedIndex): + parent, = tree.index.parents + if parent in ctx.active_indices: + ctx.index_lengths[idx] = ( + ctx.active_indices[parent].name, tree.index.lengths) with active_indices({tree.index: p.Variable(idx)}, ctx) as ctx_active: return statement(tree.children[0], ctx_active) @@ -432,6 +445,12 @@ def _expression_failure(expr, ctx): raise expr.exception +@_expression.register(gem.FactorisationAtom) +def _expression_factorisation_atom(expr, ctx): + expression_, = expr.children + return expression(expression_, ctx) + + @_expression.register(gem.Product) def _expression_product(expr, ctx): return p.Product(tuple(expression(c, ctx) for c in expr.children)) From 0d67ba5cf3009678de01725251707ec2b9c32f82 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 30 Jul 2026 22:10:16 +0100 Subject: [PATCH 21/23] Preserve compact mapped tabulations in spectral mode --- tsfc/loopy.py | 6 ++++++ tsfc/spectral.py | 4 +++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/tsfc/loopy.py b/tsfc/loopy.py index b9ae16b46b..56e2da3040 100644 --- a/tsfc/loopy.py +++ b/tsfc/loopy.py @@ -461,6 +461,12 @@ def _expression_sum(expr, ctx): return p.Sum(tuple(expression(c, ctx) for c in expr.children)) +@_expression.register(gem.FactorisationAtom) +def _expression_factorisation_atom(expr, ctx): + expression_, = expr.children + return expression(expression_, ctx) + + @_expression.register(gem.Division) def _expression_division(expr, ctx): return p.Quotient(*(expression(c, ctx) for c in expr.children)) diff --git a/tsfc/spectral.py b/tsfc/spectral.py index 6366f0a5cd..aebac160c6 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -2,7 +2,7 @@ from functools import partial from itertools import chain, zip_longest -from gem.gem import Delta, Indexed, Sum, index_sum, one +from gem.gem import Delta, FactorisationAtom, Indexed, Sum, index_sum, one from gem.node import Memoizer, MemoizerArg from gem.optimise import filtered_replace_indices from gem.optimise import delta_elimination as _delta_elimination @@ -125,6 +125,8 @@ def group_key(pair): def classify(argument_indices, expression, delta_inside): """Classifier for argument factorisation""" + if isinstance(expression, FactorisationAtom): + return ATOMIC n = len(argument_indices.intersection(expression.free_indices)) if n == 0: return OTHER From f1f0253bdbd2f4c0367efcbf823ce8a3e2812420 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 30 Jul 2026 22:54:55 +0100 Subject: [PATCH 22/23] Remove experimental ragged loop domains --- tsfc/loopy.py | 41 +++++++++++------------------------------ 1 file changed, 11 insertions(+), 30 deletions(-) diff --git a/tsfc/loopy.py b/tsfc/loopy.py index 56e2da3040..23bd7e5eaa 100644 --- a/tsfc/loopy.py +++ b/tsfc/loopy.py @@ -124,7 +124,6 @@ def __init__(self, target=None): self.active_indices = {} # gem index -> pymbolic variable self.index_extent = OrderedDict() # pymbolic variable for indices -> extent self.index_parents = {} # iname -> parent inames bounding a jagged index - self.index_lengths = {} # iname -> (parent iname, row lengths) self.gem_to_pymbolic = {} # gem node -> pymbolic variable self.name_gen = UniqueNameGenerator() self.target = target @@ -259,8 +258,7 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name instructions, event_name, preamble = profile_insns(kernel_name, instructions, log) # Create domains - domains = create_domains( - ctx.index_extent.items(), ctx.index_parents, ctx.index_lengths) + domains = create_domains(ctx.index_extent.items(), ctx.index_parents) # Create loopy kernel knl = lp.make_kernel( @@ -279,29 +277,23 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name return knl, event_name -def create_domains(indices, index_parents=None, index_lengths=None): - """Create ISL domains from rectangular and parent-bounded indices.""" +def create_domains(indices, index_parents=None): + """ Create ISL domains from indices + + :arg indices: iterable of (index_name, extent) pairs + :arg index_parents: optional mapping from index_name to a tuple of parent + index names; the domain of a jagged index is parametrized by its + parents, with upper bound extent minus the sum of the parents. + :returns: A list of ISL sets representing the iteration domain of the indices.""" + domains = [] for idx, extent in indices: - if index_lengths and idx in index_lengths: - parent, lengths = index_lengths[idx] - inames = isl.make_zero_and_vars([idx], [parent]) - domain = None - for parent_value, length in enumerate(lengths): - piece = (inames[0].le_set(inames[idx]) - & inames[idx].lt_set(inames[0] + length) - & inames[parent].eq_set( - inames[0] + parent_value)) - domain = piece if domain is None else domain.union(piece) - domains.append(domain) - continue parents = index_parents.get(idx, ()) if index_parents else () inames = isl.make_zero_and_vars([idx], parents) bound = inames[0] + extent for parent in parents: bound = bound - inames[parent] - domains.append(inames[0].le_set(inames[idx]) - & inames[idx].lt_set(bound)) + domains.append(((inames[0].le_set(inames[idx])) & (inames[idx].lt_set(bound)))) if not domains: domains = [isl.BasicSet("[] -> {[]}")] @@ -339,11 +331,6 @@ def statement_for(tree, ctx): # remains correct: jagged expressions are zero-padded. ctx.index_parents[idx] = tuple(ctx.active_indices[parent].name for parent in tree.index.parents) - elif isinstance(tree.index, gem.RaggedIndex): - parent, = tree.index.parents - if parent in ctx.active_indices: - ctx.index_lengths[idx] = ( - ctx.active_indices[parent].name, tree.index.lengths) with active_indices({tree.index: p.Variable(idx)}, ctx) as ctx_active: return statement(tree.children[0], ctx_active) @@ -445,12 +432,6 @@ def _expression_failure(expr, ctx): raise expr.exception -@_expression.register(gem.FactorisationAtom) -def _expression_factorisation_atom(expr, ctx): - expression_, = expr.children - return expression(expression_, ctx) - - @_expression.register(gem.Product) def _expression_product(expr, ctx): return p.Product(tuple(expression(c, ctx) for c in expr.children)) From eeb157d013ba1f0ad767f9396b94bcc59a3c214c Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 30 Jul 2026 22:59:36 +0100 Subject: [PATCH 23/23] WIP --- tsfc/spectral.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tsfc/spectral.py b/tsfc/spectral.py index aebac160c6..9b06827b55 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -3,7 +3,7 @@ from itertools import chain, zip_longest from gem.gem import Delta, FactorisationAtom, Indexed, Sum, index_sum, one -from gem.node import Memoizer, MemoizerArg +from gem.node import Memoizer, MemoizerArg, traversal from gem.optimise import filtered_replace_indices from gem.optimise import delta_elimination as _delta_elimination from gem.optimise import replace_division, unroll_indexsum @@ -131,7 +131,11 @@ def classify(argument_indices, expression, delta_inside): if n == 0: return OTHER elif n == 1: - if isinstance(expression, (Delta, Indexed)) and not delta_inside(expression): + mapped = any(isinstance(node, FactorisationAtom) + and node.linear_closure + for node in traversal((expression,))) + if (mapped or isinstance(expression, (Delta, Indexed))) \ + and not delta_inside(expression): return ATOMIC else: return COMPOUND