tensor_stats

BaseTensorStats

Bases: TensorStats

Base class for tensor statistics.

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

class BaseTensorStats(TensorStats):
    """
    Base class for tensor statistics.
    """

    def __init__(
        self,
        index_order: Iterable[Field] | BaseTensorStats,
        dim_sizes: MutableMapping[Field, float] | None = None,
        fill_type: AbstractFillFType | None = None,
        dim_types: MutableMapping[Field, FType | None] | None = None,
    ):
        if isinstance(index_order, BaseTensorStats):
            if dim_sizes is not None:
                raise TypeError(
                    "dim_sizes must be omitted when constructing from BaseTensorStats"
                )
            self._init_from_fields(
                index_order.index_order,
                index_order.dim_sizes,
                index_order.fill_type,
                index_order.dim_types,
            )
        else:
            if dim_sizes is None or fill_type is None:
                raise TypeError(
                    "dim_sizes and fill_type are required when constructing from fields"
                )
            self._init_from_fields(index_order, dim_sizes, fill_type, dim_types)

    def _init_from_fields(
        self,
        index_order: Iterable[Field],
        dim_sizes: MutableMapping[Field, float],
        fill_type: AbstractFillFType,
        dim_types: MutableMapping[Field, FType | None] | None = None,
    ):
        self._index_order = tuple(index_order)
        self._dim_sizes: OrderedDict[Field, float] = OrderedDict(dim_sizes)
        # A dimension type is None when it is unknown, as for a broadcast.
        self._dim_types: OrderedDict[Field, FType | None] = OrderedDict(
            dim_types if dim_types is not None else dict.fromkeys(self._dim_sizes)
        )
        # Stats inform compilation, so they keep the fill of the format, which
        # says whether kernels may specialize on the fill value.
        self._fill_type = fill_type

    @classmethod
    def from_fields(
        cls,
        index_order: Iterable[Field],
        dim_sizes: MutableMapping[Field, float],
        fill_type: AbstractFillFType,
        dim_types: MutableMapping[Field, FType | None] | None = None,
    ) -> Self:
        return cls(index_order, dim_sizes, fill_type, dim_types=dim_types)

    def copy(self) -> Self:
        new = object.__new__(type(self))
        new.__dict__ = {name: copy.copy(value) for name, value in self.__dict__.items()}
        return new

    @property
    def dim_sizes(self) -> MutableMapping[Field, float]:
        return self._dim_sizes

    @dim_sizes.setter
    def dim_sizes(self, value: MutableMapping[Field, float]):
        self._dim_sizes = OrderedDict(value)

    def get_dim_size(self, idx: Field) -> float:
        return self.dim_sizes[idx]

    @property
    def dim_types(self) -> MutableMapping[Field, FType | None]:
        return self._dim_types

    @property
    def shape_type(self) -> tuple[FType | None, ...]:
        """The type of each dimension in index order, or None when unknown."""
        return tuple(self.dim_types.get(idx) for idx in self.index_order)

    @property
    def element_type(self) -> FType:
        return self.fill_type.element_type

    @property
    def index_order(self) -> tuple[Field, ...]:
        return self._index_order

    @index_order.setter
    def index_order(self, value: Iterable[Field]):
        self._index_order = tuple(value)

    @property
    def idxs(self) -> tuple[Field, ...]:
        return self.index_order

    @property
    def fill_type(self) -> AbstractFillFType:
        """The fill of the tensor's format, which a kernel may depend on."""
        return self._fill_type

    @fill_type.setter
    def fill_type(self, value: AbstractFillFType):
        self._fill_type = value

    def get_dim_space_size(self, idx: Iterable[Field]) -> float:
        prod = 1
        for i in idx:
            prod *= int(self.dim_sizes[i])
            if prod == 0 or prod > np.iinfo(np.int64).max:
                return float("inf")
        return float(prod)

fill_type property writable

The fill of the tensor's format, which a kernel may depend on.

shape_type property

The type of each dimension in index order, or None when unknown.

BoundStats

Bases: NumericStats

Base for degree-constraint upper-bound statistics: a guaranteed bound on the number of non-fill values, built from ell_p-norms of degree sequences.

Subclasses differ only in which p values they keep (ps) and how they turn the constraints into an estimate (:meth:estimate_non_fill_values); the scan and all propagation (in :class:BoundStatsFactory) are shared.

Source code in src/finch/autoschedule/tensor_stats/bound_stats.py

class BoundStats(NumericStats):
    """
    Base for degree-constraint upper-bound statistics: a guaranteed bound on the
    number of non-fill values, built from ``ell_p``-norms of degree sequences.

    Subclasses differ only in which ``p`` values they keep (``ps``) and how they
    turn the constraints into an estimate (:meth:`estimate_non_fill_values`); the
    scan and all propagation (in :class:`BoundStatsFactory`) are shared.
    """

    default_ps: ClassVar[tuple[float, ...]] = ()

    def __init__(
        self,
        base: BaseTensorStats,
        dcs: Iterable[DC],
        ps: Iterable[float] | None = None,
    ):
        super().__init__(base)
        self.ps: tuple[float, ...] = tuple(self.default_ps if ps is None else ps)
        self.dcs = set(dcs)

    @abstractmethod
    def estimate_non_fill_values(self) -> float: ...

    def get_embedding(self) -> np.ndarray:
        sizes = [float(self.dim_sizes[field]) for field in self.index_order]
        dc_embedding = [
            dc.value
            for dc in sorted(
                self.dcs,
                key=lambda dc: (
                    tuple(sorted(str(f) for f in dc.from_indices)),
                    tuple(sorted(str(f) for f in dc.to_indices)),
                    dc.p,
                ),
            )
        ]
        return np.log2(np.array([*sizes, *dc_embedding]))

