autoschedule

Aggregate dataclass

Bases: LogicTree, LogicExpression

Represents a logical AST statement that reduces arg using op, starting with init. idxs are the dimensions to reduce. May happen in any order. init may be a literal or a tensor expression. The fields of the result are those of arg which are not reduced, followed by the fields of init which are not in arg. init and arg are broadcast over the fields of the result which they lack, as in a MapJoin.

Attributes:

Name	Type	Description
op	Literal	The reduction operation.
init	LogicExpression	The initial value for the reduction.
arg	LogicExpression	The argument to reduce.
idxs	tuple[Field, ...]	The dimensions to reduce.

Source code in src/finch/finch_logic/nodes.py

@dataclass(eq=True, frozen=True)
class Aggregate(LogicTree, LogicExpression):
    """
    Represents a logical AST statement that reduces `arg` using `op`, starting
    with `init`. `idxs` are the dimensions to reduce. May happen in any order.
    `init` may be a literal or a tensor expression. The fields of the result
    are those of `arg` which are not reduced, followed by the fields of `init`
    which are not in `arg`. `init` and `arg` are broadcast over the fields of
    the result which they lack, as in a `MapJoin`.

    Attributes:
        op: The reduction operation.
        init: The initial value for the reduction.
        arg: The argument to reduce.
        idxs: The dimensions to reduce.
    """

    op: Literal
    init: LogicExpression
    arg: LogicExpression
    idxs: tuple[Field, ...]

    @property
    def children(self):
        """Returns the children of the node."""
        return [self.op, self.init, self.arg, *self.idxs]

    def fields(self) -> tuple[Field, ...]:
        """Returns fields of the node."""
        arg_fields = [field for field in self.arg.fields() if field not in self.idxs]
        return tuple(dict.fromkeys([*arg_fields, *self.init.fields()]))

    def dimmap(
        self,
        op: Callable,
        dim_bindings: dict[Alias, tuple[T | None, ...]],
    ) -> tuple[T | None, ...]:
        idxs = self.arg.fields()
        dims = self.arg.dimmap(op, dim_bindings)
        idx_dims = {
            idx: val
            for idx, val in zip(idxs, dims, strict=True)
            if idx not in self.idxs
        }
        init_dims = self.init.dimmap(op, dim_bindings)
        for idx, dim in zip(self.init.fields(), init_dims, strict=True):
            if idx in self.idxs:
                raise ValueError(f"The init of an aggregate can't have reduced {idx}")
            idx_dims[idx] = op(idx_dims[idx], dim) if idx in idx_dims else dim
        return tuple(idx_dims.values())

    def valmap(
        self,
        f: Callable,
        g: Callable,
        h: Callable,
        bindings: dict[Alias, T],
    ) -> T:
        return g(
            self.op.val,
            self.init.valmap(f, g, h, bindings),
            self.arg.valmap(f, g, h, bindings),
        )

    @classmethod
    def from_children(cls, op, init, arg, *idxs):
        return cls(op, init, arg, idxs)

children property

Returns the children of the node.

fields()

Returns fields of the node.

Source code in src/finch/finch_logic/nodes.py

def fields(self) -> tuple[Field, ...]:
    """Returns fields of the node."""
    arg_fields = [field for field in self.arg.fields() if field not in self.idxs]
    return tuple(dict.fromkeys([*arg_fields, *self.init.fields()]))

Alias dataclass

Bases: LogicNode, NamedTerm

Represents a logical AST expression for an alias named name. Aliases are used to refer to tables in the program.

Attributes:

Name	Type	Description
name	str	The name of the alias.

Source code in src/finch/finch_logic/nodes.py

@dataclass(eq=True, frozen=True)
class Alias(LogicNode, NamedTerm):
    """
    Represents a logical AST expression for an alias named `name`. Aliases are used to
    refer to tables in the program.

    Attributes:
        name: The name of the alias.
    """

    name: str

    @property
    def symbol(self) -> str:
        return self.name

    def fields(self) -> tuple[Field, ...]:
        """Returns fields of the node."""
        raise NotImplementedError(
            "Aliases do not have fields because they are tensors not tables!"
        )

    def dimmap(
        self,
        op: Callable,
        dim_bindings: dict[Alias, tuple[T | None, ...]],
    ) -> tuple[T | None, ...]:
        if dim_bindings is None or self not in dim_bindings:
            raise NotImplementedError(f"Cannot resolve dims of Alias {self.name}")
        return dim_bindings[self]

    def valmap(
        self,
        f: Callable,
        g: Callable,
        h: Callable,
        bindings: dict[Alias, T],
    ) -> T:
        if bindings is None or self not in bindings:
            raise NotImplementedError(f"Cannot resolve value of Alias {self.name}")
        return bindings[self]

