Speed up entity_orientations by hoisting per-entity work out of the loop

[pbrubeck]

Description

AI-assisted (Claude Code, Opus 5)

Goal

#5231 took entity_orientations from 19.7 s to 2.1 s on BoxMesh(79, 79, 79) and left nogil/prange out of scope. This picks up the serial work that was still on the table: three pieces of per-entity work that do not depend on the entity, hoisted out of the loop over the cell closure. No OpenMP, no GIL change, no new build dependency.

entity_orientations runs once per mesh topology, before any solve, so it sits on the startup path of every script that builds a function space.

What's in it

• _compute_orientation_simplex() no longer allocates. It called PetscMalloc1/PetscFree twice per invocation, to hold cones of at most four points. On a tet mesh it runs for the 11 non-vertex entities of each cell's closure, so BoxMesh(79, 79, 79) made roughly 33M calls and 130M allocator round-trips. It now takes the two scratch arrays from the caller, which is exactly what its sibling _compute_orientation_interval_tensor_product() already did, and entity_orientations() allocates them once at maxConeSize alongside the three work arrays it already had.
• The polytope type is looked up once per entity, not once per (cell, entity). Every cell has the same FIAT cell, so FIAT entity e has the same type on every cell. entity_orientations() fills a small array indexed by e and passes the type in, dropping a DMPlexGetCellType() call per closure entry — about 45M of them on the same mesh. Vertices are then skipped without a call at all, since their orientation is always 0 and the output array starts zeroed.
• 2**dim is a shift. On a PetscInt, Cython compiles ** to a call to its integer power helper. The two occurrences in the tensor-product path become 1 << dim, which is exact for the dim values that reach them.

The loop order is unchanged. cell_closure is C-contiguous in e, so cell-outer/entity-inner is already the cache-friendly order; swapping it to hoist the type lookup differently would stride by numEntities.

Measurements

entity_orientations timed in isolation, minimum of 5 repeats, on one machine:

Mesh	Before	After
BoxMesh(79, 79, 79) (2,958,234 tets)	3.66 s	1.11 s
BoxMesh(40, 40, 40, hexahedral=True) (64,000 hexes)	0.282 s	0.164 s

Taken in order on the tet mesh: 3.66 s → 1.68 s for the allocation removal, → 1.11 s for the type cache. The absolute numbers are higher than #5231's because the machine is slower; the baseline above was measured on the same machine as the result.

Correctness

Output is bit-identical, not merely equivalent. np.array_equal against the values main produces, on eight meshes chosen to reach every branch of the dispatch:

UnitIntervalMesh, UnitSquareMesh (triangle and quadrilateral), BoxMesh (tetrahedral and hexahedral), UnitIcosahedralSphereMesh, UnitCubedSphereMesh, and the unstructured quadrilateral mesh from tests/firedrake/meshes. The last two carry the non-trivial quadrilateral orientations that a structured mesh does not produce.

Tests

tests/firedrake/regression/test_facet_orientation.py is the existing coverage for this code: it checks that cells sharing a node agree on it across a CG space, which is what these orientations feed through get_cell_nodes. It passes at 1 and 2 processes, over both simplex and tensor-product meshes including the unstructured quadrilateral one.

No test is added. The change alters no behaviour — the bit-identical comparison above is the claim, and a new test asserting orientation values would restate what test_facet_orientation.py already exercises end to end.

Notes for review

The one judgement call is the polytope-type cache: it assumes a single cell type across the mesh. entity_orientations() already rejects anything that is not exactly a MeshTopology, and it builds its entity-cone map from a single as_fiat_cell(mesh.ufl_cell()), so a per-cell type would already have broken the map. The cache does not weaken an assumption the function was not already making.

[JHopeCollins]

Does bitshifting here actually make a noticeable difference in the performance?
Unless it makes a very significant difference I don't think it's worth the loss in readability.

[pbrubeck]

I have made similar questions in the PR that proceeded this one #5231 and the answer for this one is yes

[pbrubeck]

This is an entity loop, ** calls a much heavier math routine

[JHopeCollins]

the answer for this one is yes

Do you have numbers?

[pbrubeck]

I lost the benchmark but I can regenerate it if you really insist. 2**N is ~150 times slower than 1<<N for small N <= 4. If you call it O(10^6) times the 2**N approach would take around 0.5 seconds, so it does dominate for this particular benchmark.