MCPcopy Create free account
hub / github.com/pytorch/executorch / _restore_segments

Function _restore_segments

exir/_serialize/_program.py:679–744  ·  view source on GitHub ↗

Moves segments from `segment_data` into `program`. This should recreate the original Program that the segments were extracted from. Args: program: The Program to restore. `program.segments` must describe the segment locations. segment_data: The data containi

(program: Program, segment_data: bytes)

Source from the content-addressed store, hash-verified

677
678
679def _restore_segments(program: Program, segment_data: bytes) -> PTEFile:
680 """Moves segments from `segment_data` into `program`.
681
682 This should recreate the original Program that the segments were extracted
683 from.
684
685 Args:
686 program: The Program to restore. `program.segments` must describe the
687 segment locations.
688 segment_data: The data containing the segments. Assumes that this data
689 begins at `segment_base_offset` from the extended header: i.e.,
690 the preceding data has been stripped off so that the first segment
691 begins at offset zero.
692 Returns:
693 PTEFile, containing the Program with delegate and constant segments restored, mutable data segment, and named data segment.
694 """
695 # Extract the list of segment data blobs, which parallel program.segments.
696 segments: List[bytes] = []
697 for i, segment in enumerate(program.segments):
698 if segment.offset + segment.size > len(segment_data):
699 raise ValueError(
700 f"Segment {i} {segment} overflows data length {len(segment_data)}"
701 )
702 segments.append(segment_data[segment.offset : segment.offset + segment.size])
703
704 # Restore delegate segments that weren't inlined previously.
705 program = _restore_delegates(program, segments)
706
707 # Replace constants from constant_segment into constant_buffer.
708 if program.constant_segment and len(program.constant_segment.offsets) > 0:
709 if program.constant_segment.segment_index >= len(segments):
710 raise ValueError(
711 f"Constant segment index {program.constant_segment.segment_index} >= num segments {len(segments)}"
712 )
713 program.constant_buffer = _restore_constant_segment(
714 program.constant_segment, segments[program.constant_segment.segment_index]
715 )
716 program.constant_segment.segment_index = 0
717 program.constant_segment.offsets = []
718
719 # Extract mutable segments.
720 mutable_data = None
721 if program.mutable_data_segments and len(program.mutable_data_segments) > 0:
722 if len(program.mutable_data_segments) > 1:
723 raise ValueError("Can't handle more than 1 mutable data segment.")
724 segment_index = program.mutable_data_segments[0].segment_index
725 if segment_index >= len(segments):
726 raise ValueError(
727 f"Mutable data segment index {segment_index} >= num segments {len(segments)}"
728 )
729 mutable_data = _restore_constant_segment(
730 program.mutable_data_segments[0],
731 segments[segment_index],
732 )
733 program.mutable_data_segments = None
734
735 # Extract named data.
736 named_data = None

Callers 1

deserialize_pte_binaryFunction · 0.85

Calls 5

_restore_delegatesFunction · 0.85
_restore_named_dataFunction · 0.85
PTEFileClass · 0.85
appendMethod · 0.45

Tested by

no test coverage detected