accelforge 0.0.1__py3-none-any.whl

accelforge/mapper/FFM/main copy.py

@@ -0,0 +1,236 @@
+from copy import deepcopy
+import inspect
+import os
+from typing import Callable
+import joblib
+import logging
+
+from accelforge import arch
+from accelforge import Spec
+from accelforge.mapper.FFM.pmappings import MultiEinsumPmappings
+from accelforge.mapper.FFM.mappings import Mappings
+import accelforge.mapper.FFM._make_pmappings.make_pmappings as pmapper
+from accelforge.frontend.workload import EinsumName
+from accelforge.mapper.FFM._join_pmappings.join_pmappings import (
+    clean_compress_and_join_pmappings,
+)
+from accelforge._accelerated_imports import pd
+
+
+logger = logging.getLogger(__name__)
+
+
+def map_workload_to_arch(
+    spec: Spec,
+    einsum_names: list[EinsumName] | None = None,
+    can_combine_multiple_runs: bool = False,
+    cache_dir: str | None = None,
+    print_number_of_pmappings: bool = True,
+    _pmapping_row_filter_function: Callable[[pd.Series], bool] | None = None,
+) -> Mappings:
+    """
+    Maps a workload to an architecture using the AccelForge FFM mapper.
+
+    Parameters
+    ----------
+    spec:
+        The Spec to map.
+    einsum_names:
+        The einsum names to map. If None, all einsums will be mapped.
+    can_combine_multiple_runs: Whether we would like to be able to combine multiple
+        make_pmappings runs. Having this as True allows you to do things like
+        `pmappings = make_pmappings(*args_a) | make_pmappings(*args_b)` but slows
+        down execution.
+    cache_dir:
+        The directory to cache pmappings in. If None, no caching will be done.
+    print_number_of_pmappings:
+        Whether to print the number of pmappings for each einsum.
+    _pmapping_row_filter_function:
+        A function that takes in a row of the pmapping dataframe and returns
+        True if the row should be included in the final mappings, and False
+        otherwise. If None, all rows will be included.
+    """
+    from accelforge.model.main import evaluate_mapping
+
+    pmappings = make_pmappings(
+        spec,
+        einsum_names=einsum_names,
+        can_combine_multiple_runs=can_combine_multiple_runs,
+        cache_dir=cache_dir,
+        print_number_of_pmappings=print_number_of_pmappings,
+    )
+    mappings = join_pmappings(
+        spec,
+        pmappings,
+        require_all_einsums=einsum_names is not None,
+        _pmapping_row_filter_function=_pmapping_row_filter_function,
+    )
+
+    new_mapping_data = []
+    for i in range(len(mappings.data)):
+        local_spec = deepcopy(spec)
+        local_spec.model.metrics = local_spec.mapper.ffm.info_metrics
+        local_spec.mapping = mappings.data.iloc[i]["Total<SEP>mapping"]()
+        # BUG: Mapping._from_pmappings create mappings that cannot be evaluated!
+        this_mapping = evaluate_mapping(local_spec)
+        new_mapping_data.append(this_mapping.data)
+
+    mappings.data = pd.concat(new_mapping_data)
+
+    return mappings
+
+
+def make_pmappings(
+    spec: Spec,
+    einsum_names: list[EinsumName] | None = None,
+    can_combine_multiple_runs: bool = False,
+    cache_dir: str | None = None,
+    print_number_of_pmappings: bool = True,
+) -> MultiEinsumPmappings:
+    """
+    Creates pmappings for a spec. Pmappings must be joined together using
+    `join_pmappings` to create a full mapping.
+
+    Parameters
+    ----------
+    spec:
+        The Spec to generate pmappings for.
+    einsum_names:
+        The einsum names to generate pmappings for. If None, all einsums will be
+        included.
+    can_combine_multiple_runs: Whether we would like to be able to combine multiple
+        make_pmappings runs. Having this as True allows you to do things like
+        `pmappings = make_pmappings(*args_a) | make_pmappings(*args_b)` but slows
+        down execution.
+    cache_dir:
+        The directory to cache pmappings in. If None, no caching will be done.
+    print_number_of_pmappings:
+        Whether to print the number of pmappings for each einsum.
+
+    Returns:
+        A MultiEinsumPmappings object.
+    """
+
+    kwargs = dict(
+        spec=spec,
+        einsum_names=einsum_names,
+        can_combine_multiple_runs=can_combine_multiple_runs,
+    )
+    assert len(kwargs) == len(inspect.signature(_make_pmappings).parameters)
+
+    if cache_dir is None:
+        result = _make_pmappings(**kwargs)
+    else:
+
+        @joblib.Memory(location=os.path.join(cache_dir), compress=True).cache
+        def _make_pmappings_cached(**kwargs) -> MultiEinsumPmappings:
+            return _make_pmappings(**kwargs)
+
+        result = _make_pmappings_cached(**kwargs)
+
+    if print_number_of_pmappings:
+        print(result.n_pmapping_string())
+
+    return result
+
+
+def join_pmappings(
+    spec: Spec,
+    pmappings: MultiEinsumPmappings,
+    require_all_einsums: bool = True,
+    _pmapping_row_filter_function: Callable[[pd.Series], bool] | None = None,
+) -> Mappings:
+    """
+    Joins pmappings into a full mappings for the entire workload. Pmappings can
+    be generated using `make_pmappings`.
+
+    Parameters
+    ----------
+    spec:
+        The complete specifications for the workload and architecture.
+    pmappings:
+        The pmappings to join.
+    require_all_einsums:
+        If True, all einsums in the workload must have pmappings. If False, only
+        einsums that have pmappings will be included in the final mappings.
+    _pmapping_row_filter_function:
+        A function that takes in a row of the pmapping dataframe and returns
+        True if the row should be included in the final mappings, and False
+        otherwise. If None, all rows will be included.
+
+    Returns
+    -------
+    A Mappings object containing all valid, optimal mappings for the workload.
+    """
+    return clean_compress_and_join_pmappings(
+        spec, pmappings, require_all_einsums, _pmapping_row_filter_function
+    )
+
+
+def _make_pmappings(
+    spec: Spec,
+    einsum_names: list[EinsumName] | None = None,
+    can_combine_multiple_runs: bool = False,
+) -> MultiEinsumPmappings:
+    if einsum_names is None:
+        einsum_names = [e.name for e in spec.workload.einsums]
+
+    einsum2symbol_table = spec.workload.get_constraint_symbol_table(
+        einsum_names, spec.renames
+    )
+    tensor2bits_per_value = spec.workload._get_bits_per_value(einsum2symbol_table)
+    for einsum_name, symbol_table in einsum2symbol_table.items():
+        for tensor in spec.workload.einsums[einsum_name].tensor_names:
+            if tensor not in tensor2bits_per_value:
+                raise ValueError(
+                    f"Tensor {tensor} not found in bits per value for Einsum {einsum_name}. "
+                    f"Bits per value:\n\t"
+                    + "\n\t".join(f"{k}: {v}" for k, v in tensor2bits_per_value.items())
+                )
+            if tensor not in symbol_table:
+                raise ValueError(
+                    f"Tensor {tensor} not found in symbol table for Einsum {einsum_name}. "
+                    f"Symbol table:\n\t"
+                    + "\n\t".join(f"{k}: {v}" for k, v in symbol_table.items())
+                )
+            symbol_table[tensor].bits_per_value = tensor2bits_per_value[tensor]
+
+    pmapping_groups, pmapping_objects, einsum2jobs = pmapper.make_pmappings(
+        parsed_spec,
+        metrics=spec.mapper.ffm.metrics,
+        einsum_names=einsum_names,
+        can_combine_multiple_runs=can_combine_multiple_runs,
+        einsum2symbol_table=einsum2symbol_table,
+        tensor2bits_per_value=tensor2bits_per_value,
+    )
+
+    resource2capacity = {}
+    who_set = {}
+    jobs_flattened = {(e, j) for e, jobs in einsum2jobs.items() for j in jobs}
+    for einsum, job in jobs_flattened:
+        memories = [m for m in job.flattened_arch if isinstance(m, arch.Memory)]
+        for m in memories:
+            resource2capacity.setdefault(m.name, m.attributes.size)
+            who_set.setdefault(m.name, einsum)
+            if resource2capacity[m.name] != m.attributes.size:
+                raise ValueError(
+                    f"Memory {m.name} has different sizes depending on which Einsum "
+                    f"is being mapped. Memory sizes should not depend on which Einsum "
+                    f"is being mapped. Einsum {who_set[m.name]} set the size to "
+                    f"{resource2capacity[m.name]}, but Einsum {einsum} set the size to "
+                    f"{m.attributes.size}."
+                )
+            resource2capacity[m.name] = m.attributes.size
+
+    m = MultiEinsumPmappings(
+        pmapping_groups,
+        pmapping_objects,
+        resource2capacity,
+        einsum2jobs,
+        can_combine_multiple_runs=can_combine_multiple_runs,
+        einsums_with_pmappings_generated=set(
+            einsum_names if einsum_names else spec.workload.einsum_names
+        ),
+    )
+
+    return m

