stages

AliasedForm

Bases: Form

AliasedForm requires that all aliases in the input are defined in the bindings or in previous queries and that all Tables are wrapping Aliases.

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

class AliasedForm(Form):
    """
    AliasedForm requires that all aliases in the input are defined
    in the bindings or in previous queries and that all Tables
    are wrapping Aliases.
    """

    @classmethod
    def validate_inputs(
        cls,
        term: Plan,
        bindings: dict[Alias, TensorFType],
        stats: dict[Alias, TensorStats],
        stats_factory: StatsFactory,
    ) -> None:
        defined_aliases = set(bindings.keys())

        def validate(node):
            match node:
                case Query(Table(Alias() as lhs, _), _):
                    defined_aliases.add(lhs)
                case QueryInto(Table(Alias() as lhs, _), _, _):
                    if lhs not in defined_aliases:
                        raise ValueError(
                            f"QueryInto updates alias {lhs.name}, which is not defined."
                        )
                case Query(lhs, _) | QueryInto(lhs, _, _):
                    raise ValueError(f"Query must write to a Table of an Alias: {lhs}")
                case Alias(name):
                    if node not in defined_aliases:
                        raise ValueError(f"Alias {name} is not defined in bindings.")
                case Table(tns, _):
                    if not isinstance(tns, Alias):
                        raise ValueError("Table nodes must wrap an Alias.")
            return node

        Rewrite(PreWalk(validate))(term)

CompilerForm

Bases: FormattedForm

CompilerForm is the input of the notation lowerer. It is a FormattedForm where every statement but the final Produces is a definition or an update. There are two valid kinds of definition: 1) copies Query(Table(lhs, lhs_idxs), Table(arg, idxs)) (lhs is reset to its fill value, and arg is written into it. A different storage order represents a transpose; notation lowering inserts equality-constrained loops to read arg in storage order. Fields absent from lhs_idxs are reduced by overwriting.) 2) reductions Query(Table(lhs, lhs_idxs), Aggregate(op, init, Reorder(arg, loop_order), )) (lhs is reset to init, a 0-D expression without aggregates or reorders, and arg is folded into it with op.) There are two valid kinds of update: 1) initializations QueryInto(Table(lhs, ), overwrite, Literal(init)) (Every element of lhs is set to init.) 2) folds QueryInto(Table(lhs, lhs_idxs), op, Reorder(arg, loop_order)) (arg is folded into lhs with op. A fold which overwrites lhs resets it to its fill value first.) In reductions and folds, arg is made of Tables, Literals, and MapJoins. The loop order contains each field once and visits all fields of lhs_idxs in order. Tables in a MapJoin must also follow loop order, but a single Table argument may have a different storage order, as in a copy. Fields absent from lhs_idxs are reduced with op. Every alias must have a TensorFType in the bindings, and a statement can't read the alias it writes.

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

