cost_models

ByKind

Bases: Combination[C], Generic[C, K]

Predict the cost of each component with the submodel of its kind, which key returns, by default the kind attribute of the component. The features of a component fill the columns of its submodel, and the other columns are zero, so the columns of a linear fit line up with the coefficients of each kind.

Source code in src/finch/autoschedule/cost_models/cost_models.py

class ByKind(Combination[C], Generic[C, K]):
    """
    Predict the cost of each component with the submodel of its kind, which
    key returns, by default the kind attribute of the component. The features
    of a component fill the columns of its submodel, and the other columns are
    zero, so the columns of a linear fit line up with the coefficients of each
    kind.
    """

    def __init__(
        self,
        models: Mapping[K, CostModel[C]],
        key: Callable[[C], K] = attrgetter("kind"),
    ):
        super().__init__(*models.values())
        self.by_kind = dict(models)
        self.key = key

    def get_component_features(
        self,
        component: C,
        machine: StatsMachine[NS],
    ) -> np.ndarray:
        kind = self.key(component)
        start = 0
        for other, model in self.by_kind.items():
            if other == kind:
                rows = model.get_component_features(component, machine)
                features = np.zeros((len(rows), self.n_features))
                features[:, start : start + model.n_features] = rows
                return features
            start += model.n_features
        raise KeyError(kind)

    def predict_component_cost(
        self,
        component: C,
        machine: StatsMachine[NS],
    ) -> float:
        model = self.by_kind[self.key(component)]
        return model.predict_component_cost(component, machine)

Concat

Bases: Combination[C]

Predict the cost of each component from the features of all submodels.

Source code in src/finch/autoschedule/cost_models/cost_models.py

class Concat(Combination[C]):
    """Predict the cost of each component from the features of all submodels."""

    def __init__(
        self,
        *models: CostModel[C],
        predictor: Callable[[np.ndarray], float],
    ):
        super().__init__(*models)
        self.predictor = predictor

CostModel

Bases: Form, Generic[C]

A cost model predicts the cost of a logic statement as a sum over its components, such as queries, loop prefixes, or tensor levels:

sum_c predict_component_cost(c)

Each stage decomposes programs into its own kind of component, and its searches price one component at a time. components yields the components of a program, advancing the StatsMachine past each statement after yielding its components. get_component_features returns X_c with shape (n_rows, n_features). By default, the predictor callable takes X_c and returns a scalar sum of independently modeled rows, including zero for an empty matrix. Like a Stage, a CostModel inherits the form of statements it accepts.

Source code in src/finch/autoschedule/cost_models/cost_models.py

class CostModel(Form, Generic[C]):
    """
    A cost model predicts the cost of a logic statement as a sum over its
    components, such as queries, loop prefixes, or tensor levels:

        sum_c predict_component_cost(c)

    Each stage decomposes programs into its own kind of component, and its
    searches price one component at a time. components yields the components
    of a program, advancing the StatsMachine past each statement after yielding
    its components. get_component_features returns X_c with shape
    (n_rows, n_features). By default, the predictor callable takes X_c and
    returns a scalar sum of independently modeled rows, including zero for an
    empty matrix. Like a Stage, a CostModel inherits the form of statements it
    accepts.
    """

    @property
    @abstractmethod
    def n_features(self) -> int: ...

    def __init__(self, predictor: Callable[[np.ndarray], float]):
        self.predictor = predictor

    @abstractmethod
    def components(
        self,
        term: LogicStatement,
        bindings: Mapping[Alias, Any],
        machine: StatsMachine[NS],
    ) -> Iterator[C]: ...

    @abstractmethod
    def get_component_features(
        self,
        component: C,
        machine: StatsMachine[NS],
    ) -> np.ndarray: ...

    def predict_component_cost(
        self,
        component: C,
        machine: StatsMachine[NS],
    ) -> float:
        return self.predictor(self.get_component_features(component, machine))

    def get_features(
        self,
        term: LogicStatement,
        bindings: Mapping[Alias, Any],
        stats: Mapping[Alias, TensorStats],
        stats_factory: StatsFactory,
    ) -> np.ndarray:
        machine = StatsMachine(stats_factory, dict(stats))
        features = [
            self.get_component_features(component, machine)
            for component in self.components(term, bindings, machine)
        ]
        return np.concatenate([np.empty((0, self.n_features)), *features])

    def predict_cost(
        self,
        term: LogicStatement,
        bindings: Mapping[Alias, Any],
        stats: Mapping[Alias, TensorStats],
        stats_factory: StatsFactory,
    ) -> float:
        machine = StatsMachine(stats_factory, dict(stats))
        return sum(
            (
                self.predict_component_cost(component, machine)
                for component in self.components(term, bindings, machine)
            ),
            0.0,
        )

