Reorder TG layers based on mladf report

src/mldebug/aie_util.py

@@ -178,13 +178,19 @@ def on_timeout():
     write(reg_map["DEBUG_CONTROL1"], pc_event << 16)
     return True
 
-  def skip_iterations_to_lock_acq(self, lock_acq_pc, count, sid):
+  def skip_iterations_to_lock_acq(self, lock_acq_pc, count, sid, is_last_layer=False):
     """
     Skip iterations without using counter
     """
     if self._is_test_mode() or count == 0:
       return True
 
+    # The last layer finishes without acquiring a next-layer lock, so there is
+    # no lock-acquire PC to break on. Just let the core run out and return.
+    if is_last_layer:
+      self.impl.continue_aie()
+      return True
+
     self.impl.set_pc_breakpoint(lock_acq_pc)
     self.impl.continue_aie()
     wait_until(self.impl.poll_core_status)

src/mldebug/batch_runner.py

@@ -385,8 +385,12 @@ def _run_stamp(self, layer, sid, target_itr, cur_it=1):
     if self.args.run_flags.skip_iter:
       self.state.error = not utl.skip_iterations(target_itr - cur_it, sid)
     elif self.args.run_flags.skip_iter2:
+      is_last_layer = self.state.get_next_layer_for_stamp(sid, idx=1) is None
       self.state.error = not utl.skip_iterations_to_lock_acq(
-        self.design_info.work_dir.stamp(sid).post_layer_lock_acq_pc, target_itr - cur_it, sid
+        self.design_info.work_dir.stamp(sid).post_layer_lock_acq_pc,
+        target_itr - cur_it,
+        sid,
+        is_last_layer,
       )
     else:
       while cur_it < target_itr:

src/mldebug/layer_info.py

@@ -28,6 +28,9 @@
   "mllib_graphs::mha_type1::mha_adf_wrapper",
   # Causes failure. TODO: investigate
   "superkernel_eltunary",
+  # Padding preamble; halting on it desyncs PC/iteration stepping on HW
+  "buffer_pad_innermost",
+  "superkernel_conv_eltbinary",
 ]
 
 
@@ -880,6 +883,59 @@ def _init_layers(self, raw_info, aie_iface, num_stamps, num_batches=1):
         info, size_shift, version, aie_iface, num_stamps, self.mladf_report, num_batches=num_batches
       )
       self.layers.append(layer)
+    self._reorder_layers_by_execution()
+
+  def _reorder_layers_by_execution(self):
+    """
+    Reorder layers into true execution order using the mladf report.
+
+    buffer_info `layer_order` is the MLIR/DAG index; the aiecompiler backend
+    reschedules layers (notably templated-graph layers) into a different
+    execution / PM-reload order. The mladf `layer_id` captures that real
+    order, so each layer is translated to its `layer_id` (via the existing
+    parent-graph map) and stable-sorted on that single scale. buffer_info
+    order is only the lookup key / tie-break -- the two numberings are never
+    compared numerically.
+
+    A TG layer with no mladf mapping is disabled (dropped later) with a
+    warning; a non-TG layer with no mapping is anchored to its previous
+    neighbour so it keeps its buffer_info position. Non-TG layers should
+    already be in execution order, so a disagreement is flagged.
+    """
+    if not self.mladf_report:
+      return
+
+    keys = []
+    last_seen = -1
+    prev_exec = None
+    disagreement = False
+    for layer in self.layers:
+      exec_order = self.mladf_report.get_exec_order_for_bilo(layer.layer_order)
+      if exec_order is None:
+        if layer.lcp.is_tg and not layer.is_unsupported:
+          LOGGER.log(
+            f"[WARNING] No mladf execution order for TG layer {layer.layer_order}; "
+            "disabling it (its kernel dumps are skipped)."
+          )
+          layer.is_unsupported = True
+        # Anchor an unmapped layer to the previous execution slot (mladf scale).
+        exec_order = last_seen
+      else:
+        last_seen = exec_order
+        if not layer.lcp.is_tg:
+          if prev_exec is not None and exec_order < prev_exec:
+            disagreement = True
+          prev_exec = exec_order
+      keys.append(exec_order)
+
+    if disagreement:
+      LOGGER.log(
+        "[WARNING] buffer_info layer_order disagrees with mladf execution order; "
+        "layer sequencing may be unreliable."
+      )
+
+    order = sorted(range(len(self.layers)), key=lambda i: (keys[i], i))
+    self.layers = [self.layers[i] for i in order]
 
   def _initialize_layers_from_workdir_x2(self, args):
     """
@@ -972,6 +1028,11 @@ def _initialize_layers_from_workdir(self, args):
             ),
             None,
           )
+          # Fall back to the mladf report's authoritative per-core ELF.
+          if elf_id is None and self.mladf_report:
+            mladf_elf = self.mladf_report.get_elfid_for_bilo(layer.layer_order, sid)
+            if mladf_elf not in (None, -1) and str(mladf_elf) in funcs_by_elf:
+              elf_id = str(mladf_elf)
         else:
           elf_id = next((e for e, fns in funcs_by_elf.items() if key in fns), None)
 

src/mldebug/mladf_report.py

@@ -50,6 +50,17 @@ def get_aiec_layers_by_bilo(self, bilo):
     aiec_layer_keys = self.bi_to_m2.get(bilo, [])
     return [self.m2_layers[k] for k in aiec_layer_keys]
 
+  def get_exec_order_for_bilo(self, bilo):
+    """
+    True execution order (mladf `layer_id`) for a buffer_info layer_order.
+
+    buffer_info `layer_order` is the MLIR/DAG index; the backend reschedules
+    layers, and the mladf `layer_id` captures the real execution/PM-reload
+    order. Returns the smallest mapped `layer_id`, or None if unmapped.
+    """
+    ids = [lyr["layer_id"] for lyr in self.get_aiec_layers_by_bilo(bilo) if "layer_id" in lyr]
+    return min(ids) if ids else None
+
   def get_skname_for_bilo(self, bilo, sid=0):
     """
     return superkernel for buffer info layer