BoundStatsFactory

Bases: BaseTensorStatsFactory[TS], StatsFactory[TS], Generic[TS]

Shared statistics propagation for :class:BoundStats subclasses, combining degree constraints keyed by (from_indices, to_indices, p). The only per-subclass knob is ps.

Source code in src/finch/autoschedule/tensor_stats/bound_stats.py

class BoundStatsFactory(BaseTensorStatsFactory[TS], StatsFactory[TS], Generic[TS]):
    """
    Shared statistics propagation for :class:`BoundStats` subclasses, combining
    degree constraints keyed by ``(from_indices, to_indices, p)``. The only
    per-subclass knob is ``ps``.
    """

    def __init__(self, stats_cls: type[TS], ps: Iterable[float] | None = None):
        super().__init__(stats_cls)
        self.ps: tuple[float, ...] = tuple(stats_cls.default_ps if ps is None else ps)

    def __call__(self, tensor: Any, fields: tuple[Field, ...]) -> TS:
        base = super().__call__(tensor, fields)
        dcs = self.compute_dcs(tensor, fields, self.ps)
        return self.stats_cls(base, dcs, self.ps)

    # Per field: the distinct-value count DC(∅, {i}) and, for each p, the ell_p
    # norm DC({i}, {*fields}, p); plus the full-tensor nnz DC(∅, {*fields}). The
    # ∅-conditioned records are length-1 sequences (p-independent), so kept once.
    @staticmethod
    def compute_dcs(
        tensor: Tensor,
        fields: Iterable[Field],
        ps: Iterable[float],
    ) -> set[DC]:
        fields = list(fields)
        if tensor.ndim == 0:
            return {DC(frozenset(), frozenset(), 1.0)}
        dcs: set[DC] = set()
        norms = get_lp_norms(tensor, fields, ps)
        for field in fields:
            for i, p in enumerate(ps):
                if p == 0:
                    dcs.add(
                        DC(frozenset(), frozenset({field}), norms[field][i], math.inf)
                    )
                elif p == 1:
                    dcs.add(
                        DC(frozenset(), frozenset(fields), norms[field][i], math.inf)
                    )
                else:
                    dcs.add(
                        DC(frozenset({field}), frozenset(fields), norms[field][i], p)
                    )
        return dcs

    def _rebuild(self, base: BaseTensorStats, dcs: Iterable[DC]) -> TS:
        return self.stats_cls(base, dcs, self.ps)

    def _mapjoin_union(self, op: FinchOperator, *union_args: TS) -> TS:
        base = super()._mapjoin_defs(op, *union_args)

        if len(union_args) == 1:
            return self._rebuild(base, union_args[0].dcs)

        dc_keys: Counter[_Key] = Counter()
        stats_dcs: list[dict[_Key, float]] = []
        for stats in union_args:
            dcs: dict[_Key, float] = {}
            extra_fields = tuple(
                x for x in base.index_order if x not in stats.index_order
            )
            extra_size = base.get_dim_space_size(extra_fields)
            for dc in stats.dcs:
                new_key = (dc.from_indices, dc.to_indices, dc.p)
                dcs[new_key] = dc.value
                dc_keys[new_key] += 1

                ext_key = (
                    dc.from_indices,
                    dc.to_indices | frozenset(extra_fields),
                    dc.p,
                )
                if ext_key not in dcs:
                    dc_keys[ext_key] += 1
                prev = dcs.get(ext_key, math.inf)
                dcs[ext_key] = min(prev, dc.value * extra_size)
            stats_dcs.append(dcs)

        new_dcs: dict[_Key, float] = {}
        for key, count in dc_keys.items():
            if count == len(union_args):
                total = sum(d.get(key, 0.0) for d in stats_dcs)
                _, to_indices, _ = key
                if to_indices.issubset(base.index_order):
                    total = min(total, base.get_dim_space_size(to_indices))
                new_dcs[key] = min(float(2**64), total)

        new_stats = {
            DC(from_indices, to_indices, value, p)
            for (from_indices, to_indices, p), value in new_dcs.items()
        }
        return self._rebuild(base, new_stats)

    def _mapjoin_join(self, op: FinchOperator, *join_args: TS) -> TS:
        base = super()._mapjoin_defs(op, *join_args)

        if len(join_args) == 1:
            return self._rebuild(base, join_args[0].dcs)

        new_dc: dict[_Key, float] = {}
        for stats in join_args:
            for dc in stats.dcs:
                key = (dc.from_indices, dc.to_indices, dc.p)
                if dc.value < new_dc.get(key, math.inf):
                    new_dc[key] = dc.value

        new_stats = {
            DC(from_indices, to_indices, value, p)
            for (from_indices, to_indices, p), value in new_dc.items()
        }
        return self._rebuild(base, new_stats)

    def aggregate(
        self,
        op: FinchOperator,
        init: Any | None,
        reduce_indices: tuple[Field, ...],
        stats: TS,
    ) -> TS:
        if len(reduce_indices) == 0:
            base: BaseTensorStats = stats.copy()
        else:
            base = self.aggregate_def(op, init, reduce_indices, stats)
        return self._rebuild(base, stats.dcs)

    def relabel(self, stats: TS, relabel_indices: tuple[Field, ...]) -> TS:
        base = self.relabel_def(stats, relabel_indices)
        mapping = dict(zip(stats.index_order, relabel_indices, strict=True))
        dcs = {
            DC(
                frozenset(mapping.get(f, f) for f in dc.from_indices),
                frozenset(mapping.get(f, f) for f in dc.to_indices),
                dc.value,
                dc.p,
            )
            for dc in stats.dcs
        }
        return self._rebuild(base, dcs)

    def reorder(self, stats: TS, reorder_indices: tuple[Field, ...]) -> TS:
        base = self.reorder_def(stats, reorder_indices)
        return self._rebuild(base, stats.dcs)