DenseLevelOption

Bases: LevelOption

Storage features are positions and slots.

Source code in src/finch/autoschedule/cost_models/format_cost_models.py

class DenseLevelOption(LevelOption):
    """Storage features are positions and slots."""

    kind = "dense"
    n_features = 2

    def next_num_pos(self, stats, stats_factory, depth, num_pos):
        return num_pos * stats.get_dim_size(stats.index_order[depth])

    def get_storage_features(
        self, stats, stats_factory, depth, num_pos, idx_size, pos_size, val_size
    ):
        return np.array(
            [num_pos, self.next_num_pos(stats, stats_factory, depth, num_pos)]
        )

ElementLevelOption

Bases: LevelOption

The element level at the rank of a tensor stores one value for each position. Storage features are values and value bytes.

Source code in src/finch/autoschedule/cost_models/format_cost_models.py

class ElementLevelOption(LevelOption):
    """
    The element level at the rank of a tensor stores one value for each
    position. Storage features are values and value bytes.
    """

    kind = "element"
    n_features = 2

    def next_num_pos(self, stats, stats_factory, depth, num_pos):
        # The element level has no child, so its positions pass through.
        return num_pos

    def get_storage_features(
        self, stats, stats_factory, depth, num_pos, idx_size, pos_size, val_size
    ):
        return np.array([num_pos, num_pos * val_size])

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

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

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}")

FormatCostModel

Bases: CostModel[LevelChoice]

A cost model where each level of each formatted alias is a component. bindings lists the level options of an alias, outermost first and ending with its element level, and the alias is formatted when it is first written.

Source code in src/finch/autoschedule/cost_models/format_cost_models.py

class FormatCostModel(CostModel[LevelChoice]):
    """
    A cost model where each level of each formatted alias is a component.
    bindings lists the level options of an alias, outermost first and ending
    with its element level, and the alias is formatted when it is first written.
    """

    def components(
        self,
        term: LogicStatement,
        bindings: Mapping[Alias, Sequence[LevelOption]],
        machine: StatsMachine[NS],
    ) -> Iterator[LevelChoice]:
        visited: set[Alias] = set()

        def visit_alias(alias: Alias, options: Sequence[LevelOption]):
            stats = machine(alias)
            # bindings name no position type, so positions are intp.
            pos_size = np_dtype(ftypes.intp).itemsize
            num_pos = 1.0
            for depth, option in enumerate(options):
                yield LevelChoice(term, alias, depth, num_pos, option, pos_size)
                num_pos = option.next_num_pos(
                    stats, machine.stats_factory, depth, num_pos
                )

        def visit_stmt(stmt: LogicStatement):
            match stmt:
                case Plan(bodies):
                    for body in bodies:
                        yield from visit_stmt(body)
                case Query(Table(Alias() as lhs, _), _) | QueryInto(
                    Table(Alias() as lhs, _), _, _
                ):
                    machine(stmt)
                    if lhs in bindings and lhs not in visited:
                        yield from visit_alias(lhs, bindings[lhs])
                    visited.add(lhs)
                case Produces():
                    return
                case _:
                    raise ValueError(f"Unsupported statement type: {stmt}")

        return visit_stmt(term)

