tg + multistamp stability fix

src/mldebug/aie_util.py

@@ -178,13 +178,22 @@ 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
 
+    # No next-layer lock to break on: run the core 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

@@ -49,6 +49,8 @@ def __init__(self, args, state, design_info, impls, aie_utls, dumper, status_han
     self.aie_utls = aie_utls
     self.dumper = dumper
     self.status_handle = status_handle
+    # Execution position of --exit_at_layer; resolved in common_init.
+    self.exit_at_index = None
 
   # ------------------------------------------------------------------ #
   # Stamp scheduling
@@ -67,6 +69,16 @@ def common_init(self):
     if self.args.run_flags.skip_iter:
       LOGGER.log("[INFO] All iterations will be skipped for this run.")
 
+    # Resolve once: layer_order is not the execution order, so -e is matched
+    # against the target's position in the execution list.
+    if self.args.exit_at_layer is not None:
+      self.exit_at_index = self.state.exec_index(self.args.exit_at_layer)
+      if self.exit_at_index is None:
+        LOGGER.log(
+          f"[WARNING] Layer {self.args.exit_at_layer} is not in the execution list; "
+          "the run will not exit early."
+        )
+
   def set_pc_breakpoint(self, pc, slot, sid=0):
     """
     Set a PC breakpoint at the given address and slot for the selected stamp.
@@ -328,7 +340,7 @@ def _process_start_breakpoint(self, layer, it, sid=0):
     if self.args.interactive:
       return
 
-    if self.args.exit_at_layer and layer.layer_order >= self.args.exit_at_layer:
+    if self.exit_at_index is not None and self.state.current_layer >= self.exit_at_index:
       LOGGER.log(f"[INFO] Exiting debugger at Layer: {layer.layer_order}")
       self._write_run_summary("SUCCESS")
       sys.exit(0)
@@ -385,8 +397,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 +451,18 @@ 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 replicas with no remaining future layer. Breakpoints must be
+    # cleared first, or the core re-halts at its still-armed start_pc.
     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/debug_state.py

@@ -104,6 +104,23 @@ def get_layer_by_order(self, order):
           return l
     return None
 
+  def exec_index(self, order, default=None):
+    """
+    Position of the layer with this buffer_info `layer_order` in execution
+    order, or `default` if no layer has it.
+
+    Args:
+      order (int): The layer_order value sought.
+      default: Returned when no layer matches.
+
+    Returns:
+      Index into self.layers, or `default`.
+    """
+    for i, l in enumerate(self.layers):
+      if l.layer_order == order:
+        return i
+    return default
+
   def get_next_layer(self):
     """
     Retrieve the layer object for the next sequential layer.
@@ -126,12 +143,13 @@ def get_last_layer(self):
 
   def add_breakpoint(self, layer, iteration):
     """
-    Add a (layer, iteration) tuple as a manual breakpoint
-    and sort the list of manual breakpoints.
+    Add a (layer, iteration) tuple as a manual breakpoint and keep the list in
+    execution order (layer_order is not the execution order).
 
     Args:
-      layer (int): Layer index (or handle).
+      layer (int): Layer order value.
       iteration (int): Iteration number within the layer.
     """
     self.manual_breakpoints.append((layer, iteration))
-    self.manual_breakpoints = sorted(self.manual_breakpoints)
+    last = len(self.layers)
+    self.manual_breakpoints.sort(key=lambda bp: (self.exec_index(bp[0], last), bp[1]))

src/mldebug/interactive_controller.py

@@ -142,7 +142,9 @@ def add_breakpoint(self, layer_num, iteration=1):
     if current_layer:
       current_layer_order = current_layer.layer_order
     final_layer_order = self.state.get_last_layer().layer_order
-    if layer_num < current_layer_order or layer_num > final_layer_order:
+    # Compare execution positions: a later layer can have a smaller layer_order.
+    target = self.state.exec_index(layer_num)
+    if target is None or target < self.state.current_layer:
       print(
         f"[ERROR] Layer Out of bounds. Current: {current_layer_order} Final: {final_layer_order}"
       )
@@ -192,9 +194,10 @@ def continue_execution(self):
       ta_layer, ta_itr = self.state.manual_breakpoints.pop(0)
 
     print(f"Goto next breakpoint at layer {ta_layer} iteration {ta_itr}")
+    ta_index = self.state.exec_index(ta_layer, len(self.state.layers))
     while True:
       cur_layer = self.state.get_current_layer()
-      if not cur_layer or not cur_layer.layer_order < ta_layer:
+      if not cur_layer or not self.state.current_layer < ta_index:
         break
       if not self.step_layer():
         print(f"Unable to continue to breakpoint at layer {ta_layer} iteration {ta_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",
 ]
 
 
@@ -235,7 +238,17 @@ class Layer:
   Contains all buffer, iteration, and kernel (stamp) mapping information.
   """
 
-  def __init__(self, info, size_shift, version, aie_iface, num_stamps, mladf_report, num_batches=1):
+  def __init__(
+    self,
+    info,
+    size_shift,
+    version,
+    aie_iface,
+    num_stamps,
+    mladf_report,
+    num_batches=1,
+    device_batch_size=1,
+  ):
     """
     Initialize a Layer object using given metadata, populating buffer and kernel/stamp lists.
 
@@ -248,6 +261,8 @@ def __init__(self, info, size_shift, version, aie_iface, num_stamps, mladf_repor
       mladf_report: Optional MladfReport for templated-graph layers.
       num_batches (int): Number of batches (B from BxSxCxR overlay). Each
         batch is a data-parallel copy of the per-batch stamps; defaults to 1.
+      device_batch_size (int): buffer_info's device_batch_size, before any
+        single-replica collapse. Gates the mladf stamp-count correction.
     """
     self.flexml_ids = []
     self.l3_ifm_buffers = []
@@ -274,6 +289,17 @@ 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 and mladf is authoritative
+    # TBD: it details only one batch replica, make it work with nBnS mode
+    if mladf_report and device_batch_size == 1 and num_stamps > 1:
+      # None or 0 means mladf could not answer; fall back to buffer_info.
+      true_n = mladf_report.get_running_stamp_count(self.layer_order, num_stamps)
+      if true_n and 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 or n_stamps
     if n_stamps and n_stamps < num_stamps:
       num_stamps = n_stamps
 
@@ -877,9 +903,57 @@ def _init_layers(self, raw_info, aie_iface, num_stamps, num_batches=1):
     for entry in raw_layers:
       info = entry[1]
       layer = Layer(
-        info, size_shift, version, aie_iface, num_stamps, self.mladf_report, num_batches=num_batches
+        info,
+        size_shift,
+        version,
+        aie_iface,
+        num_stamps,
+        self.mladf_report,
+        num_batches=num_batches,
+        device_batch_size=self.layout[0],
       )
       self.layers.append(layer)
+    self._reorder_layers_by_execution()
+
+  def _reorder_layers_by_execution(self):
+    """
+    Stable-sort layers into true execution order (mladf `layer_id`), which the
+    backend reschedules away from buffer_info's MLIR/DAG `layer_order`.
+    """
+    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 +1046,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

@@ -10,6 +10,8 @@
 
 from pathlib import Path
 
+from mldebug.utils import LOGGER
+
 
 def load_json(path):
   """
@@ -50,6 +52,45 @@ 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):
+    """
+    Number of per-batch stamps running a kernel for a buffer_info layer, or
+    None if the report does not describe all `max_stamps` stamp cores.
+    """
+    aiec_layers = self.get_aiec_layers_by_bilo(bilo)
+    running = []
+    described = 0
+    for s in range(max_stamps):
+      # Batch 0 only: per-batch stamp s is the core at column s*cps, row 0.
+      core = f"{s * self.cps}_0"
+      infos = [lyr.get("core_information", {}).get(core) for lyr in aiec_layers]
+      infos = [i for i in infos if i is not None]
+      if not infos:
+        continue
+      described += 1
+      # A core can be listed (its ELF is loaded) with an empty kernel_name,
+      # meaning it does not run this layer's compute.
+      if any(i.get("kernel_name", "") for i in infos):
+        running.append(s)
+    # An absent core makes the count untrustworthy; the caller keeps buffer_info's value.
+    if described < max_stamps:
+      return None
+    # Callers treat the count as stamps 0..n-1, so a gap would mislabel them.
+    if running != list(range(len(running))):
+      LOGGER.log(
+        f"[WARNING] Layer {bilo}: mladf running stamps {running} are not a "
+        "contiguous prefix; stamp mapping may be wrong."
+      )
+    return len(running)
+
+  def get_exec_order_for_bilo(self, bilo):
+    """
+    Smallest mladf `layer_id` (the real execution/PM-reload order) for a
+    buffer_info layer_order, 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