Move BlockIdSequence to its own file (#129)

Author: jansel

helion/autotuner/block_id_sequence.py

@@ -0,0 +1,192 @@
+from __future__ import annotations
+
+import dataclasses
+from typing import TYPE_CHECKING
+from typing import Callable
+from typing import MutableSequence
+from typing import TypeVar
+
+from torch.fx.node import map_aggregate
+
+from ..exc import InvalidConfig
+from .config_fragment import ConfigSpecFragment
+from .config_fragment import assert_integer_power_of_two
+
+if TYPE_CHECKING:
+    from . import ConfigSpec
+
+    _T = TypeVar("_T")
+    _D = TypeVar("_D")
+
+
+@dataclasses.dataclass
+class _BlockIdItem:
+    # the block_indices used in the IR
+    block_ids: list[int]
+
+    @property
+    def block_id(self) -> int:
+        """Return the first block_id for this item."""
+        return self.block_ids[0]
+
+    def _fill_missing(self) -> object:
+        """Provide a value when not provided by the user."""
+        raise NotImplementedError
+
+    def _normalize(self, name: str, value: object) -> object:
+        """Validate and normalize the value for this item."""
+        raise NotImplementedError
+
+    def _fragment(self, base: ConfigSpec) -> ConfigSpecFragment:
+        """Return the fragment used for autotunging for this item."""
+        raise NotImplementedError
+
+    def _flat_config(
+        self, base: ConfigSpec, fn: Callable[[ConfigSpecFragment], object]
+    ) -> object:
+        return fn(self._fragment(base))
+
+
+_BlockIdItemT = TypeVar("_BlockIdItemT", bound=_BlockIdItem)
+
+
+class BlockIdSequence(MutableSequence[_BlockIdItemT]):
+    """
+    A sequence of _BlockIdItem subclasses that allows for efficient
+    mapping from block_id to index in the sequence.  A generic data
+    structure used to store different types of configuration specs.
+    """
+
+    def __init__(self) -> None:
+        self._data: list[_BlockIdItemT] = []
+        self._block_id_to_index: dict[int, int] = {}
+
+    def __len__(self) -> int:
+        return len(self._data)
+
+    def _reindex(self) -> None:
+        """Rebuild the mapping from block_id to index."""
+        new_index = {}
+        for i, item in enumerate(self._data):
+            for block_id in item.block_ids:
+                new_index[block_id] = i
+        self._block_id_to_index = new_index
+
+    def __getitem__(self, index: int) -> _BlockIdItemT:
+        return self._data[index]
+
+    def __setitem__(self, index: int, value: _BlockIdItemT) -> None:
+        self._data[index] = value
+        self._reindex()  # could be faster, but uncommon case
+
+    def __delitem__(self, index: int) -> None:
+        del self._data[index]
+        self._reindex()  # could be faster, but uncommon case
+
+    def clear(self) -> None:
+        self._data.clear()
+        self._block_id_to_index.clear()
+
+    def append(self, value: _BlockIdItemT) -> None:
+        """Append a new item to the end of the sequence."""
+        index = len(self._data)
+        self._data.append(value)
+        for block_id in value.block_ids:
+            self._block_id_to_index[block_id] = index
+
+    def insert(self, index: int, value: _BlockIdItemT) -> None:
+        """Insert a new item at the given index."""
+        if index == len(self._data):
+            self.append(value)
+            return
+        self._data.insert(index, value)
+        self._reindex()  # could be faster, but uncommon case
+
+    def block_id_to_index(self, block_id: int) -> int:
+        """Return the index of the block_id in the config."""
+        return self._block_id_to_index[block_id]
+
+    def block_id_lookup(self, block_id: int) -> _BlockIdItemT:
+        """Return the index of the block_id in the config."""
+        return self._data[self._block_id_to_index[block_id]]
+
+    def disable_block_id(self, block_id: int) -> None:
+        """Remove configuration choice for the given block_id."""
+        self._data = [x for x in self._data if block_id not in x.block_ids]
+        self._reindex()
+
+    def config_get(
+        self, config: list[_T], block_id: int, default: _D = None
+    ) -> _T | _D:
+        """
+        Get the config value for the given block_id, or return default if not found.
+        """
+        index = self._block_id_to_index.get(block_id, None)
+        if index is None:
+            return default
+        return config[index]
+
+    def _flat_config(
+        self, base: ConfigSpec, fn: Callable[[ConfigSpecFragment], object]
+    ) -> list[object]:
+        """Map a flattened version of the config using the given function."""
+        return [spec._flat_config(base, fn) for spec in self._data]
+
+    def _normalize(
+        self, name: str, values: object, *, flatten: bool = False
+    ) -> list[object]:
+        """Validate and normalize the values for this config item."""
+        if flatten:
+            if values is None:
+                values = ()
+            new_values = []
+            # pyre-ignore[6]
+            map_aggregate(values, new_values.append)
+            values = new_values
+        elif not isinstance(values, (list, tuple, type(None))):
+            raise InvalidConfig(
+                f"Unexpected type for config[{name!r}], expected list or None, got {type(values).__name__}"
+            )
+        values = [*(values or ())]
+        size = len(self)
+        if len(values) > size:
+            raise InvalidConfig(
+                f"Too many values for config[{name!r}], expected {size}, got {len(values)}"
+            )
+        if len(values) < size:
+            try:
+                for spec in self._data[len(values) :]:
+                    values.append(spec._fill_missing())
+            except NotImplementedError:
+                raise InvalidConfig(
+                    f"Not enough values for config[{name!r}], expected {size}, got {len(values)}"
+                ) from None
+        for i, spec in enumerate(self._data):
+            values[i] = spec._normalize(f"config[{name}][{i}]", values[i])
+        return values
+
+    def _remove_duplicates(self) -> None:
+        new_specs = []
+        for spec in self:
+            other = self.block_id_lookup(spec.block_id)
+            if other is spec:
+                new_specs.append(spec)
+            elif len(spec.block_ids) != len(other.block_ids):
+                # this will cause invalid config errors with loop orders
+                # remove them both
+                self.disable_block_id(spec.block_id)
+                self._remove_duplicates()  # start over
+                return
+        if len(new_specs) != len(self):
+            self._data = new_specs
+            self._reindex()
+
+
+class _PowerOfTwoBlockIdItem(_BlockIdItem):
+    def _normalize(self, name: str, value: object) -> int | None:
+        try:
+            return assert_integer_power_of_two(value)
+        except InvalidConfig:
+            raise InvalidConfig(
+                f"{name} must be a power of two, got {value!r}"
+            ) from None