diff --git a/CLAUDE.md b/CLAUDE.md index 6ad3efc..edcf0ea 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -159,6 +159,8 @@ Common runtime flags (`-f`): - `multistamp` -- actually drive every replica (all batches and all per-batch stamps); default is to collapse to a single replica (B=S=1) for sanity. +- `tg_dump` (hidden) -- also dump L2 IFM buffers of templated-graph + (TG) layers, sourced from the mladf report. Off by default. ### Testing @@ -334,6 +336,16 @@ the rest of the system uses. be extended at runtime via `--unsupported_kernels`. - L3 offsets come from `flexmlrt-hsi.json`. Only a single parent spill buffer is supported today -- multi-DDR L3 is not yet handled. +- TG layers have no L1/L2 entries in `buffer_info.json`. Their L2 IFM + buffers come from the mladf report (`_initialize_tg_l2_buffers`): + entries whose last name component contains `ifm` become one + `Buffer` each (`ifm1`, `ifm2`, ..., `tg=True`), with ping and an + optional pong. OFM/wts are not handled. Unlike buffer_info, mladf + lists every replica separately with *absolute* columns (names + `templated_graph__[_]`), so TG buffers are not + mirrored by `sid` offset. Mladf memory-tile row 0 still gets the + same +1 as buffer_info. TG `Lcp` keeps its defaults (all 1), so the + IFM is dumped at every iteration. ### `work_dir.py` -- ELF disassembly @@ -369,6 +381,11 @@ L1 buffers are dumped from either the ping or the pong slot depending on iteration parity. L2 OFM dumps shift the iteration index back by one because OFM is dumped at the start of the *next* iteration. +TG L2 buffers (only with `-f tg_dump`) use their absolute mladf column +as-is; each replica dumps only the entries on its own memory tiles. +Both ping and pong are dumped (we don't know which is live), as +`batch/layer_/l2_it__.bin`. + L3 dumps are not direct reads -- they are forwarded to the FlexML runtime via `DebugServer`, which writes them to disk on FlexML's side and ACKs back. This is necessary because L3 lives in host DDR and @@ -461,6 +478,11 @@ debugger to handle templated-graph (TG) layers when `enable_tg` is on; without this report the debugger marks TG layers as unsupported and skips them. +`get_l2_buffers_for_bilo()` returns the de-duplicated +`l2_data_memory_info` entries (name, ping/pong location, offset and +size in bytes) for a buffer_info layer; it is the source of TG L2 +buffers. + ### `utils.py` and friends Cross-cutting helpers: the global `LOGGER`, a tiny `Version` class, a diff --git a/src/mldebug/input_parser.py b/src/mldebug/input_parser.py index 7c71148..cbaf530 100644 --- a/src/mldebug/input_parser.py +++ b/src/mldebug/input_parser.py @@ -51,6 +51,7 @@ class RunFlags: dump_temps: bool multistamp: bool disable_tg: bool + tg_dump: bool = False @dataclass @@ -145,6 +146,7 @@ def get_flag(s, default=False): get_flag("dump_temps"), get_flag("multistamp"), get_flag("disable_tg"), + get_flag("tg_dump"), ) # Support all stamps in standalone mode and when reporting the design's layers diff --git a/src/mldebug/layer_info.py b/src/mldebug/layer_info.py index a6c5e15..6901595 100644 --- a/src/mldebug/layer_info.py +++ b/src/mldebug/layer_info.py @@ -186,7 +186,9 @@ class Buffer: Top level Buffer e.g: ifm, ofm, wts. Handles mapping and initialization of L1 and L2 buffers. """ - def __init__(self, entry, buf_type, size_shift, aie_iface, ifm=False, ofm=False, wts=False): + def __init__( + self, entry, buf_type, size_shift, aie_iface, ifm=False, ofm=False, wts=False, tg=False + ): """ Initializes a Buffer object that stores L1 and L2 buffer mappings. @@ -198,6 +200,7 @@ def __init__(self, entry, buf_type, size_shift, aie_iface, ifm=False, ofm=False, ifm (bool, optional): Is input feature map buffer. ofm (bool, optional): Is output feature map buffer. wts (bool, optional): Is weights buffer. + tg (bool, optional): Built from the mladf report for a templated graph layer. """ self.l1 = None self.l2 = [] @@ -205,6 +208,7 @@ def __init__(self, entry, buf_type, size_shift, aie_iface, ifm=False, ofm=False, self.ifm = ifm self.ofm = ofm self.wts = wts + self.tg = tg size_shift = self._get_size_shift(entry, size_shift) # pong can be missing (no double buffering) @@ -435,8 +439,9 @@ def __init__( self._initialize_l3_buffers(info, version) - # No L2 support for templated graph layers + # Templated graph layers: only L2 IFM, sourced from the mladf report if self.lcp.is_tg: + self._initialize_tg_l2_buffers(mladf_report, aie_iface) return self._initialize_flexml_ids(info) @@ -457,6 +462,26 @@ def get_stamp(self, sid): """ return self.stamps[int(sid % self.overlay_stamps_per_batch)] + def _initialize_tg_l2_buffers(self, mladf_report, aie_iface): + """ + Build L2 IFM buffers for a TG layer from mladf entries, as buffer_info-style entries. + """ + for e in mladf_report.get_l2_buffers_for_bilo(self.layer_order, self.mladf_ids): + # mladf lists every replica with absolute columns; the dumper picks each sid's own entries + if "ifm" not in e["name"].split(".")[-1]: + continue + # Sizes are already in bytes + entry = {"dtype": "uint8_t"} + for slot in ("ping", "pong"): + if f"{slot}_location" in e: + loc = e[f"{slot}_location"] + entry[f"l2_{slot}"] = [ + [loc["column"], loc["row"], e[f"{slot}_offset_in_bytes"], e[f"{slot}_size_in_bytes"]] + ] + entry[f"l2_{slot}_buffer_names"] = [f"{e['name']}_{slot}"] + buf_type = f"ifm{len(self.in_buffers) + 1}" + self.in_buffers.append(Buffer(entry, buf_type, None, aie_iface, ifm=True, tg=True)) + def _initialize_flexml_ids(self, info): """ Populate self.flexml_ids for this layer from its metadata. diff --git a/src/mldebug/memory_dumper.py b/src/mldebug/memory_dumper.py index 325e4cd..17589af 100644 --- a/src/mldebug/memory_dumper.py +++ b/src/mldebug/memory_dumper.py @@ -123,13 +123,20 @@ def dump_memory_l2(self, buffers, it, layer_order=None, use_l2_names=False, sid= suffix = f"stamp{overlay.replica_to_stamp(sid)}" for buffer in buffers: + if buffer.tg and not self.args.run_flags.tg_dump: + continue if buffer.ofm: it -= 1 # OFM is dumped at start of next iteration for buf_id, group in groupby(buffer.l2, key=lambda x: x.buf_id): + group = list(group) buffer_name = f"l2_it{it}_{buf_id}_{suffix}" + if buffer.tg: + # e.g. layer_51/l2_it3_compute_graph.templated_graph_51_0.ifm_ping + buffer_name = f"l2_it{it}_{group[0].name}" data = [] - for l2 in list(group): - col = l2.col + sid * overlay.get_stampwidth() + for l2 in group: + # TG columns are absolute (each replica has its own entry); others are stamp-relative + col = l2.col if buffer.tg else l2.col + sid * overlay.get_stampwidth() if (col, l2.row) not in overlay.get_tiles(self.args.aie_iface.MEM_TILE_T, stamp_id=sid): continue data.extend(self.impl.dump_memory(col, l2.row, l2.address, l2.size)) @@ -137,7 +144,12 @@ def dump_memory_l2(self, buffers, it, layer_order=None, use_l2_names=False, sid= break if use_l2_names and l2.name: buffer_name = l2.name - if self.args.run_flags.l2_ifm_dump or self.args.run_flags.l2_dump_only: + if buffer.tg: + # Entry belongs to another replica + if not data: + continue + fname = f"{self.get_output_path(layer_order=layer_order, batch=batch)}/{buffer_name}" + elif self.args.run_flags.l2_ifm_dump or self.args.run_flags.l2_dump_only: fname = f"{self.get_output_path(layer_order=layer_order, batch=batch)}/{buffer.type}_{buffer_name}" else: fname = f"{self.get_output_path(buffer=buffer.type, layer_order=layer_order, batch=batch)}/{buffer_name}" diff --git a/src/mldebug/mladf_report.py b/src/mldebug/mladf_report.py index 4e06a60..346c998 100644 --- a/src/mldebug/mladf_report.py +++ b/src/mldebug/mladf_report.py @@ -99,6 +99,15 @@ def get_aiec_layers_by_bilo(self, bilo, ids=None): ids = set(ids) return [lyr for lyr in layers if lyr.get("layer_id") in ids] + def get_l2_buffers_for_bilo(self, bilo, ids=None): + """De-duplicated `l2_data_memory_info` entries for a buffer_info layer.""" + entries = {} + for lyr in self.get_aiec_layers_by_bilo(bilo, ids): + for e in lyr.get("l2_data_memory_info", []): + key = (e["name"], e["ping_location"]["column"], e["ping_offset_in_bytes"]) + entries.setdefault(key, e) + return list(entries.values()) + def get_running_stamp_count(self, bilo, max_stamps, ids=None): """ Number of per-batch stamps running a kernel for a buffer_info layer, or diff --git a/src/mldebug/mldebug_cli.py b/src/mldebug/mldebug_cli.py index c256c0e..6ebbcda 100644 --- a/src/mldebug/mldebug_cli.py +++ b/src/mldebug/mldebug_cli.py @@ -430,6 +430,7 @@ def app(): "multistamp", "disable_tg", "skip_iter2", + "tg_dump", ], help="Specify one or more runtime flags:\n" "skip_dump : Do not dump memory\n" @@ -443,6 +444,7 @@ def app(): # "dump_temps : Write intermediate (.lst) files to disk\n" "multistamp : Enable N Stamp/Batch mode\n", # "disable_tg : Disable Step to TG layers\n", + # "tg_dump : Dump L2 IFM buffers of TG layers\n", # 'mock_hang' : Simulate hang at one of the layers in test mode metavar=" ", )