Function pick_shared_obj

exir/memory_planning.py:796–889

r""" Pick the available shared object to which to assign this spec, or create a new one Algorithm details Previous: Look at every spec in chronological order. Find if previously allocated object allows it to fit in. If not, allocate a new object. New: - Sort all the specs

(
    shared_objects: List[SharedObject],
    spec: TensorSpec,
    allow_overlapping_allocations: bool = True,
)

def pick_shared_obj(
    shared_objects: List[SharedObject],
    spec: TensorSpec,
    allow_overlapping_allocations: bool = True,
) -> SharedObject:
    r"""
    Pick the available shared object to which to assign this spec,
    or create a new one
    Algorithm details
    Previous: Look at every spec in chronological order. Find if previously allocated object
    allows it to fit in. If not, allocate a new object.
    New:
    - Sort all the specs by allocation size
    - Process the specs in order
    - If the spec's size in smaller than previously allocated buckets:
        - Conditions under which previously allocated bucket can be used:
          - Lifetime of the spec does not overlap with lifetime of the bucket.
              - In this case allocate spec to that bucket and expand its lifetime.
              - Spec is allocated at offset = 0 in this bucket.
              - Add this spec to allocated object's list of specs.
          - Lifetime of the spec overlaps with lifetime of the bucket,
            partially or fully (e.g. spec's lifetime subset of bucket's lifetime)
              - If none of the specs in the bucket overlaps with spec's lifetime.
                - Allocate spec to the bucket at offset = 0.
                - Add this spec to the bucket's list of specs.
                - Expand bucket's lifetime accounting for added spec's lifetime.
              - If one or more specs in the bucket overlaps with spec's lifetime.
                - Collect offsets (at which the given overlapping spec is allocated in the bucket).
                  of all the overlapping specs, and find the max offset.
                - Allocate spec to the bucket at offset = max_offset + max_offset_spec_size.
                - Add this spec to the bucket's list of specs.
                - Expand bucket's lifetime accounting for added spec's lifetime.
        - If none of these conditions are met, allocate a new bucket.
            - Add spec to this bucket.
            - Update bucket's lifetime to that of the spec.
    - If the spec's size is larger than previously allocated buckets, allocate a new bucket.
        - Size and lifetime of this bucket is that of the spec

    Proof of correctness:
    - If allocating a new bucket, it is correct.
    - If allocating spec to an existing bucket, whose lifetime does not overlap with any
      of the previously allocated specs' lifetime, then the allocation is correct.
    Proof of correctness by induction when adding spec to an existing bucket:
    - If all previous allocations in the given bucket are correct:
        - Then the new one being added must be correct because when the requested allocation
          overlaps with one or more previous allocations, we find the largest offset among
          all the overlapping allocations, and allocate the new spec at that offset. Hence,
          the allocation at such an offset, will not overlap with any previous allocations.
    Base case: A newly added allocation within a bucket with single allocation is correct:
    because a) it must fit and b) its lifetime must not overlap with object's lifetime.
    This holds true because of the following invariants:
    - Once a bucket is created, it is never resized.
    - All the allocations within a bucket follow this:
      - Span, defined by allocation's offset + size, of two allocations can only overlap,
        if their timelines do not overlap.
    """
    picked = None
    for sobj in shared_objects: