Skip to content

simplification

Basic algebraic simplification.

This module is a standard library of IR-agnostic rewrite rules. They apply to any Finch IR whose call and literal nodes implement the CallTerm and LiteralTerm interfaces, which today means FinchLogic (MapJoin), FinchNotation (Call), and FinchAssembly (Call).

annihilate(node)

f(a..., z, b...) => z when z is an annihilator for f.

Applied over every dtype, floats included, so nan * 0 and inf * 0 fold to 0 rather than to nan. TODO: add a safe mode for nan

Source code in src/finch/symbolic/simplification.py 
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
def annihilate(node: Term) -> Term | None:
    """
    `f(a..., z, b...)` => `z` when `z` is an annihilator for `f`.

    Applied over every dtype, floats included, so `nan * 0` and `inf * 0` fold to
    `0` rather than to `nan`.
    TODO: add a safe mode for nan
    """
    match node:
        case CallTerm(
            op=ExpressionTerm(result_type=FinchOperatorFType() as op_type), args=args
        ):
            return next(
                (
                    arg
                    for arg in args
                    if isinstance(arg, LiteralTerm) and is_annihilator(op_type, arg.val)
                ),
                None,
            )
    return None

canonicalize_associative(node)

Unwraps singleton n-ary function calls. - f(x) => x

An n-ary f absorbs its immediate f-children and moves literals to the front, next to each other: - f(a..., f(b...), c...) => f(k..., rest...).

A fixed-arity f is instead rotated one step at a time so that literals bubble up and to the left, where adjacent ones merge: - f(k1, f(k2, y)) => f(f(k1, k2), y) - f(f(k, x), y) => f(k, f(x, y)) - f(x, f(k, y)) => f(k, f(x, y)) - f(x, k) => f(k, x)

Source code in src/finch/symbolic/simplification.py 
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def canonicalize_associative(node: Term) -> Term | None:
    """
    Unwraps singleton n-ary function calls.
    - `f(x)` => `x`

    An n-ary `f` absorbs its immediate `f`-children and moves literals to the
    front, next to each other:
    - `f(a..., f(b...), c...)` => `f(k..., rest...)`.

    A fixed-arity `f` is instead rotated one step at a time so that literals
    bubble up and to the left, where adjacent ones merge:
    - `f(k1, f(k2, y))` => `f(f(k1, k2), y)`
    - `f(f(k, x), y)`   => `f(k, f(x, y))`
    - `f(x, f(k, y))`   => `f(k, f(x, y))`
    - `f(x, k)`         => `f(k, x)`
    """
    match node:
        case CallTerm(
            op=ExpressionTerm(result_type=FinchOperatorFType() as op_type) as op,
            args=args,
        ) if is_associative(op_type):
            if len(args) == 1:
                return args[0]
            if math.isinf(arity(op_type)):
                flat = [
                    leaf
                    for arg in args
                    for leaf in (
                        arg.args
                        if isinstance(arg, CallTerm) and arg.op == op
                        else (arg,)
                    )
                ]
                if is_commutative(op_type):
                    flat = sorted(
                        flat, key=lambda leaf: not isinstance(leaf, LiteralTerm)
                    )
                return _call_like(node, flat) if flat != list(args) else None
            match args:
                case (
                    LiteralTerm() as k1,
                    CallTerm(op=inner, args=(LiteralTerm() as k2, y)),
                ) if inner == op and isinstance(op, LiteralTerm):
                    return _call_like(node, [_evaluate(op, (k1, k2)), y])
                case (
                    CallTerm(op=inner, args=(LiteralTerm() as k, x)),
                    y,
                ) if inner == op:
                    return _call_like(node, [k, _call_like(node, [x, y])])
                case (
                    x,
                    CallTerm(op=inner, args=(LiteralTerm() as k, y)),
                ) if inner == op and is_commutative(op_type):
                    return _call_like(node, [k, _call_like(node, [x, y])])
                case (x, LiteralTerm() as k) if not isinstance(
                    x, LiteralTerm
                ) and is_commutative(op_type):
                    return _call_like(node, [k, x])
    return None

