diff --git a/gem/coffee.py b/gem/coffee.py index e9ae2011..af9f9c84 100644 --- a/gem/coffee.py +++ b/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: diff --git a/gem/flop_count.py b/gem/flop_count.py index b9595e81..96d34c2a 100644 --- a/gem/flop_count.py +++ b/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()) diff --git a/gem/gem.py b/gem/gem.py index 845660f5..00681496 100644 --- a/gem/gem.py +++ b/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): diff --git a/gem/optimise.py b/gem/optimise.py index 906620bf..69dd178b 100644 --- a/gem/optimise.py +++ b/gem/optimise.py @@ -1,7 +1,9 @@ """A set of routines implementing various transformations on GEM expressions.""" +import math from collections import Counter, OrderedDict, defaultdict +from collections.abc import Callable, Iterable from functools import singledispatch, partial from itertools import combinations, permutations, zip_longest from numbers import Integral @@ -15,6 +17,7 @@ Product, Sum, Comparison, Conditional, Division, Index, VariableIndex, Indexed, FlexiblyIndexed, IndexSum, ComponentTensor, ListTensor, Delta, + MathFunction, MinValue, MaxValue, Power, Inverse, Solve, partial_indexed, one) @@ -723,6 +726,274 @@ def traverse_sum(expression, stop_at=None): return result +def _iteration_count(indices: Iterable[Index]) -> int: + """Count the points of a rectangular iteration space. + + Parameters + ---------- + indices + Indices an operation depends on. + + Returns + ------- + int + Number of executions of the operation. + + """ + return int(numpy.prod([index.extent for index in indices], dtype=int)) + + +def _operation_count(node: Node) -> int: + """Estimate the scalar operations performed by one GEM node. + + Parameters + ---------- + node + Scalar expression node. + + Returns + ------- + int + Operations over the node's complete iteration domain. + + Notes + ----- + Negation folds into the operation that consumes it, so a product with + ``-1`` costs nothing. + + """ + domain = _iteration_count(node.free_indices) + if isinstance(node, Product): + if any(isinstance(child, Literal) and not child.shape + and child.value == -1 for child in node.children): + return 0 + return domain + if isinstance(node, (Sum, Division, MathFunction, MinValue, MaxValue)): + return domain + if isinstance(node, Power): + _, exponent = node.children + if isinstance(exponent, Literal) and not exponent.shape: + value = exponent.value + if value > 0 and value == math.floor(value): + return math.ceil(math.log2(value)) * domain + return 5 * domain + if isinstance(node, IndexSum): + body, = node.children + return _iteration_count(body.free_indices) + if isinstance(node, Inverse): + n, _ = node.shape + return 2 * n ** 3 + if isinstance(node, Solve): + n, m = node.shape + return 2 * n * m + 2 * n ** 3 + return 0 + + +def has_arithmetic(expressions: Iterable[Node]) -> bool: + """Does a GEM DAG perform any scalar arithmetic? + + Parameters + ---------- + expressions + Roots of a scalar GEM expression DAG. + + Returns + ------- + bool + Whether any node costs arithmetic to evaluate. + + Notes + ----- + Materialising a tabulation reference buys no arithmetic, so sharing one + only adds storage. + + """ + return any(map(_operation_count, traversal(tuple(expressions)))) + + +def estimate_cost(expressions: Iterable[Node]) -> tuple[int, int, int, int]: + """Estimate arithmetic work and contraction storage for a GEM DAG. + + Each structurally shared operation counts once over the domain its free + indices induce. An :class:`~gem.gem.IndexSum` contributes one + accumulation per point of its body domain. Storage counts the result + domains of contractions, the intermediates that scheduling exposes. + + Parameters + ---------- + expressions + Roots of a scalar GEM expression DAG. + + Returns + ------- + tuple of int + Operation count, total contraction storage, largest contraction, and + expression-node count. Comparing the tuples lexicographically ranks + arithmetic first and breaks ties by storage. + + """ + nodes = tuple(traversal(tuple(expressions))) + sizes = [_iteration_count(node.free_indices) + for node in nodes if isinstance(node, IndexSum)] + return (sum(map(_operation_count, nodes)), + sum(sizes), + max(sizes, default=0), + len(nodes)) + + +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. + + """ + results = {} + active = {} + stack = [(expr, False)] + while stack: + node, expanded = stack.pop() + key = id(node) + if key in results: + continue + if not expanded: + stack.append((node, True)) + stack.extend((c, False) for c in node.children) + continue + active[key] = predicate(node) or any( + active[id(child)] for child in node.children) + if active[key] and isinstance(node, (Sum, IndexSum, Product)): + if isinstance(node, Sum): + results[key] = [ + term + for child in node.children + for term in results[id(child)]] + elif isinstance(node, IndexSum): + body, = node.children + results[key] = [ + IndexSum(term, tuple( + index for index in node.multiindex + if index in term.free_indices)) + for term in results[id(body)]] + else: # Product + a, b = node.children + ta, tb = results[id(a)], results[id(b)] + results[key] = [node] if len(ta) == 1 and len(tb) == 1 \ + else [Product(x, y) for x in ta for y in tb] + else: + results[key] = [node] + return results[id(expr)] + + +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. + + """ + return (isinstance(node, Sum) + and len(linear_indices.intersection(node.free_indices)) == 1) + + +def has_linear_maps( + expressions: Iterable[Node], + 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. + + """ + linear_indices = frozenset(linear_indices) + return any(_is_linear_map(node, linear_indices) + for node in traversal(tuple(expressions))) + + +def preserve_linear_maps( + expression: Node, + linear_indices: Iterable[Index]) -> tuple[ + tuple[Node, ...], tuple[Node, ...]]: + """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 + 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. + + """ + linear_indices = frozenset(linear_indices) + + def multilinear_sum(node: Node) -> bool: + return isinstance(node, Sum) and len( + linear_indices.intersection(node.free_indices)) > 1 + + if not has_linear_maps((expression,), linear_indices): + return (expression,), () + + terms = tuple(_distribute_sum(expression, multilinear_sum)) + groups = OrderedDict() + for term in terms: + _, factors = traverse_product(term) + for factor in factors: + if _is_linear_map(factor, linear_indices): + groups.setdefault(factor) + + if not groups: + return (expression,), () + return terms, tuple(groups) + + def repeated_contractions(expression): """Find the contractions that occur more than once in a product tree. diff --git a/gem/refactorise.py b/gem/refactorise.py index 2ca6e4cc..b0b8a914 100644 --- a/gem/refactorise.py +++ b/gem/refactorise.py @@ -2,16 +2,17 @@ refactorisation.""" from collections import Counter, OrderedDict, defaultdict, namedtuple +from collections.abc import Callable, Iterable from functools import singledispatch from itertools import product from sys import intern from gem.node import Memoizer, traversal -from gem.gem import (Node, Conditional, Zero, Product, Sum, Indexed, +from gem.gem import (Node, Conditional, Index, Zero, Product, Sum, Indexed, ListTensor, one, MathFunction) -from gem.optimise import (remove_componenttensors, sum_factorise, - traverse_product, traverse_sum, unroll_indexsum, - make_rename_map, make_renamer) +from gem.optimise import (preserve_linear_maps, remove_componenttensors, + sum_factorise, traverse_product, traverse_sum, + unroll_indexsum, make_rename_map, make_renamer) # Refactorisation labels @@ -266,23 +267,36 @@ def _collect_monomials_conditional(expression, self): return result -def collect_monomials(expressions, classifier): - """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. +def _collect_monomial_sums( + expressions: Iterable[Node], + classifier: Callable[[Node], str]) -> list[MonomialSum]: + """Collect monomial sums using the supplied node classifier. - :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. + Parameters + ---------- + expressions : iterable of Node + GEM expressions to refactorise. + classifier : callable + Function labelling each node as ``ATOMIC``, ``COMPOUND``, or + ``OTHER``. - :returns: list of :py:class:`MonomialSum`s + 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. - :raises FactorisationError: Failed to break up some "compound" - expressions with expansion. """ - # Get ComponentTensors out of the way - expressions = remove_componenttensors(expressions) # Get ListTensors out of the way must_unroll = [] # indices to unroll @@ -302,3 +316,64 @@ def collect_monomials(expressions, classifier): mapper.classifier = classifier mapper.rename_map = make_rename_map() return list(map(mapper, expressions)) + + +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. + + """ + # Remove ComponentTensors here, not in the collector. Removing them + # rebuilds nodes, and the maps found below must be the very nodes that + # the collector classifies. + expressions = remove_componenttensors(tuple(expressions)) + linear_indices = tuple(linear_indices) + if not linear_indices: + return _collect_monomial_sums(expressions, classifier) + + result = [] + for expression in expressions: + terms, linear_maps = preserve_linear_maps(expression, linear_indices) + if not linear_maps: + monomial_sum, = _collect_monomial_sums((expression,), classifier) + result.append(monomial_sum) + continue + + map_ids = frozenset(map(id, linear_maps)) + + def preserve_map(node: Node, map_ids=map_ids) -> str: + return ATOMIC if id(node) in map_ids else classifier(node) + + result.append(MonomialSum.sum( + *_collect_monomial_sums(terms, preserve_map))) + return result diff --git a/test/gem/test_sum_factorise.py b/test/gem/test_sum_factorise.py index 602d421e..27b855be 100644 --- a/test/gem/test_sum_factorise.py +++ b/test/gem/test_sum_factorise.py @@ -1,3 +1,4 @@ +from functools import partial, reduce from itertools import chain, combinations, islice import numpy @@ -8,6 +9,9 @@ from gem import optimise from gem.node import traversal from gem.optimise import sum_factorise +from gem.coffee import optimise_monomial_sum +from gem.refactorise import (ATOMIC, COMPOUND, OTHER, + collect_monomials) def contraction(nfactors, ndims, extent=2): @@ -155,3 +159,99 @@ def test_contractions_joined_by_a_shared_index(): evaluations = [numpy.einsum("ijkp,ijk->p", table, coefficient) for table, coefficient in zip(tables, coefficients)] assert numpy.allclose(result.arr, evaluations[0].dot(evaluations[1])) + + +def laplacian(ndofs=4, ndims=2): + """Build a Laplacian element tensor from a mapped gradient table. + + The physical gradient of each basis function is the reference gradient + mapped by the inverse Jacobian. Test and trial functions apply the same + map, over their own argument index. + """ + numpy.random.seed(0) + i = gem.Index(extent=ndofs) + j = gem.Index(extent=ndofs) + k = gem.Index(extent=ndims) + reference = gem.Literal(numpy.random.rand(ndofs, ndims)) + jacobian = gem.Literal(numpy.random.rand(ndims, ndims)) + + def gradient(argument): + # The pullback is a contraction over the topological dimension, + # which reaches factorisation already unrolled into a sum. + return reduce(gem.Sum, [ + gem.Product(gem.Indexed(reference, (argument, l)), + gem.Indexed(jacobian, (l, k))) + for l in range(ndims)]) + + expression = gem.IndexSum( + gem.Product(gradient(i), gradient(j)), (k,)) + expected = numpy.einsum( + "il,lk,jm,mk->ij", reference.array, jacobian.array, + reference.array, jacobian.array) + return (i, j), expression, expected + + +def monomial_sum(linear_indices): + """Collect the Laplacian into monomials, with or without preservation.""" + arguments, expression, expected = laplacian() + classifier = partial(_classify, frozenset(arguments)) + result, = collect_monomials( + [expression], classifier, + arguments if linear_indices else ()) + return arguments, result, expected + + +def _classify(arguments, expression): + shared = arguments.intersection(expression.free_indices) + if not shared: + return OTHER + if len(shared) == 1 and isinstance(expression, gem.Indexed): + return ATOMIC + return COMPOUND + + +def test_linear_map_is_preserved(): + # Distributing the map expands it into the product of its entries, so + # preserving it leaves strictly fewer monomials to factorise. + _, preserved, _ = monomial_sum(linear_indices=True) + _, expanded, _ = monomial_sum(linear_indices=False) + assert len(list(preserved)) < len(list(expanded)) + assert len(list(preserved)) == 1 + + +def test_preserved_linear_map_is_shared(): + # Test and trial apply the same map over different indices. COFFEE + # materialises it once, so one tensor carries both. + arguments, preserved, expected = monomial_sum(linear_indices=True) + expression = optimise_monomial_sum(preserved, arguments) + tensors = [node for node in traversal((expression,)) + if isinstance(node, gem.ComponentTensor)] + assert len(tensors) == 1 + + i, j = arguments + result, = evaluate([gem.ComponentTensor(expression, (i, j))]) + 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) + result, = evaluate([gem.ComponentTensor(expression, arguments)]) + assert numpy.allclose(result.arr, expected) + + +def test_has_linear_maps_detects_preservable_maps(): + arguments, expression, _ = laplacian() + assert optimise.has_linear_maps([expression], arguments) + # With no argument axis declared, nothing is a linear map. + assert not optimise.has_linear_maps([expression], ()) + + +def test_estimate_cost_counts_the_contraction(): + _, expression, _ = laplacian(ndofs=4, ndims=2) + flops, storage, largest, nodes = optimise.estimate_cost([expression]) + # Two mapped gradients over (argument, k, l) and their contraction + # over k, all counted over their own domains. + assert flops > 0 + assert storage >= largest > 0 + assert nodes > 0