@@ -74,22 +74,28 @@ class EdgeProgramToIRConverter:
7474 _default_target_spec = NeutronTargetSpec ("imxrt700" )
7575 _default_delegation_options = CustomDelegationOptions ()
7676
77+ def __init__ (self ):
78+ self .edge_to_tflite_map = {}
79+
7780 def convert_program (
7881 self ,
7982 edge_program : ExportedProgram ,
8083 conversion_config : ConversionConfig = _default_conversion_config ,
8184 neutron_target_spec : NeutronTargetSpec = _default_target_spec ,
8285 custom_delegation_options : CustomDelegationOptions = _default_delegation_options ,
83- ) -> tuple [bytes , dict [str , dict [str , DataFormat ]]]:
86+ ) -> tuple [bytes , dict [str , dict [str , DataFormat ]], dict [ int , tuple [ int , ...]] ]:
8487 """
8588 Convert ExportedProgram in Edge dialect to IR (TFLite flatbuffers) as bytes.
8689
8790 :param edge_program: Converter ExportedProgram.
8891 :param conversion_config: ConversionConfig instance.
8992 :param neutron_target_spec: Object for querying the target platform to retrieve its properties.
9093 :param custom_delegation_options: Custom user options which affect node delegation.
91- :return: TFLite flatbuffers as bytes.
94+ :return: TFLite flatbuffers as bytes, I/O formats, and edge-to-tflite mapping .
9295 """
96+ # Reset the edge to tflite map for each conversion
97+ self .edge_to_tflite_map = {}
98+
9399 parameters_mapping = self .map_inputs_to_parameters (edge_program )
94100 dim_order_map = self .map_nodes_to_dim_order (edge_program )
95101
@@ -113,14 +119,17 @@ def convert_program(
113119 # Apply optimizations and finalize the model.
114120 internal_tflite_model = cc .tflite_builder .finish ()
115121
122+ # Get the final edge to tflite mapping after optimization
123+ self .edge_to_tflite_map = cc .tflite_builder .edge_to_tflite_map
124+
116125 # Extract the formats of the model's inputs and outputs.
117126 io_formats = cc .tflite_builder .get_io_formats (edge_program .graph_signature )
118127
119128 # TFLite model generation
120129 flatbuffers_builder = flatbuffers .Builder ()
121130 internal_tflite_model .gen_tflite (flatbuffers_builder )
122131
123- return bytes (flatbuffers_builder .Output ()), io_formats
132+ return bytes (flatbuffers_builder .Output ()), io_formats , self . edge_to_tflite_map
124133
125134 @staticmethod
126135 def append_placeholders_and_tensors (nodes : list [Node ], context : ConversionContext ):
@@ -162,7 +171,6 @@ def _process_nodes(self, nodes: list[Node], conversion_context: ConversionContex
162171 exir_ops .edge .quantized_decomposed .dequantize_per_channel .default ,
163172 exir_ops .edge .quantized_decomposed .quantize_per_tensor .default ,
164173 ]
165-
166174 for node in nodes :
167175 if node .op == "call_function" :
168176 if node .target in qdq_related_functions and "cluster" in node .meta :
@@ -174,7 +182,37 @@ def _process_nodes(self, nodes: list[Node], conversion_context: ConversionContex
174182 # The node was already processed alongside the Q/DQ ops.
175183 pass
176184 elif node .target in functions_converters :
185+ # Get TFLite op count BEFORE conversion
186+ tflite_op_count_before = len (
187+ conversion_context .tflite_builder .get_operators ().vector
188+ )
189+ # Convert the node
177190 functions_converters [node .target ](conversion_context ).convert (node )
191+ # Get TFLite op count AFTER conversion
192+ tflite_op_count_after = len (
193+ conversion_context .tflite_builder .get_operators ().vector
194+ )
195+
196+ # Track the mapping - store edge debug handle in operators.
197+ # Get the edge debug handle so it can be associated with newly created operators.
198+ edge_debug_handle = node .meta .get ("debug_handle" , None )
199+ if (
200+ edge_debug_handle is not None
201+ and tflite_op_count_after > tflite_op_count_before
202+ ):
203+ operators = (
204+ conversion_context .tflite_builder .get_operators ().vector
205+ )
206+ # Node converters append new operators to the TFLite builder.
207+ # Only operators added during this conversion step (from "before" to "after")
208+ # are tagged with the current edge_debug_handle.
209+ for i in range (tflite_op_count_before , tflite_op_count_after ):
210+ # Store edge debug handle in operator's temporary attribute
211+ operators [i ].tmp_edge_debug_handle = edge_debug_handle
212+ logger .d (
213+ f"Tagged TFLite ops { list (range (tflite_op_count_before , tflite_op_count_after ))} with edge debug_handle={ edge_debug_handle } for node '{ node .name } '"
214+ )
215+
178216 else :
179217 logger .e (
180218 logger .Code .NOT_IMPLEMENTED ,
0 commit comments