Function _generate_edge_program

exir/program/_program.py:860–927

Args: config: The configuration for the edge program. program: The exported program to be converted to an edge program. core_aten_ops_exception_list: A list of aten ops that are missing decompositions to core aten. preserve_ops: A list of aten ops that should not

(
    config: EdgeCompileConfig,
    program: ExportedProgram,
    core_aten_ops_exception_list: Optional[List[torch._ops.OpOverload]] = None,
    preserve_ops: Optional[List[torch._ops.OpOverload]] = None,
)

def _generate_edge_program(
    config: EdgeCompileConfig,
    program: ExportedProgram,
    core_aten_ops_exception_list: Optional[List[torch._ops.OpOverload]] = None,
    preserve_ops: Optional[List[torch._ops.OpOverload]] = None,
) -> ExportedProgram:
    """
    Args:
        config: The configuration for the edge program.
        program: The exported program to be converted to an edge program.
        core_aten_ops_exception_list: A list of aten ops that are missing decompositions to core aten.
        preserve_ops: A list of aten ops that should not be decomposed.
    Returns:
        An ExportedProgram in edge dialect.
    """
    # Remove unused parameters
    program = remove_unused_parameters_pass(program)

    pre_op_replace_passes, post_op_replace_passes = _get_aten_to_edge_passes(config)

    passes = [
        # Remove invalid assert ops, such as _assert_tensor_metadata
        RemoveNonCoreAtenOpGraphAssertsPass(),
        # TODO move inside aten_to_edge passes after all users are migrated off v1 capture
        ReplaceViewOpsWithViewCopyOpsPass(),
    ]
    passes.extend(pre_op_replace_passes)
    if config._use_edge_ops:
        passes.append(OpReplacePass())
        if not config._skip_dim_order:
            passes.append(MemoryFormatOpsPass())

    gm = program.graph_module
    for p in passes:
        gm_res = p(gm)
        assert gm_res is not None
        gm = gm_res.graph_module

    edge_program = ExportedProgram(
        root=gm,
        graph=gm.graph,
        graph_signature=_get_updated_graph_signature(program.graph_signature, gm),
        state_dict=program.state_dict,
        range_constraints=program.range_constraints,
        module_call_graph=program.module_call_graph,
        example_inputs=program.example_inputs,
        constants=program.constants,
        verifiers=[
            EXIREdgeDialectVerifier(
                edge_compile_config=config,
                class_only=True,
                core_aten_ops_exception_list=core_aten_ops_exception_list,
                preserve_ops=preserve_ops,
            )
        ],
    )

    # Lift the tensor constants created in ScalarToTensorPass