diff --git a/grudge/array_context.py b/grudge/array_context.py index 9290314c..34aceccd 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -53,8 +53,8 @@ from pytools.tag import Tag from grudge.transform.mappers import ( - InverseMassPropagator, - InverseMassRemover, + InverseMassDistributor, + RedundantMassTimesMassInverseRemover, MassInverseTimesStiffnessSimplifier, ) from grudge.transform.metadata import ( @@ -273,10 +273,10 @@ def transform_dag(self, dag): self.dot_codes_before.append(pt.get_dot_graph(dag)) if 1: # step 1: distribute mass inverse through DAG, across index lambdas - dag = InverseMassPropagator()(dag) + dag = InverseMassDistributor()(dag) # step 2: remove mass-times-mass-inverse einsums - dag = InverseMassRemover()(dag) + dag = RedundantMassTimesMassInverseRemover()(dag) # step 3: create new operator out of inverse mass times stiffness dag = MassInverseTimesStiffnessSimplifier()(dag) @@ -289,7 +289,6 @@ def transform_dag(self, dag): # }}} - # dag = pt.transform.materialize_with_mpms(dag) dag = deduplicate_data_wrappers(dag) num_nodes_after = get_num_nodes(dag) self.dot_codes_after.append(pt.get_dot_graph(dag)) @@ -312,34 +311,10 @@ def eliminate_reshapes_of_data_wrappers(ary): dag = pt.transform.map_and_copy(dag, eliminate_reshapes_of_data_wrappers) - def materialize_all_einsums_or_reduces(expr): - from pytato.raising import ReduceOp, index_lambda_to_high_level_op - - if isinstance(expr, pt.Einsum): - return expr.tagged(pt.tags.ImplStored()) - elif (isinstance(expr, pt.IndexLambda) - and isinstance(index_lambda_to_high_level_op(expr), ReduceOp)): - return expr.tagged(pt.tags.ImplStored()) - else: - return expr - - # logger.info("transform_dag.materialize_all_einsums_or_reduces") - # dag = pt.transform.map_and_copy(dag, materialize_all_einsums_or_reduces) - logger.info("transform_dag.infer_axes_tags") from grudge.transform.metadata import unify_discretization_entity_tags dag = unify_discretization_entity_tags(dag) - def untag_loopy_call_results(expr): - from pytato.loopy import LoopyCallResult - if isinstance(expr, LoopyCallResult): - return expr.copy(tags=frozenset(), - axes=(pt.Axis(frozenset()),)*expr.ndim) - else: - return expr - - dag = pt.transform.map_and_copy(dag, untag_loopy_call_results) - def _untag_impl_stored(expr): if isinstance(expr, pt.InputArgumentBase): return expr diff --git a/grudge/transform/mappers.py b/grudge/transform/mappers.py index 35f52ada..3cf25023 100644 --- a/grudge/transform/mappers.py +++ b/grudge/transform/mappers.py @@ -74,6 +74,15 @@ def map_einsum(self, expr): # {{{ algebraic dag rewrites class MassInverseTimesStiffnessSimplifier(CopyMapperWithExtraArgs): + """ + This is part three of a three part transformation pipeline. + + Creates a new operator that is the result of a mass inverse times weak + derivative operator and replaces all instances of these two operators in + each einsum with the new operator. + + See `InverseMassDistributor` for background. + """ def map_einsum(self, expr, *args, **kwargs): iarg_stiffness = None iarg_mass_inverse = None @@ -121,7 +130,13 @@ def map_einsum(self, expr, *args, **kwargs): return expr.copy(args=tuple(new_args)) -class InverseMassRemover(CopyMapperWithExtraArgs): +class RedundantMassTimesMassInverseRemover(CopyMapperWithExtraArgs): + """ + This is part two of a three-part transformation pipeline. Removes einsum + nodes that contain both a mass inverse and mass operator. + + See `InverseMassDistributor` for more information. + """ def map_einsum(self, expr, *args, **kwargs): found_mass = False found_mass_inverse = False @@ -142,37 +157,54 @@ def map_einsum(self, expr, *args, **kwargs): return expr.copy(args=tuple(new_args)) -class InverseMassPropagator(CopyMapperWithExtraArgs): +class InverseMassDistributor(CopyMapperWithExtraArgs): r""" - Applying a full mass inverse operator when using a tensor-product - discretization results in redundant work. For example, suppose we have a 3D - tensor-product discretization. Then the weak differentiation operator - associated with the r-coordinate direction can be expressed as - - .. math:: + Implements one part of a three-part transformation pipeline to realize an + algebraic simplification arising in tensor-product discretizations. - K_r = \hat{K} \otimes \hat{M} \otimes \hat{M}, + Specifically, one can represent a weak derivative operator associated with + a tensor-product discretization as a Kronecker product of a 1D weak + derivative operator with a variable number of 1D mass matrices, which + depends on the dimension of the problem. - where :math:`\hat{M}` is a 1D mass operator and :math:`\hat{K}` is a 1D weak - differentiation operator. Similarly, the inverse mass operator can be + Let $S$ be the full weak operator, $\hat{S}$ as the 1D weak derivative + operator and $\hat{M}$ as the 1D mass operator. For example, consider a 2D + tensor-product discretization. The weak $x$-derivative operator can be expressed as - .. math:: + ..math:: + + S_x = \hat{S} \otimes \hat{M} + + Since we are using a tensor-product discretization, the mass operator can be + decomposed as a tensor-product of a variable number of mass operators. + Hence, the mass inverse operator can also be expressed via a tensor-product. - M^{-1} = \hat{M}^{-1} \otimes \hat{M}^{-1} \otimes \hat{M}^{-1}, + ..math:: - By the properties of the tensor-product, multiplying the weak derivative - operator on the left by the inverse mass operator results in + M^{-1} S_x^k = (\hat{M}^{-1} \otimes \hat{M}^{-1})(\hat{S} \otimes \hat{M}) + + By associativity of the tensor-product, .. math:: - M^{-1} K_r = \hat{M}^{-1} \hat{K}_r \otimes \hat{I} \otimes \hat{I}, + M^{-1} S_x^k = \hat{M}^{-1}\hat{S} \otimes \hat{M}^{-1}\hat{M} + + Thus, we can instead apply the operator as - where :math:`\hat{I}` is a 1D identity operator. + ..math:: - The goal of this mapper is to remove redundant mass-times-mass inverse - operations from an expression graph of operations involved with a - tensor-product discretization. + M^{-1} S_x^k = \hat{M}^{-1}\hat{S} \otimes \hat{I} + + where $\hat{I}$ is a 1D identity operator. This results in both a reduction + in the total number of operations and the required storage for the + operators. + + This transformation takes care of the distribution of the mass inverse + operator through the DAG to other tensor-product application routines + Moreover, the mass inverse is distribtued to the face mass terms (if + included in the original einsum), and properly reshapes to and from + tensor-product form to apply the 1D mass operator. """ def map_einsum(self, expr, *args, **kwargs): new_args = [] @@ -195,18 +227,9 @@ def map_einsum(self, expr, *args, **kwargs): nfaces, nelts, _ = expr.args[1].shape ndofs = expr.shape[1] - if nfaces == 2: - dim = 1 - elif nfaces == 4: - dim = 2 - elif nfaces == 6: - dim = 3 - else: - raise ValueError("Unable to determine the dimension of the", - " hypercube to apply transformations.") - from math import ceil - ndofs_1d = int(ceil(ndofs**(1/dim))) + dim = ceil(nfaces/2) + ndofs_1d = ceil(ndofs**(1/dim)) expr = expr.reshape(nelts, *(ndofs_1d,)*dim) for axis in range(dim):