diff --git a/src/mldebug/aie_util.py b/src/mldebug/aie_util.py index 6269a3c..b947a7c 100644 --- a/src/mldebug/aie_util.py +++ b/src/mldebug/aie_util.py @@ -178,13 +178,22 @@ def on_timeout(): write(reg_map["DEBUG_CONTROL1"], pc_event << 16) return True - def skip_iterations_to_lock_acq(self, lock_acq_pc, count, sid): + def skip_iterations_to_lock_acq(self, lock_acq_pc, count, sid, is_last_layer=False): """ Skip iterations without using counter """ if self._is_test_mode() or count == 0: return True + # No next-layer lock to break on: run the core out to Core_Done so it + # releases its locks / program memory for the other stamps' PM reload. + if is_last_layer: + self.impl.clear_pc_breakpoint(0) + self.impl.clear_pc_breakpoint(1) + self.impl.disable_pc_halt() + self.impl.continue_aie() + return True + self.impl.set_pc_breakpoint(lock_acq_pc) self.impl.continue_aie() wait_until(self.impl.poll_core_status) diff --git a/src/mldebug/batch_runner.py b/src/mldebug/batch_runner.py index 247a2df..2f68124 100644 --- a/src/mldebug/batch_runner.py +++ b/src/mldebug/batch_runner.py @@ -49,6 +49,8 @@ def __init__(self, args, state, design_info, impls, aie_utls, dumper, status_han self.aie_utls = aie_utls self.dumper = dumper self.status_handle = status_handle + # Execution position of --exit_at_layer; resolved in common_init. + self.exit_at_index = None # ------------------------------------------------------------------ # # Stamp scheduling @@ -67,6 +69,16 @@ def common_init(self): if self.args.run_flags.skip_iter: LOGGER.log("[INFO] All iterations will be skipped for this run.") + # Resolve once: layer_order is not the execution order, so -e is matched + # against the target's position in the execution list. + if self.args.exit_at_layer is not None: + self.exit_at_index = self.state.exec_index(self.args.exit_at_layer) + if self.exit_at_index is None: + LOGGER.log( + f"[WARNING] Layer {self.args.exit_at_layer} is not in the execution list; " + "the run will not exit early." + ) + def set_pc_breakpoint(self, pc, slot, sid=0): """ Set a PC breakpoint at the given address and slot for the selected stamp. @@ -328,7 +340,7 @@ def _process_start_breakpoint(self, layer, it, sid=0): if self.args.interactive: return - if self.args.exit_at_layer and layer.layer_order >= self.args.exit_at_layer: + if self.exit_at_index is not None and self.state.current_layer >= self.exit_at_index: LOGGER.log(f"[INFO] Exiting debugger at Layer: {layer.layer_order}") self._write_run_summary("SUCCESS") sys.exit(0) @@ -385,8 +397,12 @@ def _run_stamp(self, layer, sid, target_itr, cur_it=1): if self.args.run_flags.skip_iter: self.state.error = not utl.skip_iterations(target_itr - cur_it, sid) elif self.args.run_flags.skip_iter2: + is_last_layer = self.state.get_next_layer_for_stamp(sid, idx=1) is None self.state.error = not utl.skip_iterations_to_lock_acq( - self.design_info.work_dir.stamp(sid).post_layer_lock_acq_pc, target_itr - cur_it, sid + self.design_info.work_dir.stamp(sid).post_layer_lock_acq_pc, + target_itr - cur_it, + sid, + is_last_layer, ) else: while cur_it < target_itr: @@ -435,7 +451,8 @@ def run_layer(self, layer, target_itr=None, cur_it=None): if not res: self.state.error = True - # Unhalt right replicas that have no remaining future layer + # Unhalt replicas with no remaining future layer. Breakpoints must be + # cleared first, or the core re-halts at its still-armed start_pc. overlay = self.design_info.overlay total_replicas = len(self.state.pm_reload) if total_replicas > 1 and (target_itr is None or target_itr == layer.lcp.num_iter): @@ -443,6 +460,9 @@ def run_layer(self, layer, target_itr=None, cur_it=None): if overlay.is_leftmost_in_batch(sid): continue if not self.state.get_next_layer_for_stamp(sid, idx=1): + self.impls[sid].clear_pc_breakpoint(0) + self.impls[sid].clear_pc_breakpoint(1) + self.impls[sid].disable_pc_halt() self.impls[sid].continue_aie() if self.state.error: diff --git a/src/mldebug/debug_state.py b/src/mldebug/debug_state.py index 9ffc2dd..dbbd2f3 100644 --- a/src/mldebug/debug_state.py +++ b/src/mldebug/debug_state.py @@ -104,6 +104,23 @@ def get_layer_by_order(self, order): return l return None + def exec_index(self, order, default=None): + """ + Position of the layer with this buffer_info `layer_order` in execution + order, or `default` if no layer has it. + + Args: + order (int): The layer_order value sought. + default: Returned when no layer matches. + + Returns: + Index into self.layers, or `default`. + """ + for i, l in enumerate(self.layers): + if l.layer_order == order: + return i + return default + def get_next_layer(self): """ Retrieve the layer object for the next sequential layer. @@ -126,12 +143,13 @@ def get_last_layer(self): def add_breakpoint(self, layer, iteration): """ - Add a (layer, iteration) tuple as a manual breakpoint - and sort the list of manual breakpoints. + Add a (layer, iteration) tuple as a manual breakpoint and keep the list in + execution order (layer_order is not the execution order). Args: - layer (int): Layer index (or handle). + layer (int): Layer order value. iteration (int): Iteration number within the layer. """ self.manual_breakpoints.append((layer, iteration)) - self.manual_breakpoints = sorted(self.manual_breakpoints) + last = len(self.layers) + self.manual_breakpoints.sort(key=lambda bp: (self.exec_index(bp[0], last), bp[1])) diff --git a/src/mldebug/interactive_controller.py b/src/mldebug/interactive_controller.py index d863c94..89d51df 100644 --- a/src/mldebug/interactive_controller.py +++ b/src/mldebug/interactive_controller.py @@ -142,7 +142,9 @@ def add_breakpoint(self, layer_num, iteration=1): if current_layer: current_layer_order = current_layer.layer_order final_layer_order = self.state.get_last_layer().layer_order - if layer_num < current_layer_order or layer_num > final_layer_order: + # Compare execution positions: a later layer can have a smaller layer_order. + target = self.state.exec_index(layer_num) + if target is None or target < self.state.current_layer: print( f"[ERROR] Layer Out of bounds. Current: {current_layer_order} Final: {final_layer_order}" ) @@ -192,9 +194,10 @@ def continue_execution(self): ta_layer, ta_itr = self.state.manual_breakpoints.pop(0) print(f"Goto next breakpoint at layer {ta_layer} iteration {ta_itr}") + ta_index = self.state.exec_index(ta_layer, len(self.state.layers)) while True: cur_layer = self.state.get_current_layer() - if not cur_layer or not cur_layer.layer_order < ta_layer: + if not cur_layer or not self.state.current_layer < ta_index: break if not self.step_layer(): print(f"Unable to continue to breakpoint at layer {ta_layer} iteration {ta_itr}") diff --git a/src/mldebug/layer_info.py b/src/mldebug/layer_info.py index 94d9476..f4f0ec4 100644 --- a/src/mldebug/layer_info.py +++ b/src/mldebug/layer_info.py @@ -28,6 +28,9 @@ "mllib_graphs::mha_type1::mha_adf_wrapper", # Causes failure. TODO: investigate "superkernel_eltunary", + # Padding preamble; halting on it desyncs PC/iteration stepping on HW + "buffer_pad_innermost", + "superkernel_conv_eltbinary", ] @@ -235,7 +238,17 @@ class Layer: Contains all buffer, iteration, and kernel (stamp) mapping information. """ - def __init__(self, info, size_shift, version, aie_iface, num_stamps, mladf_report, num_batches=1): + def __init__( + self, + info, + size_shift, + version, + aie_iface, + num_stamps, + mladf_report, + num_batches=1, + device_batch_size=1, + ): """ Initialize a Layer object using given metadata, populating buffer and kernel/stamp lists. @@ -248,6 +261,8 @@ def __init__(self, info, size_shift, version, aie_iface, num_stamps, mladf_repor mladf_report: Optional MladfReport for templated-graph layers. num_batches (int): Number of batches (B from BxSxCxR overlay). Each batch is a data-parallel copy of the per-batch stamps; defaults to 1. + device_batch_size (int): buffer_info's device_batch_size, before any + single-replica collapse. Gates the mladf stamp-count correction. """ self.flexml_ids = [] self.l3_ifm_buffers = [] @@ -274,6 +289,17 @@ def __init__(self, info, size_shift, version, aie_iface, num_stamps, mladf_repor return n_stamps = info.get("no_of_stamps") + # buffer_info's no_of_stamps is sometimes wrong and mladf is authoritative + # TBD: it details only one batch replica, make it work with nBnS mode + if mladf_report and device_batch_size == 1 and num_stamps > 1: + # None or 0 means mladf could not answer; fall back to buffer_info. + true_n = mladf_report.get_running_stamp_count(self.layer_order, num_stamps) + if true_n and n_stamps and true_n != n_stamps: + LOGGER.log( + f"[WARNING] Layer {self.layer_order}: buffer_info no_of_stamps={n_stamps} " + f"disagrees with mladf ({true_n}); using {true_n}." + ) + n_stamps = true_n or n_stamps if n_stamps and n_stamps < num_stamps: num_stamps = n_stamps @@ -877,9 +903,57 @@ def _init_layers(self, raw_info, aie_iface, num_stamps, num_batches=1): for entry in raw_layers: info = entry[1] layer = Layer( - info, size_shift, version, aie_iface, num_stamps, self.mladf_report, num_batches=num_batches + info, + size_shift, + version, + aie_iface, + num_stamps, + self.mladf_report, + num_batches=num_batches, + device_batch_size=self.layout[0], ) self.layers.append(layer) + self._reorder_layers_by_execution() + + def _reorder_layers_by_execution(self): + """ + Stable-sort layers into true execution order (mladf `layer_id`), which the + backend reschedules away from buffer_info's MLIR/DAG `layer_order`. + """ + if not self.mladf_report: + return + + keys = [] + last_seen = -1 + prev_exec = None + disagreement = False + for layer in self.layers: + exec_order = self.mladf_report.get_exec_order_for_bilo(layer.layer_order) + if exec_order is None: + if layer.lcp.is_tg and not layer.is_unsupported: + LOGGER.log( + f"[WARNING] No mladf execution order for TG layer {layer.layer_order}; " + "disabling it (its kernel dumps are skipped)." + ) + layer.is_unsupported = True + # Anchor an unmapped layer to the previous execution slot (mladf scale). + exec_order = last_seen + else: + last_seen = exec_order + if not layer.lcp.is_tg: + if prev_exec is not None and exec_order < prev_exec: + disagreement = True + prev_exec = exec_order + keys.append(exec_order) + + if disagreement: + LOGGER.log( + "[WARNING] buffer_info layer_order disagrees with mladf execution order; " + "layer sequencing may be unreliable." + ) + + order = sorted(range(len(self.layers)), key=lambda i: (keys[i], i)) + self.layers = [self.layers[i] for i in order] def _initialize_layers_from_workdir_x2(self, args): """ @@ -972,6 +1046,11 @@ def _initialize_layers_from_workdir(self, args): ), None, ) + # Fall back to the mladf report's authoritative per-core ELF. + if elf_id is None and self.mladf_report: + mladf_elf = self.mladf_report.get_elfid_for_bilo(layer.layer_order, sid) + if mladf_elf not in (None, -1) and str(mladf_elf) in funcs_by_elf: + elf_id = str(mladf_elf) else: elf_id = next((e for e, fns in funcs_by_elf.items() if key in fns), None) diff --git a/src/mldebug/mladf_report.py b/src/mldebug/mladf_report.py index 6050d33..49206d2 100644 --- a/src/mldebug/mladf_report.py +++ b/src/mldebug/mladf_report.py @@ -10,6 +10,8 @@ from pathlib import Path +from mldebug.utils import LOGGER + def load_json(path): """ @@ -50,6 +52,45 @@ def get_aiec_layers_by_bilo(self, bilo): aiec_layer_keys = self.bi_to_m2.get(bilo, []) return [self.m2_layers[k] for k in aiec_layer_keys] + def get_running_stamp_count(self, bilo, max_stamps): + """ + Number of per-batch stamps running a kernel for a buffer_info layer, or + None if the report does not describe all `max_stamps` stamp cores. + """ + aiec_layers = self.get_aiec_layers_by_bilo(bilo) + running = [] + described = 0 + for s in range(max_stamps): + # Batch 0 only: per-batch stamp s is the core at column s*cps, row 0. + core = f"{s * self.cps}_0" + infos = [lyr.get("core_information", {}).get(core) for lyr in aiec_layers] + infos = [i for i in infos if i is not None] + if not infos: + continue + described += 1 + # A core can be listed (its ELF is loaded) with an empty kernel_name, + # meaning it does not run this layer's compute. + if any(i.get("kernel_name", "") for i in infos): + running.append(s) + # An absent core makes the count untrustworthy; the caller keeps buffer_info's value. + if described < max_stamps: + return None + # Callers treat the count as stamps 0..n-1, so a gap would mislabel them. + if running != list(range(len(running))): + LOGGER.log( + f"[WARNING] Layer {bilo}: mladf running stamps {running} are not a " + "contiguous prefix; stamp mapping may be wrong." + ) + return len(running) + + def get_exec_order_for_bilo(self, bilo): + """ + Smallest mladf `layer_id` (the real execution/PM-reload order) for a + buffer_info layer_order, or None if unmapped. + """ + ids = [lyr["layer_id"] for lyr in self.get_aiec_layers_by_bilo(bilo) if "layer_id" in lyr] + return min(ids) if ids else None + def get_skname_for_bilo(self, bilo, sid=0): """ return superkernel for buffer info layer