DC dataclass

A degree constraint: value bounds ||deg(to_indices | from_indices)||_p.

p = inf (the default) is the classic max-degree constraint; p = 1 is the distinct count; intermediate p capture skew between the two.

Source code in src/finch/autoschedule/tensor_stats/bound_stats.py

@dataclass(frozen=True)
class DC:
    """
    A degree constraint: ``value`` bounds ``||deg(to_indices | from_indices)||_p``.

    ``p = inf`` (the default) is the classic max-degree constraint; ``p = 1`` is
    the distinct count; intermediate ``p`` capture skew between the two.
    """

    from_indices: frozenset[Field]
    to_indices: frozenset[Field]
    value: float
    p: float = math.inf

DCStats

Bases: BoundStats

Classic degree constraints -- max degrees (p = inf) and distinct counts. Estimates size with a best-first product search over the constraints, clamped by the tensor's dense capacity.

Source code in src/finch/autoschedule/tensor_stats/bound_stats.py

class DCStats(BoundStats):
    """
    Classic degree constraints -- max degrees (``p = inf``) and distinct counts.
    Estimates size with a best-first product search over the constraints,
    clamped by the tensor's dense capacity.
    """

    default_ps: ClassVar[tuple[float, ...]] = (math.inf,)

    def estimate_non_fill_values(self) -> float:
        """Smallest degree-constraint product covering all indices, clamped by
        dense capacity."""
        idx: frozenset[Field] = frozenset(self.dim_sizes.keys())
        if len(idx) == 0:
            return 1.0

        best: dict[frozenset[Field], float] = {frozenset(): 1.0}
        frontier: set[frozenset[Field]] = {frozenset()}

        while True:
            current_bound = best.get(idx, math.inf)
            new_frontier: set[frozenset[Field]] = set()

            for node in frontier:
                for dc in self.dcs:
                    if node.issuperset(dc.from_indices):
                        y = node.union(dc.to_indices)
                        if best[node] > float(2 ** (64 - 2)) or float(dc.value) > float(
                            2 ** (64 - 2)
                        ):
                            y_weight = float(2**64)
                        else:
                            y_weight = best[node] * dc.value
                        if min(current_bound, best.get(y, math.inf)) > y_weight:
                            best[y] = y_weight
                            new_frontier.add(y)
            if len(new_frontier) == 0:
                break
            frontier = new_frontier

        min_weight = float(self.get_dim_space_size(idx))
        for node, weight in best.items():
            if node.issuperset(idx):
                min_weight = min(min_weight, weight)
        return min_weight

estimate_non_fill_values()

Smallest degree-constraint product covering all indices, clamped by dense capacity.

Source code in src/finch/autoschedule/tensor_stats/bound_stats.py

def estimate_non_fill_values(self) -> float:
    """Smallest degree-constraint product covering all indices, clamped by
    dense capacity."""
    idx: frozenset[Field] = frozenset(self.dim_sizes.keys())
    if len(idx) == 0:
        return 1.0

    best: dict[frozenset[Field], float] = {frozenset(): 1.0}
    frontier: set[frozenset[Field]] = {frozenset()}

    while True:
        current_bound = best.get(idx, math.inf)
        new_frontier: set[frozenset[Field]] = set()

        for node in frontier:
            for dc in self.dcs:
                if node.issuperset(dc.from_indices):
                    y = node.union(dc.to_indices)
                    if best[node] > float(2 ** (64 - 2)) or float(dc.value) > float(
                        2 ** (64 - 2)
                    ):
                        y_weight = float(2**64)
                    else:
                        y_weight = best[node] * dc.value
                    if min(current_bound, best.get(y, math.inf)) > y_weight:
                        best[y] = y_weight
                        new_frontier.add(y)
        if len(new_frontier) == 0:
            break
        frontier = new_frontier

    min_weight = float(self.get_dim_space_size(idx))
    for node, weight in best.items():
        if node.issuperset(idx):
            min_weight = min(min_weight, weight)
    return min_weight

DCStatsFactory

Bases: BoundStatsFactory['DCStats']

Factory for :class:DCStats -- keeps only max-degree (p = inf).

Source code in src/finch/autoschedule/tensor_stats/bound_stats.py