LevelChoice dataclass

The choice of a level option at a zero-based depth of alias, with num_pos positions entering it, in a format whose positions take pos_size bytes. term is the program which reads and writes alias.

Source code in src/finch/autoschedule/cost_models/format_cost_models.py

@dataclass(frozen=True)
class LevelChoice:
    """
    The choice of a level option at a zero-based depth of alias, with num_pos
    positions entering it, in a format whose positions take pos_size bytes.
    term is the program which reads and writes alias.
    """

    term: LogicStatement
    alias: Alias
    depth: int
    num_pos: float
    option: LevelOption
    pos_size: int

    @property
    def kind(self) -> str:
        return self.option.kind

LevelCostModel

Bases: FormatCostModel, LoopOrderedForm

The cost model of the levels of one type of option. Each level contributes an allocation row, and an access row for each query loop prefix that uses it. The features are the storage features of the option for the allocation, followed by sequential reads, scatter reads, sequential writes, and scatter writes for the accesses.

Source code in src/finch/autoschedule/cost_models/format_cost_models.py

class LevelCostModel(FormatCostModel, LoopOrderedForm):
    """
    The cost model of the levels of one type of option. Each level contributes
    an allocation row, and an access row for each query loop prefix that uses
    it. The features are the storage features of the option for the
    allocation, followed by sequential reads, scatter reads, sequential writes,
    and scatter writes for the accesses.
    """

    def __init__(
        self,
        option_type: type[LevelOption],
        predictor: Callable[[np.ndarray], float],
    ):
        super().__init__(predictor)
        self.option_type = option_type

    @property
    def n_features(self) -> int:
        return self.option_type.n_features + 4

    def get_access_features(
        self,
        alias: Alias,
        level: int,
        term: LogicStatement,
        machine: StatsMachine[NS],
    ) -> np.ndarray:
        """
        Return the accesses to alias at a zero-based depth in each query of term:
        sequential reads, scatter reads, sequential writes, and scatter writes,
        with one row per loop prefix that accesses the level in each query.
        Each access is counted by the PigeonCostModel iterations of the loops
        from the root to the current level. The element level is accessed with
        its parent.
        """
        # A sequential write is when the prefix of loops from root to current
        # level matches the ordering of the access exactly.
        # A sequential read is when the prefix of loops from root to current
        # level is a supersequence of the ordering of the access.
        features = []
        stats_interpreter = StatsMachine(
            machine.stats_factory, OrderedDict(machine.bindings)
        )
        pigeon = PigeonCostModel()

        def visit_access(idxs, loop_order, rhs, lhs_idxs, is_write, prefixes):
            if not idxs:
                return
            depth = min(level, len(idxs) - 1)
            order = tuple(idxs[: depth + 1])
            loops = tuple(loop_order[: loop_order.index(idxs[depth]) + 1])
            iterations = pigeon.get_component_features(
                Prefix(loops, rhs, lhs_idxs), stats_interpreter
            )[0, 0]
            features = prefixes.setdefault(loops, np.zeros(4))
            if is_write:
                features[2 if loops == order else 3] += iterations
            else:
                features[0 if is_subsequence(order, loops) else 1] += iterations

        def visit_stmt(stmt: LogicStatement):
            match stmt:
                case Plan(bodies):
                    for body in bodies:
                        visit_stmt(body)
                    return
                case QueryInto(Table(), _, Literal()):
                    stats_interpreter(stmt)
                    return
                case Produces():
                    return
                case Query(Table(lhs, lhs_idxs), Table() as rhs):
                    arg = rhs
                    loop_order = lhs_idxs
                case Query(
                    Table(lhs, lhs_idxs),
                    Aggregate(_, _, Reorder(arg, loop_order), _) as rhs,
                ):
                    pass
                case QueryInto(Table(lhs, lhs_idxs), _, Reorder(arg, loop_order)):
                    rhs = stmt.as_query().rhs
                case _:
                    raise ValueError(f"Unsupported statement type: {stmt}")

            prefixes = {}
            for node in PostOrderDFS(arg):
                match node:
                    case Table(tns, idxs) if tns == alias:
                        visit_access(idxs, loop_order, rhs, lhs_idxs, False, prefixes)
            if lhs == alias:
                visit_access(lhs_idxs, loop_order, rhs, lhs_idxs, True, prefixes)
            features.extend(prefixes.values())
            stats_interpreter(stmt)

        visit_stmt(term)
        return np.asarray(features, dtype=float).reshape(-1, 4)

    def get_component_features(
        self,
        component: LevelChoice,
        machine: StatsMachine[NS],
    ) -> np.ndarray:
        """
        Return the allocation row of the level, followed by its access rows.
        """
        stats = machine(component.alias)
        shape_type = stats.shape_type
        # The element level has no dimension, and a broadcast dimension is intp.
        idx_size = 0
        if component.depth < len(shape_type):
            idx_size = np_dtype(shape_type[component.depth] or ftypes.intp).itemsize
        access_features = self.get_access_features(
            component.alias, component.depth, component.term, machine
        )
        n_storage = self.option_type.n_features
        features = np.zeros((1 + len(access_features), self.n_features))
        features[0, :n_storage] = component.option.get_storage_features(
            stats,
            machine.stats_factory,
            component.depth,
            component.num_pos,
            idx_size,
            component.pos_size,
            np_dtype(stats.element_type).itemsize,
        )
        features[1:, n_storage:] = access_features
        return features

