Class _ProgramState

exir/emit/_emitter.py:103–134

State shared between all methods of a program and the graph module it represents. Initialized once within emit_program and then shared across each entry point as they are emitted.

@dataclass
class _ProgramState:
    """State shared between all methods of a program and the graph module it represents.

    Initialized once within emit_program and then shared across each entry point as they are
    emitted.
    """

    # Parallel list of specs and the buffers that backed them, have to add + 1 to any index in here
    # as index 0 in the constant_buffer is reserved.
    allocated_specs: List[TensorSpec] = field(default_factory=list)
    # Weights in any arbitrary graph_module only need to compare against weights from previously
    # emitted graph modules, not any weights emitted from itself. This should speed up the lookup,
    # from O(N) to O(1)
    cached_spec_hash_values: Dict[str, int] = field(default_factory=dict)
    cached_spec_mutable_hash_values: Dict[str, int] = field(default_factory=dict)
    # The 0 index is reserved to be pointed to by non-constant tensors, so add an empty placeholder.
    constant_buffer: List[Buffer] = field(default_factory=lambda: [Buffer(storage=b"")])
    # The 0 index is reserved to be pointed to by non-constant tensors, so add an empty placeholder.
    mutable_buffer: List[Buffer] = field(default_factory=lambda: [Buffer(storage=b"")])
    # Delegate data stored directly in the flatbuffer. Pointed to by BackendDelegateDataReference,
    # and should be copied to Program.backend_delegate_data.
    backend_delegate_data: List[BackendDelegateInlineData] = field(default_factory=list)
    # Delegate cache that is used across all entry points. Key is the hash of the delegated payload.
    backend_delegate_data_cache: Dict[str, int] = field(default_factory=dict)

    # Constants are optionally stored in external files.
    # Aggregate unique external constants into one buffer.
    external_constant_buffer: List[bytes] = field(default_factory=list)
    external_constant_hash: Dict[str, int] = field(default_factory=dict)
    # Each constant_tag groups a set of constants together.
    # {constant_tag: {fqn: index into external_constant_buffer}}
    external_constant_map: Dict[str, Dict[str, int]] = field(default_factory=dict)


@dataclass