class CompilerForm(FormattedForm):
    """
    CompilerForm is the input of the notation lowerer. It is a FormattedForm
    where every statement but the final Produces is a definition or an update.
    There are two valid kinds of definition:
    1) copies
        Query(Table(lhs, lhs_idxs), Table(arg, idxs))
    (lhs is reset to its fill value, and arg is written into it. A different
    storage order represents a transpose; notation lowering inserts
    equality-constrained loops to read arg in storage order. Fields absent from
    lhs_idxs are reduced by overwriting.)
    2) reductions
        Query(Table(lhs, lhs_idxs),
              Aggregate(op, init, Reorder(arg, loop_order), _))
    (lhs is reset to init, a 0-D expression without aggregates or reorders,
    and arg is folded into it with op.)
    There are two valid kinds of update:
    1) initializations
        QueryInto(Table(lhs, _), overwrite, Literal(init))
    (Every element of lhs is set to init.)
    2) folds
        QueryInto(Table(lhs, lhs_idxs), op, Reorder(arg, loop_order))
    (arg is folded into lhs with op. A fold which overwrites lhs resets it to
    its fill value first.)
    In reductions and folds, arg is made of Tables, Literals, and MapJoins. The
    loop order contains each field once and visits all fields of lhs_idxs in
    order. Tables in a MapJoin must also follow loop order, but a single Table
    argument may have a different storage order, as in a copy. Fields absent
    from lhs_idxs are reduced with op. Every alias must have a TensorFType in
    the bindings, and a statement can't read the alias it writes.
    """

    @classmethod
    def validate_inputs(
        cls,
        term: Plan,
        bindings: dict[Alias, TensorFType],
        stats: dict[Alias, TensorStats],
        stats_factory: StatsFactory,
    ) -> None:
        super().validate_inputs(term, bindings, stats, stats_factory)

        # FormattedForm has checked the grammar of each fold, its bindings,
        # and that it visits the lhs and its Tables in loop order.
        match term:
            case Plan((*bodies, Produces())):
                pass
            case _:
                raise ValueError("The last body of a plan must be a Produces node.")
        for body in bodies:
            match body:
                case Query(Table(Alias() as lhs, _), rhs) | QueryInto(
                    Table(Alias() as lhs, _), _, rhs
                ):
                    pass
                case _:
                    raise ValueError(f"CompilerForm only allows queries, not {body}")
            if lhs in PostOrderDFS(rhs):
                raise ValueError(
                    f"{type(body).__name__} can't both read and write {lhs}."
                )
            match body:
                case Query(_, Aggregate(_, init, _, _)) if not _is_init(init):
                    raise ValueError(
                        f"The init of an aggregate into {lhs} must be a 0-D "
                        "expression without aggregates or reorders."
                    )
                case Query(_, Aggregate(_, _, Reorder(arg, loop_order), _)) | (
                    QueryInto(_, _, Reorder(arg, loop_order))
                ):
                    if len(set(loop_order)) != len(loop_order):
                        raise ValueError("Loop order must not repeat fields.")
                    if not set(arg.fields()).issubset(loop_order):
                        raise ValueError("Loop order must include every RHS field.")
                case QueryInto(_, Literal(op), Literal()) if op != ffuncs.overwrite:
                    raise ValueError(f"Initializing {lhs} must overwrite it, not {op}.")

FormattedForm

Bases: LoopOrderedForm

FormattedForm requires that the input query has had its tensor formats set. Every alias must have a TensorFType in the bindings. The valid forms of a query are those of LoopOrderedForm.

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

class FormattedForm(LoopOrderedForm):
    """
    FormattedForm requires that the input query has had its tensor formats
    set. Every alias must have a TensorFType in the bindings. The valid forms
    of a query are those of LoopOrderedForm.
    """

    @classmethod
    def validate_inputs(
        cls,
        term: Plan,
        bindings: dict[Alias, TensorFType],
        stats: dict[Alias, TensorStats],
        stats_factory: StatsFactory,
    ) -> None:
        super().validate_inputs(term, bindings, stats, stats_factory)

        def validate(node):
            match node:
                case Plan(bodies):
                    for body in bodies[:-1]:
                        validate(body)
                case Query(lhs, rhs) | QueryInto(lhs, _, rhs):
                    validate(lhs)
                    validate(rhs)
                case Aggregate(_, _, arg, _) | Reorder(arg, _):
                    validate(arg)
                case MapJoin(_, args):
                    for arg in args:
                        validate(arg)
                case Table(tns, _):
                    if tns not in bindings:
                        raise ValueError(
                            f"Alias {tns.name} is not defined in bindings. All aliase\
                                 must have TensorFTypes specified at this stage."
                        )
                case Literal() | FillTable():
                    return
                case _:
                    raise ValueError(f"Unsupported query type: {node}")
            return

        validate(term)

LogicEinsumLowerer

Bases: FormattedForm, Stage

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