get_access_features(alias, level, term, machine)

Return the accesses to alias at a zero-based depth in each query of term: sequential reads, scatter reads, sequential writes, and scatter writes, with one row per loop prefix that accesses the level in each query. Each access is counted by the PigeonCostModel iterations of the loops from the root to the current level. The element level is accessed with its parent.

Source code in src/finch/autoschedule/cost_models/format_cost_models.py

def get_access_features(
    self,
    alias: Alias,
    level: int,
    term: LogicStatement,
    machine: StatsMachine[NS],
) -> np.ndarray:
    """
    Return the accesses to alias at a zero-based depth in each query of term:
    sequential reads, scatter reads, sequential writes, and scatter writes,
    with one row per loop prefix that accesses the level in each query.
    Each access is counted by the PigeonCostModel iterations of the loops
    from the root to the current level. The element level is accessed with
    its parent.
    """
    # A sequential write is when the prefix of loops from root to current
    # level matches the ordering of the access exactly.
    # A sequential read is when the prefix of loops from root to current
    # level is a supersequence of the ordering of the access.
    features = []
    stats_interpreter = StatsMachine(
        machine.stats_factory, OrderedDict(machine.bindings)
    )
    pigeon = PigeonCostModel()

    def visit_access(idxs, loop_order, rhs, lhs_idxs, is_write, prefixes):
        if not idxs:
            return
        depth = min(level, len(idxs) - 1)
        order = tuple(idxs[: depth + 1])
        loops = tuple(loop_order[: loop_order.index(idxs[depth]) + 1])
        iterations = pigeon.get_component_features(
            Prefix(loops, rhs, lhs_idxs), stats_interpreter
        )[0, 0]
        features = prefixes.setdefault(loops, np.zeros(4))
        if is_write:
            features[2 if loops == order else 3] += iterations
        else:
            features[0 if is_subsequence(order, loops) else 1] += iterations

    def visit_stmt(stmt: LogicStatement):
        match stmt:
            case Plan(bodies):
                for body in bodies:
                    visit_stmt(body)
                return
            case QueryInto(Table(), _, Literal()):
                stats_interpreter(stmt)
                return
            case Produces():
                return
            case Query(Table(lhs, lhs_idxs), Table() as rhs):
                arg = rhs
                loop_order = lhs_idxs
            case Query(
                Table(lhs, lhs_idxs),
                Aggregate(_, _, Reorder(arg, loop_order), _) as rhs,
            ):
                pass
            case QueryInto(Table(lhs, lhs_idxs), _, Reorder(arg, loop_order)):
                rhs = stmt.as_query().rhs
            case _:
                raise ValueError(f"Unsupported statement type: {stmt}")

        prefixes = {}
        for node in PostOrderDFS(arg):
            match node:
                case Table(tns, idxs) if tns == alias:
                    visit_access(idxs, loop_order, rhs, lhs_idxs, False, prefixes)
        if lhs == alias:
            visit_access(lhs_idxs, loop_order, rhs, lhs_idxs, True, prefixes)
        features.extend(prefixes.values())
        stats_interpreter(stmt)

    visit_stmt(term)
    return np.asarray(features, dtype=float).reshape(-1, 4)