dedup_idempotent(node)

f(a..., x, b..., x, c...) => f(a..., x, b..., c...) for idempotent f.

Source code in src/finch/symbolic/simplification.py 
108
109
110
111
112
113
114
115
116
117
118
119
120
121
def dedup_idempotent(node: Term) -> Term | None:
    """`f(a..., x, b..., x, c...)` => `f(a..., x, b..., c...)` for idempotent `f`."""
    match node:
        case CallTerm(
            op=ExpressionTerm(result_type=FinchOperatorFType() as op_type), args=args
        ) if (
            is_idempotent(op_type)
            and is_associative(op_type)
            and is_commutative(op_type)
        ):
            unique = [arg for i, arg in enumerate(args) if arg not in args[:i]]
            if len(unique) != len(args):
                return unique[0] if len(unique) == 1 else _call_like(node, unique)
    return None

drop_identities(node)

f(a..., e, b...) => f(a..., b...) when e is an identity for f, and dropping it doesn't change the type of the result.

Source code in src/finch/symbolic/simplification.py 
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
def drop_identities(node: Term) -> Term | None:
    """
    `f(a..., e, b...)` => `f(a..., b...)` when `e` is an identity for `f`, and
    dropping it doesn't change the type of the result.
    """
    match node:
        case CallTerm(
            op=ExpressionTerm(result_type=FinchOperatorFType() as op_type), args=args
        ) if is_associative(op_type):
            if len(args) == 2 and arity(op_type) == 2:
                if (
                    isinstance(args[0], LiteralTerm)
                    and is_identity(op_type, args[0].val)
                    and _keeps_type(node, args[1])
                ):
                    return args[1]
                if (
                    isinstance(args[1], LiteralTerm)
                    and is_identity(op_type, args[1].val)
                    and _keeps_type(node, args[0])
                ):
                    return args[0]
            if math.isinf(arity(op_type)):
                kept = [
                    arg
                    for arg in args
                    if not (
                        isinstance(arg, LiteralTerm) and is_identity(op_type, arg.val)
                    )
                ]
                if len(kept) == len(args):
                    return None
                new_node = _call_like(node, kept or args[-1:])
                return new_node if _keeps_type(node, new_node) else None
    return None

fold_literals(node)

Evaluate literal arguments at compile time.

  • f(x, y) => literal(f(x, y)) when every argument is a literal. The call disappears, so this is sound for any operator.
  • f(a..., x, y, b...) => f(a..., f(x, y), b...) folds adjacent literal pairs of an associative n-ary f.
Source code in src/finch/symbolic/simplification.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def fold_literals(node: Term) -> Term | None:
    """
    Evaluate literal arguments at compile time.

    - `f(x, y)` => `literal(f(x, y))` when every argument is a literal. The
      call disappears, so this is sound for any operator.
    - `f(a..., x, y, b...)` => `f(a..., f(x, y), b...)` folds adjacent
      literal pairs of an associative n-ary `f`.
    """
    match node:
        case CallTerm(op=LiteralTerm() as op, args=args) if args:
            if all(isinstance(arg, LiteralTerm) for arg in args):
                return _evaluate(op, args)
            op_type = op.result_type
            if not (
                isinstance(op_type, FinchOperatorFType)
                and math.isinf(arity(op_type))
                and is_associative(op_type)
            ):
                return None
            new_args = []
            running_literal = op.head()(None)
            on_run = False
            for x in args:
                if isinstance(x, LiteralTerm):
                    running_literal = (
                        x if not on_run else _evaluate(op, (running_literal, x))
                    )
                    on_run = True
                else:
                    if on_run:
                        new_args.append(running_literal)
                        on_run = False
                    new_args.append(x)
            if on_run:
                new_args.append(running_literal)
            if new_args != list(args):
                return _call_like(node, new_args)
    return None