accelforge 0.0.1__py3-none-any.whl

accelforge/model/_looptree/reuse/isl/mapping_to_isl/types.py

@@ -0,0 +1,188 @@
+"""
+Relevant name changes:
+- [logical] buffer/lbuf -> buffet
+- [logical] comp/lcomp -> compute_einsum
+-
+"""
+
+from abc import ABC
+
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import List, TypeAlias
+
+import islpy as isl
+
+from accelforge.frontend.mapping import Compute, MappingNode
+from accelforge.frontend.workload import TensorName
+from accelforge.model._looptree.types import Buffet
+
+
+# Mapper intermediates.
+##
+#   @brief Iteration -> Operation relation that specifies the tiling.
+#
+#   The tiling relation allows us to distribute data and operations using the
+#   skew and data distribution relations.
+#
+#   The tiling relation may have unspecified bounds which will be inferred by
+#   LoopTree. The tiling relation that goes to the nest analysis is guaranteed
+#   to be fully specified.
+EinsumName: TypeAlias = str
+"Einsum's identifier."
+Tiling: TypeAlias = isl.Map
+"Tiling of data and operations."
+BranchTiling: TypeAlias = defaultdict[MappingNode, Tiling]
+"Relation between a node and its tiling."
+BuffetTiling: TypeAlias = defaultdict[Buffet, Tiling]
+"Relation between a buffet and its tiling."
+
+
+@dataclass(frozen=True, slots=True)
+class Tag(ABC):  # pylint: disable=too-few-public-methods
+    """Associating an element with its type metadata without introspection?"""
+
+
+class TemporalTag(Tag):  # pylint: disable=too-few-public-methods
+    """The associated element is temporally spreading?"""
+
+
+@dataclass(frozen=True, slots=True)
+class SpatialTag(Tag):  # pylint: disable=too-few-public-methods
+    """The associated element is spatially spreading?"""
+
+    spatial_dim: int
+    "The spatial dim in a given buffer?"
+    buffer: MappingNode
+    "The buffer the spatial dim is across?"
+
+
+class PipelineTag(Tag):  # pylint: disable=too-few-public-methods
+    """The associated element is pipelined?"""
+
+
+class SequentialTag(Tag):  # pylint: disable=too-few-public-methods
+    """The associated element is serialized?"""
+
+
+TEMPORAL_TAGS = (TemporalTag, SequentialTag)
+BRANCH_TAGS = (PipelineTag, SequentialTag)
+LOOP_TAGS = (TemporalTag, SpatialTag)
+
+
+@dataclass(frozen=True, slots=True)
+class TaggedMap:  # pylint: disable=too-few-public-methods
+    """A :class:`isl.Map` with its dimensions tagged."""
+
+    tags: List[Tag]
+    map_: isl.Map
+
+    def __repr__(self):
+        return f"{type(self)}({self.tags}, {self.map_})"
+
+
+class Occupancy(TaggedMap):  # pylint: disable=too-few-public-methods
+    """Location of data in [logical?] hardware elements."""
+
+    def __init__(self, tags: list[Tag], map_: isl.Map):
+        assert len(tags) == map_.dim(isl.dim_type.in_), (
+            "Occupancy labels input dims with tags\n"
+            "-------------------------------------\n"
+            f"tags: {tags}\n"
+            f"map: {map_}\n"
+        )
+        super().__init__(tags, map_)
+
+
+class OperationOccupancy(TaggedMap):  # pylint: disable=too-few-public-methods
+    """Location of operations in [logical?] hardware elements."""
+
+
+class Fill(TaggedMap):
+    """Spacetime -> fill of a logical buffer"""
+
+    def __init__(self, tags: list[Tag], map_: isl.Map):
+        assert len(tags) == map_.dim(isl.dim_type.in_), (
+            "Fill labels input dims with tags\n"
+            "--------------------------------\n"
+            f"tags: {tags}\n"
+            f"map: {map_}\n"
+        )
+        super().__init__(tags, map_)
+
+
+class Skew(TaggedMap):  # pylint: disable=too-few-public-methods
+    """TODO: Figure out what this is."""
+
+
+@dataclass(frozen=True, slots=True)
+class BufferTensorEinsum:
+    """
+    A buffet relating a [logical?] hardware element storing data, a tensor it
+    contains, and the [logical?] hardware element that is requesting the tensor.
+
+    See Also:
+    ---------
+    :class:`accelforge.model._looptree.reuse.Buffet`
+    """
+
+    buffer: str
+    "The logical name of the buffer supplying the tensor."
+    tensor: TensorName
+    "The tensor being supplied."
+    einsum: Compute
+    "The leaf in mapping doing the einsum compute on tensor."
+
+
+@dataclass(frozen=True, slots=True)
+class ComputeEinsum:
+    """A logical computation the workload? needs to carry out."""
+
+    compute: str
+    """TODO: Figure out what this does."""
+    branch_leaf_node: Compute
+    """TODO: The compute element at the leaf of a :class:`BranchTiling`"""
+
+
+# Output classes.
+@dataclass(frozen=True, slots=True)
+class SkewsInfo:  # pylint: disable=too-few-public-methods
+    """TODO: Figure out what this does."""
+
+    bte_to_skew: defaultdict[BufferTensorEinsum, Skew]
+    """Relates a :class:`~.BufferTensorEinsum` to a :class:`~.Skew`"""
+    ce_unit_to_skew: defaultdict[ComputeEinsum, Skew]
+    """Relates a :class:`~.ComputeEinsum` to a :class:`~.Skew`"""
+
+
+@dataclass(frozen=True, slots=True)
+class MappingAnalysisResult:  # pylint: disable=too-few-public-methods
+    """
+    Results of mapping analysis that will become input into reuse
+    analysis.
+    """
+
+    buffet_direct_above_sequential: defaultdict[Buffet, bool]
+    """
+    Whether a buffet is right above a sequential node. This is used when calculating
+    capacity since some data can be dropped earlier than usual when using sequential
+    mapping without tiling.
+    """
+    buffet_to_occupancy: defaultdict[BufferTensorEinsum, Occupancy]
+    """The occupancy of every buffet as defined in the mapping."""
+    compute_einsum_to_occupancy: defaultdict[ComputeEinsum, OperationOccupancy]
+    """The occupancy of every compute unit."""
+    # TODO: Figure out if this is deprecated:
+    # https://github.com/NVlabs/timeloop/blob/32370826fdf1aa3c8deb0c93e6b2a2fc7cf053aa/include/loop-analysis/mapping-to-isl/fused-mapping-to-isl.hpp#L31-L35
+    # node_to_buffets
+    # Buffets found between the current root/branch node and the next one.
+    branch_tiling: BranchTiling
+    """
+    Tiling of each branch. The tiling is a relation between tiling variables and
+    operations. An uncompletely tiled branch will have multiple-valued :class:`isl.Map`.
+    """
+    compute_to_assumed_parallelism: defaultdict[MappingNode, float]
+    """
+    We can assume an amount of parallelism to quickly calculate approx. compute
+    latency by simply dividing number of operations with assumed parallelism.
+    """

