MCPcopy Create free account
hub / github.com/pytorch/executorch / preprocess

Method preprocess

backends/apple/mps/mps_preprocess.py:100–212  ·  view source on GitHub ↗
(
        edge_program: ExportedProgram,
        compile_specs: List[CompileSpec],
    )

Source from the content-addressed store, hash-verified

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,

Callers

nothing calls this directly

Calls 15

appendMethod · 0.95
_transformFunction · 0.90
EXIREdgeDialectVerifierFunction · 0.90
MPSGraphClass · 0.90
DataSegmentClass · 0.90
get_node_visitorsFunction · 0.90
convert_to_flatbufferFunction · 0.90
PreprocessResultClass · 0.90
CordClass · 0.85
debugMethod · 0.80

Tested by

no test coverage detected