formatter

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,
    )

GalleyFormatter

Bases: LogicFormatter

A formatter which picks output formats the way Galley's physical optimizer does, choosing per level between dense, bytemap, hash, and sparse list based on the estimated sparsity of the output and the loop order it is built under.

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

class GalleyFormatter(LogicFormatter):
    """
    A formatter which picks output formats the way Galley's physical optimizer
    does, choosing per level between dense, bytemap, hash, and sparse list
    based on the estimated sparsity of the output and the loop order it is
    built under.
    """

    def __init__(self, loader: LogicLoader | None = None):
        super().__init__()
        if loader is None:
            loader = MockLogicLoader()
        self.ctx = loader

    def get_tensor_ftype(
        self,
        fill_type: AbstractFillFType,
        shape_type: tuple[FType, ...],
        stats: NumericStats,
        stats_factory: StatsFactory,
        loop_order: tuple[Field, ...],
    ) -> FiberTensorFType:
        output_indices = stats.index_order
        if len(shape_type) != len(output_indices):
            raise ValueError(
                f"Got {len(shape_type)} shape dimensions for "
                f"{len(output_indices)} stats dimensions."
            )

        formats = select_output_format(stats, stats_factory, loop_order, output_indices)
        logger.debug("Galley formats for %s: %s", output_indices, formats)

        lvl = element(fill_type, fill_type.element_type)
        for level in reversed(range(len(output_indices))):
            lvl = formats[level].build(lvl, shape_type[level])
        return fiber_tensor(lvl)

    def lower(
        self,
        prgm: lgc.LogicStatement,
        bindings: dict[lgc.Alias, TensorFType],
        stats: dict[lgc.Alias, NS],
        stats_factory: StatsFactory[NS],
    ):
        bindings = bindings.copy()
        stats_bindings: OrderedDict[lgc.Alias, NS] = OrderedDict(stats)
        stats_interpreter = StatsInterpreter(stats_factory=stats_factory)
        shape_types = prgm.infer_shape_type(
            {var: val.shape_type for var, val in bindings.items()}
        )
        fill_types = prgm.infer_fill_type(
            {var: val.fill_type for var, val in bindings.items()}
        )

        def formatter(node: lgc.LogicStatement) -> lgc.LogicStatement:
            match node:
                case lgc.Plan(bodies):
                    return lgc.Plan(tuple(formatter(body) for body in bodies))
                case lgc.Query(lgc.Table(lgc.Alias() as lhs, _), rhs) | lgc.QueryInto(
                    lgc.Table(lgc.Alias() as lhs, _), _, rhs
                ):
                    # Evaluating the query gives the stats of the stored result,
                    # in the order of its left-hand table.
                    rhs_stats = stats_interpreter(node, stats_bindings)
                    if not isinstance(rhs_stats, NumericStats):
                        raise TypeError("GalleyFormatter requires NumericStats.")
                    stats_bindings[lhs] = rhs_stats  # ty: ignore[invalid-assignment]

                    if lhs not in bindings:
                        shape_type = tuple(
                            ftype(dim) if dim is not None else ftypes.intp
                            for dim in shape_types[lhs]
                        )
                        bindings[lhs] = self.get_tensor_ftype(
                            fill_types[lhs],
                            shape_type,
                            rhs_stats,
                            stats_factory,
                            query_loop_order(rhs),
                        )

                    return node
                case lgc.Produces():
                    return node
                case _:
                    raise ValueError(
                        f"Unsupported logic statement for formatting: {node}"
                    )

        prgm = formatter(prgm)

        logger.debug(prgm)

        return self.ctx(prgm, bindings, stats_bindings, stats_factory)

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