Optimise generic GEM sum factorisation

finat/physically_mapped.py

@@ -1,12 +1,40 @@
 from abc import ABCMeta, abstractmethod
-from collections.abc import Mapping
+from collections.abc import Iterable, Mapping
+from numbers import Number
 
 import gem
 import numpy
 
 from finat.citations import cite
 
 
+zero = gem.Zero()
+one = gem.Literal(1.0)
+
+
+def _as_basis_entry(value: object) -> gem.Node:
+    """Normalize a basis-transformation matrix entry.
+
+    Parameters
+    ----------
+    value
+        Scalar numerical or GEM matrix entry.
+
+    Returns
+    -------
+    gem.Node
+        Scalar GEM entry with numerical zero and one represented
+        symbolically.
+
+    """
+    if isinstance(value, Number):
+        if value == 0:
+            return zero
+        if value == 1:
+            return one
+    return gem.as_gem(value)
+
+
 class NeedsCoordinateMappingElement(metaclass=ABCMeta):
     """Abstract class for elements that require physical information
     either to map or construct their basis functions."""
@@ -16,38 +44,81 @@ 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
+    -----
+    The transformation is stored in compressed sparse row form.  One
+    constant table holds the column of each nonzero, and a second holds its
+    value.  Rows shorter than the longest are padded with zero values, which
+    contribute nothing to the sum.
+
+    Every row therefore contracts over the same number of entries.  The
+    basis axis stays one loop, so the transformation reaches the quadrature
+    contraction as a linear map over the reference tabulation.
 
-    :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)
+
+        nonzero_rows = []
+        for source_row in self.indices:
+            row = []
+            for column in range(M.shape[1]):
+                value = _as_basis_entry(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)
+        self._values = gem.ListTensor(data)
         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:
+        """Transform one reference tabulation.
+
+        Parameters
+        ----------
+        table
+            Reference tabulation with the basis axis first.
 
-        result = gem.ListTensor(exprs)
-        result, = gem.optimise.unroll_indexsum((result,), lambda index: True)
-        # result = gem.optimise.aggressive_unroll(self.M @ table)
-        return result
+        Returns
+        -------
+        gem.Node
+            Tabulation whose first axis is the transformed basis axis.
+
+        """
+        tail = gem.indices(len(table.shape) - 1)
+        row = gem.Index(extent=len(self.indices))
+        entry = gem.Index(extent=self._width)
+        column = gem.VariableIndex(gem.Indexed(self._columns, (row, entry)))
+        basis = gem.Product(gem.Indexed(self._values, (row, entry)),
+                            gem.Indexed(table, (column, *tail)))
+        mapped = gem.IndexSum(basis, (entry,))
+        return gem.ComponentTensor(mapped, (row, *tail))
 
     def __getitem__(self, alpha):
         try:
@@ -195,10 +266,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):