<commit_before>"""
The Devito symbolic engine is built on top of SymPy and provides two
classes of functions:
- for inspection of expressions
- for (in-place) manipulation of expressions
- for creation of new objects given some expressions
All exposed functions are prefixed with 'dse' (devito symbolic engine)
"""

from __future__ import absolute_import

from collections import OrderedDict, Sequence
from time import time

from sympy import (Eq, Indexed, IndexedBase, S, Symbol, cos,
                   numbered_symbols, preorder_traversal, sin)
from sympy.simplify.cse_main import tree_cse

from devito.dimension import t, x, y, z
from devito.dse.extended_sympy import bhaskara_cos, bhaskara_sin
from devito.dse.graph import temporaries_graph
from devito.dse.inspection import (collect_aliases, estimate_cost, estimate_memory,
                                   is_binary_op, is_time_invariant, terminals)
from devito.dse.manipulation import (collect_nested, flip_indices, replace_invariants,
                                     rxreplace, unevaluate_arithmetic)
from devito.interfaces import TensorFunction
from devito.logger import dse, dse_warning

__all__ = ['rewrite']

_temp_prefix = 'temp'


def rewrite(expr, mode='advanced'):
    """
    Transform expressions to reduce their operation count.

    :param expr: the target expression.
    :param mode: drive the expression transformation. Available modes are
                 'basic', 'factorize', 'approx-trigonometry' and 'advanced'
                 (default). They act as follows: ::

                     * 'noop': do nothing, but track performance metrics
                     * 'basic': apply common sub-expressions elimination.
                     * 'factorize': apply heuristic factorization of temporaries.
                     * 'approx-trigonometry': replace expensive trigonometric
                         functions with suitable polynomial approximations.
                     * 'glicm': apply heuristic hoisting of time-invariant terms.
                     * 'split': split long expressions into smaller sub-expressions
                          exploiting associativity and commutativity.
                     * 'advanced': compose all known transformations.
    """

    if isinstance(expr, Sequence):
        assert all(isinstance(e, Eq) for e in expr)
        expr = list(expr)
    elif isinstance(expr, Eq):
        expr = [expr]
    else:
        raise ValueError("Got illegal expr of type %s." % type(expr))

    if not mode:
        return State(expr)
    elif isinstance(mode, str):
        mode = set([mode])
    else:
        try:
            mode = set(mode)
        except TypeError:
            dse_warning("Arg mode must be str or tuple (got %s)" % type(mode))
            return State(expr)
    if mode.isdisjoint({'noop', 'basic', 'factorize', 'approx-trigonometry',
                        'glicm', 'advanced'}):
        dse_warning("Unknown rewrite mode(s) %s" % str(mode))
        return State(expr)
    else:
        return Rewriter(expr).run(mode)


def dse_transformation(func):

    def wrapper(self, state, **kwargs):
        if kwargs['mode'].intersection(set(self.triggers[func.__name__])):
            tic = time()
            state.update(**func(self, state))
            toc = time()

            key = <KEY> (func.__name__, len(self.timings))
            self.timings[key] = toc - tic
            if self.profile:
                self.ops[key] = estimate_cost(state.exprs)

    return wrapper


class State(object):

    def __init__(self, exprs):
        self.exprs = exprs
        self.mapper = OrderedDict()

    def update(self, exprs=None, mapper=None):
        self.exprs = exprs or self.exprs
        self.mapper = mapper or self.mapper

    @property
    def time_invariants(self):
        return [i for i in self.exprs if i.lhs in self.mapper]

    @property
    def time_varying(self):
        return [i for i in self.exprs if i not in self.time_invariants]

    @property
    def ops_time_invariants(self):
        return estimate_cost(self.time_invariants)

    @property
    def ops_time_varying(self):
        return estimate_cost(self.time_varying)

    @property
    def ops(self):
        return self.ops_time_invariants + self.ops_time_varying

    @property
    def memory_time_invariants(self):
        return estimate_memory(self.time_invariants)

    @property
    def memory_time_varying(self):
        return estimate_memory(self.time_varying)

    @property
    def memory(self):
        return self.memory_time_invariants + self.memory_time_varying

    @property
    def output_fields(self):
        return [i.lhs for i in self.exprs if isinstance(i.lhs, Indexed)]

    @property
    def clusters(self):
        """
        Clusterize the expressions in ``self.exprs``. For more information
        about clusters, refer to TemporariesGraph.clusters.
        """
        return temporaries_graph(self.exprs).clusters


