diff --git a/finat/physically_mapped.py b/finat/physically_mapped.py index da88e7d6..8027732b 100644 --- a/finat/physically_mapped.py +++ b/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( diff --git a/gem/coffee.py b/gem/coffee.py index 607439fd..2cc4e1d9 100644 --- a/gem/coffee.py +++ b/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,33 +201,70 @@ 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) @@ -235,6 +272,7 @@ def _share_linear_maps( } 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) @@ -242,10 +280,13 @@ def _share_linear_maps( 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 diff --git a/gem/gem.py b/gem/gem.py index d3d8c3a0..0953d697 100644 --- a/gem/gem.py +++ b/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) diff --git a/gem/impero_utils.py b/gem/impero_utils.py index 31f9565b..ea91ab17 100644 --- a/gem/impero_utils.py +++ b/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)] diff --git a/gem/optimise.py b/gem/optimise.py index f009618d..8f2e6914 100644 --- a/gem/optimise.py +++ b/gem/optimise.py @@ -9,7 +9,7 @@ import numpy -from gem.cost import index_space_literal +from gem.cost import estimate_cost, index_space_literal, iteration_count from gem.utils import groupby from gem.node import (Memoizer, MemoizerArg, reuse_if_untouched, reuse_if_untouched_arg, traversal, traversal_children) @@ -487,9 +487,12 @@ def size(indices): reducible[subset] = frozenset(index for index in sum_indices if not incidence[index] & ~subset) + # Rank of each contraction index in the caller's ordering, to break ties. + position = {index: n for n, index in enumerate(sum_indices)} + def order(indices): """Sum out the widest index first, breaking ties reproducibly.""" - return tuple(sorted(indices, key=lambda i: (-extents[i], i.count))) + return tuple(sorted(indices, key=lambda i: (-extents[i], position[i]))) def reduce_indices(expression, free, indices): """Sum out indices, largest extent first, costing each reduction.""" @@ -767,24 +770,11 @@ def traverse_sum(expression, stop_at=None): def distribute_sum(expr: Node, predicate: Callable[[Node], bool]) -> list[Node]: """Distribute selected sums through products and contractions. - Parameters - ---------- - expr - GEM expression to distribute. - predicate - Predicate selecting the operations to distribute. - - Returns - ------- - list of Node - Additive terms after distribution. - - Notes - ----- - Memoization uses object identity. Structurally equal GEM nodes can have - deep expression trees, while distribution only needs to reuse actual DAG - nodes. + Memoisation is by object identity, which reuses each node of the DAG once. + :arg expr: GEM expression to distribute + :arg predicate: predicate selecting the operations to distribute + :returns: the additive terms after distribution """ results = {} active = {} @@ -826,18 +816,9 @@ def distribute_sum(expr: Node, predicate: Callable[[Node], bool]) -> list[Node]: def _is_linear_map(node: Node, linear_indices: frozenset) -> bool: """Is a node a linear map into one multilinear axis? - Parameters - ---------- - node - GEM expression node. - linear_indices - Free indices identifying the multilinear axes. - - Returns - ------- - bool - Whether the node is a sum over exactly one such axis. - + :arg node: GEM expression node + :arg linear_indices: free indices identifying the multilinear axes + :returns: whether the node is a sum over exactly one such axis """ return (isinstance(node, Sum) and len(linear_indices.intersection(node.free_indices)) == 1) @@ -848,23 +829,12 @@ def has_linear_maps( linear_indices: Iterable[Index]) -> bool: """Does a GEM DAG contain a finite element linear map? - Parameters - ---------- - expressions - Roots of a multilinear GEM expression DAG. - linear_indices - Free indices identifying the multilinear axes. - - Returns - ------- - bool - Whether preserving one-axis sums can change the factorisation. - - Notes - ----- - Answering this costs one traversal, where building the preserved - factorisation to compare it costs a whole pass of monomial collection. + One traversal answers this, so a caller can gate a second factorisation + on it. + :arg expressions: roots of a multilinear GEM expression DAG + :arg linear_indices: free indices identifying the multilinear axes + :returns: whether preserving one-axis sums can change the factorisation """ linear_indices = frozenset(linear_indices) return any(_is_linear_map(node, linear_indices) @@ -878,22 +848,13 @@ def preserve_linear_maps( """Expose multilinear terms and retain each one-axis linear map. A sum that depends on one linear index represents a linear map into an - argument tabulation. A sum that depends on several linear indices - separates multilinear form terms. This function distributes the latter + argument tabulation. A sum that depends on several linear indices + separates multilinear form terms. This function distributes the latter sums and returns the former sums as factors. - Parameters - ---------- - expression - Multilinear GEM expression. - linear_indices - Free indices identifying the linear axes. - - Returns - ------- - tuple - Additive terms and the linear-map factors that they contain. - + :arg expression: multilinear GEM expression + :arg linear_indices: free indices identifying the linear axes + :returns: the additive terms, and the linear-map factors they contain """ linear_indices = frozenset(linear_indices) @@ -958,17 +919,9 @@ def _delta_axes(node: Node, self: Memoizer) -> frozenset: 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. - + :arg index: index to inspect + :returns: ``(array, indices)`` when the index is a lookup into a + :class:`~.Literal` with a plain multiindex, otherwise ``None`` """ if not isinstance(index, VariableIndex): return None @@ -1232,3 +1185,143 @@ def aggressive_unroll(expression): expression, = unroll_indexsum((expression,), predicate=lambda index: True) expression, = remove_componenttensors((expression,)) return expression + + +def factorise_scalar_sums(expression: Node) -> Node: + """Factor common products from scalar sums when this lowers GEM cost. + + Scalar geometry and basis-transformation expressions are simplified below + the indexed contraction structure. Contractions are indivisible factors: + their bound indices cannot move through an enclosing sum. Sums carrying + free indices are left to the contraction planner, whose cost model + includes their iteration domains. + + :arg expression: root of a GEM expression + :returns: the expression with profitable common product factors extracted + """ + def choose(node): + if node.free_indices: + return node + summands = traverse_sum(node) + if len(summands) < 2: + return node + + factorisations = [] + for summand in summands: + _, factors = traverse_product( + summand, + stop_at=lambda factor: isinstance(factor, IndexSum), + ) + factorisations.append(factors) + + common = Counter(factorisations[0]) + for factors in factorisations[1:]: + common &= Counter(factors) + if not common: + return node + + common_factors = list(common.elements()) + remainders = [] + for factors in factorisations: + remaining_common = common.copy() + remaining = [] + for factor in factors: + if remaining_common[factor]: + remaining_common[factor] -= 1 + else: + remaining.append(factor) + remainders.append(make_product(remaining)) + + candidate = make_product( + (*common_factors, make_sum(remainders))) + if candidate.free_indices != node.free_indices: + return node + if estimate_cost((candidate,)) < estimate_cost((node,)): + return candidate + return node + + def visit(node, self): + node = reuse_if_untouched(node, self) + if isinstance(node, Sum): + node = choose(node) + return node + + return Memoizer(visit)(expression) + + +def _indirect_gathers(expression: Node) -> OrderedDict: + """Find the maps that an expression reads tables through. + + :arg expression: the root of a scalar GEM expression + :returns: an ordered mapping from each :class:`~.VariableIndex` that + indexes a table to the number of rows it selects from + """ + gathers = OrderedDict() + for node in traversal((expression,)): + if isinstance(node, Indexed): + aggregate, = node.children + for index, extent in zip(node.multiindex, aggregate.shape): + if isinstance(index, VariableIndex) and index.expression.free_indices: + gathers.setdefault(index, extent) + return gathers + + +def tabulate_indirect_contractions(expression: Node) -> Node: + """Evaluate a contraction one time for each row of its table. + + An expression can read a table through a map. Each argument of the map + selects one row, and two arguments can select the same row. A + contraction that reads the table through the map thus repeats work. + + This function evaluates the contraction for each row, then reads those + results through the map. The rewrite is correct only if the map does not + change with the contracted indices, and no other part of the contraction + uses the argument. It is faster only if the map has more arguments than + the table has rows. + + :arg expression: the root of a scalar GEM expression + :returns: the expression with each such contraction evaluated one time + for each row + """ + gathers = _indirect_gathers(expression) + if not gathers: + return expression + + def rename(node, self, substitution): + target, replacement = substitution + if isinstance(node, Indexed): + aggregate, = node.children + multiindex = tuple(replacement if index == target else index + for index in node.multiindex) + return Indexed(self(aggregate, substitution), multiindex) + return reuse_if_untouched_arg(node, self, substitution) + + def hoist(node): + body, = node.children + free = frozenset(body.free_indices) + contracted = frozenset(node.multiindex) + + candidates = [] + for gather, nrows in gathers.items(): + arguments = frozenset(gather.expression.free_indices) + if arguments <= free and arguments.isdisjoint(contracted): + saving = iteration_count(arguments) - nrows + if saving > 0: + candidates.append((saving, gather, arguments, nrows)) + + for _, gather, arguments, nrows in sorted(candidates, key=lambda c: -c[0]): + row = Index(extent=nrows) + per_row = MemoizerArg(rename)(body, (gather, row)) + # The body must reach the arguments only through this gather. + if arguments.isdisjoint(per_row.free_indices): + table = ComponentTensor(IndexSum(per_row, node.multiindex), (row,)) + return Indexed(table, (gather,)) + return node + + def visit(node, self): + node = reuse_if_untouched(node, self) + if isinstance(node, IndexSum): + node = hoist(node) + return node + + return Memoizer(visit)(expression) diff --git a/gem/refactorise.py b/gem/refactorise.py index b0b8a914..da9f1166 100644 --- a/gem/refactorise.py +++ b/gem/refactorise.py @@ -272,30 +272,15 @@ def _collect_monomial_sums( classifier: Callable[[Node], str]) -> list[MonomialSum]: """Collect monomial sums using the supplied node classifier. - Parameters - ---------- - expressions : iterable of Node - GEM expressions to refactorise. - classifier : callable - Function labelling each node as ``ATOMIC``, ``COMPOUND``, or - ``OTHER``. - - Returns - ------- - list of MonomialSum - Polynomial representations of the expressions. - - Raises - ------ - FactorisationError - If a compound expression cannot be expanded. - - Notes - ----- ``expressions`` must already have had its ComponentTensors removed, so - that a caller identifying nodes to classify sees the nodes this - collector will visit. - + that a caller identifying nodes to classify sees the nodes this collector + will visit. + + :arg expressions: GEM expressions to refactorise + :arg classifier: a function labelling each node as ``ATOMIC``, + ``COMPOUND``, or ``OTHER`` + :returns: list of :py:class:`MonomialSum`s + :raises FactorisationError: a compound expression cannot be expanded. """ # Get ListTensors out of the way @@ -322,36 +307,25 @@ def collect_monomials( expressions: Iterable[Node], classifier: Callable[[Node], str], linear_indices: Iterable[Index] = ()) -> list[MonomialSum]: - """Collect structure-preserving sum-of-products representations. - - Parameters - ---------- - expressions - GEM expressions to refactorise. - classifier - Function that labels GEM nodes for polynomial collection. - linear_indices - Free indices identifying the multilinear axes. Sums depending on - exactly one such axis represent linear maps and remain atomic. - - Returns - ------- - list of MonomialSum - One polynomial representation for each input expression. - - Notes - ----- - A one-axis sum is a finite element linear operand: examples include a - sparse basis transformation and a tensor-product tabulation factor. - Preserving it keeps domain-specific basis structure available for - contraction optimisation. This does not make the map opaque: its GEM - expression remains available for scalar simplification and code motion. - - Sums involving several linear axes separate form monomials and are - distributed. COFFEE subsequently chooses scalar factorisations across - the resulting operands. Keeping these two stages distinct avoids a - Cartesian expansion of basis-map entries before loop placement. + """Refactorises expressions into a sum-of-products form, using + distributivity rules (i.e. a*(b + c) -> a*b + a*c). Expansion + proceeds until all "compound" expressions are broken up. + + A sum over exactly one linear index is a finite element linear map, such + as a sparse basis transformation or a tensor-product tabulation factor, + and stays atomic. A sum over several linear indices separates form + monomials, and is distributed. + :arg expressions: GEM expressions to refactorise + :arg classifier: a function that can classify any GEM expression + as ``ATOMIC``, ``COMPOUND``, or ``OTHER``. This + classification drives the factorisation. + :arg linear_indices: free indices identifying the multilinear axes + + :returns: list of :py:class:`MonomialSum`s + + :raises FactorisationError: Failed to break up some "compound" + expressions with expansion. """ # Remove ComponentTensors here, not in the collector. Removing them # rebuilds nodes, and the maps found below must be the very nodes that diff --git a/test/gem/test_simplify.py b/test/gem/test_simplify.py index 94aeac3d..4343b1a4 100644 --- a/test/gem/test_simplify.py +++ b/test/gem/test_simplify.py @@ -100,3 +100,41 @@ def test_flatten_indexsum(A): result = gem.IndexSum(gem.IndexSum(Aij, (i,)), (j,)) expected = gem.IndexSum(Aij, (i, j)) assert result == expected + + +def test_rename_index_under_variable_index(): + """Renaming a bound index must reach the lookup of an indirect gather.""" + values = gem.Literal(numpy.array([10.0, 20.0, 30.0])) + lookup = gem.Literal(numpy.array([2, 0, 1], dtype=gem.uint_type), + dtype=gem.uint_type) + + k, kp = gem.indices(2) + gather = gem.Indexed( + values, (gem.VariableIndex(gem.Indexed(lookup, (k,))),)) + assert k in gather.free_indices + + # This is how make_renamer separates two sums that bind the same index. + renamed = gem.Indexed(gem.ComponentTensor(gather, (k,)), (kp,)) + assert renamed.free_indices == (kp,) + + +def test_product_of_sums_over_one_index(): + """Two sums that bind the same index expand to a double sum, not one.""" + from gem.interpreter import evaluate + from gem.optimise import make_rename_map, make_renamer + + values = gem.Literal(numpy.array([1.0, 2.0, 3.0])) + lookup = gem.Literal(numpy.array([2, 0, 1], dtype=gem.uint_type), + dtype=gem.uint_type) + k, = gem.indices(1) + gather = gem.Indexed( + values, (gem.VariableIndex(gem.Indexed(lookup, (k,))),)) + + renamer = make_renamer(make_rename_map()) + (k0,), first = renamer((k,)) + (k1,), second = renamer((k,)) + assert k0 != k1 + + square = gem.IndexSum(gem.Product(first(gather), second(gather)), (k0, k1)) + result, = evaluate([square]) + assert numpy.isclose(result.arr, (1.0 + 2.0 + 3.0) ** 2) diff --git a/test/gem/test_sum_factorise.py b/test/gem/test_sum_factorise.py index 5dca72b9..33ac1991 100644 --- a/test/gem/test_sum_factorise.py +++ b/test/gem/test_sum_factorise.py @@ -7,7 +7,7 @@ import gem from gem.interpreter import evaluate from gem import cost, optimise -from gem.node import traversal +from gem.node import post_traversal, traversal from gem.optimise import sum_factorise from gem.coffee import optimise_monomial_sum from gem.refactorise import (ATOMIC, COMPOUND, OTHER, @@ -233,6 +233,34 @@ def test_preserved_linear_map_is_shared(): assert numpy.allclose(result.arr, expected) +def test_linear_map_sharing_normalises_order_and_scale() -> None: + i, j = gem.Index(extent=3), gem.Index(extent=3) + a_values = numpy.array([1.0, 2.0, 3.0]) + b_values = numpy.array([4.0, 5.0, 6.0]) + a = gem.Literal(a_values) + b = gem.Literal(b_values) + left_terms = (gem.Indexed(a, (i,)), gem.Indexed(b, (i,))) + right_terms = (gem.Indexed(b, (j,)), gem.Indexed(a, (j,))) + left_map = optimise.make_sum(left_terms * 2) + right_map = optimise.make_sum(right_terms * 2) + expression = gem.Product(left_map, right_map) + + arguments = (i, j) + classifier = partial(_classify, frozenset(arguments)) + preserved, = collect_monomials( + [expression], classifier, arguments) + optimised = optimise_monomial_sum(preserved, arguments) + tensors = [node for node in traversal((optimised,)) + if isinstance(node, gem.ComponentTensor)] + assert len(tensors) == 1 + body, = tensors[0].children + assert len(optimise.traverse_sum(body)) == 2 + + result, = evaluate([gem.ComponentTensor(optimised, arguments)]) + expected = 4 * numpy.outer(a_values + b_values, a_values + b_values) + assert numpy.allclose(result.arr, expected) + + def test_expanded_and_preserved_agree(): arguments, expanded, expected = monomial_sum(linear_indices=False) expression = optimise_monomial_sum(expanded, arguments) @@ -272,3 +300,55 @@ def test_contraction_counts_index_absent_from_factors() -> None: result, = evaluate([expression]) assert result.arr == 2 * numpy.arange(3.0).sum() + + +def test_common_factor_is_extracted_from_scalar_sum() -> None: + common = gem.Variable("common", ()) + left = gem.Variable("left", ()) + right = gem.Variable("right", ()) + expression = gem.Sum( + gem.Product(common, left), gem.Product(common, right)) + + expression = optimise.factorise_scalar_sums(expression) + + _, factors = optimise.traverse_product(expression) + assert common in factors + result, = evaluate([expression], { + common: numpy.asarray(2.0), + left: numpy.asarray(3.0), + right: numpy.asarray(4.0), + }) + assert result.arr == 14 + + +def test_scalar_factorisation_leaves_indexed_sum_to_planner() -> None: + i = gem.Index(extent=3) + common = gem.Variable("common", ()) + left = gem.Indexed(gem.Literal(numpy.arange(3.0)), (i,)) + right = gem.Indexed(gem.Literal(numpy.arange(3.0, 6.0)), (i,)) + expression = gem.Sum( + gem.Product(common, left), gem.Product(common, right)) + + assert optimise.factorise_scalar_sums(expression) is expression + + +def test_contraction_plan_follows_the_callers_index_order() -> None: + """A plan depends only on the caller's ordering of the sum indices.""" + def plan(reverse_creation): + if reverse_creation: + r = gem.Index(extent=4) + q = gem.Index(extent=4) + else: + q = gem.Index(extent=4) + r = gem.Index(extent=4) + # Both factors carry both indices, so the two reduce at the same + # subset and the tie decides their order within one IndexSum. + factors = [gem.Indexed(gem.Variable(name, (4, 4)), (q, r)) + for name in ("A", "B")] + expression = optimise.sum_factorise([q, r], factors) + role = {q: "q", r: "r"} + return [tuple(role[i] for i in node.multiindex) + for node in post_traversal((expression,)) + if isinstance(node, gem.IndexSum)] + + assert plan(False) == plan(True) == [("q", "r")]