tg + multistamp stability fix

src/mldebug/aie_util.py

@@ -178,13 +178,26 @@ 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. Clear this stamp's breakpoints and stop
+    # halting on PC events, otherwise continuing would only advance one
+    # iteration before re-halting at the still-armed layer start_pc; the core
+    # must run all remaining iterations out to Core_Done so it releases its
+    # locks / program memory for the other stamps' PM reload.
+    if is_last_layer:
+      self.impl.clear_pc_breakpoint(0)
+      self.impl.clear_pc_breakpoint(1)
+      self.impl.disable_pc_halt()
+      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:
@@ -435,14 +439,21 @@ def run_layer(self, layer, target_itr=None, cur_it=None):
         if not res:
           self.state.error = True
 
-    # Unhalt right replicas that have no remaining future layer
+    # Unhalt right replicas that have no remaining future layer. Clear their
+    # breakpoints and disable PC-halt first; otherwise continue_aie() only
+    # advances one iteration before the core re-halts at its still-armed
+    # start_pc, leaving the replica stuck (and blocking other stamps' PM
+    # reload, which needs every core to run out).
     overlay = self.design_info.overlay
     total_replicas = len(self.state.pm_reload)
     if total_replicas > 1 and (target_itr is None or target_itr == layer.lcp.num_iter):
       for sid in range(total_replicas):
         if overlay.is_leftmost_in_batch(sid):
           continue
         if not self.state.get_next_layer_for_stamp(sid, idx=1):
+          self.impls[sid].clear_pc_breakpoint(0)
+          self.impls[sid].clear_pc_breakpoint(1)
+          self.impls[sid].disable_pc_halt()
           self.impls[sid].continue_aie()
 
     if self.state.error:

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",
 ]
 
 
@@ -274,6 +277,18 @@ def __init__(self, info, size_shift, version, aie_iface, num_stamps, mladf_repor
       return
 
     n_stamps = info.get("no_of_stamps")
+    # buffer_info's no_of_stamps is sometimes wrong (a core can be listed with
+    # an empty kernel_name). The true count is how many stamps actually run a
+    # kernel per the mladf report; prefer it and warn on a disagreement.
+    if mladf_report:
+      true_n = mladf_report.get_running_stamp_count(self.layer_order, num_stamps)
+      if true_n:
+        if n_stamps and true_n != n_stamps:
+          LOGGER.log(
+            f"[WARNING] Layer {self.layer_order}: buffer_info no_of_stamps={n_stamps} "
+            f"disagrees with mladf ({true_n}); using {true_n}."
+          )
+        n_stamps = true_n
     if n_stamps and n_stamps < num_stamps:
       num_stamps = n_stamps
 
@@ -880,6 +895,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 +1040,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,37 @@ 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_running_stamp_count(self, bilo, max_stamps):
+    """
+    True number of stamps that actually run a kernel for a buffer_info layer.
+
+    A stamp `sid` runs the layer only if its leftmost core (`sid*cps`_0) has a
+    NON-EMPTY kernel_name in the mladf core_information. A core may be listed
+    (its ELF is loaded) with an empty kernel_name, meaning it does not run the
+    layer -- so mere core presence over-counts. Assumes stamps are contiguous
+    from 0. Returns 0 when the layer has no mladf mapping.
+    """
+    count = 0
+    for sid in range(max_stamps):
+      core = f"{sid * self.cps}_0"
+      for lyr in self.get_aiec_layers_by_bilo(bilo):
+        ci = lyr.get("core_information", {})
+        if core in ci and ci[core].get("kernel_name", ""):
+          count += 1
+          break
+    return count
+
+  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