Factor scalar maps through basis transformations

finat/physically_mapped.py

@@ -4,6 +4,7 @@
 
 import gem
 import numpy
+from gem.optimise import factorise_scalar_sums
 
 from finat.citations import cite
 
@@ -21,25 +22,17 @@ def dual_transformation(self, Q, coordinate_mapping=None):
 
 
 class MappedTabulation(Mapping):
-    """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.
+    """A lazy tabulation dict that applies the basis transformation only
+    on the requested derivatives.
 
+    Rows are padded to a common number of entries, so that a loop over the
+    basis index has an affine iteration domain.  Constant tables select the
+    reference column and one of the distinct symbolic coefficients, which
+    shares equal entries without materialising the 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__(
@@ -66,7 +59,8 @@ def __init__(
         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)
+            data[index, :len(row)] = tuple(
+                factorise_scalar_sums(gem.as_gem(value)) for _, value in row)
         self._width = width
         self._columns = gem.Literal(columns, dtype=gem.uint_type)
         values = []
@@ -97,20 +91,11 @@ def _row_index(self) -> gem.Index:
     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.
-
+        :arg r: index over the rows retained by the element
+        :arg a: index over the reference basis
+        :returns: 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 = self._row_index
         entry = gem.Indexed(

gem/coffee.py

@@ -4,17 +4,17 @@
 This file is NOT for code generation as a COFFEE AST.
 """
 
-from collections import defaultdict
+from collections import Counter, defaultdict
 from itertools import chain, repeat
 import logging
 
 import numpy
 
-from gem.gem import ComponentTensor, Index, Indexed, IndexSum, one
+from gem.gem import ComponentTensor, Index, Indexed, IndexSum, Literal, Node, one
 from gem.node import MemoizerArg
 from gem.cost import has_arithmetic
 from gem.optimise import (filtered_replace_indices,
-                          make_sum, make_product)
+                          make_sum, make_product, traverse_sum)
 from gem.refactorise import Monomial, MonomialSum
 from gem.utils import groupby
 
@@ -201,51 +201,92 @@ def group_key(monomial):
     return new_monomials
 
 
-def _share_linear_maps(
+def _additive_map_key(
+        expression: Node) -> frozenset[tuple[Node, int]]:
+    """Represent an additive linear map independently of term order."""
+    return frozenset(Counter(traverse_sum(expression)).items())
+
+
+def _extract_repeated_linear_maps(
         monomial_sum: MonomialSum,
         linear_indices: tuple[Index, ...]) -> MonomialSum:
-    """Share isomorphic maps of distinct multilinear axes.
+    """Move uniform multiplicities from linear maps into scalar factors.
 
-    Parameters
-    ----------
-    monomial_sum
-        Sum-of-products representation of a multilinear expression.
-    linear_indices
-        Free indices identifying argument axes.
+    A repeated additive map is a scalar multiple of the map formed from its
+    distinct summands.  Keeping that scalar in the monomial remainder leaves
+    the atomic factor as the finite element map that should be materialised.
 
-    Returns
-    -------
-    MonomialSum
-        Representation whose repeated linear maps access one tensor.
+    :arg monomial_sum: sum-of-products representation of a multilinear
+                       expression
+    :arg linear_indices: free indices identifying argument axes
+    :returns: representation with primitive linear maps and scalar
+              multiplicities
+    """
+    linear_set = frozenset(linear_indices)
+    result = MonomialSum()
+    for monomial in monomial_sum:
+        atomics = []
+        factors = []
+        for atomic in monomial.atomics:
+            involved = linear_set.intersection(atomic.free_indices)
+            if len(involved) == 1:
+                summands = traverse_sum(atomic)
+                counts = Counter(summands)
+                multiplicities = set(counts.values())
+                if len(multiplicities) == 1:
+                    multiplicity, = multiplicities
+                    if multiplicity > 1:
+                        atomic = make_sum(list(counts))
+                        factors.append(Literal(float(multiplicity)))
+            atomics.append(atomic)
+        result.add(
+            monomial.sum_indices,
+            atomics,
+            make_product((*factors, monomial.rest)),
+        )
+    return result
 
-    Notes
-    -----
-    Test and trial axes use distinct indices even when they apply the same
-    finite element map.  Renaming each axis to a canonical index exposes
-    that isomorphism without inspecting the element family.  Materialising
-    the canonical map is generalised code motion: the basis transformation
-    is evaluated once and both axes index its result.
 
+def _share_linear_maps(
+        monomial_sum: MonomialSum,
+        linear_indices: tuple[Index, ...]) -> MonomialSum:
+    """Share isomorphic maps of distinct multilinear axes.
+
+    Test and trial axes use distinct indices even when they apply the same
+    finite element map.  Renaming each axis to a canonical index exposes that
+    isomorphism without inspecting the element family.  Materialising the
+    canonical map is generalised code motion: the basis transformation is
+    evaluated once and both axes index its result.
+
+    :arg monomial_sum: sum-of-products representation of a multilinear
+                       expression
+    :arg linear_indices: free indices identifying argument axes
+    :returns: representation whose repeated linear maps access one tensor
     """
     linear_indices = tuple(linear_indices)
+    monomial_sum = _extract_repeated_linear_maps(monomial_sum, linear_indices)
     linear_set = frozenset(linear_indices)
     canonical = {
         index.extent: Index(extent=index.extent)
         for index in linear_indices
     }
     replacer = MemoizerArg(filtered_replace_indices)
     groups = defaultdict(list)
+    representatives = {}
     for monomial in monomial_sum:
         for atomic in monomial.atomics:
             involved = linear_set.intersection(atomic.free_indices)
             if len(involved) != 1:
                 continue
             index, = involved
             normal = replacer(atomic, ((index, canonical[index.extent]),))
-            groups[normal].append((atomic, index))
+            key = _additive_map_key(normal)
+            groups[key].append((atomic, index))
+            representatives.setdefault(key, normal)
 
     replacements = {}
-    for normal, occurrences in groups.items():
+    for key, occurrences in groups.items():
+        normal = representatives[key]
         indices = {index for _, index in occurrences}
         if len(indices) < 2 or not has_arithmetic((normal,)):
             continue

gem/gem.py

@@ -727,7 +727,14 @@ def __new__(cls, aggregate, multiindex):
                 C, = B.children
                 kk = B.multiindex
                 ff = C.free_indices
-                if not any((j in ff) for j in jj):
+                # An index reached only through a VariableIndex is not a key
+                # of kk, so the rewrite below would drop its replacement and
+                # leave the old index behind, bound by nothing.  Those belong
+                # to replace_indices, which does substitute inside the lookup.
+                hh = frozenset(chain.from_iterable(
+                    k.expression.free_indices
+                    for k in kk if isinstance(k, VariableIndex)))
+                if not any((j in ff or j in hh) for j in jj):
                     # Only replace indices that are not present in C
                     rep = dict(zip(jj, ii))
                     ll = tuple(rep.get(k, k) for k in kk)

gem/impero_utils.py

@@ -128,14 +128,26 @@ def inline_temporaries(expressions, ops):
     for op in ops:
         if isinstance(op, imp.Evaluate):
             expr = op.expression
-            if expr.shape == () and refcount[expr] == 1:
+            reduction_view = (isinstance(expr, gem.ComponentTensor)
+                              and isinstance(expr.children[0], gem.IndexSum))
+            if (expr.shape == () or reduction_view) and refcount[expr] == 1:
                 candidates.add(expr)
 
+    # A reduction view carries no name of its own, so it can only be inlined
+    # where the consumer indexes straight through it.  Inverse, Solve and
+    # FlexiblyIndexed read the tensor whole and need it materialised.
+    indexed = {child for node in traversal(expressions)
+               if isinstance(node, gem.Indexed)
+               for child in node.children}
+    candidates = {expr for expr in candidates
+                  if not isinstance(expr, gem.ComponentTensor) or expr in indexed}
+
     # Prevent inlining that pulls expressions into inner loops
     for node in traversal(expressions):
         for child in node.children:
             if child in candidates and set(child.free_indices) < set(node.free_indices):
-                candidates.remove(child)
+                if not isinstance(child, gem.ComponentTensor):
+                    candidates.remove(child)
 
     # Filter out candidates
     return [op for op in ops if not (isinstance(op, imp.Evaluate) and op.expression in candidates)]