auto detect overlay for coredump

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,14 +228,15 @@ 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
 
     Args:
       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

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.