(
edge_program: ExportedProgram,
compile_specs: List[CompileSpec],
)
| 98 | |
| 99 | @staticmethod |
| 100 | def preprocess( |
| 101 | edge_program: ExportedProgram, |
| 102 | compile_specs: List[CompileSpec], |
| 103 | ) -> PreprocessResult: |
| 104 | # The EdgeIR nodes are processed in the following order: |
| 105 | # 1. Process first the input feeds to the graph (in the same |
| 106 | # order as args from forward(*args)), and generate a unique |
| 107 | # id for each input placeholder. Each input id is appended to |
| 108 | # `input_ids` array from the FlatBuffer schema. |
| 109 | # 2. Process the nodes the graph (e.g `call_function`). For each |
| 110 | # EdgeIR node, create an equivalent MPS node in the FlatBuffer, |
| 111 | # based on which the MPSGraph is constructed at runtime. During |
| 112 | # this process, any visited constant in the EdgeIR is added to the |
| 113 | # final MPS FlatBuffer schema. Each constant id is appended to the |
| 114 | # `constant_ids` FlatBuffer schema. |
| 115 | # 3. After all the inputs, nodes and constants are added to the |
| 116 | # FlatBuffer graph, process the `output` nodes and add their id to |
| 117 | # the `output_ids` array in the schema. |
| 118 | |
| 119 | # TODO: Remove this once we have a better support for the dim-order ops. |
| 120 | # Need to override the verifier to skip the non dim-order ops from tripping the default verifier. |
| 121 | edge_program = _transform( |
| 122 | edge_program, |
| 123 | DimOrderOpsRevertPass(), |
| 124 | override_verifiers=[ |
| 125 | EXIREdgeDialectVerifier( |
| 126 | edge_compile_config=exir.EdgeCompileConfig( |
| 127 | _check_ir_validity=False, # Disable the edge dialect verifier, since we are in the mps backend. |
| 128 | ), |
| 129 | class_only=True, |
| 130 | ) |
| 131 | ], |
| 132 | ) |
| 133 | |
| 134 | mps_graph = MPSGraph( |
| 135 | version="0", |
| 136 | mps_nodes=[], |
| 137 | mps_values=[], |
| 138 | input_ids=[], |
| 139 | output_ids=[], |
| 140 | constant_ids=[], |
| 141 | graph_type=OpType.mps_graph, |
| 142 | constant_segment=DataSegment(0, 0), |
| 143 | ) |
| 144 | |
| 145 | convert_model_to_fp16 = True |
| 146 | for spec in compile_specs: |
| 147 | if spec.key == "use_fp16": |
| 148 | convert_model_to_fp16 = bool(list(bytes(spec.value))[0]) |
| 149 | |
| 150 | logging.debug(f"Convert model to FP16: {convert_model_to_fp16}") |
| 151 | |
| 152 | node_visitors = get_node_visitors(edge_program, convert_model_to_fp16) |
| 153 | if logging.DEBUG >= logging.root.level: |
| 154 | edge_program.graph.print_tabular() |
| 155 | |
| 156 | process_placeholder_nodes( |
| 157 | edge_program, |
nothing calls this directly
no test coverage detected