From 42fc5ba2041e1ad8ae1dee135160d31a3987b125 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 19 Aug 2026 09:58:37 +0100 Subject: [PATCH 1/7] Padded basis transformation for better codegen --- finat/physically_mapped.py | 100 ++++++++++++++++++++++---------- gem/interpreter.py | 34 +++++++++-- gem/optimise.py | 12 +++- test/finat/test_zany_mapping.py | 35 ++++++++++- test/gem/test_sum_factorise.py | 17 ++++++ 5 files changed, 160 insertions(+), 38 deletions(-) diff --git a/finat/physically_mapped.py b/finat/physically_mapped.py index 511d3e146..6a8d6d3ff 100644 --- a/finat/physically_mapped.py +++ b/finat/physically_mapped.py @@ -1,5 +1,5 @@ from abc import ABCMeta, abstractmethod -from collections.abc import Mapping +from collections.abc import Iterable, Mapping import gem import numpy @@ -7,6 +7,10 @@ from finat.citations import cite +zero = gem.Zero() +one = gem.Literal(1.0) + + class NeedsCoordinateMappingElement(metaclass=ABCMeta): """Abstract class for elements that require physical information either to map or construct their basis functions.""" @@ -16,38 +20,78 @@ def dual_transformation(self, Q, coordinate_mapping=None): class MappedTabulation(Mapping): - """A lazy tabulation dict that applies the basis transformation only - on the requested derivatives. + """Apply a sparse basis transformation to reference tabulations. + + Parameters + ---------- + M : gem.ListTensor + Basis-transformation matrix. + ref_tabulation : Mapping + Reference tabulations indexed by derivative order. + indices : iterable of int, optional + Rows retained by an element restriction. + + Notes + ----- + In order to generate good loopy kernels, rows are padded so that they have + the same number of entries. Constant tables select the reference column + and one of the distinct symbolic coefficients. Interning coefficients + preserves their sharing without materialising a symbolic matrix entry by + entry. - :arg M: a gem.ListTensor with the basis transformation matrix. - :arg ref_tabulation: a dict of tabulations on the reference cell. - :kwarg indices: an optional list of restriction indices on the basis functions. """ - def __init__(self, M, ref_tabulation, indices=None): - self.M = M + + def __init__( + self, M: gem.ListTensor, ref_tabulation: Mapping, + indices: Iterable[int] | None = None) -> None: self.ref_tabulation = ref_tabulation if indices is None: - indices = list(range(M.shape[0])) - self.indices = indices - # we expect M to be sparse with O(1) nonzeros per row - # for each row, get the column index of each nonzero entry - csr = [[j for j in range(M.shape[1]) if not isinstance(M.array[i, j], gem.Zero)] - for i in indices] - self.csr = csr + indices = range(M.shape[0]) + self.indices = tuple(indices) + self._space_dim = len(self.indices) + + nonzero_rows = [] + for source_row in self.indices: + row = [] + for column in range(M.shape[1]): + value = M.array[source_row, column] + if not isinstance(value, gem.Zero): + row.append((column, value)) + nonzero_rows.append(row) + width = max((len(row) for row in nonzero_rows), default=0) + nrows = len(self.indices) + columns = numpy.zeros((nrows, width), dtype=gem.uint_type) + data = numpy.full((nrows, width), zero, dtype=object) + for index, row in enumerate(nonzero_rows): + columns[index, :len(row)] = tuple(column for column, _ in row) + data[index, :len(row)] = tuple(gem.as_gem(value) for _, value in row) + self._width = width + self._columns = gem.Literal(columns, dtype=gem.uint_type) + values = [] + value_numbers = {} + value_indices = numpy.empty(data.shape, dtype=gem.uint_type) + for multiindex, value in numpy.ndenumerate(data): + try: + number = value_numbers[value] + except KeyError: + number = len(values) + value_numbers[value] = number + values.append(value) + value_indices[multiindex] = number + self._value_indices = gem.Literal(value_indices, dtype=gem.uint_type) + self._values = gem.ListTensor(values) self._tabulation_cache = {} - def matvec(self, table): - # basis recombination using hand-rolled sparse-dense matrix multiplication - ii = gem.indices(len(table.shape)-1) - phi = [gem.Indexed(table, (j, *ii)) for j in range(self.M.shape[1])] - # the sum approach is faster than calling numpy.dot or gem.IndexSum - exprs = [gem.ComponentTensor(gem.Sum(*(self.M.array[i, j] * phi[j] for j in js)), ii) - for i, js in zip(self.indices, self.csr)] + def matvec(self, table: gem.Node) -> gem.Node: + r = gem.Index(extent=self._space_dim) + k = gem.Index(extent=self._width) + i = gem.VariableIndex(gem.Indexed(self._value_indices, (r, k))) + j = gem.VariableIndex(gem.Indexed(self._columns, (r, k))) + A = gem.Indexed(self._values, (i,)) - result = gem.ListTensor(exprs) - result, = gem.optimise.unroll_indexsum((result,), lambda index: True) - # result = gem.optimise.aggressive_unroll(self.M @ table) - return result + tail = gem.indices(len(table.shape) - 1) + mapped = gem.IndexSum(gem.Product(A, gem.Indexed(table, (j, *tail))), (k,)) + return gem.ComponentTensor(mapped, (r, *tail)) def __getitem__(self, alpha): try: @@ -195,10 +239,6 @@ def physical_vertices(self): (gdim, ).""" -zero = gem.Zero() -one = gem.Literal(1.0) - - def identity(*shape): V = numpy.eye(*shape, dtype=object) for multiindex in numpy.ndindex(V.shape): diff --git a/gem/interpreter.py b/gem/interpreter.py index 13eeb44a2..b2dd609e0 100644 --- a/gem/interpreter.py +++ b/gem/interpreter.py @@ -263,8 +263,34 @@ def _evaluate_conditional(e, self): def _evaluate_indexed(e, self): """Indexing maps shape to free indices""" val = self(e.children[0]) - fids = tuple(i for i in e.multiindex if isinstance(i, gem.Index)) + variable_indices = {i: self(i.expression) for i in e.multiindex + if isinstance(i, gem.VariableIndex)} + + if any(result.fids for result in variable_indices.values()): + # Some variable index depends on free indices: gather entries + # one by one over the extent of the free indices. + fids = list(val.fids) + for i in e.multiindex: + new_fids = (i,) if isinstance(i, gem.Index) else \ + variable_indices[i].fids if isinstance(i, gem.VariableIndex) else () + fids.extend(f for f in new_fids if f not in fids) + fids = tuple(fids) + out = numpy.empty(tuple(f.extent for f in fids), dtype=val.arr.dtype) + for idx in numpy.ndindex(out.shape): + env = dict(zip(fids, idx)) + vidx = [env[f] for f in val.fids] + for i in e.multiindex: + if isinstance(i, gem.Index): + vidx.append(env[i]) + elif isinstance(i, gem.VariableIndex): + result = variable_indices[i] + vidx.append(int(result.arr[tuple(env[f] for f in result.fids)])) + else: + vidx.append(i) + out[idx] = val.arr[tuple(vidx)] + return Result(out, fids) + fids = tuple(i for i in e.multiindex if isinstance(i, gem.Index)) idx = [] # First pick up all the existing free indices for _ in val.fids: @@ -275,10 +301,10 @@ def _evaluate_indexed(e, self): # Free index, want entire extent idx.append(slice(None)) elif isinstance(i, gem.VariableIndex): - # Variable index, evaluate inner expression - result, = self(i.expression) + # Variable index, constant during kernel execution + result = variable_indices[i] assert not result.tshape - idx.append(result[()]) + idx.append(int(result.arr[()])) else: # Fixed index, just pick that value idx.append(i) diff --git a/gem/optimise.py b/gem/optimise.py index 052314162..67a526234 100644 --- a/gem/optimise.py +++ b/gem/optimise.py @@ -530,6 +530,10 @@ def _sum_factorise_connected(sum_indices, groups): :arg groups: product factors, grouped by free indices :returns: optimised GEM expression """ + if not groups: + extent = numpy.prod([index.extent for index in sum_indices], dtype=int) + return Literal(float(extent)) + if len(groups) <= _MAX_PLANNED_FACTORS: return _plan_contraction(sum_indices, groups) @@ -577,9 +581,11 @@ def sum_factorise(sum_indices, factors): :arg factors: product factors :returns: optimised GEM expression """ - if len(factors) == 0 and len(sum_indices) == 0: - # Empty product - return one + if len(factors) == 0: + # The empty product is one, so contracting it counts the tuples in the + # index space. + extent = numpy.prod([index.extent for index in sum_indices], dtype=int) + return Literal(float(extent)) # Form groups by free indices groups = groupby(factors, key=lambda f: f.free_indices) diff --git a/test/finat/test_zany_mapping.py b/test/finat/test_zany_mapping.py index 9220dca99..59f1e8883 100644 --- a/test/finat/test_zany_mapping.py +++ b/test/finat/test_zany_mapping.py @@ -1,11 +1,44 @@ import FIAT import finat +import gem import numpy as np import pytest import pprint from gem.interpreter import evaluate -from finat.physically_mapped import PhysicallyMappedElement +from gem.node import traversal +from finat.physically_mapped import MappedTabulation, PhysicallyMappedElement + + +def test_sparse_mapped_tabulation(): + """Apply a sparse basis map at the cost of its nonzeros.""" + coefficient = gem.Variable("coefficient", ()) + matrix = gem.ListTensor(np.asarray([ + [gem.Literal(1.0), gem.Zero(), coefficient], + [gem.Zero(), gem.Literal(1.0), gem.Zero()], + ], dtype=object)) + table_values = np.arange(1.0, 7.0).reshape(3, 2) + table = gem.Literal(table_values) + + mapped_tabulation = MappedTabulation(matrix, {None: table}) + # Equal symbolic entries share one coefficient slot, so the sparse map + # indexes a vector rather than materialising a symbolic matrix. + assert mapped_tabulation._values.shape == (3,) + + mapped = mapped_tabulation[None] + + # The three unit entries cost no multiplication, and the one remaining + # nonzero costs exactly one. Nothing is selected by a branch. + products = [node for node in traversal((mapped,)) + if isinstance(node, gem.Product)] + assert len(products) == 1 + assert not any(isinstance(node, gem.Conditional) + for node in traversal((mapped,))) + + actual, = evaluate([mapped], {coefficient: np.asarray(2.0)}) + expected = np.asarray([[1.0, 0.0, 2.0], [0.0, 1.0, 0.0]]) \ + @ table_values + assert np.array_equal(actual.arr, expected) def make_unisolvent_points(element, interior=False): diff --git a/test/gem/test_sum_factorise.py b/test/gem/test_sum_factorise.py index c938aa4c7..5dca72b93 100644 --- a/test/gem/test_sum_factorise.py +++ b/test/gem/test_sum_factorise.py @@ -255,3 +255,20 @@ def test_estimate_cost_counts_the_contraction(): assert flops > 0 assert storage >= largest > 0 assert nodes > 0 + + +def test_empty_product_contraction_counts_index_tuples() -> None: + i, j = gem.Index(extent=2), gem.Index(extent=3) + expression = sum_factorise((i, j), ()) + result, = evaluate([expression]) + assert result.arr == 6 + + +def test_contraction_counts_index_absent_from_factors() -> None: + i, j = gem.Index(extent=2), gem.Index(extent=3) + factor = gem.Indexed(gem.Literal(numpy.arange(3.0)), (j,)) + + expression = sum_factorise((i, j), (factor,)) + result, = evaluate([expression]) + + assert result.arr == 2 * numpy.arange(3.0).sum() From ba8be9b4c16df1bf6b8f2c2931f97f2565d86d60 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 22 Aug 2026 18:18:32 +0100 Subject: [PATCH 2/7] Select facet tabulations through indirect reductions A padded basis transformation tabulates each facet as an IndexSum, so selecting one by a variable facet index reached _select_expression with a type it could not factorise. Rewrite the summands over one shared multiindex and select inside the reduction, which the equal extents on every facet make well defined. Co-Authored-By: Claude Opus 5 --- gem/optimise.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/gem/optimise.py b/gem/optimise.py index 67a526234..76b2a163c 100644 --- a/gem/optimise.py +++ b/gem/optimise.py @@ -270,6 +270,15 @@ def child(expression): elif all(e.j == k and e.i == expr.i for k, e in enumerate(expressions)): return expr.reconstruct(expr.i, index) + if types == {IndexSum}: + extents = {tuple(i.extent for i in e.multiindex) for e in expressions} + if len(extents) == 1: + multiindex = tuple(Index(extent=extent) for extent in extents.pop()) + summands = [Indexed(ComponentTensor(e.children[0], e.multiindex), + multiindex) + for e in expressions] + return IndexSum(_select_expression(summands, index), multiindex) + if len(types) == 1: cls, = types if cls.__front__ or cls.__back__: From edc97d94a0672eb46fbbee5364661867679216d1 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sun, 23 Aug 2026 11:34:14 +0100 Subject: [PATCH 3/7] Express the basis transformation as a GEM contraction The padded transformation was applied by building its row-padded gather directly, which fixed the orientation at construction: contracting the coefficient against a mapped tabulation then costs one gather per quadrature point, six times the symmetric test-side scatter. Represent M instead as a rank-2 expression, an interned entry summed over the padded row against a Delta selecting its column. Cancelling that Delta reproduces the gather, so the mat-mat is unchanged, while contracting M's own axes first pulls a coefficient back to the reference basis once per cell. Guzman-Neilan 3D action: 147180 -> 51920 flops, largest working temporary 144 -> 24 entries. Delta now propagates the free indices of a VariableIndex operand, and substitution folds a variable index that has become constant. Co-Authored-By: Claude Opus 5 --- finat/physically_mapped.py | 44 +++++++-- gem/gem.py | 17 +++- gem/optimise.py | 191 ++++++++++++++++++++++++++++++++++++- 3 files changed, 239 insertions(+), 13 deletions(-) diff --git a/finat/physically_mapped.py b/finat/physically_mapped.py index 6a8d6d3ff..315bfd055 100644 --- a/finat/physically_mapped.py +++ b/finat/physically_mapped.py @@ -49,6 +49,7 @@ def __init__( indices = range(M.shape[0]) self.indices = tuple(indices) self._space_dim = len(self.indices) + self._value_dim = M.shape[1] nonzero_rows = [] for source_row in self.indices: @@ -82,22 +83,51 @@ def __init__( self._values = gem.ListTensor(values) self._tabulation_cache = {} - def matvec(self, table: gem.Node) -> gem.Node: - r = gem.Index(extent=self._space_dim) + def _entry(self, r: gem.Index, a: gem.Index) -> gem.Node: + """Entry ``M[r, a]`` of the basis transformation. + + Parameters + ---------- + r + Index over the rows retained by the element. + a + Index over the reference basis. + + Returns + ------- + gem.Node + A sum over the padded row of an interned entry against a Delta + selecting its column, so that contracting either axis of ``M`` + is ordinary GEM algebra. + + """ k = gem.Index(extent=self._width) - i = gem.VariableIndex(gem.Indexed(self._value_indices, (r, k))) - j = gem.VariableIndex(gem.Indexed(self._columns, (r, k))) - A = gem.Indexed(self._values, (i,)) + entry = gem.Indexed( + self._values, + (gem.VariableIndex(gem.Indexed(self._value_indices, (r, k))),)) + column = gem.VariableIndex(gem.Indexed(self._columns, (r, k))) + return gem.IndexSum(gem.Product(entry, gem.Delta(column, a)), (k,)) + + def matrix(self) -> gem.Node: + """The basis transformation as a rank-2 GEM expression.""" + r = gem.Index(extent=self._space_dim) + a = gem.Index(extent=self._value_dim) + return gem.ComponentTensor(self._entry(r, a), (r, a)) + def matmul(self, table: gem.Node) -> gem.Node: + """Apply the basis transformation to a reference tabulation.""" + r = gem.Index(extent=self._space_dim) + a = gem.Index(extent=self._value_dim) tail = gem.indices(len(table.shape) - 1) - mapped = gem.IndexSum(gem.Product(A, gem.Indexed(table, (j, *tail))), (k,)) + mapped = gem.IndexSum( + gem.Product(self._entry(r, a), gem.Indexed(table, (a, *tail))), (a,)) return gem.ComponentTensor(mapped, (r, *tail)) def __getitem__(self, alpha): try: return self._tabulation_cache[alpha] except KeyError: - result = self.matvec(self.ref_tabulation[alpha]) + result = self.matmul(self.ref_tabulation[alpha]) return self._tabulation_cache.setdefault(alpha, result) def __iter__(self): diff --git a/gem/gem.py b/gem/gem.py index 006814969..d3d8c3a08 100644 --- a/gem/gem.py +++ b/gem/gem.py @@ -681,6 +681,15 @@ def __reduce__(self): return type(self), (self.expression,) +def _index_free_indices(index): + """Return the free indices represented by an index expression.""" + if isinstance(index, Index): + return (index,) + if isinstance(index, VariableIndex): + return index.expression.free_indices + return () + + class Indexed(Scalar): __slots__ = ('children', 'multiindex', 'indirect_children') __back__ = ('multiindex',) @@ -1072,9 +1081,11 @@ def __new__(cls, i, j, dtype=None): self = super(Delta, cls).__new__(cls) self.i = i self.j = j - # Set up free indices - free_indices = [index for index in (i, j) if isinstance(index, Index)] - self.free_indices = tuple(unique(free_indices)) + # Set up free indices. A VariableIndex is not itself a free index, + # but the expression it wraps may be free in others; those propagate + # here exactly as they do through Indexed. + self.free_indices = tuple(unique(chain.from_iterable( + _index_free_indices(index) for index in (i, j)))) self._dtype = dtype return self diff --git a/gem/optimise.py b/gem/optimise.py index 76b2a163c..9958f5af9 100644 --- a/gem/optimise.py +++ b/gem/optimise.py @@ -14,7 +14,7 @@ reuse_if_untouched_arg, traversal) from gem.gem import (Node, Failure, Identity, Constant, Literal, Zero, Product, Sum, Comparison, Conditional, Division, - Index, VariableIndex, Indexed, FlexiblyIndexed, + Index, IndexBase, VariableIndex, Indexed, FlexiblyIndexed, IndexSum, ComponentTensor, ListTensor, Delta, partial_indexed, one) @@ -100,7 +100,13 @@ def replace_indices(node, self, subst): def _replace_indices_atomic(i, self, subst): if isinstance(i, VariableIndex): new_expr = self(i.expression, subst) - return i if new_expr == i.expression else VariableIndex(new_expr) + if new_expr == i.expression: + return i + # A variable index that substitution has made constant is a fixed + # index, and folding it lets the lookup itself be evaluated. + if isinstance(new_expr, Literal) and not new_expr.shape: + return int(new_expr.array) + return VariableIndex(new_expr) else: substitute = dict(subst) return substitute.get(i, i) @@ -935,6 +941,185 @@ def repeated_contractions(expression): return frozenset(expr for expr, count in counts.items() if count > 1) +def _cancellable_delta(node: Node) -> bool: + """Is there a Delta below ``node`` cancelling one of its own indices? + + Parameters + ---------- + node + An IndexSum. + + Returns + ------- + bool + Whether cancelling is possible below it. + + """ + contracted = frozenset(node.multiindex) + return any(isinstance(child, Delta) + and bool({child.i, child.j} & contracted) + for child in traversal(node.children)) + + +def _constant_map(index: IndexBase) -> tuple | None: + """The literal table behind a VariableIndex, and the indices addressing it. + + Parameters + ---------- + index + Index to inspect. + + Returns + ------- + tuple or None + ``(array, indices)`` when the index is a lookup into a Literal with a + plain multiindex, otherwise None. + + """ + if not isinstance(index, VariableIndex): + return None + expression = index.expression + if not isinstance(expression, Indexed): + return None + table, = expression.children + if not isinstance(table, Literal): + return None + if not all(isinstance(i, Index) for i in expression.multiindex): + return None + return table.array, expression.multiindex + + +def _pull_back( + delta: Delta, + sum_indices: Iterable[Index], + factors: Iterable[Node], + replacer: MemoizerArg) -> tuple | None: + """Contract a Delta's own axes before its column axis. + + ``sum_a (sum_rk v(r,k) delta(c(r,k), a)) T(a, q)`` is cancelled by + substituting ``a := c(r,k)``, which makes ``T`` depend on ``r`` and ``k`` + and so forces that contraction inside the ``q`` loop. When ``r`` and + ``k`` are contracted here and ``T`` carries indices of its own, summing + them first is cheaper: it yields a dense vector indexed by ``a``. + + Parameters + ---------- + delta + Candidate Delta, a factor of the product. + sum_indices + Indices contracted over the product. + factors + Product factors. + replacer + ``MemoizerArg(filtered_replace_indices)``. + + Returns + ------- + tuple or None + New ``(sum_indices, factors)``, or None when cancelling is better. + + """ + column = delta.j if isinstance(delta.i, VariableIndex) else delta.i + variable = delta.i if isinstance(delta.i, VariableIndex) else delta.j + if not isinstance(column, Index) or not isinstance(variable, VariableIndex): + return None + lookup = _constant_map(variable) + if lookup is None: + return None + table, source_indices = lookup + sources = frozenset(source_indices) + if column not in sum_indices or not sources <= set(sum_indices): + return None + + others = [f for f in factors if f is not delta] + spanning = [f for f in others if column in f.free_indices] + pulled = [f for f in others if column not in f.free_indices] + if not spanning or not pulled: + return None + # Cancelling couples the spanning factors to the source indices. That + # only costs anything when they carry indices of their own. + if not any(set(f.free_indices) - sources - {column} for f in spanning): + return None + + vector = numpy.empty(column.extent, dtype=object) + contributions = defaultdict(list) + for position in numpy.ndindex(table.shape): + substitution = tuple(zip(source_indices, (int(p) for p in position))) + contributions[int(table[position])].append(substitution) + for value in range(column.extent): + terms = [make_product([replacer(f, substitution) for f in pulled]) + for substitution in contributions.get(value, ())] + # A reference basis function that no row maps onto contributes nothing. + vector[value] = make_sum(terms) if terms else Zero() + + rest = tuple(i for i in sum_indices if i not in sources) + return rest, [Indexed(ListTensor(vector), (column,)), *spanning] + + +def cancel_deltas( + sum_indices: Iterable[Index], + factors: Iterable[Node], + replacer: MemoizerArg) -> tuple[list, list]: + """Cancel contracted Deltas, pulling a map back through its own axes first. + + Parameters + ---------- + sum_indices + Indices contracted over the product. + factors + Product factors. + replacer + ``MemoizerArg(filtered_replace_indices)``. + + Returns + ------- + tuple + Remaining sum indices and factors. + + """ + for delta in [f for f in factors if isinstance(f, Delta)]: + specialised = _pull_back(delta, sum_indices, factors, replacer) + if specialised is not None: + sum_indices, factors = specialised + break + return delta_elimination(sum_indices, factors, index_replacer=replacer) + + +def eliminate_deltas(expression: Node) -> Node: + """Cancel contracted Deltas that ``delta_elimination`` cannot reach. + + Parameters + ---------- + expression + Root of a scalar GEM expression. + + Returns + ------- + Node + Expression with those Deltas cancelled. + + Notes + ----- + ``delta_elimination`` only inspects top-level product factors, so a Delta + inside a preserved linear map is invisible to it. Flattening the product + tree first exposes it, and hoists the contractions it sits under so that + substituting the Delta's variable index cannot capture them. + + """ + replacer = MemoizerArg(filtered_replace_indices) + + def visit(node, self): + node = reuse_if_untouched(node, self) + if not isinstance(node, IndexSum) or not _cancellable_delta(node): + return node + sum_indices, factors = traverse_product(node, index_replacer=replacer) + sum_indices, factors = cancel_deltas(sum_indices, factors, replacer) + factors = [replacer(factor, ()) for factor in factors] + return IndexSum(make_product(factors), tuple(sum_indices)) + + return Memoizer(visit)(expression) + + def contraction(expression): """Optimise the contractions of the tensor product at the root of the expression, including: @@ -961,7 +1146,7 @@ def rebuild(expression): sum_indices, factors = traverse_product( expression, index_replacer=index_replacer, stop_at=lambda e: e is not root and e in keep) - sum_indices, factors = delta_elimination(sum_indices, factors, index_replacer=index_replacer) + sum_indices, factors = cancel_deltas(sum_indices, factors, index_replacer) factors = [index_replacer(f, ()) for f in factors] return sum_factorise(sum_indices, factors) From a5baaf4d1e9225ab5c9afb3cee811512ad11563f Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sun, 23 Aug 2026 12:21:10 +0100 Subject: [PATCH 4/7] Share the contraction indices of a mapped tabulation MappedTabulation minted a fresh index on every call, so the tabulations of different derivative orders contracted over distinct indices of equal extent. Expressions that are structurally equal then hash apart, and the scheduler gives each its own loop nest. Reuse one index per instance for the reference basis and for the padded row. Equal tabulations now share a subexpression, and their loops fuse without any change to the scheduler. The row index of a tabulation stays per call, since a ComponentTensor binds it and sharing it only forces redundant materialisation. Guzman-Neilan 3D action: 51920 -> 51296 flops, 17 -> 15 array temporaries; 2D action: 2676 -> 2586 flops. Four groups of sibling loops over equal extents collapse to one loop each, and the Argyris and Johnson-Mercier actions lose theirs likewise. Co-Authored-By: Claude Opus 5 --- finat/physically_mapped.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/finat/physically_mapped.py b/finat/physically_mapped.py index 315bfd055..847fb6595 100644 --- a/finat/physically_mapped.py +++ b/finat/physically_mapped.py @@ -1,5 +1,6 @@ from abc import ABCMeta, abstractmethod from collections.abc import Iterable, Mapping +from functools import cached_property import gem import numpy @@ -83,6 +84,16 @@ def __init__( self._values = gem.ListTensor(values) self._tabulation_cache = {} + @cached_property + def _reference_index(self) -> gem.Index: + """Contraction over the reference basis, shared by all tabulations.""" + return gem.Index(extent=self._value_dim) + + @cached_property + def _row_index(self) -> gem.Index: + """Contraction over a padded row, shared by all tabulations.""" + return gem.Index(extent=self._width) + def _entry(self, r: gem.Index, a: gem.Index) -> gem.Node: """Entry ``M[r, a]`` of the basis transformation. @@ -101,7 +112,7 @@ def _entry(self, r: gem.Index, a: gem.Index) -> gem.Node: is ordinary GEM algebra. """ - k = gem.Index(extent=self._width) + k = self._row_index entry = gem.Indexed( self._values, (gem.VariableIndex(gem.Indexed(self._value_indices, (r, k))),)) @@ -111,13 +122,13 @@ def _entry(self, r: gem.Index, a: gem.Index) -> gem.Node: def matrix(self) -> gem.Node: """The basis transformation as a rank-2 GEM expression.""" r = gem.Index(extent=self._space_dim) - a = gem.Index(extent=self._value_dim) + a = self._reference_index return gem.ComponentTensor(self._entry(r, a), (r, a)) def matmul(self, table: gem.Node) -> gem.Node: """Apply the basis transformation to a reference tabulation.""" r = gem.Index(extent=self._space_dim) - a = gem.Index(extent=self._value_dim) + a = self._reference_index tail = gem.indices(len(table.shape) - 1) mapped = gem.IndexSum( gem.Product(self._entry(r, a), gem.Indexed(table, (a, *tail))), (a,)) From a0ac495d41d1c1af7a1a795b7cc016bcf0cf1375 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 28 Aug 2026 11:10:35 +0100 Subject: [PATCH 5/7] Test sparse mapped tabulations after contraction MappedTabulation retains its selector Delta until the contraction optimizer runs, so its raw DAG legitimately contains two Product nodes. Exercise the production contraction path before asserting the sparse arithmetic structure. --- test/finat/test_zany_mapping.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/test/finat/test_zany_mapping.py b/test/finat/test_zany_mapping.py index 59f1e8883..f0742caf5 100644 --- a/test/finat/test_zany_mapping.py +++ b/test/finat/test_zany_mapping.py @@ -7,6 +7,7 @@ from gem.interpreter import evaluate from gem.node import traversal +from gem.optimise import contraction from finat.physically_mapped import MappedTabulation, PhysicallyMappedElement @@ -25,7 +26,8 @@ def test_sparse_mapped_tabulation(): # indexes a vector rather than materialising a symbolic matrix. assert mapped_tabulation._values.shape == (3,) - mapped = mapped_tabulation[None] + i, j = gem.indices(2) + mapped = contraction(gem.Indexed(mapped_tabulation[None], (i, j))) # The three unit entries cost no multiplication, and the one remaining # nonzero costs exactly one. Nothing is selected by a branch. @@ -35,7 +37,10 @@ def test_sparse_mapped_tabulation(): assert not any(isinstance(node, gem.Conditional) for node in traversal((mapped,))) - actual, = evaluate([mapped], {coefficient: np.asarray(2.0)}) + actual, = evaluate( + [gem.ComponentTensor(mapped, (i, j))], + {coefficient: np.asarray(2.0)}, + ) expected = np.asarray([[1.0, 0.0, 2.0], [0.0, 1.0, 0.0]]) \ @ table_values assert np.array_equal(actual.arr, expected) From be99e70549b90828032d22a5ae6aeef94d4d4a77 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 29 Aug 2026 15:45:02 +0100 Subject: [PATCH 6/7] Name the delta rewrites for what each one does Three functions carried names for delta cancellation and one of them was a wrapper. Fold the search for a pull-back candidate into the pull-back itself, and call it pull_back_indirect_delta. Its callers now run it and delta_elimination in sequence, which is what they always did. Rename the whole-DAG traversal to cancel_nested_deltas, so that it no longer reads as a synonym of delta_elimination. Its guard walked the subtree again at every enclosing contraction. A memoised map from a node to the Delta axes below it answers the same question once per node. Use one helper for the cardinality of an index space, and promote the traversal child rule in gem.node so that other modules can share it. Drop MappedTabulation.matrix(); nothing calls it. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k --- finat/physically_mapped.py | 6 -- gem/cost.py | 12 +++ gem/node.py | 21 ++-- gem/optimise.py | 212 +++++++++++++++---------------------- 4 files changed, 111 insertions(+), 140 deletions(-) diff --git a/finat/physically_mapped.py b/finat/physically_mapped.py index 847fb6595..da88e7d6e 100644 --- a/finat/physically_mapped.py +++ b/finat/physically_mapped.py @@ -119,12 +119,6 @@ def _entry(self, r: gem.Index, a: gem.Index) -> gem.Node: column = gem.VariableIndex(gem.Indexed(self._columns, (r, k))) return gem.IndexSum(gem.Product(entry, gem.Delta(column, a)), (k,)) - def matrix(self) -> gem.Node: - """The basis transformation as a rank-2 GEM expression.""" - r = gem.Index(extent=self._space_dim) - a = self._reference_index - return gem.ComponentTensor(self._entry(r, a), (r, a)) - def matmul(self, table: gem.Node) -> gem.Node: """Apply the basis transformation to a reference tabulation.""" r = gem.Index(extent=self._space_dim) diff --git a/gem/cost.py b/gem/cost.py index a46a22387..7451b5761 100644 --- a/gem/cost.py +++ b/gem/cost.py @@ -86,6 +86,18 @@ def iteration_count(indices: Iterable[Index]) -> int: return int(numpy.prod([index.extent for index in indices], dtype=int)) +def index_space_literal(indices: Iterable[Index]) -> Literal: + """The cardinality of a rectangular index space, as a scalar. + + The empty product is one, so contracting it counts the tuples in the + index space. + + :arg indices: indices spanning the space + :returns: the number of points, as a floating point literal + """ + return Literal(float(iteration_count(indices))) + + def operation_count(node: Node) -> int: """Estimate the scalar operations performed by one GEM node. diff --git a/gem/node.py b/gem/node.py index 190fe6d40..34e057754 100644 --- a/gem/node.py +++ b/gem/node.py @@ -103,7 +103,12 @@ def get_hash(self): return hash((type(self), *self._arguments)) -def _make_traversal_children(node): +def traversal_children(node): + """The children a DAG walk descends into, index expressions included. + + :arg node: a GEM expression + :returns: the child nodes, plus the nodes hidden in index expressions + """ if isinstance(node, (gem.Indexed, gem.FlexiblyIndexed)): # Include child nodes hidden in index expressions. return node.children + node.indirect_children @@ -117,7 +122,7 @@ def pre_traversal(expression_dags): Notes ----- This function also walks through nodes in index expressions - (e.g., `VariableIndex`s); see ``_make_traversal_children()``. + (e.g., `VariableIndex`s); see ``traversal_children()``. """ seen = set() @@ -133,7 +138,7 @@ def pre_traversal(expression_dags): while lifo: node = lifo.pop() yield node - children = _make_traversal_children(node) + children = traversal_children(node) for child in reversed(children): if child not in seen: seen.add(child) @@ -146,7 +151,7 @@ def post_traversal(expression_dags): Notes ----- This function also walks through nodes in index expressions - (e.g., `VariableIndex`s); see ``_make_traversal_children()``. + (e.g., `VariableIndex`s); see ``traversal_children()``. """ @@ -158,13 +163,13 @@ def post_traversal(expression_dags): for root in expression_dags: if root not in seen: seen.add(root) - lifo.append((root, list(_make_traversal_children(root)))) + lifo.append((root, list(traversal_children(root)))) while lifo: node, deps = lifo[-1] for i, dep in enumerate(deps): if dep is not None and dep not in seen: - lifo.append((dep, list(_make_traversal_children(dep)))) + lifo.append((dep, list(traversal_children(dep)))) deps[i] = None break else: @@ -184,12 +189,12 @@ def collect_refcount(expression_dags): ----- This function also collects reference counts of nodes in index expressions (e.g., `VariableIndex`s); see - ``_make_traversal_children()``. + ``traversal_children()``. """ result = collections.Counter(expression_dags) for node in traversal(expression_dags): - result.update(_make_traversal_children(node)) + result.update(traversal_children(node)) return result diff --git a/gem/optimise.py b/gem/optimise.py index 9958f5af9..f009618d2 100644 --- a/gem/optimise.py +++ b/gem/optimise.py @@ -9,9 +9,10 @@ import numpy +from gem.cost import index_space_literal from gem.utils import groupby from gem.node import (Memoizer, MemoizerArg, reuse_if_untouched, - reuse_if_untouched_arg, traversal) + reuse_if_untouched_arg, traversal, traversal_children) from gem.gem import (Node, Failure, Identity, Constant, Literal, Zero, Product, Sum, Comparison, Conditional, Division, Index, IndexBase, VariableIndex, Indexed, FlexiblyIndexed, @@ -546,8 +547,7 @@ def _sum_factorise_connected(sum_indices, groups): :returns: optimised GEM expression """ if not groups: - extent = numpy.prod([index.extent for index in sum_indices], dtype=int) - return Literal(float(extent)) + return index_space_literal(sum_indices) if len(groups) <= _MAX_PLANNED_FACTORS: return _plan_contraction(sum_indices, groups) @@ -597,10 +597,7 @@ def sum_factorise(sum_indices, factors): :returns: optimised GEM expression """ if len(factors) == 0: - # The empty product is one, so contracting it counts the tuples in the - # index space. - extent = numpy.prod([index.extent for index in sum_indices], dtype=int) - return Literal(float(extent)) + return index_space_literal(sum_indices) # Form groups by free indices groups = groupby(factors, key=lambda f: f.free_indices) @@ -941,24 +938,21 @@ def repeated_contractions(expression): return frozenset(expr for expr, count in counts.items() if count > 1) -def _cancellable_delta(node: Node) -> bool: - """Is there a Delta below ``node`` cancelling one of its own indices? +def _delta_axes(node: Node, self: Memoizer) -> frozenset: + """The axes compared by the Deltas below a node, including its own. - Parameters - ---------- - node - An IndexSum. - - Returns - ------- - bool - Whether cancelling is possible below it. + Memoising this over the DAG keeps the search for a cancellable Delta + linear, rather than re-walking the subtree at every enclosing + contraction. + :arg node: a GEM expression + :arg self: memoizer visiting the DAG + :returns: the indices some Delta at or below ``node`` compares """ - contracted = frozenset(node.multiindex) - return any(isinstance(child, Delta) - and bool({child.i, child.j} & contracted) - for child in traversal(node.children)) + axes = frozenset().union(*map(self, traversal_children(node))) + if isinstance(node, Delta): + axes = axes | {node.i, node.j} + return axes def _constant_map(index: IndexBase) -> tuple | None: @@ -989,131 +983,94 @@ def _constant_map(index: IndexBase) -> tuple | None: return table.array, expression.multiindex -def _pull_back( - delta: Delta, +def pull_back_indirect_delta( sum_indices: Iterable[Index], factors: Iterable[Node], - replacer: MemoizerArg) -> tuple | None: - """Contract a Delta's own axes before its column axis. + replacer: MemoizerArg) -> tuple: + """Contract an indirect Delta's own axes before its column axis. ``sum_a (sum_rk v(r,k) delta(c(r,k), a)) T(a, q)`` is cancelled by - substituting ``a := c(r,k)``, which makes ``T`` depend on ``r`` and ``k`` - and so forces that contraction inside the ``q`` loop. When ``r`` and - ``k`` are contracted here and ``T`` carries indices of its own, summing - them first is cheaper: it yields a dense vector indexed by ``a``. - - Parameters - ---------- - delta - Candidate Delta, a factor of the product. - sum_indices - Indices contracted over the product. - factors - Product factors. - replacer - ``MemoizerArg(filtered_replace_indices)``. - - Returns - ------- - tuple or None - New ``(sum_indices, factors)``, or None when cancelling is better. - + `delta_elimination` substituting ``a := c(r,k)``, which makes ``T`` + depend on ``r`` and ``k`` and so forces that contraction inside the ``q`` + loop. When ``r`` and ``k`` are contracted here and ``T`` carries indices + of its own, summing them first is cheaper: it yields a dense vector + indexed by ``a``. Run this before `delta_elimination` to take that + cheaper route where it exists. + + :arg sum_indices: indices contracted over the product + :arg factors: product factors + :arg replacer: ``MemoizerArg(filtered_replace_indices)`` + :returns: new ``(sum_indices, factors)``, unchanged when no Delta is + worth pulling back """ - column = delta.j if isinstance(delta.i, VariableIndex) else delta.i - variable = delta.i if isinstance(delta.i, VariableIndex) else delta.j - if not isinstance(column, Index) or not isinstance(variable, VariableIndex): - return None - lookup = _constant_map(variable) - if lookup is None: - return None - table, source_indices = lookup - sources = frozenset(source_indices) - if column not in sum_indices or not sources <= set(sum_indices): - return None - - others = [f for f in factors if f is not delta] - spanning = [f for f in others if column in f.free_indices] - pulled = [f for f in others if column not in f.free_indices] - if not spanning or not pulled: - return None - # Cancelling couples the spanning factors to the source indices. That - # only costs anything when they carry indices of their own. - if not any(set(f.free_indices) - sources - {column} for f in spanning): - return None - - vector = numpy.empty(column.extent, dtype=object) - contributions = defaultdict(list) - for position in numpy.ndindex(table.shape): - substitution = tuple(zip(source_indices, (int(p) for p in position))) - contributions[int(table[position])].append(substitution) - for value in range(column.extent): - terms = [make_product([replacer(f, substitution) for f in pulled]) - for substitution in contributions.get(value, ())] - # A reference basis function that no row maps onto contributes nothing. - vector[value] = make_sum(terms) if terms else Zero() - - rest = tuple(i for i in sum_indices if i not in sources) - return rest, [Indexed(ListTensor(vector), (column,)), *spanning] - - -def cancel_deltas( - sum_indices: Iterable[Index], - factors: Iterable[Node], - replacer: MemoizerArg) -> tuple[list, list]: - """Cancel contracted Deltas, pulling a map back through its own axes first. - - Parameters - ---------- - sum_indices - Indices contracted over the product. - factors - Product factors. - replacer - ``MemoizerArg(filtered_replace_indices)``. - - Returns - ------- - tuple - Remaining sum indices and factors. + for delta in factors: + if not isinstance(delta, Delta): + continue + column = delta.j if isinstance(delta.i, VariableIndex) else delta.i + variable = delta.i if isinstance(delta.i, VariableIndex) else delta.j + if not isinstance(column, Index) or not isinstance(variable, VariableIndex): + continue + lookup = _constant_map(variable) + if lookup is None: + continue + table, source_indices = lookup + sources = frozenset(source_indices) + if column not in sum_indices or not sources <= set(sum_indices): + continue - """ - for delta in [f for f in factors if isinstance(f, Delta)]: - specialised = _pull_back(delta, sum_indices, factors, replacer) - if specialised is not None: - sum_indices, factors = specialised - break - return delta_elimination(sum_indices, factors, index_replacer=replacer) + others = [f for f in factors if f is not delta] + spanning = [f for f in others if column in f.free_indices] + pulled = [f for f in others if column not in f.free_indices] + if not spanning or not pulled: + continue + # Cancelling couples the spanning factors to the source indices. That + # only costs anything when they carry indices of their own. + if not any(set(f.free_indices) - sources - {column} for f in spanning): + continue + vector = numpy.empty(column.extent, dtype=object) + contributions = defaultdict(list) + for position in numpy.ndindex(table.shape): + substitution = tuple(zip(source_indices, (int(p) for p in position))) + contributions[int(table[position])].append(substitution) + for value in range(column.extent): + terms = [make_product([replacer(f, substitution) for f in pulled]) + for substitution in contributions.get(value, ())] + # A reference basis function that no row maps onto contributes + # nothing. + vector[value] = make_sum(terms) if terms else Zero() + + rest = tuple(i for i in sum_indices if i not in sources) + return rest, [Indexed(ListTensor(vector), (column,)), *spanning] -def eliminate_deltas(expression: Node) -> Node: - """Cancel contracted Deltas that ``delta_elimination`` cannot reach. + return sum_indices, factors - Parameters - ---------- - expression - Root of a scalar GEM expression. - Returns - ------- - Node - Expression with those Deltas cancelled. +def cancel_nested_deltas(expression: Node) -> Node: + """Apply `delta_elimination` at every contraction of a whole DAG. - Notes - ----- - ``delta_elimination`` only inspects top-level product factors, so a Delta + `delta_elimination` only inspects top-level product factors, so a Delta inside a preserved linear map is invisible to it. Flattening the product tree first exposes it, and hoists the contractions it sits under so that substituting the Delta's variable index cannot capture them. + :arg expression: root of a scalar GEM expression + :returns: the expression with those Deltas cancelled """ replacer = MemoizerArg(filtered_replace_indices) + delta_axes = Memoizer(_delta_axes) def visit(node, self): node = reuse_if_untouched(node, self) - if not isinstance(node, IndexSum) or not _cancellable_delta(node): + if not isinstance(node, IndexSum): + return node + if not delta_axes(node).intersection(node.multiindex): return node sum_indices, factors = traverse_product(node, index_replacer=replacer) - sum_indices, factors = cancel_deltas(sum_indices, factors, replacer) + sum_indices, factors = pull_back_indirect_delta( + sum_indices, factors, replacer) + sum_indices, factors = delta_elimination( + sum_indices, factors, index_replacer=replacer) factors = [replacer(factor, ()) for factor in factors] return IndexSum(make_product(factors), tuple(sum_indices)) @@ -1146,7 +1103,10 @@ def rebuild(expression): sum_indices, factors = traverse_product( expression, index_replacer=index_replacer, stop_at=lambda e: e is not root and e in keep) - sum_indices, factors = cancel_deltas(sum_indices, factors, index_replacer) + sum_indices, factors = pull_back_indirect_delta( + sum_indices, factors, index_replacer) + sum_indices, factors = delta_elimination( + sum_indices, factors, index_replacer=index_replacer) factors = [index_replacer(f, ()) for f in factors] return sum_factorise(sum_indices, factors) From 5934db9a1c8d7882b2a6e4fc87f974159efb3746 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 8 Sep 2026 16:49:41 +0100 Subject: [PATCH 7/7] Cancel only the Deltas that no later pass can lower `cancel_nested_deltas` rewrote every Delta it reached, and so made a cost decision it had no way to cost. Substituting a Delta between two plain indices makes the gather it feeds depend on an argument, and monomial collection then expands the contraction that gather sits in, one monomial per basis function. H(div)/H(curl) and tensor element interpolation lose their sum factorisation that way. The pass exists for the Delta a padded basis transformation emits, which compares a `VariableIndex`. That is also the only kind nothing downstream can lower: monomial collection cancels the Deltas that surface as factors of a monomial, and one buried in a preserved linear map never does, so it would reach code generation. Cancel that kind, and leave a Delta between two plain indices to monomial collection, which cancels it knowing what the substitution costs there. Narrowing the memoised search to those Deltas keeps the pass off the contractions it has no business flattening, and a contraction where nothing cancelled is returned untouched rather than flattened into a single product. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01M6QZ26z2p9o7sYzxZXV1B3 --- gem/optimise.py | 61 ++++++++++++++++++++++++++++++++----------------- 1 file changed, 40 insertions(+), 21 deletions(-) diff --git a/gem/optimise.py b/gem/optimise.py index f009618d2..9e1066874 100644 --- a/gem/optimise.py +++ b/gem/optimise.py @@ -325,12 +325,15 @@ def select_expression(expressions, index): return ComponentTensor(selected, alpha) -def delta_elimination(sum_indices, factors, index_replacer=None): +def delta_elimination(sum_indices, factors, index_replacer=None, indirect_only=False): """IndexSum-Delta cancellation. :arg sum_indices: free indices for contractions :arg factors: product factors :kwarg index_replacer: MemoizerArg(filtered_replace_indices) + :kwarg indirect_only: only cancel a Delta that compares a + :class:`~.VariableIndex`, leaving one between two + plain indices to a later pass :returns: optimised (sum_indices, factors) """ @@ -347,9 +350,15 @@ def substitute(expression, from_, to_): else: return Indexed(ComponentTensor(expression, (from_,)), (to_,)) - delta_queue = [(f, index) - for f in factors if isinstance(f, Delta) - for index in (f.i, f.j) if index in sum_indices] + def cancellable(factors): + return [(f, index) + for f in factors if isinstance(f, Delta) + if not indirect_only + or isinstance(f.i, VariableIndex) or isinstance(f.j, VariableIndex) + for index in (f.i, f.j) + if index in sum_indices] + + delta_queue = cancellable(factors) while delta_queue: delta, from_ = delta_queue[0] to_, = list({delta.i, delta.j} - {from_}) @@ -358,9 +367,7 @@ def substitute(expression, from_, to_): factors = [substitute(f, from_, to_) for f in factors] - delta_queue = [(f, index) - for f in factors if isinstance(f, Delta) - for index in (f.i, f.j) if index in sum_indices] + delta_queue = cancellable(factors) return sum_indices, factors @@ -938,19 +945,21 @@ def repeated_contractions(expression): return frozenset(expr for expr, count in counts.items() if count > 1) -def _delta_axes(node: Node, self: Memoizer) -> frozenset: - """The axes compared by the Deltas below a node, including its own. +def _indirect_delta_axes(node: Node, self: Memoizer) -> frozenset: + """The axes compared by the indirect Deltas below a node, including its own. - Memoising this over the DAG keeps the search for a cancellable Delta - linear, rather than re-walking the subtree at every enclosing - contraction. + An indirect Delta compares a :class:`~.VariableIndex`, and is the only + kind `cancel_nested_deltas` cancels. Memoising this over the DAG keeps + the search for one linear, rather than re-walking the subtree at every + enclosing contraction. :arg node: a GEM expression :arg self: memoizer visiting the DAG - :returns: the indices some Delta at or below ``node`` compares + :returns: the indices some indirect Delta at or below ``node`` compares """ axes = frozenset().union(*map(self, traversal_children(node))) - if isinstance(node, Delta): + if isinstance(node, Delta) and any(isinstance(i, VariableIndex) + for i in (node.i, node.j)): axes = axes | {node.i, node.j} return axes @@ -1047,18 +1056,23 @@ def pull_back_indirect_delta( def cancel_nested_deltas(expression: Node) -> Node: - """Apply `delta_elimination` at every contraction of a whole DAG. + """Cancel the indirect Deltas at every contraction of a whole DAG. `delta_elimination` only inspects top-level product factors, so a Delta inside a preserved linear map is invisible to it. Flattening the product tree first exposes it, and hoists the contractions it sits under so that substituting the Delta's variable index cannot capture them. + A Delta comparing a :class:`~.VariableIndex` is the only kind handled + here. It is the only kind that has to be: nothing downstream can lower + one. A Delta between two plain indices is left to monomial collection, + which cancels it knowing what the substitution costs there. + :arg expression: root of a scalar GEM expression :returns: the expression with those Deltas cancelled """ replacer = MemoizerArg(filtered_replace_indices) - delta_axes = Memoizer(_delta_axes) + delta_axes = Memoizer(_indirect_delta_axes) def visit(node, self): node = reuse_if_untouched(node, self) @@ -1067,12 +1081,17 @@ def visit(node, self): if not delta_axes(node).intersection(node.multiindex): return node sum_indices, factors = traverse_product(node, index_replacer=replacer) - sum_indices, factors = pull_back_indirect_delta( + cancelled, new_factors = pull_back_indirect_delta( sum_indices, factors, replacer) - sum_indices, factors = delta_elimination( - sum_indices, factors, index_replacer=replacer) - factors = [replacer(factor, ()) for factor in factors] - return IndexSum(make_product(factors), tuple(sum_indices)) + cancelled, new_factors = delta_elimination( + cancelled, new_factors, index_replacer=replacer, indirect_only=True) + if tuple(cancelled) == tuple(sum_indices) and tuple(new_factors) == tuple(factors): + # Nothing cancelled, so rebuilding would only flatten the + # contractions this node nests into a single product, and sum + # factorisation needs them nested. + return node + factors = [replacer(factor, ()) for factor in new_factors] + return IndexSum(make_product(factors), tuple(cancelled)) return Memoizer(visit)(expression)