class LogicEinsumLowerer(FormattedForm, Stage):
    @abstractmethod
    def lower(
        self,
        term: LogicStatement,
        bindings: dict[Alias, TensorFType],
        stats: dict[Alias, TensorStats],
        stats_factory: StatsFactory,
    ) -> tuple[ein.EinsumStatement, dict[ein.Alias, TensorFType]]:
        """
        Generate Einsum Notation from the given logic and input types.  Also
        return a dictionary including additional tables needed to run the kernel.
        """

lower(term, bindings, stats, stats_factory) abstractmethod

Generate Einsum Notation from the given logic and input types. Also return a dictionary including additional tables needed to run the kernel.

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

@abstractmethod
def lower(
    self,
    term: LogicStatement,
    bindings: dict[Alias, TensorFType],
    stats: dict[Alias, TensorStats],
    stats_factory: StatsFactory,
) -> tuple[ein.EinsumStatement, dict[ein.Alias, TensorFType]]:
    """
    Generate Einsum Notation from the given logic and input types.  Also
    return a dictionary including additional tables needed to run the kernel.
    """

LogicFactorizer

Bases: AliasedForm, LogicLoader

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

class LogicFactorizer(AliasedForm, LogicLoader):
    @abstractmethod
    def lower(
        self,
        term: LogicStatement,
        bindings: dict[Alias, TensorFType],
        stats: dict[Alias, TensorStats],
        stats_factory: StatsFactory,
    ) -> tuple[
        AssemblyLibrary,
        dict[Alias, TensorFType],
        dict[Alias, tuple[Field | None, ...]],
        LogicStatement,
    ]:
        """
        Optimize the aggregate structure of the given logic statement and
        make decisions about materialization.
        """

lower(term, bindings, stats, stats_factory) abstractmethod

Optimize the aggregate structure of the given logic statement and make decisions about materialization.

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

@abstractmethod
def lower(
    self,
    term: LogicStatement,
    bindings: dict[Alias, TensorFType],
    stats: dict[Alias, TensorStats],
    stats_factory: StatsFactory,
) -> tuple[
    AssemblyLibrary,
    dict[Alias, TensorFType],
    dict[Alias, tuple[Field | None, ...]],
    LogicStatement,
]:
    """
    Optimize the aggregate structure of the given logic statement and
    make decisions about materialization.
    """

LogicFormatter

Bases: LoopOrderedForm, LogicLoader

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

class LogicFormatter(LoopOrderedForm, LogicLoader):
    @abstractmethod
    def lower(
        self,
        term: LogicStatement,
        bindings: dict[Alias, TensorFType],
        stats: dict[Alias, TensorStats],
        stats_factory: StatsFactory,
    ) -> tuple[
        AssemblyLibrary,
        dict[Alias, TensorFType],
        dict[Alias, tuple[Field | None, ...]],
        LogicStatement,
    ]:
        """
        Optimize the tensor formats and output orders for each query.
        """

lower(term, bindings, stats, stats_factory) abstractmethod

Optimize the tensor formats and output orders for each query.

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

@abstractmethod
def lower(
    self,
    term: LogicStatement,
    bindings: dict[Alias, TensorFType],
    stats: dict[Alias, TensorStats],
    stats_factory: StatsFactory,
) -> tuple[
    AssemblyLibrary,
    dict[Alias, TensorFType],
    dict[Alias, tuple[Field | None, ...]],
    LogicStatement,
]:
    """
    Optimize the tensor formats and output orders for each query.
    """

LogicLoopOrderer

Bases: SingleAggregateForm, LogicLoader

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

class LogicLoopOrderer(SingleAggregateForm, LogicLoader):
    @abstractmethod
    def lower(
        self,
        term: LogicStatement,
        bindings: dict[Alias, TensorFType],
        stats: dict[Alias, TensorStats],
        stats_factory: StatsFactory,
    ) -> tuple[
        AssemblyLibrary,
        dict[Alias, TensorFType],
        dict[Alias, tuple[Field | None, ...]],
        LogicStatement,
    ]:
        """
        Optimize the loop order of each query and add transposes where
        necessary.
        """

lower(term, bindings, stats, stats_factory) abstractmethod

