Repository navigation
Optimize tensor contractions and cell geometry evaluation - #919
Draft
garth-wells wants to merge 26 commits into
Draft
garth-wells wants to merge 26 commits into
garth-wells wants to merge 26 commits into
Conversation
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>
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.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 optimised 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.
The main and release baselines are pinned to ff8ad409 and v0.11.0. PR/main use Basix 0.12.0.dev0 and UFL 2026.2.0.dev0; the release uses Basix 0.11.0 and UFL 2026.1.0. Each revision uses its own UFCx header.
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.The table includes the latest release, v0.11.0, the earlier PR revision
30757e95(before), and current kernel source9e736b9f(PR). The before revision already includes local block accumulation and point-major evaluation; “vs before” isolates the subsequent row-blocking, symmetric-diagonal split and correctness fixes. “vs release” compares the full PR with v0.11.0. Both ratios are speedups: baseline time divided by PR time.All six displayed P2/P3 tetrahedral cases were re-timed together for release, before and PR with both compilers. Quadrature, coordinates and coefficients match exactly across these variants; the release uses its own Basix 0.11.0, UFL 2026.1.0 and UFCx header. Earlier measurements of all 18 simplex Poisson, elasticity and hyperelastic cases remain in the results file, including Poisson and P1 controls.
Source, inputs, validation and the matching memory-traffic-only reference are described in
benchmarks/streaming/roofline.py; samples are inbenchmarks/streaming/roofline-results.json. The six-row release comparison maps to thestreaming_releaserun inbenchmarks/refresh-current-results.json. 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 reported value is the minimum of nine approximately 0.1-second samples; variant order rotates between samples. All variants in each table were re-timed together, with one nine-sample run per compiler.All four performance tables were refreshed on 20 October 2026 using kernel source
9e736b9f. Regeneration after the main merge at8f231a04confirms identical C kernel bodies for all measured form families. Every displayed column, including main, release and the earlier streaming revision, was re-timed. The main and release generated sources and dependency versions remain pinned as specified above; the streaming before/after sources use matching dependencies. Source revisions, full samples, errors and source hashes are recorded inbenchmarks/refresh-current-results.json.GCC 16
Clang 22
Hot-cache Poisson samples versus main and release
The same compiler flags and timing method apply. These affine forms retain each revision's default quadrature, so release comparisons include dependency and quadrature-selection differences.
Validation and reproducibility
Benchmark checks compare every tensor entry on all 32 cells against independent quadrature or reference kernels. 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 and interior-facet block-equivalence chains. Hyperelasticity tests include the undeformed linear limit, tangent symmetry and finite differences.
The wider form suites cover loads, mass, diffusion, weighted and anisotropic operators, transport, reaction–diffusion, vector mass, mixed divergence, H(div) and H(curl) on simplex and tensor-product cells. Their recorded screening results are specific to the tested cases and compiler flags; they do not establish a universal no-regression guarantee.
Reproduction drivers are
benchmarks/elasticity.py,benchmarks/forms.pyandbenchmarks/streaming/roofline.py. The current samples, environments, source hashes and run-to-table mappings are inbenchmarks/refresh-current-results.json;benchmarks/refresh-results.jsonretains the earlier refresh. The other results files retain the earlier affine, non-affine, coefficient, fast-math and contraction investigations. Experimental variants inbenchmarks/streaming/remain separate from default code generation.