class DCStatsFactory(BoundStatsFactory["DCStats"]):
    """Factory for :class:`DCStats` -- keeps only max-degree (``p = inf``)."""

    def __init__(self):
        super().__init__(DCStats, ps=(0, 1, math.inf))

LPStats

Bases: BoundStats

LpBound statistics: keeps ell_p-norms of degree sequences for a configurable set of p (not just max degree), and bounds size via the LPflow network-flow LP (:func:_flow_bound).

Source code in src/finch/autoschedule/tensor_stats/bound_stats.py

class LPStats(BoundStats):
    """
    LpBound statistics: keeps ``ell_p``-norms of degree sequences for a
    configurable set of ``p`` (not just max degree), and bounds size via the
    ``LPflow`` network-flow LP (:func:`_flow_bound`).
    """

    default_ps: ClassVar[tuple[float, ...]] = DEFAULT_PS

    def estimate_non_fill_values(self) -> float:
        idx: frozenset[Field] = frozenset(self.dim_sizes.keys())
        if len(idx) == 0:
            return 1.0

        # An empty (all-fill) input has no non-fill values.
        for dc in self.dcs:
            if len(dc.from_indices) == 0 and dc.value == 0:
                return 0.0

        positive = [dc for dc in self.dcs if dc.value > 0]
        dense_cap = float(self.get_dim_space_size(idx))
        sol = _flow_bound(positive, list(self.dim_sizes.keys()))
        if sol is None:
            return dense_cap
        try:
            bound = math.pow(2.0, sol)
        except OverflowError:
            return dense_cap
        return min(bound, dense_cap)

LPStatsFactory

Bases: BoundStatsFactory['LPStats']

Factory for :class:LPStats with a configurable set of p norms.

Source code in src/finch/autoschedule/tensor_stats/bound_stats.py

class LPStatsFactory(BoundStatsFactory["LPStats"]):
    """Factory for :class:`LPStats` with a configurable set of ``p`` norms."""

    def __init__(self, ps: Iterable[float] | None = None):
        super().__init__(LPStats, ps)

SamplingStats

Bases: NumericStats

sketch : materialized table over bound dimensions remainder_size : product of sizes of dimensions aggregated out of the sketch remainder_prob : product of their sampling probabilities sample_probs : per-dimension probabilities in index_order

Source code in src/finch/autoschedule/tensor_stats/sampling_stats.py

