Add Graph Execution to host_build_graph

simpler_setup/goldens/qwen3_14b_decode.py

@@ -115,26 +115,32 @@ def _paged_block_table_slot_mapping(seq_lens: torch.Tensor) -> tuple[torch.Tenso
     return block_table, slot_mapping
 
 
-def generate_inputs(seed: int = 1234, seq_len: int = DEFAULT_SEQ_LEN) -> TaskArgsBuilder:
-    """Deterministic fixture for decode_fwd_layers (N=40), stacked x40 along dim 0.
+def generate_inputs(
+    seed: int = 1234,
+    seq_len: int = DEFAULT_SEQ_LEN,
+    n_layers: int = N_LAYERS,
+) -> TaskArgsBuilder:
+    """Deterministic fixture for decode_fwd_layers, stacked along dim 0.
 
     Every lane uses sequence length ``seq_len`` (default 3500, the stress prompt).
     Per-layer weights are replicated (stack0) so every layer reuses layer 0's
     weights, matching the lib's const-layer-0 stacked-fwd reference; each layer
     still has its own KV pool.
 
-    The stacks are the bulk of the fixture: ~24.6 GiB of weights plus ~13.4 GiB
-    of paged KV at this regime.
+    At the default 40 layers, the stacks are the bulk of the fixture:
+    ~24.6 GiB of weights plus ~13.4 GiB of paged KV at this regime.
     """
     if not (1 <= seq_len <= MAX_SEQ):
         raise ValueError(f"seq_len must be in [1, {MAX_SEQ}], got {seq_len}")
+    if n_layers <= 0:
+        raise ValueError(f"n_layers must be positive, got {n_layers}")
     g = torch.Generator().manual_seed(seed)
 
     def rn(shape, std=1.0, bias=0.0):
         return torch.empty(shape).normal_(0.0, std, generator=g) + bias
 
     def s0(t):  # replicate along dim 0 (one slice per layer)
-        return torch.cat([t] * N_LAYERS, dim=0).contiguous()
+        return torch.cat([t] * n_layers, dim=0).contiguous()
 
     seq_lens = torch.full([BATCH], seq_len, dtype=torch.int32)
     block_table, slot_mapping = _paged_block_table_slot_mapping(seq_lens)
@@ -265,9 +271,11 @@ def _one_layer(args, layer: int, x: torch.Tensor) -> torch.Tensor:
     return down + h1  # FP32
 
 
-def compute_golden(args: TaskArgsBuilder) -> None:
-    """Fill ``args.out`` (and INOUT k_cache/v_cache) for the 40-layer decode chunk."""
+def compute_golden(args: TaskArgsBuilder, n_layers: int = N_LAYERS) -> None:
+    """Fill ``args.out`` and the INOUT KV caches for a decoder stack."""
+    if n_layers <= 0:
+        raise ValueError(f"n_layers must be positive, got {n_layers}")
     cur = args.hidden_states.float()  # copy_hidden: bf16 input embedded as FP32
-    for layer in range(N_LAYERS):
+    for layer in range(n_layers):
         cur = _one_layer(args, layer, cur)
     args.out[:] = cur.to(torch.bfloat16)  # copy_out: single FP32->bf16 round

simpler_setup/tools/swimlane_converter.py

@@ -119,6 +119,99 @@ def format_task_display(task_id):
     return f"r{ring}t{local}"
 
 
+def _decode_graph_node_task_id(task_id):
+    """Decode Scheduler-owned ring-1 Graph-node ids.
+
+    Graph nodes use ``local=(outer_local << 10) | node_index`` while the
+    stream-visible outer Graph task remains on ring 0.
+    """
+    tid = normalize_pto2_task_id_int(task_id)
+    if tid is None or ((tid >> 32) & 0xFFFFFFFF) != 1:
+        return None
+    local = tid & 0xFFFFFFFF
+    return local >> 10, local & 0x3FF
+
+
+def _collect_graph_execution_instances(tasks, scheduler_phases):  # noqa: PLR0912
+    """Join Graph-node rows to their outer GraphPrepare records."""
+    prepare_by_outer = defaultdict(list)
+    dummy_rows = []
+    for thread_idx, records in enumerate(scheduler_phases or []):
+        for record in records:
+            phase = record.get("phase")
+            if phase == "graph_prepare":
+                outer_task_id = normalize_pto2_task_id_int(record.get("task_id"))
+                if outer_task_id is not None and (outer_task_id >> 32) == 0:
+                    prepare_by_outer[outer_task_id].append(record)
+            elif phase == "dummy_task":
+                dummy_rows.append((record, thread_idx))
+
+    rows_by_outer = defaultdict(list)
+    for task in tasks:
+        decoded = _decode_graph_node_task_id(task.get("task_id"))
+        if decoded is not None:
+            outer_task_id, node_index = decoded
+            rows_by_outer[outer_task_id].append((task, node_index))
+
+    dummy_by_outer = defaultdict(list)
+    for record, thread_idx in dummy_rows:
+        decoded = _decode_graph_node_task_id(record.get("task_id"))
+        if decoded is not None:
+            outer_task_id, node_index = decoded
+            dummy_by_outer[outer_task_id].append((record, node_index, thread_idx))
+
+    instances = []
+    for outer_task_id, prepare_records in prepare_by_outer.items():
+        rows = rows_by_outer.get(outer_task_id, [])
+        aicpu_rows = dummy_by_outer.get(outer_task_id, [])
+        if not rows and not aicpu_rows:
+            continue
+        node_indices = {node_index for _, node_index in rows}
+        node_indices.update(node_index for _, node_index, _ in aicpu_rows)
+        starts = [
+            task.get("dispatch_time_us", _task_slice_start_us(task))
+            if task.get("dispatch_time_us", -1) >= 0
+            else _task_slice_start_us(task)
+            for task, _ in rows
+        ]
+        starts.extend(record["start_time_us"] for record, _, _ in aicpu_rows)
+        ends = [
+            task.get("finish_time_us", 0) if task.get("finish_time_us", 0) > 0 else task["end_time_us"]
+            for task, _ in rows
+        ]
+        ends.extend(record["end_time_us"] for record, _, _ in aicpu_rows)
+        prepare_start_us = min(record["start_time_us"] for record in prepare_records)
+        instances.append(
+            {
+                "outer_task_id": outer_task_id,
+                "rows": rows,
+                "aicpu_rows": aicpu_rows,
+                "visible_node_indices": sorted(node_indices),
+                "execution_start_us": min(starts),
+                "execution_end_us": max(ends),
+                "prepare_start_us": prepare_start_us,
+                "prepare_end_us": max(record["end_time_us"] for record in prepare_records),
+                "prepare_duration_us": sum(
+                    record["end_time_us"] - record["start_time_us"] for record in prepare_records
+                ),
+                "prepare_slice_count": len(prepare_records),
+            }
+        )
+
+    instances.sort(key=lambda instance: instance["prepare_start_us"])
+    lane_finish_us = []
+    for instance_idx, instance in enumerate(instances):
+        start_us = instance["prepare_start_us"]
+        lane_idx = next((idx for idx, finish_us in enumerate(lane_finish_us) if finish_us <= start_us), -1)
+        if lane_idx < 0:
+            lane_idx = len(lane_finish_us)
+            lane_finish_us.append(0.0)
+        lane_finish_us[lane_idx] = instance["execution_end_us"]
+        instance["instance_idx"] = instance_idx
+        instance["lane_idx"] = lane_idx
+    return instances
+
+
 def read_perf_data(filepath):  # noqa: PLR0912, PLR0915
     """Read performance data from a swimlane JSON file.
 
@@ -1113,6 +1206,7 @@ def generate_chrome_trace_json(  # noqa: PLR0912, PLR0913, PLR0915
         core_to_thread: Optional list mapping core_id (index) to scheduler thread index (-1 = unassigned)
 
     Generates processes in the trace:
+        - pid=5 "Graph Execution": one end-to-end envelope per Graph task
         - pid=1 "AICPU Orchestrator": orchestrator phase bars (chip_swimlane_level >= 4)
         - pid=2 "AICPU Scheduler": scheduler phase bars (chip_swimlane_level >= 3)
         - pid=3 "Scheduler View": dispatch_time_us to finish_time_us (AICPU perspective)
@@ -1150,12 +1244,15 @@ def generate_chrome_trace_json(  # noqa: PLR0912, PLR0913, PLR0915
                 if resolved >= 0:
                     task["func_id"] = resolved
 
+    graph_instances = _collect_graph_execution_instances(tasks, scheduler_phases)
+
     # Step 2: Generate JSON events
     events = []
 
     # Metadata event: Process names and sort order.
     # pid is renumbered in pipeline order (top → bottom in Perfetto):
     #   pid=1  AICPU Orchestrator  (submits tasks — earliest)
+    #   pid=5  Graph Execution     (Scheduler-local expansion + execution)
     #   pid=2  AICPU Scheduler     (pops ready, dispatches, completes)
     #   pid=3  Scheduler View      (AICPU-eye view of each worker's dispatch→finish)
     #   pid=4  Worker View         (physical AIC/AIV execution rows)
@@ -1165,6 +1262,51 @@ def generate_chrome_trace_json(  # noqa: PLR0912, PLR0913, PLR0915
         task_map[t["task_id"]].append(t)
     spmd_task_ids = _identify_spmd_task_ids(task_map, deps_block_map)
 
+    if graph_instances:
+        events.append(
+            {"args": {"name": "Graph Execution"}, "cat": "__metadata", "name": "process_name", "ph": "M", "pid": 5}
+        )
+        events.append(
+            {"args": {"sort_index": 1}, "cat": "__metadata", "name": "process_sort_index", "ph": "M", "pid": 5}
+        )
+        for lane_idx in sorted({instance["lane_idx"] for instance in graph_instances}):
+            events.append(
+                {
+                    "args": {"name": f"Graph_{lane_idx}"},
+                    "cat": "__metadata",
+                    "name": "thread_name",
+                    "ph": "M",
+                    "pid": 5,
+                    "tid": 5000 + lane_idx,
+                }
+            )
+        for instance in graph_instances:
+            outer_display = format_task_display(instance["outer_task_id"])
+            node_indices = instance["visible_node_indices"]
+            events.append(
+                {
+                    "args": {
+                        "outer_task_id": instance["outer_task_id"],
+                        "visible_node_count": len(node_indices),
+                        "visible_node_index_min": min(node_indices),
+                        "visible_node_index_max": max(node_indices),
+                        "prepare_slice_count": instance["prepare_slice_count"],
+                        "prepare_duration_us": instance["prepare_duration_us"],
+                        "execution_start_us": instance["execution_start_us"],
+                        "execution_duration_us": instance["execution_end_us"] - instance["execution_start_us"],
+                        "synthetic_id_layout": "ring1:(outer_task_id << 10) | node_index",
+                    },
+                    "cat": "graph_execution",
+                    "cname": "rail_animation",
+                    "name": f"GraphExecution({outer_display}, {len(node_indices)} visible nodes)",
+                    "ph": "X",
+                    "pid": 5,
+                    "tid": 5000 + instance["lane_idx"],
+                    "ts": instance["prepare_start_us"],
+                    "dur": instance["execution_end_us"] - instance["prepare_start_us"],
+                }
+            )
+
     events.append({"args": {"name": "Worker View"}, "cat": "__metadata", "name": "process_name", "ph": "M", "pid": 4})
     events.append({"args": {"sort_index": 4}, "cat": "__metadata", "name": "process_sort_index", "ph": "M", "pid": 4})
 
@@ -1433,6 +1575,7 @@ def sched_lane_tid(thread_idx, lane=0):
             "drain": "cq_build_running",  # handle_drain_mode outer
             "drain_prepare": "cq_build_attempt_runnable",  # inner: cluster scan + build_payload
             "drain_publish": "cq_build_attempt_passed",  # inner: MMIO write_reg per subtask (the cohort launch)
+            "graph_prepare": "rail_animation",  # bounded Scheduler-side Definition expansion
             # Inner phase — nests inside Complete or Dummy via time containment
             "resolve": "vsync_highlight_color",  # on_task_complete: walk consumer list
             # Separate-lane (Worker View AICPU_N) — fallback color if it ever lands on Sched
@@ -1593,6 +1736,7 @@ def _find_containing_complete(thread_idx: int, finish_us: float):
                     "drain",
                     "drain_prepare",
                     "drain_publish",
+                    "graph_prepare",
                 ):
                     continue
                 start_us = record["start_time_us"]

src/a2a3/platform/sim/host/device_runner.cpp

@@ -768,6 +768,7 @@ int DeviceRunner::finalize() {
     release_callable_state();
 
     unload_executor_binaries();
+    release_graph_execution_buffers();
 
     // Release the three per-Worker pooled arenas. Must precede mem_alloc_.finalize()
     // so the arenas free through the still-live allocator, not after it.

src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp

@@ -279,7 +279,6 @@ int32_t AicpuExecutor::run(Runtime *runtime) {
             runtime->set_slot_states_ptr(nullptr);
 
             sched_ctx_.bind_runtime(rt);
-
             // Latch the host-built task count (on_orchestration_done sets total_tasks_)
             // BEFORE the runtime_init_ready_ release below — that store is the barrier
             // that unblocks the scheduler threads. Otherwise they would acquire
@@ -288,13 +287,8 @@ int32_t AicpuExecutor::run(Runtime *runtime) {
             // to keep them alive).
             // NOTE: do NOT call rt_orchestration_done(rt) here. The HOST already
             // called it in run_host_orchestration; the orchestrator's own
-            // task-allocator pointers are intentionally NOT relocated (only the
-            // SM cross-task pointers and the host-built fanout adjacency —
-            // dep_pool / ready queues / fanout_head — were), so they still hold
-            // host addresses and mark_done()'s active_count() read would
-            // dereference host memory and fault the AICPU. on_orchestration_done
-            // only needs total_tasks and the scalar
-            // orchestrator.inline_completed_tasks, both already valid.
+            // task-allocator pointers are intentionally NOT relocated, so they
+            // still hold host addresses and mark_done() would fault the AICPU.
             sched_ctx_.on_orchestration_done(runtime, rt, thread_idx, runtime->host_total_tasks);
             LOG_INFO("Thread %d: host-orch boot complete (%d tasks)", thread_idx, runtime->host_total_tasks);
         }
@@ -309,8 +303,7 @@ int32_t AicpuExecutor::run(Runtime *runtime) {
     // orchestration setup, seeds its disjoint slice of the whole graph's ready
     // set + wake lists, then barriers. Only once all slices are done does the
     // leader publish runtime_init_ready_, so no thread dispatches against a
-    // half-seeded graph. This replaces the O(total_tasks) serial classify the
-    // leader used to run alone while the others idle-waited.
+    // half-seeded graph.
     while (!classify_ready_.load(std::memory_order_acquire)) {
         SPIN_WAIT_HINT();
     }
@@ -364,6 +357,9 @@ int32_t AicpuExecutor::run(Runtime *runtime) {
         if (rt != nullptr) {
             // Clear g_current_runtime in this DSO before destroying rt.
             framework_bind_runtime(nullptr);
+            // Graph execution blocks are host-owned GM retained by the
+            // DeviceRunner. This run only drops its submission references;
+            // the blocks remain reusable until Worker finalization.
             runtime_destroy(rt, runtime_arena_);
             rt = nullptr;
         }