fields()

Returns fields of the node.

Source code in src/finch/finch_logic/nodes.py

def fields(self) -> tuple[Field, ...]:
    """Returns fields of the node."""
    raise NotImplementedError(
        "Aliases do not have fields because they are tensors not tables!"
    )

BufferizedNDArrayFormatter

Bases: MonoLogicFormatter

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

class BufferizedNDArrayFormatter(MonoLogicFormatter):
    def get_tensor_ftype(
        self, fill_type: AbstractFillFType, shape_type: tuple[FType, ...]
    ):
        """
        Return the FType of the output tensor produced within the
        autoscheduler.
        """
        fill_ftype = fill_type.element_type
        return BufferizedNDArrayFType(
            buffer_type=NumpyBufferFType(fill_ftype),
            ndim=len(shape_type),
            dimension_type=TupleFType.from_tuple(shape_type),
            fill_type=fill_type,
        )

get_tensor_ftype(fill_type, shape_type)

Return the FType of the output tensor produced within the autoscheduler.

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

def get_tensor_ftype(
    self, fill_type: AbstractFillFType, shape_type: tuple[FType, ...]
):
    """
    Return the FType of the output tensor produced within the
    autoscheduler.
    """
    fill_ftype = fill_type.element_type
    return BufferizedNDArrayFType(
        buffer_type=NumpyBufferFType(fill_ftype),
        ndim=len(shape_type),
        dimension_type=TupleFType.from_tuple(shape_type),
        fill_type=fill_type,
    )

CompilerFormLowerer

Bases: FormattedForm, LogicLoader

Pass a program in FormattedForm to a loader which expects CompilerForm. CompilerForm keeps the definitions and updates of FormattedForm, so the program is unchanged.

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

class CompilerFormLowerer(FormattedForm, LogicLoader):
    """
    Pass a program in FormattedForm to a loader which expects CompilerForm.
    CompilerForm keeps the definitions and updates of FormattedForm, so the
    program is unchanged.
    """

    def __init__(self, ctx: LogicLoader):
        self.ctx = ctx

    def lower(
        self,
        prgm: lgc.LogicStatement,
        bindings: dict[lgc.Alias, TensorFType],
        stats: dict[lgc.Alias, TensorStats],
        stats_factory: StatsFactory,
    ) -> tuple[
        AssemblyLibrary,
        dict[lgc.Alias, TensorFType],
        dict[lgc.Alias, tuple[lgc.Field | None, ...]],
        lgc.LogicStatement,
    ]:
        return self.ctx(prgm, bindings, stats, stats_factory)

Field dataclass

Bases: LogicNode, NamedTerm

Represents a logical AST expression for a field named name. Fields are used to name the dimensions of a tensor. The named tensor is referred to as a "table".

Attributes:

Name	Type	Description
name	str	The name of the field.

Source code in src/finch/finch_logic/nodes.py

@dataclass(eq=True, frozen=True)
class Field(LogicNode, NamedTerm):
    """
    Represents a logical AST expression for a field named `name`.
    Fields are used to name the dimensions of a tensor. The named
    tensor is referred to as a "table".

    Attributes:
        name: The name of the field.
    """

    name: str

    @property
    def symbol(self) -> str:
        return self.name

Literal dataclass

Bases: LogicExpression, LiteralTerm, ExpressionTerm

Represents a logical AST expression for the literal value val.

A literal behaves as a zero-dimensional tensor: it has no fields, and its element type and fill value are those of val itself.

Attributes:

Name	Type	Description
val	T	The literal value.

Source code in src/finch/finch_logic/nodes.py

