Function emit_program

exir/emit/_emit_program.py:118–216

Given a exported program, it returns the program in the format of the Python version of the flatbuffer Program schema. Args: methods: Either the exported program (Exported_Program) that we want to emit into the flatbuffer, or a dictionary of method names to

(
    methods: Union[ExportedProgram, Dict[str, ExportedProgram]],
    emit_stacktrace: bool = False,
    prim_getters: Optional[Dict[str, Any]] = None,
    emit_mutable_buffer_names: bool = False,
)

def emit_program(
    methods: Union[ExportedProgram, Dict[str, ExportedProgram]],
    emit_stacktrace: bool = False,
    prim_getters: Optional[Dict[str, Any]] = None,
    emit_mutable_buffer_names: bool = False,
) -> EmitterOutput:
    """
    Given a exported program, it returns the program in the format
    of the Python version of the flatbuffer Program schema.

    Args:
        methods: Either the exported program (Exported_Program) that we want to
            emit into the flatbuffer, or a dictionary of method names to
            ExportedPrograms.
        emit_stacktrace: Flag to enable emission of a stacktrace for each
           instruction for debugging purposes

    Return:
        The program in a Python class which mimics the flatbuffer schema
    """

    if isinstance(methods, ExportedProgram):
        methods = {"forward": methods}

    # validation
    bad_methods = []
    for name, exported_program in methods.items():
        if not isinstance(exported_program, ExportedProgram):
            bad_methods.append(name)
    if len(bad_methods) != 0:
        raise ExportError(
            ExportErrorType.INVALID_INPUT_TYPE,
            f"Did not receive ExportedProgram for the following methods {str(bad_methods)}",
        )

    plans = []
    debug_handle_map = {}
    method_to_delegate_debug_id_map = {}
    instruction_id_to_num_outs_map = {}
    program_state = _ProgramState()

    # emit each entry point in order according to name.
    for name, exported_program in sorted(methods.items()):
        # create empty state
        emitter_state = _EmitterState(
            values=[],
            operators=[],
            delegates=[],
            operator_cache={},
            emit_stacktrace=emit_stacktrace,
            emit_mutable_buffer_names=emit_mutable_buffer_names,
        )

        gm = _remove_non_user_outputs(exported_program)

        emitter = _TopLevelEmitter(
            name, exported_program, gm, program_state, emitter_state
        )