accelforge/model/_looptree/reuse/isl/spatial.py

@@ -0,0 +1,260 @@
+"""
+Handles the ISL spatial reuse functions.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Optional
+
+import islpy as isl
+
+from accelforge.frontend.mapping import MappingNode
+from accelforge.model._looptree.reuse.isl.isl_functions import (
+    insert_equal_dims_map,
+    reorder_projector,
+)
+from accelforge.model._looptree.reuse.isl.mapping_to_isl.types import (
+    TEMPORAL_TAGS,
+    Fill,
+    Occupancy,
+    SpatialTag,
+    Tag,
+    TaggedMap,
+)
+
+
+class Transfers(TaggedMap):
+    """Transfers between regions in spacetime."""
+
+
+class Reads(TaggedMap):
+    """Reads between regions in spacetime."""
+
+
+@dataclass(frozen=True, slots=True)
+class TransferInfo:
+    """Data transfer information about a certain [subset] of the chip."""
+
+    # Crucial information to transfer info.
+    fulfilled_fill: Transfers
+    """Fills done by peer-to-peer transfers."""
+    unfulfilled_fill: Fill
+    """Fills not performed."""
+    parent_reads: Reads
+    """Fills done by parent-to-child transfers."""
+    hops: isl.PwQPolynomial
+    """Peer-to-peer transfer cost metric across spacetime."""
+
+    # Metadata on what is occurring.
+    link_transfer: bool
+
+
+class TransferModel(ABC):
+    """
+    A peer-to-peer/multicast transfer model for spatial analysis.
+    """
+
+    @abstractmethod
+    def apply(self, buff: MappingNode, fills: Fill, occs: Occupancy) -> TransferInfo:
+        """
+        Given a buffer, its fills across time, and its occupancies across time,
+        calculate the spatial transfers.
+
+        Parameters
+        ----------
+        buff:
+            The buffer whose spatial analysis is being considered.
+        fills:
+            The fill of `buffer` across time from parents.
+        occs:
+            The occupancy of `buffer` across time.
+
+        Returns
+        -------
+        Fills that were fulfilled, Fills that were unfilled, and parent reads per
+        position in spacetime. Then, gets hops per timestep.
+        """
+        raise NotImplementedError(
+            f"{type(self)} has not implemented `apply(self, MappingNode, Fill, Occupancy)`"
+        )
+
+    def __repr__(self):
+        """Returns what transfer model it is."""
+        return f"{type(self)}"
+
+
+class SimpleLinkTransferModel(TransferModel):
+    """
+    Basic link transfer model.
+    """
+
+    def apply(self, buff: MappingNode, fills: Fill, occs: Occupancy) -> TransferInfo:
+        # Sanity check the fill is for the same occupancy. Necessary but insufficient proof.
+        assert fills.tags == occs.tags, (
+            "Fill and Occupancy mismatch"
+            "---------------------------"
+            f"Fill: {fills}\n"
+            f"Occs: {occs}\n"
+        )
+
+        # Gets number of input dimensions, along with spatial and temporal indices.
+        n: int = fills.map_.dim(isl.dim_type.in_)
+        spatial_dims: list[int] = get_spatial_tags_idxs(fills.tags, buff)
+        last_temporal: Optional[int] = get_last_temporal_tag_idx(fills.tags)
+
+        # No temporal or no spatial dims, you're just not moving data across time
+        # so no transfers occurring.
+        if last_temporal is None or len(spatial_dims) == 0:
+            return TransferInfo(
+                fulfilled_fill=Transfers(
+                    fills.tags, fills.map_.subtract(fills.map_)
+                ),  # Empty map
+                unfulfilled_fill=fills,  # No fulfilled_fills, so only unfulfilled_fills
+                parent_reads=Reads(
+                    occs.tags, occs.map_.subtract(occs.map_)
+                ),  # Empty map
+                hops=isl.PwQPolynomial.from_qpolynomial(
+                    isl.QPolynomial.zero_on_domain(fills.map_.domain().get_space())
+                ),
+                link_transfer=True,
+            )
+        # Gets the connectivity between points in space.
+        connectivity: isl.Map = make_mesh_connectivity(
+            len(spatial_dims), occs.map_.get_tuple_name(isl.dim_type.in_)
+        )
+        padded_connectivity: isl.Map = insert_equal_dims_map(
+            connectivity, 0, 0, n - len(spatial_dims) - 1
+        )
+        permutation: list[int] = make_connectivity_permutation(spatial_dims, n)
+        reorder_map: isl.Map = reorder_projector(
+            permutation, occs.map_.get_tuple_name(isl.dim_type.in_)
+        )
+        complete_connectivity: isl.Map = reorder_map.apply_range(
+            padded_connectivity
+        ).apply_range(reorder_map.reverse())
+
+        # Gets data available from neighbors at each point in space per time.
+        available_from_neighbors: isl.Map = complete_connectivity.apply_range(occs.map_)
+        # Prunes data that does not need to be fetched from a higher in the mem hierarchy.
+        neighbor_filled: isl.Map = fills.map_.intersect(available_from_neighbors)
+
+        return TransferInfo(
+            fulfilled_fill=Transfers(fills.tags, neighbor_filled.coalesce()),
+            unfulfilled_fill=Fill(
+                fills.tags, fills.map_.subtract(neighbor_filled).coalesce()
+            ),
+            # Empty, since only p2p analyzed.
+            parent_reads=Reads(
+                fills.tags, neighbor_filled.subtract(neighbor_filled).coalesce()
+            ),
+            hops=isl.PwQPolynomial.from_qpolynomial(
+                isl.QPolynomial.one_on_domain(neighbor_filled.wrap().get_space())
+            )
+            .intersect_domain(neighbor_filled.wrap())
+            .coalesce(),
+            link_transfer=True,
+        )
+
+
+def make_mesh_connectivity(n: int, spacetime: str) -> isl.Map:
+    """
+    Makes a neighbor-to-neighbor mesh connection given a number of spatial dims.
+
+    Parameters
+    ----------
+    n:
+        The number of spatial dimensions.
+    spacetime:
+        The name of the spacetime the mesh is operating on.
+
+    Returns
+    -------
+    A direct orthogonal adjacency map on the space `spacetime[t, x_1, x_2, ..., x_n]`
+    """
+    mesh: isl.Map
+    match (n):
+        case 2:
+            mesh = isl.Map.read_from_str(
+                isl.DEFAULT_CONTEXT,
+                "{ [t, x, y] -> [t-1, x', y'] : "
+                " (y'=y and x'=x-1) or (y'=y and x'=x+1) "
+                " or (x'=x and y'=y-1) or (x'=x and y'=y+1) }",
+            )
+        case 1:
+            mesh = isl.Map.read_from_str(
+                isl.DEFAULT_CONTEXT,
+                "{ [t, x] -> [t-1, x'] : (x'=x-1) or (x'=x+1) }",
+            )
+        case _:
+            raise ValueError(f"Cannot make mesh with {n} spatial dims")
+
+    mesh = mesh.set_tuple_name(isl.dim_type.in_, spacetime).set_tuple_name(
+        isl.dim_type.out, spacetime
+    )
+
+    return mesh
+
+
+def make_connectivity_permutation(spatial_idxs: list[int], dims: int) -> list[int]:
+    """TODO: Figure out what this is doing."""
+    permutation: list[int] = []
+
+    cur_spatial_idx: int = 0
+    for i in range(dims):
+        if cur_spatial_idx < len(spatial_idxs) and i == spatial_idxs[cur_spatial_idx]:
+            cur_spatial_idx += 1
+        else:
+            permutation.append(i)
+
+    for spatial_idx in spatial_idxs:
+        permutation.append(spatial_idx)
+
+    return permutation
+
+
+def get_spatial_tags_idxs(tags: list[Tag], buffer: MappingNode) -> list[int]:
+    """
+    Given a list if tags, identify the spatial dimensions belong to a given `buffer`.
+
+    Parameters
+    ----------
+    tags:
+        The `Occupancy` or `Fill` domain dimension tags.
+    buffer:
+        The `MappingNode` which is the logical-memory we're looking for spatial
+        dims over.
+
+    Returns
+    -------
+    A list of the spatial_dim_idxs in order.
+    """
+    spatial_dim_idxs: list[int] = [
+        i
+        for i, tag in enumerate(tags)
+        if isinstance(tag, SpatialTag) and tag.buffer == buffer
+    ]
+
+    return spatial_dim_idxs
+
+
+def get_last_temporal_tag_idx(tags: list[Tag]) -> Optional[int]:
+    """
+    Returns the idx of the deepest temporal tag in the list.
+
+    Parameters
+    ----------
+    tags:
+        A list of `Tags`.
+
+    Returns
+    -------
+    The index of the last tag that is a `TEMPORAL_TAGS`.
+    """
+    if len(tags) == 0:
+        return None
+
+    for idx, tag in reversed(list(enumerate(tags))):
+        if isinstance(tag, TEMPORAL_TAGS):
+            return idx
+
+    return None

accelforge/model/_looptree/reuse/isl/temporal.py

@@ -0,0 +1,182 @@
+"""
+Handles the ISL temporal reuse functions.
+"""
+
+from dataclasses import dataclass
+
+import islpy as isl
+
+from accelforge.model._looptree.reuse.isl.isl_functions import map_to_shifted
+from accelforge.model._looptree.reuse.isl.mapping_to_isl.types import (
+    TEMPORAL_TAGS,
+    Fill,
+    Occupancy,
+    Tag,
+)
+
+
+@dataclass(frozen=True)
+class TemporalReuse:
+    """Results for an temporal reuse analysis."""
+
+    effective_occupancy: Occupancy
+    """TODO: Figure this out."""
+    fill: Fill
+    """Data deliveries to locations in spacetime that need to be made."""
+
+
+def analyze_temporal_reuse(
+    occ: Occupancy, exploit_reuse: bool = True, multi_loop_reuse: bool = True
+) -> TemporalReuse:
+    """
+    Computes the required fill to satisfy the buffer occupancy.
+    If the buffer can `exploit_reuse`, then the fill will only consist
+    of data not currently resident in buffer.
+
+    Parameters
+    ----------
+    occ:
+        The logical occupancy to be temporally analyzed.
+    exploit_reuse:
+        Temporally exploits reuse through persisting data currently in buffer
+        to the next time step.
+    multi_loop_reuse:
+        Whether when this loop, or one above it in the memory hierarchy, loops,
+        does the buffer flush.
+
+    Returns
+    -------
+    A struct containing a `..types.Fill` which is how to load data into the buffer
+    across time and a `..types.Occupancy` describing the effective_occupancy across
+    time (i.e., what data needs to be persisted in the buffer per time step and what
+    can be ignored/purged).
+
+    TODO: Make sure spaces are named properly
+    """
+    if exploit_reuse:
+        return fill_from_occupancy(occ, multi_loop_reuse)
+    return TemporalReuse(occ, Fill(occ.tags, occ.map_))
+
+
+def fill_from_occupancy(
+    occupancy: Occupancy, multiple_loop_reuse: bool
+) -> TemporalReuse:
+    """
+    Given an occupancy and if you're allowed to reuse across loops, calculate the
+    `fill` and the `effective_occupancy` per time step.
+
+    Parameters
+    ----------
+    occupancy:
+        The logical occupancy of data in logical buffers.
+    multi_loop_reuse:
+        If you are allowed to use data between loop iterations.
+
+    Returns
+    -------
+    A `TemporalReuse` object that contains the `fill` and `effective_occpancy`
+    of the lowest buffer level.
+    """
+    # Iterates through each dimension in reverse order (i.e., deepest loop first)
+    occ = occupancy.map_.copy()
+    tags = occupancy.tags.copy()
+    for dim_idx, tag in reversed(list(enumerate(occupancy.tags))):
+        if not isinstance(tag, TEMPORAL_TAGS):
+            continue
+        # Check if temporal dimension is "trivial," i.e., equals a singular value
+        proj_occ: isl.Map = occ.project_out(
+            isl.dim_type.in_, dim_idx, 1
+        ).set_tuple_name(
+            isl.dim_type.in_, f"{occ.get_tuple_name(isl.dim_type.in_)}_abridged"
+        )
+        reinserted_occ: isl.Map = (
+            proj_occ.insert_dims(isl.dim_type.in_, dim_idx, 1).set_tuple_name(
+                isl.dim_type.in_,
+                occ.get_tuple_name(isl.dim_type.in_).removesuffix("_abridged"),
+            )
+        ).intersect_domain(occ.domain())
+
+        if occ.plain_is_equal(reinserted_occ) or occ.is_equal(reinserted_occ):
+            occ = proj_occ
+            tags.pop(dim_idx)
+            continue
+
+        # Nontrivial analysis
+        time_shift: isl.Map
+        if not multiple_loop_reuse:
+            # TODO: Verify space names are preserved and/or replace.
+            time_shift = map_to_shifted(occ.domain().get_space(), dim_idx, -1)
+        # Calculates the time_shift assuming no cache flushing for loops.
+        else:
+            # TODO: this is a better way of getting time_shift. Use method to
+            # replace the other branch (!multi_loop_reuse)
+            time_shift = construct_time_shift(occ, tags)
+
+        # Gets the fill (i.e., feeds data not currently in buffer).
+        occ_before: isl.Map = time_shift.apply_range(occ)
+        fill: isl.Map = occ.subtract(occ_before)
+
+        return TemporalReuse(Occupancy(tags, occ), Fill(tags, fill))
+
+    return TemporalReuse(Occupancy(tags, occ.coalesce()), Fill(tags, occ.coalesce()))
+
+
+def construct_time_shift(occ: isl.Map, tags: list[Tag]):
+    """
+    Given an occupancy and its input dimension tags, create the proper spatial
+    and temporal separation objects.
+
+    Parameters
+    ----------
+    occ:
+        The occupancy map we're analyzing the reuse for.
+    tags:
+        The tags of what an input represents.
+
+    Returns
+    -------
+    time_shift:
+        Relation of the current time step to the previous one across loops.
+    """
+    # Creates the spacetime deconstruction to the two separate components.
+    spacetime: isl.Set = occ.domain()
+    spacetime_to_time: isl.Map = isl.Map.identity(spacetime.get_space().map_from_set())
+    spacetime_to_space: isl.Map = isl.Map.identity(spacetime.get_space().map_from_set())
+    # Prunes out the output dimensions that do not correspond to the
+    # correct mapping into a generic space-to-space relation.
+    for idx, t in reversed(list(enumerate(tags))):
+        if not isinstance(t, TEMPORAL_TAGS):
+            spacetime_to_time = spacetime_to_time.project_out(isl.dim_type.out, idx, 1)
+        else:
+            spacetime_to_space = spacetime_to_space.project_out(
+                isl.dim_type.out, idx, 1
+            )
+    # Gets the names correct after transformations.
+    spacetime_to_time = spacetime_to_time.set_tuple_name(
+        isl.dim_type.out, f"{spacetime.get_tuple_name()}_time"
+    )
+    spacetime_to_space = spacetime_to_space.set_tuple_name(
+        isl.dim_type.out, f"{spacetime.get_tuple_name()}_space"
+    )
+
+    # Properly constrains the spacetime_to_time's domain.
+    spacetime_to_time = spacetime_to_time.intersect_domain(spacetime)
+    time_: isl.Set = spacetime_to_time.range()
+    # Creates a map of time_ to previous regions of time_.
+    time_to_past: isl.Map = (
+        isl.Map.lex_gt(time_.get_space()).intersect_domain(time_).intersect_range(time_)
+    )
+    # Restricts the relation to only the most recent previous region of time_.
+    time_to_most_recent_past = time_to_past.lexmax()
+    # Relates the current spacetime to its direct predecessor in time.
+    time_shift: isl.Map = spacetime_to_time.apply_range(
+        time_to_most_recent_past.apply_range(spacetime_to_time.reverse())
+    )
+
+    # Prunes spatial relations to only ones that are valid.
+    spacetime_space_preserver: isl.Map = spacetime_to_space.apply_range(
+        spacetime_to_space.reverse()
+    )
+    # Intersects with time_shift as space information is lost with the compression of
+    # spacetime to time_ and then rexpansion to past time_.
+    return time_shift.intersect(spacetime_space_preserver)

accelforge/model/_looptree/reuse/symbolic/__init__.py

@@ -0,0 +1 @@
+from .symbolic import *