Optimize the loop order of each query and add transposes where necessary.

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

@abstractmethod
def lower(
    self,
    term: LogicStatement,
    bindings: dict[Alias, TensorFType],
    stats: dict[Alias, TensorStats],
    stats_factory: StatsFactory,
) -> tuple[
    AssemblyLibrary,
    dict[Alias, TensorFType],
    dict[Alias, tuple[Field | None, ...]],
    LogicStatement,
]:
    """
    Optimize the loop order of each query and add transposes where
    necessary.
    """

LogicNotationLowerer

Bases: CompilerForm, Stage

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

class LogicNotationLowerer(CompilerForm, Stage):
    @abstractmethod
    def lower(
        self,
        term: LogicStatement,
        bindings: dict[Alias, TensorFType],
        stats: dict[Alias, TensorStats],
        stats_factory: StatsFactory,
    ) -> ntn.Module:
        """
        Generate Finch Notation from the given logic and input types.  Also
        return a dictionary including additional tables needed to run the kernel.
        """

lower(term, bindings, stats, stats_factory) abstractmethod

Generate Finch Notation from the given logic and input types. Also return a dictionary including additional tables needed to run the kernel.

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

@abstractmethod
def lower(
    self,
    term: LogicStatement,
    bindings: dict[Alias, TensorFType],
    stats: dict[Alias, TensorStats],
    stats_factory: StatsFactory,
) -> ntn.Module:
    """
    Generate Finch Notation from the given logic and input types.  Also
    return a dictionary including additional tables needed to run the kernel.
    """

LoopOrderedForm

Bases: SingleAggregateForm

LoopOrderedForm assumes that the input query has had its loop order set. There are four valid forms for a query in LoopOrderedForm: 1) transpose queries Query(Table(, output_order), Table(, )) 2) aggregate queries Query(Table(, output_order), Aggregate(, , Reorder(arg, loop_order), )) 3) in-place queries QueryInto(Table(, lhs_idxs), , Reorder(arg, loop_order)) 4) in-place initializations QueryInto(Table(, ), , Literal(_)) For aggregate and in-place queries, the loop order includes every lhs field and visits those fields in order. The Tables of arg follow the loop order, except that an in-place query of a single Table may read it in another order, representing a transpose.

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

class LoopOrderedForm(SingleAggregateForm):
    """
    LoopOrderedForm assumes that the input query has had its loop order set.
    There are four valid forms for a query in LoopOrderedForm:
        1) transpose queries
            Query(Table(_, output_order), Table(_, _))
        2) aggregate queries
            Query(Table(_, output_order), Aggregate(_, _, Reorder(arg, loop_order), _))
        3) in-place queries
            QueryInto(Table(_, lhs_idxs), _, Reorder(arg, loop_order))
        4) in-place initializations
            QueryInto(Table(_, _), _, Literal(_))
    For aggregate and in-place queries, the loop order includes every lhs
    field and visits those fields in order. The Tables of arg follow the loop
    order, except that an in-place query of a single Table may read it in
    another order, representing a transpose.
    """

    @staticmethod
    def _check_loop_order(idxs, loop_order):
        rel_loop_order = [idx for idx in loop_order if idx in idxs]
        return tuple(rel_loop_order) == tuple(idxs)

    @classmethod
    def validate_inputs(
        cls,
        term: Plan,
        bindings: dict[Alias, TensorFType],
        stats: dict[Alias, TensorStats],
        stats_factory: StatsFactory,
    ) -> None:
        super().validate_inputs(term, bindings, stats, stats_factory)

        def validate(node, loop_order):
            match node:
                case Plan(bodies):
                    for body in bodies[:-1]:
                        validate(body, loop_order)
                case Query(Table(), Table()) | QueryInto(Table(), _, Literal()):
                    return None
                case QueryInto(Table(_, lhs_idxs), _, Reorder(Table(), idxs)):
                    if not cls._check_loop_order(lhs_idxs, idxs):
                        raise ValueError("Table index order does not match loop order.")
                    return None
                case Query(
                    Table(_, lhs_idxs), Aggregate(_, _, Reorder(arg, idxs), _)
                ) | QueryInto(Table(_, lhs_idxs), _, Reorder(arg, idxs)):
                    if not cls._check_loop_order(lhs_idxs, idxs):
                        raise ValueError("Table index order does not match loop order.")
                    return validate(arg, idxs)
                case Query(Table(), Aggregate(_, _, arg, _)):
                    raise ValueError(
                        "All aggregates must wrap a Reorder node specifying\
                             the loop order."
                    )
                case QueryInto():
                    raise ValueError("In-place queries must have a loop order!")
                case MapJoin(_, args):
                    for arg in args:
                        validate(arg, loop_order)
                case Table(_, idxs):
                    if not cls._check_loop_order(idxs, loop_order):
                        raise ValueError("Table index order does not match loop order.")
                case Reorder(arg, _):
                    raise ValueError("Reorder nodes should only appear in loop orders!")
                case Literal() | FillTable():
                    return None
                case _:
                    raise ValueError(f"Unsupported query type: {node}")
            return None

        validate(term, None)

