Function _extract_delegate_segments

exir/_serialize/_program.py:279–346

Extracts the delegate segments inlined in the program into a list of buffers. The program is modified in-place to remove the delegate data. Args: program: The program to extract segments from. Modified in-place. segments: A list of buffers to append extracted segments to

(
    program: Program,
    segments: List[AlignedData],
)

def _extract_delegate_segments(
    program: Program,
    segments: List[AlignedData],
) -> None:
    """Extracts the delegate segments inlined in the program into a list of buffers.
        The program is modified in-place to remove the delegate data.

    Args:
        program: The program to extract segments from. Modified in-place.
        segments: A list of buffers to append extracted segments to. Modified in-place.
    """
    remaining_inline: List[BackendDelegateInlineData] = []
    inline_indices_seen: set[int] = set()
    segment_index_map: dict[bytes, int] = {}
    for plan in program.execution_plan:
        for delegate in plan.delegates:
            if delegate.processed.location != DataLocation.INLINE:
                raise ValueError(
                    "Program must only contain inline delegate data, "
                    + f"saw {repr(delegate)}"
                )
            # TODO(T144120904): Don't extract small blobs into segments;
            # have a cutoff. Or callers could provide a callback that
            # returns true/false for a given BackendDelegate, letting them
            # use their own logic.
            try:
                inline: BackendDelegateInlineData = program.backend_delegate_data[
                    delegate.processed.index
                ]
            except IndexError:
                raise ValueError(
                    f"Delegate processed index {delegate.processed.index} "
                    + ">= len(Program.backend_delegate_data) "
                    + f"{len(program.backend_delegate_data)} "
                    + f"in {repr(delegate)}"
                )
            inline_indices_seen.add(delegate.processed.index)
            if inline.data:
                # Move the delegate data out of the program.
                segment_index = segment_index_map.get(inline.data)
                if segment_index is None:
                    segment_index = len(segments)
                    segments.append(AlignedData(Cord(inline.data)))
                    segment_index_map[inline.data] = segment_index
                delegate.processed = BackendDelegateDataReference(
                    location=DataLocation.SEGMENT,
                    index=segment_index,
                )
            else:
                # Not moving into a segment. Keep it inline, but update the
                # index.
                new_index = len(remaining_inline)
                remaining_inline.append(inline)
                delegate.processed.index = new_index

    # Make sure we visited all entries in backend_delegate_data, so that it's
    # safe to overwrite it.
    remaining_indices: set[int] = set(