@dataclass(eq=True, frozen=True)
class Literal(LogicExpression, LiteralTerm, ExpressionTerm):
    """
    Represents a logical AST expression for the literal value `val`.

    A literal behaves as a zero-dimensional tensor: it has no fields, and its
    element type and fill value are those of `val` itself.

    Attributes:
        val: The literal value.
    """

    @property
    def result_type(self) -> FType:
        return ftype(self.val)

    def __hash__(self):
        try:
            return hash((type(self.val), self.val))
        except TypeError:
            return hash(id(self.val))

    def __eq__(self, value):
        if not isinstance(value, Literal):
            return False
        if type(self.val) is not type(value.val):
            return False
        # For consistency with __hash__, we fall back to pointer equality
        # when the value is unhashable
        try:
            hash(value.val)
            hash(self.val)
            return self.val == value.val
        except TypeError:
            return id(self.val) == id(value.val)

    def __repr__(self) -> str:
        return literal_repr(type(self).__name__, {"val": self.val})

    def fields(self) -> tuple[Field, ...]:
        """Returns fields of the node."""
        return ()

    def dimmap(
        self,
        op: Callable,
        dim_bindings: dict[Alias, tuple[T | None, ...]],
    ) -> tuple[T | None, ...]:
        return ()

    def valmap(
        self,
        f: Callable,
        g: Callable,
        h: Callable,
        bindings: dict[Alias, T],
    ) -> T:
        # A literal is a scalar, so its fill value is itself.
        return h(StaticFill(self.val))

fields()

Returns fields of the node.

Source code in src/finch/finch_logic/nodes.py

