accelforge 0.0.1__py3-none-any.whl

accelforge/frontend/_binding.py

@@ -0,0 +1,129 @@
+"""
+Defines Pydantic models to handle Binding Specifications that relate logical to
+physical architectures.
+"""
+
+from abc import abstractmethod
+from typing import Dict, Tuple
+
+from pydantic import StrictFloat
+import islpy as isl
+
+from accelforge.util._basetypes import ParsableDict, ParsableList, ParsableModel
+
+
+class Domain(ParsableModel):
+    """
+    Represents an architecture dangling reference of the binding.
+    """
+
+    name: str
+
+    @property
+    @abstractmethod
+    def isl_space(self) -> isl.Space:
+        """Gets the domain as an isl.Space"""
+        raise NotImplementedError(f"{type(self)} has not implemented isl_space")
+
+    @property
+    @abstractmethod
+    def isl_universe(self) -> isl.Set:
+        """Gets the domain as an isl.Set"""
+        raise NotImplementedError(f"{type(self)} has not implemented isl_universe")
+
+
+class LogicalDomain(Domain):
+    """
+    Represents the logical architecture domain space of logical dims × tensor ranks.
+    """
+
+    ranks: Tuple[str] = ("c", "h", "w", "p", "q", "r", "s")
+    l_dims: ParsableList[str]
+
+    @property
+    def isl_space(self) -> isl.Space:
+        return isl.Space.create_from_names(
+            isl.DEFAULT_CONTEXT, in_=self.ranks, out=self.l_dims
+        ).set_tuple_name(isl.dim_type.out, f"l_{self.name}_dims")
+
+    @property
+    def isl_universe(self) -> isl.Map:
+        return isl.Map.universe(self.isl_space)
+
+
+class PhysicalDomain(Domain):
+    """
+    Represents the logical architecture domain space of physical dims.
+    The physical space is defined as the physical architecture dims.
+    """
+
+    p_dims: ParsableList[str]
+
+    @property
+    def isl_space(self) -> isl.Space:
+        return isl.Space.create_from_names(
+            isl.DEFAULT_CONTEXT, set=self.p_dims
+        ).set_tuple_name(isl.dim_type.set, f"p_{self.name}_dims")
+
+    @property
+    def isl_universe(self) -> isl.Set:
+        return isl.Set.universe(self.isl_space)
+
+
+class BindingNode(ParsableModel):
+    """
+    How a logical architecture is implemented on a particular physical architecture
+    for a particular hardware level. Represents a injection relation between points
+    in logical to physical space.
+
+    The logical space is defined as logical architecture dims × tensor dims.
+    The physical space is defined as the physical architecture dims.
+    """
+
+    logical: LogicalDomain
+    physical: PhysicalDomain
+    relations: ParsableDict[str, str]
+
+    @property
+    def isl_relations(self) -> Dict[str, isl.Map]:
+        """
+        Converts the logical, physical, and binding relation strings into an
+        isl.Map representing the bindings at this binding node.
+        """
+
+        def islify_relation(key: str) -> isl.Map:
+            """Converts a relation at a given key into isl"""
+            relation: str = self.relations[key]
+            logical_space: isl.Space = self.logical.isl_space.set_tuple_name(
+                isl.dim_type.in_, f"{key}_ranks"
+            )
+
+            binding_space: isl.Space = logical_space.wrap().map_from_domain_and_range(
+                range=self.physical.isl_space,
+            )
+
+            # Simple bodge to get the binding space into a real space
+            binding_str: str = binding_space.to_str()
+            binding_str: str = f"{binding_str[:-1]}: {relation} {binding_str[-1]}"
+
+            binding: isl.Map = isl.Map.read_from_str(
+                ctx=isl.DEFAULT_CONTEXT, str=binding_str
+            )
+
+            return binding
+
+        isl_relations: Dict[str, isl.Map] = {
+            key: islify_relation(key) for key in self.relations
+        }
+
+        return isl_relations
+
+
+class Binding(ParsableModel):
+    """
+    A collection of binding nodes that fully specifies a relation between the
+    logical and physical space.
+    """
+
+    # version: StrictFloat
+    nodes: ParsableList[BindingNode]