get_component_features(component, machine)

Return the allocation row of the level, followed by its access rows.

Source code in src/finch/autoschedule/cost_models/format_cost_models.py

def get_component_features(
    self,
    component: LevelChoice,
    machine: StatsMachine[NS],
) -> np.ndarray:
    """
    Return the allocation row of the level, followed by its access rows.
    """
    stats = machine(component.alias)
    shape_type = stats.shape_type
    # The element level has no dimension, and a broadcast dimension is intp.
    idx_size = 0
    if component.depth < len(shape_type):
        idx_size = np_dtype(shape_type[component.depth] or ftypes.intp).itemsize
    access_features = self.get_access_features(
        component.alias, component.depth, component.term, machine
    )
    n_storage = self.option_type.n_features
    features = np.zeros((1 + len(access_features), self.n_features))
    features[0, :n_storage] = component.option.get_storage_features(
        stats,
        machine.stats_factory,
        component.depth,
        component.num_pos,
        idx_size,
        component.pos_size,
        np_dtype(stats.element_type).itemsize,
    )
    features[1:, n_storage:] = access_features
    return features

LevelOption

Bases: ABC

A kind of level, which describes the storage of a level with n_features storage features, and the number of positions it passes to its child. Both depend on the stats of the tensor, the zero-based depth of the level, and the num_pos positions entering it. Storage features also depend on the bytes of an index of the level, and of a position and a value of the tensor.

Source code in src/finch/autoschedule/cost_models/format_cost_models.py

class LevelOption(ABC):
    """
    A kind of level, which describes the storage of a level with n_features
    storage features, and the number of positions it passes to its child. Both
    depend on the stats of the tensor, the zero-based depth of the level, and
    the num_pos positions entering it. Storage features also depend on the
    bytes of an index of the level, and of a position and a value of the tensor.
    """

    kind: ClassVar[str]
    n_features: ClassVar[int]

    @abstractmethod
    def next_num_pos(
        self, stats: NS, stats_factory: StatsFactory[NS], depth: int, num_pos
    ): ...

    @abstractmethod
    def get_storage_features(
        self,
        stats: NS,
        stats_factory: StatsFactory[NS],
        depth: int,
        num_pos,
        idx_size: int,
        pos_size: int,
        val_size: int,
    ) -> np.ndarray: ...

LoopCostModel

Bases: CostModel[Prefix]

A cost model where each loop prefix, including the empty prefix, is a component. components yields prefixes in query order, from the shortest prefix to the longest, and get_component_features returns one row for each prefix. Predictors used with branch-and-bound loop search must predict nonnegative component costs.

Source code in src/finch/autoschedule/cost_models/loop_cost_models.py