def fields(self) -> tuple[Field, ...]:
    """Returns fields of the node."""
    return ()

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.
    """

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.
    """

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.
    """

MapJoin dataclass

Bases: LogicTree, LogicExpression, CallTerm

Represents a logical AST expression for mapping the function op across args.... Dimensions which are not present are broadcasted. Dimensions which are present must match. The order of fields in the mapjoin is unique(vcat(map(getfields, args)...))

Attributes:

Name	Type	Description
op	Literal	The function to map.
args	tuple[LogicExpression, ...]	The arguments to map the function across.

Source code in src/finch/finch_logic/nodes.py

@dataclass(eq=True, frozen=True)
class MapJoin(LogicTree, LogicExpression, CallTerm):
    """
    Represents a logical AST expression for mapping the function `op` across `args...`.
    Dimensions which are not present are broadcasted. Dimensions which are
    present must match.  The order of fields in the mapjoin is
    `unique(vcat(map(getfields, args)...))`

    Attributes:
        op: The function to map.
        args: The arguments to map the function across.
    """

    op: Literal
    args: tuple[LogicExpression, ...]

    @property
    def children(self):
        """Returns the children of the node."""
        return [self.op, *self.args]

    def fields(self) -> tuple[Field, ...]:
        """Returns fields of the node."""
        args_fields = [x.fields() for x in self.args]
        return tuple(dict.fromkeys([f for fs in args_fields for f in fs]))

    def dimmap(
        self,
        op: Callable,
        dim_bindings: dict[Alias, tuple[T | None, ...]],
    ) -> tuple[T | None, ...]:
        arg_dims: dict[Field, T | None] = {}
        for arg in self.args:
            dims = arg.dimmap(op, dim_bindings)
            fields = arg.fields()
            for idx, dim in zip(fields, dims, strict=True):
                if idx in arg_dims:
                    arg_dims[idx] = op(arg_dims[idx], dim)
                else:
                    arg_dims[idx] = dim
        return tuple(arg_dims[f] for f in self.fields())

    def valmap(
        self,
        f: Callable,
        g: Callable,
        h: Callable,
        bindings: dict[Alias, T],
    ) -> T:
        return f(self.op.val, *[arg.valmap(f, g, h, bindings) for arg in self.args])

    @classmethod
    def from_children(cls, op, *args):
        return cls(op, args)

children property

Returns the children of the node.

fields()

Returns fields of the node.

Source code in src/finch/finch_logic/nodes.py

def fields(self) -> tuple[Field, ...]:
    """Returns fields of the node."""
    args_fields = [x.fields() for x in self.args]
    return tuple(dict.fromkeys([f for fs in args_fields for f in fs]))

Plan dataclass

Bases: LogicTree, LogicStatement

Represents a logical AST statement that executes a sequence of statements bodies.... Returns the last statement.

Attributes:

Name	Type	Description
bodies	tuple[LogicStatement, ...]	The sequence of statements to execute.

Source code in src/finch/finch_logic/nodes.py

@dataclass(eq=True, frozen=True)
class Plan(LogicTree, LogicStatement):
    """
    Represents a logical AST statement that executes a sequence of statements
    `bodies...`. Returns the last statement.

    Attributes:
        bodies: The sequence of statements to execute.
    """

    bodies: tuple[LogicStatement, ...] = ()

    def infer_dimmap(
        self,
        op: Callable,
        dim_bindings: dict[Alias, tuple[T | None, ...]],
    ) -> dict[Alias, tuple[T | None, ...]]:
        """Infers dimmaps for all aliases defined in the statement. The results
        will be stored in the dictionary passed to the method."""
        for body in self.bodies:
            body.infer_dimmap(op, dim_bindings)
        return dim_bindings

    def infer_valmap(
        self,
        f: Callable,
        g: Callable,
        h: Callable,
        bindings: dict[Alias, T],
    ) -> dict[Alias, T]:
        """Infers valmaps for all aliases defined in the statement. The results
        will be stored in the dictionary passed to the method."""
        for body in self.bodies:
            body.infer_valmap(f, g, h, bindings)
        return bindings

    @property
    def children(self):
        """Returns the children of the node."""
        return [*self.bodies]

    @classmethod
    def from_children(cls, *bodies):
        return cls(bodies)

children property

Returns the children of the node.

infer_dimmap(op, dim_bindings)

Infers dimmaps for all aliases defined in the statement. The results will be stored in the dictionary passed to the method.

Source code in src/finch/finch_logic/nodes.py

def infer_dimmap(
    self,
    op: Callable,
    dim_bindings: dict[Alias, tuple[T | None, ...]],
) -> dict[Alias, tuple[T | None, ...]]:
    """Infers dimmaps for all aliases defined in the statement. The results
    will be stored in the dictionary passed to the method."""
    for body in self.bodies:
        body.infer_dimmap(op, dim_bindings)
    return dim_bindings

infer_valmap(f, g, h, bindings)

Infers valmaps for all aliases defined in the statement. The results will be stored in the dictionary passed to the method.

Source code in src/finch/finch_logic/nodes.py

def infer_valmap(
    self,
    f: Callable,
    g: Callable,
    h: Callable,
    bindings: dict[Alias, T],
) -> dict[Alias, T]:
    """Infers valmaps for all aliases defined in the statement. The results
    will be stored in the dictionary passed to the method."""
    for body in self.bodies:
        body.infer_valmap(f, g, h, bindings)
    return bindings

PostWalk

A rewriter which recursively rewrites the arguments of each node using rw, then rewrites the resulting node. If all rewriters return nothing, returns nothing.

Attributes:

Name	Type	Description
rw	RwCallable	The rewriter function to apply.

Source code in src/finch/symbolic/rewriters.py

class PostWalk:
    """
    A rewriter which recursively rewrites the arguments of each node using
    `rw`, then rewrites the resulting node. If all rewriters return `nothing`,
    returns `nothing`.

    Attributes:
        rw (RwCallable): The rewriter function to apply.
    """

    def __init__(self, rw: RwCallable):
        self.rw = rw

    def __call__(self, x: T) -> T | None:
        if isinstance(x, TermTree):
            args = x.children
            new_args = list(map(self, args))
            if all(arg is None for arg in new_args):
                return self.rw(x)
            y = x.make_term(
                x.head(), *map(lambda x1, x2: default_rewrite(x1, x2), new_args, args)
            )
            assert isinstance(y, type(x))
            return default_rewrite(self.rw(y), y)
        return self.rw(x)

PreWalk

A rewriter which recursively rewrites each node using rw, then rewrites the arguments of the resulting node. If all rewriters return nothing, returns nothing.

Attributes:

Name	Type	Description
rw	RwCallable	The rewriter function to apply.

Source code in src/finch/symbolic/rewriters.py

class PreWalk:
    """
    A rewriter which recursively rewrites each node using `rw`, then rewrites
    the arguments of the resulting node. If all rewriters return `nothing`,
    returns `nothing`.

    Attributes:
        rw (RwCallable): The rewriter function to apply.
    """

    def __init__(self, rw: RwCallable):
        self.rw = rw

    def __call__(self, x: T) -> T | None:
        y = self.rw(x)
        if y is not None:
            if isinstance(y, TermTree):
                args = y.children
                return y.make_term(
                    y.head(), *[default_rewrite(self(arg), arg) for arg in args]
                )
            return y
        if isinstance(x, TermTree):
            args = x.children
            new_args = list(map(self, args))
            if not all(arg is None for arg in new_args):
                ret = x.make_term(
                    x.head(),
                    *map(lambda x1, x2: default_rewrite(x1, x2), new_args, args),
                )
                assert isinstance(ret, type(x))
                return ret
        return None

Produces dataclass

Bases: LogicTree, LogicStatement

Represents a logical AST statement that returns args... from the current plan. Halts execution of the program.

Attributes:

Name	Type	Description
args	tuple[Alias | LogicExpression, ...]	The arguments to return.

Source code in src/finch/finch_logic/nodes.py

@dataclass(eq=True, frozen=True)
class Produces(LogicTree, LogicStatement):
    """
    Represents a logical AST statement that returns `args...` from the current plan.
    Halts execution of the program.

    Attributes:
        args: The arguments to return.
    """

    args: tuple[Alias | LogicExpression, ...]

    @property
    def children(self):
        """Returns the children of the node."""
        return [*self.args]

    @classmethod
    def from_children(cls, *args):
        return cls(args)

    def infer_dimmap(
        self,
        op: Callable,
        dim_bindings: dict[Alias, tuple[T | None, ...]],
    ) -> dict[Alias, tuple[T | None, ...]]:
        """Infers dimmaps for all aliases defined in the statement. The results
        will be stored in the dictionary passed to the method."""
        return dim_bindings

    def infer_valmap(
        self,
        f: Callable,
        g: Callable,
        h: Callable,
        bindings: dict[Alias, T],
    ) -> dict[Alias, T]:
        """Infers valmaps for all aliases defined in the statement. The results
        will be stored in the dictionary passed to the method."""
        return bindings

children property

Returns the children of the node.

infer_dimmap(op, dim_bindings)

Infers dimmaps for all aliases defined in the statement. The results will be stored in the dictionary passed to the method.

Source code in src/finch/finch_logic/nodes.py

def infer_dimmap(
    self,
    op: Callable,
    dim_bindings: dict[Alias, tuple[T | None, ...]],
) -> dict[Alias, tuple[T | None, ...]]:
    """Infers dimmaps for all aliases defined in the statement. The results
    will be stored in the dictionary passed to the method."""
    return dim_bindings

infer_valmap(f, g, h, bindings)

Infers valmaps for all aliases defined in the statement. The results will be stored in the dictionary passed to the method.

Source code in src/finch/finch_logic/nodes.py

def infer_valmap(
    self,
    f: Callable,
    g: Callable,
    h: Callable,
    bindings: dict[Alias, T],
) -> dict[Alias, T]:
    """Infers valmaps for all aliases defined in the statement. The results
    will be stored in the dictionary passed to the method."""
    return bindings

Query dataclass

Bases: LogicTree, LogicStatement

Represents a logical AST statement that evaluates rhs, storing the result in the table lhs. The alias lhs.tns is bound to a tensor whose dimensions are ordered as lhs.idxs, so a query behaves as though its right-hand side were wrapped in Reorder(rhs, lhs.idxs). A query is not in place: the previous value of lhs.tns is replaced, and is only read if rhs refers to it. See QueryInto for in-place updates.

Attributes:

Name	Type	Description
lhs	Table	The table to write, a Table wrapping an Alias.
rhs	LogicExpression	The right-hand side to evaluate.

Source code in src/finch/finch_logic/nodes.py

@dataclass(eq=True, frozen=True)
class Query(LogicTree, LogicStatement):
    """
    Represents a logical AST statement that evaluates `rhs`, storing the result
    in the table `lhs`. The alias `lhs.tns` is bound to a tensor whose
    dimensions are ordered as `lhs.idxs`, so a query behaves as though its
    right-hand side were wrapped in `Reorder(rhs, lhs.idxs)`. A query is not
    in place: the previous value of `lhs.tns` is replaced, and is only read if
    `rhs` refers to it. See `QueryInto` for in-place updates.

    Attributes:
        lhs: The table to write, a `Table` wrapping an `Alias`.
        rhs: The right-hand side to evaluate.
    """

    lhs: Table
    rhs: LogicExpression

    @property
    def children(self):
        """Returns the children of the node."""
        return [self.lhs, self.rhs]

    def infer_dimmap(
        self,
        op: Callable,
        dim_bindings: dict[Alias, tuple[T | None, ...]],
    ) -> dict[Alias, tuple[T | None, ...]]:
        """Infers dimmaps for all aliases defined in the statement. The results
        will be stored in the dictionary passed to the method."""
        var = self.lhs.tns
        assert isinstance(var, Alias)
        dims = Reorder(self.rhs, self.lhs.idxs).dimmap(op, dim_bindings)
        if var in dim_bindings:
            for dim1, dim2 in zip(dims, dim_bindings[var], strict=True):
                op(dim1, dim2)
        else:
            dim_bindings[var] = dims
        return dim_bindings

    def infer_valmap(
        self,
        f: Callable,
        g: Callable,
        h: Callable,
        bindings: dict[Alias, T],
    ) -> dict[Alias, T]:
        """Infers valmaps for all aliases defined in the statement. The results
        will be stored in the dictionary passed to the method."""
        var = self.lhs.tns
        assert isinstance(var, Alias)
        if var in bindings:
            val = self.rhs.valmap(f, g, h, bindings)
            prev = bindings[var]
            # A dynamic value is compatible with any value of its dtype.
            if not (is_dynamic(val) or is_dynamic(prev)) and val != prev:
                raise ValueError(f"Cannot rebind alias {var} to a different values")
        else:
            bindings[var] = self.rhs.valmap(f, g, h, bindings)
        return bindings

children property

Returns the children of the node.

infer_dimmap(op, dim_bindings)

Infers dimmaps for all aliases defined in the statement. The results will be stored in the dictionary passed to the method.

Source code in src/finch/finch_logic/nodes.py

def infer_dimmap(
    self,
    op: Callable,
    dim_bindings: dict[Alias, tuple[T | None, ...]],
) -> dict[Alias, tuple[T | None, ...]]:
    """Infers dimmaps for all aliases defined in the statement. The results
    will be stored in the dictionary passed to the method."""
    var = self.lhs.tns
    assert isinstance(var, Alias)
    dims = Reorder(self.rhs, self.lhs.idxs).dimmap(op, dim_bindings)
    if var in dim_bindings:
        for dim1, dim2 in zip(dims, dim_bindings[var], strict=True):
            op(dim1, dim2)
    else:
        dim_bindings[var] = dims
    return dim_bindings

infer_valmap(f, g, h, bindings)

Infers valmaps for all aliases defined in the statement. The results will be stored in the dictionary passed to the method.

Source code in src/finch/finch_logic/nodes.py

def infer_valmap(
    self,
    f: Callable,
    g: Callable,
    h: Callable,
    bindings: dict[Alias, T],
) -> dict[Alias, T]:
    """Infers valmaps for all aliases defined in the statement. The results
    will be stored in the dictionary passed to the method."""
    var = self.lhs.tns
    assert isinstance(var, Alias)
    if var in bindings:
        val = self.rhs.valmap(f, g, h, bindings)
        prev = bindings[var]
        # A dynamic value is compatible with any value of its dtype.
        if not (is_dynamic(val) or is_dynamic(prev)) and val != prev:
            raise ValueError(f"Cannot rebind alias {var} to a different values")
    else:
        bindings[var] = self.rhs.valmap(f, g, h, bindings)
    return bindings

QueryInto dataclass

Bases: LogicTree, LogicStatement

Represents a logical AST statement that updates the table lhs in place, using the reduction operator op to fold each element of rhs into the matching element of lhs. The alias lhs.tns must already be bound. The fields of rhs which are not in lhs.idxs are reduced with op, and rhs is broadcast over the fields of lhs which it lacks, so a QueryInto is equivalent to the aggregate which starts from lhs, Query(lhs, Aggregate(op, lhs, rhs, setdiff(rhs.fields(), lhs.idxs))).

Attributes:

Name	Type	Description
lhs	Table	The table to update, a Table wrapping an Alias.
op	Literal	The reduction operator used to combine old and new values.
rhs	LogicExpression	The right-hand side to evaluate.

Source code in src/finch/finch_logic/nodes.py

@dataclass(eq=True, frozen=True)
class QueryInto(LogicTree, LogicStatement):
    """
    Represents a logical AST statement that updates the table `lhs` in place,
    using the reduction operator `op` to fold each element of `rhs` into the
    matching element of `lhs`. The alias `lhs.tns` must already be bound. The
    fields of `rhs` which are not in `lhs.idxs` are reduced with `op`, and `rhs`
    is broadcast over the fields of `lhs` which it lacks, so a `QueryInto` is
    equivalent to the aggregate which starts from `lhs`,
    `Query(lhs, Aggregate(op, lhs, rhs, setdiff(rhs.fields(), lhs.idxs)))`.

    Attributes:
        lhs: The table to update, a `Table` wrapping an `Alias`.
        op: The reduction operator used to combine old and new values.
        rhs: The right-hand side to evaluate.
    """

    lhs: Table
    op: Literal
    rhs: LogicExpression

    @property
    def children(self):
        """Returns the children of the node."""
        return [self.lhs, self.op, self.rhs]

    def as_query(self) -> Query:
        """The equivalent aggregate query, which is not in place."""
        idxs = tuple(idx for idx in self.rhs.fields() if idx not in self.lhs.idxs)
        return Query(self.lhs, Aggregate(self.op, self.lhs, self.rhs, idxs))

    def infer_dimmap(
        self,
        op: Callable,
        dim_bindings: dict[Alias, tuple[T | None, ...]],
    ) -> dict[Alias, tuple[T | None, ...]]:
        """Infers dimmaps for all aliases defined in the statement. The results
        will be stored in the dictionary passed to the method."""
        return self.as_query().infer_dimmap(op, dim_bindings)

    def infer_valmap(
        self,
        f: Callable,
        g: Callable,
        h: Callable,
        bindings: dict[Alias, T],
    ) -> dict[Alias, T]:
        """Infers valmaps for all aliases defined in the statement. The results
        will be stored in the dictionary passed to the method."""
        return self.as_query().infer_valmap(f, g, h, bindings)

children property

Returns the children of the node.

as_query()

The equivalent aggregate query, which is not in place.

Source code in src/finch/finch_logic/nodes.py

def as_query(self) -> Query:
    """The equivalent aggregate query, which is not in place."""
    idxs = tuple(idx for idx in self.rhs.fields() if idx not in self.lhs.idxs)
    return Query(self.lhs, Aggregate(self.op, self.lhs, self.rhs, idxs))

infer_dimmap(op, dim_bindings)

Infers dimmaps for all aliases defined in the statement. The results will be stored in the dictionary passed to the method.

Source code in src/finch/finch_logic/nodes.py

def infer_dimmap(
    self,
    op: Callable,
    dim_bindings: dict[Alias, tuple[T | None, ...]],
) -> dict[Alias, tuple[T | None, ...]]:
    """Infers dimmaps for all aliases defined in the statement. The results
    will be stored in the dictionary passed to the method."""
    return self.as_query().infer_dimmap(op, dim_bindings)

infer_valmap(f, g, h, bindings)

Infers valmaps for all aliases defined in the statement. The results will be stored in the dictionary passed to the method.

Source code in src/finch/finch_logic/nodes.py

def infer_valmap(
    self,
    f: Callable,
    g: Callable,
    h: Callable,
    bindings: dict[Alias, T],
) -> dict[Alias, T]:
    """Infers valmaps for all aliases defined in the statement. The results
    will be stored in the dictionary passed to the method."""
    return self.as_query().infer_valmap(f, g, h, bindings)

Relabel dataclass

Bases: LogicTree, LogicExpression

Represents a logical AST statement that relabels the dimensions of arg to be idxs.... Number of fields of arg and idxs must match.

Attributes:

Name	Type	Description
arg	LogicExpression	The argument to relabel.
idxs	tuple[Field, ...]	The new labels for dimensions.

Source code in src/finch/finch_logic/nodes.py

@dataclass(eq=True, frozen=True)
class Relabel(LogicTree, LogicExpression):
    """
    Represents a logical AST statement that relabels the dimensions of `arg` to be
    `idxs...`. Number of fields of `arg` and `idxs` must match.

    Attributes:
        arg: The argument to relabel.
        idxs: The new labels for dimensions.
    """

    arg: LogicExpression
    idxs: tuple[Field, ...]

    def fields(self) -> tuple[Field, ...]:
        """Returns fields of the node."""
        return self.idxs

    def dimmap(
        self, op: Callable, dim_bindings: dict[Alias, tuple[T | None, ...]]
    ) -> tuple[T | None, ...]:
        return self.arg.dimmap(op, dim_bindings)

    def valmap(
        self,
        f: Callable,
        g: Callable,
        h: Callable,
        bindings: dict[Alias, T],
    ) -> T:
        return self.arg.valmap(f, g, h, bindings)

    @property
    def children(self):
        """Returns the children of the node."""
        return [self.arg, *self.idxs]

    @classmethod
    def from_children(cls, arg, *idxs):
        return cls(arg, idxs)

children property

Returns the children of the node.

fields()

Returns fields of the node.

Source code in src/finch/finch_logic/nodes.py

def fields(self) -> tuple[Field, ...]:
    """Returns fields of the node."""
    return self.idxs

Reorder dataclass

Bases: LogicTree, LogicExpression

Represents a logical AST statement that reorders the dimensions of arg to be idxs.... Dimensions known to be length 1 may be dropped. Dimensions that do not exist in arg may be added.

Attributes:

Name	Type	Description
arg	LogicExpression	The argument to reorder.
idxs	tuple[Field, ...]	The new order of dimensions.

Source code in src/finch/finch_logic/nodes.py

@dataclass(eq=True, frozen=True)
class Reorder(LogicTree, LogicExpression):
    """
    Represents a logical AST statement that reorders the dimensions of `arg` to be
    `idxs...`. Dimensions known to be length 1 may be dropped. Dimensions that do not
    exist in `arg` may be added.

    Attributes:
        arg: The argument to reorder.
        idxs: The new order of dimensions.
    """

    arg: LogicExpression
    idxs: tuple[Field, ...]

    @property
    def children(self):
        """Returns the children of the node."""
        return [self.arg, *self.idxs]

    def fields(self) -> tuple[Field, ...]:
        """Returns fields of the node."""
        return self.idxs

    def dimmap(
        self, op: Callable, dim_bindings: dict[Alias, tuple[T | None, ...]]
    ) -> tuple[T | None, ...]:
        idxs = self.arg.fields()
        dims = self.arg.dimmap(op, dim_bindings)
        idx_dims = dict(zip(idxs, dims, strict=True))
        for idx in idxs:
            if idx not in self.idxs:
                # when squeezing a dimension, we combine with None
                op(idx_dims[idx], None)
        return tuple(idx_dims.get(f) for f in self.idxs)

    def valmap(
        self,
        f: Callable,
        g: Callable,
        h: Callable,
        bindings: dict[Alias, T],
    ) -> T:
        return self.arg.valmap(f, g, h, bindings)

    @classmethod
    def from_children(cls, arg, *idxs):
        return cls(arg, idxs)

children property

Returns the children of the node.

fields()

Returns fields of the node.

Source code in src/finch/finch_logic/nodes.py

def fields(self) -> tuple[Field, ...]:
    """Returns fields of the node."""
    return self.idxs

Table dataclass

Bases: LogicTree, LogicExpression

Represents a logical AST expression for a tensor object tns, indexed by fields idxs.... A table is a tensor with named dimensions.

Attributes:

Name	Type	Description
tns	Literal | Alias	The tensor object.
idxs	tuple[Field, ...]	The fields indexing the tensor.

Source code in src/finch/finch_logic/nodes.py

@dataclass(eq=True, frozen=True)
class Table(LogicTree, LogicExpression):
    """
    Represents a logical AST expression for a tensor object `tns`, indexed by fields
    `idxs...`. A table is a tensor with named dimensions.

    Attributes:
        tns: The tensor object.
        idxs: The fields indexing the tensor.
    """

    tns: Literal | Alias
    idxs: tuple[Field, ...]

    @property
    def children(self):
        """Returns the children of the node."""
        return [self.tns, *self.idxs]

    def fields(self) -> tuple[Field, ...]:
        """Returns fields of the node."""
        return self.idxs

    def dimmap(
        self, op: Callable, dim_bindings: dict[Alias, tuple[T | None, ...]]
    ) -> tuple[T | None, ...]:
        match self.tns:
            case Alias():
                return self.tns.dimmap(op, dim_bindings)
        raise NotImplementedError(f"Cannot resolve dims of {type(self.tns).__name__}")

    def valmap(
        self,
        f: Callable,
        g: Callable,
        h: Callable,
        bindings: dict[Alias, T],
    ) -> T:
        match self.tns:
            case Alias():
                return self.tns.valmap(f, g, h, bindings)
        raise NotImplementedError("Cannot resolve value of Tables")

    @classmethod
    def from_children(cls, tns, *idxs):
        return cls(tns, idxs)

children property

Returns the children of the node.

fields()

Returns fields of the node.

Source code in src/finch/finch_logic/nodes.py

def fields(self) -> tuple[Field, ...]:
    """Returns fields of the node."""
    return self.idxs

Value dataclass

Bases: LogicNode

Represents a logical AST expression for an expression ex of type type, yet to be evaluated.

Attributes:

Name	Type	Description
ex	Any	The expression to be evaluated.
type_	Any	The type of the expression.

Source code in src/finch/finch_logic/nodes.py

@dataclass(eq=True, frozen=True)
class Value(LogicNode):
    """
    Represents a logical AST expression for an expression `ex` of type `type`,
    yet to be evaluated.

    Attributes:
        ex: The expression to be evaluated.
        type_: The type of the expression.
    """

    ex: Any
    type_: Any

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

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