Skip to content

Optimize tensor contractions and cell geometry evaluation - #919

Draft
garth-wells wants to merge 26 commits into
mainfrom
perf/elasticity-kernel-contractions
Draft

garth-wells wants to merge 26 commits into
mainfrom
perf/elasticity-kernel-contractions

Conversation

@garth-wells

@garth-wells garth-wells commented Oct 17, 2026 •

Copy link
Copy Markdown
Member

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 A are preserved.

The implementation:

  • Preintegrates linear and bilinear blocks with cellwise-constant scalar factors, merges reference tensors sharing a factor and output map, and flattens eligible contiguous contractions. Reference tables have a 64 KiB limit; partial preintegration retains quadrature for the remaining blocks.
  • Accumulates eligible point-varying component blocks in bounded local arrays with contiguous trial-index loops and zero-padded tables. One final pass interleaves component rows and adds them to A. Small or unsupported contractions retain their existing layout.
  • Proves block equality, transposed equality and symmetry from factor values and argument-table mappings. Equivalent blocks reuse an existing local array. Eligible symmetric diagonal blocks accumulate only on and above SIMD-sized diagonal groups, then mirror the omitted entries. Short unsplit blocks retain their update order.
  • Evaluates eligible coefficient and geometry expressions across quadrature points using point-major table copies and a bounded intermediate cache. Point padding is limited to 25%; contractions use only the original quadrature points. Cheap evaluations and unsupported cases remain fused.
  • Factors scalar arithmetic, shares reciprocals and geometry-weight scaling, collapses point-independent quadrature, and prunes unused scalars and tables. Coordinate reductions combine zero/sign-paired columns at tabulation round-off precision. Dense coefficient reductions of at most eight dofs are expanded selectively on varying geometry.
  • Explains these stages, local block layouts and padded loops in generated comments. Shared section descriptions cover C and Numba. Literal weighted factors on point facets and chained block equivalences have regression coverage.

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 -fPIC and -lm.

The table includes the latest release, v0.11.0, the earlier PR revision 30757e95 (before), and current kernel source 9e736b9f (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.

Kernel GCC v0.11.0 GCC before GCC PR vs before vs release Clang v0.11.0 Clang before Clang PR vs before vs release
Linear elasticity P2 2,111.92 1,034.39 914.31 1.13× 2.31× 3,703.05 1,316.90 900.40 1.46× 4.11×
Linear elasticity P3 15,507.68 3,421.01 3,208.52 1.07× 4.83× 29,473.26 3,777.13 3,374.54 1.12× 8.73×
Hyperelastic residual P2 460.71 161.67 163.69 0.99× 2.81× 293.74 176.96 176.75 1.00× 1.66×
Hyperelastic residual P3 1,207.84 485.26 489.34 0.99× 2.47× 1,246.50 551.81 547.18 1.01× 2.28×
Hyperelastic tangent P2 6,071.94 2,415.94 2,310.90 1.05× 2.63× 10,279.88 2,225.77 2,259.05 0.99× 4.55×
Hyperelastic tangent P3 33,008.15 11,151.38 9,481.76 1.18× 3.48× 65,465.76 11,946.45 10,672.64 1.12× 6.13×

Source, inputs, validation and the matching memory-traffic-only reference are described in benchmarks/streaming/roofline.py; samples are in benchmarks/streaming/roofline-results.json. The six-row release comparison maps to the streaming_release run in benchmarks/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 at 8f231a04 confirms 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 in benchmarks/refresh-current-results.json.

GCC 16

Kernel v0.11.0 (ns) main (ns) PR (ns) Speedup vs main Speedup vs release
Linear elasticity P1 291.54 139.31 90.33 1.54× 3.23×
Linear elasticity P2 3,097.67 1,123.46 399.50 2.81× 7.75×
Linear elasticity P3 25,916.30 8,537.26 1,526.78 5.59× 16.97×
Hyperelastic residual P1 79.94 67.21 32.98 2.04× 2.42×
Hyperelastic residual P2 645.68 673.68 375.84 1.79× 1.72×
Hyperelastic residual P3 1,946.54 2,068.38 1,290.59 1.60× 1.51×
Hyperelastic tangent P1 1,068.18 469.12 125.99 3.72× 8.48×
Hyperelastic tangent P2 10,515.09 4,513.10 2,338.08 1.93× 4.50×
Hyperelastic tangent P3 58,616.50 21,453.71 8,900.72 2.41× 6.59×

Clang 22

Kernel v0.11.0 (ns) main (ns) PR (ns) Speedup vs main Speedup vs release
Linear elasticity P1 207.43 128.36 75.05 1.71× 2.76×
Linear elasticity P2 3,264.38 2,266.34 359.40 6.31× 9.08×
Linear elasticity P3 26,619.58 8,411.60 1,657.13 5.08× 16.06×
Hyperelastic residual P1 60.78 52.80 32.66 1.62× 1.86×
Hyperelastic residual P2 495.01 498.93 352.08 1.42× 1.41×
Hyperelastic residual P3 1,430.43 1,441.66 1,033.66 1.39× 1.38×
Hyperelastic tangent P1 843.99 437.23 119.30 3.67× 7.07×
Hyperelastic tangent P2 9,478.50 7,583.66 2,198.74 3.45× 4.31×
Hyperelastic tangent P3 57,819.66 21,738.61 8,724.89 2.49× 6.63×

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.

Form GCC release GCC main GCC PR Clang release Clang main Clang PR
Tri P2 Poisson 35.67 37.28 8.95 58.79 26.98 9.02
Tri P3 Poisson 137.24 96.44 30.13 176.88 319.41 29.62
Tet P2 Poisson 233.51 111.12 60.47 203.73 245.85 60.14
Tet P3 Poisson 1,815.11 557.56 150.99 2,907.86 1,924.77 147.46

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.py and benchmarks/streaming/roofline.py. The current samples, environments, source hashes and run-to-table mappings are in benchmarks/refresh-current-results.json; benchmarks/refresh-results.json retains the earlier refresh. The other results files retain the earlier affine, non-affine, coefficient, fast-math and contraction investigations. Experimental variants in benchmarks/streaming/ remain separate from default code generation.

@garth-wells garth-wells changed the title Optimize scalar contractions and affine-cell tensor assembly Optimize tensor contractions and cell geometry evaluation Oct 17, 2026
@garth-wells
garth-wells marked this pull request as draft October 18, 2026 14:51
garth-wells and others added 11 commits October 18, 2026 14:57
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>
garth-wells and others added 5 commits October 18, 2026 18:57
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

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant