Function to_edge

exir/program/_program.py:1451–1530

:func:`to_edge` constructs an EdgeProgramManager from a set of exported programs in ATen dialect. Upon construction those programs are transformed into edge dialect. Args: programs: Can be a single ExportedProgram or a dictionary mapping function names to their corresponding Ex

(
    programs: Union[ExportedProgram, Dict[str, ExportedProgram]],
    constant_methods: Optional[Dict[str, Any]] = None,
    compile_config: Optional[EdgeCompileConfig] = None,
    generate_etrecord: bool = False,
)

@et_logger("to_edge")
def to_edge(
    programs: Union[ExportedProgram, Dict[str, ExportedProgram]],
    constant_methods: Optional[Dict[str, Any]] = None,
    compile_config: Optional[EdgeCompileConfig] = None,
    generate_etrecord: bool = False,
) -> "EdgeProgramManager":
    """
    :func:`to_edge` constructs an EdgeProgramManager from a set of exported programs in
    ATen dialect. Upon construction those programs are transformed into edge dialect.

    Args:
        programs: Can be a single ExportedProgram or a dictionary mapping function names to their corresponding ExportedPrograms. If only a single ExportedProgram is provided it will be assigned the name "forward".

        constant_methods: An optional dictionary of method name to the constant value returned by that method in eager mode. Often used to store config information on Edge models.

        compile_config: An optional argument used to provide greater control over the transformation to edge dialect process.

        generate_etrecord: An optional argument used to generate an etrecord for debugging purposes. Default is False.

    Returns:
        EdgeProgramManager
    """
    assert not isinstance(constant_methods, EdgeCompileConfig)
    config = compile_config or EdgeCompileConfig()
    if not isinstance(programs, dict):
        aten_programs = {"forward": programs}
    else:
        aten_programs = programs

    edge_programs: Dict[str, ExportedProgram] = {}

    for name, program in aten_programs.items():
        # Decompose to Core ATen
        table = _default_decomposition_table()
        preserve_ops = []
        if compile_config:
            preserve_ops = compile_config.preserve_ops
            for op in compile_config.preserve_ops:
                table.pop(op, None)
        program = program.run_decompositions(table)

        if config._check_ir_validity:
            # Remove invalid assert ops, such as _assert_tensor_metadata.
            # This pass is run in _generate_edge_program; it is required here to
            # ensure the graph is in ATen dialect before verification.
            gm = program.graph_module
            gm_res = RemoveNonCoreAtenOpGraphAssertsPass()(gm)
            assert gm_res is not None
            gm = gm_res.graph_module
            try:
                EXIRATenDialectVerifier(
                    edge_compile_config=config,
                    class_only=False,
                )(gm)
            except ExportError as e:
                logging.info(f"Input program {name} is not in ATen dialect.")
                raise e