accelforge 0.0.1__py3-none-any.whl

accelforge/mapper/FFM/_make_pmappings/make_pmappings.py

@@ -0,0 +1,411 @@
+import logging
+from math import prod
+
+from typing import Callable, Optional
+import uuid
+import copy
+
+from joblib import delayed
+from tqdm import tqdm
+
+
+from accelforge.frontend import arch
+from accelforge.frontend.spec import Spec
+from accelforge.frontend.mapping import Loop, Mapping, TensorHolder
+from accelforge.frontend._workload_isl._isl import (
+    get_rank_variable_bounds,
+    get_tensor_size,
+    get_operation_space_size,
+)
+from accelforge.frontend.workload import EinsumName, SymbolTable, TensorName
+
+from accelforge.mapper.FFM._make_pmappings.make_pmapping_templates import (
+    make_pmapping_templates,
+)
+from accelforge.frontend.mapper.metrics import Metrics
+from accelforge.mapper.FFM._make_pmappings.make_pmappings_from_templates import (
+    make_pmappings_from_templates,
+)
+from accelforge.mapper.FFM._join_pmappings.compatibility import Compatibility
+from accelforge.mapper.FFM._join_pmappings.pmapping_group import PmappingGroup
+from accelforge.util.parallel import (
+    parallel,
+    _memmap_read,
+    get_n_parallel_jobs,
+    is_using_parallel_processing,
+)
+from accelforge.mapper.FFM._make_pmappings.pmapper_job import (
+    Job,
+    SameCompatibilityJobs,
+)
+
+
+def get_rank_variable_bounds_for_all_einsums(spec: Spec):
+    rank_variable_bounds = {
+        einsum_name: get_rank_variable_bounds(spec.workload, einsum_name)
+        for einsum_name in spec.workload.einsum_names
+    }
+    result = {}
+    for e1, rv1 in rank_variable_bounds.items():
+        result.update(rv1)
+        for e2, rv2 in rank_variable_bounds.items():
+            for r in set(rv1.keys()) & set(rv2.keys()):
+                if rv1[r] != rv2[r]:
+                    raise ValueError(
+                        f"Rank variable {r} has different bounds for "
+                        f"einsum {e1} and {e2}: {rv1[r]} and {rv2[r]}"
+                    )
+    return result
+
+
+def get_num_computes(spec: Spec, einsum_name: EinsumName | None = None) -> int:
+    einsums = spec.workload.einsums
+    einsums = [einsum_name] if einsum_name is not None else spec.workload.einsum_names
+    return sum(get_operation_space_size(spec.workload, e) for e in einsums)
+
+
+def get_per_tensor_size(spec: Spec) -> dict[TensorName, int]:
+    return {
+        tensor: get_tensor_size(spec.workload, tensor)
+        for tensor in spec.workload.tensor_names
+    }
+
+
+def get_jobs(
+    spec: Spec,
+    metrics: Metrics,
+    einsum_names: list[EinsumName],
+    fail_if_no_pmappings_for_einsum: bool,
+) -> dict[EinsumName, dict[Compatibility, SameCompatibilityJobs]]:
+
+    spec = spec
+
+    einsum2jobs = {}
+    fusable_tensors = spec.workload.tensor_names_used_in_multiple_einsums
+    rank_variable_bounds = get_rank_variable_bounds_for_all_einsums(spec)
+
+    einsum2spec: dict[EinsumName, Spec] = {}
+    s = f"Getting energy, latency, and leak power for components running "
+    pbar = tqdm(einsum_names, desc=s)
+    for einsum_name in pbar:
+        pbar.set_description(s + einsum_name)
+        einsum2spec[einsum_name] = spec._spec_parse_expressions(
+            einsum_name=einsum_name,
+            _parse_arch=True,
+            _parse_non_arch=False,
+        ).calculate_component_area_energy_latency_leak(
+            einsum_name=einsum_name,
+            area=False,
+        )
+        einsum2spec[einsum_name] = _memmap_read(einsum2spec[einsum_name])
+
+    def make_jobs_for_einsum(einsum_name: EinsumName, spec: Spec):
+        jobs = {}
+        workload_einsum = spec.workload.einsums[einsum_name]
+        for flattened_arch in spec._get_flattened_architecture():
+            # Create jobs for each Einsum
+            job = Job(
+                spec=spec,
+                einsum_name=einsum_name,
+                metrics=metrics,
+                rank_variable_bounds=rank_variable_bounds,
+                flattened_arch=_memmap_read(flattened_arch),
+                job_id=uuid.uuid4(),
+                fusable_tensors=fusable_tensors & workload_einsum.tensor_names,
+            )
+            for j in make_pmapping_templates(job):
+                jobs.setdefault(j.compatibility, SameCompatibilityJobs()).append(j)
+
+        return einsum_name, jobs
+
+    for einsum_name, jobs in parallel(
+        [
+            delayed(make_jobs_for_einsum)(einsum_name, spec)
+            for einsum_name, spec in einsum2spec.items()
+        ],
+        pbar="Generating jobs",
+        return_as="generator",
+    ):
+        einsum2jobs.setdefault(einsum_name, {})
+        for compatibility, job_list in jobs.items():
+            einsum2jobs[einsum_name].setdefault(
+                compatibility, SameCompatibilityJobs()
+            ).extend(job_list)
+
+    if fail_if_no_pmappings_for_einsum:
+        for einsum_name, jobs in einsum2jobs.items():
+            if len(jobs) == 0:
+                raise ValueError(
+                    f"No pmappings for {einsum_name}. Was the mapspace overconstrained?"
+                )
+
+    total_jobs = sum(len(jobs) for jobs in einsum2jobs.values())
+    n_procs = get_n_parallel_jobs()
+    memory_limit = min(
+        spec.mapper.ffm.memory_limit, spec.mapper.ffm.memory_limit_per_process / n_procs
+    )
+    time_limit = min(
+        spec.mapper.ffm.time_limit * n_procs / max(total_jobs, 1),
+        spec.mapper.ffm.time_limit_per_pmapping_template,
+    )
+    for einsum_name, compatibility_jobs in einsum2jobs.items():
+        total_jobs = sum(len(j) for j in compatibility_jobs.values())
+        logging.warning(f"Einsum {einsum_name} has {total_jobs} pmapping templates:")
+        for job_list in compatibility_jobs.values():
+            for job in job_list:
+                logging.warning(f"\t{job.mapping.compact_str()}")
+                job.memory_limit = memory_limit
+                job.time_limit = time_limit
+
+    return einsum2jobs
+
+
+def get_memories_to_track(
+    spec: Spec,
+    einsum2jobs: dict[EinsumName, list[Job]],
+    metrics: Metrics,
+    can_combine_multiple_runs: bool,
+) -> tuple[list[str], list[str]]:
+
+    memories_track_all = set()
+    for einsum, jobs in einsum2jobs.items():
+        for job in jobs:
+            memories_track_all.update(
+                m.name for m in job.flattened_arch if isinstance(m, arch.Memory)
+            )
+
+    memories_track_pmappings_only = []
+    ignored_resources = set()
+
+    # If we're combining the pmappings from multiple runs, we can't conclude anything
+    # about the metrics to track
+    if can_combine_multiple_runs:
+        ignored_resources = memories_track_all
+        return (
+            memories_track_all,
+            memories_track_pmappings_only,
+            ignored_resources,
+        )
+
+    if metrics.RESOURCE_USAGE in metrics:
+        ignored_resources = memories_track_all
+        return (
+            memories_track_all,
+            memories_track_pmappings_only,
+            ignored_resources,
+        )
+
+    tensor_sizes = {}
+    for tensor, size in get_per_tensor_size(spec).items():
+        scale = 1
+        for einsum in spec.workload.einsums_with_tensor(tensor):
+            if einsum.tensor_accesses[tensor].persistent:
+                scale = max(scale, spec.workload.n_instances * einsum.n_instances)
+        tensor_sizes[tensor] = size * scale
+
+    # If the memory is big enough to hold all the tensors then we don't need to consider
+    # it
+    for memory in list(memories_track_all):
+        usage = 0
+        for einsum in einsum2jobs.keys():
+            job = einsum2jobs[einsum][0]
+            try:
+                mem: arch.Memory = job.spec.arch.find(memory)
+            except ValueError:
+                continue
+            for tensor in spec.workload.einsums[einsum].tensor_names:
+                if mem.size == 0:
+                    usage = 2  # FAIL
+                else:
+                    scale = mem.bits_per_value_scale[tensor] / mem.size
+                    usage += tensor_sizes[tensor] * scale
+
+        if usage <= 1:
+            ignored_resources.add(memory)
+            print(
+                f"Not tracking memory {memory}. It is big enough to hold "
+                f"every workload tensor that may be stored in it. Max possible "
+                f"usage: {usage * 100:.2f}%"
+            )
+            memories_track_all.remove(memory)
+
+    # If the memory is below every backing tensor holder node, then we need it for the
+    # pmapping exploration but can drop it immediately
+    for m in list(memories_track_all):
+        must_track = False
+        for job in jobs:
+            seen = False
+            for node in job.mapping.nodes:
+                if isinstance(node, TensorHolder) and node.component == m:
+                    seen = True
+                    if node.persistent:
+                        ignored_resources.add(m)
+                    if node._backing:
+                        must_track = True
+                if isinstance(node, Loop) and node._fused and seen:
+                    must_track = True
+
+        if not must_track:
+            memories_track_all.remove(m)
+            memories_track_pmappings_only.append(m)
+            print(
+                f"Not tracking memory {m} across joining stages. It is never "
+                f"reserved across fused loop iterations."
+            )
+
+    return memories_track_all, memories_track_pmappings_only, ignored_resources
+
+
+def make_pmappings(
+    spec: Spec,
+    can_combine_multiple_runs: bool,
+    metrics: Metrics = Metrics.ENERGY | Metrics.LATENCY,
+    einsum_names: Optional[list[EinsumName]] = None,
+    fail_if_no_pmappings_for_einsum: bool | None = None,
+) -> tuple[
+    dict[EinsumName, list[PmappingGroup]],
+    dict[EinsumName, dict[uuid.UUID, Mapping]],
+    dict[EinsumName, list[Job]],
+]:
+    """
+    Explores pmapspace of `einsum_names` (default: all Einsums in workload).
+    """
+    spec = copy.deepcopy(spec)
+
+    if einsum_names is None:
+        einsum_names = spec.workload.einsum_names
+
+    if fail_if_no_pmappings_for_einsum is None:
+        fail_if_no_pmappings_for_einsum = not can_combine_multiple_runs
+
+    spec = spec._spec_parse_expressions(
+        _parse_arch=False,
+        _parse_non_arch=True,
+    )
+
+    einsum2jobs = {}
+    new_einsum2jobs = get_jobs(
+        spec,
+        metrics,
+        einsum_names,
+        fail_if_no_pmappings_for_einsum,
+    )
+    _fill_jobs_with_memories_to_track(
+        new_einsum2jobs, spec, metrics, can_combine_multiple_runs
+    )
+    for einsum_name, jobs in new_einsum2jobs.items():
+        einsum2jobs.setdefault(einsum_name, {})
+        for compatibility, job_list in jobs.items():
+            einsum2jobs[einsum_name].setdefault(
+                compatibility, SameCompatibilityJobs()
+            ).extend(job_list)
+
+    calls = _allocate_jobs(einsum2jobs)
+
+    # Sort the calls by the length of the longest mapping in each job. We get long
+    # poles with the long mappings, so we want to get them done early so we don't
+    # have one or two procs slowing us down at the end.
+    def get_longest_mapping_length(call):
+        j: SameCompatibilityJobs = call[2]["jobs_with_similar_compatibilities"]
+        return max([len(j2.mapping.nodes) for j2 in j])
+
+    calls = sorted(calls, key=get_longest_mapping_length, reverse=True)
+    # # Randomly permute the calls
+    # import random
+    # random.shuffle(calls)
+
+    pmapping_objects = {}
+    pmapping_groups = {einsum_name: [] for einsum_name in spec.workload.einsum_names}
+    return_jobs = {}
+    for (
+        einsum_name,
+        new_pmapping_groups,
+        pmappings,
+        jobs_with_similar_compatibilities,
+    ) in parallel(
+        calls,
+        pbar=f"Generating pmappings",
+        return_as="generator_unordered",
+    ):
+        pmapping_groups[einsum_name].extend(new_pmapping_groups)
+        pmapping_objects.setdefault(einsum_name, {}).update(pmappings)
+        return_jobs.setdefault(einsum_name, []).extend(
+            jobs_with_similar_compatibilities
+        )
+
+    for einsum_name in list(pmapping_groups.keys()):
+        pmapping_groups[einsum_name] = PmappingGroup.combine_combineable(
+            pmapping_groups[einsum_name],
+            "All",
+            pbar_postfix=f" for {einsum_name}",
+        )
+
+    return pmapping_groups, pmapping_objects, return_jobs
+
+
+def _raise_error_if_no_pmappings(einsum2jobs):
+    for einsum_name, jobs in einsum2jobs.items():
+        if len(jobs) == 0:
+            raise ValueError(
+                f"No pmappings for {einsum_name}. " f"Was the mapspace overconstrained?"
+            )
+
+
+def _allocate_jobs(einsum2jobs):
+    calls = []
+    for einsum_name, jobs in einsum2jobs.items():
+        calls.extend(
+            delayed(make_pmappings_from_templates)(
+                jobs_with_similar_compatibilities=job_list,
+            )
+            for job_list in jobs.values()
+        )
+
+    split = False
+    if (
+        not split
+        and is_using_parallel_processing()
+        and len(calls) < get_n_parallel_jobs() * 4
+    ):
+        logging.warning(
+            f"Insufficient jobs available to utilize available threads. "
+            f"Splitting jobs into smaller chunks."
+        )
+        split = True
+
+    if split:
+        calls = []
+        for einsum_name, jobs in einsum2jobs.items():
+            for job_list in jobs.values():
+                calls.extend(
+                    delayed(make_pmappings_from_templates)(
+                        jobs_with_similar_compatibilities=job,
+                    )
+                    for job in job_list.split()
+                )
+    return calls
+
+
+def _fill_jobs_with_memories_to_track(
+    einsum2jobs: dict[EinsumName, dict[Compatibility, SameCompatibilityJobs]],
+    spec,
+    metrics,
+    can_combine_multiple_runs,
+):
+    einsum2jobs_flattened = {
+        e: [j for jobs in v.values() for j in jobs] for e, v in einsum2jobs.items()
+    }
+
+    memories_track_all, memories_track_pmappings_only, ignored_resources = (
+        get_memories_to_track(
+            spec,
+            einsum2jobs_flattened,
+            metrics,
+            can_combine_multiple_runs,
+        )
+    )
+    for jobs in einsum2jobs_flattened.values():
+        for j in jobs:
+            j.memories_track_all = memories_track_all
+            j.memories_track_pmappings_only = memories_track_pmappings_only
+            j.ignored_resources = ignored_resources

accelforge/mapper/FFM/_make_pmappings/make_pmappings_from_templates/__init__.py

@@ -0,0 +1 @@
+from .make_pmappings_from_templates import make_pmappings_from_templates