Repository navigation
Optimize tensor contractions and cell geometry evaluation - #919
Open
garth-wells wants to merge 53 commits into
Open
garth-wells wants to merge 53 commits into
garth-wells wants to merge 53 commits into
Conversation
garth-wells
changed the title
garth-wells
marked this pull request as draft
October 18, 2026 14:51
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
For point-varying contractions that update three or more component blocks of the element tensor, accumulate each block in a zero-initialised local array with the trial index innermost and contiguous, padding the loop to a 32-byte multiple with zero-padded tables. A single pass after the quadrature loop adds the blocks into A. This removes the strided, interleaved updates of A at every quadrature point and moves the arithmetic ahead of the first access to A. Blocks whose terms equal those of an earlier block, or its transpose with transposed entry maps, are not computed. Factors are compared through value ids that expand small point-varying polynomials, so equal values built in different orders are recognised. Symmetric tangents compute six of nine 3D blocks. In the split evaluation loop, read element tables through point-major copies so coefficient and geometry evaluation is contiguous across points, and run the loop to a SIMD multiple of points when this adds at most a quarter more points. Streaming benchmark (Xeon w7-3445, one core, 1 GiB of cell records, -O3 -march=native -ffast-math), time per cell, gcc-16 / clang-22: hyperelastic tangent P2 4.02 / 6.78 us -> 2.41 / 2.30 us hyperelastic tangent P3 20.9 / 20.2 us -> 11.3 / 12.0 us hyperelastic residual P2 324 / 256 ns -> 166 / 191 ns hyperelastic residual P3 985 / 941 ns -> 492 / 529 ns Simplex Poisson and linear elasticity kernels are unchanged. Across 173 forms on affine and non-affine cells, L1-resident timings improve by a geometric mean of 9% (gcc) and 5% (clang) on affine cells, and 22% and 16% on non-affine cells. Add benchmarks/streaming/roofline.py, which times kernels against a memory-traffic-only kernel giving the per-core DRAM bound. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Comment each local element tensor block with the entries of A it accumulates, and the entries its values or transpose also give when a block is reused by symmetry. Note padded entries that never reach A, padded quadrature points that are evaluated but not contracted, and the dimension order of point-major tables. Contraction comments no longer state that A is updated inside the quadrature loop, and localised contraction sections list the local blocks as their outputs. Generated code is otherwise unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Pre-integrated contractions of blocked (vector) arguments stored to A at a stride equal to the block size. Each test dof now evaluates its component block rows contiguously into a small local array, interleaves them in a row buffer and adds the row to A contiguously. When the summed tables prove the blocks symmetric, only the upper blocks are computed, from flattened table copies. Point-varying local blocks reuse the same row buffer scatter. Value keys now expand polynomials separately for piecewise and varying values, so equal blocks are recognised when their factors differ only in ordering. Forms with fewer than three blocks keep the strided update, which is cheaper there. Streaming linear elasticity on tetrahedra improves from 57-69% to 68-81% of the per-core DRAM bound, and hyperelastic P2/P3 kernels are about twice as fast. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
MSVC has no C99 _Complex, so mark the complex128 cases as expected failures there, as elsewhere in the test suite. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Re-time every variant of the description tables that report the current revision, using the stored sources and original method, so the tables reflect the latest contraction changes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A diagonal component block of a symmetric blocked tangent equals its own transpose, which the block equivalence proof now detects. Its local array is accumulated only from the SIMD column group holding the diagonal onwards, with rows taken in groups of one vector width so every loop keeps constant bounds and all blocks still share their table loads. The omitted entries are mirrored once after the quadrature loop. The split runs after loop-invariant hoisting and only for blocks with at least four row groups; smaller blocks do not recover the cost of the extra loops. Hyperelastic tangents on P3 tetrahedra are about 10% faster. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Blocks too small for row groups keep their original statement order, so their kernels are unchanged by the symmetric split. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A constant factor on a point facet folds with its unit weight to a literal, which has no symbol to record as a section input. In the jump of a vector argument on interior facets, a block can equal another block that is later found to be the transpose of a third; it now refers to the third block directly instead of to a block without a local array. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
garth-wells
marked this pull request as ready for review
October 25, 2026 19:55
Preserve optimized coefficient evaluation with the section-list API and retain invariant quadrature handling alongside sum-factorised contractions.
Share helpers for repeated code: literal substitution and loop unrolling (optimizer.substitute, unroll_loop), element tensor entries, SIMD width, column sign normalisation, point-independent tensor factors, unrolled coefficient sections and loop trip-count work estimates. Compute the blocked row layout and the loop-split work estimates once, merge the zero-padding of tables and weights, and build block entry comments from the block layout. Generated code is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Replace the SECTION_DESCRIPTIONS lookup by section name with an optional description on L.Section, set by the code that builds the section and kept when sections are rebuilt or fused. Formatters print it directly, so renaming a section no longer drops its comment. Generated code is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Move the sub-node walk to optimizer.children and build symbol reference collection, scalar use counting and the integral generator's tree queries on it, instead of separate per-node-type walkers. Loop bounds now count as reads, and uses inside array indices and conditionals are counted. Generated code is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
FFCXBackendAccess.table_access returns the value of a table constant over entities and points when it is read at a literal dof. The coefficient and coordinate evaluations and the unrolled block contraction then go through their general paths instead of reading table values themselves. Literal dof indices of the element tensor are folded, and unit factors now fold out of products. Machine code is unchanged under gcc and clang. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Register table symbols first, generate the quadrature loops, then declare each table once with the zero padding and point-major copies the loops recorded. This replaces patching declared table values in place and searching the code for where to insert the copies. Weights padding shares the table padding map. Generated code is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ry evaluation The sum-factorised evaluation of a linear combination of dofs tracks its stage shape, so directions can be contracted in any order. Point-independent directions, such as the derivative direction of a Q1 coordinate basis, go first: their differences of dofs are formed before interpolating the others. This replaces a fully unrolled special case limited to rules of at most three points, and also applies to higher-degree forms on Q1 geometry, halving their Jacobian stage arrays. Hot-loop timings of quadrilateral and hexahedron kernels at degrees 1 to 3 are unchanged or up to 14% faster with gcc and clang once code alignment is fixed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Keep the tensor computation description in a module constant beside the integrand factors one. Describe a rank-two block that is not a row and column layout by its flat entries, naming every argument index. Update docstrings for the point-independent-first contraction order and the stack budget estimate, which is an upper bound. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Share the evaluation order of a linear combination of dofs between the stack budget estimate and code generation: point-independent directions stored at one point and contracted first, then the others last direction first, with stages keyed by direction and table. The estimate previously counted every direction at every point, last direction first, so it overstated the evaluation arrays and left the contraction less of the stack limit than it could use. Fewer batches are now needed under the same limit; lower the limit of one batching test so it still splits into batches. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Mypy infers the dictionary keys from UFL's metaclass, which a type looked up from a terminal instance does not match. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
GCC and Clang completely unroll the short constant loops of the directional contractions before loop vectorisation, which left the stage updates of, for example, a four-point rule as scalar code. Loops whose length fills whole SIMD vectors and whose contracted direction has at least a vector's worth of points are now emitted as one plain loop marked to stay rolled (FFCX_KEEP_ROLLED, a GCC or Clang unroll pragma defined in the file preamble and empty for other compilers). Integrand caches and stage arrays also store the contiguous last direction padded to whole vectors when that costs at most a quarter more storage, e.g. seven points to eight; padding entries are zero and never reach A. Loops that do not fill whole vectors keep their previous form, so Q1, Q2 and Q4 kernels in the benchmark are unchanged. On distorted hexahedra with GLL Q3 elements, in-cache times improve by 1.3-2.2x (gcc) and 1.3-1.4x (clang): Laplace 8.7 to 4.9 us and elasticity 160 to 74 us with gcc. Quadrilateral Q3 kernels improve by 1.1-1.6x. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Without -ffast-math, a call such as log(J) keeps the whole evaluation loop scalar: errno must be preserved, and glibc only declares its SIMD variants for GCC under __FAST_MATH__. Each library call now gets loops of its own before the evaluation loop: the first computes its arguments at every point, storing the values it computes, and the second applies the function. The evaluation loop loads those values and the call results, so it vectorises under any flags; only the short call loop may stay scalar. The stored values share the fw cache stack budget. Hyperelastic residuals on tetrahedra built without -ffast-math are 2.2-2.5x faster (P3: 1261 to 499 ns with gcc) and now within 10-20% of a -ffast-math build, which also improves slightly. Loads with sin, exp and log improve by 1.2-1.6x on geometric mean without -ffast-math; with it they are 3-8% faster on average, with individual changes between -11% and +23%. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Clang can vectorise a loop across its inner loops, unrolling them. For a quadrature loop around contractions of 20 dofs this turns contiguous table reads and element tensor updates into gathers and scatters, and for test-function loops that update A at a stride along both loop indices, as for the velocity blocks of a Taylor-Hood form, it scatters across rows. Such loops are now marked FFCX_NO_VECTORIZE, a Clang pragma defined in the file preamble and empty for other compilers, so the inner loops vectorise instead. Quadrature loops are marked only when an inner loop has at least three SIMD vectors of iterations; across shorter inner loops Clang's choice is better. Loops meant to vectorise across points (split evaluation and library call loops) and sum-factorised kernels, where the effect was mixed, are left alone. With Clang and -ffast-math, a P3 tetrahedral load vector drops from 133 to 66 ns, level with GCC, and the Taylor-Hood Stokes form from 454 to 307 ns. Across 346 benchmark forms the geometric mean improves by 6-7%, with no form more than 10% slower. GCC output is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Two changes to sum-factorised kernels on quadrilaterals and hexahedra. Rows of the final-direction contraction that fill whole SIMD vectors now accumulate in a loop kept rolled when the rule has more points in that direction than a vector holds. Written out per entry over such a quadrature loop, the compiler kept the accumulators in memory as scalars, and the vector update of A then could not forward their stores; this took over 40% of the time of a GLL Q3 weighted Laplacian on hexahedra. When integrand factors are stored unpadded and the rule has more points than a vector, they are evaluated in one loop over all points rather than a nest over the directions, whose innermost loop of two to five points was too short to vectorise the geometry. The nest is kept if anything in the loop reads a single direction's point index. On distorted hexahedra, the weighted Laplacian is 1.6x faster at Q3 (28.3 to 16.3 us with gcc) and 1.2x at Q1 (494 to 394 ns); other Q1 and Q2 kernels gain up to 1.4x. A restructured contraction over all dof pairs of the last two directions was also tried and rejected: it loses the combination of terms sharing outer tables and was up to 2x slower at Q3. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Lagrange and DG elements on quadrilaterals and hexahedra created with basix.ufl.element number their dofs by entity rather than in tensor product order, so only elements from create_tp_element were sum factorised. Factorise these elements too, using the 1D factors of the tensor-product-ordered element of the same family and variants, and the element dof of each product of 1D basis functions, found by matching basis functions at points in the cell. Sum-factorised evaluation copies the dofs of such a coefficient or of the coordinates into tensor product order before the first stage, and the element tensor is updated through a table of the element dofs. Ordinary quadrilateral and hexahedral meshes now sum-factorise their geometry. The flat point loop test uses a quadratic mesh, whose derivatives vary in every direction. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
With few dofs, contracting one direction at a time saves little work and its loops cost more per operation, so low-order kernels such as a degree 1 Poisson operator on quadrilaterals were slower when sum factorised. Estimate the multiply-adds contracting the argument tables each way, and sum factorise a bilinear form only if it cuts them by at least sum_factorization_min_speedup (default 1.5). Linear forms are always sum factorised, as evaluating their coefficients one direction at a time pays even at low degree. Document that elements numbering their dofs by entity are sum factorised. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The evaluation loop of the test's small contraction is only split from the contraction on x86-64 Linux, where the compiler can call vector math functions, so the test failed on macOS. Set the host as the neighbouring platform test does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The option said whether to sum factorise, but its default now also weighs the cost. Make it a string: "auto", the default, uses sum factorisation where it is supported and, for a bilinear form, estimated to pay; "always" uses it wherever it is supported, whatever the estimated cost; and "never" never. Other values raise an error. Lower the default sum_factorization_min_speedup from 1.5 to 1.33. With the shortest exact rules, Laplace and elasticity forms of degree 2 on quadrilaterals are estimated at 1.38, and Clang's ordinary quadrature kernels for them are 4-5x slower. Degree 1 Poisson forms, estimated at 1.29, still use ordinary quadrature. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The derivative of a degree 1 basis is the same at every point in its own direction. A stage contracting over that direction summed the integrand against the weighted product of the argument tables once for each dof of every argument. Leave the constant tables out: sum against the weights and the varying tables once, for each dof of the arguments whose tables vary, before that direction's dof loops, and scale by the constant tables where the stage is read. A weighted Laplacian on degree 1 hexahedra with a three-point rule does 15% fewer floating point operations and takes 350 ns rather than 411 ns, and a Laplacian on degree 1 hexahedra 448 ns rather than 526 ns (GCC). The cost estimate choosing between sum factorisation and ordinary quadrature is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Member
|
This PR is very hard to beat (unlike the current |
This branch has not been deployed
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
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.
Generated kernels repeat basis contractions and geometry calculations, and blocked vector forms update interleaved output entries at every quadrature point. This change reduces that work and accumulates eligible component blocks locally before adding them to the element tensor. The existing kernel interface and accumulation into nonzero
Aare preserved.The implementation:
A. Small or unsupported contractions retain their existing layout.log) out of the quadrature evaluation loop into separate argument, call and result loops, so the remaining evaluation vectorises without-ffast-math.basix.ufl.elementon quadrilaterals and hexahedra, not only elements with tensor-product dof ordering. The 1D factors come from the tensor-product-ordered element of the same family and variants; a permutation, found by matching basis functions, maps each tensor-product index to the element's dof. The kernel gathers coefficient and coordinate dofs into tensor-product order and writesAthrough a static dof map, so its interface and the element's dof numbering are unchanged. Ordinary Q1 meshes now sum-factorise their geometry.sum_factorizationa choice of"auto"(the new default),"always"and"never". With"auto", a bilinear form uses ordinary quadrature when sum factorisation is estimated to cut the multiply-adds contracting the argument tables by less than the newsum_factorization_min_speedupoption (default 1.33). This keeps low-order forms with several terms, such as Q1 Poisson on quadrilaterals, on the faster ordinary path. Linear forms are always sum-factorised."always"sum-factorises wherever it is supported, whatever the estimate; other values, includingTrueandFalse, raise an error.A, to stay scalar. The hint is empty for other compilers.These transformations use compiler auto-vectorisation and reassociate floating-point arithmetic. Selection rules and storage budgets limit the extra local arrays, cached factors and loop padding.
Performance comparisons
All times below are nanoseconds per cell, including output zeroing; lower is better. Speedup is baseline time divided by PR time. Kernels and drivers are separate translation units without LTO. Measurements use logical CPU 2 of an Intel Xeon w7-3445, GCC 16.0.1 (experimental r16-8246) and Clang 22.1.2. Hot-cache and DRAM-streaming results are reported separately. All tables were re-timed on 31 October 2026, using PR revision
4a6aa405except the quadrilateral and hexahedral table, which uses719b29ea.The baselines are current main, feb3cc41 (the merge base, including sum factorisation), and the latest release, v0.11.0. PR and main use Basix 0.12.0.dev0 and UFL 2026.3.0.dev0; the release uses Basix 0.11.0 and UFL 2026.1.0. Each revision uses its own UFCx header. Every tensor entry on all 32 benchmark cells is checked across variants before timing.
DRAM-streaming improvement
Each cell owns a distinct output tensor, coordinates and coefficients in a 64-byte-aligned record. Each variant sweeps a 1 GiB allocation; seven interleaved samples are taken after initialisation and warm-up. Flags are
-O3 -march=native -ffp-contract=fast -ffast-math, with-shared -fPICand-lm. Quadrature, coordinates and coefficients are identical across variants.The table compares v0.11.0, an earlier PR revision
30757e95(before) and the current PR. The before revision already includes local block accumulation and point-major evaluation; “vs before” isolates the later row blocking, symmetric-diagonal split, library-call isolation and Clang loop hints.The v0.11.0 P3 elasticity kernel is very large, and its time varies by 20–40% between runs, so its “vs release” ratios are approximate. These are single-core local tensor assembly measurements; global assembly and socket-wide peak bandwidth are outside their scope.
Hot-cache elasticity and hyperelasticity versus main and release
These runs cycle through 32 affine tetrahedra and reuse a local output tensor. They use equispaced vector Lagrange P1–P3 elements, isotropic linear elasticity and compressible neo-Hookean residuals/tangents with mu = 1 and lambda = 1.25. All variants share the same prescribed quadrature. Flags are
-O3 -march=native -ffp-contract=fast, without-ffast-math. Each value is the minimum over two runs of nine approximately 0.1-second samples; variant order rotates between samples. Clang P3 times varied by up to 35% between runs with identical kernel sources.GCC 16
Clang 22
Hot-cache Poisson samples versus main and release
The same compiler flags and timing method apply, with one nine-sample run. These affine forms use each revision's default quadrature, so release comparisons include dependency and quadrature-selection differences.
Hot-cache quadrilateral and hexahedral samples versus main and release
The same method applies to equispaced Lagrange Q1–Q3 on affine quadrilaterals and Q1–Q2 on affine hexahedra (which FFCx does not treat as affine), with each revision's default quadrature. These elements are built with
basix.ufl.element, so main does not sum-factorise them. The PR sum-factorises all of them except the quadrilateral Q1 Poisson forms, for which the cost estimate selects ordinary quadrature. The weighted forms multiply by a coefficient in the same space.The PR is faster than main on every case here: 1.4–10.5×. Q2/Q3 kernels are 1.9–10.5× faster. Under Clang, the quadrilateral Q2/Q3 Poisson kernels are now 2.1–2.9× faster than v0.11.0, where main is 1.6–2.5× slower than it.
The cost estimate does not always choose the faster path. GCC hex Q1 Poisson is 0.80× the speed of ordinary quadrature when sum-factorised (449 ns against 358 ns), while under Clang it is 2.3× faster, so it stays sum-factorised. Forcing either path for every form (75 quadrilateral and hexahedral forms, both compilers) showed bilinear forms below the 1.33 threshold were 13–29% slower when sum-factorised, and nearly all above it were faster; the exception nearest the threshold, quadrilateral Q1 reaction–diffusion (estimate 1.46), is 4–18% slower. With the shortest exact quadrature rules, Q1 Laplace on hexahedra and Q1 elasticity on quadrilaterals are estimated below the threshold but are 1.1–2.7× faster sum-factorised, while Q1 elasticity on hexahedra, with a similar estimate, is 2.0–2.5× faster with ordinary quadrature. Linear forms are neutral on average when sum-factorised, with losses of up to 30% on a few small kernels, and are up to 6.9× faster with a coefficient to evaluate.
Validation
Benchmark checks compare every tensor entry on all 32 cells against independent quadrature or the other variants. Tests cover nonzero output accumulation, signed custom quadrature, rectangular and blocked tensors, real/complex single/double precision, curved simplex and warped tensor-product geometry, H(div)/H(curl) maps, literal point-facet factors, interior-facet block-equivalence chains, rolled sum-factorised loops, flattened point loops, sum factorisation of entity-ordered elements at degrees 1–3 and the cost-based choice of path. Hyperelasticity tests include the undeformed linear limit, tangent symmetry and finite differences.
The results are specific to the tested cases, compilers and flags; they do not establish a universal no-regression guarantee.
🤖 Generated with Claude Code