From e3f8b7c04ede2e07230210922f6a97bd7d2cfd32 Mon Sep 17 00:00:00 2001
From: anurag <anurag@amd.com>
Date: Wed, 5 Aug 2026 17:04:30 -0600
Subject: [PATCH 1/7] reorder execution based on mladf

Signed-off-by: anurag <anurag@amd.com>
---
 src/mldebug/layer_info.py   | 53 +++++++++++++++++++++++++++++++++++++
 src/mldebug/mladf_report.py | 11 ++++++++
 2 files changed, 64 insertions(+)

diff --git a/src/mldebug/layer_info.py b/src/mldebug/layer_info.py
index 94d9476..3d3a849 100644
--- a/src/mldebug/layer_info.py
+++ b/src/mldebug/layer_info.py
@@ -880,6 +880,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 "
+        "for non-TG layers; 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):
     """
diff --git a/src/mldebug/mladf_report.py b/src/mldebug/mladf_report.py
index 6050d33..9a26e88 100644
--- a/src/mldebug/mladf_report.py
+++ b/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

From 44b91c14cee9a7fa06076796efa0b4efa3485941 Mon Sep 17 00:00:00 2001
From: anurag <anurag@amd.com>
Date: Wed, 5 Aug 2026 17:55:04 -0600
Subject: [PATCH 2/7] fix minor issues

Signed-off-by: anurag <anurag@amd.com>
---
 src/mldebug/aie_util.py     | 8 +++++++-
 src/mldebug/batch_runner.py | 6 +++++-
 src/mldebug/layer_info.py   | 2 ++
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/src/mldebug/aie_util.py b/src/mldebug/aie_util.py
index 6269a3c..4819e87 100644
--- a/src/mldebug/aie_util.py
+++ b/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)
diff --git a/src/mldebug/batch_runner.py b/src/mldebug/batch_runner.py
index 247a2df..3c1842b 100644
--- a/src/mldebug/batch_runner.py
+++ b/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:
diff --git a/src/mldebug/layer_info.py b/src/mldebug/layer_info.py
index 3d3a849..254e880 100644
--- a/src/mldebug/layer_info.py
+++ b/src/mldebug/layer_info.py
@@ -28,6 +28,8 @@
   "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",
 ]
 
 

From b80a58186242254cf57fb4d9d06c04a8ea6935d2 Mon Sep 17 00:00:00 2001
From: anurag <anurag@amd.com>
Date: Wed, 5 Aug 2026 17:59:25 -0600
Subject: [PATCH 3/7] blacklist kernel

Signed-off-by: anurag <anurag@amd.com>
---
 src/mldebug/layer_info.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/mldebug/layer_info.py b/src/mldebug/layer_info.py
index 254e880..619a220 100644
--- a/src/mldebug/layer_info.py
+++ b/src/mldebug/layer_info.py
@@ -30,6 +30,7 @@
   "superkernel_eltunary",
   # Padding preamble; halting on it desyncs PC/iteration stepping on HW
   "buffer_pad_innermost",
+  "superkernel_conv_eltbinary",
 ]
 
 

From a12f0566a3c8da98e3db0e8960eb99db471deae5 Mon Sep 17 00:00:00 2001
From: anurag <anurag@amd.com>
Date: Thu, 6 Aug 2026 11:29:38 -0600
Subject: [PATCH 4/7] fix multistamp failure

Signed-off-by: anurag <anurag@amd.com>
---
 src/mldebug/layer_info.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/src/mldebug/layer_info.py b/src/mldebug/layer_info.py
index 619a220..4db1952 100644
--- a/src/mldebug/layer_info.py
+++ b/src/mldebug/layer_info.py
@@ -930,8 +930,8 @@ def _reorder_layers_by_execution(self):
 
     if disagreement:
       LOGGER.log(
-        "[WARNING] buffer_info layer_order disagrees with mladf execution order "
-        "for non-TG layers; layer sequencing may be unreliable."
+        "[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))
@@ -1028,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)
 

From 49733c9ec504710a467beeb088a590451169dca9 Mon Sep 17 00:00:00 2001
From: anurag <anurag@amd.com>
Date: Fri, 7 Aug 2026 14:26:33 -0600
Subject: [PATCH 5/7] autodetect disparity between bi and mladf

Signed-off-by: anurag <anurag@amd.com>
---
 src/mldebug/aie_util.py     |  9 ++++++++-
 src/mldebug/batch_runner.py |  9 ++++++++-
 src/mldebug/layer_info.py   | 12 ++++++++++++
 src/mldebug/mladf_report.py | 20 ++++++++++++++++++++
 4 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/src/mldebug/aie_util.py b/src/mldebug/aie_util.py
index 4819e87..daf05ed 100644
--- a/src/mldebug/aie_util.py
+++ b/src/mldebug/aie_util.py
@@ -186,8 +186,15 @@ def skip_iterations_to_lock_acq(self, lock_acq_pc, count, sid, is_last_layer=Fal
       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.
+    # 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
 
diff --git a/src/mldebug/batch_runner.py b/src/mldebug/batch_runner.py
index 3c1842b..ecbbb06 100644
--- a/src/mldebug/batch_runner.py
+++ b/src/mldebug/batch_runner.py
@@ -439,7 +439,11 @@ 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):
@@ -447,6 +451,9 @@ def run_layer(self, layer, target_itr=None, cur_it=None):
         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:
diff --git a/src/mldebug/layer_info.py b/src/mldebug/layer_info.py
index 4db1952..0e98c59 100644
--- a/src/mldebug/layer_info.py
+++ b/src/mldebug/layer_info.py
@@ -277,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
 
diff --git a/src/mldebug/mladf_report.py b/src/mldebug/mladf_report.py
index 9a26e88..0a5a02f 100644
--- a/src/mldebug/mladf_report.py
+++ b/src/mldebug/mladf_report.py
@@ -50,6 +50,26 @@ 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.

From 1e4e4fd26a7dc8e8819a3c574bbc26d583db48fa Mon Sep 17 00:00:00 2001
From: anurag <anurag@amd.com>
Date: Mon, 10 Aug 2026 11:58:27 -0600
Subject: [PATCH 6/7] cleanup comments and ignore mladf for nbatch design

Signed-off-by: anurag <anurag@amd.com>
---
 src/mldebug/aie_util.py               |  8 +---
 src/mldebug/batch_runner.py           | 21 +++++++---
 src/mldebug/debug_state.py            | 26 ++++++++++--
 src/mldebug/interactive_controller.py |  7 +++-
 src/mldebug/layer_info.py             | 60 +++++++++++++++------------
 src/mldebug/mladf_report.py           | 52 +++++++++++++----------
 6 files changed, 108 insertions(+), 66 deletions(-)

diff --git a/src/mldebug/aie_util.py b/src/mldebug/aie_util.py
index daf05ed..b947a7c 100644
--- a/src/mldebug/aie_util.py
+++ b/src/mldebug/aie_util.py
@@ -185,12 +185,8 @@ def skip_iterations_to_lock_acq(self, lock_acq_pc, count, sid, is_last_layer=Fal
     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.
+    # 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)
diff --git a/src/mldebug/batch_runner.py b/src/mldebug/batch_runner.py
index ecbbb06..2f68124 100644
--- a/src/mldebug/batch_runner.py
+++ b/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)
@@ -439,11 +451,8 @@ 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. 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).
+    # 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):
diff --git a/src/mldebug/debug_state.py b/src/mldebug/debug_state.py
index 9ffc2dd..dbbd2f3 100644
--- a/src/mldebug/debug_state.py
+++ b/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]))
diff --git a/src/mldebug/interactive_controller.py b/src/mldebug/interactive_controller.py
index d863c94..89d51df 100644
--- a/src/mldebug/interactive_controller.py
+++ b/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}")
diff --git a/src/mldebug/layer_info.py b/src/mldebug/layer_info.py
index 0e98c59..34a3055 100644
--- a/src/mldebug/layer_info.py
+++ b/src/mldebug/layer_info.py
@@ -238,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.
 
@@ -251,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 = []
@@ -277,18 +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 (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:
+    # buffer_info's no_of_stamps is sometimes wrong; mladf is authoritative,
+    # but it details only one batch replica so it is unusable when B > 1.
+    if mladf_report and device_batch_size == 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:
-        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 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
 
@@ -892,27 +903,22 @@ 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):
     """