class Rewriter(object):

    """
    Transform expressions to reduce their operation count.
    """

    """
    Track what options trigger a given transformation.
    """
    triggers = {
        '_cse': ('basic', 'advanced'),
        '_factorize': ('factorize', 'advanced'),
        '_optimize_trigonometry': ('approx-trigonometry', 'advanced'),
        '_replace_time_invariants': ('glicm', 'advanced'),
        '_split_expressions': ('split', 'devito3.0')  # TODO: -> 'advanced' upon release
    }

    """
    Bag of thresholds, to be used to trigger or prevent certain transformations.
    """
    thresholds = {
        'expensive_expression': 100,
        'max_operands': 40,
    }

    def __init__(self, exprs, profile=False):
        self.exprs = exprs

        self.profile = profile
        self.ops = OrderedDict()
        self.timings = OrderedDict()

    def run(self, mode):
        state = State(self.exprs)

        self._cse(state, mode=mode)
        self._optimize_trigonometry(state, mode=mode)
        self._replace_time_invariants(state, mode=mode)
        self._split_expressions(state, mode=mode)
        self._factorize(state, mode=mode)

        self._finalize(state)

        self._summary(mode)

        return state

    @dse_transformation
    def _factorize(self, state, **kwargs):
        """
        Collect terms in each expr in exprs based on the following heuristic:

            * Collect all literals;
            * Collect all temporaries produced by CSE;
            * If the expression has an operation count higher than
              self.threshold, then this is applied recursively until
              no more factorization opportunities are available.
        """

        processed = []
        for expr in state.exprs:
            handle = collect_nested(expr)
            cost_handle = estimate_cost(handle)

            if cost_handle >= self.thresholds['expensive_expression']:
                handle_prev = handle
                cost_prev = estimate_cost(expr)<commit_msg>dse: Anticipate factorize() over split()<commit_after>"""
The Devito symbolic engine is built on top of SymPy and provides two
classes of functions:
- for inspection of expressions
- for (in-place) manipulation of expressions
- for creation of new objects given some expressions
All exposed functions are prefixed with 'dse' (devito symbolic engine)
"""

from __future__ import absolute_import

from collections import OrderedDict, Sequence
from time import time

from sympy import (Eq, Indexed, IndexedBase, S, Symbol, cos,
                   numbered_symbols, preorder_traversal, sin)
from sympy.simplify.cse_main import tree_cse

from devito.dimension import t, x, y, z
from devito.dse.extended_sympy import bhaskara_cos, bhaskara_sin
from devito.dse.graph import temporaries_graph
from devito.dse.inspection import (collect_aliases, estimate_cost, estimate_memory,
                                   is_binary_op, is_time_invariant, terminals)
from devito.dse.manipulation import (collect_nested, flip_indices, replace_invariants,
                                     rxreplace, unevaluate_arithmetic)
from devito.interfaces import TensorFunction
from devito.logger import dse, dse_warning

__all__ = ['rewrite']

_temp_prefix = 'temp'


def rewrite(expr, mode='advanced'):
    """
    Transform expressions to reduce their operation count.

    :param expr: the target expression.
    :param mode: drive the expression transformation. Available modes are
                 'basic', 'factorize', 'approx-trigonometry' and 'advanced'
                 (default). They act as follows: ::

                     * 'noop': do nothing, but track performance metrics
                     * 'basic': apply common sub-expressions elimination.
                     * 'factorize': apply heuristic factorization of temporaries.
                     * 'approx-trigonometry': replace expensive trigonometric
                         functions with suitable polynomial approximations.
                     * 'glicm': apply heuristic hoisting of time-invariant terms.
                     * 'split': split long expressions into smaller sub-expressions
                          exploiting associativity and commutativity.
                     * 'advanced': compose all known transformations.
    """

    if isinstance(expr, Sequence):
        assert all(isinstance(e, Eq) for e in expr)
        expr = list(expr)
    elif isinstance(expr, Eq):
        expr = [expr]
    else:
        raise ValueError("Got illegal expr of type %s." % type(expr))

    if not mode:
        return State(expr)
    elif isinstance(mode, str):
        mode = set([mode])
    else:
        try:
            mode = set(mode)
        except TypeError:
            dse_warning("Arg mode must be str or tuple (got %s)" % type(mode))
            return State(expr)
    if mode.isdisjoint({'noop', 'basic', 'factorize', 'approx-trigonometry',
                        'glicm', 'advanced'}):
        dse_warning("Unknown rewrite mode(s) %s" % str(mode))
        return State(expr)
    else:
        return Rewriter(expr).run(mode)


