From b3e3edd3ea4028e1c08265de89a8ef559b3761da Mon Sep 17 00:00:00 2001 From: anurag Date: Mon, 27 Jul 2026 10:48:25 -0600 Subject: [PATCH 1/2] cleanup overlay interface Signed-off-by: anurag --- mldebuglib.py | 52 +++++++++++++++++--------------------- src/mldebug/aie_overlay.py | 16 +++++++----- src/mldebug/layer_info.py | 4 ++- 3 files changed, 36 insertions(+), 36 deletions(-) diff --git a/mldebuglib.py b/mldebuglib.py index bd8058d..b7374e8 100644 --- a/mldebuglib.py +++ b/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: diff --git a/src/mldebug/aie_overlay.py b/src/mldebug/aie_overlay.py index 9cb0444..5049d26 100644 --- a/src/mldebug/aie_overlay.py +++ b/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) diff --git a/src/mldebug/layer_info.py b/src/mldebug/layer_info.py index 28120c7..41b06cf 100644 --- a/src/mldebug/layer_info.py +++ b/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() From 544e6db44dfb3c6221fa820a56e64627d6bf4edb Mon Sep 17 00:00:00 2001 From: anurag Date: Mon, 27 Jul 2026 10:58:42 -0600 Subject: [PATCH 2/2] support stamp group size Signed-off-by: anurag --- src/mldebug/layer_info.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/mldebug/layer_info.py b/src/mldebug/layer_info.py index 41b06cf..163b6a1 100644 --- a/src/mldebug/layer_info.py +++ b/src/mldebug/layer_info.py @@ -831,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")