accelforge/mapper/FFM/main.py

@@ -0,0 +1,208 @@
+from copy import deepcopy
+import inspect
+import os
+from typing import Callable
+import joblib
+import logging
+
+from accelforge import arch
+from accelforge import Spec
+from accelforge.mapper.FFM.pmappings import MultiEinsumPmappings
+from accelforge.mapper.FFM.mappings import Mappings
+import accelforge.mapper.FFM._make_pmappings.make_pmappings as pmapper
+from accelforge.frontend.workload import EinsumName
+from accelforge.mapper.FFM._join_pmappings.join_pmappings import (
+    clean_compress_and_join_pmappings,
+)
+from accelforge._accelerated_imports import pd
+
+
+logger = logging.getLogger(__name__)
+
+
+def map_workload_to_arch(
+    spec: Spec,
+    einsum_names: list[EinsumName] | None = None,
+    can_combine_multiple_runs: bool = False,
+    cache_dir: str | None = None,
+    print_number_of_pmappings: bool = True,
+    _pmapping_row_filter_function: Callable[[pd.Series], bool] | None = None,
+) -> Mappings:
+    """
+    Maps a workload to an architecture using the AccelForge FFM mapper.
+
+    Parameters
+    ----------
+    spec:
+        The Spec to map.
+    einsum_names:
+        The einsum names to map. If None, all einsums will be mapped.
+    can_combine_multiple_runs: Whether we would like to be able to combine multiple
+        make_pmappings runs. Having this as True allows you to do things like
+        `pmappings = make_pmappings(*args_a) | make_pmappings(*args_b)` but slows
+        down execution.
+    cache_dir:
+        The directory to cache pmappings in. If None, no caching will be done.
+    print_number_of_pmappings:
+        Whether to print the number of pmappings for each einsum.
+    _pmapping_row_filter_function:
+        A function that takes in a row of the pmapping dataframe and returns
+        True if the row should be included in the final mappings, and False
+        otherwise. If None, all rows will be included.
+    """
+    from accelforge.model.main import evaluate_mapping
+
+    pmappings = make_pmappings(
+        spec,
+        einsum_names=einsum_names,
+        can_combine_multiple_runs=can_combine_multiple_runs,
+        cache_dir=cache_dir,
+        print_number_of_pmappings=print_number_of_pmappings,
+    )
+    mappings = join_pmappings(
+        spec,
+        pmappings,
+        require_all_einsums=einsum_names is not None,
+        _pmapping_row_filter_function=_pmapping_row_filter_function,
+    )
+
+    new_mapping_data = []
+    for i in range(len(mappings.data)):
+        local_spec = deepcopy(spec)
+        local_spec.model.metrics = local_spec.mapper.ffm.info_metrics
+        local_spec.mapping = mappings.data.iloc[i]["Total<SEP>mapping"]()
+        # BUG: Mapping._from_pmappings create mappings that cannot be evaluated!
+        this_mapping = evaluate_mapping(
+            local_spec,
+            flattened_arches=mappings.flattened_arches,
+            parsed_specs=mappings.parsed_specs,
+        )
+        new_mapping_data.append(this_mapping.data)
+
+    mappings.data = pd.concat(new_mapping_data)
+
+    return mappings
+
+
+def make_pmappings(
+    spec: Spec,
+    einsum_names: list[EinsumName] | None = None,
+    can_combine_multiple_runs: bool = False,
+    cache_dir: str | None = None,
+    print_number_of_pmappings: bool = True,
+) -> MultiEinsumPmappings:
+    """
+    Creates pmappings for a spec. Pmappings must be joined together using
+    `join_pmappings` to create a full mapping.
+
+    Parameters
+    ----------
+    spec:
+        The Spec to generate pmappings for.
+    einsum_names:
+        The einsum names to generate pmappings for. If None, all einsums will be
+        included.
+    can_combine_multiple_runs: Whether we would like to be able to combine multiple
+        make_pmappings runs. Having this as True allows you to do things like
+        `pmappings = make_pmappings(*args_a) | make_pmappings(*args_b)` but slows
+        down execution.
+    cache_dir:
+        The directory to cache pmappings in. If None, no caching will be done.
+    print_number_of_pmappings:
+        Whether to print the number of pmappings for each einsum.
+
+    Returns:
+        A MultiEinsumPmappings object.
+    """
+    kwargs = dict(
+        spec=spec,
+        einsum_names=einsum_names,
+        can_combine_multiple_runs=can_combine_multiple_runs,
+    )
+    assert len(kwargs) == len(inspect.signature(_make_pmappings).parameters)
+
+    if cache_dir is None:
+        result = _make_pmappings(**kwargs)
+    else:
+
+        @joblib.Memory(location=os.path.join(cache_dir), compress=True).cache
+        def _make_pmappings_cached(**kwargs) -> MultiEinsumPmappings:
+            return _make_pmappings(**kwargs)
+
+        result = _make_pmappings_cached(**kwargs)
+
+    if print_number_of_pmappings:
+        print(result.n_pmapping_string())
+
+    return result
+
+
+def join_pmappings(
+    spec: Spec,
+    pmappings: MultiEinsumPmappings,
+    require_all_einsums: bool = True,
+    _pmapping_row_filter_function: Callable[[pd.Series], bool] | None = None,
+) -> Mappings:
+    """
+    Joins pmappings into a full mappings for the entire workload. Pmappings can
+    be generated using `make_pmappings`.
+
+    Parameters
+    ----------
+    spec:
+        The complete specifications for the workload and architecture.
+    pmappings:
+        The pmappings to join.
+    require_all_einsums:
+        If True, all einsums in the workload must have pmappings. If False, only
+        einsums that have pmappings will be included in the final mappings.
+    _pmapping_row_filter_function:
+        A function that takes in a row of the pmapping dataframe and returns
+        True if the row should be included in the final mappings, and False
+        otherwise. If None, all rows will be included.
+
+    Returns
+    -------
+    A Mappings object containing all valid, optimal mappings for the workload.
+    """
+    return clean_compress_and_join_pmappings(
+        spec, pmappings, require_all_einsums, _pmapping_row_filter_function
+    )
+
+
+def _make_pmappings(
+    spec: Spec,
+    einsum_names: list[EinsumName] | None = None,
+    can_combine_multiple_runs: bool = False,
+) -> MultiEinsumPmappings:
+    if einsum_names is None:
+        einsum_names = [e.name for e in spec.workload.einsums]
+
+    pmapping_groups, pmapping_objects, einsum2jobs = pmapper.make_pmappings(
+        spec,
+        metrics=spec.mapper.ffm.metrics,
+        einsum_names=einsum_names,
+        can_combine_multiple_runs=can_combine_multiple_runs,
+    )
+
+    flattened_arches = {}
+    parsed_specs = {}
+    for einsum_name, jobs in einsum2jobs.items():
+        for job in jobs:
+            compute_name = job.flattened_arch[-1].name
+            flattened_arches[(einsum_name, compute_name)] = job.flattened_arch
+            parsed_specs[einsum_name] = job.spec
+
+    m = MultiEinsumPmappings(
+        pmapping_groups,
+        pmapping_objects,
+        einsum2jobs,
+        can_combine_multiple_runs=can_combine_multiple_runs,
+        einsums_with_pmappings_generated=set(
+            einsum_names if einsum_names else spec.workload.einsum_names
+        ),
+        flattened_arches=flattened_arches,
+        parsed_specs=parsed_specs,
+    )
+
+    return m