Padded basis transformation for better codegen

finat/physically_mapped.py

@@ -1,12 +1,17 @@
 from abc import ABCMeta, abstractmethod
-from collections.abc import Mapping
+from collections.abc import Iterable, Mapping
+from functools import cached_property
 
 import gem
 import numpy
 
 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,44 +21,118 @@ 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)
+        self._value_dim = M.shape[1]
+
+        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)]
+    @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.
+
+        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.
 
-        result = gem.ListTensor(exprs)
-        result, = gem.optimise.unroll_indexsum((result,), lambda index: True)
-        # result = gem.optimise.aggressive_unroll(self.M @ table)
-        return result
+        """
+        k = self._row_index
+        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 matmul(self, table: gem.Node) -> gem.Node:
+        """Apply the basis transformation to a reference tabulation."""
+        r = gem.Index(extent=self._space_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,))
+        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):
@@ -195,10 +274,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):

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.
 

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
 

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)

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