From 9e1047300f9eb0e9607d884a38165d9e8099061e Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sun, 19 Jul 2026 00:33:48 +0100 Subject: [PATCH 1/6] 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 2/6] 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 3/6] 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 4/6] 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 5/6] 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 6/6] 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",