diff --git a/src/mldebug/backend/core_dump_impl.py b/src/mldebug/backend/core_dump_impl.py index e66a9db..ec81be6 100644 --- a/src/mldebug/backend/core_dump_impl.py +++ b/src/mldebug/backend/core_dump_impl.py @@ -44,6 +44,7 @@ def __init__(self, core_dump_file, dev_name, no_header=False): self.metadata = DEVICE_CONFIGS[self.dev_name] self.header_size = 256 # Default header size + self.total_cols = 0 # Column count from the header; 0 when there is no header # Open the binary dump file if not Path(self.filename).exists(): @@ -134,7 +135,7 @@ def peek_device(filename): def _parse_header(self): """ - Parse the core dump file header to learn header_size. + Parse the core dump file header to learn header_size and total_cols. Device detection is handled earlier (see ``peek_device`` and ``set_device``); this only validates the magic number and reads the @@ -170,6 +171,11 @@ def _parse_header(self): f"Invalid header size in core dump: {self.header_size} bytes (expected 18-1048576)" ) + # hwGen, coreRowStart, memRowStart, memTileRows, totalNumRows, totalNumCols + geometry = self.file_handle.read(6) + if len(geometry) == 6: + self.total_cols = geometry[5] + except (ValueError, RuntimeError) as e: raise ValueError("I/O error while reading core dump header") from e except OSError as e: @@ -222,7 +228,7 @@ def _calculate_file_position(self, col, row, offset): file_position = self.header_size + tile_pos_index + offset return file_position - def read_register(self, col, row, offset): + def read_register(self, col, row, offset, strict=False): """ Read a 32-bit register from the core dump file @@ -230,6 +236,7 @@ def read_register(self, col, row, offset): col (int): Column index row (int): Row index offset (int): Register offset + strict (bool): Raise on read failure instead of returning 0 Returns: int: 32-bit register value @@ -255,9 +262,29 @@ def read_register(self, col, row, offset): return value except Exception as e: + if strict: + raise print(f"[ERROR] Failed to read register at col={col}, row={row}, offset=0x{offset:x}: {e}") return 0 + def detect_num_cols(self): + """ + Return the number of columns captured in the core dump. + + Reads the same register in successive columns until the read raises, + bounded by the header's column count (the device's, when the header + carried none). + """ + max_cols = self.total_cols or self.metadata["numcols"] + # Probe the last row so a column only counts when its whole tower is present. + row = self.metadata["numrows"] - 1 + for col in range(max_cols): + try: + self.read_register(col, row, 0, strict=True) + except (ValueError, RuntimeError, OSError): + return col + return max_cols + def dump_buffer(self, col, row, offset, size): """ Read a buffer from memory (L1/L2) and return as list of 32-bit words diff --git a/src/mldebug/input_parser.py b/src/mldebug/input_parser.py index 48af1f0..7c71148 100644 --- a/src/mldebug/input_parser.py +++ b/src/mldebug/input_parser.py @@ -22,6 +22,7 @@ AIE_DEV_PHX, AIE_DEV_STX, AIE_DEV_TEL, + DEVICE_CONFIGS, ) from mldebug.backend.core_dump_impl import CoreDumpFallbackReader from mldebug.backend.factory import BackendConfig, create_backend @@ -99,6 +100,7 @@ def create_run_flags(args, subgraph_path: str, fsp: str, fsp_execution_order: li args.subgraph_name = None set_device(args) + set_coredump_overlay(args) set_aie_scc(args, subgraph_path) # Metadata check @@ -335,6 +337,38 @@ def set_device(args) -> None: print(f"[INFO] Using AIE Device: {args.device}.", end=endmsg) +def set_coredump_overlay(args) -> None: + """ + Size the overlay to a core dump: columns are probed from the file, rows come + from the device's core-tile count. + + Only runs for core dumps when the user did not pass -o. Must be called after + ``set_device``, which resolves the device and ``args.no_header``. + """ + if not getattr(args, "core_dump", None) or args.overlay: + return + + dev_name = getattr(args, "sub_device", None) or args.device + try: + reader = CoreDumpFallbackReader( + args.core_dump, dev_name, no_header=getattr(args, "no_header", False) + ) + ncols = reader.detect_num_cols() + except (ValueError, OSError, RuntimeError) as e: + print(f"[WARNING] Could not detect core dump columns: {e}") + return + + if ncols < 1: + print("[WARNING] Core dump has no readable columns. Keeping the default overlay.") + return + + # Overlay rows count core tiles only; shim/memtile rows are added back by Overlay. + cfg = DEVICE_CONFIGS[dev_name] + nrows = cfg["numrows"] - cfg["core_row_start"] + args.overlay = f"{ncols}x{nrows}" + print(f"[INFO] Core dump contains {ncols} column(s). Using overlay {args.overlay}.") + + def print_hw_context_table(current_contexts: dict[str, dict[str, str]]) -> None: """ Prints the current hardware contexts in a table format.