accelforge/frontend/_workload_isl/__init__.py

@@ -0,0 +1,2 @@
+from accelforge.frontend.workload import *
+from ._isl import get_rank_variable_bounds

accelforge/frontend/_workload_isl/_isl.py

@@ -0,0 +1,149 @@
+import math
+import islpy as isl
+
+from accelforge.frontend.renames import RankVariable
+from accelforge.frontend.workload import Workload, TensorName, Einsum, EinsumName
+
+
+def get_einsum_operation_space(workload: Workload, einsum_name: str) -> isl.Set:
+    """Return isl.Set of all operations in an einsum."""
+    einsum_shape = workload.get_iteration_space_shape_isl_string(einsum_name)
+    rank_variable_names = ",".join(
+        map(str, workload.einsums[einsum_name].rank_variables)
+    )
+    try:
+        return isl.Set(
+            f"{{ {einsum_name}_operation[{rank_variable_names}] : {einsum_shape} }}"
+        )
+    except:
+        raise Exception(f"Error creating isl.Set for {einsum_name}: {einsum_shape}")
+
+
+def get_dim_bounds(isl_set: isl.Set) -> list[int]:
+    bounds = []
+    for i in range(isl_set.dim(isl.dim_type.set)):
+        max_val = isl_set.dim_max_val(i)
+        min_val = isl_set.dim_min_val(i)
+        shape = max_val - min_val + 1  # max is inclusive
+        try:
+            bounds.append(shape.to_python())
+        except:
+            raise Exception(
+                f"Shape is not an integer. Are all rank variables bounded? "
+                f"Shape {shape} for rank variable {i} in {isl_set}"
+            )
+    return bounds
+
+
+def get_rank_variable_bounds(
+    workload: Workload, einsum_name: EinsumName
+) -> dict[RankVariable, int]:
+    """Return dictionary mapping rank variable name to bound."""
+    operation_space = get_einsum_operation_space(workload, einsum_name)
+    dim_shapes = get_dim_bounds(operation_space)
+    return {
+        rank_var: shape
+        for rank_var, shape in zip(
+            workload.einsums[einsum_name].rank_variables, dim_shapes
+        )
+    }
+
+
+def get_projection_multi_aff(einsum: Einsum, tensor: TensorName) -> isl.MultiAff:
+    """Return isl.MultiAff of projection from einsum to tensor."""
+    rank_variables = einsum.rank_variables
+    projection = einsum.tensor_accesses[tensor].projection
+
+    rank_variables_str = ",".join(map(str, rank_variables))
+
+    projection_str = ", ".join(
+        f"{rank_name}={rank_projection}"
+        for rank_name, rank_projection in projection.items()
+    )
+
+    return isl.MultiAff(
+        f"{{ {einsum.name}_operation[{rank_variables_str}] -> "
+        f"{tensor}[{projection_str}] }}"
+    )
+
+
+def get_projection_map(einsum: Einsum, tensor: TensorName) -> isl.Map:
+    """Return isl.Map of projection from einsum to tensor."""
+    return get_projection_multi_aff(einsum, tensor).as_map()
+
+
+def get_tensor_data_space(workload: Workload, tensor: TensorName) -> isl.Set:
+    """
+    Get tensor data space based on the operation spaces of (for lack of
+    a better term)'canonical' Einsums.
+
+    Canonical Einsums (for this purpose) are all reader Einsums if the
+    tensor is only ever read or all writer EInsums if the tensor is ever
+    an output tensor.
+    """
+    writer_einsums = workload.einsums_with_tensor_as_output(tensor)
+    if len(writer_einsums) == 0:
+        reader_einsums = workload.einsums_with_tensor_as_input(tensor)
+        canonical_einsums = reader_einsums
+    else:
+        canonical_einsums = writer_einsums
+
+    tensor_data_space = None
+    for einsum in canonical_einsums:
+        operation_space = get_einsum_operation_space(workload, einsum.name)
+        projection_map = get_projection_map(einsum, tensor)
+        if tensor_data_space is None:
+            tensor_data_space = operation_space.apply(projection_map)
+        else:
+            tensor_data_space = tensor_data_space.intersect(
+                operation_space.apply(projection_map)
+            )
+
+    return tensor_data_space
+
+
+def _card_box(data_space: isl.Set) -> int:
+    dims = []
+    for i in range(data_space.dim(isl.dim_type.set)):
+        dim_min = data_space.dim_min(i)
+        dim_max = data_space.dim_max(i)
+
+        if dim_min.is_cst() and dim_max.is_cst():
+            min_val = dim_min.as_aff().get_constant_val().to_python()
+            max_val = dim_max.as_aff().get_constant_val().to_python()
+        else:
+            raise ValueError(f"Data space is not rectangular: {data_space}")
+
+        dims.append(max_val - min_val + 1)
+
+    return math.prod(dims)
+
+
+ERRMSG = """ Non-box-shaped sets are not supported. This happens if ISL is installed without
+Barvinok support. Please install ISL with Barvinok support, or use workloads with
+rectangular data spaces and operation spaces. Non-rectangular spaces occur when the
+workload contains complex expressions involving the rank variables, such as
+multi-rank-variable inequalities.
+Offending space: {space}.
+"""
+
+
+def get_tensor_size(workload: Workload, tensor: TensorName):
+    """Get the size (num. of elements) of a tensor."""
+    data_space = get_tensor_data_space(workload, tensor)
+    if data_space.is_box():
+        return _card_box(data_space)
+    if not hasattr(data_space, "card"):
+        raise RuntimeError(ERRMSG.format(space=str(data_space)))
+    card_pwqp = isl.PwQPolynomial.card(data_space)
+    return card_pwqp.eval(card_pwqp.domain().sample_point()).to_python()
+
+
+def get_operation_space_size(workload: Workload, einsum_name: str):
+    operation_space = get_einsum_operation_space(workload, einsum_name)
+    if operation_space.is_box():
+        return _card_box(operation_space)
+    if not hasattr(operation_space, "card"):
+        raise RuntimeError(ERRMSG.format(space=str(operation_space)))
+    card_pwqp = isl.PwQPolynomial.card(operation_space)
+    return card_pwqp.eval(card_pwqp.domain().sample_point()).to_python()

