Function serialize_pte_binary

exir/_serialize/_program.py:422–592

Returns the runtime binary representation of the given Program. Args: pte_file: PTEFile class containing the program and segments. extract_delegate_segments: Whether to move delegate data blobs from the Program into separate segments, rather than encoding those blobs

(
    pte_file: PTEFile,
    *,
    extract_delegate_segments: bool = False,
    segment_alignment: int = 128,
    constant_tensor_alignment: Optional[int] = None,
    delegate_alignment: Optional[int] = None,
)

def serialize_pte_binary(
    pte_file: PTEFile,
    *,
    extract_delegate_segments: bool = False,
    segment_alignment: int = 128,
    constant_tensor_alignment: Optional[int] = None,
    delegate_alignment: Optional[int] = None,
) -> Cord:
    """Returns the runtime binary representation of the given Program.

    Args:
        pte_file: PTEFile class containing the program and segments.
        extract_delegate_segments: Whether to move delegate data blobs from the
            Program into separate segments, rather than encoding those blobs
            in the flatbuffer data. When true, will also:
            - Add an extended header to the output, containing the program size
              and the starting segment offset.
            - Update the Program.segments field with the offsets and lengths
              of each segment.
        segment_alignment: Alignment in bytes. The starting offset of each
            segment will be aligned to this value in the output data.
        constant_tensor_alignment: The minimum alignment of tensor
            buffers in the program. Must be a power of 2. Defaults to ALIGNMENT.
        delegate_alignment: If provided, the minimum alignment of delegate data
            in the program. Must be a power of 2. If not provided, uses the
            value in the schema file.
    Returns:
        The serialized form of the Program, ready for execution by the runtime.
    """
    # Default tensor alignment.
    if constant_tensor_alignment is None:
        constant_tensor_alignment = ALIGNMENT

    # Don't modify the original program.
    # TODO(T144120904): Could avoid yet more huge copies with a more shallow
    # copy, reusing the actual data blobs.
    program = copy.deepcopy(pte_file.program)

    # Store extracted segment data, with any buffer-specific alignment.
    # This may be constant data, delegate data or named data.
    segments: List[AlignedData] = []

    constant_segment_data, constant_segment_offsets = _extract_constant_segment(
        program.constant_buffer, tensor_alignment=constant_tensor_alignment
    )

    # If there are no constants, len(constant_segment_data) = 0. However, there may
    # be non-constants, in which case len(constant_segment_offsets) = 1, containing
    # the placeholder value 0. Ensure the placeholder value is put into
    # program.constant_segment.offsets.
    if len(constant_segment_offsets) > 0:
        # Update program.constant_segment with constant subsegment offset information.
        program.constant_segment = SubsegmentOffsets(
            segment_index=len(segments), offsets=constant_segment_offsets
        )
        # Clear the constant buffer, as constant data will be stored in segments.
        program.constant_buffer = []
        # Add to the aggregate segments cord.