Preserve finite element linear maps during factorisation

gem/coffee.py

@@ -4,14 +4,17 @@
 This file is NOT for code generation as a COFFEE AST.
 """
 
+from collections import defaultdict
 from itertools import chain, repeat
 import logging
 
 import numpy
 
-from gem.gem import IndexSum, one
-from gem.optimise import make_sum, make_product
-from gem.refactorise import Monomial
+from gem.gem import ComponentTensor, Index, Indexed, IndexSum, one
+from gem.node import MemoizerArg
+from gem.optimise import (filtered_replace_indices, has_arithmetic,
+                          make_sum, make_product)
+from gem.refactorise import Monomial, MonomialSum
 from gem.utils import groupby
 
 
@@ -197,6 +200,74 @@ def group_key(monomial):
     return new_monomials
 
 
+def _share_linear_maps(
+        monomial_sum: MonomialSum,
+        linear_indices: tuple[Index, ...]) -> MonomialSum:
+    """Share isomorphic maps of distinct multilinear axes.
+
+    Parameters
+    ----------
+    monomial_sum
+        Sum-of-products representation of a multilinear expression.
+    linear_indices
+        Free indices identifying argument axes.
+
+    Returns
+    -------
+    MonomialSum
+        Representation whose repeated linear maps access one tensor.
+
+    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.
+
+    """
+    linear_indices = tuple(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)
+    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))
+
+    replacements = {}
+    for normal, occurrences in groups.items():
+        indices = {index for _, index in occurrences}
+        if len(indices) < 2 or not has_arithmetic((normal,)):
+            continue
+        index = canonical[next(iter(indices)).extent]
+        tensor = ComponentTensor(normal, (index,))
+        replacements.update(
+            (atomic, Indexed(tensor, (original,)))
+            for atomic, original in occurrences)
+
+    if not replacements:
+        return monomial_sum
+
+    result = MonomialSum()
+    for monomial in monomial_sum:
+        result.add(
+            monomial.sum_indices,
+            tuple(replacements.get(atomic, atomic)
+                  for atomic in monomial.atomics),
+            monomial.rest,
+        )
+    return result
+
+
 def optimise_monomial_sum(monomial_sum, linear_indices):
     """Choose optimal common atomic subexpressions and factorise a
     :class:`MonomialSum` object to create a GEM expression.
@@ -206,6 +277,7 @@ def optimise_monomial_sum(monomial_sum, linear_indices):
 
     :returns: factorised GEM expression
     """
+    monomial_sum = _share_linear_maps(monomial_sum, linear_indices)
     groups = groupby(monomial_sum, key=lambda m: frozenset(m.sum_indices))
     new_monomials = []
     for _, monomials in groups:

gem/flop_count.py

@@ -5,6 +5,7 @@
 
 import gem.gem as gem
 import gem.impero as imp
+from contextvars import ContextVar
 from functools import singledispatch
 import numpy
 import math
@@ -26,7 +27,11 @@ def statement_for(tree, temporaries):
     extent = tree.index.extent
     assert extent is not None
     child, = tree.children
-    flops = statement(child, temporaries)
+    token = _active_indices.set(_active_indices.get() | {tree.index})
+    try:
+        flops = statement(child, temporaries)
+    finally:
+        _active_indices.reset(token)
     return flops * extent
 
 
@@ -168,7 +173,14 @@ def flops_solve(expr, temporaries):
 
 @flops.register(gem.ComponentTensor)
 def flops_componenttensor(expr, temporaries):
-    raise ValueError("Not expecting ComponentTensor")
+    body, = expr.children
+    # Scheduling emits an assignment over the indices that no enclosing For
+    # already iterates.  Count those extents here; the loops that carry the
+    # rest count them.
+    implicit = tuple(index for index in expr.multiindex
+                     if index not in _active_indices.get())
+    extent = numpy.prod([index.extent for index in implicit], dtype=int)
+    return extent * expression_flops(body, temporaries)
 
 
 def expression_flops(expression, temporaries, top=False):
@@ -192,6 +204,14 @@ def count_flops(impero_c):
     :returns: approximate flop count for the tree.
     """
     try:
-        return statement(impero_c.tree, set(impero_c.temporaries))
+        token = _active_indices.set(frozenset())
+        try:
+            return statement(impero_c.tree, set(impero_c.temporaries))
+        finally:
+            _active_indices.reset(token)
     except (ValueError, NotImplementedError):
         return 0
+
+
+_active_indices = ContextVar("flop_count_active_indices",
+                             default=frozenset())

gem/gem.py

@@ -311,10 +311,13 @@ def is_equal(self, other):
             return False
         if self.shape != other.shape:
             return False
+        if self.dtype != other.dtype:
+            return False
         return numpy.array_equal(self.array, other.array)
 
     def get_hash(self):
-        return hash((type(self), self.shape, tuple(self.array.flat)))
+        return hash((type(self), self.shape, self.dtype,
+                     tuple(self.array.flat)))
 
     @property
     def value(self):