Skip to content

factorize_cost_models

FactorizeCostModel

Bases: CostModel[Statement]

A cost model where each query statement is a component, in program order. get_component_features returns one row for each statement.

Source code in src/finch/autoschedule/cost_models/factorize_cost_models.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
class FactorizeCostModel(CostModel[Statement]):
    """
    A cost model where each query statement is a component, in program order.
    get_component_features returns one row for each statement.
    """

    def components(
        self,
        term: LogicStatement,
        bindings: Mapping[Alias, TensorFType],
        machine: StatsMachine[NS],
    ) -> Iterator[Statement]:
        match term:
            case Plan(bodies):
                for body in bodies:
                    yield from self.components(body, bindings, machine)
            case Query() | QueryInto():
                yield Statement(term)
                machine(term)
            case Produces():
                return
            case _:
                raise ValueError(f"Unsupported statement type: {term}")

FlopsCostModel

Bases: FactorizeCostModel, AliasedForm

FlopsCostModel counts the non-fill values in the iteration space of each query, and in the output it materializes. Each query is a component. By default, allocations are weighted 10x iterations.

Source code in src/finch/autoschedule/cost_models/factorize_cost_models.py 
 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
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
class FlopsCostModel(FactorizeCostModel, AliasedForm):
    """
    FlopsCostModel counts the non-fill values in the iteration space of each
    query, and in the output it materializes. Each query is a component.
    By default, allocations are weighted 10x iterations.
    """

    @property
    def n_features(self) -> int:
        return 2

    def __init__(self, predictor: Callable[[np.ndarray], float] | None = None):
        super().__init__(
            predictor
            if predictor is not None
            else LinearPredictor(np.array([1.0, 10.0]))
        )

    def get_component_features(
        self,
        component: Statement,
        machine: StatsMachine[NS],
    ) -> np.ndarray:
        stats_factory = machine.stats_factory

        def visit_expr(expr: LogicExpression) -> tuple[NS, NS]:
            """Return the stats of the output of expr, and of its iteration space."""
            match expr:
                case Aggregate(Literal(op), Literal(init), arg, idxs):
                    iter_stats, _ = visit_expr(arg)
                    out_stats = stats_factory.aggregate(op, init, idxs, iter_stats)
                    return out_stats, iter_stats
                case Aggregate(Literal(op), init, arg, idxs):
                    # A tensor init is folded into the reduction.
                    iter_stats, _ = visit_expr(arg)
                    init_stats, _ = visit_expr(init)
                    out_stats = stats_factory.aggregate(op, None, idxs, iter_stats)
                    out_stats = stats_factory.mapjoin(op, init_stats, out_stats)
                    return out_stats, iter_stats
                case MapJoin(Literal(op), (arg,)):
                    return visit_expr(arg)
                case MapJoin(Literal(op), args):
                    out_stats = stats_factory.mapjoin(
                        op, *(visit_expr(arg)[0] for arg in args)
                    )
                    return out_stats, out_stats
                case Reorder(arg, idxs):
                    out_stats, iter_stats = visit_expr(arg)
                    return stats_factory.reorder(out_stats, idxs), iter_stats
                case Relabel(arg, idxs):
                    out_stats, iter_stats = visit_expr(arg)
                    return stats_factory.relabel(out_stats, idxs), iter_stats
                case Table() | Literal():
                    out_stats = machine(expr)
                    return out_stats, out_stats
                case _:
                    raise ValueError(f"Unsupported expression type: {expr}")

        match component.stmt:
            case Query(Table(Alias()), rhs):
                out_stats, iter_stats = visit_expr(rhs)
                return np.array(
                    [
                        [
                            iter_stats.estimate_non_fill_values(),
                            out_stats.estimate_non_fill_values(),
                        ]
                    ],
                    dtype=float,
                )
            case QueryInto() as stmt:
                return self.get_component_features(Statement(stmt.as_query()), machine)
            case stmt:
                raise ValueError(f"Unsupported statement type: {stmt}")

Statement dataclass

A Query or QueryInto of a program.

Source code in src/finch/autoschedule/cost_models/factorize_cost_models.py 
28
29
30
31
32
@dataclass(frozen=True)
class Statement:
    """A Query or QueryInto of a program."""

    stmt: LogicStatement