class LoopCostModel(CostModel[Prefix]):
    """
    A cost model where each loop prefix, including the empty prefix, is a component.
    components yields prefixes in query order, from the shortest prefix to the
    longest, and get_component_features returns one row for each prefix.
    Predictors used with branch-and-bound loop search must predict nonnegative
    component costs.
    """

    def components(
        self,
        term: LogicStatement,
        bindings: Mapping[Alias, TensorFType],
        machine: StatsMachine[NS],
    ) -> Iterator[Prefix]:
        match term:
            case Plan(bodies):
                for body in bodies:
                    yield from self.components(body, bindings, machine)
                return
            case QueryInto(Table(), _, Literal()):
                machine(term)
                return
            case Produces():
                return
            case Query(Table(_, lhs_idxs), Table() as expr):
                idxs = lhs_idxs
            case Query(
                Table(_, lhs_idxs),
                Aggregate(_, _, Reorder(_, idxs), _) as expr,
            ):
                pass
            case QueryInto(Table(_, lhs_idxs), _, Reorder(_, idxs)):
                expr = term.as_query().rhs
            case _:
                raise ValueError(f"Unsupported statement type: {term}")

        for n in range(len(idxs) + 1):
            yield Prefix(idxs[:n], expr, lhs_idxs)
        machine(term)

PigeonCostModel

Bases: LoopCostModel, LoopOrderedForm

PigeonCostModel counts iteration, operation, and output costs across loop prefixes. By default, allocations are weighted 10x reads and operations.

Source code in src/finch/autoschedule/cost_models/loop_cost_models.py

class PigeonCostModel(LoopCostModel, LoopOrderedForm):
    """
    PigeonCostModel counts iteration, operation, and output costs across loop
    prefixes. By default, allocations are weighted 10x reads and operations.
    """

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

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

    def get_component_features(
        self,
        component: Prefix,
        machine: StatsMachine[NS],
    ) -> np.ndarray:
        """Return the incremental features for an outermost-first loop prefix.

        A nonempty prefix charges its last index, projecting execution and
        output onto prefix[:-1], as in the inside-out calculation. The empty
        prefix charges the output's non-fill values once. Summing all prefixes,
        including the empty one, gives the query's features.
        """
        prefix, expr, output_vars = (
            component.prefix,
            component.expr,
            component.output_vars,
        )
        stats_factory = machine.stats_factory
        match expr:
            case Aggregate(arg=arg):
                pass
            case _:
                arg = expr
        out_space = machine(expr)
        if output_vars is not None:
            reduce_idxs = tuple(
                idx for idx in out_space.index_order if idx not in output_vars
            )
            if reduce_idxs:
                out_space = stats_factory.aggregate(
                    ffuncs.or_, False, reduce_idxs, stats_factory.support(out_space)
                )
            out_space = stats_factory.reorder(out_space, output_vars)
        if not prefix:
            return np.array([[0.0, 0.0, out_space.estimate_non_fill_values()]])

        def project(stat: NS) -> NS:
            reduce_idxs = tuple(idx for idx in stat.index_order if idx not in parent)
            return stats_factory.aggregate(
                ffuncs.or_, False, reduce_idxs, stats_factory.support(stat)
            )

        parent = set(prefix[:-1])
        idx = prefix[-1]
        exec_space = project(machine(arg))
        inputs = []
        for node in PostOrderDFS(arg):
            match node:
                case Table(_, idxs) if idx in idxs:
                    inputs.append(stats_factory.support(machine(node)))
        iter_space = exec_space
        if inputs:
            iter_space = stats_factory.mapjoin(
                ffuncs.and_, exec_space, stats_factory.mapjoin(ffuncs.or_, *inputs)
            )
        out_cost = 0.0
        if idx in out_space.index_order:
            out_cost = project(out_space).estimate_non_fill_values()
        return np.array(
            [
                [
                    iter_space.estimate_non_fill_values(),
                    exec_space.estimate_non_fill_values() * len(inputs),
                    out_cost,
                ]
            ]
        )

get_component_features(component, machine)

Return the incremental features for an outermost-first loop prefix.

A nonempty prefix charges its last index, projecting execution and output onto prefix[:-1], as in the inside-out calculation. The empty prefix charges the output's non-fill values once. Summing all prefixes, including the empty one, gives the query's features.

Source code in src/finch/autoschedule/cost_models/loop_cost_models.py