class SamplingStats(NumericStats):
    """
    sketch : materialized table over bound dimensions
    remainder_size : product of sizes of dimensions aggregated out of the sketch
    remainder_prob : product of their sampling probabilities
    sample_probs : per-dimension probabilities in index_order
    """

    sketch: LogicExpression
    remainder_size: float
    remainder_prob: float
    sample_probs: list[float]

    def __init__(
        self,
        base: BaseTensorStats,
        sketch: LogicExpression,
        sample_probs: list[float],
        estimator: str = "uj1",
        remainder_size: float = 1.0,
        remainder_prob: float = 1.0,
        seeds_ref: dict[tuple[Field, int], int] | None = None,
    ):
        super().__init__(base)
        self.sketch = sketch
        if len(sample_probs) != len(self.index_order):
            raise ValueError("Expected one sampling probability per dimension")
        self.sample_probs = list(sample_probs)
        self.remainder_size = float(remainder_size)
        self.remainder_prob = float(remainder_prob)
        self.estimator = estimator
        self.seeds_ref = seeds_ref if seeds_ref is not None else {}
        self.scan_cache: tuple[float, float, float, dict | None] | None = None

    def scan(self, needs_freq: bool) -> tuple[float, float, float, dict | None]:
        from finch.compile_jl.runtime import JuliaOwnedTensor

        if self.scan_cache is None or (needs_freq and self.scan_cache[3] is None):
            match self.sketch:
                case Table(Literal(tensor), _):
                    pass
                case _:
                    raise TypeError("scan requires a materialized sketch")
            match tensor:
                case JuliaOwnedTensor():
                    tensor = tensor._as_tensor()
            match tensor:
                case FiberTensor(lvl=level):
                    while True:
                        match level:
                            case ElementLevel():
                                values = level.val.arr
                                break
                            case (
                                DenseLevel(lvl=child)
                                | SparseListLevel(lvl=child)
                                | SparseCOOLevel(lvl=child)
                                | SparseHashLevel(lvl=child)
                                | SparseByteMapLevel(lvl=child)
                            ):
                                level = child
                            case _:
                                raise TypeError(
                                    f"Unsupported sketch level: {type(level).__name__}"
                                )
                case _:
                    values = np.asarray(tensor).reshape(-1)
            vals, cts = np.unique_counts(values)
            n = float(np.dot(vals.astype(float), cts))
            d_n = float(cts[vals > 0].sum())
            f_1 = float(cts[vals == 1].sum())
            frequencies = (
                {int(v): float(c) for v, c in zip(vals, cts, strict=True) if v > 0}
                if needs_freq
                else None
            )
            self.scan_cache = (n, d_n, f_1, frequencies or None)
        return self.scan_cache

    def coverage_correction(self) -> float:
        from finch.autoschedule.default_schedulers import NON_RECURSIVE_SCHEDULER

        _, d_n_raw, _, _ = self.scan(needs_freq=False)
        coverage = 1.0
        for field, prob in zip(self.index_order, self.sample_probs, strict=True):
            size = int(self.dim_sizes[field])
            seed = self.seeds_ref.get((field, size))
            if seed is None or size == 0:
                continue
            mask = RandomMaskTensor(size, prob, seed=seed, dtype=np.intp)
            out = Alias("sampled_count")
            query = Query(
                Table(out, ()),
                Aggregate(
                    Literal(ffuncs.add),
                    Literal(np.intp(0)),
                    mask_table(field, mask),
                    (field,),
                ),
            )
            (count,) = NON_RECURSIVE_SCHEDULER(Plan((query, Produces((out,)))))
            actual_fraction = float(np.asarray(count)[()]) / size
            if actual_fraction <= 0:
                continue
            coverage *= actual_fraction
        if coverage <= 0:
            return d_n_raw
        return d_n_raw / coverage

    def estimate_non_fill_values(self) -> float:
        """Correct reduced-dimension sampling, then scale to all output coordinates."""
        q_output = math.prod(self.sample_probs)
        q = self.remainder_prob
        needs_freq = (
            self.estimator in ("uj2", "schlosser", "sh2", "sh3", "good1")
            and 0.0 < q < 1.0
        )
        n, d_n, f_1, frequencies = self.scan(needs_freq)
        if d_n == 0 or q_output <= 0:
            return 0.0
        bound_size = math.prod(int(self.dim_sizes[f]) for f in self.index_order)
        # The formulas estimate classes among the retained output coordinates.
        # Their population consists of nonzero contributions, not all tensor cells.
        N = n / q if q > 0 else n

        if self.estimator == "uj1":
            formula_est = _duj1(d_n, f_1, q, n)

        elif self.estimator == "good1":
            formula_est = _dgood1(d_n, frequencies, n, N)

        elif self.estimator == "sj1":
            formula_est = _dsj1(d_n, q, N)

        elif self.estimator == "uj2":
            formula_est = _duj2(d_n, f_1, frequencies, q, n, N)
        elif self.estimator == "schlosser":
            formula_est = _dsh(d_n, f_1, frequencies, q, n)

        elif self.estimator == "sh2":
            formula_est = _dsh2(d_n, f_1, frequencies, q, n, N)

        elif self.estimator == "sh3":
            formula_est = _dsh3(d_n, f_1, frequencies, q, n)

        elif self.estimator == "silly":
            formula_est = d_n

        else:
            raise ValueError(
                f"Unknown estimator: {self.estimator!r}."
                f"Choose from: {', '.join(SAMPLING_ESTIMATORS)}"
            )

        return float(min(bound_size, max(d_n, formula_est) / q_output))

    def get_embedding(self) -> np.ndarray:
        sizes = [float(self.dim_sizes[f]) for f in self.index_order]
        nnz = self.estimate_non_fill_values()
        size_part = np.log2(np.array(sizes))
        nnz_part = np.log2(np.array(nnz) + 1)
        return np.concatenate([size_part, nnz_part])

estimate_non_fill_values()

Correct reduced-dimension sampling, then scale to all output coordinates.

Source code in src/finch/autoschedule/tensor_stats/sampling_stats.py

def estimate_non_fill_values(self) -> float:
    """Correct reduced-dimension sampling, then scale to all output coordinates."""
    q_output = math.prod(self.sample_probs)
    q = self.remainder_prob
    needs_freq = (
        self.estimator in ("uj2", "schlosser", "sh2", "sh3", "good1")
        and 0.0 < q < 1.0
    )
    n, d_n, f_1, frequencies = self.scan(needs_freq)
    if d_n == 0 or q_output <= 0:
        return 0.0
    bound_size = math.prod(int(self.dim_sizes[f]) for f in self.index_order)
    # The formulas estimate classes among the retained output coordinates.
    # Their population consists of nonzero contributions, not all tensor cells.
    N = n / q if q > 0 else n

    if self.estimator == "uj1":
        formula_est = _duj1(d_n, f_1, q, n)

    elif self.estimator == "good1":
        formula_est = _dgood1(d_n, frequencies, n, N)

    elif self.estimator == "sj1":
        formula_est = _dsj1(d_n, q, N)

    elif self.estimator == "uj2":
        formula_est = _duj2(d_n, f_1, frequencies, q, n, N)
    elif self.estimator == "schlosser":
        formula_est = _dsh(d_n, f_1, frequencies, q, n)

    elif self.estimator == "sh2":
        formula_est = _dsh2(d_n, f_1, frequencies, q, n, N)

    elif self.estimator == "sh3":
        formula_est = _dsh3(d_n, f_1, frequencies, q, n)

    elif self.estimator == "silly":
        formula_est = d_n

    else:
        raise ValueError(
            f"Unknown estimator: {self.estimator!r}."
            f"Choose from: {', '.join(SAMPLING_ESTIMATORS)}"
        )

    return float(min(bound_size, max(d_n, formula_est) / q_output))

SamplingStatsFactory

Bases: BaseTensorStatsFactory['SamplingStats'], StatsFactory['SamplingStats']

Source code in src/finch/autoschedule/tensor_stats/sampling_stats.py