-    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.
+    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
diff --git a/src/mldebug/mladf_report.py b/src/mldebug/mladf_report.py
index 0a5a02f..49206d2 100644
--- a/src/mldebug/mladf_report.py
+++ b/src/mldebug/mladf_report.py
@@ -10,6 +10,8 @@
 
 from pathlib import Path
 
+from mldebug.utils import LOGGER
+
 
 def load_json(path):
   """
@@ -52,31 +54,39 @@ def get_aiec_layers_by_bilo(self, bilo):
 
   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.
+    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.
     """
-    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
+    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):
     """
-    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.
+    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

From cd7209860df58818c2aa2fc146ceb08859d902b0 Mon Sep 17 00:00:00 2001
From: anurag <anurag@amd.com>
Date: Mon, 10 Aug 2026 12:23:02 -0600
Subject: [PATCH 7/7] don't check mladf for single stamp

Signed-off-by: anurag <anurag@amd.com>
---
 src/mldebug/layer_info.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/mldebug/layer_info.py b/src/mldebug/layer_info.py
index 34a3055..f4f0ec4 100644
--- a/src/mldebug/layer_info.py
+++ b/src/mldebug/layer_info.py
@@ -289,9 +289,9 @@ def __init__(
       return
 
     n_stamps = info.get("no_of_stamps")
-    # buffer_info's no_of_stamps is sometimes wrong; mladf is authoritative,
-    # but it details only one batch replica so it is unusable when B > 1.
-    if mladf_report and device_batch_size == 1:
+    # 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: