Support ministamps in overlay

mldebuglib.py

@@ -10,7 +10,7 @@
 from types import SimpleNamespace
 
 from mldebug.aie_status import AIEStatus as _AIES
-from mldebug.arch import AIE_DEV_PHX, AIE_DEV_STX, AIE_DEV_TEL
+from mldebug.arch import AIE_DEV_PHX, AIE_DEV_STX
 from mldebug.arch import load_aie_arch as _load_aie
 from mldebug.aie_overlay import Overlay as _OL
 from mldebug.input_parser import RunFlags as _RunFlags
@@ -37,27 +37,27 @@ class MLDebug:
   # Parse overlay string to create layout array [stamps, ncol, nrow]
   @staticmethod
   def parse_overlay_string(overlay_str):
-      """
-      Parse overlay string and convert to layout array.
-      Examples:
-        "4x4" -> [1, 4, 4]     (default 1 stamp)
-        "2x4x4" -> [2, 4, 4]   (2 stamps, 4 cols, 4 rows)
-        "1x8x8" -> [1, 8, 8]   (1 stamp, 8 cols, 8 rows)
-      """
-      parts = overlay_str.split('x')
-
-      if len(parts) == 2:
-          # Format: "4x4" (cols x rows, default to 1 stamp)
-          cols, rows = map(int, parts)
-          return [1, cols, rows]
-      elif len(parts) == 3:
-          # Format: "2x4x4" (stamps x cols x rows)
-          stamps, cols, rows = map(int, parts)
-          return [stamps, cols, rows]
-      else:
-          # Fallback to default if format is unexpected
-          print(f"Warning: Unexpected overlay format '{overlay_str}', using default [1, 4, 4]")
-          return [1, 4, 4]
+    """
+    Parse overlay string and convert to layout array.
+    Examples:
+      "4x4" -> [1, 4, 4]     (default 1 stamp)
+      "2x4x4" -> [2, 4, 4]   (2 stamps, 4 cols, 4 rows)
+      "1x8x8" -> [1, 8, 8]   (1 stamp, 8 cols, 8 rows)
+    """
+    parts = overlay_str.split('x')
+
+    if len(parts) == 2:
+      # Format: "4x4" (cols x rows, default to 1 stamp)
+      cols, rows = map(int, parts)
+      return [1, cols, rows]
+    elif len(parts) == 3:
+      # Format: "2x4x4" (stamps x cols x rows)
+      stamps, cols, rows = map(int, parts)
+      return [stamps, cols, rows]
+    else:
+      # Fallback to default if format is unexpected
+      print(f"Warning: Unexpected overlay format '{overlay_str}', using default [1, 4, 4]")
+      return [1, 4, 4]
 
   def __init__(self, device=AIE_DEV_STX, overlay="4x4", ctxid=None, pid=None, backend=BACKEND_XRT):
     if backend not in (BACKEND_XRT, BACKEND_TEST):
@@ -73,14 +73,8 @@ def __init__(self, device=AIE_DEV_STX, overlay="4x4", ctxid=None, pid=None, back
       if ctxid is None or pid is None:
         ctxid, pid = _check_hwc(args)
 
-    class OverlayArgs:
-      def __init__(self, aie_iface, overlay_string):
-        self.aie_iface = aie_iface
-        self.overlay = overlay_string
-
-    overlay_args = OverlayArgs(self.aie_iface, overlay)
     layout = self.parse_overlay_string(overlay)
-    self._ov_hdl = _OL(overlay_args, layout)
+    self._ov_hdl = _OL(self.aie_iface, layout, overlay=overlay)
     tiles = self._ov_hdl.get_tiles(self.aie_iface.AIE_TILE_T, stamp_id=0)
 
     if backend == BACKEND_TEST:

src/mldebug/aie_overlay.py

@@ -21,21 +21,25 @@ class Overlay:
   refers to as "sid" (stamp id).
   """
 
-  def __init__(self, args, layout):
+  def __init__(self, aie_iface, layout, overlay=None, multistamp=False):
     """
     Initialize the Overlay with layout and tile information.
 
     Args:
-      args: Argument object containing configuration options, including
-        aie_iface and overlay string.
+      aie_iface: AIE architecture interface providing tile-type constants and
+        helpers (row offset, filter_tiles).
       layout: Tuple representing the layout from buffer_info. Either
         (batches, stamps, nrow, ncol) (new 4-element form) or
         (stamps, nrow, ncol) (legacy; treated as batches=1).
+      overlay (str, optional): User-specified overlay override string (e.g.
+        '2x4x4' or '4x4'). Forces batches=1 when supplied.
+      multistamp (bool, optional): When False, collapse to a single active
+        replica. Defaults to False.
     """
-    self.aie_iface = args.aie_iface
+    self.aie_iface = aie_iface
     self.stamps = {}
     self.impls = {}
-    batches, stamps_per_batch, ncol, nrow = self._get_layout(args.overlay, layout)
+    batches, stamps_per_batch, ncol, nrow = self._get_layout(overlay, layout)
 
     # Materialize tiles for every physical replica so dropped ones stay quiescible.
     for b in range(batches):
@@ -50,7 +54,7 @@ def __init__(self, args, layout):
 
     # Without `multistamp`, collapse to one active replica so LayerInfo/DebugState/
     # backends size to it; extras stay in self.stamps (see get_inactive_tiles).
-    if args.run_flags.multistamp:
+    if multistamp:
       self.layout = (batches, stamps_per_batch, ncol, nrow)
     else:
       self.layout = (1, 1, ncol, nrow)

src/mldebug/layer_info.py

@@ -521,7 +521,9 @@ def __init__(self, args):
     if has_bi:
       data = self._read_buffer_info(args.buffer_info)
     # 2. Initialize Overlay from Layout
-    self.overlay = Overlay(args, self.layout)
+    self.overlay = Overlay(
+      args.aie_iface, self.layout, overlay=args.overlay, multistamp=args.run_flags.multistamp
+    )
     # Re-sync local view in case Overlay applied -o overrides.
     num_batches = self.overlay.get_batch_count()
     num_stamps = self.overlay.get_stamps_per_batch()
@@ -829,12 +831,18 @@ def _read_buffer_info(self, buffer_info_file):
     # S (per-batch stamps used) comes from max_stamps_used, with sensible
     # fallbacks: layer hints, then the overlay's nominal stamp count.
     stamps = data[".meta"].get("max_stamps_used")
+
+    # Stamp group results in one stamp being used as N stamps
+    max_stamp_group_size = data[".meta"].get("max_stamp_group_size", 1)
+
     if not stamps:
       if data.get("layers"):
         stamps = max(lyr.get("no_of_stamps", 1) for _, lyr in data["layers"].items())
       else:
         stamps = overlay_stamps
 
+    stamps = stamps // max_stamp_group_size
+
     self.layout = (batches, stamps, nrow, ncol)
     if batches > 1:
       LOGGER.log("Batched design detected")