class SamplingStatsFactory(
    BaseTensorStatsFactory["SamplingStats"], StatsFactory["SamplingStats"]
):
    def __init__(self, sample_nnz: int = 10_000, estimator: str = "sh3"):
        super().__init__(SamplingStats)
        if sample_nnz < 1:
            raise ValueError("sample_nnz must be positive")
        if estimator not in SAMPLING_ESTIMATORS:
            raise ValueError(f"Unknown estimator: {estimator!r}")
        self.sample_nnz = sample_nnz
        self.estimator = estimator
        self._seeds: dict[tuple[Field, int], int] = {}
        self._rng = np.random.default_rng()

    def _get_mask(self, field: Field, size: int, prob: float) -> RandomMaskTensor:
        seed_key = (field, size)
        if seed_key not in self._seeds:
            self._seeds[seed_key] = int(self._rng.integers(0, 1 << 64, dtype=np.uint64))
        return RandomMaskTensor(size, prob, seed=self._seeds[seed_key], dtype=np.intp)

    def __call__(
        self,
        tensor: Any,
        fields: tuple[Field, ...],
        sample_probs: list[float] | None = None,
    ) -> SamplingStats:
        base = super().__call__(tensor, fields)
        fill = tensor.fill_value
        sample_probs = (
            [1.0] * len(fields) if sample_probs is None else list(sample_probs)
        )

        # Reuse each field's seed across tensors so joins sample the same
        # coordinates. An entry survives only if every dimension is kept.
        masks = [
            self._get_mask(field, int(base.dim_sizes[field]), prob)
            for field, prob in zip(fields, sample_probs, strict=True)
        ]
        non_fill = MapJoin(
            Literal(ffuncs.ne), (Table(Literal(tensor), fields), Literal(fill))
        )
        mask_tables = [
            mask_table(field, mask) for field, mask in zip(fields, masks, strict=True)
        ]

        sketch = compute_sketch(MapJoin(Literal(ffuncs.mul), (non_fill, *mask_tables)))

        return self.remask(
            SamplingStats(
                base,
                sketch=sketch,
                sample_probs=sample_probs,
                estimator=self.estimator,
                seeds_ref=self._seeds,
            )
        )

    def _mask_sketch(self, stats: SamplingStats, probs: list[float]) -> SamplingStats:
        result = self.copy(stats)
        masks = []
        for field, old_prob, prob in zip(
            stats.index_order, stats.sample_probs, probs, strict=True
        ):
            if prob == old_prob:
                continue
            size = int(stats.dim_sizes[field])
            seed = stats.seeds_ref.get((field, size))
            if seed is None:
                mask = self._get_mask(field, size, prob)
                result.seeds_ref[field, size] = self._seeds[field, size]
            else:
                mask = RandomMaskTensor(size, prob, seed=seed, dtype=np.intp)
            masks.append(mask_table(field, mask))
        if masks:
            result.sketch = compute_sketch(
                MapJoin(Literal(ffuncs.mul), (stats.sketch, *masks))
            )
            result.sample_probs = list(probs)
            result.scan_cache = None
        return result

    def remask(self, stats: SamplingStats) -> SamplingStats:
        """Halve the largest nonzero projection until the sketch fits the budget."""
        if (
            math.prod(int(stats.dim_sizes[f]) for f in stats.index_order)
            <= self.sample_nnz
        ):
            return stats
        while stats.scan(needs_freq=False)[1] > self.sample_nnz:
            match stats.sketch:
                case Table(Literal(tensor), fields):
                    norms = get_lp_norms(tensor, fields, (0,))
                case _:
                    raise TypeError("remask requires a materialized sketch")
            dim = max(
                range(len(stats.index_order)),
                key=lambda idx: norms[stats.index_order[idx]][0],
            )
            probs = list(stats.sample_probs)
            probs[dim] *= 0.5
            stats = self._mask_sketch(stats, probs)
        return stats

    def _align_mapjoin(
        self, *args: SamplingStats
    ) -> tuple[tuple[SamplingStats, ...], dict[Field, float]]:
        probs: dict[Field, float] = {}
        for arg in args:
            for field, prob in zip(arg.index_order, arg.sample_probs, strict=True):
                probs[field] = min(probs.get(field, prob), prob)
        aligned = tuple(
            self._mask_sketch(arg, [probs[f] for f in arg.index_order]) for arg in args
        )
        return aligned, probs

    def _mapjoin_join(
        self, op: FinchOperator, *join_args: SamplingStats
    ) -> SamplingStats:
        """
        N(C)_i =  N(A)_j * N(B)_k
        """

        if len(join_args) == 1:
            return self.remask(self.copy(join_args[0]))

        base_stats = super()._mapjoin_defs(op, *join_args)
        join_args, probs = self._align_mapjoin(*join_args)
        result_sketch = compute_sketch(
            MapJoin(Literal(ffuncs.mul), tuple(arg.sketch for arg in join_args))
        )

        return self.remask(
            SamplingStats(
                base_stats,
                sketch=result_sketch,
                sample_probs=[probs[f] for f in base_stats.index_order],
                remainder_size=math.prod(arg.remainder_size for arg in join_args),
                remainder_prob=math.prod(arg.remainder_prob for arg in join_args),
                estimator=self.estimator,
                seeds_ref=self._seeds,
            )
        )

    def _mapjoin_union(self, op: FinchOperator, *union_args: SamplingStats):
        """
        N(C)_i = sum_{juk\\i}[N(A)_j *prod_{l in k\\j}n(B)_l +
        N(B)_k *prod_{l in j\\k}n(A)_l ] - N(A)_j*N(B)_k
        """
        base_stats = super()._mapjoin_defs(op, *union_args)
        union_args, probs = self._align_mapjoin(*union_args)
        output_indices = set(base_stats.index_order)

        terms = []
        for arg in union_args:
            arg_indices = set(arg.index_order)
            other_free_size = 1.0
            for other in union_args:
                if other is arg:
                    continue
                for f in other.index_order:
                    if f not in arg_indices and f not in output_indices:
                        other_free_size *= other.dim_sizes.get(f, 1.0)
                other_free_size *= other.remainder_size
            missing_masks = [
                mask_table(
                    field,
                    self._get_mask(
                        field, int(base_stats.dim_sizes[field]), probs[field]
                    ),
                )
                for field in base_stats.index_order
                if field not in arg_indices
            ]
            terms.append(
                MapJoin(
                    Literal(ffuncs.mul),
                    (arg.sketch, Literal(other_free_size), *missing_masks),
                )
            )

        result = terms[0]
        for term in terms[1:]:
            result = MapJoin(Literal(ffuncs.add), (result, term))

        if len(union_args) >= 2:
            inter = MapJoin(
                Literal(ffuncs.mul), tuple(arg.sketch for arg in union_args)
            )

            result = MapJoin(Literal(ffuncs.sub), (result, inter))

        return self.remask(
            SamplingStats(
                base_stats,
                sketch=compute_sketch(result),
                sample_probs=[probs[f] for f in base_stats.index_order],
                remainder_size=math.prod(arg.remainder_size for arg in union_args),
                remainder_prob=math.prod(arg.remainder_prob for arg in union_args),
                estimator=self.estimator,
                seeds_ref=self._seeds,
            )
        )

    def aggregate(
        self,
        op: FinchOperator,
        init: Any | None,
        reduce_indices: tuple[Field, ...],
        stats: SamplingStats,
    ):
        """
        op is identity on fill : N(B)_i = sum_j N(A)_j
        op annihilates fill: N(B)_i = prod(l in k)n_l * min_k 1[N(A)_j > 0]
        otherwise : N(B)_i = prod(l in k)n_l * exists(N(A)_j)
        """

        base_stats = self.aggregate_def(op, init, reduce_indices, stats)
        reduce_set = set(reduce_indices) & set(stats.index_order)
        reduce_fields = tuple(reduce_set)

        # check is_annihilator
        if is_identity(op.ftype, stats.fill_type):
            new_sketch = (
                Aggregate(
                    Literal(ffuncs.add),
                    Literal(np.intp(0)),
                    stats.sketch,
                    reduce_fields,
                )
                if reduce_fields
                else stats.sketch
            )

        elif is_annihilator(op.ftype, stats.fill_type):
            exists: LogicExpression = MapJoin(
                Literal(ffuncs.gt), (stats.sketch, Literal(0.0))
            )
            if reduce_fields:
                exists = Aggregate(
                    Literal(ffuncs.min), Literal(np.float64(1.0)), exists, reduce_fields
                )
            prod_n = math.prod(int(stats.dim_sizes[f]) for f in reduce_set)
            new_sketch = MapJoin(Literal(ffuncs.mul), (exists, Literal(prod_n)))
        else:
            prod_n = math.prod(int(stats.dim_sizes[f]) for f in reduce_set)
            exists = MapJoin(Literal(ffuncs.gt), (stats.sketch, Literal(0.0)))
            if reduce_fields:
                exists = Aggregate(
                    Literal(ffuncs.max), Literal(np.float64(0.0)), exists, reduce_fields
                )
            new_sketch = MapJoin(Literal(ffuncs.mul), (exists, Literal(prod_n)))

        probs = dict(zip(stats.index_order, stats.sample_probs, strict=True))

        return self.remask(
            SamplingStats(
                base_stats,
                sketch=compute_sketch(new_sketch),
                sample_probs=[probs[f] for f in base_stats.index_order],
                remainder_size=stats.remainder_size
                * math.prod(stats.dim_sizes[f] for f in reduce_fields),
                remainder_prob=stats.remainder_prob
                * math.prod(probs[f] for f in reduce_fields),
                estimator=self.estimator,
                seeds_ref=stats.seeds_ref,
            )
        )

    def relabel(
        self, stats: SamplingStats, relabel_indices: tuple[Field, ...]
    ) -> SamplingStats:
        base_stats = self.relabel_def(stats, relabel_indices)
        seeds = dict(stats.seeds_ref)
        for old, new in zip(stats.index_order, relabel_indices, strict=True):
            size = int(stats.dim_sizes[old])
            if (old, size) in stats.seeds_ref:
                seeds[new, size] = stats.seeds_ref[old, size]
        return self.remask(
            SamplingStats(
                base_stats,
                sketch=compute_sketch(Relabel(stats.sketch, relabel_indices)),
                sample_probs=stats.sample_probs,
                remainder_size=stats.remainder_size,
                remainder_prob=stats.remainder_prob,
                estimator=self.estimator,
                seeds_ref=seeds,
            )
        )

    def reorder(
        self, stats: SamplingStats, reorder_indices: tuple[Field, ...]
    ) -> SamplingStats:
        base_stats = self.reorder_def(stats, reorder_indices)
        probs = dict(zip(stats.index_order, stats.sample_probs, strict=True))
        return self.remask(
            SamplingStats(
                base_stats,
                sketch=compute_sketch(Reorder(stats.sketch, reorder_indices)),
                sample_probs=[probs.get(f, 1.0) for f in base_stats.index_order],
                remainder_size=stats.remainder_size,
                remainder_prob=stats.remainder_prob,
                estimator=self.estimator,
                seeds_ref=stats.seeds_ref,
            )
        )