def get_component_features(
    self,
    component: Prefix,
    machine: StatsMachine[NS],
) -> np.ndarray:
    """Return the incremental features for an outermost-first loop prefix.

    A nonempty prefix charges its last index, projecting execution and
    output onto prefix[:-1], as in the inside-out calculation. The empty
    prefix charges the output's non-fill values once. Summing all prefixes,
    including the empty one, gives the query's features.
    """
    prefix, expr, output_vars = (
        component.prefix,
        component.expr,
        component.output_vars,
    )
    stats_factory = machine.stats_factory
    match expr:
        case Aggregate(arg=arg):
            pass
        case _:
            arg = expr
    out_space = machine(expr)
    if output_vars is not None:
        reduce_idxs = tuple(
            idx for idx in out_space.index_order if idx not in output_vars
        )
        if reduce_idxs:
            out_space = stats_factory.aggregate(
                ffuncs.or_, False, reduce_idxs, stats_factory.support(out_space)
            )
        out_space = stats_factory.reorder(out_space, output_vars)
    if not prefix:
        return np.array([[0.0, 0.0, out_space.estimate_non_fill_values()]])

    def project(stat: NS) -> NS:
        reduce_idxs = tuple(idx for idx in stat.index_order if idx not in parent)
        return stats_factory.aggregate(
            ffuncs.or_, False, reduce_idxs, stats_factory.support(stat)
        )

    parent = set(prefix[:-1])
    idx = prefix[-1]
    exec_space = project(machine(arg))
    inputs = []
    for node in PostOrderDFS(arg):
        match node:
            case Table(_, idxs) if idx in idxs:
                inputs.append(stats_factory.support(machine(node)))
    iter_space = exec_space
    if inputs:
        iter_space = stats_factory.mapjoin(
            ffuncs.and_, exec_space, stats_factory.mapjoin(ffuncs.or_, *inputs)
        )
    out_cost = 0.0
    if idx in out_space.index_order:
        out_cost = project(out_space).estimate_non_fill_values()
    return np.array(
        [
            [
                iter_space.estimate_non_fill_values(),
                exec_space.estimate_non_fill_values() * len(inputs),
                out_cost,
            ]
        ]
    )

Prefix dataclass

An outermost-first loop prefix of expr, the right-hand side of a query. output_vars orders the output of the query, when there is one.

Source code in src/finch/autoschedule/cost_models/loop_cost_models.py

@dataclass(frozen=True)
class Prefix:
    """
    An outermost-first loop prefix of expr, the right-hand side of a query.
    output_vars orders the output of the query, when there is one.
    """

    prefix: tuple[Field, ...]
    expr: LogicExpression
    output_vars: tuple[Field, ...] | None = None

QuadraticPredictor

Sum row-wise quadratic forms, with flattened row-major coefficients.

Source code in src/finch/autoschedule/cost_models/cost_models.py

class QuadraticPredictor:
    """Sum row-wise quadratic forms, with flattened row-major coefficients."""

    def __init__(self, coefficients: np.ndarray):
        self.coefficients: np.ndarray = coefficients

    def __call__(self, features: np.ndarray) -> float:
        x = np.einsum("ij,ik->ijk", features, features)
        return float(x.sum(axis=0).ravel() @ self.coefficients)

SparseHashLevelOption

Bases: LevelOption

Storage features are positions, entries, position bytes, and index bytes.

Source code in src/finch/autoschedule/cost_models/format_cost_models.py

class SparseHashLevelOption(LevelOption):
    """Storage features are positions, entries, position bytes, and index bytes."""

    kind = "sparse_hash"
    n_features = 4

    def next_num_pos(self, stats, stats_factory, depth, num_pos):
        return nnz_after(stats, stats_factory, depth)

    def get_storage_features(
        self, stats, stats_factory, depth, num_pos, idx_size, pos_size, val_size
    ):
        nnz_l = nnz_after(stats, stats_factory, depth)
        return np.array([num_pos, nnz_l, (num_pos + 1) * pos_size, nnz_l * idx_size])

