Skip to content

normalize

inline_constant_scalars(prgm)

Replace each reference to a bound ConstantScalar with its value.

Source code in src/finch/autoschedule/normalize.py 
 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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
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
def inline_constant_scalars(prgm: LogicNode) -> LogicNode:
    """
    Replace each reference to a bound `ConstantScalar` with its value.
    """
    constants: dict[Alias, Any] = {}
    produced: set[Alias] = set()

    def gather(node: LogicNode) -> None:
        match node:
            case Query(Table(Alias() as lhs, ()), Table(Literal(val), ())) if (
                isinstance(val, ConstantScalar)
            ):
                constants[lhs] = val.val
            case Produces(args):
                produced.update(a for a in args if isinstance(a, Alias))
        return

    Rewrite(PostWalk(gather))(prgm)
    if not constants:
        return prgm

    def constant_ref(node: LogicNode) -> Alias | None:
        """The constant `node` refers to, if it is exactly such a reference."""
        match node:
            case (
                Table(Alias() as tns, ())
                | Reorder(Table(Alias() as tns, ()), ())
                | Relabel(Table(Alias() as tns, ()), ())
            ) if tns in constants:
                return tns
            case _:
                return None

    def inline(node: LogicNode) -> LogicNode | None:
        match node:
            case Table(Alias() as tns, ()) if tns in constants:
                return Literal(constants[tns])
            case _:
                return None

    # Constants that keep their binding because inlining could not reach them.
    live: set[Alias] = set()

    def inline_stmt(stmt: LogicStatement) -> LogicStatement:
        match stmt:
            case Plan(bodies):
                return Plan(tuple(inline_stmt(body) for body in bodies))
            case Query(Table(Alias() as lhs, _), rhs) if lhs in produced and (
                ref := constant_ref(rhs)
            ):
                # `Query(Table(a, ()), Literal(v))` names no tensor to produce, so a
                # produced query that is *only* a constant reference is left
                # alone -- and the binding it reads has to stay.
                live.add(ref)
                return stmt
            case Query(lhs, rhs):
                return Query(lhs, cast(LogicExpression, Rewrite(PostWalk(inline))(rhs)))
            case QueryInto(lhs, op, rhs):
                return QueryInto(
                    lhs, op, cast(LogicExpression, Rewrite(PostWalk(inline))(rhs))
                )
            case _:
                return stmt

    inlined = inline_stmt(cast(LogicStatement, prgm))

    def drop(node: LogicNode) -> LogicNode | None:
        match node:
            case Plan(bodies):
                kept = tuple(
                    body
                    for body in bodies
                    if not (
                        isinstance(body, Query)
                        and body.lhs.tns in constants
                        and body.lhs.tns not in live
                    )
                )
                return Plan(kept) if kept != bodies else None
            case _:
                return None

    return Rewrite(PostWalk(drop))(inlined)

normalize_names(prgm, bindings) 

normalize_names(
    prgm: LogicStatement, bindings: dict[Alias, T]
) -> tuple[LogicStatement, dict[Alias, T]]

normalize_names(
    prgm: LogicExpression, bindings: dict[Alias, T]
) -> tuple[LogicExpression, dict[Alias, T]]

normalize_names(
    prgm: LogicNode, bindings: dict[Alias, T]
) -> tuple[LogicNode, dict[Alias, T]]

Normalizes names of aliases and fields in the logic program to avoid conflicts.

Source code in src/finch/autoschedule/normalize.py 
41
42
43
44
45
46
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
def normalize_names(
    prgm: LogicNode, bindings: dict[Alias, T]
) -> tuple[LogicNode, dict[Alias, T]]:
    """
    Normalizes names of aliases and fields in the logic program to avoid conflicts.
    """
    if bindings is None:
        bindings = {}
    spc = Namespace()
    renames: dict[str, str] = {}

    def rule_0(node: LogicNode) -> LogicNode | None:
        match node:
            case Alias(name):
                if name in renames:
                    return Alias(renames[name])
                new_name = spc.freshen("A")
                renames[name] = new_name
                return Alias(new_name)
            case Field(name):
                if name in renames:
                    return Field(renames[name])
                new_name = spc.freshen("i")
                renames[name] = new_name
                return Field(new_name)
            case _:
                return None

    bindings = {Rewrite(rule_0)(var): tns for var, tns in bindings.items()}
    root = Rewrite(PostWalk(rule_0))(prgm)

    return root, bindings