aggregate(op, init, reduce_indices, stats)

op is identity on fill : N(B)_i = sum_j N(A)_j op annihilates fill: N(B)_i = prod(l in k)n_l * min_k 1[N(A)_j > 0] otherwise : N(B)_i = prod(l in k)n_l * exists(N(A)_j)

Source code in src/finch/autoschedule/tensor_stats/sampling_stats.py

def aggregate(
    self,
    op: FinchOperator,
    init: Any | None,
    reduce_indices: tuple[Field, ...],
    stats: SamplingStats,
):
    """
    op is identity on fill : N(B)_i = sum_j N(A)_j
    op annihilates fill: N(B)_i = prod(l in k)n_l * min_k 1[N(A)_j > 0]
    otherwise : N(B)_i = prod(l in k)n_l * exists(N(A)_j)
    """

    base_stats = self.aggregate_def(op, init, reduce_indices, stats)
    reduce_set = set(reduce_indices) & set(stats.index_order)
    reduce_fields = tuple(reduce_set)

    # check is_annihilator
    if is_identity(op.ftype, stats.fill_type):
        new_sketch = (
            Aggregate(
                Literal(ffuncs.add),
                Literal(np.intp(0)),
                stats.sketch,
                reduce_fields,
            )
            if reduce_fields
            else stats.sketch
        )

    elif is_annihilator(op.ftype, stats.fill_type):
        exists: LogicExpression = MapJoin(
            Literal(ffuncs.gt), (stats.sketch, Literal(0.0))
        )
        if reduce_fields:
            exists = Aggregate(
                Literal(ffuncs.min), Literal(np.float64(1.0)), exists, reduce_fields
            )
        prod_n = math.prod(int(stats.dim_sizes[f]) for f in reduce_set)
        new_sketch = MapJoin(Literal(ffuncs.mul), (exists, Literal(prod_n)))
    else:
        prod_n = math.prod(int(stats.dim_sizes[f]) for f in reduce_set)
        exists = MapJoin(Literal(ffuncs.gt), (stats.sketch, Literal(0.0)))
        if reduce_fields:
            exists = Aggregate(
                Literal(ffuncs.max), Literal(np.float64(0.0)), exists, reduce_fields
            )
        new_sketch = MapJoin(Literal(ffuncs.mul), (exists, Literal(prod_n)))

    probs = dict(zip(stats.index_order, stats.sample_probs, strict=True))

    return self.remask(
        SamplingStats(
            base_stats,
            sketch=compute_sketch(new_sketch),
            sample_probs=[probs[f] for f in base_stats.index_order],
            remainder_size=stats.remainder_size
            * math.prod(stats.dim_sizes[f] for f in reduce_fields),
            remainder_prob=stats.remainder_prob
            * math.prod(probs[f] for f in reduce_fields),
            estimator=self.estimator,
            seeds_ref=stats.seeds_ref,
        )
    )

