Add golden-master numeric test for HyperElasticity's cell kernels - #867
Draft
garth-wells wants to merge 28 commits into
Draft
Add golden-master numeric test for HyperElasticity's cell kernels#867garth-wells wants to merge 28 commits into
garth-wells wants to merge 28 commits into
Conversation
The fused quadrature loop interleaves evaluation of the varying quantities with the tensor contraction. The strided reads of the basis tables in the contraction stop GCC vectorising the loop, and because they share a loop body with the evaluation, math functions of the spatial coordinate are left scalar: a P2 load vector with a sin(pi x) sin(pi y) sin(pi z) source spends 64% of its time in scalar __sin_fma, where the P1 kernel of the same form calls libmvec's four-wide _ZGVdN4v_sin. Store the fw intermediates in an array indexed by quadrature point and contract in a second loop, leaving the evaluation loop free of the strided accesses. On a P2 Poisson load vector (tetrahedra, 1.57M cells) the kernel goes from 98.6 to 42 ns/cell and end-to-end RHS assembly in DOLFINx from 0.183 s to 0.098 s, with bit-identical output. P1 and matrix assembly are unchanged to within run-to-run noise. Applies to single-index quadrature rules with fw intermediates; other cases keep the fused loop. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Loop-invariant code motion in the tensor contraction hoisted each term of an entry separately, creating one temporary array and one accumulation into the element tensor per term. A 3D Poisson operator has one term per pair of derivative directions, so the inner loop carried nine temporaries and nine read-modify-writes of A where three suffice: terms that share the factor left in the inner loop differ only in what was hoisted, so one temporary can hold their sum. Group the terms of each entry by the factors that remain in the inner loop, sum the hoisted parts into a single temporary per group, and emit one accumulation per group. For P1 Poisson on tetrahedra the contraction goes from nine temporaries and nine statements to three of each. Measured on DOLFINx matrix re-assembly (unit cube, tetrahedra, serial), interleaved A/B against this branch point, Frobenius norm bit-identical in every run: P1, 6M cells default flags 2.84 s -> 1.79 s 1.6x P3, 1.57M cells default flags 179 s -> 65 s 2.8x The gain is at FFCx's default compile flags, where the generated code is built with the interpreter's CFLAGS (-O2). With -O3 -ffast-math -march=native the change is neutral (1.00-1.04x), as the compiler is then free to perform the same regrouping itself. The benefit grows with the number of degrees of freedom per cell. Note that grouping is a distributive regrouping, a*b + a*c -> a*(b+c), so results may differ in the last ulp; every case measured here was bit-identical. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
After grouping, each entry of the element tensor was still updated once per group: three read-modify-writes of the same address in the inner loop of a 3D Poisson operator. Sum the group contributions and assign once. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The tensor indices were reversed before building the loop nest, putting the row index of the element tensor in the inner loop. Each inner iteration then stepped over a whole row -- a stride of 4 for P1 tetrahedra and 20 for P3 -- which stops the accumulation vectorising and leaves the hoisted temporaries indexed by the inner loop. Nest in the natural order instead, so the last tensor index varies fastest. The element tensor is then walked contiguously and the temporaries are invariant in the inner loop. Measured on DOLFINx matrix re-assembly at FFCx's default compile flags (unit cube, tetrahedra, serial), interleaved A/B, together with the preceding single-accumulation commit, Frobenius norm bit-identical: P1, 6M cells 1.78 s -> 1.61 s 1.09x P3, 1.57M cells 66.3 s -> 44.1 s 1.50x Cumulative for the three contraction commits, same conditions: P1 2.84 s -> 1.61 s 1.8x P3 179 s -> 44.1 s 4.1x Neutral under -O3 -ffast-math -march=native, as with the grouping commit. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
#852 splits the quadrature loop into an evaluation pass and a contraction pass whenever any fw intermediate exists, so that vectorised math-library calls (sin, exp, ...) can escape the strided contraction loop. But an fw intermediate exists for almost every integral -- even a bare mass matrix has one (the quadrature weight times |detJ|) -- and when there is nothing to unblock, the split is pure overhead: an extra loop plus a store/load of every fw through a cache array. Benchmarked across ten forms on top of this branch (mass matrix, vector Poisson, elasticity, Taylor-Hood Stokes, N1curl mass, SIPG DG, HyperElasticity), the unconditional split regresses math-function-free forms by up to 11-13% at -O2 while leaving the forms #852 targets fully intact. Gate the split on the evaluation code actually containing a MathFunction call, matching the open question raised in #852's own description. Closes the regression on math-function-free forms (mass matrix: 0.89x -> 0.99x at -O2) at no cost to the cases #852 targets. `pytest test/` passes: 46 passed, 1 skipped (unchanged from this branch). AI assistance: I used Claude Code (Opus 5) to draft this change, building on #852. I reviewed and take responsibility for the final contribution.
Every affine cell's pseudo-inverse Jacobian divides each of its nine cofactors by the same determinant. FFCx's CSE already recognises the shared divisor but still emits nine separate hardware divisions; this was measured by hand as a real win in an earlier benchmarking pass but never actually wired into the compiler -- the redundant divisions live in generate_partition's `intermediates` list, which optimize() never touched (only the separate `definitions` list was optimised). Fixed by running a reciprocal-hoisting transform directly on `intermediates`, for both the piecewise (once-per-cell) and varying (per-quadrature- point) partitions. Separately, symmetric-gradient forms (elasticity_p1) showed a second, related waste: a value doubled in one place and halved in another, landing back on the original value, computed twice under different names because ordinary CSE only recognises syntactically identical expressions, not algebraic identities. Multiplying and dividing by an exact power of two never rounds in IEEE-754 -- (x + x) * 0.5 is bit-identical to x for any finite, non-subnormal x -- so this is always safe to fold, but GCC won't do it itself without -ffast-math. Added power_of_two_cse_statements, which tracks each declared scalar's value as (base symbol, power-of-two exponent) through chains of +self and multiply/divide-by-literal-power-of-two, and aliases a statement whose value is already known under an earlier name instead of recomputing it. Must run before reciprocal-CSE, which would otherwise hide the literal power-of-two divisors this looks for behind a reciprocal symbol. Verified: full test suite unchanged (46 passed, 1 skipped). Numerically cross-checked old vs. new generated code for elasticity_p1's bilinear kernel via ctypes on non-degenerate tetrahedron coordinates -- max absolute difference 2.22e-16 (a single extra rounding step from x * (1/d) vs x / d, the same order of difference TSFC's unconditional version of the reciprocal trick already produces). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…etry Avoids GCC auto-vectorizing a tiny fixed-trip-count reduction loop into an expensive permute-then-horizontal-reduce sequence. Falls back to the runtime loop when the coordinate element has a tensor-product factorisation (large dof counts).
licm() emitted one separate ForRange per hoisted temp, all sharing the same (index, begin, end), and never re-fused them. For a form with many block entries (elasticity_p1's 3x3 block structure, HyperElasticity's larger tensors) this produced dozens of independent tiny loops back to back -- each individually too small for the compiler to do much with, where a single loop computing all of them per iteration is a much better vectorisation target and has far less loop overhead. Also drops the zero-initialiser on the hoisted temp arrays, which is dead: every entry is unconditionally overwritten by the loop immediately below. Measured (GCC 14.2, -O3 -march=native -mprefer-vector-width=512 -ffp-contract=fast): elasticity_p1's bilinear kernel drops from 2097 to 1621 compiled instructions (-23%), 199.1ns -> 153.0ns per cell (-23%). HyperElasticity (6 kernels) drops 7.3% overall. No regressions on kernels this pass doesn't touch; kernels with only 1-3 hoisted temps (poisson_p1, dg_sipg, load_sin) see a small (3-7%) slowdown, likely a longer dependency chain in the single fused loop outweighing the reduced loop count when there's almost nothing to fuse.
Vector-valued forms with uncoupled, identical per-component blocks (e.g. vector Poisson's diagonal Laplacian blocks) produced the exact same hoisted expression once per component -- recomputed and stored from scratch each time instead of once and reused. Cache by the hoisted expression's structural key (new _key() helper, hashable identity for LNode expressions that aren't otherwise hashable). vector_poisson_p2: -8%, ratio 1.17x -> 1.07x. stokes_th: -7.5%, ratio 1.29x -> 1.19x. Small wins elsewhere (dg_sipg -2%, load_sin -6%, mass_p2/poisson_p1/poisson_p3 flat to slightly better). HyperElasticity +2.4% (noise-level; its blocks aren't identical, so no dedup fires -- this is a nonlinear form separately known to be noisy in this survey). 212/212 tests still pass.
For an affine cell, the per-block scale factor (a metric-tensor entry scaled by |det J|) is computed once per cell in the piecewise partition, then each entry gets multiplied by the quadrature weight separately inside the quadrature loop -- 2 multiplies per entry. Firedrake/TSFC instead fuses |det J| * weight into one scalar first, then multiplies each entry by that once -- for N entries and Q quadrature points, N + Q multiplies against FFCx's 2N. FFCx's strategy is the right one whenever Q > N (the common case for anything above the lowest quadrature orders, since the shared scale is reused across every point); it's specifically wrong for the low-point-count rules low-order affine geometry uses. Detects when every block entry going through generate_block_parts has a piecewise value that's a plain two-factor product sharing one factor in common (recorded when generate_partition builds it), and only when there are more sharing entries than quadrature points, computes the fused factor once per quadrature point and reuses it. Falls back to the existing per-entry multiply otherwise -- no risk to any case this doesn't fire for. The original, now sometimes-unused, piecewise scale-factor declarations are left in place and cleaned up generically: a small text-level pass in C/integral.py drops any scalar VariableDecl (always side-effect-free in this codegen) whose declared name is never read anywhere in the rendered function body, which is exactly what -Wunused-variable was catching in the JIT test suite once this fusion made some of them dead. Measured (GCC 14.2, -O3 -march=native -mprefer-vector-width=512 -ffp-contract=fast, both sides measured against Firedrake together): poisson_p1 194.1ns/47ns/1.22x -> 42.6ns/1.07x (-12%), essentially closing the remaining gap from PR #862+#863. elasticity_p1 and stokes_th (same pattern, but with a much larger tensor-accumulation phase dominating their runtime) move only slightly since the scale computation is a smaller fraction of their total work. No regressions on forms that don't trigger this path (mass_p2 unaffected, confirmed across repeated runs). 212/212 tests pass, including a numeric spot-check of the P1 stiffness matrix via ctypes.
power_of_two_cse_statements can reroute a statement straight back to an earlier value (e.g. collapsing a double-then-halve round trip to the original operand), skipping the statement it used to depend on. When nothing else needs that same intermediate value, it's left with no remaining reader -- dead code that GCC rejects under -Werror=unused-variable, which was failing every "Run FFCx demos" CI job for HyperElasticity (a symmetric-gradient form, exactly what this pass targets). Whether such a statement is genuinely dead depends on code outside its own partition's statement list (a different quadrature rule's varying partition consuming a piecewise value, or the tensor contraction reading it directly), so power_of_two_cse_statements no longer decides on its own: it records symbols it keeps only speculatively, and generate() prunes any that end up unreferenced once the whole kernel body -- definitions, intermediates, and tensor contraction together -- is assembled. Separately, reciprocal_cse_statements used the divisor LExpr directly as a dict key to count repeated divisions. That's fine for a Symbol divisor (the Jacobian-determinant case this was written for) but crashes for a literal divisor, since LiteralFloat defines __eq__ without __hash__ (unlike LiteralInt). This was crashing every demo compiled with `--scalar_type complex64/complex128` (SpatialCoordinates) in CI. Key on the divisor's structural identity (via the existing _key() helper, already used by licm()) instead of the LExpr object itself. Verified: full pytest suite passes (212 passed, 1 skipped -- unrelated, pygraphviz not installed). All 30 demos regenerated and recompiled with -std=c17 -Wunused-variable -Werror for every scalar type (216 passed, 16 skipped for unsupported type/demo combinations) -- previously 8 failing outright and the rest failing to compile. Confirmed both CSE passes still fire on affected forms (VectorPoisson's Jacobian reciprocal-CSE, HyperElasticity's power-of-two aliasing) after the fix. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Resolves conflicts with the licm() grouping/accumulate-once rewrite brought in by #853/#861: keeps the grouped, single-read-modify-write structure and layers #863's shared hoist loop and cross-group temp dedup on top of it, keyed on each group's already-summed hoisted expression rather than on a single ungrouped term. Also reconciles two independent dead-scalar cleanup passes that now coexist: prune_dead_scalars() (LNode-level, from #861's CSE work) and the textual _eliminate_dead_scalar_decls() pass in C/integral.py (from #863, needed for its piecewise scale-factor fusion). The new _piecewise_mul_operands bookkeeping in generate_piecewise_partition() now runs after the power-of-two/reciprocal CSE passes so it also recognises products that CSE rewrote from divisions.
definitions.py: the unrolled coordinate-dofs sum (#862) shadowed the outer-scope 'tables' name with a fresh annotated list, which mypy flags as a redefinition -- rename to 'coord_tables'. Also assert tabledata.offset is not None once, narrowing 'begin' from int | None for both the loop and unrolled code paths (the loop path's index expression built through the LExpr type happened not to trigger the same check). integral_generator.py: annotate 'decomps' (empty list gives mypy nothing to infer from) and check 'bases is not None' alongside 'shared is not None' before zipping, since narrowing 'shared' alone left 'bases' as list[LExpr] | None at the zip() call.
Investigated a ~2x mass-matrix regression found while benchmarking garth/perf-stack against main under -O3 -march=native -ffp-contract=fast (no fast-math). This zero-fill is value-dead (every element is written unconditionally by the loop below before any read), and restoring it does not fix the regression -- confirmed by A/B testing, the actual cause is PR #853's removal of integral_generator.py's B_indices reversal, which trades a large win on Poisson/elasticity/Stokes-shaped forms for a loss on plain single-block forms like a mass matrix, via GCC's vectoriser heuristics under strict FP. That trade-off is kept as-is (reverting it recovers the mass matrix but roughly halves the gains on 3 of the other 6 benchmarked forms) and documented on PR #865 rather than papered over here. Restoring the zero-fill has no measured effect either way; kept only so this line matches #853's original diff as closely as possible.
Bisected the ~2x mass_p2 regression found while benchmarking this branch against main to a single commit within PR #853 (52e7fa0, "Nest contraction loops with the contiguous tensor index innermost"): tested it in isolation with the other two #853 commits absent, and it alone reproduces the full regression. That commit's row-contiguous loop order is exactly what drives its own large wins on poisson_p3/vector_poisson_p2/stokes_th, so reverting it outright (tested) recovers mass_p2 but roughly halves those gains. Neither a restrict-qualified local row pointer nor explicit ivdep/unroll pragmas on the accumulation loop changed anything either (tested directly on the generated mass_p2.c). What actually distinguishes the losing case: mass_p2's entries each sum exactly one term (a plain product, no derivatives), while every form that wins big from the contiguous order sums multiple terms per entry (one per spatial direction, from a gradient-gradient or divergence pairing). Picking the loop order per block on that signal -- already available at this point via the existing 'keep' dict, no new bookkeeping -- gets both: mass_p2 recovers to parity with main (1.06x-1.07x instead of 0.54x-0.55x, strict FP and fast-math alike), and every other benchmarked form keeps its win. stokes_th improves further still (2.4x-3.6x to 2.4x-4.5x), since its velocity-pressure coupling blocks are themselves single-term and now pick up the same fix. Verified against all 7 forms benchmarked on PR #865, strict FP and with -ffast-math, 3 interleaved rounds each; full test suite still 212 passed, 1 skipped.
The demo's residual and Jacobian forms (P2 vector displacement, nonlinear strain energy, Jacobian built via ufl.derivative()) are the deepest-nested tensor composition anywhere in the test/demo corpus, and had zero numeric correctness coverage -- demo/test_demos.py only checks that the demo compiles, never asserts a result. Any future change to how nested tensor expressions get scalarized or factorized had nothing to catch a silent wrong-value regression against. Freezes the current compiler's output for both cell kernels against fixed, deterministic, non-degenerate coefficient/constant/coordinate data (stored in test/data/hyperelasticity_golden.npz alongside the inputs, so the fixture is self-contained and reproducible). Verified the check actually catches a regression by perturbing one frozen value and confirming the test fails with a clear mismatch report, then restoring it.
|
coverage: 85.484% (+0.7%) from 84.786% — garth/hyperelasticity-golden-master into main |
garth-wells
marked this pull request as draft
August 20, 2026 06:58
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
demo/HyperElasticity.py's residual (a_F) and Jacobian (a_J) forms (P2 vector displacement, nonlinear strain energy, Jacobian built viaufl.derivative()) are the deepest-nested tensor composition anywhere in the test/demo corpus, and had zero numeric correctness coverage --demo/test_demos.pyonly checks that the demo compiles, never asserts a result. Any change to how nested tensor expressions get scalarized or factorized inffcx/ir/analysishad nothing to catch a silent wrong-value regression against.Fix
Adds
test/test_hyperelasticity_golden.py, which JIT-compiles both forms, tabulates both cell kernels against fixed, deterministic, non-degenerate coefficient/constant/coordinate data, and compares to frozen values stored intest/data/hyperelasticity_golden.npz(which also stores the input data itself, so the fixture is fully self-contained).Verified the check actually catches a regression: perturbed one frozen value, confirmed the test fails with a clear mismatch report, then restored it.
Test plan
ruff check/ruff format --checkclean.🤖 Generated with Claude Code