def dse_transformation(func):

    def wrapper(self, state, **kwargs):
        if kwargs['mode'].intersection(set(self.triggers[func.__name__])):
            tic = time()
            state.update(**func(self, state))
            toc = time()

            key = '%s%d' % (func.__name__, len(self.timings))
            self.timings[key] = toc - tic
            if self.profile:
                self.ops[key] = estimate_cost(state.exprs)

    return wrapper


class State(object):

    def __init__(self, exprs):
        self.exprs = exprs
        self.mapper = OrderedDict()

    def update(self, exprs=None, mapper=None):
        self.exprs = exprs or self.exprs
        self.mapper = mapper or self.mapper

    @property
    def time_invariants(self):
        return [i for i in self.exprs if i.lhs in self.mapper]

    @property
    def time_varying(self):
        return [i for i in self.exprs if i not in self.time_invariants]

    @property
    def ops_time_invariants(self):
        return estimate_cost(self.time_invariants)

    @property
    def ops_time_varying(self):
        return estimate_cost(self.time_varying)

    @property
    def ops(self):
        return self.ops_time_invariants + self.ops_time_varying

    @property
    def memory_time_invariants(self):
        return estimate_memory(self.time_invariants)

    @property
    def memory_time_varying(self):
        return estimate_memory(self.time_varying)

    @property
    def memory(self):
        return self.memory_time_invariants + self.memory_time_varying

    @property
    def output_fields(self):
        return [i.lhs for i in self.exprs if isinstance(i.lhs, Indexed)]

    @property
    def clusters(self):
        """
        Clusterize the expressions in ``self.exprs``. For more information
        about clusters, refer to TemporariesGraph.clusters.
        """
        return temporaries_graph(self.exprs).clusters


class Rewriter(object):

    """
    Transform expressions to reduce their operation count.
    """

    """
    Track what options trigger a given transformation.
    """
    triggers = {
        '_cse': ('basic', 'advanced'),
        '_factorize': ('factorize', 'advanced'),
        '_optimize_trigonometry': ('approx-trigonometry', 'advanced'),
        '_replace_time_invariants': ('glicm', 'advanced'),
        '_split_expressions': ('split', 'devito3.0')  # TODO: -> 'advanced' upon release
    }

    """
    Bag of thresholds, to be used to trigger or prevent certain transformations.
    """
    thresholds = {
        'expensive_expression': 100,
        'max_operands': 40,
    }

    def __init__(self, exprs, profile=False):
        self.exprs = exprs

        self.profile = profile
        self.ops = OrderedDict()
        self.timings = OrderedDict()

    def run(self, mode):
        state = State(self.exprs)

        self._cse(state, mode=mode)
        self._optimize_trigonometry(state, mode=mode)
        self._replace_time_invariants(state, mode=mode)
        self._factorize(state, mode=mode)
        self._split_expressions(state, mode=mode)

        self._finalize(state)

        self._summary(mode)

        return state

    @dse_transformation
    def _factorize(self, state, **kwargs):
        """
        Collect terms in each expr in exprs based on the following heuristic:

            * Collect all literals;
            * Collect all temporaries produced by CSE;
            * If the expression has an operation count higher than
              self.threshold, then this is applied recursively until
              no more factorization opportunities are available.
        """

        processed = []
        for expr in state.exprs:
            handle = collect_nested(expr)
            cost_handle = estimate_cost(handle)

            if cost_handle >= self.thresholds['expensive_expression']:
                handle_prev = handle
                cost_prev = estimate_cost(expr)
{'fraction_non_alphanumeric': 0.09105202973127502, 'fraction_numerical': 0.0010005717552887365, 'mean_word_length': 2.4414658140678798, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1966944', 'n_tokens_mistral': 4041, 'n_tokens_neox': 3794, 'n_words': 1121}