SingleAggregateForm

Bases: AliasedForm

SingleAggregateForm assumes that the fusion strategy has already been optimized for this query. There are three valid kinds of input query: 1) transpose queries Query(Table(, output_order), Table(, )) 2) aggregate queries Query(Table(, output_order), Aggregate(, init, arg, )) (init is a 0-D expression without aggregates or reorders.) 3) in-place queries QueryInto(Table(_, output_order), op, arg) (Here, arg has no aggregates. The fields of arg which are not in output_order are reduced with op.)

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

class SingleAggregateForm(AliasedForm):
    """
    SingleAggregateForm assumes that the fusion strategy has
    already been optimized for this query. There are three valid kinds of input query:
    1) transpose queries
        Query(Table(_, output_order), Table(_, _))
    2) aggregate queries
        Query(Table(_, output_order), Aggregate(_, init, arg, _))
    (init is a 0-D expression without aggregates or reorders.)
    3) in-place queries
        QueryInto(Table(_, output_order), op, arg)
    (Here, arg has no aggregates. The fields of arg which are not in
    output_order are reduced with op.)
    """

    @classmethod
    def validate_inputs(
        cls,
        term: Plan,
        bindings: dict[Alias, TensorFType],
        stats: dict[Alias, TensorStats],
        stats_factory: StatsFactory,
    ) -> None:
        super().validate_inputs(term, bindings, stats, stats_factory)

        def validate(node, agg_allowed):
            match node:
                case Plan(bodies):
                    if not isinstance(bodies[-1], Produces):
                        raise ValueError(
                            "The last body of a plan must be a Produces node."
                        )
                    for body in bodies[:-1]:
                        validate(body, True)
                case Query(Table(), Table()):
                    return None
                case Query(Table(), Aggregate(_, init, arg, _)) if _is_init(init):
                    return validate(arg, False)
                case Query(Table(), Aggregate(_, init, _, _)):
                    raise ValueError(
                        f"Aggregate queries must start from a 0-D expression "
                        f"without aggregates or reorders, not {init}. "
                        "Copy the init into the output and update it with a "
                        "QueryInto instead."
                    )
                case QueryInto(Table(), _, arg):
                    return validate(arg, False)
                case Query(_, rhs):
                    raise ValueError(f"Unsupported query right-hand side: {rhs}")
                case Aggregate(_, _, arg, _):
                    if not agg_allowed:
                        raise ValueError("Nested aggregates are not supported.")
                    return validate(arg, False)
                case Reorder(arg, _):
                    return validate(arg, agg_allowed)
                case MapJoin(_, args):
                    for arg in args:
                        validate(arg, agg_allowed)
                    return None
                case Literal() | FillTable() | Alias() | Table():
                    return None
                case _:
                    raise ValueError(f"Unsupported query type: {node}")
            return None

        validate(term, True)