diff --git a/DESIGN.md b/DESIGN.md new file mode 100644 index 0000000000..5913158386 --- /dev/null +++ b/DESIGN.md @@ -0,0 +1,257 @@ +# 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 — 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 +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/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", 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 8d0bc79655..efc8ba6453 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,122 @@ 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) + + +@pytest.mark.parametrize("cellname,degree", [("triangle", 3), ("tetrahedron", 2)]) +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. + """ + from FIAT.reference_element import UFCTetrahedron, UFCTriangle + from finat.quadrature import make_quadrature + from finat.spectral import Legendre + 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) + 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: 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_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_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/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..8a6bad3e1a 100644 --- a/tsfc/fem.py +++ b/tsfc/fem.py @@ -11,9 +11,10 @@ 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 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 @@ -707,6 +708,45 @@ 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) @@ -745,45 +785,51 @@ 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_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) 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) 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)