diff --git a/FIAT/bernstein.py b/FIAT/bernstein.py index f960c64f..41ef4ccd 100644 --- a/FIAT/bernstein.py +++ b/FIAT/bernstein.py @@ -11,9 +11,134 @@ from FIAT.finite_element import FiniteElement from FIAT.dual_set import DualSet -from FIAT.polynomial_set import mis +from FIAT.expansions import ExpansionSet +from FIAT.polynomial_set import PolynomialSet, mis from FIAT.pointwise_dual import compute_pointwise_dual -from FIAT.reference_element import make_lattice +from FIAT.reference_element import make_lattice, multiindex_equal + + +def _bernstein_factors(n, eta): + """Tabulate univariate Bernstein factors for a collapsed simplex axis. + + Parameters + ---------- + n : int + The total polynomial degree. + eta : numpy.ndarray + Points on ``[-1, 1]``. + + Returns + ------- + tuple[numpy.ndarray, ...] + Value, degree-lowered, and degree-and-index-lowered tables. + + """ + eta = numpy.asarray(eta) + z = 0.5 * (1.0 + eta) + shape = (n + 1, n + 1, len(eta)) + values = numpy.zeros(shape, dtype=z.dtype) + lower = numpy.zeros(shape, dtype=z.dtype) + shifted = numpy.zeros(shape, dtype=z.dtype) + for m in range(n + 1): + degree = n - m + for i in range(degree + 1): + values[m, i] = math.comb(degree, i) * z**i * (1.0 - z)**(degree-i) + for i in range(degree): + lower[m, i] = math.comb(degree - 1, i) * z**i * (1.0 - z)**(degree-1-i) + shifted[m, i+1] = lower[m, i] + return values, lower, shifted + + +class BernsteinExpansionSet(ExpansionSet): + """Bernstein polynomial expansion set on a simplex.""" + + def __init__(self, ref_el): + if not ref_el.is_simplex(): + raise ValueError("Bernstein expansion sets require a simplex") + super().__init__(ref_el, scale=1.0) + + def get_duffy_permutation(self, n): + """Map Duffy lattice positions to Bernstein expansion members.""" + dim = self.ref_el.get_spatial_dimension() + return numpy.arange(math.comb(n + dim, dim)) + + def _tabulate_on_cell(self, n, pts, order=0, cell=0, direction=None): + """Tabulate the expansion set and its derivatives on one cell.""" + if direction is not None: + raise NotImplementedError("directional Bernstein tabulation is not implemented") + + ref_el = self.ref_el + dim = ref_el.get_spatial_dimension() + topology = ref_el.get_topology() + vertices = ref_el.get_vertices_of_subcomplex(topology[dim][cell]) + + B2R = numpy.vstack([numpy.asarray(vertices).T, numpy.ones(len(vertices))]) + R2B = numpy.linalg.inv(B2R) + points = numpy.asarray(pts) + B = numpy.concatenate([points, numpy.ones((*points.shape[:-1], 1))], + axis=-1).dot(R2B.T) + + raw_result = { + (derivative, i): vec + for i, alpha in enumerate(multiindex_equal(dim+1, n)) + for o in range(order + 1) + for derivative, vec in bernstein_Dx( + B, alpha, o, R2B + ).items() + } + num_members = math.comb(n + dim, dim) + dtype = numpy.array(list(raw_result.values())).dtype + result = { + alpha: numpy.zeros((num_members, *points.shape[:-1]), dtype=dtype) + for o in range(order + 1) + for alpha in mis(dim, o) + } + for (alpha, i), vec in raw_result.items(): + result[alpha][i] = vec + return result + + def tabulate_duffy(self, n, eta_pts, order=0, cell=0): + """Tabulate Bernstein polynomials in separable collapsed form. + + The raw lattice multi-index ``j`` represents the barycentric + exponent tuple ``(n - sum(j), *reversed(j))``. Reversing the + collapsed axes gives the product from Ainsworth et al., + ``prod_t B[j_t, n - sum(j[:t])]``. + """ + if order > 1: + raise NotImplementedError("tabulate_duffy is limited to first derivatives") + + dim = self.ref_el.get_spatial_dimension() + assert len(eta_pts) == dim + A, _ = self.affine_mappings[cell] + axes = tuple(reversed(range(dim))) + tables = tuple(_bernstein_factors(n, eta_pts[axis]) for axis in axes) + values = tuple(table[0] for table in tables) + result = {(0,) * dim: [(1.0, tuple(zip(axes, values)))]} + + if order: + lower = tuple(table[1] for table in tables) + # On the default (-1, 1) simplex, + # d/dxi_l B_alpha^n = n/2 * (B_{alpha-e_{l+1}}^{n-1} + # - B_{alpha-e_0}^{n-1}). + # Both degree-lowered terms retain the separable product form. + for k in range(dim): + terms = [] + for ell in range(dim): + coeff = 0.5 * n * A[ell, k] + if coeff == 0.0: + continue + s = dim - 1 - ell + shifted = tuple(table[1] if t < s else + table[2] if t == s else table[0] + for t, table in enumerate(tables)) + terms.extend(((coeff, tuple(zip(axes, shifted))), + (-coeff, tuple(zip(axes, lower))))) + if not terms: + terms.append((0.0, tuple(zip(axes, values)))) + alpha = tuple(int(i == k) for i in range(dim)) + result[alpha] = terms + return result class BernsteinDualSet(DualSet): @@ -33,7 +158,7 @@ def __init__(self, ref_el, degree): # Generate triangular barycentric indices dim = ref_el.get_spatial_dimension() - kss = mis(dim + 1, degree) + kss = multiindex_equal(dim+1, degree) # Fill data structures nodes = [] @@ -55,6 +180,13 @@ def __init__(self, ref_el, degree): dual = BernsteinDualSet(ref_el, degree) k = 0 # 0-form super().__init__(ref_el, dual, degree, k) + + expansion_set = BernsteinExpansionSet(ref_el) + size = math.comb(degree + ref_el.get_spatial_dimension(), + ref_el.get_spatial_dimension()) + coeffs = numpy.eye(size) + self.poly_set = PolynomialSet(ref_el, degree, degree, expansion_set, coeffs) + pts = make_lattice(ref_el.vertices, degree, variant="gll") newdual = compute_pointwise_dual(self, pts) self.dual = newdual @@ -63,6 +195,10 @@ def degree(self): """The degree of the polynomial space.""" return self.get_order() + def get_nodal_basis(self): + """Return the Bernstein basis encoded as a polynomial set.""" + return self.poly_set + def value_shape(self): """The value shape of the finite element functions.""" return () @@ -90,30 +226,8 @@ def tabulate(self, order, points, entity=None): points = numpy.asarray(points) cell_points = entity_transform(points) - # Construct Cartesian to Barycentric coordinate mapping - vs = numpy.asarray(ref_el.get_vertices()) - B2R = numpy.vstack([vs.T, numpy.ones(len(vs))]) - R2B = numpy.linalg.inv(B2R) - - B = numpy.concatenate([cell_points, numpy.ones((*cell_points.shape[:-1], 1)) - ], axis=-1).dot(R2B.T) - - # Evaluate everything - deg = self.degree() - raw_result = {(alpha, i): vec - for i, ks in enumerate(mis(dim + 1, deg)) - for o in range(order + 1) - for alpha, vec in bernstein_Dx(B, ks, o, R2B).items()} - - # Rearrange result - space_dim = self.space_dimension() - dtype = numpy.array(list(raw_result.values())).dtype - result = {alpha: numpy.zeros((space_dim, *points.shape[:-1]), dtype=dtype) - for o in range(order + 1) - for alpha in mis(dim, o)} - for (alpha, i), vec in raw_result.items(): - result[alpha][i] = vec - return result + return self.poly_set.get_expansion_set()._tabulate( + self.degree(), cell_points, order=order) def bernstein_db(points, ks, alpha=None): @@ -140,9 +254,9 @@ def bernstein_db(points, ks, alpha=None): ls = ks - alpha if any(k < 0 for k in ls): - return numpy.zeros(len(points)) + return numpy.zeros(points.shape[:-1]) elif all(k == 0 for k in ls): - return numpy.ones(len(points)) + return numpy.ones(points.shape[:-1]) else: # Calculate coefficient coeff = math.factorial(ks.sum()) diff --git a/FIAT/expansions.py b/FIAT/expansions.py index 33f3ed00..e7a3a2b2 100644 --- a/FIAT/expansions.py +++ b/FIAT/expansions.py @@ -24,14 +24,21 @@ def morton_index3(p, q=0, r=0): def jrc(a, b, n): """Jacobi recurrence coefficients""" an = (2*n+1+a+b)*(2*n+2+a+b) / (2*(n+1)*(n+1+a+b)) - bn = (a+b)*(a-b)*(2*n+1+a+b) / (2*(n+1)*(n+1+a+b)*(2*n+a+b)) - cn = (n+a)*(n+b)*(2*n+2+a+b) / ((n+1)*(n+1+a+b)*(2*n+a+b)) + if n == 0: + bn = 0.5 * (a - b) + cn = 0.0 + else: + bn = (a+b)*(a-b)*(2*n+1+a+b) / (2*(n+1)*(n+1+a+b)*(2*n+a+b)) + cn = (n+a)*(n+b)*(2*n+2+a+b) / ((n+1)*(n+1+a+b)*(2*n+a+b)) return an, bn, cn def integrated_jrc(a, b, n): """Integrated Jacobi recurrence coefficients""" - if n == 1: + if n == 0: + an = bn = -0.5 + cn = 0.0 + elif n == 1: an = (a + b + 2) / 2 bn = (a - 3*b - 2) / 2 cn = 0.0 @@ -40,6 +47,20 @@ def integrated_jrc(a, b, n): return an, bn, cn +def dubiner_jacobi_weights(axis, m, variant): + """Return the Jacobi weights for a given Dubiner axis.""" + if variant == "bubble": + alpha = 2 * m + beta = 0 + else: + alpha = 2 * m + axis + beta = 0 + if variant == "dual": + alpha += 1 + axis + beta = 1 + return alpha, beta + + def pad_coordinates(ref_pts, embedded_dim): """Pad reference coordinates by appending -1.0.""" return tuple(ref_pts) + (-1.0, )*(embedded_dim - len(ref_pts)) @@ -51,6 +72,19 @@ def pad_jacobian(A, embedded_dim): return tuple(row[..., None] for row in A) +def dubiner_norm2(d: int, m: int, i: int, variant: str | None) -> float: + """Return the squared normalization for one Dubiner principal function.""" + if variant is not None: + shift = 1 if variant == "dual" else 0 + p = i + shift + alpha = 2 * (m + d * shift) - 1 + norm2 = (0.5 + d) / d + if p > 0 and p + alpha > 0: + norm2 *= (p + alpha) * (2*p + alpha) / p + return norm2 + return (2*(m + i) + d) / d + + def jacobi_factors(x, y, z, dx, dy, dz): fb = 0.5 * (y + z) fa = x + (fb + 1.0) @@ -137,13 +171,7 @@ def _product_derivative(factor: numpy.ndarray, return result -def dubiner_recurrence(dim: int, - n: int, - order: int, - ref_pts: numpy.ndarray, - Jinv: numpy.ndarray, - scale: float, - variant: str | None = None) -> list[numpy.ndarray]: +def dubiner_recurrence(dim, n, order, ref_pts, Jinv, scale, variant=None): """Tabulate a Dubiner expansion set using the recurrence from (Kirby 2010). Parameters @@ -195,29 +223,20 @@ def dubiner_recurrence(dim: int, if dim > 3 or dim < 0: raise ValueError("Invalid number of spatial dimensions") - beta = 1 if variant == "dual" else 0 coefficients = integrated_jrc if variant == "bubble" else jrc X = pad_coordinates(ref_pts, pad_dim) idx = (lambda p: p, morton_index2, morton_index3)[dim-1] - for codim in range(dim): - # Extend the basis from codim to codim + 1 - fa, fb, fc, dfa, dfb, dfc = jacobi_factors(*X[codim:codim+3], *dX[codim:codim+3]) + for axis in range(dim): + # Extend the basis from axis to axis + 1 + fa, fb, fc, dfa, dfb, dfc = jacobi_factors(*X[axis:axis+3], *dX[axis:axis+3]) ddfc = 2 * numpy.outer(dfb, dfb) - for sub_index in reference_element.lattice_iter(0, n, codim): + for sub_index in reference_element.lattice_iter(0, n, axis): # handle i = 0 icur = idx(*sub_index, 0) inext = idx(*sub_index, 1) - if variant == "bubble": - alpha = 2 * sum(sub_index) - a = b = -0.5 - else: - alpha = 2 * sum(sub_index) + len(sub_index) - if variant == "dual": - alpha += 1 + len(sub_index) - a = 0.5 * (alpha + beta) + 1.0 - b = 0.5 * (alpha - beta) - + alpha, beta = dubiner_jacobi_weights(axis, sum(sub_index), variant) + a, b, c = coefficients(alpha, beta, 0) fcur = a * fa - b * fb phi[inext] = phi[icur] * fcur if order: @@ -249,24 +268,87 @@ def dubiner_recurrence(dim: int, results[k][inext] += _product_derivative(fprev, dfprev, ddfprev, prev, k) # normalize - d = codim + 1 - shift = 1 if variant == "dual" else 0 + d = axis + 1 for index in reference_element.lattice_iter(0, n+1, d): icur = idx(*index) - if variant is not None: - p = index[-1] + shift - alpha = 2 * (sum(index[:-1]) + d * shift) - 1 - norm2 = (0.5 + d) / d - if p > 0 and p + alpha > 0: - norm2 *= (p + alpha) * (2*p + alpha) / p - else: - norm2 = (2*sum(index) + d) / d - scale = math.sqrt(norm2) + scale = math.sqrt(dubiner_norm2(d, sum(index[:-1]), index[-1], variant)) for result in results: result[icur] *= scale return results +def principal_functions(n, eta, axis, order=0, variant=None): + """Tabulate one axis of the Dubiner basis in collapsed coordinates. + + Returns value (``V``), derivative (``D``), affine derivative (``tD``), + and one-power-lowered (``W``) tables indexed by ``(m, i, point)``. + """ + coefficients = integrated_jrc if variant == "bubble" else jrc + eta = numpy.asarray(eta) + w = 0.5 * (1.0 - eta) + npts = eta.shape[0] + num_m = 1 if axis == 1 else n + 1 + modes = ("V",) if order == 0 else ("V", "D", "W", "tD") + tables = {mode: numpy.zeros((num_m, n + 1, npts), dtype=eta.dtype) + for mode in modes} + powers = w[None, :] ** numpy.arange(num_m)[:, None] + for m in range(num_m): + count = n + 1 - m + g = tables["V"][m, :count] + dg = numpy.zeros_like(g) if order else None + alpha, beta = dubiner_jacobi_weights(axis-1, m, variant) + a, b, c = coefficients(alpha, beta, 0) + g[0] = 1.0 + if count > 1: + g[1] = a * eta + b + if order: + dg[1] = a + for i in range(1, count - 1): + a, b, c = coefficients(alpha, beta, i) + factor = a * eta + b + g[i + 1] = factor * g[i] - c * g[i - 1] + if order: + dg[i + 1] = (factor * dg[i] + a * g[i] + - c * dg[i - 1]) + + if n > 0: + normalisation = numpy.fromiter( + (math.sqrt(dubiner_norm2(axis, m, i, variant)) + for i in range(count)), dtype=float, count=count) + g *= normalisation[:, None] + if order: + dg *= normalisation[:, None] + + if order and m: + tables["W"][m, :count] = powers[m - 1] * g + if order: + tables["D"][m, :count] = ( + powers[m] * dg - 0.5 * m * tables["W"][m, :count]) + tables["tD"][m, :count] = ( + 0.5 * (1.0 + eta) * tables["D"][m, :count]) + g *= powers[m] + return tables + + +def _duffy_derivative_factors(tables, direction): + """Factorize one physical derivative of the Duffy pullback. + + For the lower-triangular Duffy map, ``d/dxi_direction`` is a sum over + collapsed axes up to ``direction``. Axes before the differentiated axis + contribute values. Axes after it contribute one lower power of their + collapse weight. A non-final differentiated axis also contributes its + affine collapsed coordinate. + """ + terms = [] + for axis in range(direction + 1): + derivative = "D" if axis == direction else "tD" + factors = tuple(table["V"] for table in tables[:axis]) + factors += (tables[axis][derivative],) + factors += tuple(table["W"] for table in tables[axis + 1:]) + terms.append(factors) + return tuple(terms) + + def C0_basis(dim, n, tabulations): """Modify a tabulation of a hierarchical basis to enforce C0-continuity. @@ -446,6 +528,79 @@ def distance(alpha, beta): result[alpha] = vals return result + def get_duffy_permutation(self, n: int) -> numpy.ndarray: + """Map lexicographic Duffy lattice positions to expansion members.""" + dim = self.ref_el.get_spatial_dimension() + index = (lambda p: p, morton_index2, morton_index3)[dim - 1] + return numpy.fromiter( + (index(*alpha[::-1]) + for alpha in reference_element.lattice_iter(0, n+1, dim)), + dtype=int, count=math.comb(n + dim, dim)) + + def tabulate_duffy(self, n: int, eta_pts: tuple, order: int = 0, + cell: int = 0) -> dict: + """Tabulate the expansion set in separable Duffy form. + + Parameters + ---------- + n : int + Polynomial degree. + eta_pts : tuple + One collapsed-coordinate point array per spatial dimension. + order : int, optional + Highest derivative order. + cell : int, optional + Cell in the reference complex. + + Returns + ------- + dict + Separable terms for each derivative multi-index. + + Notes + ----- + The Duffy map sends the tensor-product grid ``eta_pts`` to the + reference cell. Each result term is ``(coefficient, factors)``. + ``factors`` is an ordered tuple of ``(axis, table)`` pairs such that:: + + D^alpha phi_(i_1..i_d) = sum_terms coeff * prod_q F_q[m_q, i_q, pt_axis_q] + + The basis lattice satisfies ``i_1 + ... + i_d <= n``. Its prefix + index is ``m_1 = 0`` and ``m_t = i_1 + ... + i_{t-1}``. Factor order + determines the lattice axes. Each axis label selects a point + coordinate. ``morton_index2`` and ``morton_index3`` enumerate the + lattice members. Consumers can deduplicate shared tables by identity. + + The result contains the raw tabulation for every continuity variant. + The ``C0_basis`` recombination mixes lattice multi-indices. Callers + compose this recombination with their expansion coefficients. + + """ + if order > 1: + raise NotImplementedError("tabulate_duffy is limited to first derivatives") + sd = self.ref_el.get_spatial_dimension() + assert len(eta_pts) == sd + A, b = self.affine_mappings[cell] + scale = self.get_scale(n, cell=cell) + if self.variant == "bubble": + scale = -scale + tables = [principal_functions(n, eta_pts[t], t + 1, order=order, variant=self.variant) + for t in range(sd)] + + result = {(0,) * sd: [ + (scale, tuple(enumerate(table["V"] for table in tables))) + ]} + if order > 0: + xi_terms = tuple(_duffy_derivative_factors(tables, k) + for k in range(sd)) + # Push forward to the cell coordinates: d/dx_k = sum_l A[l, k] d/dxi_l. + for k in range(sd): + alpha = tuple(int(u == k) for u in range(sd)) + result[alpha] = [(scale * A[l, k], tuple(enumerate(factors))) + for l in range(sd) if A[l, k] != 0.0 + for factors in xi_terms[l]] + return result + def _tabulate(self, n, pts, order=0): """A version of tabulate() that also works for a single point.""" pts = numpy.asarray(pts) diff --git a/finat/duffy.py b/finat/duffy.py new file mode 100644 index 00000000..753b36e3 --- /dev/null +++ b/finat/duffy.py @@ -0,0 +1,112 @@ +from functools import reduce + +import numpy + +import gem + +from FIAT.expansions import C0_basis +from finat.physically_mapped import MappedTabulation +from finat.point_set import CollapsedTensorProductPointSet + + +class DuffyElement: + """Mixin for sum-factorized tabulation on collapsed simplex coordinates.""" + + def basis_evaluation(self, order, ps, entity=None, coordinate_mapping=None): + """Tabulate on a collapsed point set when the element supports it.""" + sd = self.cell.get_dimension() + if not (isinstance(ps, CollapsedTensorProductPointSet) + and order <= 1 + and (entity is None or entity == (sd, 0)) + and not self.complex.is_macrocell()): + return super().basis_evaluation( + order, ps, entity=entity, + coordinate_mapping=coordinate_mapping) + return self.duffy_evaluation(order, ps, entity) + + def get_coefficient_matrix(self) -> gem.Literal: + """Return the map from lattice-ordered expansions to dofs. + + Returns + ------- + gem.Literal + Coefficient matrix in Duffy lattice order. + + """ + degree = self.degree + sd = self.cell.get_spatial_dimension() + poly_set = self._element.get_nodal_basis() + coeffs = numpy.array(poly_set.get_coeffs(), copy=True) + expansion_set = poly_set.get_expansion_set() + if expansion_set.continuity == "C0": + recombination, = C0_basis(sd, degree, + [numpy.eye(coeffs.shape[1])]) + coeffs = coeffs @ recombination + + coeffs = coeffs[:, expansion_set.get_duffy_permutation(degree)] + if numpy.allclose(coeffs, 0.0): + raise ValueError("empty Duffy coefficient matrix") + return gem.Literal(coeffs) + + def duffy_evaluation(self, order, ps, entity=None): + """Return a sum-factorized, dof-indexed tabulation.""" + assert isinstance(ps, CollapsedTensorProductPointSet) + cell_dim = self.cell.get_dimension() + if entity is not None and entity != (cell_dim, 0): + raise NotImplementedError( + "duffy_evaluation is only supported on the cell interior") + if self.complex.is_macrocell(): + raise NotImplementedError("duffy_evaluation is not supported on split cells") + + degree = self.degree + sd = self.cell.get_spatial_dimension() + poly_set = self._element.get_nodal_basis() + expansion_set = poly_set.get_expansion_set() + etas = tuple(2.0 * factor.points.ravel() - 1.0 + for factor in ps.factors) + duffy = expansion_set.tabulate_duffy(degree, etas, order=order) + + def lookup_index(index_table, multiindex): + index_table = gem.Literal(index_table, dtype=gem.uint_type) + return gem.VariableIndex(gem.Indexed(index_table, multiindex)) + + multiindex = [] + for _ in range(sd): + multiindex.append(gem.JaggedIndex(extent=degree + 1, parents=tuple(multiindex))) + multiindex = tuple(multiindex) + # The first table axis is the sum of all preceding lattice indices, + # not an independent iteration axis, so it is an indirect index. + duffy_indices = [0, *multiindex[:1]] + for t in range(2, sd): + index_table = reduce(numpy.add.outer, (numpy.arange(degree + 1),) * t) + index_table = numpy.minimum(index_table, degree) + duffy_indices.append(lookup_index(index_table, multiindex[:t])) + + literals = {} + + def as_gem(table): + key = id(table) + try: + return literals[key] + except KeyError: + return literals.setdefault(key, gem.Literal(table)) + + result = {} + for alpha, terms in duffy.items(): + exprs = [] + for coeff, factors in terms: + expr = gem.Product(*( + gem.Indexed(as_gem(table), + (index_expr, i, ps.indices[axis])) + for (axis, table), index_expr, i + in zip(factors, duffy_indices, multiindex))) + if coeff != 1.0: + expr = gem.Product(gem.Literal(coeff), expr) + exprs.append(expr) + result[alpha] = gem.Sum(*exprs) + + coefficients = self.get_coefficient_matrix() + tabulation = { + alpha: gem.FlattenedTensor(expr, multiindex) + for alpha, expr in result.items()} + return MappedTabulation(coefficients, tabulation) diff --git a/finat/fiat_elements.py b/finat/fiat_elements.py index 22408cf0..695099bd 100644 --- a/finat/fiat_elements.py +++ b/finat/fiat_elements.py @@ -3,6 +3,7 @@ import numpy as np from gem.utils import cached_property +from finat.duffy import DuffyElement from finat.finiteelementbase import FiniteElementBase from finat.point_set import PointSet, PointSingleton @@ -296,9 +297,8 @@ def value_shape(self): return () -class Bernstein(ScalarFiatElement): - # TODO: Replace this with a smarter implementation - def __init__(self, cell, degree): +class Bernstein(DuffyElement, ScalarFiatElement): + def __init__(self, cell: FIAT.reference_element.SimplicialComplex, degree: int) -> None: super().__init__(FIAT.Bernstein(cell, degree)) diff --git a/finat/physically_mapped.py b/finat/physically_mapped.py index 8027732b..2b2aceee 100644 --- a/finat/physically_mapped.py +++ b/finat/physically_mapped.py @@ -1,6 +1,7 @@ from abc import ABCMeta, abstractmethod from collections.abc import Iterable, Mapping from functools import cached_property +from numbers import Number import gem import numpy @@ -50,7 +51,8 @@ def __init__( row = [] for column in range(M.shape[1]): value = M.array[source_row, column] - if not isinstance(value, gem.Zero): + if not (isinstance(value, gem.Zero) + or isinstance(value, Number) and value == 0): row.append((column, value)) nonzero_rows.append(row) width = max((len(row) for row in nonzero_rows), default=0) diff --git a/finat/point_set.py b/finat/point_set.py index 068f7659..ce92af5f 100644 --- a/finat/point_set.py +++ b/finat/point_set.py @@ -1,7 +1,8 @@ import abc import hashlib -from functools import cached_property +from functools import cached_property, reduce from itertools import chain, product +from operator import mul import numpy @@ -233,6 +234,47 @@ def almost_equal(self, other, tolerance=1e-12): for s, o in zip(self.factors, other.factors)) +def _collapsed_coordinates(eta): + """Map unit tensor-product coordinates to the unit simplex.""" + return tuple( + eta[t] * reduce( + mul, (1 - eta[u] for u in range(t + 1, len(eta))), 1) + for t in range(len(eta))) + + +class CollapsedTensorProductPointSet(TensorPointSet): + r"""A tensor point set mapped to the simplex by collapsed coordinates. + + The factors are one-dimensional point sets on the ``[0, 1]`` reference + interval. Their tensor product is mapped by + + .. math:: x_t = \eta_t \prod_{u=t+1}^{d-1}(1 - \eta_u). + + Parameters + ---------- + factors : tuple of AbstractPointSet + One-dimensional point sets of collapsed coordinates, one per + spatial dimension of the simplex. + + """ + + def __init__(self, factors: tuple[AbstractPointSet, ...]) -> None: + super().__init__(factors) + assert all(ps.dimension == 1 for ps in self.factors) + + @cached_property + def points(self): + etas = [ps.points.ravel() for ps in self.factors] + grids = list(numpy.meshgrid(*etas, indexing="ij")) + return numpy.stack([x.ravel() for x in _collapsed_coordinates(grids)], + axis=-1) + + @cached_property + def expression(self): + etas = [gem.Indexed(ps.expression, (0,)) for ps in self.factors] + return gem.ListTensor(_collapsed_coordinates(etas)) + + class FacetPointSet(AbstractPointSet): """A point set on facets. diff --git a/finat/quadrature.py b/finat/quadrature.py index 858eef94..d2c827ea 100644 --- a/finat/quadrature.py +++ b/finat/quadrature.py @@ -1,6 +1,7 @@ import hashlib from abc import ABCMeta, abstractmethod from functools import cached_property +from math import factorial import gem import numpy @@ -8,28 +9,38 @@ from FIAT.quadrature_schemes import create_quadrature as fiat_scheme from FIAT.reference_element import HEXAHEDRON, LINE, QUADRILATERAL, TENSORPRODUCT from gem.utils import safe_repr +from recursivenodes.quadrature import gaussjacobi -from finat.point_set import (GaussLegendrePointSet, GaussLobattoLegendrePointSet, +from finat.point_set import (CollapsedTensorProductPointSet, + GaussLegendrePointSet, GaussLobattoLegendrePointSet, KMVPointSet, PointSet, TensorPointSet) def make_quadrature(ref_el, degree, scheme="default"): - """ - Generate quadrature rule for given reference element - that will integrate an polynomial of order 'degree' exactly. - - For low-degree polynomials on triangles (<=50) and tetrahedra (<=15), this - uses hard-coded rules, otherwise it falls back to a collapsed - Gauss scheme on simplices. On tensor-product cells, it is a - tensor-product quadrature rule of the subcells. - - :arg ref_el: The FIAT cell to create the quadrature for. - :arg degree: The degree of polynomial that the rule should - integrate exactly. - :kwarg scheme: The quadrature scheme, can be choosen from ["default", "canonical", "KMV"] - "default" -> hard-coded scheme for low degree and collapsed Gauss scheme for high degree, - "canonical" -> collapsed Gauss scheme, - "KMV" -> spectral lumped scheme for low degree (<=6 on triangles, <=3 on tetrahedra). + """Create a quadrature rule for a reference element. + + Parameters + ---------- + ref_el : FIAT.reference_element.Cell + Reference cell. + degree : int or tuple of int + Exact polynomial degree. + scheme : str, optional + Quadrature scheme. + + Returns + ------- + QuadratureRule + Quadrature points, weights, and reference cell. + + Notes + ----- + The ``default`` scheme uses tabulated simplex rules at low degree. It + uses collapsed Gauss rules at high degree. The ``canonical`` scheme uses + collapsed Gauss rules with flat points. The ``KMV`` scheme uses spectral + lumped rules. The ``collapsed`` scheme retains the tensor-product Duffy + structure for sum factorization. + """ if ref_el.get_shape() == TENSORPRODUCT: try: @@ -48,6 +59,9 @@ def make_quadrature(ref_el, degree, scheme="default"): if degree < 0: raise ValueError("Need positive degree, not %d" % degree) + if scheme.lower() == "collapsed": + return collapsed_gauss_jacobi_quadrature(ref_el, degree) + if scheme.lower() in {"kmv", "lump"}: fiat_rule = fiat_scheme(ref_el, degree, "KMV") if ref_el.get_shape() == LINE: @@ -69,6 +83,47 @@ def make_quadrature(ref_el, degree, scheme="default"): return QuadratureRule(point_set, fiat_rule.get_weights(), ref_el=ref_el, io_ornt_map_tuple=fiat_rule._intrinsic_orientation_permutation_map_tuple) +def collapsed_gauss_jacobi_quadrature(ref_el, degree): + """Create a structured collapsed Gauss-Jacobi quadrature rule. + + Parameters + ---------- + ref_el : FIAT.reference_element.Cell + The simplex to create the quadrature rule on. + degree : int + The degree of polynomial that the rule should integrate exactly. + + Returns + ------- + CollapsedTensorProductQuadratureRule + The structured quadrature rule. + + Notes + ----- + The Duffy map is + ``x_t = eta_t * prod(u > t) (1 - eta_u)``. The Jacobi weight + ``(1 - eta_u)**u`` absorbs the Duffy Jacobian on axis ``u``. The simplex + weights are products of the one-dimensional weights. + + """ + if ref_el.is_macrocell(): + raise NotImplementedError("Collapsed quadrature is not supported on split cells") + dim = ref_el.get_spatial_dimension() + num_points = (degree + 1 + 1) // 2 # exact integration + factors = [] + for axis in range(dim): + xs, ws = gaussjacobi(num_points, axis, 0.0) + # Map from the biunit to the unit interval, folding the change of + # measure of the Jacobi weight into the quadrature weights + xs = (1.0 + xs) / 2.0 + ws = ws / 2.0 ** (axis + 1) + if axis == 0: + # The Duffy map produces points on the unit simplex + ws = ws * (ref_el.volume() * factorial(dim)) + factors.append(QuadratureRule(PointSet(xs[:, None]), ws)) + return CollapsedTensorProductQuadratureRule(factors, ref_el=ref_el) + + class AbstractQuadratureRule(metaclass=ABCMeta): """Abstract class representing a quadrature rule as point set and a corresponding set of weights.""" @@ -176,3 +231,35 @@ def point_set(self): @cached_property def weight_expression(self): return gem.Product(*(q.weight_expression for q in self.factors)) + + +class CollapsedTensorProductQuadratureRule(AbstractQuadratureRule): + """Simplex quadrature rule with tensor-product structure in collapsed + coordinates, following Karniadakis & Sherwin. + + Parameters + ---------- + factors : tuple of QuadratureRule + One-dimensional quadrature rules of collapsed coordinates on the + unit interval, one per spatial dimension of the simplex, with the + Duffy Jacobian folded into the weights. + ref_el : FIAT.reference_element.Cell + The simplex the quadrature rule integrates over. + + """ + + def __init__(self, factors, ref_el=None): + self.ref_el = ref_el + self.factors = tuple(factors) + self._intrinsic_orientation_permutation_map_tuple = (None,) + + def __repr__(self): + return f"{type(self).__name__}({self.factors!r}, {self.ref_el!r})" + + @cached_property + def point_set(self): + return CollapsedTensorProductPointSet([q.point_set for q in self.factors]) + + @cached_property + def weight_expression(self): + return gem.Product(*(q.weight_expression for q in self.factors)) diff --git a/finat/spectral.py b/finat/spectral.py index ab402265..a27345c5 100644 --- a/finat/spectral.py +++ b/finat/spectral.py @@ -4,8 +4,10 @@ from abc import ABCMeta, abstractmethod from finat.citations import cite +from finat.duffy import DuffyElement from finat.fiat_elements import ScalarFiatElement, Lagrange, DiscontinuousLagrange -from finat.point_set import GaussLobattoLegendrePointSet, GaussLegendrePointSet, KMVPointSet +from finat.point_set import (GaussLobattoLegendrePointSet, GaussLegendrePointSet, + KMVPointSet) class SpectralElement(metaclass=ABCMeta): @@ -64,7 +66,7 @@ def __init__(self, cell, degree): cite("Geevers2018new") -class Legendre(ScalarFiatElement): +class Legendre(DuffyElement, ScalarFiatElement): """DG element with Legendre polynomials.""" def __init__(self, cell, degree, variant=None): @@ -72,7 +74,7 @@ def __init__(self, cell, degree, variant=None): super().__init__(fiat_element) -class IntegratedLegendre(ScalarFiatElement): +class IntegratedLegendre(DuffyElement, ScalarFiatElement): """CG element with integrated Legendre polynomials.""" def __init__(self, cell, degree, variant=None): diff --git a/gem/__init__.py b/gem/__init__.py index f1e77203..703c5ea9 100644 --- a/gem/__init__.py +++ b/gem/__init__.py @@ -1,2 +1,3 @@ from gem.gem import * # noqa +from gem.jagged import compact_index_layout, simplex_lattice_ranks # noqa from gem.optimise import select_expression # noqa diff --git a/gem/coffee.py b/gem/coffee.py index 2cc4e1d9..692c5014 100644 --- a/gem/coffee.py +++ b/gem/coffee.py @@ -10,8 +10,9 @@ import numpy -from gem.gem import ComponentTensor, Index, Indexed, IndexSum, Literal, Node, one -from gem.node import MemoizerArg +from gem.gem import (ComponentTensor, FlattenedTensor, Index, Indexed, + IndexSum, Literal, Node, one) +from gem.node import MemoizerArg, traversal from gem.cost import has_arithmetic from gem.optimise import (filtered_replace_indices, make_sum, make_product, traverse_sum) @@ -278,6 +279,9 @@ def _share_linear_maps( involved = linear_set.intersection(atomic.free_indices) if len(involved) != 1: continue + if any(isinstance(node, FlattenedTensor) + for node in traversal((atomic,))): + continue index, = involved normal = replacer(atomic, ((index, canonical[index.extent]),)) key = _additive_map_key(normal) @@ -310,15 +314,46 @@ def _share_linear_maps( return result -def optimise_monomial_sum(monomial_sum, linear_indices): +def optimise_monomial_sum( + monomial_sum, linear_indices, contraction_order=()): """Choose optimal common atomic subexpressions and factorise a :class:`MonomialSum` object to create a GEM expression. :arg monomial_sum: a :class:`MonomialSum` object :arg linear_indices: tuple of linear indices + :arg contraction_order: contraction indices, outermost first :returns: factorised GEM expression """ + if contraction_order: + grouped = defaultdict(MonomialSum) + order = {} + remaining = frozenset(contraction_order) + for monomial in monomial_sum: + inner_indices = tuple(index for index in monomial.sum_indices + if index in remaining) + involved = frozenset(inner_indices) + inner_atomics = tuple( + atomic for atomic in monomial.atomics + if involved.intersection(atomic.free_indices)) + outer_indices = tuple(index for index in monomial.sum_indices + if index not in remaining) + outer_atomics = tuple( + atomic for atomic in monomial.atomics + if atomic not in inner_atomics) + key = outer_indices, outer_atomics + order.setdefault(key, None) + grouped[key].add( + inner_indices, inner_atomics, monomial.rest) + + outer_sum = MonomialSum() + for outer_indices, outer_atomics in order: + inner = optimise_monomial_sum( + grouped[(outer_indices, outer_atomics)], + linear_indices, contraction_order[1:]) + outer_sum.add(outer_indices, outer_atomics, inner) + monomial_sum = outer_sum + monomial_sum = _share_linear_maps(monomial_sum, linear_indices) groups = groupby(monomial_sum, key=lambda m: frozenset(m.sum_indices)) new_monomials = [] diff --git a/gem/driver.py b/gem/driver.py new file mode 100644 index 00000000..25439363 --- /dev/null +++ b/gem/driver.py @@ -0,0 +1,196 @@ +"""Composite optimisation pipelines over GEM. + +The passes in `gem.optimise`, `gem.coffee`, `gem.refactorise` and +`gem.jagged` are primitives; each rewrites a DAG in one way. This module is +the layer above them, where a pipeline may use all four. Keeping the +pipelines here is what lets the primitive modules stay free of each other. +""" + +from collections import OrderedDict +from collections.abc import Iterable +from itertools import zip_longest + +import numpy + +from gem.coffee import optimise_monomial_sum +from gem.gem import (ComponentTensor, Delta, FlattenedTensor, Indexed, + IndexSum, ListTensor, Node, Product) +from gem.jagged import unflatten_free_indices, unflatten +from gem.node import MemoizerArg, traversal +from gem.optimise import (cancel_nested_deltas, delta_elimination, + filtered_replace_indices, + pull_back_indirect_delta, repeated_contractions, + sum_factorise, traverse_product) +from gem.refactorise import (ATOMIC, COMPOUND, OTHER, FactorisationError, + collect_monomials) + + +def contraction(expression): + """Optimise the contractions of the tensor product at the root of + the expression, including: + + - IndexSum-Delta cancellation + - Sum factorisation + + This routine was designed with finite element coefficient + evaluation in mind. + """ + + # Common memoizer to remove ComponentTensors + index_replacer = MemoizerArg(filtered_replace_indices) + + # Eliminate annoying ComponentTensors + expression = index_replacer(expression, ()) + + def flatten(root): + """Break the product tree at ``root`` up and cancel its Deltas.""" + # The contraction at the root is always broken up, as that is the + # one being optimised + keep = repeated_contractions(root) + sum_indices, factors = traverse_product( + root, index_replacer=index_replacer, + stop_at=lambda e: e is not root and e in keep) + sum_indices, factors = pull_back_indirect_delta( + sum_indices, factors, index_replacer) + sum_indices, factors = delta_elimination( + sum_indices, factors, index_replacer=index_replacer) + return sum_indices, [index_replacer(f, ()) for f in factors] + + def rebuild(expression): + sum_indices, factors = flatten(expression) + flattened = IndexSum(Product(*factors), sum_indices) + unflattened = unflatten(flattened) + if unflattened is not flattened: + sum_indices, factors = flatten(unflattened) + return sum_factorise(sum_indices, factors) + + # Sometimes the value shape is composed as a ListTensor, which + # could get in the way of decomposing factors. In particular, + # this is the case for H(div) and H(curl) conforming tensor + # product elements. So if ListTensors are used, they are pulled + # out to be outermost, so we can straightforwardly factorise each + # of its entries. + lt_fis = OrderedDict() # ListTensor free indices + for node in traversal((expression,)): + if isinstance(node, Indexed): + child, = node.children + if isinstance(child, ListTensor): + lt_fis.update(zip_longest(node.multiindex, ())) + lt_fis = tuple(index for index in lt_fis if index in expression.free_indices) + + if lt_fis: + # Rebuild each split component + tensor = ComponentTensor(expression, lt_fis) + entries = [Indexed(tensor, zeta) for zeta in numpy.ndindex(tensor.shape)] + entries = [index_replacer(e, ()) for e in entries] + return Indexed(ListTensor( + numpy.array(list(map(rebuild, entries))).reshape(tensor.shape) + ), lt_fis) + else: + # Rebuild whole expression at once + return rebuild(expression) + + +def _refactor_unflattened_outputs( + outputs: list[tuple[Node, Node]]) -> list[tuple[Node, Node]] | None: + r"""Recover sum factorisation after exposing every argument lattice. + + Consider a bilinear form whose argument tabulations are transformed by + sparse matrices, + + .. math:: + + A_{ij} = \sum_q (S B(q))_i\,G(q)\,(T C(q))_j. + + Sparse-delta cancellation moves rows of ``S`` and ``T`` into the return + scatter. Both flat argument indices must then be replaced by their + jagged lattice multiindices *before* expanding derivative sums. Expanding + after only one replacement duplicates the second lattice once per + summand, producing many equivalent loop nests. + + The refactorisation is valid when the non-tabulation part of every + monomial is independent of the argument indices. This condition says + exactly that the sparse transforms have been absorbed by the scatter. + Dense or otherwise residual transforms use the conservative legacy path. + """ + candidates = [] + has_nontrivial_sum = False + for variable, expression in outputs: + argument_indices = frozenset(variable.free_indices) + contraction_indices = frozenset( + index + for node in traversal((expression,)) + if isinstance(node, IndexSum) + for index in node.multiindex) + + def classify(node): + involved = argument_indices.intersection(node.free_indices) + if not involved: + return OTHER + if isinstance(node, Indexed): + return ATOMIC if contraction_indices.intersection( + node.free_indices) else OTHER + return COMPOUND + + try: + monomial_sum, = collect_monomials([expression], classify) + except FactorisationError: + return None + + monomials = tuple(monomial_sum) + if any(argument_indices.intersection(monomial.rest.free_indices) + for monomial in monomials): + return None + has_nontrivial_sum |= len(monomials) > 1 + sum_indices = tuple(OrderedDict.fromkeys( + index + for monomial in monomials + for index in monomial.sum_indices)) + candidates.append((variable, monomial_sum, sum_indices)) + + if not has_nontrivial_sum: + return None + return [ + (variable, optimise_monomial_sum( + monomial_sum, variable.index_ordering(), sum_indices)) + for variable, monomial_sum, sum_indices in candidates + ] + + +def unflatten_returns( + pairs: Iterable[tuple[Node, Node]]) -> list[tuple[Node, Node]]: + """Unflatten free argument indices in assignment pairs. + + Every compatible argument lattice is exposed jointly. Bilinear + expressions are then refactorised as monomial sums so that sparse + argument transforms remain in the output scatter while quadrature + contractions recover their one-dimensional factors. Expressions with + residual argument-dependent transforms retain the established local + distribution strategy. + """ + pairs = list(pairs) + if not any(isinstance(node, FlattenedTensor) + for _, expression in pairs + for node in traversal((expression,))): + return pairs + + result = [] + for variable, expression in pairs: + expression = cancel_nested_deltas(expression) + joint_outputs, changed = unflatten_free_indices( + variable, expression, split_separable_sums=False) + refactored = _refactor_unflattened_outputs( + joint_outputs) if changed else None + if refactored is not None: + result.extend(refactored) + continue + + outputs, changed = unflatten_free_indices( + variable, expression, split_separable_sums=True) + if changed or any(isinstance(node, Delta) + for _, output in outputs + for node in traversal((output,))): + outputs = [(output_variable, contraction(output)) + for output_variable, output in outputs] + result.extend(outputs) + return result diff --git a/gem/flop_count.py b/gem/flop_count.py index 1f311c04..e30cb572 100644 --- a/gem/flop_count.py +++ b/gem/flop_count.py @@ -27,13 +27,28 @@ def statement_block(tree, temporaries): def statement_for(tree, temporaries): extent = tree.index.extent assert extent is not None - child, = tree.children - token = _active_indices.set(_active_indices.get() | {tree.index}) + active_token = _active_indices.set( + _active_indices.get() | {tree.index}) try: + index_values = _index_values.get() + if tree.index.parents and all( + parent in index_values for parent in tree.index.parents): + extent = tree.index.iteration_extent(index_values) + child, = tree.children + if tree.index in _control_indices.get(): + flops = 0 + for value in range(extent): + token = _index_values.set( + index_values | {tree.index: value}) + try: + flops += statement(child, temporaries) + finally: + _index_values.reset(token) + return flops flops = statement(child, temporaries) + return flops * extent finally: - _active_indices.reset(token) - return flops * extent + _active_indices.reset(active_token) @statement.register(imp.Initialise) @@ -143,13 +158,36 @@ def flops_dense_linear_algebra(expr, temporaries): @flops.register(gem.ComponentTensor) def flops_componenttensor(expr, temporaries): 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) + implicit_indices = tuple( + index for index in expr.multiindex + if index not in _active_indices.get()) + if not implicit_indices: + return expression_flops(body, temporaries) + control = _control_indices.get().intersection(implicit_indices) + if not control and not any( + index.parents for index in implicit_indices): + extent = numpy.prod( + [index.extent for index in implicit_indices], dtype=int) + return extent * expression_flops(body, temporaries) + + def count(position): + if position == len(implicit_indices): + return expression_flops(body, temporaries) + index = implicit_indices[position] + values = _index_values.get() + extent = index.extent + if index.parents: + extent = index.iteration_extent(values) + total = 0 + for value in range(extent): + token = _index_values.set(values | {index: value}) + try: + total += count(position + 1) + finally: + _index_values.reset(token) + return total + + return count(0) def expression_flops(expression, temporaries, top=False): @@ -173,14 +211,35 @@ def count_flops(impero_c): :returns: approximate flop count for the tree. """ try: - token = _active_indices.set(frozenset()) + control_token = _control_indices.set( + frozenset(_find_control_indices(impero_c.tree))) + index_token = _index_values.set({}) + active_token = _active_indices.set(frozenset()) try: return statement(impero_c.tree, set(impero_c.temporaries)) finally: - _active_indices.reset(token) + _active_indices.reset(active_token) + _index_values.reset(index_token) + _control_indices.reset(control_token) except (ValueError, NotImplementedError): return 0 +_index_values = ContextVar("flop_count_index_values", default={}) _active_indices = ContextVar("flop_count_active_indices", default=frozenset()) +_control_indices = ContextVar("flop_count_control_indices", + default=frozenset()) + + +def _find_control_indices(tree): + """Find loop indices controlling dependent loop bounds.""" + result = set() + if isinstance(tree, imp.For): + if tree.index.parents: + result.update(tree.index.parents) + result.update(_find_control_indices(tree.children[0])) + elif isinstance(tree, imp.Block): + for child in tree.children: + result.update(_find_control_indices(child)) + return result diff --git a/gem/gem.py b/gem/gem.py index 0953d697..1f41d2ba 100644 --- a/gem/gem.py +++ b/gem/gem.py @@ -16,7 +16,7 @@ from abc import ABCMeta from itertools import chain, repeat -from functools import partial, reduce +from functools import lru_cache, partial, reduce from operator import attrgetter from numbers import Integral, Number @@ -32,8 +32,9 @@ 'Variable', 'Sum', 'Product', 'Division', 'FloorDiv', 'Remainder', 'Power', 'MathFunction', 'MinValue', 'MaxValue', 'Comparison', 'LogicalNot', 'LogicalAnd', 'LogicalOr', 'Conditional', - 'Index', 'VariableIndex', 'Indexed', 'ComponentTensor', - 'IndexSum', 'ListTensor', 'Concatenate', 'Delta', 'OrientationVariableIndex', + 'Index', 'JaggedIndex', 'VariableIndex', 'Indexed', 'ComponentTensor', + 'FlattenedTensor', 'IndexSum', 'ListTensor', 'Concatenate', 'Delta', + 'OrientationVariableIndex', 'index_sum', 'partial_indexed', 'reshape', 'view', 'indices', 'as_gem', 'FlexiblyIndexed', 'Inverse', 'Solve', 'extract_type', 'uint_type', 'Piecewise'] @@ -609,6 +610,10 @@ class Index(IndexBase): # Not true object count, just for naming purposes _count = 0 + # Indices whose values reduce the iteration bound; JaggedIndex overrides + # this with a slot, so every index answers to `parents`. + parents = () + __slots__ = ('name', 'extent', 'count') def __init__(self, name=None, extent=None): @@ -617,6 +622,22 @@ def __init__(self, name=None, extent=None): self.count = Index._count self.extent = extent + def iteration_extent(self, parent_values: dict) -> int: + """Return the loop extent at fixed parent-index values. + + Parameters + ---------- + parent_values + Values of indices controlling this index. + + Returns + ------- + int + Number of admissible values. + + """ + return self.extent + def set_extent(self, value): # Set extent, check for consistency if self.extent is None: @@ -645,17 +666,66 @@ def __setstate__(self, state): self.name, self.extent, self.count = state +class JaggedIndex(Index): + """Free index whose effective iteration bound depends on the values of + other (parent) free indices. + + The iteration bound is ``0 <= i < extent - (p_1 + ... + p_k)`` for + parent indices ``p_1, ..., p_k``. The ``extent`` attribute is the static + upper bound. Every parent index is zero at this bound. Consumers can + treat this as a plain :class:`Index` of extent ``extent``. Expressions + indexed by a :class:`JaggedIndex` must evaluate to zero outside the + jagged bounds. The jagged bounds only optimize the generated loops. + + Parameters + ---------- + name : str, optional + Name of the index. + extent : int, optional + Static (rectangular) upper bound of the index. + parents : tuple of Index + The indices whose values reduce the iteration bound. Loops over + this index must nest inside the loops over its parents. + + """ + + __slots__ = ('parents',) + + def __init__(self, name: str | None = None, extent: int | None = None, + parents: tuple = ()): + super().__init__(name=name, extent=extent) + parents = tuple(parents) + assert all(isinstance(p, Index) for p in parents) + self.parents = parents + + def iteration_extent(self, parent_values: dict) -> int: + return self.extent - sum( + parent_values[parent] for parent in self.parents) + + def __getstate__(self): + return super().__getstate__() + (self.parents,) + + def __setstate__(self, state): + super().__setstate__(state[:-1]) + self.parents = state[-1] + + class VariableIndex(IndexBase): """An index that is constant during a single execution of the kernel, but whose value is not known at compile time.""" __slots__ = ('expression',) - def __init__(self, expression): + def __new__(cls, expression): assert isinstance(expression, Node) assert not expression.shape if expression.dtype != uint_type: raise ValueError(f"expression.dtype ({expression.dtype}) != uint_type ({uint_type})") + if isinstance(expression, Constant): + return int(expression.value) + return super().__new__(cls) + + def __init__(self, expression): self.expression = expression def __eq__(self, other): @@ -731,10 +801,11 @@ def __new__(cls, aggregate, multiindex): # 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( + nested = set(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): + safe = not set(jj).intersection(set(ff) | nested) + if safe: # Only replace indices that are not present in C rep = dict(zip(jj, ii)) ll = tuple(rep.get(k, k) for k in kk) @@ -753,25 +824,14 @@ def __new__(cls, aggregate, multiindex): self.multiindex = multiindex self.indirect_children = tuple(i.expression for i in self.multiindex if isinstance(i, VariableIndex)) - new_indices = [] - for i in multiindex: - if isinstance(i, Index): - new_indices.append(i) - elif isinstance(i, VariableIndex): - new_indices.extend(i.expression.free_indices) + new_indices = tuple(chain.from_iterable(map(_index_free_indices, multiindex))) self.free_indices = unique(aggregate.free_indices + tuple(new_indices)) return self def index_ordering(self): """Running indices in the order of indexing in this node.""" - free_indices = [] - for i in self.multiindex: - if isinstance(i, Index): - free_indices.append(i) - elif isinstance(i, VariableIndex): - free_indices.extend(i.expression.free_indices) - return tuple(free_indices) + return tuple(chain.from_iterable(map(_index_free_indices, self.multiindex))) class FlexiblyIndexed(Scalar): @@ -912,6 +972,71 @@ def __new__(cls, expression, multiindex): return self +class FlattenedTensor(Node): + """Lexicographically flattened view of a jagged tensor. + + Parameters + ---------- + expression : Node + Scalar expression indexed by ``multiindex``. + multiindex : tuple of Index + Rectangular or jagged tensor axes, in flattening order. + + """ + + __slots__ = ('children', 'multiindex', 'shape') + __back__ = ('multiindex',) + + def __init__(self, expression: Node, + multiindex: tuple[Index, ...]) -> None: + assert not expression.shape + multiindex = tuple(multiindex) + assert set(multiindex) <= set(expression.free_indices) + self.children = (expression,) + self.multiindex = multiindex + self.shape = (len(jagged_lattice(multiindex)),) + self.free_indices = unique(set(expression.free_indices) - set(multiindex)) + + def lattice_points(self) -> numpy.ndarray: + """Return the in-bounds lattice multi-indices in flat order.""" + return jagged_lattice(self.multiindex) + + +def jagged_layout(multiindex: tuple[Index, ...]) -> tuple: + """Return a structural description of a jagged iteration domain. + + Two multiindices with the same layout enumerate the same lattice, so the + layout is the cache key of `lattice_points`. + """ + positions = {} + layout = [] + for position, index in enumerate(multiindex): + parents = tuple(positions[parent] for parent in index.parents) + layout.append((index.extent, parents)) + positions[index] = position + return tuple(layout) + + +@lru_cache(maxsize=128) +def lattice_points(layout: tuple) -> numpy.ndarray: + """Enumerate one structural jagged iteration domain.""" + points = [] + for alpha in numpy.ndindex(*(extent for extent, _ in layout)): + if all(alpha[position] < extent + - sum(alpha[parent] for parent in parents) + for position, (extent, parents) in enumerate(layout)): + points.append(alpha) + points = numpy.asarray(points).reshape(len(points), len(layout)) + points.flags.writeable = False + return points + + +def jagged_lattice(multiindex: tuple[Index, ...]) -> numpy.ndarray: + """All lattice points of ``multiindex``'s iteration domain, honouring + `JaggedIndex` bounds, as an integer array of shape (npoint, dim).""" + return lattice_points(jagged_layout(multiindex)) + + class IndexSum(Scalar): __slots__ = ('children', 'multiindex') __back__ = ('multiindex',) @@ -923,7 +1048,9 @@ def __new__(cls, summand, multiindex): return summand # Unroll singleton sums - unroll = tuple(index for index in multiindex if index.extent <= 1) + unroll = tuple( + index for index in multiindex + if index.extent <= 1 and not index.parents) if unroll: assert numpy.prod([index.extent for index in unroll]) == 1 summand = Indexed(ComponentTensor(summand, unroll), @@ -1065,6 +1192,28 @@ def shape(self): return (int(sum(numpy.prod(child.shape, dtype=int) for child in self.children)),) +@lru_cache(maxsize=128) +def _permutation_source(index): + """Return the source index and table for a tabulated permutation.""" + if not isinstance(index, VariableIndex): + return None + expression = index.expression + if not isinstance(expression, Indexed) or len(expression.multiindex) != 1: + return None + source, = expression.multiindex + table, = expression.children + if (not isinstance(source, Index) or source.extent is None + or not isinstance(table, Literal)): + return None + entries = table.array + if (entries.shape != (source.extent,) + or not numpy.all(entries < source.extent) + or not numpy.all(numpy.bincount( + entries, minlength=source.extent) == 1)): + return None + return source, table + + class Delta(Scalar, Terminal): __slots__ = ('i', 'j') __front__ = ('i', 'j') @@ -1081,6 +1230,13 @@ def __new__(cls, i, j, dtype=None): if i == j: return one + # A shared bijection preserves equality: delta(p(i), p(j)) = delta(i, j). + source_i = _permutation_source(i) + source_j = _permutation_source(j) + if (source_i is not None and source_j is not None + and source_i[1] == source_j[1]): + return Delta(source_i[0], source_j[0], dtype=dtype) + # Fixed indices if isinstance(i, Integral) and isinstance(j, Integral): return one if i == j else Zero() @@ -1169,10 +1325,15 @@ def __rmod__(self, other): def unique(indices): """Sorts free indices and eliminates duplicates. + Indices are ordered by creation, not by :func:`id`. Object addresses + depend on what else the process has allocated, so an address ordering + would make every ``free_indices`` tuple, and hence the loop order + chosen downstream, depend on which forms were compiled beforehand. + :arg indices: iterable of indices :returns: sorted tuple of unique free indices """ - return tuple(sorted(set(indices), key=id)) + return tuple(sorted(set(indices), key=attrgetter("count"))) def index_sum(expression, indices): diff --git a/gem/impero_utils.py b/gem/impero_utils.py index ea91ab17..65790d64 100644 --- a/gem/impero_utils.py +++ b/gem/impero_utils.py @@ -10,7 +10,7 @@ from itertools import chain, groupby from gem.node import traversal, collect_refcount -from gem import gem, impero as imp, optimise, scheduling +from gem import gem, impero as imp, jagged, optimise, scheduling # ImperoC is named tuple for C code generation. @@ -30,6 +30,7 @@ class NoopError(Exception): def preprocess_gem(expressions, replace_delta=True, remove_componenttensors=True): """Lower GEM nodes that cannot be translated to C directly.""" + expressions = jagged.replace_flattened(expressions) if remove_componenttensors: expressions = optimise.remove_componenttensors(expressions) if replace_delta: @@ -72,8 +73,36 @@ def nonzero(assignment): get_indices = lambda expr: apply_ordering(expr.free_indices) + def get_loop_indices(expr: gem.Node) -> tuple[gem.Index, ...]: + """Return every explicit loop axis used to evaluate an expression. + + Parameters + ---------- + expr + GEM expression being scheduled. + + Returns + ------- + tuple of gem.Index + Free indices followed by bound value indices in global loop + order. + + Notes + ----- + A ``ComponentTensor`` binds its multi-index in GEM, but evaluating + the tensor still executes that index as a value loop. Exposing the + loop to Impero lets several tensor outputs share scalar work within + one fused loop instead of materializing that work as arrays. + + """ + indices = expr.free_indices + if isinstance(expr, gem.ComponentTensor): + indices = (*indices, *expr.multiindex) + return apply_ordering(indices) + # Build operation ordering - ops = scheduling.emit_operations(assignments, get_indices, emit_return_accumulate) + ops = scheduling.emit_operations( + assignments, get_loop_indices, emit_return_accumulate) # Empty kernel if len(ops) == 0: @@ -83,7 +112,7 @@ def nonzero(assignment): ops = inline_temporaries(expressions, ops) # Build Impero AST - tree = make_loop_tree(ops, get_indices) + tree = make_loop_tree(ops, get_loop_indices) # Collect temporaries temporaries = collect_temporaries(tree) @@ -97,9 +126,25 @@ def nonzero(assignment): def make_prefix_ordering(indices, prefix_ordering): """Creates an ordering of ``indices`` which starts with those - indices in ``prefix_ordering``.""" + indices in ``prefix_ordering``. A `gem.JaggedIndex` is placed after + its parents, so that its loop nests inside theirs and the jagged + bound can be tightened.""" # Need to return deterministically ordered indices - return tuple(prefix_ordering) + tuple(k for k in indices if k not in prefix_ordering) + ordering = tuple(prefix_ordering) + tuple(k for k in indices if k not in prefix_ordering) + result = [] + seen = set() + + def visit(k): + if k not in seen: + seen.add(k) + for parent in getattr(k, 'parents', ()): + if parent in ordering: + visit(parent) + result.append(k) + + for k in ordering: + visit(k) + return tuple(result) def make_index_orderer(index_ordering): @@ -203,11 +248,18 @@ def place_declarations(tree, temporaries, get_indices): numbering = {t: n for n, t in enumerate(temporaries)} assert len(numbering) == len(temporaries) - # Collect the total number of temporary references + # Collect the total number of temporary references. Impero is a + # tree, so structurally equal subtrees still represent distinct + # executions and every occurrence must be visited. The generic GEM + # traversal is DAG-oriented and deliberately skips equal nodes. total_refcount = collections.Counter() - for node in traversal((tree,)): + pending = [tree] + while pending: + node = pending.pop() if isinstance(node, imp.Terminal): total_refcount.update(temp_refcount(numbering, node)) + else: + pending.extend(reversed(node.children)) assert set(total_refcount) == set(temporaries) # Result diff --git a/gem/interpreter.py b/gem/interpreter.py index b2dd609e..5b23694a 100644 --- a/gem/interpreter.py +++ b/gem/interpreter.py @@ -8,6 +8,7 @@ import itertools from gem import gem, node +from gem.jagged import replace_flattened from gem.optimise import replace_delta __all__ = ("evaluate", ) @@ -383,6 +384,7 @@ def evaluate(expressions, bindings=None): exprs = tuple(expressions) except TypeError: exprs = (expressions, ) + exprs = replace_flattened(exprs) mapper = node.Memoizer(_evaluate) mapper.bindings = bindings if bindings is not None else {} return list(map(mapper, exprs)) diff --git a/gem/jagged.py b/gem/jagged.py new file mode 100644 index 00000000..8a34c754 --- /dev/null +++ b/gem/jagged.py @@ -0,0 +1,465 @@ +"""Jagged tensors: the geometry of a `JaggedIndex` lattice, and the +rewrites that trade a flat axis for the lattice it enumerates. + +`gem.gem` defines what a jagged domain *is* and enumerates its points. This +module builds on that: it compacts a product of simplex lattices into +storage, and it rewrites contractions and assignments that gather a +`FlattenedTensor` so that the ordinary contraction optimiser can see the +tensor-product factors again. +""" + +from collections import OrderedDict, defaultdict +from collections.abc import Iterable +from functools import lru_cache, partial + +import numpy + +from gem.gem import (ComponentTensor, Delta, FlattenedTensor, Index, IndexSum, + Indexed, JaggedIndex, Literal, Node, Sum, VariableIndex, + jagged_lattice, jagged_layout, lattice_points, uint_type) +from gem.node import Memoizer, MemoizerArg, reuse_if_untouched, traversal +from gem.optimise import (delta_elimination, distribute_sum, + filtered_replace_indices, make_sum, sum_factorise, + traverse_product, traverse_sum) + + +def _index_components(indices: tuple[Index, ...]) -> tuple[tuple, ...]: + """Partition indices into connected parent domains.""" + index_set = frozenset(indices) + neighbours = {index: set() for index in indices} + for index in indices: + for parent in index.parents: + if parent in index_set: + neighbours[index].add(parent) + neighbours[parent].add(index) + + components = [] + unseen = set(indices) + for seed in indices: + if seed not in unseen: + continue + component = {seed} + pending = [seed] + unseen.remove(seed) + while pending: + for index in neighbours[pending.pop()] & unseen: + unseen.remove(index) + component.add(index) + pending.append(index) + components.append(tuple(index for index in indices + if index in component)) + return tuple(components) + + +def _is_simplex_lattice(component: tuple[Index, ...]) -> bool: + """Check whether indices describe a nested simplex lattice.""" + return all( + isinstance(index, JaggedIndex) + and index.extent == component[0].extent + and index.parents == component[:position] + for position, index in enumerate(component)) + + +def compact_index_layout( + indices: tuple[Index, ...]) -> tuple[tuple[int, ...], tuple]: + """Compact a product of independent simplex lattices. + + Parameters + ---------- + indices + Indices in loop order. + + Returns + ------- + shape + Compact storage shape. + layout + Scalar indices and compact simplex components. + + Notes + ----- + A simplex lattice is stored along one compact dimension holding just + its lattice points, in the same lexicographic order that + `FlattenedTensor` flattens a jagged tensor. Rectangular padding is + exponential in the lattice dimension, so it is avoided whenever the + lattice is smaller than the box enclosing it. + + """ + shape = [] + layout = [] + for component in _index_components(indices): + points = _compact_extent(component) + if points is None: + shape.extend(index.extent for index in component) + layout.extend(component) + else: + shape.append(points) + layout.append(component) + return tuple(shape), tuple(layout) + + +def _compact_extent(component: tuple[Index, ...]) -> int | None: + """Return the compact extent of a simplex lattice, or None to pad. + + A lattice that already fills its box gains nothing from compaction: + the rank lookup would just be the identity. + + """ + if not _is_simplex_lattice(component): + return None + points = len(jagged_lattice(component)) + if points >= numpy.prod([index.extent for index in component]): + return None + return points + + +@lru_cache(maxsize=128) +def _lattice_ranks(layout: tuple) -> numpy.ndarray: + """Rank of every point of one structural jagged iteration domain.""" + points = lattice_points(layout) + ranks = numpy.zeros(tuple(extent for extent, _ in layout), dtype=uint_type) + ranks[tuple(points.T)] = numpy.arange(len(points), dtype=uint_type) + ranks.flags.writeable = False + return ranks + + +def simplex_lattice_ranks(component: tuple[Index, ...]) -> numpy.ndarray: + """Tabulate the compact rank of every point of a simplex lattice. + + Parameters + ---------- + component + Nested simplex lattice indices, in loop order. + + Returns + ------- + numpy.ndarray + Rectangular table of the lexicographic rank of each lattice + point, indexed by the lattice indices themselves. Entries + outside the jagged bounds are never read and are set to zero. + + """ + return _lattice_ranks(jagged_layout(component)) + + +def _clone_multiindex(multiindex: Iterable[Index]) -> tuple[Index, ...]: + """Clone an index tuple while preserving its internal jagged parents.""" + clones = {} + for index in multiindex: + if isinstance(index, JaggedIndex): + parents = tuple(clones[parent] for parent in index.parents) + clones[index] = JaggedIndex(extent=index.extent, parents=parents) + else: + clones[index] = Index(extent=index.extent) + return tuple(clones[index] for index in multiindex) + + +def _replace_gathers(node: Node, self, subst: tuple) -> Node: + """Replace selected flat gathers, then apply ordinary index substitution.""" + try: + return self.replacements[node] + except KeyError: + return filtered_replace_indices(node, self, subst) + + +def _flattened_layout(gather: Indexed) -> tuple: + """Return a structural key for a flattened tensor's iteration lattice.""" + tensor, = gather.children + positions = {index: position + for position, index in enumerate(tensor.multiindex)} + return tuple((type(index), index.extent, + tuple(positions[parent] for parent in index.parents)) + for index in tensor.multiindex) + + +def _flat_index_bijection( + index, extent: int) -> tuple[Index, tuple[int, ...] | None] | None: + """Identify a direct or compile-time bijective flat index. + + Parameters + ---------- + index + Index of a flattened tensor. + extent + Length of the flattened tensor. + + Returns + ------- + tuple or None + Source index and its forward permutation. A direct index has no + permutation. + + """ + if isinstance(index, Index): + return (index, None) if index.extent == extent else None + if not isinstance(index, VariableIndex): + return None + + expression = index.expression + if not (isinstance(expression, Indexed) + and len(expression.multiindex) == 1 + and isinstance(expression.multiindex[0], Index) + and isinstance(expression.children[0], Literal)): + return None + source, = expression.multiindex + table, = expression.children + if table.shape != (source.extent,) or source.extent != extent: + return None + + permutation = tuple(map(int, table.array)) + if tuple(sorted(permutation)) != tuple(range(extent)): + return None + return source, permutation + + +def _find_unflattenable_index( + nodes: Iterable[Node], + indices: Iterable[Index]) -> tuple | None: + """Find compatible flat gathers at one unconstrained index. + + An index occurring in a :class:`Delta` is not a candidate. Cancelling + that delta is cheaper than replacing the flat index with a full lattice + loop and an indirect comparison. + """ + indices = tuple(indices) + index_set = frozenset(indices) + constrained = set() + gathers = defaultdict(lambda: defaultdict(OrderedDict)) + for node in nodes: + if isinstance(node, Delta): + constrained.update(node.free_indices) + elif (isinstance(node, Indexed) + and len(node.multiindex) == 1 + and isinstance(node.children[0], FlattenedTensor)): + tensor, = node.children + bijection = _flat_index_bijection( + node.multiindex[0], tensor.shape[0]) + if bijection is None: + continue + source, permutation = bijection + if source in index_set: + key = _flattened_layout(node), permutation + gathers[source][key].setdefault(node) + + for index in indices: + groups = gathers[index] + if len(groups) == 1 and index not in constrained: + layout, candidates = next(iter(groups.items())) + return index, layout, tuple(candidates) + return None + + +def _prepare_unflattening( + gathers: tuple[Indexed, ...], + source: Index) -> tuple[MemoizerArg, tuple, VariableIndex]: + """Prepare one joint rewrite of compatible flat gathers. + + Each flattened tensor is inlined on the same fresh lattice multiindex. + The returned index maps each lattice point to the original source index. + A compile-time permutation is inverted before the index is used in the + return variable. + """ + gather = gathers[0] + tensor, = gather.children + assert all(_flattened_layout(other) == _flattened_layout(gather) + for other in gathers) + bijection = _flat_index_bijection( + gather.multiindex[0], tensor.shape[0]) + assert bijection is not None and bijection[0] == source + permutation = bijection[1] + assert all(_flat_index_bijection( + other.multiindex[0], other.children[0].shape[0]) == bijection + for other in gathers) + multiindex = _clone_multiindex(tensor.multiindex) + replacer = MemoizerArg(filtered_replace_indices) + mapper = MemoizerArg(_replace_gathers) + mapper.replacements = {} + for other in gathers: + tensor, = other.children + mapper.replacements[other] = replacer( + tensor.children[0], tuple(zip(tensor.multiindex, multiindex))) + + shape = tuple(index.extent for index in multiindex) + points = gather.children[0].lattice_points() + ordering = numpy.zeros(shape, dtype=uint_type) + ordering[tuple(points.T)] = numpy.arange(len(points)) + if permutation is not None: + inverse = numpy.empty(len(permutation), dtype=uint_type) + inverse[numpy.asarray(permutation)] = numpy.arange( + len(permutation), dtype=uint_type) + ordering = inverse[ordering] + source_index = VariableIndex(Indexed( + Literal(ordering, dtype=uint_type), multiindex)) + return mapper, multiindex, source_index + + +def _separable_sum(node: Node, indices: frozenset[Index]) -> bool: + """Whether distributing a sum exposes smaller index dependencies.""" + if not isinstance(node, Sum): + return False + involved = indices.intersection(node.free_indices) + return bool(involved) and any( + any(indices.intersection(factor.free_indices) < involved + for factor in traverse_product(summand)[1]) + for summand in traverse_sum(node)) + + +def _unflatten_contracted_terms( + summand: Node, index: Index) -> tuple[list, list[Node]]: + """Unflatten additive terms containing a gather at a contracted index.""" + rewritten = [] + leftover = [] + for term in traverse_sum(summand): + candidate = _find_unflattenable_index( + traversal((term,)), (index,)) + if candidate is None: + leftover.append(term) + continue + _, _, gathers = candidate + mapper, multiindex, source_index = _prepare_unflattening( + gathers, index) + term = mapper(term, ((index, source_index),)) + own = frozenset(multiindex) + predicate = partial(_separable_sum, indices=own) + rewritten.extend( + (multiindex, piece) + for piece in distribute_sum(term, predicate=predicate)) + return rewritten, leftover + + +def _unflatten_contractions(node: Node, self) -> Node: + """Memoizer callback for flat indices bound by an :class:`IndexSum`.""" + node = reuse_if_untouched(node, self) + if not isinstance(node, IndexSum): + return node + summand, = node.children + for index in node.multiindex: + rewritten, leftover = _unflatten_contracted_terms(summand, index) + if not rewritten: + continue + rest = tuple(other for other in node.multiindex if other != index) + pieces = [] + for own, term in rewritten: + term = self(IndexSum( + term, own + tuple(i for i in rest if i in term.free_indices))) + indices, factors = traverse_product(term) + indices, factors = delta_elimination(indices, factors) + pieces.append(sum_factorise(indices, factors)) + if leftover: + residual = make_sum(leftover) + indices = tuple(i for i in (index,) + rest + if i in residual.free_indices) + pieces.append(self(IndexSum(residual, indices))) + return make_sum(pieces) + return node + + +def unflatten(expression: Node) -> Node: + """Replace flat contractions by loops over their jagged lattices. + + A contraction over the flat axis of a :class:`FlattenedTensor` hides its + tensor-product factors. This rewrite substitutes the tensor's own + (possibly jagged) multiindex for that flat axis, enabling the ordinary + contraction optimizer to recover sum factorisation. + """ + if not any(isinstance(node, FlattenedTensor) + for node in traversal((expression,))): + return expression + return Memoizer(_unflatten_contractions)(expression) + + +def _has_flat_gather(nodes: Iterable[Node], + indices: Iterable[Index]) -> bool: + """Whether ``nodes`` gather a flattened tensor at one of ``indices``.""" + indices = frozenset(indices) + return any(isinstance(node, Indexed) + and len(node.multiindex) == 1 + and node.multiindex[0] in indices + and isinstance(node.children[0], FlattenedTensor) + for node in nodes) + + +def unflatten_free_indices( + variable: Node, expression: Node, *, + split_separable_sums: bool) -> tuple[list[tuple[Node, Node]], bool]: + """Replace flattened gathers at free indices of one assignment. + + Compatible gathers are rewritten together. Distribution is restricted + to deltas until a lattice has been exposed; the optional legacy path then + splits separable sums immediately. + """ + pending = [(variable, expression)] + outputs = [] + changed = False + while pending: + current_variable, current_expression = pending.pop() + nodes = tuple(traversal((current_expression,))) + candidate = _find_unflattenable_index( + nodes, current_variable.free_indices) + if candidate is not None: + index, layout, gathers = candidate + groups = OrderedDict([ + ((index, layout), ([current_expression], list(gathers)))]) + elif not _has_flat_gather(nodes, current_variable.free_indices): + outputs.append((current_variable, current_expression)) + continue + else: + predicate = (lambda node: isinstance(node, Delta)) \ + if any(isinstance(node, Delta) for node in nodes) else None + groups = OrderedDict() + for term in distribute_sum( + current_expression, predicate=predicate): + candidate = _find_unflattenable_index( + traversal((term,)), current_variable.free_indices) + key = candidate[:2] if candidate is not None else None + terms, gathers = groups.setdefault(key, ([], [])) + terms.append(term) + if candidate is not None: + gathers.extend(candidate[2]) + + for key, (terms, gathers) in groups.items(): + term = make_sum(terms) + if key is None: + outputs.append((current_variable, term)) + continue + + changed = True + index, _ = key + gathers = tuple(OrderedDict.fromkeys(gathers)) + mapper, multiindex, source_index = _prepare_unflattening( + gathers, index) + substitution = ((index, source_index),) + new_variable = MemoizerArg(filtered_replace_indices)( + current_variable, substitution) + new_expression = mapper(term, substitution) + if split_separable_sums: + lattice_indices = frozenset(multiindex) + predicate = partial(_separable_sum, indices=lattice_indices) + pending.extend( + (new_variable, piece) + for piece in distribute_sum( + new_expression, predicate=predicate)) + else: + pending.append((new_variable, new_expression)) + return outputs, changed + + +def _replace_flattened(node, self): + node = reuse_if_untouched(node, self) + if not isinstance(node, FlattenedTensor): + return node + expression, = node.children + points = node.lattice_points() + index = Index(extent=node.shape[0]) + subst = tuple( + (axis, VariableIndex(Indexed( + Literal(points[:, position], dtype=uint_type), (index,)))) + for position, axis in enumerate(node.multiindex)) + body = MemoizerArg(filtered_replace_indices)(expression, subst) + return ComponentTensor(body, (index,)) + + +def replace_flattened(expressions): + """Lower remaining flattened tensors to indirect flat-index gathers.""" + mapper = Memoizer(_replace_flattened) + return [mapper(expression) for expression in expressions] diff --git a/gem/optimise.py b/gem/optimise.py index 8f2e6914..437949f7 100644 --- a/gem/optimise.py +++ b/gem/optimise.py @@ -4,7 +4,7 @@ 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 itertools import combinations, permutations from numbers import Integral import numpy @@ -17,7 +17,7 @@ Product, Sum, Comparison, Conditional, Division, Index, IndexBase, VariableIndex, Indexed, FlexiblyIndexed, IndexSum, ComponentTensor, ListTensor, Delta, - partial_indexed, one) + jagged_lattice, partial_indexed, one) @singledispatch @@ -767,6 +767,33 @@ def traverse_sum(expression, stop_at=None): return result +def _distributed_indexsum_term( + term: Node, indices: tuple[Index, ...]) -> Node: + """Preserve a joint contraction domain after distributing one term.""" + active = tuple(index for index in indices if index in term.free_indices) + missing = tuple(index for index in indices if index not in active) + if not missing: + return IndexSum(term, indices) + + if any(index.parents for index in indices): + points = jagged_lattice(indices) + if active: + positions = tuple(indices.index(index) for index in active) + multiplicity = numpy.zeros( + tuple(index.extent for index in active)) + numpy.add.at( + multiplicity, + tuple(points[:, position] for position in positions), + 1, + ) + factor = Indexed(Literal(multiplicity), active) + else: + factor = Literal(float(len(points))) + else: + factor = index_space_literal(missing) + return IndexSum(Product(term, factor), active) + + def distribute_sum(expr: Node, predicate: Callable[[Node], bool]) -> list[Node]: """Distribute selected sums through products and contractions. @@ -799,9 +826,7 @@ def distribute_sum(expr: Node, predicate: Callable[[Node], bool]) -> list[Node]: elif isinstance(node, IndexSum): body, = node.children results[key] = [ - IndexSum(term, tuple( - index for index in node.multiindex - if index in term.free_indices)) + _distributed_indexsum_term(term, node.multiindex) for term in results[id(body)]] else: # Product a, b = node.children @@ -1030,66 +1055,6 @@ def visit(node, self): return Memoizer(visit)(expression) -def contraction(expression): - """Optimise the contractions of the tensor product at the root of - the expression, including: - - - IndexSum-Delta cancellation - - Sum factorisation - - This routine was designed with finite element coefficient - evaluation in mind. - """ - - # Common memoizer to remove ComponentTensors - index_replacer = MemoizerArg(filtered_replace_indices) - - # Eliminate annoying ComponentTensors - expression = index_replacer(expression, ()) - - # Flatten product tree, eliminate deltas, sum factorise - def rebuild(expression): - root = expression - # The contraction at the root is always broken up, as that is the - # one being optimised - keep = repeated_contractions(expression) - sum_indices, factors = traverse_product( - expression, index_replacer=index_replacer, - stop_at=lambda e: e is not root and e in keep) - sum_indices, factors = pull_back_indirect_delta( - sum_indices, factors, index_replacer) - sum_indices, factors = delta_elimination( - sum_indices, factors, index_replacer=index_replacer) - factors = [index_replacer(f, ()) for f in factors] - return sum_factorise(sum_indices, factors) - - # Sometimes the value shape is composed as a ListTensor, which - # could get in the way of decomposing factors. In particular, - # this is the case for H(div) and H(curl) conforming tensor - # product elements. So if ListTensors are used, they are pulled - # out to be outermost, so we can straightforwardly factorise each - # of its entries. - lt_fis = OrderedDict() # ListTensor free indices - for node in traversal((expression,)): - if isinstance(node, Indexed): - child, = node.children - if isinstance(child, ListTensor): - lt_fis.update(zip_longest(node.multiindex, ())) - lt_fis = tuple(index for index in lt_fis if index in expression.free_indices) - - if lt_fis: - # Rebuild each split component - tensor = ComponentTensor(expression, lt_fis) - entries = [Indexed(tensor, zeta) for zeta in numpy.ndindex(tensor.shape)] - entries = [index_replacer(e, ()) for e in entries] - return Indexed(ListTensor( - numpy.array(list(map(rebuild, entries))).reshape(tensor.shape) - ), lt_fis) - else: - # Rebuild whole expression at once - return rebuild(expression) - - @singledispatch def _replace_delta(node, self): raise AssertionError("cannot handle type %s" % type(node)) diff --git a/gem/scheduling.py b/gem/scheduling.py index 7e361cd8..c03b543c 100644 --- a/gem/scheduling.py +++ b/gem/scheduling.py @@ -135,9 +135,11 @@ def handle(ops, push, decref, node): elif isinstance(node, impero.Return): ops.append(node) decref(node.expression) + decref(node.variable) elif isinstance(node, impero.ReturnAccumulate): ops.append(node) decref(node.indexsum.children[0]) + decref(node.variable) else: raise AssertionError("no handler for node type %s" % type(node)) @@ -157,8 +159,10 @@ def emit_operations(assignments, get_indices, emit_return_accumulate=True): :returns: list of Impero terminals correctly ordered to evaluate the assignments """ - # Prepare reference counts - refcount = collect_refcount([e for v, e in assignments]) + # Prepare reference counts. Return variables participate too: their + # index expressions (e.g. `gem.VariableIndex` gather tables) may need + # evaluating just like the right-hand sides. + refcount = collect_refcount([node for pair in assignments for node in pair]) # Stage return operations staging = [] diff --git a/test/FIAT/unit/test_bernstein.py b/test/FIAT/unit/test_bernstein.py index 2b4714a8..da40e41e 100644 --- a/test/FIAT/unit/test_bernstein.py +++ b/test/FIAT/unit/test_bernstein.py @@ -20,9 +20,10 @@ import numpy import pytest -from FIAT.reference_element import ufc_simplex from FIAT.bernstein import Bernstein +from FIAT.polynomial_set import mis from FIAT.quadrature_schemes import create_quadrature +from FIAT.reference_element import multiindex_equal, ufc_simplex D02 = numpy.array([ @@ -74,10 +75,13 @@ def test_bernstein_2nd_derivatives(): points = rule.get_points() actual = elem.tabulate(2, points) + barycentric_indices = list(mis(3, degree)) + ordering = [barycentric_indices.index(alpha) + for alpha in multiindex_equal(3, degree)] - assert numpy.allclose(D02, actual[(0, 2)]) - assert numpy.allclose(D11, actual[(1, 1)]) - assert numpy.allclose(D20, actual[(2, 0)]) + assert numpy.allclose(D02[ordering], actual[(0, 2)]) + assert numpy.allclose(D11[ordering], actual[(1, 1)]) + assert numpy.allclose(D20[ordering], actual[(2, 0)]) if __name__ == '__main__': diff --git a/test/finat/test_dual_basis.py b/test/finat/test_dual_basis.py index c995cfa9..e8c345d8 100644 --- a/test/finat/test_dual_basis.py +++ b/test/finat/test_dual_basis.py @@ -2,6 +2,7 @@ import numpy import finat import gem +import gem.driver from FIAT import ufc_simplex from gem.interpreter import evaluate @@ -62,7 +63,7 @@ def coefficient_evaluation(element, ps, dofs): table = element.basis_evaluation(0, ps)[(0,) * dim] dofs = gem.Literal(dofs.reshape([index.extent for index in beta])) value = gem.Product(gem.Indexed(table, beta + zeta), gem.Indexed(dofs, beta)) - return gem.ComponentTensor(gem.optimise.contraction(gem.IndexSum(value, beta)), zeta) + return gem.ComponentTensor(gem.driver.contraction(gem.IndexSum(value, beta)), zeta) def nodal_values(element, fn): diff --git a/test/finat/test_point_evaluation.py b/test/finat/test_point_evaluation.py index 72bfae2a..2efa20e2 100644 --- a/test/finat/test_point_evaluation.py +++ b/test/finat/test_point_evaluation.py @@ -69,6 +69,48 @@ def test_point_evaluation_zany(ref_to_phys, element, degree): assert numpy.allclose(val, expected[alpha][:num_dof]) +@pytest.mark.parametrize("make_element", [ + pytest.param(lambda cell, degree: finat.Legendre(cell, degree, variant="integral"), + id="integral"), + pytest.param(finat.IntegratedLegendre, id="integrated"), + pytest.param(finat.Bernstein, id="bernstein"), +]) +@pytest.mark.parametrize('degree', [1, 4]) +def test_duffy_evaluation(cell, degree, make_element): + from finat.point_set import PointSet, CollapsedTensorProductPointSet + + dim = cell.get_spatial_dimension() + element = make_element(cell, degree) + + # Unequal point counts per axis, including the collapsed vertex eta=1 + factors = [PointSet(numpy.linspace(0, 1, 3 + axis)[:, None]) for axis in range(dim)] + ps = CollapsedTensorProductPointSet(factors) + duffy = element.duffy_evaluation(1, ps) + + dense_ps = PointSet(ps.points) + expected = element.basis_evaluation(1, dense_ps) + assert expected.keys() == duffy.keys() + + # Both tabulations use a flat dof axis. + ndof = element.space_dimension() + for alpha, table in expected.items(): + exp, = gem.interpreter.evaluate([table]) + exp = exp.broadcast(dense_ps.indices) + act, = gem.interpreter.evaluate([duffy[alpha]]) + act = act.broadcast(ps.indices).reshape(-1, ndof) + assert numpy.allclose(act, exp, rtol=1E-10, atol=1E-12) + + coefficients = numpy.random.default_rng(1).random(ndof) + index = gem.Index(extent=ndof) + contraction = gem.IndexSum( + gem.Product(gem.Indexed(gem.Literal(coefficients), (index,)), + gem.Indexed(duffy[alpha], (index,))), (index,)) + value, = gem.interpreter.evaluate([contraction]) + value = value.broadcast(ps.indices) + assert numpy.allclose(value.reshape(-1), exp @ coefficients, + rtol=1E-10, atol=1E-12) + + if __name__ == '__main__': import os pytest.main(os.path.abspath(__file__)) diff --git a/test/finat/test_quadrature.py b/test/finat/test_quadrature.py index 4d1acb9f..bae6b46e 100644 --- a/test/finat/test_quadrature.py +++ b/test/finat/test_quadrature.py @@ -1,6 +1,9 @@ +import numpy import pytest +import gem from FIAT import ufc_cell +from FIAT.quadrature_schemes import create_quadrature as fiat_scheme from finat.quadrature import make_quadrature @@ -17,3 +20,24 @@ def test_quadrature_rules_are_hashable(cell_name): assert hash(quadrature1) == hash(quadrature2) assert repr(quadrature1) == repr(quadrature2) assert quadrature1 == quadrature2 + + +@pytest.mark.parametrize("cell_name", ["interval", "triangle", "tetrahedron"]) +@pytest.mark.parametrize("degree", [3, 8]) +def test_collapsed_quadrature(cell_name, degree): + ref_cell = ufc_cell(cell_name) + dim = ref_cell.get_spatial_dimension() + rule = make_quadrature(ref_cell, degree, scheme="collapsed") + ps = rule.point_set + result, = gem.interpreter.evaluate([rule.weight_expression]) + weights = result.broadcast(ps.indices).ravel() + + reference = fiat_scheme(ref_cell, degree, "canonical") + ref_points = reference.get_points() + ref_weights = reference.get_weights() + for alpha in numpy.ndindex((degree + 1,) * dim): + if sum(alpha) > degree: + continue + monomial = lambda pts: numpy.prod(pts ** numpy.asarray(alpha), axis=-1) + exact = numpy.dot(ref_weights, monomial(ref_points)) + assert numpy.allclose(numpy.dot(weights, monomial(ps.points)), exact) diff --git a/test/finat/test_zany_mapping.py b/test/finat/test_zany_mapping.py index f0742caf..3596f447 100644 --- a/test/finat/test_zany_mapping.py +++ b/test/finat/test_zany_mapping.py @@ -7,10 +7,20 @@ from gem.interpreter import evaluate from gem.node import traversal -from gem.optimise import contraction +from gem.driver import contraction from finat.physically_mapped import MappedTabulation, PhysicallyMappedElement +def test_numeric_zero_mapped_tabulation_is_sparse() -> None: + """Discard numeric zeros when constructing mapped rows.""" + matrix = gem.Literal(np.eye(3)) + mapped = MappedTabulation(matrix, {None: gem.Literal(np.eye(3))}) + + assert mapped._width == 1 + assert np.array_equal(mapped._columns.array[:, 0], np.arange(3)) + assert mapped._values.shape == (1,) + + def test_sparse_mapped_tabulation(): """Apply a sparse basis map at the cost of its nonzeros.""" coefficient = gem.Variable("coefficient", ()) diff --git a/test/gem/test_simplify.py b/test/gem/test_simplify.py index 4343b1a4..fb74710c 100644 --- a/test/gem/test_simplify.py +++ b/test/gem/test_simplify.py @@ -2,6 +2,23 @@ import gem import numpy +from gem import impero +from gem.coffee import monomial_sum_to_expression +from gem.flop_count import count_flops +from gem.impero_utils import (collect_temporaries, compile_gem, + place_declarations) +from gem.node import traversal +from gem.interpreter import evaluate +from gem.driver import contraction, unflatten_returns +from gem.optimise import ( + cancel_nested_deltas, + distribute_sum, + preserve_linear_maps, +) +from gem.refactorise import (ATOMIC, COMPOUND, OTHER, + collect_monomials) +from gem.gem import jagged_lattice + @pytest.fixture def A(): @@ -19,6 +36,26 @@ def X(): return gem.Variable("X", (2, 2)) +def test_compact_simplex_lattice_product(): + """Compact independent simplex lattices and preserve lexicographic rank.""" + p = gem.JaggedIndex(extent=4) + q = gem.JaggedIndex(extent=4, parents=(p,)) + r = gem.JaggedIndex(extent=4) + s = gem.JaggedIndex(extent=4, parents=(r,)) + + shape, layout = gem.compact_index_layout((p, q, r, s)) + + assert shape == (10, 10) + assert layout == ((p, q), (r, s)) + # A lattice filling its box gains nothing from a rank lookup. + assert gem.compact_index_layout((p,)) == ((4,), (p,)) + + # Ranks agree with the order in which FlattenedTensor flattens. + ranks = gem.simplex_lattice_ranks((p, q)) + points = jagged_lattice((p, q)) + assert [int(ranks[tuple(point)]) for point in points] == list(range(10)) + + def test_listtensor_from_indexed(X): k = gem.Index() elems = [gem.Indexed(X, (k, *i)) for i in numpy.ndindex(X.shape[1:])] @@ -71,6 +108,38 @@ def test_componenttensor_from_diagonal(): assert result == gem.Sum(gem.Indexed(a, (0,)), gem.Indexed(b, (1,))) +def test_componenttensor_flop_count(): + i = gem.Index(extent=3) + j = gem.Index(extent=3) + x = gem.Variable("x", (3,)) + result = gem.Variable("result", (3,)) + tensor = gem.ComponentTensor(2 * gem.Indexed(x, (i,)), (i,)) + expression = gem.Indexed(tensor, (j,)) + impero_c = compile_gem( + [(gem.Indexed(result, (j,)), expression)], (j,)) + + assert count_flops(impero_c) == 6 + + +def test_componenttensor_sharing_uses_scalar_temporary(): + """Keep shared work inside a component tensor's value loop.""" + i = gem.Index(extent=3) + j = gem.Index(extent=3) + x = gem.Variable("x", (3,)) + result = gem.Variable("result", (3,)) + shared = 2 * gem.Indexed(x, (i,)) + positive = gem.ComponentTensor(shared + 1, (i,)) + negative = gem.ComponentTensor(shared - 1, (i,)) + expression = gem.Indexed(positive, (j,)) \ + + gem.Indexed(negative, (j,)) + + impero_c = compile_gem( + [(gem.Indexed(result, (j,)), expression)], (j,)) + + assert shared in impero_c.temporaries + assert impero_c.indices[shared] == () + + def test_indexed_transpose(A): i, j = gem.indices(2) ATij = gem.Indexed(A.T, (i, j)) @@ -102,6 +171,396 @@ def test_flatten_indexsum(A): assert result == expected +def test_selective_distribution(): + a = gem.Variable("a", ()) + b = gem.Variable("b", ()) + c = gem.Variable("c", ()) + i = gem.Index(extent=2) + p = gem.Index(extent=1) + row = gem.VariableIndex(gem.Indexed( + gem.Literal([0], dtype=gem.uint_type), (p,))) + delta = gem.Delta(i, row) + common = gem.Sum(a, b) + expression = gem.Product(common, gem.Sum(c, delta)) + + terms = distribute_sum( + expression, predicate=lambda node: isinstance(node, gem.Delta)) + + assert len(terms) == 2 + assert all(common in set(traversal((term,))) for term in terms) + + +def test_preserve_linear_maps_early_exit(): + """Keep a multilinear sum that contains no separate linear maps.""" + i = gem.Index(extent=2) + j = gem.Index(extent=2) + variables = [gem.Variable(f"a{k}", (2, 2)) for k in range(4)] + expression = gem.Sum(*( + gem.Indexed(variable, (i, j)) for variable in variables)) + + terms, linear_maps = preserve_linear_maps(expression, (i, j)) + + assert terms == (expression,) + assert linear_maps == () + + +def test_collect_monomials_preserves_linear_maps(): + """Keep finite element linear maps intact during factorization.""" + i = gem.Index(extent=2) + j = gem.Index(extent=2) + left = gem.Sum( + gem.Indexed(gem.Literal([1.0, 2.0]), (i,)), + gem.Indexed(gem.Literal([3.0, 5.0]), (i,))) + right = gem.Sum( + gem.Indexed(gem.Literal([7.0, 11.0]), (j,)), + gem.Indexed(gem.Literal([13.0, 17.0]), (j,))) + expression = left * right + linear_indices = frozenset((i, j)) + + def classifier(node: gem.Node) -> str: + support = linear_indices.intersection(node.free_indices) + if not support: + return OTHER + if isinstance(node, gem.Indexed): + return ATOMIC + return COMPOUND + + monomial_sum, = collect_monomials( + (expression,), classifier, linear_indices) + + monomial, = tuple(monomial_sum) + assert frozenset(monomial.atomics) == frozenset((left, right)) + expected, = evaluate([gem.ComponentTensor(expression, (i, j))]) + actual, = evaluate([gem.ComponentTensor( + monomial_sum_to_expression(monomial_sum), (i, j))]) + assert numpy.array_equal(actual.arr, expected.arr) + + +def test_constant_variable_index(): + index = gem.VariableIndex(gem.Literal(1, dtype=gem.uint_type)) + assert index == 1 + + +def test_place_declarations_counts_equal_impero_subtrees(): + """Equal Impero nodes are distinct occurrences in the loop tree.""" + expression = gem.Variable("a", ()) * gem.Variable("b", ()) + tree = impero.Block([ + impero.Evaluate(expression), + impero.Evaluate(expression), + ]) + temporaries = collect_temporaries(tree) + + declare, indices = place_declarations( + tree, temporaries, lambda node: node.free_indices) + + assert declare[tree] == [] + assert indices[expression] == () + assert all(declare[statement] for statement in tree.children) + + +def test_delta_elimination_preserves_indirect_free_index(): + i = gem.Index(extent=4) + k = gem.Index(extent=2) + entries = numpy.array([1, 3], dtype=gem.uint_type) + indirect = gem.VariableIndex(gem.Indexed( + gem.Literal(entries, dtype=gem.uint_type), (k,))) + values = gem.Literal([2.0, 3.0, 5.0, 7.0]) + expression = gem.IndexSum( + gem.Delta(i, indirect) * gem.Indexed(values, (i,)), (i,)) + + result = cancel_nested_deltas(expression) + assert result.free_indices == (k,) + actual, = evaluate([result]) + assert numpy.array_equal(actual.arr, values.array[entries]) + + +def test_unflatten_compatible_returns_together(): + extent = 3 + p = gem.Index(extent=extent) + q = gem.JaggedIndex(extent=extent, parents=(p,)) + X = gem.Variable("X", (extent, extent)) + Y = gem.Variable("Y", (extent, extent)) + ft_x = gem.FlattenedTensor(gem.Indexed(X, (p, q)), (p, q)) + ft_y = gem.FlattenedTensor(gem.Indexed(Y, (p, q)), (p, q)) + + r = gem.Index(extent=6) + result = gem.Variable("result", (6,)) + pairs = unflatten_returns([ + (gem.Indexed(result, (r,)), + gem.Sum(gem.Indexed(ft_x, (r,)), gem.Indexed(ft_y, (r,)))) + ]) + + assert len(pairs) == 1 + variable, expression = pairs[0] + assert variable.free_indices == expression.free_indices + assert len(variable.free_indices) == 2 + assert not any(isinstance(node, gem.FlattenedTensor) + for node in traversal((expression,))) + + +def test_unflatten_bijective_return_index(): + """Invert a compile-time permutation before exposing a return lattice.""" + extent = 3 + p = gem.Index(extent=extent) + q = gem.JaggedIndex(extent=extent, parents=(p,)) + X = gem.Variable("X", (extent, extent)) + table = gem.FlattenedTensor( + gem.Indexed(X, (p, q)), (p, q)) + + r = gem.Index(extent=6) + permutation = numpy.asarray([2, 0, 5, 1, 4, 3], dtype=gem.uint_type) + mapped = gem.Indexed(table, (gem.VariableIndex(gem.Indexed( + gem.Literal(permutation, dtype=gem.uint_type), (r,))),)) + output = gem.Variable("output", (6,)) + variable, optimized = unflatten_returns([ + (gem.Indexed(output, (r,)), mapped) + ])[0] + + assert len(variable.free_indices) == 2 + scatter_expression = variable.multiindex[0].expression + assert scatter_expression.children[0].shape == (extent, extent) + assert not any(isinstance(node, gem.FlattenedTensor) + for node in traversal((optimized,))) + + values = numpy.arange(extent * extent, dtype=float).reshape( + extent, extent) + expected, = evaluate([mapped], {X: values}) + actual, scatter = evaluate( + [optimized, variable.multiindex[0].expression], {X: values}) + points = table.lattice_points() + coordinates = { + index: points[:, position] + for position, index in enumerate(variable.free_indices) + } + + def sample(value): + return value.arr[tuple(coordinates[index] + for index in value.fids)] + + dense = numpy.empty(6) + dense[sample(scatter)] = sample(actual) + assert numpy.array_equal(dense, expected.broadcast((r,))) + + +def test_unflatten_factorises_local_sum(): + extent = 3 + p = gem.JaggedIndex(extent=extent) + q = gem.JaggedIndex(extent=extent, parents=(p,)) + ip, iq = gem.indices(2) + A = gem.Variable("A", (extent, 2)) + B = gem.Variable("B", (extent, extent, 2)) + C = gem.Variable("C", (extent, 2)) + D = gem.Variable("D", (extent, extent, 2)) + lattice = gem.Sum( + gem.Product(gem.Indexed(A, (p, ip)), gem.Indexed(B, (p, q, iq))), + gem.Product(gem.Indexed(C, (p, ip)), gem.Indexed(D, (p, q, iq))), + ) + table = gem.FlattenedTensor(lattice, (p, q)) + + r = gem.Index(extent=6) + w = gem.Variable("w", (6,)) + expression = gem.IndexSum( + gem.Product(gem.Indexed(table, (r,)), gem.Indexed(w, (r,))), (r,)) + result = contraction(expression) + + sums = [node for node in traversal((result,)) + if isinstance(node, gem.IndexSum)] + assert sums + assert all(len(node.multiindex) == 1 for node in sums) + assert not any(isinstance(node, gem.FlattenedTensor) + for node in traversal((result,))) + + +def test_unflatten_factorises_bilinear_arguments_together(): + """Two argument lattices are exposed before their local sums expand.""" + extent = 3 + p = gem.JaggedIndex(extent=extent) + q = gem.JaggedIndex(extent=extent, parents=(p,)) + r = gem.JaggedIndex(extent=extent) + s = gem.JaggedIndex(extent=extent, parents=(r,)) + ip, iq = gem.indices(2) + + variables = tuple( + gem.Variable(name, shape) + for name, shape in [ + ("A", (extent, 2)), + ("B", (extent, extent, 2)), + ("C", (extent, 2)), + ("D", (extent, extent, 2)), + ("E", (extent, 2)), + ("F", (extent, extent, 2)), + ("G", (extent, 2)), + ("H", (extent, extent, 2)), + ]) + A, B, C, D, E, F, G, H = variables + left = gem.FlattenedTensor(gem.Sum( + gem.Product(gem.Indexed(A, (p, ip)), + gem.Indexed(B, (p, q, iq))), + gem.Product(gem.Indexed(C, (p, ip)), + gem.Indexed(D, (p, q, iq))), + ), (p, q)) + right = gem.FlattenedTensor(gem.Sum( + gem.Product(gem.Indexed(E, (r, ip)), + gem.Indexed(F, (r, s, iq))), + gem.Product(gem.Indexed(G, (r, ip)), + gem.Indexed(H, (r, s, iq))), + ), (r, s)) + + i, j = gem.indices(2) + output = gem.Variable("output", (6, 6)) + expression = gem.IndexSum( + gem.Product(gem.Indexed(left, (i,)), + gem.Indexed(right, (j,))), + (ip, iq)) + pairs = unflatten_returns([ + (gem.Indexed(output, (i, j)), expression) + ]) + + assert len(pairs) == 1 + variable, optimized = pairs[0] + assert len(variable.free_indices) == 4 + assert all(isinstance(index, gem.JaggedIndex) + for index in variable.free_indices) + assert not any(isinstance(node, gem.FlattenedTensor) + for node in traversal((optimized,))) + assert all(len(node.multiindex) == 1 + for node in traversal((optimized,)) + if isinstance(node, gem.IndexSum)) + + rng = numpy.random.default_rng(2) + bindings = { + variable_: rng.random(variable_.shape) + for variable_ in variables + } + expected, = evaluate([expression], bindings) + actual, = evaluate([optimized], bindings) + points = left.lattice_points() + row_map, column_map = evaluate([ + index.expression for index in variable.multiindex + ]) + row = row_map.arr[points[:, 0], points[:, 1]] + column = column_map.arr[points[:, 0], points[:, 1]] + row_indices = { + index: points[:, position, None] + for position, index in enumerate(row_map.fids) + } + column_indices = { + index: points[None, :, position] + for position, index in enumerate(column_map.fids) + } + indices = tuple((row_indices | column_indices)[index] + for index in actual.fids) + values = actual.arr[indices] + dense = numpy.empty((6, 6)) + dense[row[:, None], column[None, :]] = values + assert numpy.allclose(dense, expected.broadcast((i, j))) + + +def test_distribute_sum_preserves_rectangular_multiplicity(): + """Preserve rectangular contraction multiplicity after distribution.""" + indices = tuple(gem.Index(extent=2) for _ in range(2)) + extra = gem.Index(extent=2) + + def unit(index): + return gem.Indexed(gem.Literal(numpy.ones(2)), (index,)) + + factor = gem.Sum(gem.IndexSum(unit(indices[0]) * unit(extra), (extra,)), + unit(indices[1])) + expression = gem.IndexSum(factor, indices) + terms = distribute_sum( + expression, predicate=lambda node: isinstance(node, gem.Sum)) + expression = gem.Sum(*terms) + value, = evaluate([expression]) + assert value.arr == 12 + + +def test_sum_factorise_jagged_distribution(): + """Preserve the joint jagged domain after distribution.""" + parent = gem.JaggedIndex(extent=3) + child = gem.JaggedIndex(extent=3, parents=(parent,)) + + def unit(index: gem.Index) -> gem.Node: + """Return a unit vector carrying one free index. + + Parameters + ---------- + index + Free index of the vector. + + Returns + ------- + gem.Node + Indexed unit vector. + """ + return gem.Indexed(gem.Literal(numpy.ones(3)), (index,)) + + triangle = numpy.fromfunction( + lambda i, j: j < 3 - i, (3, 3), dtype=int) + factor = gem.Sum( + unit(parent), gem.Indexed(gem.Literal(triangle), (parent, child))) + expression = gem.IndexSum(factor, (parent, child)) + terms = distribute_sum( + expression, predicate=lambda node: isinstance(node, gem.Sum)) + expression = gem.Sum(*terms) + value, = evaluate([expression]) + assert value.arr == 12 + + +def test_literal_distinguishes_dtypes(): + """Tell an index literal apart from a value literal. + + An index table holds unsigned integers and a coefficient table holds + floats. GEM memoizes on node identity, so the two must not compare + equal when they happen to hold the same number. + """ + index = gem.Literal(numpy.uint32(3), dtype=gem.uint_type) + value = gem.Literal(3.0) + + assert index.dtype != value.dtype + assert index != value + assert hash(index) != hash(value) + assert {index: "index"}.get(value) is None + + +def _churned_indices(count): + """Create indices whose addresses disagree with their creation order. + + CPython reuses freed object slots, so allocating a pool, releasing half + of it, and allocating again places later-created indices below + earlier-created ones in memory. + """ + pool = [gem.Index(extent=2) for _ in range(4 * count)] + indices = pool[::2][:count] + del pool + indices += [gem.Index(extent=2) for _ in range(count)] + return indices + + +def test_free_indices_ignore_allocation_addresses(): + """Order free indices by creation, not by memory address. + + Sorting by :func:`id` makes every ``free_indices`` tuple depend on + where the allocator happened to place each index, so compiling one form + changes the loop order chosen for the next one in the same process. + """ + indices = _churned_indices(16) + assert any(a.count < b.count and id(a) > id(b) + for a in indices for b in indices), \ + "precondition: allocator did not invert creation order" + + by_creation = lambda ids: tuple(sorted(ids, key=lambda i: i.count)) + assert gem.gem.unique(reversed(indices)) == by_creation(indices) + + expression = gem.Indexed(gem.Variable("A", (2,) * len(indices)), + tuple(indices)) + assert expression.free_indices == by_creation(indices) + + bound = tuple(indices[::2]) + free = by_creation(set(indices) - set(bound)) + assert gem.IndexSum(expression, bound).free_indices == free + assert gem.ComponentTensor(expression, bound).free_indices == free + + 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])) diff --git a/test/gem/test_sum_factorise.py b/test/gem/test_sum_factorise.py index 33ac1991..48690460 100644 --- a/test/gem/test_sum_factorise.py +++ b/test/gem/test_sum_factorise.py @@ -6,7 +6,7 @@ import gem from gem.interpreter import evaluate -from gem import cost, optimise +from gem import cost, driver, optimise from gem.node import post_traversal, traversal from gem.optimise import sum_factorise from gem.coffee import optimise_monomial_sum @@ -111,7 +111,7 @@ def test_contraction_preserves_repeated_contractions(): table = gem.Indexed(gem.Literal(numpy.random.rand(3, 3, 3, 4)), ijk + (p,)) dofs = numpy.random.rand(3, 3, 3) - evaluation = optimise.contraction(gem.IndexSum(gem.Product(table, gem.Indexed(gem.Literal(dofs), ijk)), ijk)) + evaluation = driver.contraction(gem.IndexSum(gem.Product(table, gem.Indexed(gem.Literal(dofs), ijk)), ijk)) assert isinstance(evaluation, gem.IndexSum) weights = numpy.random.rand(4, 4) @@ -119,7 +119,7 @@ def test_contraction_preserves_repeated_contractions(): gem.Product(evaluation, evaluation)) expression = gem.IndexSum(cubed, (p,)) - optimised = optimise.contraction(expression) + optimised = driver.contraction(expression) assert evaluation in set(traversal([optimised])) # The evaluation is contracted once and reused, so the result holds