accelforge/frontend/_workload_isl/_symbolic.py

@@ -0,0 +1,141 @@
+from dataclasses import dataclass
+from typing import Any
+from functools import reduce
+from operator import mul
+
+import sympy
+
+from accelforge.frontend.workload import (
+    TensorName,
+    Einsum,
+    EinsumName,
+    Workload,
+    Rank,
+    RankVariable,
+)
+from ._isl import get_rank_variable_bounds
+
+
+def get_projection_expr(einsum: Einsum, tensor: TensorName) -> dict[Rank, sympy.Expr]:
+    projection = einsum.tensor_accesses[tensor].projection
+    return {
+        rank_name: sympy.parsing.sympy_parser.parse_expr(proj_str)
+        for rank_name, proj_str in projection.items()
+    }
+
+
+class Irrelevant:
+    pass
+
+
+@dataclass
+class Relevant:
+    rank: Any
+
+
+@dataclass
+class PartiallyRelevant:
+    rank: Any
+
+
+def get_rank_variable_relevancy(einsum: Einsum, tensor: TensorName):
+    relevancy = {}
+    projection = einsum.tensor_accesses[tensor].projection
+    for rank_variable in einsum.rank_variables:
+        relevancy[rank_variable] = Irrelevant()
+        for rank_name, projection_str in projection.items():
+            projection_expr = sympy.parsing.sympy_parser.parse_expr(projection_str)
+            is_simple = len(sympy.Add.make_args(projection_expr)) == 1
+            is_relevant = (
+                sympy.symbols(f"{rank_variable}") in projection_expr.free_symbols
+            )
+
+            if not is_relevant:
+                continue
+
+            if is_simple:
+                relevancy[rank_variable] = Relevant(rank=rank_name)
+            else:
+                relevancy[rank_variable] = PartiallyRelevant(rank=rank_name)
+
+            break
+    return relevancy
+
+
+def compute_dense_tile_occupancy(
+    projection_expr: dict[str, sympy.Expr], rank_variable_shapes: dict
+):
+    substitutions = [
+        (rank_variable, rank_variable_shape - 1)
+        for rank_variable, rank_variable_shape in rank_variable_shapes.items()
+    ]
+    return reduce(
+        mul,
+        [index_expr.subs(substitutions) + 1 for index_expr in projection_expr.values()],
+    )
+
+
+def compute_rank_occupancy(projection_expr: sympy.Expr, rank_variable_shapes: dict):
+    substitutions = [
+        (rank_variable, rank_variable_shape - 1)
+        for rank_variable, rank_variable_shape in rank_variable_shapes.items()
+    ]
+    return projection_expr.subs(substitutions) + 1
+
+
+def get_stride_and_halo_of_einsum(
+    einsum_name: str,
+    workload: Workload,
+    rank_variable_bounds: dict[RankVariable, int] | None = None,
+) -> dict[TensorName, dict[tuple[Rank, RankVariable]], tuple[int, int]]:
+    """
+    Get stride and halo (initial delta) for an Einsum in workload.
+
+    Returns dictionary mapping tensor to another dictionary mapping
+    (rank, rank_var) to the stride and halo.
+    """
+    stride_and_halo = {}
+    einsum = workload.einsums[einsum_name]
+    if rank_variable_bounds is None:
+        shape = get_rank_variable_bounds(workload, einsum_name)
+    else:
+        shape = rank_variable_bounds
+    for tensor in einsum.tensor_names:
+        stride_and_halo[tensor] = {}
+        tensor_stride_and_halo = stride_and_halo[tensor]
+
+        projection = get_projection_expr(einsum, tensor)
+        tensor_accesses = einsum.tensor_accesses[tensor]
+        for rank, rank_vars in tensor_accesses.rank2rank_variables.items():
+            rank_projection = projection[rank]
+            for rank_var in rank_vars:
+                stride = rank_projection.coeff(rank_var)
+
+                # Careful: in-place mutation of cons_shape
+                original_shape = shape[rank_var]
+                shape[rank_var] = 1
+                halo = compute_rank_occupancy(rank_projection, shape) - 1
+                shape[rank_var] = original_shape
+
+                tensor_stride_and_halo[(rank, rank_var)] = (stride, halo)
+    return stride_and_halo
+
+
+def get_stride_and_halo(
+    workload: Workload,
+) -> dict[
+    tuple[EinsumName, TensorName],
+    dict[tuple[Rank, RankVariable], tuple[int, int]],
+]:
+    """
+    Get stride and halo (initial delta) for Einsums in workload.
+
+    Returns dictionary mapping (Einsum, tensor) to another dictionary mapping
+    (rank, rank_var) to the stride and halo.
+    """
+    stride_and_halo = {}
+    for einsum in workload.einsums:
+        stride_and_halo_of_einsum = get_stride_and_halo_of_einsum(einsum.name, workload)
+        for tensor, ranks2stride_and_halo in stride_and_halo_of_einsum.items():
+            stride_and_halo[(einsum.name, tensor)] = ranks2stride_and_halo
+    return stride_and_halo