remask(stats)

Halve the largest nonzero projection until the sketch fits the budget.

Source code in src/finch/autoschedule/tensor_stats/sampling_stats.py

def remask(self, stats: SamplingStats) -> SamplingStats:
    """Halve the largest nonzero projection until the sketch fits the budget."""
    if (
        math.prod(int(stats.dim_sizes[f]) for f in stats.index_order)
        <= self.sample_nnz
    ):
        return stats
    while stats.scan(needs_freq=False)[1] > self.sample_nnz:
        match stats.sketch:
            case Table(Literal(tensor), fields):
                norms = get_lp_norms(tensor, fields, (0,))
            case _:
                raise TypeError("remask requires a materialized sketch")
        dim = max(
            range(len(stats.index_order)),
            key=lambda idx: norms[stats.index_order[idx]][0],
        )
        probs = list(stats.sample_probs)
        probs[dim] *= 0.5
        stats = self._mask_sketch(stats, probs)
    return stats

TensorStats

Bases: ABC

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

class TensorStats(ABC):
    @property
    @abstractmethod
    def idxs(self) -> tuple[Field, ...]: ...

    @property
    def index_order(self) -> tuple[Field, ...]:
        return self.idxs

    @property
    @abstractmethod
    def dim_sizes(self) -> MutableMapping[Field, float]: ...

    @property
    @abstractmethod
    def fill_type(self) -> AbstractFillFType:
        """The fill of the tensor's format, which a kernel may depend on."""

    @fill_type.setter
    @abstractmethod
    def fill_type(self, value: AbstractFillFType) -> None: ...

fill_type abstractmethod property writable

The fill of the tensor's format, which a kernel may depend on.