Statement dataclass

A Query or QueryInto of a program.

Source code in src/finch/autoschedule/cost_models/factorize_cost_models.py

@dataclass(frozen=True)
class Statement:
    """A Query or QueryInto of a program."""

    stmt: LogicStatement

Sum

Bases: Combination[C]

Sum the costs that submodels predict for each component. The coefficients of linear submodels concatenate like their features.

Source code in src/finch/autoschedule/cost_models/cost_models.py

class Sum(Combination[C]):
    """
    Sum the costs that submodels predict for each component. The coefficients
    of linear submodels concatenate like their features.
    """

    def predict_component_cost(
        self,
        component: C,
        machine: StatsMachine[NS],
    ) -> float:
        return sum(
            (model.predict_component_cost(component, machine) for model in self.models),
            0.0,
        )

TransposeCostModel

Bases: LoopCostModel, LoopOrderedForm

Count dense and sparse transposed entries, charging each input once per query.

Source code in src/finch/autoschedule/cost_models/loop_cost_models.py

class TransposeCostModel(LoopCostModel, LoopOrderedForm):
    """
    Count dense and sparse transposed entries, charging each input once per query.
    """

    @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([0.5, 60.0]))
        )

    def get_component_features(
        self,
        component: Prefix,
        machine: StatsMachine[NS],
    ) -> np.ndarray:
        """Charge only the input transpositions newly forced by this prefix.

        Deduplicate input tables within this prefix. Comparing with the parent
        prefix avoids charging them again at deeper loop levels. Empty prefixes
        have no transposition cost.
        """
        prefix, expr = component.prefix, component.expr
        features = np.zeros((1, self.n_features))
        if not prefix:
            return features

        match expr:
            case Aggregate(arg=arg):
                pass
            case _:
                arg = expr
        for node in set(PostOrderDFS(arg)):
            match node:
                case Table(_, idxs) if prefix[-1] in idxs:
                    stat = machine(node)
                case _:
                    continue
            order = stat.index_order
            parent = tuple(idx for idx in prefix[:-1] if idx in order)
            if parent == order[: len(parent)] and prefix[-1] != order[len(parent)]:
                nnz = stat.estimate_non_fill_values()
                space = stat.get_dim_space_size(order)
                dense = 0 < space < float("inf") and nnz / space > 0.1
                features[0, 0 if dense else 1] += nnz
        return features

get_component_features(component, machine)

Charge only the input transpositions newly forced by this prefix.

Deduplicate input tables within this prefix. Comparing with the parent prefix avoids charging them again at deeper loop levels. Empty prefixes have no transposition cost.

Source code in src/finch/autoschedule/cost_models/loop_cost_models.py

def get_component_features(
    self,
    component: Prefix,
    machine: StatsMachine[NS],
) -> np.ndarray:
    """Charge only the input transpositions newly forced by this prefix.

    Deduplicate input tables within this prefix. Comparing with the parent
    prefix avoids charging them again at deeper loop levels. Empty prefixes
    have no transposition cost.
    """
    prefix, expr = component.prefix, component.expr
    features = np.zeros((1, self.n_features))
    if not prefix:
        return features

    match expr:
        case Aggregate(arg=arg):
            pass
        case _:
            arg = expr
    for node in set(PostOrderDFS(arg)):
        match node:
            case Table(_, idxs) if prefix[-1] in idxs:
                stat = machine(node)
            case _:
                continue
        order = stat.index_order
        parent = tuple(idx for idx in prefix[:-1] if idx in order)
        if parent == order[: len(parent)] and prefix[-1] != order[len(parent)]:
            nnz = stat.estimate_non_fill_values()
            space = stat.get_dim_space_size(order)
            dense = 0 < space < float("inf") and nnz / space > 0.1
            features[0, 0 if dense else 1] += nnz
    return features