diff --git a/simpler_setup/goldens/qwen3_14b_decode.py b/simpler_setup/goldens/qwen3_14b_decode.py index b56c94a57a..ca0079aa76 100644 --- a/simpler_setup/goldens/qwen3_14b_decode.py +++ b/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 diff --git a/simpler_setup/tools/swimlane_converter.py b/simpler_setup/tools/swimlane_converter.py index 540252db66..969549bcdb 100644 --- a/simpler_setup/tools/swimlane_converter.py +++ b/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. @@ -138,6 +231,8 @@ def read_perf_data(filepath): # noqa: PLR0912, PLR0915 "aicpu_tasks": [[core_id, reg_task_id, dispatch_cycles, finish_cycles], ...], "aicpu_scheduler_phases": [ [ {kind, start_cycles, end_cycles, ...}, ... ], ... ], "aicpu_orchestrator_phases": [ [ {submit_idx, task_id, start_cycles, end_cycles}, ... ], ... ] + "host_orchestrator": {start_cycles, end_cycles, + first_publish_cycles, records: [...]} } aicore_tasks columns (v3 schema): the trailing receive_to_start_cycles @@ -191,6 +286,7 @@ def read_perf_data(filepath): # noqa: PLR0912, PLR0915 aicpu_rows = data.get("aicpu_tasks") or [] sched_phases_raw = data.get("aicpu_scheduler_phases") or [] orch_phases_raw = data.get("aicpu_orchestrator_phases") or [] + host_orchestrator_raw = data.get("host_orchestrator") or {} # AICore lookup keyed by (core_id, reg_task_id). Two dispatches of the # same PTO2 task_token_raw to the same core (SPMD over-subscription, MIX @@ -370,6 +466,37 @@ def _phase_us(pr): converted.append(out) aicpu_orchestrator_phases.append(converted) + host_orchestrator = None + host_start_cycles = int(host_orchestrator_raw.get("start_cycles", 0)) + host_end_cycles = int(host_orchestrator_raw.get("end_cycles", 0)) + first_publish_cycles = int(host_orchestrator_raw.get("first_publish_cycles", 0)) + if host_end_cycles > host_start_cycles: + # Host CLOCK_MONOTONIC and AICPU syscnt do not share an epoch. Streaming + # host-build-graph provides a causal cross-clock anchor: the scheduler's + # first device record cannot precede the first released task prefix, so + # align that publication to device t=0. Archived batch-mode captures do + # not carry the anchor and retain the old end-to-zero fallback. + has_publish_anchor = host_start_cycles <= first_publish_cycles <= host_end_cycles + host_anchor_cycles = first_publish_cycles if has_publish_anchor else host_end_cycles + clock_alignment = ( + "host_first_publish_aligned_to_first_device_event" + if has_publish_anchor + else "host_orch_end_aligned_to_device_zero" + ) + host_orchestrator = { + "start_time_us": (host_start_cycles - host_anchor_cycles) * cycles_to_us_factor, + "end_time_us": (host_end_cycles - host_anchor_cycles) * cycles_to_us_factor, + "clock_alignment": clock_alignment, + "records": [], + } + for record in host_orchestrator_raw.get("records", []): + converted = dict(record) + converted["start_time_us"] = (int(record["start_cycles"]) - host_anchor_cycles) * cycles_to_us_factor + converted["end_time_us"] = (int(record["end_cycles"]) - host_anchor_cycles) * cycles_to_us_factor + converted.pop("start_cycles", None) + converted.pop("end_cycles", None) + host_orchestrator["records"].append(converted) + out = { "l2_swimlane_level": level, "tasks": tasks, @@ -378,6 +505,8 @@ def _phase_us(pr): out["aicpu_scheduler_phases"] = aicpu_scheduler_phases if aicpu_orchestrator_phases: out["aicpu_orchestrator_phases"] = aicpu_orchestrator_phases + if host_orchestrator is not None: + out["host_orchestrator"] = host_orchestrator if core_to_thread: out["core_to_thread"] = core_to_thread return out @@ -1096,6 +1225,7 @@ def generate_chrome_trace_json( # noqa: PLR0912, PLR0913, PLR0915 deps_kernel_map=None, deps_block_map=None, emit_overhead=False, + host_orchestrator=None, ): """Generate Chrome Trace Event Format JSON from task data. @@ -1113,7 +1243,8 @@ 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=1 "AICPU Orchestrator": orchestrator phase bars (l2_swimlane_level >= 4) + - pid=5 "Graph Execution": one end-to-end envelope per Graph task + - pid=1 "Host Orchestrator" or "AICPU Orchestrator": orchestration bars - pid=2 "AICPU Scheduler": scheduler phase bars (l2_swimlane_level >= 3) - pid=3 "Scheduler View": dispatch_time_us to finish_time_us (AICPU perspective) - pid=4 "Worker View": per-subtask kernel execution on physical cores @@ -1150,12 +1281,16 @@ 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) + graph_outer_task_ids = {instance["outer_task_id"] for instance in graph_instances} + # 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 +1300,103 @@ 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 host_orchestrator: + events.append( + {"args": {"name": "Host Orchestrator"}, "cat": "__metadata", "name": "process_name", "ph": "M", "pid": 1} + ) + events.append( + {"args": {"sort_index": 0}, "cat": "__metadata", "name": "process_sort_index", "ph": "M", "pid": 1} + ) + events.append( + { + "args": {"name": "Host_Orch"}, + "cat": "__metadata", + "name": "thread_name", + "ph": "M", + "pid": 1, + "tid": 4000, + } + ) + events.append( + { + "args": {"clock_alignment": host_orchestrator["clock_alignment"]}, + "cat": "host_orchestrator", + "cname": "rail_animation", + "name": "host_orchestration", + "ph": "X", + "pid": 1, + "tid": 4000, + "ts": host_orchestrator["start_time_us"], + "dur": host_orchestrator["end_time_us"] - host_orchestrator["start_time_us"], + } + ) + for record in host_orchestrator["records"]: + task_id = normalize_pto2_task_id_int(record.get("task_id")) + submit_kind = "graph_submit" if task_id in graph_outer_task_ids else "task_submit" + events.append( + { + "args": { + "submit_idx": record.get("submit_idx", 0), + "task_id": task_id, + "submit_kind": submit_kind, + "clock_alignment": host_orchestrator["clock_alignment"], + }, + "cat": "host_orchestrator", + "cname": "rail_animation", + "name": f"{submit_kind}({format_task_display(task_id)})", + "ph": "X", + "pid": 1, + "tid": 4000, + "ts": record["start_time_us"], + "dur": record["end_time_us"] - record["start_time_us"], + } + ) + + 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 +1665,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 @@ -1594,6 +1827,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"] @@ -1713,7 +1947,7 @@ def _find_containing_complete(thread_idx: int, finish_us: float): # orchestrator timing. There is no separate aggregate summary — the # device-side LOG_INFO "orch_start=… orch_end=… orch_cost=…" log # line covers the run-window envelope for debugging without swimlane. - if orchestrator_phases: + if orchestrator_phases and not host_orchestrator: # Process metadata orch_process_label = f"AICPU {orchestrator_name}" if orchestrator_name else "AICPU Orchestrator" events.append( @@ -2438,6 +2672,20 @@ def _find_containing_complete(thread_idx: int, finish_us: float): if verbose: print(f" Overhead Analysis: {sum(1 for e in oh if e.get('ph') == 'C')} counter points (8 tracks)") + # Host orchestration happens before device execution, but its monotonic + # clock has no common epoch with AICPU syscnt. read_perf_data() represents + # that ordering by aligning Host Orch end to device t=0, which leaves the + # host slice at a negative timestamp. Perfetto's default viewport starts at + # zero and therefore hides that otherwise valid slice. Shift every timed + # event by the Host Orch duration: Host Orch becomes [0, duration] and the + # device timeline starts at duration, preserving all relative ordering. + if host_orchestrator: + trace_origin_shift_us = max(0.0, -float(host_orchestrator["start_time_us"])) + if trace_origin_shift_us > 0.0: + for event in events: + if "ts" in event: + event["ts"] = float(event["ts"]) + trace_origin_shift_us + with open(output_path, "w") as f: json.dump({"traceEvents": events}, f, indent=2) @@ -2664,6 +2912,7 @@ def main(): deps_kernel_map=deps_kernel_map, deps_block_map=deps_block_map, emit_overhead=args.overhead, + host_orchestrator=data.get("host_orchestrator"), ) if args.overhead and deps_edges is None: print( diff --git a/src/a2a3/platform/onboard/host/device_runner.cpp b/src/a2a3/platform/onboard/host/device_runner.cpp index 8409eca5fe..6fc8f33890 100644 --- a/src/a2a3/platform/onboard/host/device_runner.cpp +++ b/src/a2a3/platform/onboard/host/device_runner.cpp @@ -509,8 +509,9 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { rc = launch_run(runtime, num_aicore, launch_aicpu_num, selected_pipeline_slot); if (rc != 0) return rc; + runtime.notify_deferred_host_execution_started(); - rc = reap_run(selected_pipeline_slot); + rc = reap_run(runtime, selected_pipeline_slot); if (rc != 0) return rc; // The run owns its AICore stream, so a destroy this run cannot complete is @@ -614,7 +615,7 @@ int DeviceRunner::launch_run(Runtime &runtime, int num_aicore, int launch_aicpu_ return 0; } -int DeviceRunner::reap_run(unsigned slot) { +int DeviceRunner::reap_run(Runtime &runtime, unsigned slot) { if (!run_stream_slots_.ready(slot)) { LOG_ERROR("reap_run: invalid stream set %u", slot); return -1; @@ -641,6 +642,16 @@ int DeviceRunner::reap_run(unsigned slot) { read_device_wall_ns(); + // The host orchestration can run concurrently with the device. Device + // completion implies it has published EOS, so its profile snapshot is now + // immutable and safe to hand to the collector before export. + if (enable_l2_swimlane_ && l2_swimlane_level_ >= L2SwimlaneLevel::ORCH_PHASES) { + l2_swimlane_collector_.set_host_orch_records( + runtime.get_host_orch_phase_records(), runtime.get_host_orch_start_cycles(), + runtime.get_host_orch_end_cycles(), runtime.get_host_orch_first_publish_cycles() + ); + } + // Tear down collectors. stop() joins mgmt then collector in the only safe // order (mgmt's final-drain pass into L2 has poll as its consumer). teardown_shared_collectors_after_run(); diff --git a/src/a2a3/platform/onboard/host/device_runner.h b/src/a2a3/platform/onboard/host/device_runner.h index dfd35984cf..0ee12f4801 100644 --- a/src/a2a3/platform/onboard/host/device_runner.h +++ b/src/a2a3/platform/onboard/host/device_runner.h @@ -269,7 +269,7 @@ class DeviceRunner : public DeviceRunnerBase { // The kernel submission boundary is separate from the stream wait and the // post-run teardown; run() invokes the two back-to-back. int launch_run(Runtime &runtime, int num_aicore, int launch_aicpu_num, unsigned slot); - int reap_run(unsigned slot); + int reap_run(Runtime &runtime, unsigned slot); // On an AICore launch/sync error, best-effort drain the device so a later // run() on the same DeviceRunner can recover in place; if the drain itself diff --git a/src/a2a3/platform/sim/host/device_runner.cpp b/src/a2a3/platform/sim/host/device_runner.cpp index dcd4822e1b..b8808315e4 100644 --- a/src/a2a3/platform/sim/host/device_runner.cpp +++ b/src/a2a3/platform/sim/host/device_runner.cpp @@ -556,6 +556,7 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { // Both simulated kernel thread groups now exist. This is the sim's real // launch boundary: publish before joining either group. publish_task_accepted(); + runtime.notify_deferred_host_execution_started(); for (auto &t : aicpu_threads) { t.join(); @@ -608,6 +609,13 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { return runtime_rc; } + if (enable_l2_swimlane_ && l2_swimlane_level_ >= L2SwimlaneLevel::ORCH_PHASES) { + l2_swimlane_collector_.set_host_orch_records( + runtime.get_host_orch_phase_records(), runtime.get_host_orch_start_cycles(), + runtime.get_host_orch_end_cycles(), runtime.get_host_orch_first_publish_cycles() + ); + } + // Tear down collectors. stop() joins mgmt then collector in the only safe // order (mgmt's final-drain pass into L2 has poll as its consumer). if (enable_l2_swimlane_) { diff --git a/src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp b/src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp index ccc76b6b4c..3e4291d08b 100644 --- a/src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp +++ b/src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp @@ -29,6 +29,7 @@ // Runtime headers (full struct definition for create/destroy + PTO2_SCOPE) #include "pto_runtime2.h" #include "pto_runtime2_types.h" +#include "graph_execution.h" #include "pto_shared_memory.h" // Performance profiling headers @@ -94,14 +95,6 @@ struct AicpuExecutor { std::atomic hs_arrived_{0}; std::atomic hs_thread_seq_{0}; - // Parallel-boot-classify coordination (see AicpuExecutor::run). classify_ready_ - // is published by the boot leader once its leader-only orchestration setup is - // visible; classify_arrived_ is the barrier counting threads that finished - // their slice of the initial classify. Both are one-shot per run and reset in - // deinit(). - std::atomic classify_ready_{false}; - std::atomic classify_arrived_{0}; - int32_t aicpu_thread_num_{0}; // ===== Task queue state (managed by scheduler ready queues) ===== @@ -226,16 +219,10 @@ int32_t AicpuExecutor::run(Runtime *runtime) { } int32_t run_rc = 0; - // Boot: the last AICPU thread (aicpu_thread_num_ - 1) performs the one-time - // host-orch attach. host_build_graph's orchestrator already ran on the host, - // which also relocated every cross-task pointer to its final device address - // before H2D — so the SM/arena this thread sees are already fully - // device-addressed. This thread attaches the prebuilt arena, points the SM - // handle's ring-header pointers at the device SM WITHOUT resetting the - // host-populated data, hands the host-computed task count to the scheduler, - // and releases the other threads. It then falls through and schedules its own - // cores like every other thread — host_build_graph has no device-side - // orchestrator, so there is no orch/sched split. + // Boot: the last AICPU thread performs the one-time host-orch attach. The + // device image starts empty; the host publishes device-addressed task + // prefixes concurrently after launch. Attach without resetting the SM, + // then release every scheduler thread to classify published prefixes. if (thread_idx == aicpu_thread_num_ - 1) { void *prebuilt_arena = runtime->get_prebuilt_arena_base(); size_t off_runtime = runtime->get_prebuilt_runtime_offset(); @@ -279,55 +266,21 @@ 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 - // runtime_init_ready_ with total_tasks_=0 and race to an early exit before - // the host task count is visible (host-orch has no concurrent orchestrator - // 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. - 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); + // The host publishes immutable task prefixes while these scheduler + // threads run. Start with an empty classifier cursor; the dispatch + // loop acquires current_task_index before consuming each prefix. + sched_ctx_.on_host_orchestration_stream_start(runtime, rt, thread_idx); + LOG_INFO("Thread %d: host-orch streaming boot complete", thread_idx); } - // Publish "leader setup done" (SM attached, task count latched, queues - // allocated). Every thread then classifies its slice below before any of - // them may dispatch — the leader holds runtime_init_ready_ until then. - classify_ready_.store(true, std::memory_order_release); + // Publish the completed attach even on boot failure so peer threads do + // not spin forever; rt remains null and they skip dispatch in that case. + runtime_init_ready_.store(true, std::memory_order_release); } - // Parallel initial classify. Every AICPU thread waits for the leader's - // 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. - while (!classify_ready_.load(std::memory_order_acquire)) { + while (!runtime_init_ready_.load(std::memory_order_acquire)) { SPIN_WAIT_HINT(); } - if (!sched_ctx_.is_completed() && rt != nullptr) { - sched_ctx_.classify_partition(thread_idx, aicpu_thread_num_); - } - classify_arrived_.fetch_add(1, std::memory_order_acq_rel); - if (thread_idx == aicpu_thread_num_ - 1) { - while (classify_arrived_.load(std::memory_order_acquire) < aicpu_thread_num_) { - SPIN_WAIT_HINT(); - } - runtime_init_ready_.store(true, std::memory_order_release); - } else { - while (!runtime_init_ready_.load(std::memory_order_acquire)) { - SPIN_WAIT_HINT(); - } - } // Every AICPU thread schedules its assigned cores. if (!sched_ctx_.is_completed()) { @@ -367,6 +320,10 @@ int32_t AicpuExecutor::run(Runtime *runtime) { if (rt != nullptr) { // Clear g_current_runtime in this DSO before destroying rt. framework_bind_runtime(nullptr); + // Graph nodes are expanded into AICPU-local storage. Reclaim all + // completed executions into the bounded graph-affine pool before + // the HBM submission descriptors for this run are released. + graph_execution_collect_retired(); runtime_destroy(rt, runtime_arena_); rt = nullptr; } @@ -399,8 +356,6 @@ void AicpuExecutor::deinit(Runtime *runtime) { hs_setup_done_.store(false, std::memory_order_release); hs_arrived_.store(0, std::memory_order_release); hs_thread_seq_.store(0, std::memory_order_release); - classify_ready_.store(false, std::memory_order_release); - classify_arrived_.store(0, std::memory_order_release); thread_idx_.store(0, std::memory_order_release); finished_.store(false, std::memory_order_release); diff --git a/src/a2a3/runtime/host_build_graph/docs/GRAPH_EXECUTION.md b/src/a2a3/runtime/host_build_graph/docs/GRAPH_EXECUTION.md new file mode 100644 index 0000000000..6d1f8f920c --- /dev/null +++ b/src/a2a3/runtime/host_build_graph/docs/GRAPH_EXECUTION.md @@ -0,0 +1,332 @@ +# Graph Execution + +Graph Execution is available only in the `host_build_graph` runtime. A Graph is +a composite incore task: it is submitted and completed once like an AIC, AIV, +MIX, or SPMD task, but contains a recorded task DAG. + +The first invocation executes normally and records the DAG. A later invocation +places one `GRAPH` task in the host task window. The device Scheduler expands +the saved topology and dispatches its internal nodes; the Host Orchestrator does +not submit those nodes again. + +## Step-1 API + +A Graph uses `L0TaskArgs`, the existing incore argument type: + +```cpp +void graph_function(const L0TaskArgs &args, int variant) { + const Tensor &input = args.tensor(0).ref(); + const Tensor &weight = args.tensor(1).ref(); + const Tensor &output = args.tensor(2).ref(); + + const std::array shape{input.shapes[0]}; + TensorCreateInfo intermediate( + shape.data(), static_cast(shape.size()), input.dtype + ); + + L0TaskArgs matmul_args; + matmul_args.add_input(input, weight); + matmul_args.add_output(intermediate); + matmul_args.add_scalar(uint32_t{16}); // fixed Definition data + TaskOutputTensors matmul = rt_submit_aic_task( + variant == 0 ? FUNC_MATMUL : FUNC_MATMUL_TRANSPOSED, + matmul_args + ); + + L0TaskArgs activation_args; + activation_args.add_input(matmul.get_ref(0)); + activation_args.add_output(output); + rt_submit_aiv_task(FUNC_ACTIVATION, activation_args); +} + +void submit_layer(const L0TaskArgs &args) { + rt_submit_graph(&graph_function, args, /*variant=*/0); +} +``` + +The function pointer is the default Graph identity. Trailing integral, +`float`, `double`, and `bool` construction parameters are forwarded to the +Graph function and hashed by value into the cache key. They are separate from +execution scalars in `L0TaskArgs`: changing a construction parameter selects a +different Definition rather than patching an existing one. + +An explicit identity is available for call sites that need a stable name: + +```cpp +rt_submit_graph( + GRAPH_KEY("qwen_decoder_layer_v1"), + &graph_function, + args, + /*variant=*/0 +); +``` + +There are no public `GraphArgs`, `GraphBindings`, `Patch`, or `ScalarRef` +types. The boundary is represented by `L0TaskArgs`. + +## Supported dynamic and static data + +Step 1 deliberately supports a narrow, safe contract: + +- Boundary Tensor addresses may change for every invocation. +- Construction parameters are part of Graph identity and may control the + function's task count, kernel selection, or other structural choices. +- Boundary Tensor shape, stride, dtype, size, direction, contiguity, and alias + partition must match the first invocation. +- Scalars inside internal task args are fixed Definition data. +- Scalars in the boundary `L0TaskArgs` are not cacheable yet. Such a call uses + the ordinary task-submit path. +- Boundary storage is caller-owned. `INPUT`, `INOUT`, `OUTPUT_EXISTING`, and + `NO_DEP` are supported. A boundary `TensorCreateInfo` tagged `OUTPUT` is not. +- Early-resolve hints apply while recording the first invocation. Replayed + internal nodes use the saved completion topology without the hint. +- A recorded task may depend on a Graph-external producer when that producer + is the creator of a boundary Tensor. The outer Graph owns that dependency on + replay; arbitrary cross-boundary explicit dependencies remain unsupported. + +Structural or alias mismatch logs a warning and executes the Graph function +normally for that invocation. It never reuses heap offsets recorded for a +different shape. Debug builds also assert at these unsupported boundaries so +development catches a violated fixed-shape contract immediately; the ordinary +path remains the defensive release-build behavior. + +## Qwen decoder-layer example + +The upper layer packages all Tensor I/O in `L0TaskArgs`; the wrapper has no +separate `hidden`, `weight`, or `output` parameters: + +```cpp +void qwen_decoder_layer(const L0TaskArgs &args) { + const Tensor &hidden = args.tensor(0).ref(); + const Tensor &attention_weight = args.tensor(1).ref(); + const Tensor &mlp_weight = args.tensor(2).ref(); + const Tensor &output = args.tensor(3).ref(); + + const std::array hidden_shape{hidden.shapes[0]}; + TensorCreateInfo attention_out( + hidden_shape.data(), static_cast(hidden_shape.size()), hidden.dtype + ); + + L0TaskArgs attention_args; + attention_args.add_input(hidden, attention_weight); + attention_args.add_output(attention_out); + attention_args.add_scalar(uint32_t{16}); // fixed model configuration + TaskOutputTensors attention = + rt_submit_aic_task(FUNC_ATTENTION, attention_args); + + MixedKernels mlp; + mlp.aic_kernel_id = FUNC_MLP_AIC; + mlp.aiv0_kernel_id = FUNC_MLP_AIV; + + L0TaskArgs mlp_args; + mlp_args.add_input(attention.get_ref(0), mlp_weight); + mlp_args.add_output(output); + rt_submit_task(mlp, mlp_args); +} + +void submit_qwen_decoder_layer(const L0TaskArgs &args) { + rt_submit_graph(&qwen_decoder_layer, args); +} + +void decode_three_layers( + const std::array &hidden, + const std::array &attention_weight, + const std::array &mlp_weight, + const std::array &output +) { + for (std::size_t layer = 0; layer < hidden.size(); ++layer) { + L0TaskArgs args; + args.add_input( + hidden[layer], + attention_weight[layer], + mlp_weight[layer] + ); + args.add_output(output[layer]); + submit_qwen_decoder_layer(args); + } +} +``` + +The first layer records ordinary task submissions. Layers two and three submit +one Graph task each when their Tensor metadata matches. A per-layer or +per-token scalar is not dynamic in step 1; use ordinary submission or a +different fixed Graph function/key until dynamic scalar support is added. + +## Definition + +Recording uses host-only C++ state: + +- `std::vector` for nodes, tensors, scalars, fanins, and pending uploads; +- `std::unordered_map` for the per-run Definition cache; +- `std::unique_ptr` for the active recording. + +The cache stores at most 16 Definitions and allocates each entry to its actual +serialized size. No fixed maximum-size recording array is copied on a cache +hit. + +At `graph_end`, recording is compacted into one contiguous, pointer-free POD +Definition. It contains: + +- node order and AIC/AIV/MIX/SPMD kernel metadata; +- `root_indices` plus both directions of the immutable topology: + fanin CSR and fanout CSR; +- one packed-heap offset per node; +- each node's Tensor source: + `BOUNDARY_EXACT`, `BOUNDARY_VIEW`, `INTERNAL`, or `OWN_OUTPUT`; +- fixed scalar values; +- fixed boundary signatures and alias representatives. + +The header also carries a content hash of the complete Definition image. The +device execution pool requires this hash, the Graph key, and the node count to +all match before reusing a resident Definition. A new run may record different +metadata under the same function identity, so key-only reuse is not safe. + +All references are 32-bit offsets from the Definition base. Cross-boundary +Tensors use the fixed-width `GraphTensor` wire POD rather than the +64-byte-aligned C++ `Tensor` object. The upload is therefore one contiguous +copy with no raw Host pointers and no relocation pass. + +Before materialization, the Scheduler recomputes the Definition content hash +and validates section ranges, topology indices, node heap offsets, the outer +heap extent, Tensor metadata, and Tensor-source bounds. Invalid wire data is +rejected before an offset participates in pointer arithmetic. + +There is no cache schema version. The cache is per run and starts empty, so a +persistent-format version would currently have no effect. + +## Cache hit and memory + +For a cache hit, the Host Orchestrator: + +1. validates the fixed boundary contract; +2. reserves one task-window slot; +3. reserves one heap block large enough for every internal intermediate; +4. computes only external fanin and boundary tensormap effects; +5. emits one outer `GRAPH` task; +6. uploads the exact-size POD submission image. + +Internal nodes consume no ring task-window slots. Their descriptor, payload, +and slot state are built in an AICPU-local execution block. Active and pooled +blocks share one hard budget of 16 MiB and 64 blocks. The pool prefers a block +last used by the same Graph key and Definition content hash; an exact match can +skip the compact Definition copy. + +## Scheduler flow + +Host orchestration is streamed rather than built as one batch. Device +execution launches against an empty shared-memory image, and every completed +`rt_submit_graph` call is a commit boundary. The Host publishes a newly +committed prefix in this order: + +1. exact-size Graph POD images; +2. task descriptors and payloads; +3. relocated slot-state POD bytes and completion bytes; +4. a release update of the 32-bit `current_task_index`. + +Scheduler thread 0 acquires that index and classifies every newly visible task +exactly once. Other scheduler threads may immediately dispatch the resulting +ready work while the Host continues later orchestration calls. For the +three-layer Qwen example, layer 0's ordinary tasks therefore execute while +layers 1 and 2 are still submitting their outer Graph tasks. The Host publishes +`orchestrator_done` only as an end-of-stream marker; the Scheduler exits after +EOS, after every published task has been classified, and after all published +tasks have completed. + +A Graph is placed in two independent control flows: + +- `graph_prepare_queue`: materialize the saved nodes even while external fanin + is still pending; +- `graph_ready_queue`: signal that the outer Graph's external fanin is ready. + +Core-owning Scheduler threads pop at most one item from each queue per loop. A +prepare call expands at most four nodes and requeues unfinished work, +interleaving Graph expansion with normal scheduling. + +Preparation and external readiness set two bits in one atomic activation gate. +Whichever operation sets the second bit activates the saved root nodes exactly +once. + +Internal dependency readiness borrows the completion-state polling idea, but +dependency wiring remains an Orchestrator responsibility: + +- recording constructs both fanin and fanout CSR in the immutable Definition; +- materialization patches runnable node state and registers each non-root on + one producer selected from its saved fanin CSR; +- a node's release/acquire `task_state` is its Graph-local completion flag, so + internal nodes need neither ring completion flags nor task-window slots; +- producer completion closes and drains only its current wake-list rather than + traversing the saved fanout CSR; +- a woken consumer scans its saved fanin CSR and either enters its shape queue + or registers on the next incomplete producer; +- `WAKE_LIST_SENTINEL` closes the completion/registration race: a failed + registration observes completion and immediately rescans. + +The runtime wake-list registration is a transient polling subscription, not +dependency discovery or Graph rewiring. Fanout CSR remains in the Definition +as part of the complete recorded topology and for DFX, but readiness does not +walk it. + +```text +outer GRAPH + -> activate root_indices[] + -> producer completion drains its current wake-list + -> each waiter polls saved fanin completion state + -> ready waiter enters its ordinary shape queue + or registers on another incomplete producer + -> final internal completion completes the outer GRAPH +``` + +Internal nodes count as zero submitted stream tasks. The final node completes +the one outer Graph task, publishes the outer ring completion flag, wakes +external consumers, and contributes one to the host-visible completion count. + +Localization or materialization failure is fail-fast: the Scheduler latches an +error instead of leaving an already-submitted outer Graph unable to complete. + +## Current unsupported cases + +These cases assert in debug builds and execute through the ordinary path in a +release build: + +- dynamic boundary scalars; +- variable Tensor shape or metadata; +- changed boundary aliasing; +- runtime-allocated boundary outputs; +- nested Graph recording; +- dispatch predicates; +- cross-boundary explicit dependencies that are not represented by a boundary + Tensor's creator; +- an unclassifiable internal Tensor source; +- more than 16 Definitions, 1024 internal nodes, or 32 boundary Tensors; +- insufficient task-window or heap capacity detected before outer submission. + +An AICPU execution-pool or materialization failure happens after the outer +Graph has already been submitted. It therefore latches a Scheduler fatal error +instead of falling back; leaving the outer task pending would otherwise wedge +completion. + +Explicit dependencies between recorded internal nodes are preserved when they +are otherwise supported; ordinary Tensor dependencies are always preserved. + +## DFX + +With L2 swimlane level 4: + +- `Host Orchestrator` shows cache-miss `task_submit` records and one + `graph_submit` record for every cache hit; +- `Graph Execution` spans an outer Graph execution; +- `AICPU Scheduler` shows bounded `graph_prepare` slices separately from normal + dispatch; +- existing Scheduler and Worker lanes show the expanded internal tasks. + +Host `CLOCK_MONOTONIC` and device syscnt have different epochs. Streaming +captures record the first released task prefix as a causal cross-clock anchor; +the converter aligns that publication with the first device event. This makes +the overlap with later Host Graph submissions visible without assuming that +Host orchestration ended before device scheduling began. + +The scene coverage under `tests/st/a2a3/host_build_graph/graph_execution` +includes an AIV fanin/fanout DAG, a Qwen-style AIV/AIC decoder-layer DAG, a +three-slot multi-block MIX/SPMD Graph, and a manual Qwen3-14B three-layer +decode. Every scene invokes the same fixed Graph three times: one recording +execution followed by two outer-Graph submissions. diff --git a/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md b/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md index cd10a5409f..49a92dc7b1 100644 --- a/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md +++ b/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md @@ -74,13 +74,14 @@ Two runtime backends exist under `src/runtime/`, each representing a different o ### 1.1 host_build_graph The host-orchestration variant of `tensormap_and_ringbuffer`: it shares the same -ring-buffer task storage, GM heap, and TensorMap dependency tracking, and runs -the orchestration SO **on the host CPU** to build the complete task graph before -launching device execution. The device then boots scheduler-only. +ring-buffer task storage, GM heap, and TensorMap dependency tracking. Device +execution launches against an empty image, then the orchestration SO runs on a +host thread and publishes immutable committed task prefixes while the AICPU +scheduler consumes them. The device remains scheduler-only. - **Task storage**: `PTO2TaskDescriptor[]` in shared memory ring buffer (same as 1.2) - **Dependencies**: automatically derived from tensor read/write patterns via TensorMap (same as 1.2) -- **Scheduling**: AICPU attaches the host-populated, already-device-addressed SM and dispatches the pre-built graph +- **Scheduling**: AICPU attaches an empty device-addressed SM, incrementally classifies each published prefix, and dispatches it without waiting for orchestration EOS - **Use case**: host-side graph construction; device runs no orchestrator thread ### 1.2 tensormap_and_ringbuffer (PTO2) @@ -91,8 +92,8 @@ The primary production runtime. Uses ring buffers for task slots and output memo - **Memory**: GM Heap ring for output buffer allocation - **Dependencies**: automatically derived from tensor read/write patterns via TensorMap - **Thread model**: 3 scheduler threads + 1 orchestrator thread on AICPU -- **Single ring**: host_build_graph builds the whole graph on the host with no - execution-time reclaim, so HeapRing and TaskRing are single +- **Single ring**: host_build_graph incrementally publishes the graph from the + host with no execution-time reclaim, so HeapRing and TaskRing are single whole-graph-resident instances (`PTO2_MAX_RING_DEPTH == 1`); all scope depths map to ring 0. - **Use case**: production workloads; supports streaming, flow control, and large batch sizes @@ -588,9 +589,11 @@ Each scheduler thread runs a tight loop with two main phases: its **first unmet** producer's wake list (`register_wake`); that producer's completion re-drives the classification. The decision is terminal — tasks are never re-polled — because `completion_flags` are monotonic. This wake machinery -is seeded by the device **boot classify** (`on_orchestration_done`), which scans -the submitted tasks once and either pushes the fanin-free ones to the ready -queue or registers each remaining task on its first unmet producer. +is seeded incrementally by `classify_published_tasks`. Thread S0 acquires +`current_task_index`, classifies every newly committed task exactly once, and +either routes a fanin-free task or registers it on its first unmet producer. +The release publication of `current_task_index` is the host→device visibility +edge for the corresponding descriptor/payload prefix. **Early staging status.** Early producer propagation is currently disabled in HBG: `propagate_dispatch_fanin` is a stub, so the polling path does not populate @@ -635,9 +638,9 @@ gates on `completed_watermark >= producer.last_consumer_local_id` to observe Slot reclaim is inert: host_build_graph is whole-graph-resident, so `last_task_alive` is never advanced at runtime and there is no -`advance_ring_pointers` step. `reset_for_reuse()` runs **once at init** -(`pto_shared_memory.cpp`) to zero each slot before the host orchestrator -populates it — it is not a runtime recycle hook. +`advance_ring_pointers` step. `reset_for_reuse()` runs once while the host +mirror is initialized, before the host orchestrator populates and publishes a +slot — it is not a runtime recycle hook. ### 8.5 SchedulerContext @@ -647,11 +650,12 @@ Public surface (called from `AicpuExecutor::init/run/deinit`): | Method | Phase | Purpose | | ------ | ----- | ------- | -| `init(runtime, aicpu_thread_num, regs_base)` | once per run | Handshake + assign cores, reset counters, latch `regs_base`, bind `func_id_to_addr_` | +| `pre_handshake_init` / `handshake_partition` / `post_handshake_init` | once per run | Reset state, handshake disjoint core slices, then assign cores and initialize profiling | | `bind_runtime(rt)` | boot thread | Wire `sched_` to `rt->scheduler` once the boot thread attaches the host-built `rt` | | `resolve_and_dispatch(runtime, thread_idx)` | per scheduler thread | Main dispatch loop | +| `run_resolution_thread(runtime, thread_idx)` | P thread | Resolve completed tasks and publish dependency wakeups | | `shutdown(thread_idx)` | per thread on exit | `platform_deinit_aicore_regs` for this thread's cores; PMU finalize when enabled | -| `on_orchestration_done(runtime, rt, thread_idx, total_tasks)` | boot thread | Publish core assignments, latch task count, fold inline-completed tasks, flip `orchestrator_done_` (or `emergency_shutdown` on fatal) | +| `on_host_orchestration_stream_start(runtime, rt, thread_idx)` | boot thread | Initialize streaming counters, resolution queues, and core-assignment DFX before the first prefix arrives | | `deinit()` | once per run | Reset every scheduler-owned field to its post-construction default | | Read-only accessors | various | `aic_count()` / `aiv_count()` / `is_completed()` / `completed_tasks_count()` | @@ -659,9 +663,11 @@ Private internals are split across three .cpp files by responsibility: - `scheduler_completion.cpp` — completion polling, drain protocol - `scheduler_dispatch.cpp` — task dispatch loop and helpers -- `scheduler_cold_path.cpp` — exit checks, stall diagnostics, profiling, lifecycle (`init/deinit`), core management (`handshake_all_cores` / `assign_cores_to_threads` / `reassign_cores_for_all_threads` / `emergency_shutdown`), and `on_orchestration_done` +- `scheduler_cold_path.cpp` — exit checks, stall diagnostics, profiling, lifecycle, core management, streaming bootstrap, and published-prefix classification -`AicpuExecutor` calls neither `handshake_*`, `assign_*`, `reassign_*`, nor `emergency_shutdown` directly — they are private, invoked only by `init` and `on_orchestration_done`. +`AicpuExecutor` owns the cross-thread initialization barriers; scheduler state, +prefix classification, completion resolution, and emergency shutdown stay +encapsulated in `SchedulerContext`. --- diff --git a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp index fa6c6aba1e..a322d96386 100644 --- a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp @@ -12,13 +12,14 @@ * Runtime Builder - rt2 Implementation (host_build_graph: Host Orchestration) * * Provides init_runtime_impl and validate_runtime_impl functions for rt2 runtime. - * The HOST runs the orchestrator to completion, populates shared memory + the - * prebuilt arena, and H2Ds the image; the device boots scheduler-only. + * The device boots scheduler-only from an empty prebuilt image. After launch, + * the HOST runs orchestration concurrently and publishes one immutable task + * prefix after each Graph submission. * * init_runtime_impl: * - Converts host tensor pointers to device pointers (all inputs copied H2D; * only OUTPUT/INOUT tensors are copied back D2H) - * - dlopens the orchestration SO on the host and runs it to build the graph + * - dlopens the orchestration SO and stages deferred host orchestration * - Sets up runtime state for host orchestration * * validate_runtime_impl: @@ -32,9 +33,12 @@ #include #include #include +#include #include #include +#include +#include #include #include #include @@ -42,7 +46,12 @@ #include #include #include +#include +#include #include +#include +#include +#include #include #include #include @@ -50,6 +59,7 @@ #include "../common/pto_runtime_status.h" #include "../runtime/common.h" #include "../runtime/dep_gen_host_graph.h" +#include "../runtime/graph_host_state.h" #include "../runtime/pto_orchestrator.h" #include "../runtime/pto_runtime2.h" #include "../runtime/pto_shared_memory.h" @@ -102,6 +112,13 @@ static int64_t _now_ms() { return static_cast(tv.tv_sec) * 1000 + tv.tv_usec / 1000; } +static uint64_t host_prof_cycles() { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return static_cast(ts.tv_sec) * PLATFORM_PROF_SYS_CNT_FREQ + + static_cast(ts.tv_nsec) * PLATFORM_PROF_SYS_CNT_FREQ / 1000000000ull; +} + static bool is_power_of_2_u64(uint64_t value) { return value != 0 && (value & (value - 1)) == 0; } template @@ -294,13 +311,9 @@ static int32_t pto2_read_runtime_status(Runtime *runtime, const HostApi *api, PT namespace { -// host_build_graph is host-orchestration-first: the HOST dlopens the -// orchestration .so and runs it to completion. The shared memory + arena carry -// host-DDR cross-task pointers (slot_state.task/payload, -// payload.fanin_inline_slot_states[], dep_pool/ready queues); the host relocates them to -// their final device addresses (relocate_host_orch_image, below) BEFORE the H2D -// copy, so the device receives a fully device-addressed image and schedules -// only — no on-device pointer fixup. +// host_build_graph launches scheduler-only device execution first, then runs +// the orchestration .so on a host thread. Each committed SM prefix is converted +// to device-addressed POD bytes before its publication counter is released. bool write_all_bytes(int fd, const uint8_t *data, size_t size) { size_t total = 0; @@ -354,190 +367,290 @@ struct HostOrchEntryPoints { OrchestrationBindFunc bind{nullptr}; }; -// Run the orchestrator on the host. `rt` was built with its scheduler half -// pointing at the device SM; here we re-point ONLY the orchestrator half at a -// host SM mirror, run the orchestration entry against it, latch the submitted -// task count, and H2D the populated SM to the device (the device scheduler -// reads task descriptors from there). The device never dereferences the -// orchestrator's SM pointers, so leaving them host-side is safe. Returns the -// total task count (>= 0) on success, or -1 on failure. -// host_build_graph host-orch: the orchestrator built the task graph in a host -// SM mirror and (when wiring is folded into submit) the fanout adjacency in the -// host arena, storing host-DDR addresses into the cross-task pointers. Relocate -// them to their FINAL device addresses here on the host, BEFORE the SM/arena are -// copied to the device — so the device receives a fully device-addressed image -// and boots scheduler-only with no on-device pointer fixup. -// -// Relocated pointers span TWO regions with DIFFERENT deltas: the SM block -// (slot_state.task/.payload, fanin_inline_slot_states[], dep-entry.slot_state, -// ready-queue slot.slot_state) and the arena block (slot_state.fanout_head, -// dep-entry.next point into the SM but live in the arena). -// Rather than track which delta each field needs, reloc() classifies every -// pointer by the region it points INTO and applies that region's delta; foreign -// and null pointers pass through untouched. The fanout adjacency is wired inline -// during host submit, so dep_pool/ready are already populated here. -// -// The orchestrator's own task-allocator pointers are intentionally NOT relocated -// (the device runs scheduler-only and never dereferences them, and must not call -// rt_orchestration_done — the host already did). Multi-fanin spill is not yet -// relocated; a task exceeding PTO2_FANIN_INLINE_CAP producers latches fatal here -// (returns false) rather than shipping un-relocated host pointers to the device. -// Returns false on any unrelocatable pointer so the caller can fail the prepare. -static bool relocate_host_orch_image( - PTO2SharedMemoryHandle &host_sm_handle, [[maybe_unused]] PTO2Runtime *rt, uint64_t host_sm, uint64_t sm_size, - int64_t sm_delta, uint64_t host_arena, uint64_t arena_size, int64_t arena_delta -) { - // host_build_graph is single-ring; the loops below iterate the lone ring and - // index header->ring (singular). If the ring depth ever grows, those loops - // would relocate the same ring N times (applying the delta repeatedly = - // corruption), so pin the assumption here. - static_assert(PTO2_MAX_RING_DEPTH == 1, "relocate_host_orch_image assumes a single ring"); - - // SM and arena windows must not overlap — reloc classifies a pointer by - // which window it falls in, so an overlap would misclassify and apply the - // wrong delta. Both are independent malloc-backed host buffers in practice; - // assert it so a future shared-buffer layout can't silently corrupt. - if (!(host_sm + sm_size <= host_arena || host_arena + arena_size <= host_sm)) { +// Host orchestration owns a private C++ build image while the device consumes +// immutable committed prefixes from the pooled SM. Only this host object holds +// STL state; every host-device transfer remains exact-size POD bytes. +struct DeferredHostOrchestration { + Runtime *runtime{nullptr}; + const HostApi *api{nullptr}; + DeviceArena host_arena; + PTO2RuntimeArenaLayout layout{}; + PTO2Runtime *rt{nullptr}; + std::vector host_sm; + PTO2SharedMemoryHandle host_sm_handle{}; + GraphHostStatePtr graph_state; + const HostOrchEntryPoints *entry_points{nullptr}; + void *device_sm{nullptr}; + void *device_arena{nullptr}; + void *gm_heap{nullptr}; + uint64_t sm_size{0}; + std::array task_window_sizes{}; + std::array heap_sizes{}; + int32_t published_tasks{0}; + size_t published_uploads{0}; + std::vector device_graph_submissions; + bool capture_host_orch{false}; + std::vector phase_records; + uint64_t host_start_cycles{0}; + uint64_t host_end_cycles{0}; + uint64_t first_publish_cycles{0}; + std::mutex start_mutex; + std::condition_variable start_cv; + bool device_execution_started{false}; + std::atomic finished{false}; + bool started{false}; + int result{-1}; +}; + +bool upload_new_graph_executions(DeferredHostOrchestration &state) { + PTO2OrchestratorState &orch = state.rt->orchestrator; + if (orch.graph_host_state == nullptr) return false; + const size_t count = graph_host_upload_count(*orch.graph_host_state); + for (size_t i = state.published_uploads; i < count; ++i) { + std::optional upload = graph_host_upload(*orch.graph_host_state, i); + if (!upload.has_value() || upload->outer_slot->task_kind != PTO2TaskKind::GRAPH || + upload->outer_slot->task == nullptr) { + LOG_ERROR("host-orch: invalid pending Graph POD image"); + return false; + } + void *device_submission = state.api->device_malloc(upload->bytes); + if (device_submission == nullptr) { + LOG_ERROR("host-orch: failed to allocate %zu bytes for Graph submission", upload->bytes); + return false; + } + if (state.api->copy_to_device(device_submission, upload->data, upload->bytes) != 0) { + LOG_ERROR("host-orch: failed to upload Graph submission POD image"); + state.api->device_free(device_submission); + return false; + } + upload->outer_slot->graph_context = device_submission; + state.device_graph_submissions.push_back(device_submission); + state.published_uploads = i + 1; + } + return true; +} + +bool publish_committed_prefix(DeferredHostOrchestration &state) { + static_assert(PTO2_MAX_RING_DEPTH == 1, "incremental host publication assumes one ring"); + if (!upload_new_graph_executions(state)) return false; + + PTO2SharedMemoryRingHeader &ring = state.host_sm_handle.header->ring; + const int32_t committed = ring.fc.current_task_index.load(std::memory_order_acquire); + if (committed < state.published_tasks || committed > static_cast(ring.task_window_size)) { LOG_ERROR( - "host-orch: SM window [%#lx,+%#lx) overlaps arena window [%#lx,+%#lx); cannot relocate", host_sm, sm_size, - host_arena, arena_size + "host-orch: invalid committed prefix [%d, %d) for window=%lu", state.published_tasks, committed, + static_cast(ring.task_window_size) ); return false; } + const int32_t count = committed - state.published_tasks; + if (count == 0) return true; + + const int32_t first = state.published_tasks; + const size_t n = static_cast(count); + const uint64_t window = ring.task_window_size; + const auto offsets = pto2_sm_layout::ring_segment_offsets(window); + auto *device_base = static_cast(state.device_sm); + + if (state.api->copy_to_device( + device_base + offsets.descriptors + static_cast(first) * sizeof(PTO2TaskDescriptor), + ring.task_descriptors + first, n * sizeof(PTO2TaskDescriptor) + ) != 0 || + state.api->copy_to_device( + device_base + offsets.payloads + static_cast(first) * sizeof(PTO2TaskPayload), + ring.task_payloads + first, n * sizeof(PTO2TaskPayload) + ) != 0) { + LOG_ERROR("host-orch: failed to upload committed task descriptor/payload prefix"); + return false; + } - bool ok = true; - auto reloc = [&](auto *&p) { - using Ptr = std::remove_reference_t; - uint64_t v = reinterpret_cast(p); - if (v == 0) { - return; - } - if (v >= host_sm && v < host_sm + sm_size) { - p = reinterpret_cast(static_cast(v + sm_delta)); - } else if (v >= host_arena && v < host_arena + arena_size) { - p = reinterpret_cast(static_cast(v + arena_delta)); - } else { - // A non-null pointer in neither window is an external/host address - // the device would dereference verbatim after H2D. No field should - // legitimately carry one; latch fatal rather than ship a host VA to - // the device (silent AICPU corruption otherwise). - LOG_ERROR("host-orch: pointer %#lx is outside both SM and arena windows; cannot relocate for device", v); - ok = false; - } - }; + // Slot-state contains two SM pointers. Patch their byte representation in + // an ordinary contiguous buffer so host code does not construct/copy an + // array of atomic-bearing PTO2TaskSlotState objects. + std::vector slot_bytes(n * sizeof(PTO2TaskSlotState)); + for (int32_t id = first; id < committed; ++id) { + const size_t relative = static_cast(id - first) * sizeof(PTO2TaskSlotState); + std::memcpy(slot_bytes.data() + relative, ring.slot_states + id, sizeof(PTO2TaskSlotState)); + const uintptr_t device_payload_addr = reinterpret_cast( + device_base + offsets.payloads + static_cast(id) * sizeof(PTO2TaskPayload) + ); + const uintptr_t device_task_addr = reinterpret_cast( + device_base + offsets.descriptors + static_cast(id) * sizeof(PTO2TaskDescriptor) + ); + std::memcpy( + slot_bytes.data() + relative + offsetof(PTO2TaskSlotState, payload), &device_payload_addr, + sizeof(device_payload_addr) + ); + std::memcpy( + slot_bytes.data() + relative + offsetof(PTO2TaskSlotState, task), &device_task_addr, + sizeof(device_task_addr) + ); + } + if (state.api->copy_to_device( + device_base + offsets.slot_states + static_cast(first) * sizeof(PTO2TaskSlotState), + slot_bytes.data(), slot_bytes.size() + ) != 0 || + state.api->copy_to_device( + device_base + offsets.completion_flags + static_cast(first) * sizeof(std::atomic), + ring.completion_flags + first, n * sizeof(std::atomic) + ) != 0) { + LOG_ERROR("host-orch: failed to upload committed slot-state/completion prefix"); + return false; + } - PTO2SharedMemoryHeader *header = host_sm_handle.header; - if (header != nullptr) { - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - PTO2SharedMemoryRingHeader &ring = header->ring; - int32_t count = ring.fc.current_task_index.load(std::memory_order_acquire); - for (int32_t slot = 0; slot < count; slot++) { - PTO2TaskSlotState *ss = &ring.slot_states[slot]; - // Polling: fanin is a flat array of position-independent local-id - // integers on the payload, so only the two per-slot arena/SM - // pointers need relocating. There is no fanout_head/dep_pool graph - // and no host-seeded ready queue (the device boot scan classifies), - // so those relocation passes are gone. - reloc(ss->task); - reloc(ss->payload); - } - } + // This is the sole visibility edge: the Scheduler acquires this count before + // reading any byte in the newly committed prefix. + if (state.api->publish_i32(pto2_sm_layout::ring_current_task_index_addr(state.device_sm), committed) != 0) { + LOG_ERROR("host-orch: failed to publish committed task count %d", committed); + return false; } - return ok; + if (state.first_publish_cycles == 0) state.first_publish_cycles = host_prof_cycles(); + state.published_tasks = committed; + LOG_DEBUG("host-orch: published task prefix [0, %d)", committed); + return true; } -int32_t run_host_orchestration( - Runtime *runtime, const HostApi *api, PTO2Runtime *rt, DeviceArena &host_arena, - const PTO2RuntimeArenaLayout &layout, void *device_sm, uint64_t sm_size, void *device_arena, void *gm_heap, - const uint64_t eff_heap_sizes[PTO2_MAX_RING_DEPTH], const uint64_t eff_task_window_sizes[PTO2_MAX_RING_DEPTH], - void *host_orch_func_ptr, const L2TaskArgs &orch_l2 -) { - // The dep_gen graph belongs to the orchestration that is about to run. - dep_gen_host_graph_begin_capture(); +bool graph_commit_callback(void *context) { + return context != nullptr && publish_committed_prefix(*static_cast(context)); +} - std::vector host_sm_buf(sm_size, 0); - void *host_sm = host_sm_buf.data(); +} // namespace - // Re-point the orchestrator half at the host SM (scheduler keeps device SM). - // init_data_from_layout resets the orchestrator state, so this is safe. - if (!rt->orchestrator.init_data_from_layout( - layout.orch, host_arena, host_sm, gm_heap, eff_heap_sizes[0], eff_task_window_sizes[0] - )) { - LOG_ERROR("host-orch: orchestrator re-init against host SM failed"); - return -1; +void Runtime::notify_deferred_host_execution_started() { + auto *state = static_cast(deferred_host_orchestration_); + if (state == nullptr) return; + { + std::lock_guard lock(state->start_mutex); + state->device_execution_started = true; } - rt->orchestrator.wire_arena_pointers(layout.orch, host_arena, &rt->scheduler); + state->start_cv.notify_one(); +} - // Initialize the host SM header (ring flow control) so submit_task can run. - PTO2SharedMemoryHandle host_sm_handle; - if (!host_sm_handle.init_per_ring(host_sm, sm_size, eff_task_window_sizes, eff_heap_sizes)) { - LOG_ERROR("host-orch: host SM init_per_ring failed"); - return -1; +int Runtime::run_deferred_host_orchestration(const HostApi *api) { + auto *state = static_cast(deferred_host_orchestration_); + if (state == nullptr || api == nullptr || api != state->api || state->started) return -1; + state->started = true; + { + std::unique_lock lock(state->start_mutex); + state->start_cv.wait(lock, [state]() { + return state->device_execution_started; + }); } + dep_gen_host_graph_begin_capture(); - // Install the ops table (host s_runtime_ops) and latch this run's cluster - // counts. worker_count is published by DeviceRunner::prepare_launch_shape - // before this bind, so the host orchestrator sees the same geometry the - // AICPU re-derives from the handshake at boot. - const int32_t block_dim = runtime->get_worker_count() / PLATFORM_CORES_PER_BLOCKDIM; - if (block_dim < 1) { - LOG_ERROR("host-orch: worker_count %d yields no clusters", runtime->get_worker_count()); - return -1; + PTO2Runtime *rt = state->rt; + const HostOrchEntryPoints *eps = state->entry_points; + bool ok = rt != nullptr && eps != nullptr && eps->entry != nullptr && eps->bind != nullptr; + if (!ok) { + LOG_ERROR("host-orch: deferred orchestration state is incomplete"); } - runtime_finalize_after_wire( - rt, block_dim * PLATFORM_AIC_CORES_PER_BLOCKDIM, block_dim * PLATFORM_AIV_CORES_PER_BLOCKDIM - ); - rt->mode = PTO2_MODE_EXECUTE; - // get_tensor_data/set_tensor_data dereference buffer.addr directly: the - // input tensors were mapped into host address space at staging time - // (HostApi::register_device_memory_to_host), so the host orchestrator can - // read control tensors (e.g. paged_attention's context_lens/block_table) in - // place. - - // Bind both framework_current_runtime instances: the host library's (used by - // rt_scope_* / rt_orchestration_done) and the orch .so's own copy (used by - // its inline rt_submit_* -> current_runtime()). - const HostOrchEntryPoints *eps = reinterpret_cast(host_orch_func_ptr); - framework_bind_runtime(rt); - if (eps->bind != nullptr) { + + if (ok) { + L2TaskArgs orch_l2; + orch_l2.create_from_chip_args(get_orch_args()); + rt->active_callable_hash = reinterpret_cast(eps->entry); + framework_bind_runtime(rt); eps->bind(rt); - } else { - LOG_ERROR("host-orch: orch .so framework_bind_runtime was not resolved"); - return -1; + +#if SIMPLER_DFX + if (state->capture_host_orch) { + rt->orchestrator.l2_swimlane_level = L2SwimlaneLevel::ORCH_PHASES; + size_t capacity = static_cast(state->task_window_sizes[0]); + if (capacity > PLATFORM_PHASE_RECORDS_PER_THREAD) capacity = PLATFORM_PHASE_RECORDS_PER_THREAD; + state->phase_records.resize(capacity); + rt->orchestrator.host_orch_phase_records = state->phase_records.data(); + rt->orchestrator.host_orch_phase_capacity = capacity; + rt->orchestrator.host_orch_phase_count = 0; + state->host_start_cycles = host_prof_cycles(); + } +#endif + + try { + rt_scope_begin(rt); + eps->entry(orch_l2); + rt_scope_end(rt); + rt_orchestration_done(rt); + ok = publish_committed_prefix(*state); + if (!ok && !rt->orchestrator.fatal) { + rt->orchestrator.report_fatal( + PTO2_ERROR_EXPLICIT_ORCH_FATAL, __FUNCTION__, "failed to publish final orchestration prefix" + ); + } + } catch (const std::exception &error) { + LOG_ERROR("host-orch: orchestration threw an exception: %s", error.what()); + ok = false; + } catch (...) { + LOG_ERROR("host-orch: orchestration threw an unknown exception"); + ok = false; + } + +#if SIMPLER_DFX + if (state->capture_host_orch) { + state->host_end_cycles = host_prof_cycles(); + size_t count = rt->orchestrator.host_orch_phase_count; + if (count > state->phase_records.size()) count = state->phase_records.size(); + state->phase_records.resize(count); + rt->orchestrator.host_orch_phase_records = nullptr; + rt->orchestrator.host_orch_phase_capacity = 0; + rt->orchestrator.host_orch_phase_count = 0; + } +#endif } - rt_scope_begin(rt); - eps->entry(orch_l2); - rt_scope_end(rt); - rt_orchestration_done(rt); - - int32_t total_tasks = pto2_sm_layout::ring_current_task_index_addr(host_sm)->load(std::memory_order_acquire); - - // Relocate the host-DDR cross-task pointers to their final DEVICE addresses - // on the host, before the SM and arena leave for the device. Pointers into - // the SM shift by sm_delta; pointers into the arena (fanout adjacency, wiring - // queue) shift by arena_delta. After this both the SM and arena carry device - // addresses, so the device boots scheduler-only. - const int64_t sm_delta = static_cast(reinterpret_cast(device_sm)) - - static_cast(reinterpret_cast(host_sm)); - const int64_t arena_delta = static_cast(reinterpret_cast(device_arena)) - - static_cast(reinterpret_cast(host_arena.base())); - if (!relocate_host_orch_image( - host_sm_handle, rt, reinterpret_cast(host_sm), sm_size, sm_delta, - reinterpret_cast(host_arena.base()), layout.arena_size, arena_delta - )) { - LOG_ERROR("host-orch: relocation failed; refusing to H2D an image with unrelocated host pointers"); - return -1; + if (!ok && rt != nullptr && !rt->orchestrator.fatal) { + rt->orchestrator.report_fatal(PTO2_ERROR_EXPLICIT_ORCH_FATAL, __FUNCTION__, "host orchestration failed"); + rt->orchestrator.mark_done(); } - if (api->copy_to_device(device_sm, host_sm, sm_size) != 0) { - LOG_ERROR("host-orch: H2D of populated SM failed"); - return -1; + // Publish the immutable host profile before EOS. The device cannot finish + // before observing EOS, and the host collector acquires finished afterwards. + state->result = ok && rt != nullptr && !rt->orchestrator.fatal ? 0 : -1; + state->finished.store(true, std::memory_order_release); + + int32_t orch_error = rt == nullptr ? PTO2_ERROR_EXPLICIT_ORCH_FATAL : + rt->orchestrator.sm_header->orch_error_code.load(std::memory_order_acquire); + if (api->publish_i32(pto2_sm_layout::orch_error_code_addr(state->device_sm), orch_error) != 0 || + api->publish_i32(pto2_sm_layout::orchestrator_done_addr(state->device_sm), 1) != 0) { + LOG_ERROR("host-orch: failed to publish orchestration EOS"); + state->result = -1; } - return total_tasks; + LOG_INFO("host-orch: finished after publishing %d tasks", state->published_tasks); + return state->result; } -} // namespace +void Runtime::release_deferred_host_orchestration(const HostApi *api) { + auto *state = static_cast(deferred_host_orchestration_); + if (state == nullptr) return; + if (api != nullptr) { + for (void *ptr : state->device_graph_submissions) + api->device_free(ptr); + } + delete state; + deferred_host_orchestration_ = nullptr; +} + +const std::vector &Runtime::get_host_orch_phase_records() const { + auto *state = static_cast(deferred_host_orchestration_); + if (state != nullptr && state->finished.load(std::memory_order_acquire)) return state->phase_records; + return host_orch_phase_records_; +} + +uint64_t Runtime::get_host_orch_start_cycles() const { + auto *state = static_cast(deferred_host_orchestration_); + return state != nullptr && state->finished.load(std::memory_order_acquire) ? state->host_start_cycles : + host_orch_start_cycles_; +} + +uint64_t Runtime::get_host_orch_end_cycles() const { + auto *state = static_cast(deferred_host_orchestration_); + return state != nullptr && state->finished.load(std::memory_order_acquire) ? state->host_end_cycles : + host_orch_end_cycles_; +} + +uint64_t Runtime::get_host_orch_first_publish_cycles() const { + auto *state = static_cast(deferred_host_orchestration_); + return state != nullptr && state->finished.load(std::memory_order_acquire) ? state->first_publish_cycles : 0; +} /** * Stage the per-callable resources (kernel binaries + orchestration SO) into @@ -658,7 +771,7 @@ register_callable_impl(const ChipCallable *callable, uint64_t (*upload_fn)(const extern "C" int bind_callable_to_runtime_impl( Runtime *runtime, const HostApi *api, const ChipStorageTaskArgs *orch_args, void *host_orch_func_ptr, const ArgDirection *signature, int sig_count, const uint64_t *ring_task_window, const uint64_t *ring_heap, - [[maybe_unused]] const uint64_t *ring_dep_pool // polling has no dep_pool; kept for ABI stability + [[maybe_unused]] const uint64_t *ring_dep_pool, int32_t l2_swimlane_level ) { if (runtime == nullptr) { LOG_ERROR("Runtime pointer is null"); @@ -782,8 +895,18 @@ extern "C" int bind_callable_to_runtime_impl( uint64_t sm_size = PTO2SharedMemoryHandle::calculate_size_per_ring(eff_task_window_sizes); int64_t t_prebuilt_start = _now_ms(); - DeviceArena host_arena; // libc malloc backend by default + auto deferred = std::make_unique(); + deferred->runtime = runtime; + deferred->api = api; + deferred->sm_size = sm_size; + deferred->capture_host_orch = l2_swimlane_level >= static_cast(L2SwimlaneLevel::ORCH_PHASES); + for (int r = 0; r < PTO2_MAX_RING_DEPTH; ++r) { + deferred->task_window_sizes[static_cast(r)] = eff_task_window_sizes[r]; + deferred->heap_sizes[static_cast(r)] = eff_heap_sizes[r]; + } + DeviceArena &host_arena = deferred->host_arena; PTO2RuntimeArenaLayout layout = runtime_reserve_layout(host_arena, eff_task_window_sizes, eff_heap_sizes); + deferred->layout = layout; if (host_arena.commit(DeviceArena::kDefaultBaseAlign) == nullptr) { LOG_ERROR("Failed to commit host arena for prebuilt runtime image"); return -1; @@ -841,30 +964,46 @@ extern "C" int bind_callable_to_runtime_impl( } runtime_wire_arena_pointers(host_arena, layout, rt); - // host_build_graph host-orch: run the orchestrator on the host now, against - // a host SM mirror, and ship the populated SM to the device. The arena - // (copied to the device below) carries the resulting orchestrator/scheduler - // state; the device boots scheduler-only. register_callable_impl guarantees - // host_orch_func_ptr is non-null on success (it fails the whole prepare - // otherwise), so this is an assertion-style guard, not a fallback path. - if (host_orch_func_ptr == nullptr) { + if (host_orch_func_ptr == nullptr || api->publish_i32 == nullptr) { LOG_ERROR("host-orch: orchestration entry points were not resolved"); return -1; } - { - L2TaskArgs orch_l2; - orch_l2.create_from_chip_args(device_args); - int32_t total_tasks = run_host_orchestration( - runtime, api, rt, host_arena, layout, sm_ptr, sm_size, runtime_arena_dev, gm_heap, eff_heap_sizes, - eff_task_window_sizes, host_orch_func_ptr, orch_l2 - ); - if (total_tasks < 0) { - LOG_ERROR("host-orch: orchestration run failed"); - return -1; - } - runtime->host_total_tasks = total_tasks; - LOG_INFO("host-orch: submitted %d tasks on host", total_tasks); + deferred->rt = rt; + deferred->entry_points = reinterpret_cast(host_orch_func_ptr); + deferred->device_sm = sm_ptr; + deferred->device_arena = runtime_arena_dev; + deferred->gm_heap = gm_heap; + deferred->host_sm.resize(static_cast(sm_size), 0); + if (!deferred->host_sm_handle.init_per_ring( + deferred->host_sm.data(), sm_size, eff_task_window_sizes, eff_heap_sizes + )) { + LOG_ERROR("host-orch: host SM init_per_ring failed"); + return -1; + } + if (!rt->orchestrator.init_data_from_layout( + layout.orch, host_arena, deferred->host_sm.data(), gm_heap, eff_heap_sizes[0], eff_task_window_sizes[0] + )) { + LOG_ERROR("host-orch: orchestrator re-init against host SM failed"); + return -1; + } + rt->orchestrator.wire_arena_pointers(layout.orch, host_arena, &rt->scheduler); + deferred->graph_state = make_graph_host_state(); + if (!deferred->graph_state) { + LOG_ERROR("host-orch: failed to allocate Graph host state"); + return -1; } + graph_host_set_commit_callback(*deferred->graph_state, graph_commit_callback, deferred.get()); + rt->orchestrator.graph_host_state = deferred->graph_state.get(); + + const int32_t block_dim = runtime->get_worker_count() / PLATFORM_CORES_PER_BLOCKDIM; + if (block_dim < 1) { + LOG_ERROR("host-orch: worker_count %d yields no clusters", runtime->get_worker_count()); + return -1; + } + runtime_finalize_after_wire( + rt, block_dim * PLATFORM_AIC_CORES_PER_BLOCKDIM, block_dim * PLATFORM_AIV_CORES_PER_BLOCKDIM + ); + rt->mode = PTO2_MODE_EXECUTE; // Stash the layout inside the PTO2Runtime image so the AICPU can recover // every arena-internal offset after rtMemcpy. The runtime arena's device @@ -873,12 +1012,24 @@ extern "C" int bind_callable_to_runtime_impl( // *before* it can dereference the image. rt->prebuilt_layout = layout; + // Upload an empty, scheduler-ready image before host orchestration starts. + // The host-only Graph cache pointer must not cross the boundary; restore it + // immediately after the synchronous copy so the deferred host thread keeps + // its C++ state. + GraphHostState *host_graph_state = rt->orchestrator.graph_host_state; + rt->orchestrator.graph_host_state = nullptr; int rc_upload = api->copy_to_device(runtime_arena_dev, host_arena.base(), layout.arena_size); + rt->orchestrator.graph_host_state = host_graph_state; if (rc_upload != 0) { LOG_ERROR("Failed to rtMemcpy prebuilt runtime arena to device (rc=%d)", rc_upload); return -1; } + if (api->copy_to_device(sm_ptr, deferred->host_sm.data(), static_cast(sm_size)) != 0) { + LOG_ERROR("host-orch: failed to initialize empty device SM"); + return -1; + } runtime->set_prebuilt_arena(runtime_arena_dev, layout.off_runtime); + runtime->set_deferred_host_orchestration(deferred.release()); int64_t t_prebuilt_end = _now_ms(); LOG_INFO("Device orchestration ready: %d tensors + %d scalars", tensor_count, scalar_count); @@ -983,12 +1134,19 @@ extern "C" int validate_runtime_impl(Runtime *runtime, const HostApi *api, int e // Release the SVM host mapping installed at staging time before // freeing the device buffer (unregister-before-free, as the HAL // requires). No-op on sim. Keyed by dev_ptr. - api->unregister_device_memory_from_host(tensor_pairs[i].dev_ptr); + if (tensor_pairs[i].host_ptr != nullptr) { + api->unregister_device_memory_from_host(tensor_pairs[i].dev_ptr); + } api->device_free(tensor_pairs[i].dev_ptr); } } LOG_INFO("Freed %d device allocations", tensor_pair_count); + // Graph submission PODs are owned by the deferred host publisher rather + // than tensor_pairs_ so the device Runtime can be copied concurrently with + // host construction without racing an STL vector mutation. + runtime->release_deferred_host_orchestration(api); + // Clear the per-run dispatch-table entries staged by register_callable_impl. // The underlying chip-callable device buffer is pool-managed by // DeviceRunner (keyed by content hash) and bulk-freed in diff --git a/src/a2a3/runtime/host_build_graph/orchestration/pto_orchestration_api.h b/src/a2a3/runtime/host_build_graph/orchestration/pto_orchestration_api.h index 92171d11e4..7dcdae0bd7 100644 --- a/src/a2a3/runtime/host_build_graph/orchestration/pto_orchestration_api.h +++ b/src/a2a3/runtime/host_build_graph/orchestration/pto_orchestration_api.h @@ -30,10 +30,12 @@ #include #include +#include #include // Type headers needed by orchestration #include "common.h" // framework_bind_runtime / framework_current_runtime +#include "graph_cache.h" // Graph Execution key and result helpers #include "pto_runtime2_types.h" // PTO2_ERROR_* #include "pto_submit_types.h" // MixedKernels, INVALID_KERNEL_ID, subtask slots #include "pto_types.h" // Arg, TaskOutputTensors, TensorArgType @@ -91,6 +93,9 @@ typedef struct PTO2RuntimeOps { // (one AIC each) and standalone AIV cores. int32_t (*available_cluster_count)(PTO2Runtime *rt); int32_t (*available_aiv_count)(PTO2Runtime *rt); + GraphScopeResult (*graph_begin)(PTO2Runtime *rt, uint64_t graph_key, const L0TaskArgs &args); + void (*graph_end)(PTO2Runtime *rt); + void (*graph_commit)(PTO2Runtime *rt); // Stash the call-site of the next PTO2ScopeGuard so the [ScopeStats] // collector can log it. Always present to keep ops-table layout stable @@ -208,6 +213,30 @@ static inline TaskOutputTensors rt_submit_dummy_task(const L0TaskArgs &args) { return rt->ops->submit_dummy_task(rt, args); } +static inline GraphScopeResult rt_graph_begin(uint64_t graph_key, const L0TaskArgs &args) { + PTO2Runtime *rt = current_runtime(); + if (rt->ops->is_fatal(rt) || rt->ops->graph_begin == nullptr) { + return GraphScopeResult{}; + } + return rt->ops->graph_begin(rt, graph_key, args); +} + +static inline void rt_graph_end() { + PTO2Runtime *rt = current_runtime(); + if (rt->ops->is_fatal(rt) || rt->ops->graph_end == nullptr) { + return; + } + rt->ops->graph_end(rt); +} + +static inline void rt_graph_commit() { + PTO2Runtime *rt = current_runtime(); + if (rt->ops->is_fatal(rt) || rt->ops->graph_commit == nullptr) { + return; + } + rt->ops->graph_commit(rt); +} + static inline void rt_scope_begin(PTO2ScopeMode mode = PTO2ScopeMode::AUTO) { PTO2Runtime *rt = current_runtime(); if (rt->ops->is_fatal(rt)) { @@ -354,6 +383,84 @@ class PTO2ScopeGuard { PTO2Runtime *rt_; }; +// Define or submit a Graph Execution. On a cache miss the function executes +// normally and its sub-DAG is recorded. On a hit the function is skipped and +// one Graph task is submitted; Scheduler expands the cached topology with the +// current invocation's L0TaskArgs. +using GraphFunction = void (*)(const L0TaskArgs &); + +template +static inline uint64_t rt_graph_function_id(Function function) { + static_assert(std::is_pointer_v, "Graph function identity requires a function pointer"); + static_assert(sizeof(function) <= sizeof(uint64_t), "Graph function pointer must fit in a 64-bit identity"); + uint64_t function_id = 0; + std::memcpy(&function_id, &function, sizeof(function)); + return function_id; +} + +template +static inline GraphSubmitResult rt_submit_graph_impl(uint64_t graph_key, const L0TaskArgs &args, Invoke invoke) { + debug_assert(!args.has_error && "Graph boundary L0TaskArgs construction failed"); + debug_assert( + args.tensor_count() <= static_cast(GRAPH_MAX_TENSOR_ARGS) && + "Graph boundary exceeds the step-1 tensor limit" + ); + debug_assert(args.scalar_count() == 0 && "Dynamic Graph boundary scalars are not supported in step 1"); + debug_assert( + args.explicit_dep_count() == 0 && "Explicit dependencies crossing the Graph boundary are not supported" + ); + for (int32_t i = 0; i < args.tensor_count(); ++i) { + debug_assert( + args.tag(i) != TensorArgType::OUTPUT && + "Runtime-allocated TensorCreateInfo is not supported at the Graph boundary" + ); + } + if (!rt_graph_args_cacheable(args)) { + invoke(); + rt_graph_commit(); + return GraphSubmitResult{}; + } + GraphScopeResult result = rt_graph_begin(graph_key, args); + if (result.execute_block) invoke(); + if (result.recording) { + rt_graph_end(); + } + rt_graph_commit(); + return result; +} + +static inline GraphSubmitResult rt_submit_graph(uint64_t graph_id, GraphFunction function, const L0TaskArgs &args) { + debug_assert(function != nullptr && "Graph function must not be null"); + if (function == nullptr) return GraphSubmitResult{}; + return rt_submit_graph_impl(rt_graph_make_key(graph_id), args, [&]() { + function(args); + }); +} + +static inline GraphSubmitResult rt_submit_graph(GraphFunction function, const L0TaskArgs &args) { + return rt_submit_graph(rt_graph_function_id(function), function, args); +} + +template +using GraphFunctionWithConfig = void (*)(const L0TaskArgs &, Config...); + +template +static inline GraphSubmitResult rt_submit_graph( + uint64_t graph_id, GraphFunctionWithConfig function, const L0TaskArgs &args, Config... config +) { + debug_assert(function != nullptr && "Graph function must not be null"); + if (function == nullptr) return GraphSubmitResult{}; + return rt_submit_graph_impl(rt_graph_make_key(graph_id, config...), args, [&]() { + function(args, config...); + }); +} + +template +static inline GraphSubmitResult +rt_submit_graph(GraphFunctionWithConfig function, const L0TaskArgs &args, Config... config) { + return rt_submit_graph(rt_graph_function_id(function), function, args, config...); +} + #define _PTO2_CONCATENATE_IMPL(x, y) x##y #define _PTO2_CONCATENATE(x, y) _PTO2_CONCATENATE_IMPL(x, y) diff --git a/src/a2a3/runtime/host_build_graph/runtime/dep_gen_host_graph.h b/src/a2a3/runtime/host_build_graph/runtime/dep_gen_host_graph.h index c4d286267f..d4c8345e7e 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/dep_gen_host_graph.h +++ b/src/a2a3/runtime/host_build_graph/runtime/dep_gen_host_graph.h @@ -23,7 +23,7 @@ * host can only replay a ring of captured submits. * * Capture surface: - * begin_capture() — once per orchestration, from run_host_orchestration + * begin_capture() — once per deferred host orchestration * begin_task() — one per submit, before its dependency steps * add_explicit_edge() — STEP 1, per declared dependency * add_creator_edge() — STEP 3 Step A, per creator-retention producer diff --git a/src/a2a3/runtime/host_build_graph/runtime/graph_cache.h b/src/a2a3/runtime/host_build_graph/runtime/graph_cache.h new file mode 100644 index 0000000000..121b094fe5 --- /dev/null +++ b/src/a2a3/runtime/host_build_graph/runtime/graph_cache.h @@ -0,0 +1,88 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#pragma once + +#include +#include + +#include + +#include "pto_task_id.h" +#include "pto_types.h" + +inline constexpr uint32_t GRAPH_MAX_TENSOR_ARGS = 32; + +struct GraphScopeResult { + bool execute_block{true}; + bool recording{false}; + PTO2TaskId task_id{PTO2TaskId::invalid()}; +}; + +using GraphSubmitResult = GraphScopeResult; + +constexpr uint64_t graph_hash_byte(uint64_t h, uint8_t b) { return (h ^ static_cast(b)) * 1099511628211ULL; } + +inline uint64_t graph_hash_bytes(uint64_t h, const void *data, size_t bytes) { + const auto *p = static_cast(data); + for (size_t i = 0; i < bytes; ++i) { + h = graph_hash_byte(h, p[i]); + } + return h; +} + +constexpr uint64_t graph_const_hash_impl(const char *s, uint64_t h) { + return (*s == '\0') ? h : graph_const_hash_impl(s + 1, graph_hash_byte(h, static_cast(*s))); +} + +constexpr uint64_t GRAPH_KEY(const char *s) { return graph_const_hash_impl(s, 1469598103934665603ULL); } + +inline bool rt_graph_args_cacheable(const L0TaskArgs &args) { + // Step 1 supports dynamic tensor addresses only. Kernel scalars are + // literals inside the Graph function and become immutable Definition data. + if (args.has_error || args.tensor_count() > static_cast(GRAPH_MAX_TENSOR_ARGS) || + args.scalar_count() != 0) { + return false; + } + for (int32_t i = 0; i < args.tensor_count(); ++i) { + // A Graph boundary is caller-owned storage. Runtime-allocated + // TensorCreateInfo outputs remain on the ordinary submit path. + if (args.tag(i) == TensorArgType::OUTPUT) return false; + } + return true; +} + +inline uint64_t rt_graph_make_key(uint64_t graph_id) { return graph_id; } + +template +inline uint64_t graph_hash_config_value(uint64_t hash, T value) { + using Value = std::remove_cv_t>; + static_assert( + std::is_integral_v || std::is_same_v || std::is_same_v, + "Graph construction parameters must be integral, float, or double values" + ); + constexpr uint8_t category = std::is_same_v ? 1 : + std::is_integral_v ? (std::is_signed_v ? 2 : 3) : + 4; + constexpr uint8_t width = sizeof(Value); + hash = graph_hash_byte(hash, category); + hash = graph_hash_byte(hash, width); + return graph_hash_bytes(hash, &value, sizeof(value)); +} + +template +inline uint64_t rt_graph_make_key(uint64_t graph_id, Config... config) { + uint64_t hash = graph_hash_bytes(1469598103934665603ULL, &graph_id, sizeof(graph_id)); + const uint32_t count = sizeof...(Config); + hash = graph_hash_bytes(hash, &count, sizeof(count)); + ((hash = graph_hash_config_value(hash, config)), ...); + return hash; +} diff --git a/src/a2a3/runtime/host_build_graph/runtime/graph_execution.h b/src/a2a3/runtime/host_build_graph/runtime/graph_execution.h new file mode 100644 index 0000000000..e6ccde8928 --- /dev/null +++ b/src/a2a3/runtime/host_build_graph/runtime/graph_execution.h @@ -0,0 +1,329 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#pragma once + +#include +#include + +#include +#include + +#include "pto_runtime2_types.h" +#include "tensor.h" + +inline constexpr uint32_t GRAPH_MAX_NODES = 1024; +inline constexpr int32_t GRAPH_MATERIALIZE_SLICE_NODES = 4; + +enum class GraphTensorSource : uint8_t { + BOUNDARY_EXACT = 0, + BOUNDARY_VIEW = 1, + INTERNAL = 2, + OWN_OUTPUT = 3, +}; + +// Wire representation of Tensor. Tensor itself is a host/runtime C++ type with +// 64-byte alignment and helper methods; placing it inside vector +// would not guarantee that alignment. Keep the boundary image C-compatible and +// copy only semantic fields into this naturally 8-byte-aligned POD. +struct GraphTensor { + uint64_t buffer_addr; + uint64_t buffer_size; + uint64_t owner_task_id; + uint64_t start_offset; + uint64_t extent_elem; + int32_t version; + uint32_t shapes[MAX_TENSOR_DIMS]; + uint32_t strides[MAX_TENSOR_DIMS]; + uint8_t ndims; + uint8_t dtype; + uint8_t manual_dep; + uint8_t is_contiguous; + uint8_t child_memory; + uint8_t reserved[3]; +}; + +// Everything from GraphTensorSourceRef through GraphSubmission is copied +// across the host-device boundary. Keep it pointer-free, fixed-width and +// position-independent: every reference is an offset from its owning header. +struct GraphTensorSourceRef { + uint8_t source; + uint8_t reserved; + uint16_t source_index; + uint32_t reserved2; + uint64_t packed_offset; +}; + +struct GraphNodeDefinition { + int32_t kernel_id[PTO2_SUBTASK_SLOT_COUNT]; + uint8_t active_mask; + uint8_t task_attrs; + int16_t logical_block_num; + int16_t total_required_subtasks; + uint16_t reserved; + int32_t tensor_count; + int32_t scalar_count; + int32_t total_output_size; + uint32_t tensor_offset; + uint32_t scalar_offset; +}; + +struct GraphBoundarySignature { + uint64_t buffer_size; + uint32_t shapes[MAX_TENSOR_DIMS]; + uint32_t strides[MAX_TENSOR_DIMS]; + uint16_t alias_rep; + uint8_t ndims; + uint8_t dtype; + uint8_t tag; + uint8_t manual_dep; + uint8_t is_contiguous; + uint8_t reserved; +}; + +struct GraphDefinition { + uint64_t full_key; + uint64_t content_hash; + uint64_t required_heap; + uint32_t total_bytes; + uint32_t task_count; + uint32_t edge_count; + uint32_t root_count; + uint32_t boundary_count; + uint32_t tensor_arg_count; + uint32_t scalar_arg_count; + uint32_t off_fanout_offsets; + uint32_t off_fanout_indices; + uint32_t off_fanin_offsets; + uint32_t off_fanin_indices; + uint32_t off_root_indices; + uint32_t off_node_offsets; + uint32_t off_nodes; + uint32_t off_tensors; + uint32_t off_tensor_sources; + uint32_t off_scalars; + uint32_t off_boundary_signatures; +}; + +struct GraphSubmission { + uint64_t graph_key; + uint64_t local_execution; + uint32_t activation_gate; + uint32_t total_bytes; + uint32_t definition_offset; + uint32_t tensors_offset; + uint32_t tensor_count; + uint32_t reserved; +}; + +static_assert(std::is_trivially_copyable_v); +static_assert(std::is_standard_layout_v); +static_assert(std::is_trivially_copyable_v); +static_assert(std::is_standard_layout_v); +static_assert(std::is_trivially_copyable_v); +static_assert(std::is_standard_layout_v); +static_assert(std::is_trivially_copyable_v); +static_assert(std::is_standard_layout_v); +static_assert(std::is_trivially_copyable_v); +static_assert(std::is_standard_layout_v); +static_assert(std::is_trivially_copyable_v); +static_assert(std::is_standard_layout_v); + +inline GraphTensor graph_tensor_pack(const Tensor &tensor) { + GraphTensor packed{}; + packed.buffer_addr = tensor.buffer.addr; + packed.buffer_size = tensor.buffer.size; + packed.owner_task_id = tensor.owner_task_id.raw; + packed.start_offset = tensor.start_offset; + packed.extent_elem = tensor.extent_elem_cache; + packed.version = tensor.version; + for (uint32_t i = 0; i < tensor.ndims; ++i) { + packed.shapes[i] = tensor.shapes[i]; + packed.strides[i] = tensor.strides[i]; + } + packed.ndims = static_cast(tensor.ndims); + packed.dtype = static_cast(tensor.dtype); + packed.manual_dep = tensor.manual_dep ? 1 : 0; + packed.is_contiguous = tensor.is_contiguous ? 1 : 0; + packed.child_memory = tensor.child_memory; + return packed; +} + +inline void graph_tensor_unpack(const GraphTensor &packed, Tensor *tensor) { + tensor->buffer = PTOBufferHandle{packed.buffer_addr, packed.buffer_size}; + tensor->owner_task_id = PTO2TaskId{packed.owner_task_id}; + tensor->start_offset = packed.start_offset; + tensor->extent_elem_cache = packed.extent_elem; + tensor->version = packed.version; + tensor->ndims = packed.ndims; + tensor->dtype = static_cast(packed.dtype); + tensor->manual_dep = packed.manual_dep != 0; + tensor->is_contiguous = packed.is_contiguous != 0; + tensor->child_memory = packed.child_memory; + for (uint32_t i = 0; i < MAX_TENSOR_DIMS; ++i) { + tensor->shapes[i] = packed.shapes[i]; + tensor->strides[i] = packed.strides[i]; + } + for (uint8_t &byte : tensor->_pad_cl2) + byte = 0; +} + +inline bool graph_tensor_wire_valid(const GraphTensor &tensor) { + if (tensor.buffer_addr == 0 || tensor.ndims == 0 || tensor.ndims > MAX_TENSOR_DIMS || + tensor.dtype >= static_cast(DataType::DATA_TYPE_NUM) || tensor.manual_dep > 1 || + tensor.is_contiguous > 1 || tensor.child_memory > 1) { + return false; + } + + uint64_t extent = 1; + uint64_t expected_stride = 1; + bool contiguous = true; + for (int32_t i = static_cast(tensor.ndims) - 1; i >= 0; --i) { + const uint64_t shape = tensor.shapes[i]; + const uint64_t stride = tensor.strides[i]; + if (shape == 0 || stride == 0) return false; + contiguous &= stride == expected_stride; + if (shape - 1 > (UINT64_MAX - extent) / stride || expected_stride > UINT64_MAX / shape) return false; + extent += (shape - 1) * stride; + expected_stride *= shape; + } + if (extent != tensor.extent_elem || contiguous != (tensor.is_contiguous != 0)) return false; + + const uint64_t element_size = get_element_size(static_cast(tensor.dtype)); + const uint64_t buffer_elements = tensor.buffer_size / element_size; + return tensor.start_offset <= buffer_elements && tensor.extent_elem <= buffer_elements - tensor.start_offset; +} + +template +inline const T *graph_definition_array(const GraphDefinition &definition, uint32_t offset, uint32_t count) { + if (offset == 0 || offset > definition.total_bytes || offset % alignof(T) != 0) return nullptr; + const size_t remaining = static_cast(definition.total_bytes - offset); + if (count > remaining / sizeof(T)) return nullptr; + return reinterpret_cast(reinterpret_cast(&definition) + offset); +} + +template +inline const T *graph_definition_ptr(const GraphDefinition &definition, uint32_t offset) { + return graph_definition_array(definition, offset, 1); +} + +inline GraphSubmission *graph_submission_from_slot(PTO2TaskSlotState &slot) { + return slot.task_kind == PTO2TaskKind::GRAPH ? static_cast(slot.graph_context) : nullptr; +} + +inline const GraphDefinition *graph_submission_definition(const GraphSubmission &submission) { + if (submission.definition_offset == 0 || submission.definition_offset % alignof(GraphDefinition) != 0 || + submission.definition_offset > submission.total_bytes || + sizeof(GraphDefinition) > submission.total_bytes - submission.definition_offset) { + return nullptr; + } + return reinterpret_cast( + reinterpret_cast(&submission) + submission.definition_offset + ); +} + +inline const GraphTensor *graph_submission_tensors(const GraphSubmission &submission) { + if (submission.tensors_offset == 0 || submission.tensors_offset % alignof(GraphTensor) != 0 || + submission.tensors_offset > submission.total_bytes || + submission.tensor_count > (submission.total_bytes - submission.tensors_offset) / sizeof(GraphTensor)) { + return nullptr; + } + return reinterpret_cast( + reinterpret_cast(&submission) + submission.tensors_offset + ); +} + +enum class GraphExecutionState : uint8_t { + SUBMITTED = 0, + MATERIALIZING = 1, + PREPARED = 2, + ACTIVE = 3, + COMPLETED = 4, +}; + +enum class GraphMaterializeResult : uint8_t { + INVALID = 0, + BUSY = 1, + PENDING = 2, + PREPARED = 3, +}; + +struct alignas(64) GraphNodeStorage { + PTO2TaskDescriptor task; + PTO2TaskPayload payload; + PTO2TaskSlotState slot; +}; + +struct GraphExecution { + std::atomic state{GraphExecutionState::SUBMITTED}; + std::atomic materialize_busy{0}; + std::atomic remaining_nodes{0}; + std::atomic retired_nodes{0}; + int32_t node_count{0}; + int32_t node_capacity{0}; + int32_t materialized_nodes{0}; + int32_t materialized_node_count{0}; + int32_t constructed_nodes{0}; + size_t allocation_bytes{0}; + size_t definition_capacity{0}; + uint64_t graph_key{0}; + uint64_t definition_hash{0}; + uint64_t materialized_graph_key{0}; + uint64_t materialized_definition_hash{0}; + bool definition_affine_reuse{false}; + PTO2TaskSlotState *outer_slot{nullptr}; + GraphNodeStorage *nodes{nullptr}; + GraphNodeStorage *node_storage{nullptr}; + void *definition_storage{nullptr}; + const GraphDefinition *definition{nullptr}; + const uint32_t *fanin_offsets{nullptr}; + const uint16_t *fanin_indices{nullptr}; + const GraphTensor *boundary_tensors{nullptr}; + uint32_t boundary_tensor_count{0}; + GraphExecution *next{nullptr}; +}; + +GraphExecution * +graph_execution_create(int32_t node_count, uint64_t graph_key, uint64_t definition_hash, size_t definition_bytes); +void graph_execution_discard(GraphExecution *execution); +void graph_execution_publish(GraphExecution *execution); +void graph_execution_collect_retired(); +GraphExecution *graph_execution_localize(PTO2TaskSlotState &outer_slot); +GraphMaterializeResult graph_execution_materialize_slice( + PTO2TaskSlotState &outer_slot, GraphExecution &execution, int32_t max_nodes, int32_t *nodes_materialized = nullptr +); + +inline GraphExecution *graph_execution_from_slot(PTO2TaskSlotState &slot) { + return slot.task_kind == PTO2TaskKind::GRAPH_NODE ? static_cast(slot.graph_context) : nullptr; +} + +inline bool graph_execution_complete_node(GraphExecution &execution) { + return execution.remaining_nodes.fetch_sub(1, std::memory_order_acq_rel) == 1; +} + +inline void graph_execution_mark_completed(GraphExecution &execution) { + execution.state.store(GraphExecutionState::COMPLETED, std::memory_order_release); +} + +inline void graph_execution_retire_node(GraphExecution &execution) { + execution.retired_nodes.fetch_add(1, std::memory_order_release); +} + +inline bool graph_submission_signal(GraphSubmission &submission, uint32_t bit) { + constexpr uint32_t BOTH = 0x3; + uint32_t observed = __atomic_fetch_or(&submission.activation_gate, bit, __ATOMIC_ACQ_REL); + return (observed | bit) == BOTH; +} + +inline GraphExecution *graph_submission_local_execution(GraphSubmission &submission) { + uint64_t raw = __atomic_load_n(&submission.local_execution, __ATOMIC_ACQUIRE); + return reinterpret_cast(static_cast(raw)); +} diff --git a/src/a2a3/runtime/host_build_graph/runtime/graph_host_state.h b/src/a2a3/runtime/host_build_graph/runtime/graph_host_state.h new file mode 100644 index 0000000000..b738bedd79 --- /dev/null +++ b/src/a2a3/runtime/host_build_graph/runtime/graph_host_state.h @@ -0,0 +1,41 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#pragma once + +#include +#include +#include + +struct PTO2TaskSlotState; +struct GraphHostState; + +inline constexpr size_t GRAPH_MAX_DEFINITIONS = 16; + +struct GraphHostStateDeleter { + void operator()(GraphHostState *state) const noexcept; +}; + +using GraphHostStatePtr = std::unique_ptr; + +struct GraphHostUpload { + PTO2TaskSlotState *outer_slot; + const std::byte *data; + size_t bytes; +}; + +using GraphHostCommitCallback = bool (*)(void *context); + +GraphHostStatePtr make_graph_host_state(); +size_t graph_host_upload_count(const GraphHostState &state); +std::optional graph_host_upload(const GraphHostState &state, size_t index); +void graph_host_set_commit_callback(GraphHostState &state, GraphHostCommitCallback callback, void *context); +bool graph_host_commit(GraphHostState &state); diff --git a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp index df2831d891..eb3d0c9e10 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp @@ -27,10 +27,24 @@ #include #include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + #include "common/platform_config.h" #include "common/unified_log.h" #include "dep_gen_host_graph.h" #include "pto_dep_compute.h" +#include "graph_execution.h" +#include "graph_host_state.h" #include "pto_runtime2_types.h" #include "pto_shared_memory.h" #include "pto_tensormap.h" @@ -85,6 +99,27 @@ __attribute__((weak, visibility("hidden"))) volatile uint32_t *get_reg_ptr(uint6 // ============================================================================= // Orchestrator Profiling (compile-time toggle) // ============================================================================= +#if SIMPLER_DFX +namespace { + +void record_host_orch_phase( + PTO2OrchestratorState *orch, uint64_t start_time, uint64_t end_time, uint64_t task_id, uint32_t submit_idx +) { + if (orch == nullptr || orch->host_orch_phase_records == nullptr || + orch->host_orch_phase_count >= orch->host_orch_phase_capacity) { + return; + } + L2SwimlaneAicpuOrchPhaseRecord &record = orch->host_orch_phase_records[orch->host_orch_phase_count++]; + record.start_time = start_time; + record.end_time = end_time; + record.task_id = task_id; + record.submit_idx = submit_idx; + record._pad = 0; +} + +} // namespace +#endif + #if SIMPLER_ORCH_PROFILING #include "aicpu/device_time.h" #include "aicpu/l2_swimlane_collector_aicpu.h" @@ -151,11 +186,11 @@ uint64_t g_orch_scope_end_atomic_count = 0; acc += (_t1 - _t0); \ _t0 = _t1; \ } while (0) -#define CYCLE_COUNT_ORCH_SUBMIT_RECORD(tid) \ - do { \ - if (_prof_active) { \ - l2_swimlane_aicpu_record_orch_phase(_submit_start_ts, _t1, (tid), g_orch_submit_idx); \ - } \ +#define CYCLE_COUNT_ORCH_SUBMIT_RECORD(tid) \ + do { \ + if (_prof_active) { \ + record_host_orch_phase(orch, _submit_start_ts, _t1, (tid), g_orch_submit_idx); \ + } \ } while (0) #elif SIMPLER_DFX #include "aicpu/device_time.h" @@ -183,12 +218,12 @@ static uint32_t g_orch_submit_idx = 0; #define CYCLE_COUNT_LAP(acc) \ do { \ } while (0) -#define CYCLE_COUNT_ORCH_SUBMIT_RECORD(tid) \ - do { \ - if (_prof_active) { \ - _t1 = get_sys_cnt_aicpu(); \ - l2_swimlane_aicpu_record_orch_phase(_submit_start_ts, _t1, (tid), g_orch_submit_idx); \ - } \ +#define CYCLE_COUNT_ORCH_SUBMIT_RECORD(tid) \ + do { \ + if (_prof_active) { \ + _t1 = get_sys_cnt_aicpu(); \ + record_host_orch_phase(orch, _submit_start_ts, _t1, (tid), g_orch_submit_idx); \ + } \ } while (0) #else #define CYCLE_COUNT_START() @@ -232,13 +267,13 @@ orch_report_fatal_v(PTO2OrchestratorState *orch, int32_t error_code, const char return; } - char message[1024]; - vsnprintf(message, sizeof(message), fmt, args); + std::array message{}; + vsnprintf(message.data(), message.size(), fmt, args); if (latched_code != PTO2_ERROR_NONE && latched_code != error_code) { - unified_log_error(func, "FATAL(code=%d, latched=%d): %s", error_code, latched_code, message); + unified_log_error(func, "FATAL(code=%d, latched=%d): %s", error_code, latched_code, message.data()); return; } - unified_log_error(func, "FATAL(code=%d): %s", error_code, message); + unified_log_error(func, "FATAL(code=%d): %s", error_code, message.data()); } void PTO2OrchestratorState::report_fatal(int32_t error_code, const char *func, const char *fmt, ...) { @@ -249,6 +284,452 @@ void PTO2OrchestratorState::report_fatal(int32_t error_code, const char *func, c va_end(args); } +enum class GraphRecordedTensorSource : uint8_t { + BOUNDARY_EXACT, + BOUNDARY_VIEW, + INTERNAL, + OWN_OUTPUT, +}; + +struct GraphRecordedTensorSourceRef { + GraphRecordedTensorSource source{GraphRecordedTensorSource::BOUNDARY_EXACT}; + size_t source_index{0}; + uint64_t packed_offset{0}; +}; + +struct GraphRecordedNode { + std::array kernel_ids{}; + ActiveMask active_mask{}; + TaskAttrs task_attrs{}; + int16_t logical_block_num{1}; + int16_t total_required_subtasks{0}; + size_t total_output_size{0}; + uintptr_t record_packed_base{0}; + std::vector tensors; + std::vector tensor_sources; + std::vector scalars; + std::vector internal_fanins; +}; + +struct GraphRecording { + uint64_t full_key{0}; + int32_t start_local_task_id{0}; + std::optional current_task_index; + bool unsupported{false}; + std::vector current_fanins; + std::vector boundary_tensors; + std::vector boundary_types; + std::vector nodes; +}; + +struct GraphPendingUpload { + PTO2TaskSlotState *outer_slot{nullptr}; + std::vector image; +}; + +struct GraphHostState { + std::unordered_map> definitions; + std::unique_ptr recording; + std::vector pending_uploads; + GraphHostCommitCallback commit_callback{nullptr}; + void *commit_context{nullptr}; +}; + +namespace { + +GraphHostState *graph_state_from(PTO2OrchestratorState *orch) { + return orch == nullptr ? nullptr : static_cast(orch->graph_host_state); +} + +uint64_t graph_full_key(uint64_t callable_hash, uint64_t graph_key) { + uint64_t h = 1469598103934665603ULL; + h = graph_hash_bytes(h, &callable_hash, sizeof(callable_hash)); + return graph_hash_bytes(h, &graph_key, sizeof(graph_key)); +} + +bool graph_tensor_exact(const Tensor &lhs, const Tensor &rhs) { + if (lhs.ndims > MAX_TENSOR_DIMS || rhs.ndims > MAX_TENSOR_DIMS || lhs.buffer.addr != rhs.buffer.addr || + lhs.buffer.size != rhs.buffer.size || lhs.start_offset != rhs.start_offset || lhs.version != rhs.version || + lhs.ndims != rhs.ndims || lhs.dtype != rhs.dtype || lhs.manual_dep != rhs.manual_dep || + lhs.is_contiguous != rhs.is_contiguous || lhs.child_memory != rhs.child_memory) { + return false; + } + return std::equal(std::begin(lhs.shapes), std::begin(lhs.shapes) + lhs.ndims, std::begin(rhs.shapes)) && + std::equal(std::begin(lhs.strides), std::begin(lhs.strides) + lhs.ndims, std::begin(rhs.strides)); +} + +bool graph_tensor_from_boundary( + const GraphRecording &recording, const Tensor &tensor, GraphRecordedTensorSourceRef *source +) { + for (size_t i = 0; i < recording.boundary_tensors.size(); ++i) { + if (!graph_tensor_exact(tensor, recording.boundary_tensors[i])) continue; + source->source = GraphRecordedTensorSource::BOUNDARY_EXACT; + source->source_index = i; + source->packed_offset = 0; + return true; + } + for (size_t i = 0; i < recording.boundary_tensors.size(); ++i) { + const Tensor &boundary = recording.boundary_tensors[i]; + if (tensor.buffer.addr != boundary.buffer.addr || tensor.buffer.size != boundary.buffer.size || + tensor.start_offset < boundary.start_offset) { + continue; + } + source->source = GraphRecordedTensorSource::BOUNDARY_VIEW; + source->source_index = i; + source->packed_offset = tensor.start_offset - boundary.start_offset; + return true; + } + return false; +} + +bool graph_classify_tensor( + const GraphRecording &recording, const GraphRecordedNode ¤t, int32_t task_index, const Tensor &tensor, + GraphRecordedTensorSourceRef *source +) { + if (graph_tensor_from_boundary(recording, tensor, source)) return true; + const uint64_t tensor_addr = tensor.buffer.addr; + for (int32_t producer_index = task_index; producer_index >= 0; --producer_index) { + const GraphRecordedNode &producer = + producer_index == task_index ? current : recording.nodes[static_cast(producer_index)]; + if (producer.record_packed_base == 0 || producer.total_output_size == 0 || + producer.total_output_size > UINTPTR_MAX - producer.record_packed_base) { + continue; + } + const uintptr_t begin = producer.record_packed_base; + const uintptr_t end = begin + producer.total_output_size; + if (tensor_addr < begin || tensor_addr >= end) continue; + source->source = + producer_index == task_index ? GraphRecordedTensorSource::OWN_OUTPUT : GraphRecordedTensorSource::INTERNAL; + source->source_index = static_cast(producer_index); + source->packed_offset = tensor_addr - begin; + return true; + } + return false; +} + +void graph_record_begin_task(PTO2OrchestratorState *orch, PTO2TaskId task_id) { + GraphHostState *state = graph_state_from(orch); + if (state == nullptr || state->recording == nullptr || state->recording->unsupported) return; + GraphRecording &recording = *state->recording; + const int32_t index = static_cast(task_id.local()) - recording.start_local_task_id; + if (index < 0 || index >= static_cast(GRAPH_MAX_NODES) || + index != static_cast(recording.nodes.size())) { + recording.unsupported = true; + return; + } + recording.current_task_index = static_cast(index); + recording.current_fanins.clear(); +} + +void graph_record_note_fanin(PTO2OrchestratorState *orch, PTO2TaskSlotState *producer) { + GraphHostState *state = graph_state_from(orch); + if (state == nullptr || state->recording == nullptr || state->recording->unsupported) return; + GraphRecording &recording = *state->recording; + if (producer == nullptr || producer->task == nullptr || !recording.current_task_index.has_value()) { + recording.unsupported = true; + return; + } + const int32_t producer_index = + static_cast(producer->task->task_id.local()) - recording.start_local_task_id; + if (producer_index >= 0 && static_cast(producer_index) >= *recording.current_task_index) { + recording.unsupported = true; + return; + } + if (producer_index >= 0) recording.current_fanins.push_back(static_cast(producer_index)); +} + +void graph_record_mark_unsupported(PTO2OrchestratorState *orch) { + GraphHostState *state = graph_state_from(orch); + if (state != nullptr && state->recording != nullptr) state->recording->unsupported = true; +} + +void graph_record_task( + PTO2OrchestratorState *orch, PTO2TaskId task_id, const PTO2TaskDescriptor &task, const PTO2TaskPayload &payload, + const PTO2TaskSlotState &slot, const L0TaskArgs &args +) { + GraphHostState *state = graph_state_from(orch); + if (state == nullptr || state->recording == nullptr || state->recording->unsupported) return; + GraphRecording &recording = *state->recording; + const int32_t task_index = static_cast(task_id.local()) - recording.start_local_task_id; + if (task_index < 0 || !recording.current_task_index.has_value() || + static_cast(task_index) != *recording.current_task_index || + static_cast(task_index) != recording.nodes.size() || args.predicate().op != PredicateOp::NONE) { + recording.unsupported = true; + return; + } + for (uint32_t i = 0; i < args.explicit_dep_count(); ++i) { + const PTO2TaskId dep = args.explicit_dep(i); + const int32_t dep_index = static_cast(dep.local()) - recording.start_local_task_id; + if (!dep.is_valid() || dep.ring() != 0 || dep_index >= task_index) { + recording.unsupported = true; + return; + } + if (dep_index < 0) { + const bool represented_by_boundary = std::any_of( + recording.boundary_tensors.begin(), recording.boundary_tensors.end(), [dep](const Tensor &tensor) { + return tensor.owner_task_id == dep; + } + ); + if (!represented_by_boundary) { + recording.unsupported = true; + return; + } + } + } + + GraphRecordedNode node; + std::copy_n(std::begin(task.kernel_id), PTO2_SUBTASK_SLOT_COUNT, node.kernel_ids.begin()); + node.active_mask = slot.active_mask; + node.task_attrs = slot.task_attrs; + node.task_attrs.set_early_resolve(false); + node.logical_block_num = slot.logical_block_num; + node.total_required_subtasks = slot.total_required_subtasks; + const uintptr_t packed_base = reinterpret_cast(task.packed_buffer_base); + const uintptr_t packed_end = reinterpret_cast(task.packed_buffer_end); + if (packed_end < packed_base) { + recording.unsupported = true; + return; + } + node.total_output_size = packed_end - packed_base; + node.record_packed_base = packed_base; + node.tensors.assign(payload.tensors, payload.tensors + payload.tensor_count); + node.tensor_sources.resize(static_cast(payload.tensor_count)); + node.scalars.assign(payload.scalars, payload.scalars + payload.scalar_count); + for (int32_t i = 0; i < payload.tensor_count; ++i) { + if (!graph_classify_tensor( + recording, node, task_index, payload.tensors[i], &node.tensor_sources[static_cast(i)] + )) { + recording.unsupported = true; + return; + } + } + for (size_t producer : recording.current_fanins) { + if (producer >= static_cast(task_index)) { + recording.unsupported = true; + return; + } + node.internal_fanins.push_back(producer); + } + recording.nodes.push_back(std::move(node)); + recording.current_task_index.reset(); + recording.current_fanins.clear(); +} + +GraphBoundarySignature graph_boundary_signature(const Tensor &tensor, TensorArgType type, uint16_t alias_rep) { + GraphBoundarySignature signature{}; + signature.buffer_size = tensor.buffer.size; + std::copy(std::begin(tensor.shapes), std::end(tensor.shapes), std::begin(signature.shapes)); + std::copy(std::begin(tensor.strides), std::end(tensor.strides), std::begin(signature.strides)); + signature.alias_rep = alias_rep; + signature.ndims = static_cast(tensor.ndims); + signature.dtype = static_cast(tensor.dtype); + signature.tag = static_cast(type); + signature.manual_dep = tensor.manual_dep ? 1 : 0; + signature.is_contiguous = tensor.is_contiguous ? 1 : 0; + return signature; +} + +template +uint32_t graph_append_section(std::vector *image, const std::vector &values) { + if (values.empty()) return 0; + if (image->size() > UINT32_MAX || values.size() > UINT32_MAX / sizeof(T)) return 0; + const size_t aligned = PTO2_ALIGN_UP(image->size(), alignof(T)); + const size_t bytes = values.size() * sizeof(T); + if (aligned > UINT32_MAX || bytes > UINT32_MAX - aligned) return 0; + image->resize(aligned + bytes); + std::memcpy(image->data() + aligned, values.data(), bytes); + return static_cast(aligned); +} + +std::optional graph_pack_tensor_source(const GraphRecordedTensorSourceRef &source) { + if (source.source_index > UINT16_MAX) return std::nullopt; + + GraphTensorSourceRef packed{}; + switch (source.source) { + case GraphRecordedTensorSource::BOUNDARY_EXACT: + packed.source = static_cast(GraphTensorSource::BOUNDARY_EXACT); + break; + case GraphRecordedTensorSource::BOUNDARY_VIEW: + packed.source = static_cast(GraphTensorSource::BOUNDARY_VIEW); + break; + case GraphRecordedTensorSource::INTERNAL: + packed.source = static_cast(GraphTensorSource::INTERNAL); + break; + case GraphRecordedTensorSource::OWN_OUTPUT: + packed.source = static_cast(GraphTensorSource::OWN_OUTPUT); + break; + } + packed.source_index = static_cast(source.source_index); + packed.packed_offset = source.packed_offset; + return packed; +} + +bool graph_build_definition(const GraphRecording &recording, std::vector *image) { + if (image == nullptr || recording.unsupported || recording.nodes.empty() || + recording.nodes.size() > GRAPH_MAX_NODES || recording.boundary_tensors.size() > UINT16_MAX || + recording.boundary_tensors.size() != recording.boundary_types.size() || + std::any_of(recording.boundary_tensors.begin(), recording.boundary_tensors.end(), [](const Tensor &tensor) { + return tensor.ndims > MAX_TENSOR_DIMS; + })) { + return false; + } + + std::vector fanout_counts(recording.nodes.size(), 0); + std::vector fanin_offsets(recording.nodes.size() + 1, 0); + std::vector fanin_indices; + std::vector roots; + std::vector node_offsets(recording.nodes.size(), 0); + std::vector nodes(recording.nodes.size()); + std::vector tensors; + std::vector tensor_sources; + std::vector scalars; + + uint64_t required_heap = 0; + uint32_t edge_count = 0; + for (size_t i = 0; i < recording.nodes.size(); ++i) { + const GraphRecordedNode &source = recording.nodes[i]; + if (source.total_output_size > static_cast(INT32_MAX) || + source.tensors.size() > static_cast(INT32_MAX) || + source.scalars.size() > static_cast(INT32_MAX) || source.internal_fanins.size() > UINT16_MAX || + source.tensors.size() != source.tensor_sources.size() || + tensors.size() > UINT32_MAX - source.tensors.size() || + tensor_sources.size() > UINT32_MAX - source.tensor_sources.size() || + scalars.size() > UINT32_MAX - source.scalars.size() || + std::any_of(source.tensors.begin(), source.tensors.end(), [](const Tensor &tensor) { + return tensor.ndims > MAX_TENSOR_DIMS; + })) { + return false; + } + node_offsets[i] = required_heap; + const uint64_t output_bytes = PTO2_ALIGN_UP(source.total_output_size, PTO2_ALIGN_SIZE); + if (required_heap > UINT64_MAX - output_bytes) return false; + required_heap += output_bytes; + + fanin_offsets[i + 1] = fanin_offsets[i] + static_cast(source.internal_fanins.size()); + if (source.internal_fanins.empty()) roots.push_back(static_cast(i)); + for (size_t producer : source.internal_fanins) { + if (producer >= i) return false; + fanout_counts[producer]++; + fanin_indices.push_back(static_cast(producer)); + edge_count++; + } + + GraphNodeDefinition &node = nodes[i]; + std::copy(source.kernel_ids.begin(), source.kernel_ids.end(), std::begin(node.kernel_id)); + node.active_mask = source.active_mask.raw(); + node.task_attrs = source.task_attrs.raw(); + node.logical_block_num = source.logical_block_num; + node.total_required_subtasks = source.total_required_subtasks; + node.tensor_count = static_cast(source.tensors.size()); + node.scalar_count = static_cast(source.scalars.size()); + node.total_output_size = static_cast(source.total_output_size); + node.tensor_offset = static_cast(tensors.size()); + node.scalar_offset = static_cast(scalars.size()); + for (const Tensor &tensor : source.tensors) + tensors.push_back(graph_tensor_pack(tensor)); + for (const GraphRecordedTensorSourceRef &tensor_source : source.tensor_sources) { + std::optional packed_source = graph_pack_tensor_source(tensor_source); + if (!packed_source.has_value()) return false; + tensor_sources.push_back(*packed_source); + } + scalars.insert(scalars.end(), source.scalars.begin(), source.scalars.end()); + } + + std::vector fanout_offsets(recording.nodes.size() + 1, 0); + for (size_t i = 0; i < recording.nodes.size(); ++i) + fanout_offsets[i + 1] = fanout_offsets[i] + fanout_counts[i]; + std::vector fanout_indices(edge_count); + std::vector cursors(fanout_offsets.begin(), fanout_offsets.end() - 1); + for (size_t consumer = 0; consumer < recording.nodes.size(); ++consumer) { + for (size_t producer : recording.nodes[consumer].internal_fanins) { + fanout_indices[cursors[producer]++] = static_cast(consumer); + } + } + + std::vector signatures; + signatures.reserve(recording.boundary_tensors.size()); + for (size_t i = 0; i < recording.boundary_tensors.size(); ++i) { + uint16_t alias_rep = static_cast(i); + for (size_t j = 0; j < i; ++j) { + if (recording.boundary_tensors[j].buffer.addr == recording.boundary_tensors[i].buffer.addr && + recording.boundary_tensors[j].buffer.size == recording.boundary_tensors[i].buffer.size) { + alias_rep = static_cast(j); + break; + } + } + signatures.push_back( + graph_boundary_signature(recording.boundary_tensors[i], recording.boundary_types[i], alias_rep) + ); + } + + image->assign(sizeof(GraphDefinition), std::byte{0}); + GraphDefinition definition{}; + definition.full_key = recording.full_key; + definition.required_heap = required_heap; + definition.task_count = static_cast(nodes.size()); + definition.edge_count = edge_count; + definition.root_count = static_cast(roots.size()); + definition.boundary_count = static_cast(signatures.size()); + definition.tensor_arg_count = static_cast(tensors.size()); + definition.scalar_arg_count = static_cast(scalars.size()); + definition.off_fanout_offsets = graph_append_section(image, fanout_offsets); + definition.off_fanout_indices = graph_append_section(image, fanout_indices); + definition.off_fanin_offsets = graph_append_section(image, fanin_offsets); + definition.off_fanin_indices = graph_append_section(image, fanin_indices); + definition.off_root_indices = graph_append_section(image, roots); + definition.off_node_offsets = graph_append_section(image, node_offsets); + definition.off_nodes = graph_append_section(image, nodes); + definition.off_tensors = graph_append_section(image, tensors); + definition.off_tensor_sources = graph_append_section(image, tensor_sources); + definition.off_scalars = graph_append_section(image, scalars); + definition.off_boundary_signatures = graph_append_section(image, signatures); + if (definition.off_fanout_offsets == 0 || definition.off_fanin_offsets == 0 || definition.off_node_offsets == 0 || + definition.off_nodes == 0 || definition.off_boundary_signatures == 0 || + (!tensors.empty() && definition.off_tensors == 0) || + (!tensor_sources.empty() && definition.off_tensor_sources == 0) || + (!scalars.empty() && definition.off_scalars == 0) || + (!fanout_indices.empty() && definition.off_fanout_indices == 0) || + (!fanin_indices.empty() && definition.off_fanin_indices == 0) || + (!roots.empty() && definition.off_root_indices == 0)) { + return false; + } + definition.total_bytes = static_cast(image->size()); + std::memcpy(image->data(), &definition, sizeof(definition)); + definition.content_hash = graph_hash_bytes(1469598103934665603ULL, image->data(), image->size()); + std::memcpy(image->data(), &definition, sizeof(definition)); + return true; +} + +const GraphDefinition *graph_definition(const std::vector &image) { + if (image.size() < sizeof(GraphDefinition)) return nullptr; + const auto *definition = reinterpret_cast(image.data()); + return definition->total_bytes == image.size() ? definition : nullptr; +} + +} // namespace + +GraphHostStatePtr make_graph_host_state() { return GraphHostStatePtr{new (std::nothrow) GraphHostState{}}; } + +void GraphHostStateDeleter::operator()(GraphHostState *state) const noexcept { delete state; } + +size_t graph_host_upload_count(const GraphHostState &state) { return state.pending_uploads.size(); } + +std::optional graph_host_upload(const GraphHostState &state, size_t index) { + if (index >= state.pending_uploads.size()) return std::nullopt; + const GraphPendingUpload &upload = state.pending_uploads[index]; + if (upload.outer_slot == nullptr || upload.image.empty()) return std::nullopt; + return GraphHostUpload{upload.outer_slot, upload.image.data(), upload.image.size()}; +} + +void graph_host_set_commit_callback(GraphHostState &state, GraphHostCommitCallback callback, void *context) { + state.commit_callback = callback; + state.commit_context = context; +} + +bool graph_host_commit(GraphHostState &state) { + return state.commit_callback == nullptr || state.commit_callback(state.commit_context); +} + static uint32_t next_fanin_seen_epoch(PTO2OrchestratorState *orch) { uint32_t next = orch->fanin_seen_current_epoch + 1; if (next == 0) { @@ -309,6 +790,7 @@ static bool append_fanin_or_fail( if (fanin_builder->mark_seen(prod_ring, prod_slot)) { return true; } + graph_record_note_fanin(orch, prod_state); if (fanin_builder->count >= PTO2_MAX_FANIN) { orch_mark_fatal(orch, PTO2_ERROR_DEP_POOL_OVERFLOW); return false; @@ -402,6 +884,8 @@ static bool prepare_task( out->task = &orch->sm_header->ring.task_descriptors[out->alloc_result.slot]; out->payload = &orch->sm_header->ring.task_payloads[out->alloc_result.slot]; + graph_record_begin_task(orch, out->task_id); + out->payload->prefetch(args.tensor_count(), args.scalar_count()); // Re-bind payload/task pointers each submit. Value is per-slot constant @@ -433,6 +917,7 @@ static bool prepare_task( out->slot_state->logical_block_num = block_num; out->slot_state->active_mask = active_mask; out->slot_state->task_attrs = task_attrs; + out->slot_state->task_kind = active_mask ? PTO2TaskKind::KERNEL : PTO2TaskKind::DUMMY; // Reclaim gate: seed last_consumer to self, so a producer with no consumers // is retirable once completed_watermark >= its own id. Each fanin edge bumps // it in append_fanin_or_fail. completion_flags for this slot are already 0 @@ -677,9 +1162,13 @@ static TaskOutputTensors submit_task_common( // fanout now that the swimlane hot path no longer records it. const bool capture_dep_graph = dep_gen_host_graph_enabled(); if (capture_dep_graph) { - const int32_t kernel_ids_capture[3] = {aic_kernel_id, aiv0_kernel_id, aiv1_kernel_id}; + const std::array kernel_ids_capture{ + aic_kernel_id, + aiv0_kernel_id, + aiv1_kernel_id, + }; dep_gen_host_graph_begin_task( - task_id.raw, orch->in_manual_scope(), args.allow_early_resolve(), kernel_ids_capture, + task_id.raw, orch->in_manual_scope(), args.allow_early_resolve(), kernel_ids_capture.data(), args.launch_spec.block_num(), args.tensor_count(), args.tensor_data(), args.tag_data() ); } @@ -849,6 +1338,7 @@ static TaskOutputTensors submit_task_common( // of position-independent integers, none of this needs host->device pointer // relocation. payload.fanin_count = fanin_builder.count; + graph_record_task(orch, task_id, task, payload, *prepared.slot_state, args); (void)sched; CYCLE_COUNT_LAP(g_orch_fanin_cycle); @@ -864,6 +1354,275 @@ static TaskOutputTensors submit_task_common( return result; } +namespace { + +bool graph_boundary_matches(const GraphDefinition &definition, const L0TaskArgs &args) { + if (args.scalar_count() != 0 || args.explicit_dep_count() != 0 || + args.tensor_count() != static_cast(definition.boundary_count)) { + LOG_WARN( + "[GraphExecution] fixed boundary contract mismatch: tensors=%d/%u scalars=%d explicit_deps=%u", + args.tensor_count(), definition.boundary_count, args.scalar_count(), args.explicit_dep_count() + ); + return false; + } + const auto *signatures = graph_definition_array( + definition, definition.off_boundary_signatures, definition.boundary_count + ); + if (signatures == nullptr) return false; + + bool alias_mismatch = false; + for (int32_t i = 0; i < args.tensor_count(); ++i) { + const Tensor &tensor = args.tensor(i).ref(); + const GraphBoundarySignature &signature = signatures[i]; + if (tensor.ndims > MAX_TENSOR_DIMS) { + debug_assert(tensor.ndims <= MAX_TENSOR_DIMS && "Graph boundary Tensor rank is not supported"); + LOG_WARN("[GraphExecution] Tensor rank %u exceeds the fixed Graph boundary limit", tensor.ndims); + return false; + } + const auto shape_end = std::begin(tensor.shapes) + tensor.ndims; + const auto stride_end = std::begin(tensor.strides) + tensor.ndims; + const bool metadata_match = tensor.buffer.size == signature.buffer_size && tensor.ndims == signature.ndims && + static_cast(tensor.dtype) == signature.dtype && + static_cast(args.tag(i)) == signature.tag && + static_cast(tensor.manual_dep ? 1 : 0) == signature.manual_dep && + static_cast(tensor.is_contiguous ? 1 : 0) == signature.is_contiguous && + std::equal(std::begin(tensor.shapes), shape_end, std::begin(signature.shapes)) && + std::equal(std::begin(tensor.strides), stride_end, std::begin(signature.strides)); + if (!metadata_match) { + debug_assert(metadata_match && "Variable Graph boundary tensor shape/metadata is not supported"); + LOG_WARN( + "[GraphExecution] fixed tensor shape/metadata mismatch at boundary arg %d; using ordinary path", i + ); + return false; + } + uint16_t alias_rep = static_cast(i); + for (int32_t j = 0; j < i; ++j) { + const Tensor &other = args.tensor(j).ref(); + if (other.buffer.addr == tensor.buffer.addr && other.buffer.size == tensor.buffer.size) { + alias_rep = static_cast(j); + break; + } + } + alias_mismatch |= alias_rep != signature.alias_rep; + } + if (alias_mismatch) { + debug_assert(!alias_mismatch && "Changing the Graph boundary alias partition is not supported"); + LOG_WARN("%s", "[GraphExecution] boundary alias partition differs from recording; using ordinary path"); + return false; + } + return true; +} + +void graph_reset_outer_payload(PTO2TaskPayload &payload) { + payload.tensor_count = 0; + payload.scalar_count = 0; + payload.fanin_count = 0; + payload.predicate = DispatchPredicate{}; + payload.early_dispatch_state.store(PTO2_EARLY_DISPATCH_NONE, std::memory_order_relaxed); + for (auto &word : payload.staged_core_mask) + word.store(0, std::memory_order_relaxed); + payload.dispatch_fanin.store(0, std::memory_order_relaxed); + payload.dispatch_propagated.store(0, std::memory_order_relaxed); + payload.published_block_count.store(0, std::memory_order_relaxed); + payload.early_dispatch_launch_state.store(PTO2_EARLY_DISPATCH_LAUNCH_NONE, std::memory_order_relaxed); + payload.running_slot_count.store(0, std::memory_order_relaxed); + payload.early_sync_drain_state.store(PTO2_EARLY_SYNC_DRAIN_NONE, std::memory_order_relaxed); +} + +bool graph_build_submission_image( + const std::vector &definition_image, const L0TaskArgs &args, std::vector *submission_image +) { + if (submission_image == nullptr || graph_definition(definition_image) == nullptr) return false; + const size_t definition_offset = PTO2_ALIGN_UP(sizeof(GraphSubmission), alignof(GraphDefinition)); + const size_t tensors_offset = PTO2_ALIGN_UP(definition_offset + definition_image.size(), alignof(GraphTensor)); + const size_t tensor_bytes = static_cast(args.tensor_count()) * sizeof(GraphTensor); + if (definition_offset > UINT32_MAX || tensors_offset > UINT32_MAX || tensors_offset > UINT32_MAX - tensor_bytes) { + return false; + } + submission_image->assign(tensors_offset + tensor_bytes, std::byte{0}); + std::memcpy(submission_image->data() + definition_offset, definition_image.data(), definition_image.size()); + auto *tensors = reinterpret_cast(submission_image->data() + tensors_offset); + for (int32_t i = 0; i < args.tensor_count(); ++i) + tensors[i] = graph_tensor_pack(args.tensor(i).ref()); + + const GraphDefinition &definition = *graph_definition(definition_image); + GraphSubmission submission{}; + submission.graph_key = definition.full_key; + submission.total_bytes = static_cast(submission_image->size()); + submission.definition_offset = static_cast(definition_offset); + submission.tensors_offset = static_cast(tensors_offset); + submission.tensor_count = static_cast(args.tensor_count()); + std::memcpy(submission_image->data(), &submission, sizeof(submission)); + return true; +} + +bool graph_submit_definition( + PTO2OrchestratorState *orch, GraphHostState *state, const std::vector &definition_image, + const L0TaskArgs &args, PTO2TaskId *submitted_id +) { + const GraphDefinition *definition = graph_definition(definition_image); + if (definition == nullptr || !graph_boundary_matches(*definition, args) || + definition->required_heap > static_cast(INT32_MAX)) { + return false; + } + auto &allocator = orch->ring.task_allocator; + if (allocator.active_count() + 1 >= allocator.window_size() || + definition->required_heap > allocator.heap_available()) { + LOG_WARN("%s", "[GraphExecution] task-window/heap preflight failed; using ordinary path"); + return false; + } + + GraphPendingUpload pending; + if (!graph_build_submission_image(definition_image, args, &pending.image)) return false; + + DepInputs boundary_inputs{ + args.tensor_count(), args.tensor_data(), args.tag_data(), 0, nullptr, + }; + const int32_t tensormap_needed = count_registrable_outputs(boundary_inputs, orch->in_manual_scope()); + if (tensormap_needed > 0 && !ensure_tensormap_capacity(orch, tensormap_needed)) return false; + if (!check_scope_can_accept_task(orch, allocator, 0)) return false; + + const PTO2TaskAllocResult allocation = allocator.alloc(static_cast(definition->required_heap)); + if (allocation.failed()) { + orch_mark_fatal(orch, PTO2_ERROR_HEAP_RING_DEADLOCK); + return false; + } + const PTO2TaskId task_id = PTO2TaskId::make(0, static_cast(allocation.task_id)); + PTO2SharedMemoryRingHeader &ring = orch->sm_header->ring; + PTO2TaskDescriptor &task = ring.task_descriptors[allocation.slot]; + PTO2TaskPayload &payload = ring.task_payloads[allocation.slot]; + PTO2TaskSlotState &slot = ring.get_slot_state_by_slot(allocation.slot); + + slot.bind_buffers(&payload, &task); + slot.task_state.store(PTO2_TASK_PENDING, std::memory_order_relaxed); + slot.last_consumer_local_id = static_cast(task_id.local()); + slot.active_mask = ActiveMask{}; + slot.task_attrs = TaskAttrs{}; + slot.total_required_subtasks = 0; + slot.logical_block_num = 1; + slot.task_kind = PTO2TaskKind::GRAPH; + slot.graph_context = nullptr; + scope_tasks_push(orch, &slot); + + task.task_id = task_id; + std::fill(std::begin(task.kernel_id), std::end(task.kernel_id), INVALID_KERNEL_ID); + task.packed_buffer_base = allocation.packed_base; + task.packed_buffer_end = allocation.packed_end; + graph_reset_outer_payload(payload); + + PTO2FaninBuilder fanin_builder(orch, &payload, static_cast(task_id.local()), next_fanin_seen_epoch(orch)); + orch->tensor_map.sync_tensormap(task_id, ring.fc.last_task_alive.load(std::memory_order_acquire)); + auto emit = [&](PTO2TaskId producer_id) -> bool { + const int32_t producer_local = static_cast(producer_id.local()); + const int32_t producer_slot = ring.get_slot_by_task_id(producer_local); + PTO2TaskSlotState *producer = &ring.get_slot_state_by_slot(producer_slot); + return append_fanin_or_fail(orch, producer_id.ring(), producer_slot, producer, producer_id, &fanin_builder); + }; + if (!compute_task_fanin(boundary_inputs, orch->tensor_map, orch->in_manual_scope(), emit)) return false; + register_task_outputs(boundary_inputs, task_id, orch->tensor_map, orch->in_manual_scope()); + payload.fanin_count = fanin_builder.count; + + pending.outer_slot = &slot; + state->pending_uploads.push_back(std::move(pending)); + if (submitted_id != nullptr) *submitted_id = task_id; +#if SIMPLER_DFX + orch->tasks_submitted++; +#endif + return true; +} + +} // namespace + +GraphScopeResult +PTO2OrchestratorState::graph_begin(uint64_t graph_key, const L0TaskArgs &args, uint64_t callable_hash) { + auto *orch = this; + GraphScopeResult result; + GraphHostState *state = graph_state_from(orch); + if (state == nullptr || !rt_graph_args_cacheable(args) || args.explicit_dep_count() != 0) { + debug_assert(args.scalar_count() == 0 && "Graph execution scalars are not supported in step 1"); + debug_assert(args.explicit_dep_count() == 0 && "Graph boundary explicit dependencies are not supported"); + return result; + } + if (state->recording != nullptr) { + state->recording->unsupported = true; + debug_assert(state->recording == nullptr && "Nested Graph recording is not supported"); + LOG_WARN("%s", "[GraphExecution] nested Graph recording is not supported"); + return result; + } + + const uint64_t full_key = graph_full_key(callable_hash, graph_key); + auto definition_it = state->definitions.find(full_key); + if (definition_it != state->definitions.end()) { +#if SIMPLER_DFX + const uint64_t start = orch->host_orch_phase_records == nullptr ? 0 : get_sys_cnt_aicpu(); +#endif + PTO2TaskId submitted = PTO2TaskId::invalid(); + if (graph_submit_definition(orch, state, definition_it->second, args, &submitted)) { + result.execute_block = false; + result.task_id = submitted; +#if SIMPLER_DFX + if (start != 0) { + record_host_orch_phase(orch, start, get_sys_cnt_aicpu(), submitted.raw, g_orch_submit_idx); + } + g_orch_submit_idx++; +#if SIMPLER_ORCH_PROFILING + g_orch_submit_count++; +#endif +#endif + } + return result; + } + if (state->definitions.size() >= GRAPH_MAX_DEFINITIONS) { + debug_assert( + state->definitions.size() < GRAPH_MAX_DEFINITIONS && + "Graph Definition cache exceeds the supported per-worker limit" + ); + LOG_WARN( + "[GraphExecution] Definition cache is full (%zu entries); using ordinary path", state->definitions.size() + ); + return result; + } + + auto recording = std::make_unique(); + recording->full_key = full_key; + recording->start_local_task_id = orch->ring.task_allocator.active_count(); + recording->boundary_tensors.reserve(static_cast(args.tensor_count())); + recording->boundary_types.reserve(static_cast(args.tensor_count())); + for (int32_t i = 0; i < args.tensor_count(); ++i) { + recording->boundary_tensors.push_back(args.tensor(i).ref()); + recording->boundary_types.push_back(args.tag(i)); + } + state->recording = std::move(recording); + result.recording = true; + return result; +} + +void PTO2OrchestratorState::graph_end() { + GraphHostState *state = graph_state_from(this); + if (state == nullptr || state->recording == nullptr) return; + std::unique_ptr recording = std::move(state->recording); + std::vector definition; + if (!graph_build_definition(*recording, &definition)) { + debug_assert(false && "The recorded Graph contains a construct that Graph Execution does not support"); + LOG_WARN("%s", "[GraphExecution] unsupported construct observed; definition was not cached"); + return; + } + const GraphDefinition *header = graph_definition(definition); + if (header == nullptr) return; + LOG_DEBUG( + "[GraphExecution] define key=0x%llx nodes=%u bytes=%u", static_cast(header->full_key), + header->task_count, header->total_bytes + ); + state->definitions.emplace(header->full_key, std::move(definition)); +} + +void PTO2OrchestratorState::graph_commit() { + GraphHostState *state = graph_state_from(this); + if (state != nullptr && !graph_host_commit(*state)) { + report_fatal(PTO2_ERROR_EXPLICIT_ORCH_FATAL, __FUNCTION__, "failed to publish committed Graph prefix"); + } +} + TaskOutputTensors PTO2OrchestratorState::submit_task(const MixedKernels &mixed_kernels, const L0TaskArgs &args) { auto *orch = this; @@ -976,6 +1735,7 @@ TaskOutputTensors PTO2OrchestratorState::submit_dummy_task(const L0TaskArgs &arg TaskOutputTensors PTO2OrchestratorState::alloc_tensors(const L0TaskArgs &args) { auto *orch = this; + graph_record_mark_unsupported(orch); // Orchestration API should short-circuit after fatal, but keep this entry // robust as a no-op in case a caller reaches it directly. if (orch->fatal) { diff --git a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp index 41e7396cdf..c105f89291 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp @@ -78,6 +78,19 @@ static TaskOutputTensors submit_dummy_task_impl(PTO2Runtime *rt, const L0TaskArg return rt->orchestrator.submit_dummy_task(args); } +static GraphScopeResult graph_begin_impl(PTO2Runtime *rt, uint64_t graph_key, const L0TaskArgs &args) { + if (rt == nullptr) return GraphScopeResult{}; + return rt->orchestrator.graph_begin(graph_key, args, rt->active_callable_hash); +} + +static void graph_end_impl(PTO2Runtime *rt) { + if (rt != nullptr) rt->orchestrator.graph_end(); +} + +static void graph_commit_impl(PTO2Runtime *rt) { + if (rt != nullptr) rt->orchestrator.graph_commit(); +} + void rt_scope_begin(PTO2Runtime *rt) { PTO2ScopeMode mode = rt->pending_scope_mode; rt->pending_scope_mode = PTO2ScopeMode::AUTO; @@ -305,6 +318,9 @@ static const PTO2RuntimeOps s_runtime_ops = { .submit_dummy_task = submit_dummy_task_impl, .available_cluster_count = available_cluster_count_impl, .available_aiv_count = available_aiv_count_impl, + .graph_begin = graph_begin_impl, + .graph_end = graph_end_impl, + .graph_commit = graph_commit_impl, #if SIMPLER_DFX .scope_set_site = scope_set_site_impl, #else diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_orchestrator.h b/src/a2a3/runtime/host_build_graph/runtime/pto_orchestrator.h index 3582a8875d..aa1d75dd7c 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_orchestrator.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_orchestrator.h @@ -31,6 +31,7 @@ #include "common/l2_swimlane_profiling.h" #include "utils/device_arena.h" #include "pto_ring_buffer.h" +#include "graph_cache.h" #include "pto_runtime2_types.h" #include "pto_submit_types.h" #include "scheduler/pto_scheduler.h" @@ -38,6 +39,8 @@ #include "pto_tensormap.h" #include "pto_types.h" +struct GraphHostState; + /** * Layout descriptor produced by PTO2OrchestratorState::reserve_layout(). Holds * arena offsets for every sub-region the orchestrator owns (per-ring fanin @@ -113,11 +116,23 @@ struct PTO2OrchestratorState { // after orchestration finishes so shutdown/profiling totals remain closed. int64_t inline_completed_tasks{0}; + // This host-only state is cleared before the arena/shared-memory image + // crosses to the device. + GraphHostState *graph_host_state{nullptr}; + // === STATISTICS === #if SIMPLER_DFX int64_t tasks_submitted; int64_t buffers_allocated; int64_t bytes_allocated; + + // Host-build-graph runs submit_task through this state before the device + // collector exists. Point directly at the per-run host capture vector so + // submit records cannot be lost through ELF symbol interposition between + // the host runtime and simulator/AICPU DSOs. + L2SwimlaneAicpuOrchPhaseRecord *host_orch_phase_records{nullptr}; + size_t host_orch_phase_capacity{0}; + size_t host_orch_phase_count{0}; #endif bool in_manual_scope() const { return scope_stack_top >= manual_begin_depth; } @@ -154,6 +169,9 @@ struct PTO2OrchestratorState { TaskOutputTensors submit_task(const MixedKernels &mixed_kernels, const L0TaskArgs &args); TaskOutputTensors submit_dummy_task(const L0TaskArgs &args); TaskOutputTensors alloc_tensors(const L0TaskArgs &args); + GraphScopeResult graph_begin(uint64_t graph_key, const L0TaskArgs &args, uint64_t callable_hash); + void graph_end(); + void graph_commit(); void mark_done(); }; diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2.h b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2.h index 1fc66c88cb..2f9002826f 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2.h @@ -36,6 +36,7 @@ #include "utils/device_arena.h" #include "pto_runtime2_types.h" +#include "graph_cache.h" #include "pto_submit_types.h" #include "pto_shared_memory.h" #include "pto_ring_buffer.h" @@ -89,11 +90,13 @@ struct PTO2RuntimeOps { ); TaskOutputTensors (*alloc_tensors)(PTO2Runtime *rt, const L0TaskArgs &args); TaskOutputTensors (*submit_dummy_task)(PTO2Runtime *rt, const L0TaskArgs &args); - // This-run core geometry from runtime_finalize_after_wire: MIX clusters // (one AIC each) and standalone AIV cores. int32_t (*available_cluster_count)(PTO2Runtime *rt); int32_t (*available_aiv_count)(PTO2Runtime *rt); + GraphScopeResult (*graph_begin)(PTO2Runtime *rt, uint64_t graph_key, const L0TaskArgs &args); + void (*graph_end)(PTO2Runtime *rt); + void (*graph_commit)(PTO2Runtime *rt); // Stash the call-site captured by PTO2ScopeGuard into the [ScopeStats] // collector. Always present in the struct to keep ops-table layout stable // across SIMPLER_DFX settings; set to nullptr at SIMPLER_DFX=0. @@ -150,6 +153,9 @@ struct PTO2Runtime { // Statistics int64_t total_cycles; + // Graph definitions are process-local host cache entries. The callable + // identity prevents two orchestration DSOs from sharing the same key. + uint64_t active_callable_hash; // Prebuilt-arena fast path metadata. Carries every offset // wire_arena_pointers needs at AICPU boot so the AICPU can reconstruct diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h index 933c14b1a4..111f853fd8 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h @@ -72,9 +72,9 @@ // Use pto2_task_slot(sched, task_id) for slot calculation. #define PTO2_TASK_WINDOW_SIZE 16384 // Default per-ring task window size (power of 2) -// Single ring. host_build_graph is host-orch: the whole graph is built on the -// host, fits one ring, and the device runs it once without reclaim (see stages -// 1-2 — execution-time recycle removed). The multi-ring design existed only to +// Single ring. host_build_graph publishes the graph incrementally from the +// host, keeps the complete run resident in one ring, and executes it without +// reclaim. The multi-ring design existed only to // let inner scopes reclaim independently under small rings; with no reclaim and // a whole-graph-resident ring, per-depth isolation is moot, so all scope depths // map to the single ring 0 (0 == 0). @@ -165,6 +165,13 @@ struct PTO2TaskAllocResult { bool failed() const { return task_id < 0; } }; +enum class PTO2TaskKind : uint8_t { + KERNEL = 0, + DUMMY = 1, + GRAPH = 2, + GRAPH_NODE = 3, +}; + struct PTO2OutputLayout { uint64_t offsets[MAX_TENSOR_ARGS] = {}; uint64_t buffer_sizes[MAX_TENSOR_ARGS] = {}; @@ -480,7 +487,7 @@ struct alignas(64) PTO2TaskSlotState { // MPSC-deferred completion. The release write is sequenced before // on_subtask_complete's acq_rel fetch_add and the acquire read after. std::atomic any_subtask_deferred{false}; - uint8_t _async_pad{0}; + PTO2TaskKind task_kind{PTO2TaskKind::KERNEL}; std::atomic completed_subtasks{0}; // Each core completion increments by 1 int16_t total_required_subtasks{0}; // = logical_block_num * popcount(active_mask) @@ -490,6 +497,14 @@ struct alignas(64) PTO2TaskSlotState { // ranges through claim_block_range(). std::atomic next_block_idx{0}; + // Graph-only scheduling metadata occupies the former tail padding, keeping + // the slot state at one cache line and preserving the 40-byte descriptor + // ABI consumed by AICore. Readiness uses the shared intrusive wake-list + // fields above; this index identifies the node in the saved fanin CSR. + // Ordinary ring tasks leave both Graph fields -1/null. + int32_t graph_node_index{-1}; + void *graph_context{nullptr}; + int32_t claim_block_range(int32_t block_limit, int32_t max_count, int32_t &start) { int16_t current = next_block_idx.load(std::memory_order_relaxed); while (current < block_limit && max_count > 0) { @@ -535,6 +550,9 @@ struct alignas(64) PTO2TaskSlotState { any_subtask_deferred.store(false, std::memory_order_relaxed); completed_subtasks.store(0, std::memory_order_relaxed); next_block_idx.store(0, std::memory_order_relaxed); + graph_node_index = -1; + graph_context = nullptr; + task_kind = PTO2TaskKind::KERNEL; // Note: active_mask and task_attrs are per-submit-constant fields // rewritten in prepare_task on every reuse, so they are not reset here. // last_consumer_local_id is seeded in prepare_task once the id is known. diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h b/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h index bd404c3bc2..cc1337b45c 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h @@ -281,6 +281,12 @@ inline std::atomic *orch_error_code_addr(void *sm_dev_base) noexcept { ); } +inline std::atomic *orchestrator_done_addr(void *sm_dev_base) noexcept { + return reinterpret_cast *>( + static_cast(sm_dev_base) + offsetof(PTO2SharedMemoryHeader, orchestrator_done) + ); +} + inline PTO2SharedMemoryRingHeader *ring_header_addr(void *sm_dev_base) noexcept { return reinterpret_cast( static_cast(sm_dev_base) + offsetof(PTO2SharedMemoryHeader, ring) diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_submit_types.h b/src/a2a3/runtime/host_build_graph/runtime/pto_submit_types.h index 9ec8b647b8..92b14508fc 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_submit_types.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_submit_types.h @@ -183,6 +183,10 @@ static_assert(sizeof(ActiveMask) == 1, "ActiveMask must be exactly 1 byte"); class TaskAttrs { public: constexpr TaskAttrs() = default; + constexpr explicit TaskAttrs(uint8_t raw) : + raw_(raw) {} + + uint8_t raw() const { return raw_; } bool allow_early_resolve() const { return (raw_ & BIT_EARLY_RESOLVE) != 0; } void set_early_resolve(bool v) { diff --git a/src/a2a3/runtime/host_build_graph/runtime/runtime.h b/src/a2a3/runtime/host_build_graph/runtime/runtime.h index c6e28784e7..1322bafb33 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/runtime.h +++ b/src/a2a3/runtime/host_build_graph/runtime/runtime.h @@ -177,11 +177,8 @@ class Runtime { // NOTE: Made public for direct access from aicore code uint64_t func_id_to_addr_[RUNTIME_MAX_FUNC_ID]; - // Total tasks submitted by the host orchestrator — handed to the scheduler - // (SchedulerContext::on_orchestration_done) in place of latching the SM ring - // head on device. host_build_graph builds the whole graph on the host, so - // the boot thread reads this instead of counting SM ring heads. - int32_t host_total_tasks; + // Legacy task-count field retained in the shared platform ABI. Streaming + // host orchestration publishes the live count through the SM ring head. private: // Kernel binary tracking for cleanup @@ -212,6 +209,11 @@ class Runtime { char device_orch_func_name_[RUNTIME_MAX_ORCH_SYMBOL_NAME]; char device_orch_config_name_[RUNTIME_MAX_ORCH_SYMBOL_NAME]; + // Host-only owner for the deferred orchestration/publisher state. The raw + // pointer is copied as inert bytes in the device Runtime image and is never + // dereferenced outside the host runtime. + void *deferred_host_orchestration_{nullptr}; + public: /** * Constructor - zero-initialize all arrays @@ -276,6 +278,13 @@ class Runtime { void set_device_orch_func_name(const char *name); void set_device_orch_config_name(const char *name); + bool has_deferred_host_orchestration() const { return deferred_host_orchestration_ != nullptr; } + void notify_deferred_host_execution_started(); + int run_deferred_host_orchestration(const HostApi *api); + void release_deferred_host_orchestration(const HostApi *api); + void set_deferred_host_orchestration(void *state) { deferred_host_orchestration_ = state; } + void *get_deferred_host_orchestration() const { return deferred_host_orchestration_; } + uint64_t get_function_bin_addr(int func_id) const; void set_function_bin_addr(int func_id, uint64_t addr); /** @@ -308,6 +317,22 @@ class Runtime { // garbage, identical to host_api above. No fixed cap — grows with the // chip-level entry-tensor count. std::vector tensor_pairs_; + + // Host-build-graph runs orchestration before the device collector starts. + // Preserve its host-clock envelope and per-submit records here, then hand + // both to the collector during DeviceRunner::run. These fields are + // host-only like tensor_pairs_. Keep them in the ABI even when a + // translation unit is built without DFX: + // platform and runtime objects are compiled with different profiling + // defines but share this placement-new'd Runtime object. + std::vector host_orch_phase_records_; + uint64_t host_orch_start_cycles_{0}; + uint64_t host_orch_end_cycles_{0}; + + const std::vector &get_host_orch_phase_records() const; + uint64_t get_host_orch_start_cycles() const; + uint64_t get_host_orch_end_cycles() const; + uint64_t get_host_orch_first_publish_cycles() const; }; // Number of bytes of the Runtime image that must be copied to the device. diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/graph_execution.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/graph_execution.cpp new file mode 100644 index 0000000000..beb3f3bec6 --- /dev/null +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/graph_execution.cpp @@ -0,0 +1,618 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include "graph_execution.h" + +#include +#include +#include +#include +#include + +#include "graph_cache.h" + +namespace { + +GraphExecution *g_graph_executions = nullptr; +GraphExecution *g_graph_execution_pool = nullptr; +size_t g_graph_execution_pool_bytes = 0; +int32_t g_graph_execution_pool_blocks = 0; +size_t g_graph_execution_total_bytes = 0; +int32_t g_graph_execution_total_blocks = 0; +std::atomic_flag g_graph_execution_lock = ATOMIC_FLAG_INIT; + +constexpr size_t GRAPH_EXECUTION_POOL_MAX_BYTES = 16ULL * 1024 * 1024; +constexpr int32_t GRAPH_EXECUTION_POOL_MAX_BLOCKS = 64; + +struct GraphExecutionLockGuard { + GraphExecutionLockGuard() { + while (g_graph_execution_lock.test_and_set(std::memory_order_acquire)) {} + } + ~GraphExecutionLockGuard() { g_graph_execution_lock.clear(std::memory_order_release); } +}; + +bool checked_align_up(size_t value, size_t alignment, size_t *result) { + if (alignment == 0 || value > std::numeric_limits::max() - (alignment - 1)) return false; + *result = (value + alignment - 1) & ~(alignment - 1); + return true; +} + +bool graph_execution_layout( + int32_t node_capacity, size_t definition_capacity, size_t *nodes_offset, size_t *definition_offset, + size_t *allocation_bytes +) { + if (node_capacity <= 0 || + static_cast(node_capacity) > std::numeric_limits::max() / sizeof(GraphNodeStorage)) { + return false; + } + const size_t nodes_bytes = static_cast(node_capacity) * sizeof(GraphNodeStorage); + if (!checked_align_up(sizeof(GraphExecution), alignof(GraphNodeStorage), nodes_offset) || + *nodes_offset > std::numeric_limits::max() - nodes_bytes || + !checked_align_up(*nodes_offset + nodes_bytes, alignof(GraphDefinition), definition_offset) || + *definition_offset > std::numeric_limits::max() - definition_capacity) { + return false; + } + return checked_align_up(*definition_offset + definition_capacity, alignof(GraphNodeStorage), allocation_bytes); +} + +void destroy_execution_nodes(GraphExecution *execution) { + for (int32_t i = 0; i < execution->constructed_nodes; ++i) { + execution->node_storage[i].~GraphNodeStorage(); + } + execution->constructed_nodes = 0; +} + +void reset_execution(GraphExecution *execution, int32_t node_count, uint64_t graph_key, uint64_t definition_hash) { + size_t nodes_offset = 0; + size_t definition_offset = 0; + size_t bytes = 0; + if (!graph_execution_layout( + execution->node_capacity, execution->definition_capacity, &nodes_offset, &definition_offset, &bytes + )) { + return; + } + execution->definition_affine_reuse = graph_key != 0 && execution->materialized_graph_key == graph_key && + execution->materialized_definition_hash == definition_hash && + execution->materialized_node_count == node_count && + execution->constructed_nodes >= node_count; + execution->state.store(GraphExecutionState::SUBMITTED, std::memory_order_relaxed); + execution->materialize_busy.store(0, std::memory_order_relaxed); + execution->remaining_nodes.store(node_count, std::memory_order_relaxed); + execution->retired_nodes.store(0, std::memory_order_relaxed); + execution->node_count = node_count; + execution->materialized_nodes = 0; + execution->allocation_bytes = bytes; + execution->graph_key = graph_key; + execution->definition_hash = definition_hash; + execution->outer_slot = nullptr; + execution->nodes = nullptr; + execution->node_storage = + reinterpret_cast(reinterpret_cast(execution) + nodes_offset); + execution->definition_storage = reinterpret_cast(execution) + definition_offset; + if (!execution->definition_affine_reuse) execution->definition = nullptr; + execution->fanin_offsets = nullptr; + execution->fanin_indices = nullptr; + execution->boundary_tensors = nullptr; + execution->boundary_tensor_count = 0; + execution->next = nullptr; +} + +void destroy_execution(GraphExecution *execution) { + if (execution == nullptr) return; + destroy_execution_nodes(execution); + execution->~GraphExecution(); + std::free(execution); +} + +void destroy_execution_locked(GraphExecution *execution) { + if (execution == nullptr) return; + const size_t bytes = execution->allocation_bytes; + if (g_graph_execution_total_blocks > 0) g_graph_execution_total_blocks--; + if (bytes <= g_graph_execution_total_bytes) g_graph_execution_total_bytes -= bytes; + destroy_execution(execution); +} + +void recycle_execution_locked(GraphExecution *execution) { + if (execution == nullptr) return; + const size_t bytes = execution->allocation_bytes; + if (bytes == 0 || bytes > GRAPH_EXECUTION_POOL_MAX_BYTES || + g_graph_execution_pool_blocks >= GRAPH_EXECUTION_POOL_MAX_BLOCKS || + g_graph_execution_pool_bytes > GRAPH_EXECUTION_POOL_MAX_BYTES - bytes) { + destroy_execution_locked(execution); + return; + } + execution->outer_slot = nullptr; + execution->nodes = nullptr; + execution->boundary_tensors = nullptr; + execution->next = g_graph_execution_pool; + g_graph_execution_pool = execution; + g_graph_execution_pool_bytes += bytes; + g_graph_execution_pool_blocks++; +} + +GraphExecution *take_pooled_execution_locked( + int32_t node_count, uint64_t graph_key, uint64_t definition_hash, size_t definition_bytes +) { + GraphExecution **best_link = nullptr; + GraphExecution *best = nullptr; + bool best_is_affine = false; + for (GraphExecution **link = &g_graph_execution_pool; *link != nullptr; link = &(*link)->next) { + GraphExecution *candidate = *link; + if (candidate->node_capacity < node_count || candidate->definition_capacity < definition_bytes) continue; + const bool affine = graph_key != 0 && candidate->materialized_graph_key == graph_key && + candidate->materialized_definition_hash == definition_hash && + candidate->materialized_node_count == node_count && + candidate->constructed_nodes >= node_count; + if (best == nullptr || (affine && !best_is_affine) || + (affine == best_is_affine && candidate->allocation_bytes < best->allocation_bytes)) { + best = candidate; + best_link = link; + best_is_affine = affine; + } + } + if (best == nullptr) return nullptr; + *best_link = best->next; + g_graph_execution_pool_bytes -= best->allocation_bytes; + g_graph_execution_pool_blocks--; + reset_execution(best, node_count, graph_key, definition_hash); + return best; +} + +void reset_graph_payload(PTO2TaskPayload &payload) { + payload.fanin_count = 0; + payload.predicate = DispatchPredicate{}; + payload.early_dispatch_state.store(PTO2_EARLY_DISPATCH_NONE, std::memory_order_relaxed); + for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; ++w) { + payload.staged_core_mask[w].store(0, std::memory_order_relaxed); + } + payload.dispatch_fanin.store(0, std::memory_order_relaxed); + payload.dispatch_propagated.store(0, std::memory_order_relaxed); + payload.published_block_count.store(0, std::memory_order_relaxed); + payload.early_dispatch_launch_state.store(PTO2_EARLY_DISPATCH_LAUNCH_NONE, std::memory_order_relaxed); + payload.running_slot_count.store(0, std::memory_order_relaxed); + payload.early_sync_drain_state.store(PTO2_EARLY_SYNC_DRAIN_NONE, std::memory_order_relaxed); +} + +bool bind_graph_topology(GraphExecution &execution) { + if (execution.definition == nullptr) return false; + const GraphDefinition &definition = *execution.definition; + const uint32_t *fanin_offsets = + graph_definition_array(definition, definition.off_fanin_offsets, definition.task_count + 1); + const uint16_t *fanin_indices = + definition.edge_count == 0 ? + nullptr : + graph_definition_array(definition, definition.off_fanin_indices, definition.edge_count); + const uint32_t *fanout_offsets = + graph_definition_array(definition, definition.off_fanout_offsets, definition.task_count + 1); + const uint16_t *fanout_indices = + definition.edge_count == 0 ? + nullptr : + graph_definition_array(definition, definition.off_fanout_indices, definition.edge_count); + const uint16_t *roots = + graph_definition_array(definition, definition.off_root_indices, definition.root_count); + const GraphNodeDefinition *nodes = + graph_definition_array(definition, definition.off_nodes, definition.task_count); + const uint64_t *node_offsets = + graph_definition_array(definition, definition.off_node_offsets, definition.task_count); + if (fanin_offsets == nullptr || fanout_offsets == nullptr || roots == nullptr || nodes == nullptr || + node_offsets == nullptr || + (definition.edge_count != 0 && (fanin_indices == nullptr || fanout_indices == nullptr)) || + fanin_offsets[0] != 0 || fanout_offsets[0] != 0 || + fanin_offsets[definition.task_count] != definition.edge_count || + fanout_offsets[definition.task_count] != definition.edge_count) { + return false; + } + + uint64_t required_heap = 0; + constexpr uint8_t VALID_ACTIVE_MASK = (1U << PTO2_SUBTASK_SLOT_COUNT) - 1U; + for (uint32_t i = 0; i < definition.task_count; ++i) { + const GraphNodeDefinition &node = nodes[i]; + if (node_offsets[i] != required_heap || node.total_output_size < 0 || node.tensor_count < 0 || + node.tensor_count > MAX_TENSOR_ARGS || node.scalar_count < 0 || node.scalar_count > MAX_SCALAR_ARGS || + node.tensor_offset > definition.tensor_arg_count || + static_cast(node.tensor_count) > definition.tensor_arg_count - node.tensor_offset || + node.scalar_offset > definition.scalar_arg_count || + static_cast(node.scalar_count) > definition.scalar_arg_count - node.scalar_offset || + (node.active_mask & ~VALID_ACTIVE_MASK) != 0 || node.logical_block_num <= 0 || + node.total_required_subtasks < 0) { + return false; + } + for (int32_t slot = 0; slot < PTO2_SUBTASK_SLOT_COUNT; ++slot) { + const bool active = (node.active_mask & (1U << slot)) != 0; + if (active != (node.kernel_id[slot] != INVALID_KERNEL_ID)) return false; + } + const uint64_t output_bytes = PTO2_ALIGN_UP(static_cast(node.total_output_size), PTO2_ALIGN_SIZE); + if (output_bytes > definition.required_heap - required_heap) return false; + required_heap += output_bytes; + } + if (required_heap != definition.required_heap) return false; + + uint32_t observed_roots = 0; + for (uint32_t consumer = 0; consumer < definition.task_count; ++consumer) { + const uint32_t begin = fanin_offsets[consumer]; + const uint32_t end = fanin_offsets[consumer + 1]; + if (begin > end || end > definition.edge_count) return false; + if (begin == end) observed_roots++; + for (uint32_t edge = begin; edge < end; ++edge) { + if (fanin_indices[edge] >= consumer) return false; + } + } + if (observed_roots != definition.root_count) return false; + for (uint32_t i = 0; i < definition.root_count; ++i) { + const uint16_t root = roots[i]; + if (root >= definition.task_count || fanin_offsets[root] != fanin_offsets[root + 1]) return false; + } + for (uint32_t producer = 0; producer < definition.task_count; ++producer) { + const uint32_t begin = fanout_offsets[producer]; + const uint32_t end = fanout_offsets[producer + 1]; + if (begin > end || end > definition.edge_count) return false; + for (uint32_t edge = begin; edge < end; ++edge) { + if (fanout_indices[edge] <= producer || fanout_indices[edge] >= definition.task_count) return false; + } + } + + execution.fanin_offsets = fanin_offsets; + execution.fanin_indices = fanin_indices; + return true; +} + +bool graph_definition_hash_matches(const GraphDefinition &definition) { + if (definition.content_hash == 0 || definition.total_bytes < sizeof(GraphDefinition)) return false; + constexpr size_t HASH_OFFSET = offsetof(GraphDefinition, content_hash); + constexpr size_t HASH_END = HASH_OFFSET + sizeof(GraphDefinition::content_hash); + const auto *bytes = reinterpret_cast(&definition); + uint64_t hash = graph_hash_bytes(1469598103934665603ULL, bytes, HASH_OFFSET); + const uint64_t zero_hash = 0; + hash = graph_hash_bytes(hash, &zero_hash, sizeof(zero_hash)); + hash = graph_hash_bytes(hash, bytes + HASH_END, definition.total_bytes - HASH_END); + return hash == definition.content_hash; +} + +bool register_initial_graph_waiter(GraphExecution &execution, int32_t consumer_index) { + const uint32_t begin = execution.fanin_offsets[consumer_index]; + const uint32_t end = execution.fanin_offsets[consumer_index + 1]; + if (begin == end) return true; + + const uint16_t producer_index = execution.fanin_indices[begin]; + PTO2TaskSlotState &producer = execution.node_storage[producer_index].slot; + PTO2TaskSlotState &consumer = execution.node_storage[consumer_index].slot; + + // Orch has already wired the dependency as a producer index in fanin CSR. + // This only creates the execution-local polling subscription. Activation + // is gated on PREPARED, so no producer can close its list concurrently. + PTO2TaskSlotState *head = producer.wake_list_head.load(std::memory_order_relaxed); + if (head == WAKE_LIST_SENTINEL) return false; + consumer.next_in_wake_list = head; + producer.wake_list_head.store(&consumer, std::memory_order_relaxed); + return true; +} + +} // namespace + +GraphExecution * +graph_execution_create(int32_t node_count, uint64_t graph_key, uint64_t definition_hash, size_t definition_bytes) { + if (node_count <= 0 || node_count > static_cast(GRAPH_MAX_NODES) || definition_bytes == 0) { + return nullptr; + } + { + GraphExecutionLockGuard guard; + if (GraphExecution *pooled = + take_pooled_execution_locked(node_count, graph_key, definition_hash, definition_bytes)) { + return pooled; + } + } + + size_t nodes_offset = 0; + size_t definition_offset = 0; + size_t allocation_bytes = 0; + if (!graph_execution_layout(node_count, definition_bytes, &nodes_offset, &definition_offset, &allocation_bytes)) { + return nullptr; + } + { + GraphExecutionLockGuard guard; + if (allocation_bytes > GRAPH_EXECUTION_POOL_MAX_BYTES || + g_graph_execution_total_blocks >= GRAPH_EXECUTION_POOL_MAX_BLOCKS || + g_graph_execution_total_bytes > GRAPH_EXECUTION_POOL_MAX_BYTES - allocation_bytes) { + return nullptr; + } + g_graph_execution_total_bytes += allocation_bytes; + g_graph_execution_total_blocks++; + } + void *storage = nullptr; + if (::posix_memalign(&storage, alignof(GraphNodeStorage), allocation_bytes) != 0) { + GraphExecutionLockGuard guard; + g_graph_execution_total_bytes -= allocation_bytes; + g_graph_execution_total_blocks--; + return nullptr; + } + auto *execution = new (storage) GraphExecution{}; + execution->node_capacity = node_count; + execution->definition_capacity = definition_bytes; + execution->allocation_bytes = allocation_bytes; + reset_execution(execution, node_count, graph_key, definition_hash); + return execution; +} + +void graph_execution_discard(GraphExecution *execution) { + GraphExecutionLockGuard guard; + recycle_execution_locked(execution); +} + +void graph_execution_publish(GraphExecution *execution) { + if (execution == nullptr) return; + GraphExecutionLockGuard guard; + execution->next = g_graph_executions; + g_graph_executions = execution; +} + +void graph_execution_collect_retired() { + GraphExecutionLockGuard guard; + GraphExecution **link = &g_graph_executions; + while (*link != nullptr) { + GraphExecution *execution = *link; + const bool completed = execution->state.load(std::memory_order_acquire) == GraphExecutionState::COMPLETED; + const bool retired = execution->retired_nodes.load(std::memory_order_acquire) >= execution->node_count; + if (!completed || !retired) { + link = &execution->next; + continue; + } + *link = execution->next; + recycle_execution_locked(execution); + } +} + +GraphExecution *graph_execution_localize(PTO2TaskSlotState &outer_slot) { + GraphSubmission *submission = graph_submission_from_slot(outer_slot); + if (submission == nullptr) return nullptr; + if (GraphExecution *existing = graph_submission_local_execution(*submission)) return existing; + + const GraphDefinition *definition = graph_submission_definition(*submission); + const GraphTensor *boundary_tensors = graph_submission_tensors(*submission); + if (definition == nullptr || definition->total_bytes == 0 || definition->task_count == 0 || + definition->task_count > GRAPH_MAX_NODES || + definition->total_bytes > submission->total_bytes - submission->definition_offset || + submission->graph_key != definition->full_key || submission->tensor_count != definition->boundary_count || + boundary_tensors == nullptr || + submission->tensors_offset < submission->definition_offset + definition->total_bytes || + !graph_definition_hash_matches(*definition) || outer_slot.task == nullptr || + outer_slot.task->packed_buffer_base == nullptr || outer_slot.task->packed_buffer_end == nullptr) { + return nullptr; + } + const uintptr_t outer_base = reinterpret_cast(outer_slot.task->packed_buffer_base); + const uintptr_t outer_end = reinterpret_cast(outer_slot.task->packed_buffer_end); + if (outer_end < outer_base || definition->required_heap > outer_end - outer_base) return nullptr; + for (uint32_t i = 0; i < submission->tensor_count; ++i) { + if (!graph_tensor_wire_valid(boundary_tensors[i])) return nullptr; + } + + graph_execution_collect_retired(); + GraphExecution *execution = graph_execution_create( + static_cast(definition->task_count), submission->graph_key, definition->content_hash, + definition->total_bytes + ); + if (execution == nullptr) return nullptr; + + if (!execution->definition_affine_reuse) { + std::memcpy(execution->definition_storage, definition, definition->total_bytes); + execution->definition = static_cast(execution->definition_storage); + } + if (!bind_graph_topology(*execution)) { + graph_execution_discard(execution); + return nullptr; + } + execution->boundary_tensors = boundary_tensors; + execution->boundary_tensor_count = submission->tensor_count; + execution->outer_slot = &outer_slot; + + uint64_t expected = 0; + const uint64_t desired = static_cast(reinterpret_cast(execution)); + if (!__atomic_compare_exchange_n( + &submission->local_execution, &expected, desired, false, __ATOMIC_RELEASE, __ATOMIC_ACQUIRE + )) { + graph_execution_discard(execution); + return reinterpret_cast(static_cast(expected)); + } + graph_execution_publish(execution); + return execution; +} + +GraphMaterializeResult graph_execution_materialize_slice( + PTO2TaskSlotState &outer_slot, GraphExecution &execution, int32_t max_nodes, int32_t *nodes_materialized +) { + if (nodes_materialized != nullptr) *nodes_materialized = 0; + if (outer_slot.task_kind != PTO2TaskKind::GRAPH || max_nodes <= 0 || execution.definition == nullptr || + execution.node_storage == nullptr) { + return GraphMaterializeResult::INVALID; + } + + GraphExecutionState state = execution.state.load(std::memory_order_acquire); + if (state >= GraphExecutionState::PREPARED) return GraphMaterializeResult::PREPARED; + + uint8_t expected_busy = 0; + if (!execution.materialize_busy.compare_exchange_strong( + expected_busy, 1, std::memory_order_acq_rel, std::memory_order_acquire + )) { + return GraphMaterializeResult::BUSY; + } + + state = execution.state.load(std::memory_order_acquire); + if (state == GraphExecutionState::SUBMITTED) { + GraphExecutionState expected = GraphExecutionState::SUBMITTED; + if (!execution.state.compare_exchange_strong( + expected, GraphExecutionState::MATERIALIZING, std::memory_order_acq_rel, std::memory_order_acquire + )) { + execution.materialize_busy.store(0, std::memory_order_release); + return GraphMaterializeResult::BUSY; + } + } else if (state != GraphExecutionState::MATERIALIZING) { + execution.materialize_busy.store(0, std::memory_order_release); + return GraphMaterializeResult::INVALID; + } + + const GraphDefinition &definition = *execution.definition; + const auto *nodes = + graph_definition_array(definition, definition.off_nodes, definition.task_count); + const auto *node_offsets = + graph_definition_array(definition, definition.off_node_offsets, definition.task_count); + const auto *definition_tensors = + definition.tensor_arg_count == 0 ? + nullptr : + graph_definition_array(definition, definition.off_tensors, definition.tensor_arg_count); + const auto *tensor_sources = definition.tensor_arg_count == 0 ? + nullptr : + graph_definition_array( + definition, definition.off_tensor_sources, definition.tensor_arg_count + ); + const auto *definition_scalars = + definition.scalar_arg_count == 0 ? + nullptr : + graph_definition_array(definition, definition.off_scalars, definition.scalar_arg_count); + if (nodes == nullptr || node_offsets == nullptr || + (definition.tensor_arg_count != 0 && (definition_tensors == nullptr || tensor_sources == nullptr)) || + (definition.scalar_arg_count != 0 && definition_scalars == nullptr)) { + execution.materialize_busy.store(0, std::memory_order_release); + return GraphMaterializeResult::INVALID; + } + + const int32_t first = execution.materialized_nodes; + const int32_t last = std::min(execution.node_count, first + max_nodes); + for (int32_t i = first; i < last; ++i) { + const GraphNodeDefinition &source = nodes[i]; + GraphNodeStorage *storage = &execution.node_storage[i]; + if (i >= execution.constructed_nodes) { + storage = new (storage) GraphNodeStorage; + execution.constructed_nodes++; + } + PTO2TaskDescriptor &task = storage->task; + PTO2TaskPayload &payload = storage->payload; + PTO2TaskSlotState &slot = storage->slot; + + const uint32_t synthetic_local = + (static_cast(outer_slot.task->task_id.local()) << 10) | static_cast(i); + task.task_id = PTO2TaskId::make(1, synthetic_local); + for (int k = 0; k < PTO2_SUBTASK_SLOT_COUNT; ++k) + task.kernel_id[k] = source.kernel_id[k]; + task.packed_buffer_base = static_cast(outer_slot.task->packed_buffer_base) + node_offsets[i]; + task.packed_buffer_end = static_cast(task.packed_buffer_base) + + PTO2_ALIGN_UP(static_cast(source.total_output_size), PTO2_ALIGN_SIZE); + + slot.reset_for_reuse(); + slot.bind_buffers(&payload, &task); + slot.task_state.store(PTO2_TASK_PENDING, std::memory_order_relaxed); + slot.active_mask = ActiveMask(source.active_mask); + slot.task_attrs = TaskAttrs(source.task_attrs); + slot.total_required_subtasks = source.total_required_subtasks; + slot.logical_block_num = source.logical_block_num; + slot.graph_node_index = i; + slot.task_kind = PTO2TaskKind::GRAPH_NODE; + slot.graph_context = &execution; + + payload.tensor_count = source.tensor_count; + payload.scalar_count = source.scalar_count; + if (source.tensor_count < 0 || source.tensor_count > MAX_TENSOR_ARGS || source.scalar_count < 0 || + source.scalar_count > MAX_SCALAR_ARGS || + static_cast(source.tensor_count) > definition.tensor_arg_count || + static_cast(source.scalar_count) > definition.scalar_arg_count || + source.tensor_offset > definition.tensor_arg_count - static_cast(source.tensor_count) || + source.scalar_offset > definition.scalar_arg_count - static_cast(source.scalar_count)) { + execution.materialize_busy.store(0, std::memory_order_release); + return GraphMaterializeResult::INVALID; + } + for (int32_t j = 0; j < source.tensor_count; ++j) { + const uint32_t tensor_index = source.tensor_offset + static_cast(j); + Tensor &tensor = payload.tensors[j]; + GraphTensor rebound = definition_tensors[tensor_index]; + if (!graph_tensor_wire_valid(rebound)) { + execution.materialize_busy.store(0, std::memory_order_release); + return GraphMaterializeResult::INVALID; + } + const GraphTensorSourceRef &ref = tensor_sources[tensor_index]; + if (ref.source == static_cast(GraphTensorSource::BOUNDARY_EXACT)) { + if (ref.source_index >= execution.boundary_tensor_count || ref.packed_offset != 0) { + execution.materialize_busy.store(0, std::memory_order_release); + return GraphMaterializeResult::INVALID; + } + rebound = execution.boundary_tensors[ref.source_index]; + } else if (ref.source == static_cast(GraphTensorSource::BOUNDARY_VIEW)) { + if (ref.source_index >= execution.boundary_tensor_count) { + execution.materialize_busy.store(0, std::memory_order_release); + return GraphMaterializeResult::INVALID; + } + const GraphTensor &boundary = execution.boundary_tensors[ref.source_index]; + if (ref.packed_offset > UINT64_MAX - boundary.start_offset) { + execution.materialize_busy.store(0, std::memory_order_release); + return GraphMaterializeResult::INVALID; + } + rebound.buffer_addr = boundary.buffer_addr; + rebound.buffer_size = boundary.buffer_size; + rebound.owner_task_id = boundary.owner_task_id; + rebound.start_offset = boundary.start_offset + ref.packed_offset; + rebound.version = boundary.version; + rebound.child_memory = boundary.child_memory; + } else if (ref.source == static_cast(GraphTensorSource::INTERNAL) || + ref.source == static_cast(GraphTensorSource::OWN_OUTPUT)) { + const bool own_output = ref.source == static_cast(GraphTensorSource::OWN_OUTPUT); + const int32_t producer_index = own_output ? i : static_cast(ref.source_index); + if (producer_index < 0 || producer_index > i || (own_output && ref.source_index != i) || + (!own_output && producer_index == i)) { + execution.materialize_busy.store(0, std::memory_order_release); + return GraphMaterializeResult::INVALID; + } + PTO2TaskDescriptor &producer = execution.node_storage[producer_index].task; + const uint64_t producer_bytes = static_cast(nodes[producer_index].total_output_size); + const uintptr_t producer_base = reinterpret_cast(producer.packed_buffer_base); + if (ref.packed_offset > producer_bytes || rebound.buffer_size > producer_bytes - ref.packed_offset || + ref.packed_offset > UINTPTR_MAX - producer_base) { + execution.materialize_busy.store(0, std::memory_order_release); + return GraphMaterializeResult::INVALID; + } + rebound.buffer_addr = producer_base + ref.packed_offset; + rebound.owner_task_id = producer.task_id.raw; + } else { + execution.materialize_busy.store(0, std::memory_order_release); + return GraphMaterializeResult::INVALID; + } + if (!graph_tensor_wire_valid(rebound)) { + execution.materialize_busy.store(0, std::memory_order_release); + return GraphMaterializeResult::INVALID; + } + graph_tensor_unpack(rebound, &tensor); + } + if (source.scalar_count > 0) { + std::memcpy( + payload.scalars, definition_scalars + source.scalar_offset, + static_cast(source.scalar_count) * sizeof(uint64_t) + ); + } + reset_graph_payload(payload); + if (!register_initial_graph_waiter(execution, i)) { + execution.materialize_busy.store(0, std::memory_order_release); + return GraphMaterializeResult::INVALID; + } + } + execution.materialized_nodes = last; + if (nodes_materialized != nullptr) *nodes_materialized = last - first; + + if (last < execution.node_count) { + execution.materialize_busy.store(0, std::memory_order_release); + return GraphMaterializeResult::PENDING; + } + + // nodes is published before PREPARED. An activator that acquires the state + // may therefore route saved roots without observing partially built nodes. + execution.nodes = execution.node_storage; + execution.materialized_graph_key = execution.graph_key; + execution.materialized_definition_hash = execution.definition_hash; + execution.materialized_node_count = execution.node_count; + execution.state.store(GraphExecutionState::PREPARED, std::memory_order_release); + execution.materialize_busy.store(0, std::memory_order_release); + return GraphMaterializeResult::PREPARED; +} diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h index 3f20a188fd..5a4b58be50 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h @@ -21,9 +21,10 @@ * advancing the per-ring completed_watermark (consumer-retirement signal) * 4. Two-stage mixed-task completion (subtask done bits -> mixed-task complete) * - * The Scheduler runs on Device AI_CPU. host_build_graph is scheduler-only (the - * orchestrator runs to completion on the host); there is no on-device slot - * reclaim (whole-graph-resident), so last_task_alive is not advanced here. + * The Scheduler runs on Device AI_CPU. host_build_graph is scheduler-only; a + * host thread publishes committed orchestration prefixes concurrently. There + * is no on-device slot reclaim (whole-run-resident), so last_task_alive is not + * advanced here. * * Based on: docs/RUNTIME_LOGIC.md */ @@ -37,6 +38,7 @@ #include "utils/device_arena.h" #include "aicpu/platform_regs.h" // get_reg_ptr / RegId for the early-dispatch doorbell #include "pto_async_wait.h" +#include "graph_execution.h" #include "pto_ring_buffer.h" #include "pto_runtime2_types.h" #include "pto_shared_memory.h" @@ -418,6 +420,8 @@ struct PTO2SchedulerLayout { size_t off_ready_queue_slots[PTO2_NUM_RESOURCE_SHAPES]; size_t off_ready_sync_queue_slots[PTO2_NUM_RESOURCE_SHAPES]; size_t off_dummy_ready_queue_slots; + size_t off_graph_ready_queue_slots; + size_t off_graph_prepare_queue_slots; size_t off_early_dispatch_queue_slots[PTO2_NUM_RESOURCE_SHAPES]; size_t off_early_sync_start_queue_slots; uint64_t ready_queue_capacity; @@ -463,6 +467,12 @@ struct PTO2SchedulerState { // the dispatch loop and completed inline -- never goes to AICore. PTO2ReadyQueue dummy_ready_queue; + // An outer Graph is control work, never an AICore task. External dependency + // readiness and bounded materialization progress independently and meet at + // the submission's single atomic activation gate. + PTO2ReadyQueue graph_ready_queue; + PTO2ReadyQueue graph_prepare_queue; + alignas(64) AsyncWaitList async_wait_list; // Statistics (cold path, isolated from hot-path fields) @@ -480,6 +490,10 @@ struct PTO2SchedulerState { // the per-shape ready_sync_queues[] (drained as Tier-0); everything else to // ready_queues[]. void push_ready_routed(PTO2TaskSlotState *slot_state) { + if (slot_state->task_kind == PTO2TaskKind::GRAPH) { + graph_ready_queue.push(slot_state); + return; + } PTO2ResourceShape shape = slot_state->active_mask.to_shape(); if (shape == PTO2ResourceShape::DUMMY || (slot_state->task_attrs.has_predicate() && !slot_state->payload->predicate.pass())) { @@ -822,6 +836,162 @@ struct PTO2SchedulerState { // scope reference. Kept as a no-op so the orchestrator call site is unchanged. void on_scope_end(PTO2TaskSlotState ** /*task_slot_states*/, int32_t /*count*/) {} + // Orch owns dependency discovery and saves the immutable fanin wire. + // Scheduler polling only chooses which already-wired producer a consumer + // waits on at this instant; it never recomputes producer relationships. + int32_t graph_first_unmet_producer(const GraphExecution &execution, const PTO2TaskSlotState &consumer) const { + const uint32_t node_index = static_cast(consumer.graph_node_index); + const uint32_t begin = execution.fanin_offsets[node_index]; + const uint32_t end = execution.fanin_offsets[node_index + 1]; + for (uint32_t edge = begin; edge < end; ++edge) { + const uint16_t producer_index = execution.fanin_indices[edge]; + const PTO2TaskSlotState &producer = execution.nodes[producer_index].slot; + if (producer.task_state.load(std::memory_order_acquire) != PTO2_TASK_COMPLETED) { + return static_cast(producer_index); + } + } + return -1; + } + + void register_graph_wake(GraphExecution &execution, PTO2TaskSlotState *producer, PTO2TaskSlotState *consumer) { + while (true) { + PTO2TaskSlotState *expected = producer->wake_list_head.load(std::memory_order_relaxed); + while (expected != WAKE_LIST_SENTINEL) { + consumer->next_in_wake_list = expected; + if (producer->wake_list_head.compare_exchange_weak( + expected, consumer, std::memory_order_acq_rel, std::memory_order_relaxed + )) { + return; + } + } + + // The producer completed between fanin classification and the CAS. + // Its release task_state store is visible after observing the + // sentinel; rescan the saved wire and either route or retarget. + const int32_t unmet_producer = graph_first_unmet_producer(execution, *consumer); + if (unmet_producer < 0) { + push_ready_routed(consumer); + return; + } + producer = &execution.nodes[unmet_producer].slot; + } + } + + uint32_t drain_graph_wake_list(GraphExecution &execution, PTO2TaskSlotState &producer) { + uint32_t consumers_rescanned = 0; + PTO2TaskSlotState *waiter = producer.wake_list_head.exchange(WAKE_LIST_SENTINEL, std::memory_order_acq_rel); + while (waiter != nullptr && waiter != WAKE_LIST_SENTINEL) { + PTO2TaskSlotState *next = waiter->next_in_wake_list; + const int32_t unmet_producer = graph_first_unmet_producer(execution, *waiter); + if (unmet_producer < 0) { + push_ready_routed(waiter); + } else { + register_graph_wake(execution, &execution.nodes[unmet_producer].slot, waiter); + } + consumers_rescanned++; + waiter = next; + } + return consumers_rescanned; + } + + int32_t activate_prepared_graph(GraphExecution &execution) { + GraphExecutionState expected = GraphExecutionState::PREPARED; + if (!execution.state.compare_exchange_strong( + expected, GraphExecutionState::ACTIVE, std::memory_order_acq_rel, std::memory_order_acquire + )) { + return 0; + } + const GraphDefinition &definition = *execution.definition; + const uint16_t *roots = + graph_definition_array(definition, definition.off_root_indices, definition.root_count); + if (roots == nullptr) return 0; + int32_t routed = 0; + for (uint32_t i = 0; i < definition.root_count; ++i) { + const uint16_t node_index = roots[i]; + if (node_index >= static_cast(execution.node_count)) continue; + // Roots have zero internal fanin and are routed only here. Every + // non-root is registered on one saved producer at a time. + push_ready_routed(&execution.nodes[node_index].slot); + routed++; + } + return routed; + } + + GraphMaterializeResult prepare_graph_task( + PTO2TaskSlotState &outer_slot, int32_t max_nodes = GRAPH_MATERIALIZE_SLICE_NODES, + int32_t *nodes_materialized = nullptr + ) { + GraphSubmission *submission = graph_submission_from_slot(outer_slot); + if (submission == nullptr) return GraphMaterializeResult::INVALID; + GraphExecution *execution = graph_execution_localize(outer_slot); + if (execution == nullptr) return GraphMaterializeResult::INVALID; + const GraphMaterializeResult result = + graph_execution_materialize_slice(outer_slot, *execution, max_nodes, nodes_materialized); + if (result == GraphMaterializeResult::PREPARED && graph_submission_signal(*submission, 0x1)) { + activate_prepared_graph(*execution); + } + return result; + } + + int32_t activate_graph_task(PTO2TaskSlotState &outer_slot) { + GraphSubmission *submission = graph_submission_from_slot(outer_slot); + if (submission == nullptr || !graph_submission_signal(*submission, 0x2)) return 0; + GraphExecution *execution = graph_submission_local_execution(*submission); + return execution == nullptr ? 0 : activate_prepared_graph(*execution); + } + + struct TaskCompletionOutcome { + uint32_t fanout_edges{0}; + int32_t stream_tasks_completed{0}; + }; + + TaskCompletionOutcome complete_task( + PTO2TaskSlotState &slot_state +#if SIMPLER_SCHED_PROFILING + , + int thread_idx +#endif + ) { + TaskCompletionOutcome outcome; + if (slot_state.task_kind != PTO2TaskKind::GRAPH_NODE) { +#if SIMPLER_SCHED_PROFILING + CompletionStats stats = on_task_complete(slot_state, thread_idx); + outcome.fanout_edges = static_cast(stats.fanout_edges); +#else + outcome.fanout_edges = on_task_complete(slot_state); +#endif + outcome.stream_tasks_completed = 1; + return outcome; + } + + GraphExecution *execution = graph_execution_from_slot(slot_state); + if (execution == nullptr || execution->definition == nullptr || execution->nodes == nullptr) return outcome; + const int32_t saved_node_index = slot_state.graph_node_index; + if (saved_node_index < 0) return outcome; + const uint32_t node_index = static_cast(saved_node_index); + if (node_index >= static_cast(execution->node_count)) return outcome; + + // Publish completion before closing the wake list. A consumer that + // loses registration to the sentinel acquires this state when it + // rescans the Orch-built fanin wire, so no wakeup can be lost. + slot_state.mark_completed(); + outcome.fanout_edges = drain_graph_wake_list(*execution, slot_state); + + const bool graph_completed = graph_execution_complete_node(*execution); + graph_execution_retire_node(*execution); + if (!graph_completed) return outcome; + + // Internal nodes count as zero stream tasks. The final node publishes + // the outer ring task exactly once, waking external consumers and + // contributing the one task the host actually submitted. + if (execution->outer_slot != nullptr) { + on_mixed_task_complete(*execution->outer_slot); + outcome.stream_tasks_completed = 1; + } + graph_execution_mark_completed(*execution); + return outcome; + } + /** * Subtask completion: atomic counter model. * Called when a single subtask (AIC, AIV0, or AIV1) finishes on any block. @@ -923,11 +1093,11 @@ AsyncWaitList::try_inline_complete_locked(AsyncWaitList::DrainCompletionSink &si // Return value (CompletionStats / consumer-walk count) discarded: // async-wait drain path has no Resolve swimlane bar attached. #if SIMPLER_SCHED_PROFILING - (void)sink.sched->on_task_complete(slot_state, sink.thread_idx); + PTO2SchedulerState::TaskCompletionOutcome outcome = sink.sched->complete_task(slot_state, sink.thread_idx); #else - (void)sink.sched->on_task_complete(slot_state); + PTO2SchedulerState::TaskCompletionOutcome outcome = sink.sched->complete_task(slot_state); #endif - sink.inline_completed++; + sink.inline_completed += outcome.stream_tasks_completed; return true; } @@ -988,12 +1158,12 @@ inline AsyncPollResult AsyncWaitList::poll_and_complete( // Return value (CompletionStats / consumer-walk count) discarded: // deferred-completion drain has no Resolve swimlane bar attached. #if SIMPLER_SCHED_PROFILING - (void)sched->on_task_complete(*entry.slot_state, thread_idx); + PTO2SchedulerState::TaskCompletionOutcome outcome = sched->complete_task(*entry.slot_state, thread_idx); #else - (void)sched->on_task_complete(*entry.slot_state); + PTO2SchedulerState::TaskCompletionOutcome outcome = sched->complete_task(*entry.slot_state); #endif // Polling: completion is fully published inline; no deferred release. - result.completed++; + result.completed += outcome.stream_tasks_completed; int32_t last = count - 1; if (i != last) entries[i] = entries[last]; diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp index a28b26141c..e0782e10a5 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp @@ -56,7 +56,8 @@ LoopAction SchedulerContext::handle_orchestrator_exit( LOG_ERROR( "Thread %d: Fatal error (code=%d), sending EXIT_SIGNAL to all cores. " "completed_tasks=%d, total_tasks=%d", - thread_idx, orch_err, completed_tasks_.load(std::memory_order_relaxed), total_tasks_ + thread_idx, orch_err, completed_tasks_.load(std::memory_order_relaxed), + total_tasks_.load(std::memory_order_relaxed) ); if (!completed_.exchange(true, std::memory_order_acq_rel)) { emergency_shutdown(runtime); @@ -72,12 +73,15 @@ LoopAction SchedulerContext::handle_orchestrator_exit( return LoopAction::BREAK_LOOP; } - task_count = total_tasks_; - if (task_count > 0 && completed_tasks_.load(std::memory_order_relaxed) >= task_count) { + task_count = total_tasks_.load(std::memory_order_acquire); + const int32_t published = header->ring.fc.current_task_index.load(std::memory_order_acquire); + const bool host_done = header->orchestrator_done.load(std::memory_order_acquire) != 0; + const bool all_classified = classified_tasks_.load(std::memory_order_acquire) >= published; + if (host_done && all_classified && completed_tasks_.load(std::memory_order_relaxed) >= published) { completed_.store(true, std::memory_order_release); LOG_INFO( "Thread %d: PTO2 completed tasks %d/%d", thread_idx, completed_tasks_.load(std::memory_order_relaxed), - task_count + published ); return LoopAction::BREAK_LOOP; } @@ -357,7 +361,9 @@ void SchedulerContext::log_shutdown_stall_snapshot( thread_count = thread_count < 0 ? 0 : MAX_AICPU_THREADS; } for (int32_t t = 0; t < thread_count; t++) { - log_stall_diagnostics(t, total_tasks_, trigger_idle_iterations, trigger_last_progress_count); + log_stall_diagnostics( + t, total_tasks_.load(std::memory_order_relaxed), trigger_idle_iterations, trigger_last_progress_count + ); } } @@ -911,9 +917,9 @@ int32_t SchedulerContext::post_handshake_init(Runtime *runtime) { int32_t ring_tasks = header->ring.fc.current_task_index.load(std::memory_order_acquire); if (ring_tasks > 0 && ring_tasks <= PTO2_SCOPE_TASKS_CAP) pto2_count += ring_tasks; } - total_tasks_ = static_cast(pto2_count); + total_tasks_.store(static_cast(pto2_count), std::memory_order_relaxed); } else { - total_tasks_ = 0; + total_tasks_.store(0, std::memory_order_relaxed); } completed_tasks_.store(0, std::memory_order_release); @@ -1009,7 +1015,8 @@ void SchedulerContext::deinit() { // Reset task counters and orchestrator state completed_tasks_.store(0, std::memory_order_release); - total_tasks_ = 0; + total_tasks_.store(0, std::memory_order_relaxed); + classified_tasks_.store(0, std::memory_order_relaxed); completed_.store(false, std::memory_order_release); // Reset core discovery and assignment state @@ -1034,15 +1041,11 @@ void SchedulerContext::bind_runtime(PTO2Runtime *rt) { } // ============================================================================= -// Post-orchestration bookkeeping. Runs once on the boot leader after the -// host-built image is attached; latches total_tasks_ and folds inline-completed -// tasks (or shuts down on a fatal orchestration error). The caller publishes -// runtime_init_ready_ (release) after this returns — that store is what makes -// total_tasks_ visible to the scheduler threads, which acquire it before -// dispatching. +// Streaming host-orchestration bootstrap. The SM and scheduler arena are ready, +// but task prefixes and EOS can still arrive while dispatch is running. // ============================================================================= -void SchedulerContext::on_orchestration_done( - Runtime *runtime, PTO2Runtime *rt, [[maybe_unused]] int32_t thread_idx, int32_t total_tasks +void SchedulerContext::on_host_orchestration_stream_start( + Runtime *runtime, PTO2Runtime *rt, [[maybe_unused]] int32_t thread_idx ) { #if SIMPLER_DFX if (l2_swimlane_level_ >= L2SwimlaneLevel::ORCH_PHASES) { @@ -1053,21 +1056,17 @@ void SchedulerContext::on_orchestration_done( } #endif - total_tasks_ = total_tasks; + completed_tasks_.store(0, std::memory_order_release); + total_tasks_.store(0, std::memory_order_release); + classified_tasks_.store(0, std::memory_order_release); // Allocate the per-S CompletedTaskQueues here on the boot leader, before it // releases runtime_init_ready_ — no scheduler thread can push until then. - // Completed-but-unresolved tasks in flight are bounded by BOTH the total task - // count and the ring's task window (a task must occupy a ring slot to run and - // complete), so size to the tighter of the two, rounded up to a power of two - // and floored at 256. The window already caps this, so there is no artificial - // ceiling and a producer never has to spin on a full queue. - uint64_t sp_bound = static_cast(total_tasks); + // The final published count is not known yet, so size each queue to the + // task window: that is the hard bound on completed-but-unresolved tasks. + uint64_t sp_bound = 0; if (sched_->ring_sched_state.ring != nullptr) { - uint64_t window = static_cast(sched_->ring_sched_state.ring->task_window_mask) + 1; - if (window < sp_bound) { - sp_bound = window; - } + sp_bound = static_cast(sched_->ring_sched_state.ring->task_window_mask) + 1; } uint64_t sp_cap = 256; while (sp_cap < sp_bound) { @@ -1078,16 +1077,8 @@ void SchedulerContext::on_orchestration_done( sp_queues_[t].init(sp_cap); } - // Fold tasks completed inline during orchestration - int32_t inline_completed = static_cast(rt->orchestrator.inline_completed_tasks); - if (inline_completed > 0) { - completed_tasks_.fetch_add(inline_completed, std::memory_order_relaxed); -#if SIMPLER_SCHED_PROFILING - rt->scheduler.tasks_completed.fetch_add(inline_completed, std::memory_order_relaxed); -#endif - } - - // Check for fatal error from orchestration; if so, shut down immediately. + // Detect a boot-time fatal marker. Normal host orchestration errors arrive + // later and are observed by the dispatch/resolution loops. int32_t orch_err = 0; if (sched_->sm_header) { orch_err = sched_->sm_header->orch_error_code.load(std::memory_order_relaxed); @@ -1097,12 +1088,7 @@ void SchedulerContext::on_orchestration_done( emergency_shutdown(runtime); } } - - // The polling initial classify (seed the ready queues + wake lists for the - // whole graph) runs AFTER this, partitioned across all AICPU threads in - // classify_partition() — see AicpuExecutor::run. It is kept out of this - // leader-only setup so the O(total_tasks) scan is not serial on one thread - // while the others idle-wait for runtime_init_ready_. + (void)rt; #if SIMPLER_DFX // Write the core-to-thread mapping so the profiling data reflects the @@ -1118,42 +1104,38 @@ void SchedulerContext::on_orchestration_done( #endif } -// Polling initial classify (device boot), partitioned across all AICPU threads. -// The host built the whole graph and no producer has executed yet — every -// completion_flags byte is 0 except the hidden-alloc tasks the host completed -// inline (pre-set to 1). Each thread classifies its contiguous slice of the -// submitted-task range exactly once: route roots (all fanin met) to the ready -// queues and register the rest on their first unmet producer's wake list. -// -// This is the same work the wiring model deferred to a device queue, now run -// N-way parallel. push_ready_routed (MPMC ready queues) and register_wake -// (lock-free wake-list CAS) are the same concurrency-safe primitives the -// scheduler threads use during the run, and at boot no producer has completed -// (wake_list heads are nullptr, never SENTINEL), so registration never -// re-classifies. The caller barriers all threads here BEFORE any of them -// publishes runtime_init_ready_, so the whole ready-set / wake-list graph is -// fully seeded before the first dispatch. -void SchedulerContext::classify_partition(int32_t thread_idx, int32_t nthreads) { - if (completed_.load(std::memory_order_acquire) || sched_->ring_sched_state.ring == nullptr) { - return; - } +bool SchedulerContext::classify_published_tasks(int32_t thread_idx) { + if (thread_idx != 0 || sched_ == nullptr || sched_->ring_sched_state.ring == nullptr) return false; PTO2SharedMemoryRingHeader &ring = *sched_->ring_sched_state.ring; - const int32_t submitted = ring.fc.current_task_index.load(std::memory_order_acquire); - // Disjoint contiguous slices covering [0, submitted): thread t owns - // [submitted*t/nthreads, submitted*(t+1)/nthreads). int64 math avoids overflow. - const int32_t lo = static_cast((static_cast(submitted) * thread_idx) / nthreads); - const int32_t hi = static_cast((static_cast(submitted) * (thread_idx + 1)) / nthreads); - for (int32_t id = lo; id < hi; id++) { - if (ring.is_completion_flag_set(id)) { - continue; // completed on the host (hidden alloc); nothing to dispatch - } - PTO2TaskSlotState &s = ring.get_slot_state_by_task_id(id); - int32_t state = sched_->classify_fanin_state(&s); - if (state < 0) { - sched_->push_ready_routed(&s); + int32_t next = classified_tasks_.load(std::memory_order_relaxed); + const int32_t published = ring.fc.current_task_index.load(std::memory_order_acquire); + if (published < next || published > static_cast(ring.task_window_size)) { + return false; + } + const bool made_progress = next < published; + while (next < published) { + PTO2TaskSlotState &slot = ring.get_slot_state_by_task_id(next); + if (ring.completion_flags[next & ring.task_window_mask].load(std::memory_order_acquire) != 0) { + completed_tasks_.fetch_add(1, std::memory_order_relaxed); +#if SIMPLER_SCHED_PROFILING + sched_->tasks_completed.fetch_add(1, std::memory_order_relaxed); +#endif } else { - int32_t prod_local = s.payload->fanin_local_ids[state]; - sched_->register_wake(&ring.get_slot_state_by_task_id(prod_local), &s); + if (slot.task_kind == PTO2TaskKind::GRAPH) { + while (!sched_->graph_prepare_queue.push_tagged(&slot, slot.task->task_id.raw)) + SPIN_WAIT_HINT(); + } + const int32_t fanin_state = sched_->classify_fanin_state(&slot); + if (fanin_state < 0) { + sched_->push_ready_routed(&slot); + } else { + const int32_t producer = slot.payload->fanin_local_ids[fanin_state]; + sched_->register_wake(&ring.get_slot_state_by_task_id(producer), &slot); + } } + ++next; + classified_tasks_.store(next, std::memory_order_release); } + total_tasks_.store(published, std::memory_order_release); + return made_progress; } diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h index 680b3046d7..2bdae7bc2b 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h @@ -151,21 +151,10 @@ class SchedulerContext { // Orchestrator threads (core_trackers_[thread_idx].core_num() == 0) are a no-op. int32_t shutdown(int32_t thread_idx); - // Run all post-orchestration scheduler bookkeeping: - // - publishes core assignments to the perf collector (SIMPLER_DFX) - // - latches submitted task count from PTO2 shared memory - // - folds inline_completed_tasks into completed_tasks_ - // (skipped on fatal error — emergency_shutdown runs instead) - // Callers must invoke rt_orchestration_done(rt) before this — that - // step belongs to the orchestrator lifecycle, not the scheduler. - void on_orchestration_done(Runtime *runtime, PTO2Runtime *rt, int32_t thread_idx, int32_t total_tasks); - - // Seed the ready queues + wake lists for the whole graph at boot. Called by - // every AICPU thread on a disjoint slice of the submitted-task range, after - // on_orchestration_done and before runtime_init_ready_ (the caller barriers - // all threads between the two). Concurrency-safe: push_ready_routed and - // register_wake are the same lock-free primitives used during the run. - void classify_partition(int32_t thread_idx, int32_t nthreads); + // Initialize scheduler-side counters for an incrementally published host + // orchestration stream. Task classification happens in the dispatch loop as + // current_task_index advances; orchestrator_done is the EOS marker. + void on_host_orchestration_stream_start(Runtime *runtime, PTO2Runtime *rt, int32_t thread_idx); // Bind the PTO2Runtime scheduler pointer. void bind_runtime(PTO2Runtime *rt); @@ -217,7 +206,8 @@ class SchedulerContext { // --- Task-execution tracking --- std::atomic completed_tasks_{0}; - int32_t total_tasks_{0}; + std::atomic total_tasks_{0}; + std::atomic classified_tasks_{0}; std::atomic completed_{false}; uint64_t *func_id_to_addr_{nullptr}; @@ -520,6 +510,8 @@ class SchedulerContext { __attribute__((noinline, cold)) LoopAction handle_orchestrator_exit(int32_t thread_idx, PTO2SharedMemoryHeader *header, Runtime *runtime, int32_t &task_count); + bool classify_published_tasks(int32_t thread_idx); + __attribute__((noinline, cold)) LoopAction check_idle_fatal_error(int32_t thread_idx, PTO2SharedMemoryHeader *header, Runtime *runtime); diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp index 4a8c467fab..9ab6ceb94c 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp @@ -908,26 +908,22 @@ int32_t SchedulerContext::run_resolution_thread(Runtime *runtime, int32_t thread while (true) { if (completed_.load(std::memory_order_acquire)) break; - // Propagate a fatal error latched by the orchestrator (host) or a - // scheduler thread; mirror resolve_and_dispatch's exit behavior. - if (header->orch_error_code.load(std::memory_order_acquire) != PTO2_ERROR_NONE || - header->sched_error_code.load(std::memory_order_acquire) != PTO2_ERROR_NONE) { - if (!completed_.exchange(true, std::memory_order_acq_rel)) { - emergency_shutdown(runtime); - } + int32_t published_task_count = 0; + if (handle_orchestrator_exit(thread_idx, header, runtime, published_task_count) == LoopAction::BREAK_LOOP) break; - } int32_t resolved_this_pass = 0; + bool resolved_any = false; for (int32_t s = 0; s < active_sched_threads_; s++) { PTO2TaskSlotState *slot; while ((slot = sp_queues_[s].pop()) != nullptr) { #if SIMPLER_SCHED_PROFILING - (void)sched_->on_task_complete(*slot, thread_idx); + PTO2SchedulerState::TaskCompletionOutcome outcome = sched_->complete_task(*slot, thread_idx); #else - (void)sched_->on_task_complete(*slot); + PTO2SchedulerState::TaskCompletionOutcome outcome = sched_->complete_task(*slot); #endif - resolved_this_pass++; + resolved_this_pass += outcome.stream_tasks_completed; + resolved_any = true; } } @@ -955,6 +951,7 @@ int32_t SchedulerContext::run_resolution_thread(Runtime *runtime, int32_t thread break; } resolved_this_pass += poll_result.completed; + resolved_any = resolved_any || poll_result.completed > 0; } // Dependency-only tasks (empty active_mask, or a predicate that failed) @@ -968,29 +965,27 @@ int32_t SchedulerContext::run_resolution_thread(Runtime *runtime, int32_t thread while ((dummy_got = sched_->dummy_ready_queue.pop_batch(dummy_batch, DUMMY_DRAIN_BATCH)) > 0) { for (int di = 0; di < dummy_got; di++) { #if SIMPLER_SCHED_PROFILING - (void)sched_->on_task_complete(*dummy_batch[di], thread_idx); + PTO2SchedulerState::TaskCompletionOutcome outcome = + sched_->complete_task(*dummy_batch[di], thread_idx); #else - (void)sched_->on_task_complete(*dummy_batch[di]); + PTO2SchedulerState::TaskCompletionOutcome outcome = sched_->complete_task(*dummy_batch[di]); #endif - resolved_this_pass++; + resolved_this_pass += outcome.stream_tasks_completed; + resolved_any = true; } } } - if (resolved_this_pass > 0) { - int32_t new_total = - completed_tasks_.fetch_add(resolved_this_pass, std::memory_order_relaxed) + resolved_this_pass; + if (resolved_any) { + if (resolved_this_pass > 0) { + completed_tasks_.fetch_add(resolved_this_pass, std::memory_order_relaxed); #if SIMPLER_SCHED_PROFILING - // P owns the completion accounting, so it owns the profiling mirror too - // (the S threads' completed_this_turn no longer feeds it in P mode). - sched_->tasks_completed.fetch_add(resolved_this_pass, std::memory_order_relaxed); + // P owns the completion accounting, so it owns the profiling mirror too + // (the S threads' completed_this_turn no longer feeds it in P mode). + sched_->tasks_completed.fetch_add(resolved_this_pass, std::memory_order_relaxed); #endif - last_progress_ts = get_sys_cnt_aicpu(); - if (total_tasks_ > 0 && new_total >= total_tasks_) { - completed_.store(true, std::memory_order_release); - LOG_INFO("Thread %d: P resolved all tasks %d/%d", thread_idx, new_total, total_tasks_); - break; } + last_progress_ts = get_sys_cnt_aicpu(); continue; // fast re-drain while work keeps arriving } @@ -1003,11 +998,12 @@ int32_t SchedulerContext::run_resolution_thread(Runtime *runtime, int32_t thread // task — a genuine forward-progress stall / pre-dispatch deadlock. uint64_t now = get_sys_cnt_aicpu(); if (now - last_progress_ts > scheduler_timeout_cycles) { - bool outstanding = total_tasks_ > 0 && completed_tasks_.load(std::memory_order_relaxed) < total_tasks_; + const int32_t total = total_tasks_.load(std::memory_order_acquire); + bool outstanding = total > 0 && completed_tasks_.load(std::memory_order_relaxed) < total; if (outstanding && no_thread_owns_running_task()) { LOG_ERROR( "Thread %d: P resolution stall (%d/%d resolved)", thread_idx, - completed_tasks_.load(std::memory_order_relaxed), total_tasks_ + completed_tasks_.load(std::memory_order_relaxed), total ); int32_t expected = PTO2_ERROR_NONE; header->sched_error_code.compare_exchange_strong( @@ -1188,6 +1184,7 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ // (or skip entirely on iters with no phase emit). iter_shared_sampled = false; #endif + if (classify_published_tasks(thread_idx)) made_progress = true; int32_t task_count = 0; if (!tracker.has_any_running_cores()) { LoopAction action = handle_orchestrator_exit(thread_idx, header, runtime, task_count); @@ -1288,9 +1285,187 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ continue; } - // Phase 3 (dependency-only dummy / predicate-failed retirement) runs on - // the resolution thread P, not here — see run_resolution_thread. The - // scheduler loop goes straight from completion detection to dispatch. + // Graph control work never consumes an AICore. External dependency + // readiness and bounded definition materialization progress + // independently, then meet at GraphSubmission::activation_gate. + // + // Keep this ahead of dummy/regular dispatch so a ready Graph can expose + // its root nodes without waiting for an otherwise unrelated dispatch + // pass. Limiting the work to one activation and one bounded prepare + // slice per loop prevents a large definition from monopolizing a + // scheduler thread. + if (thread_idx < active_sched_threads_) { + PTO2TaskSlotState *graph_slot = sched_->graph_ready_queue.pop(); + if (graph_slot != nullptr) { + if (graph_slot->task != nullptr && graph_slot->task_kind == PTO2TaskKind::GRAPH) { + (void)sched_->activate_graph_task(*graph_slot); + made_progress = true; + } else { + int32_t expected = PTO2_ERROR_NONE; + if (header->sched_error_code.compare_exchange_strong( + expected, PTO2_ERROR_INVALID_ARGS, std::memory_order_acq_rel, std::memory_order_acquire + )) { + header->sched_error_thread.store(thread_idx, std::memory_order_release); + } + header->sched_error_bitmap.fetch_or( + 1U << static_cast(thread_idx), std::memory_order_acq_rel + ); + completed_.store(true, std::memory_order_release); + break; + } + } + + uint64_t prepare_task_id = 0; + PTO2TaskSlotState *prepare_slot = sched_->graph_prepare_queue.pop_tagged(&prepare_task_id); + if (prepare_slot != nullptr) { + const bool valid_slot = prepare_slot->task != nullptr && + prepare_slot->task_kind == PTO2TaskKind::GRAPH && + prepare_slot->task->task_id.raw == prepare_task_id; + if (!valid_slot) { + int32_t expected = PTO2_ERROR_NONE; + if (header->sched_error_code.compare_exchange_strong( + expected, PTO2_ERROR_INVALID_ARGS, std::memory_order_acq_rel, std::memory_order_acquire + )) { + header->sched_error_thread.store(thread_idx, std::memory_order_release); + } + header->sched_error_bitmap.fetch_or( + 1U << static_cast(thread_idx), std::memory_order_acq_rel + ); + completed_.store(true, std::memory_order_release); + break; + } +#if SIMPLER_DFX + uint64_t graph_prepare_t0 = + l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES ? get_sys_cnt_aicpu() : 0; +#endif + int32_t nodes_materialized = 0; + GraphMaterializeResult result = + sched_->prepare_graph_task(*prepare_slot, GRAPH_MATERIALIZE_SLICE_NODES, &nodes_materialized); + if (result == GraphMaterializeResult::PENDING || result == GraphMaterializeResult::BUSY) { + while (!sched_->graph_prepare_queue.push_tagged(prepare_slot, prepare_task_id)) { + SPIN_WAIT_HINT(); + } + } else if (result == GraphMaterializeResult::INVALID) { + int32_t expected = PTO2_ERROR_NONE; + if (header->sched_error_code.compare_exchange_strong( + expected, PTO2_ERROR_INVALID_ARGS, std::memory_order_acq_rel, std::memory_order_acquire + )) { + header->sched_error_thread.store(thread_idx, std::memory_order_release); + } + header->sched_error_bitmap.fetch_or( + 1U << static_cast(thread_idx), std::memory_order_acq_rel + ); + completed_.store(true, std::memory_order_release); + break; + } + if (nodes_materialized > 0 || result == GraphMaterializeResult::PREPARED) { + made_progress = true; + } +#if SIMPLER_DFX + if (graph_prepare_t0 != 0) { + uint64_t graph_prepare_t1 = get_sys_cnt_aicpu(); + l2_swimlane_aicpu_record_graph_prepare( + thread_idx, graph_prepare_t0, graph_prepare_t1, l2_swimlane.sched_loop_count, prepare_task_id, + static_cast(nodes_materialized) + ); + _t0_phase = graph_prepare_t1; + } +#endif + } + } + + // Phase 3: Drain dummy ready queue (S0/S1/S2). + // + // Dependency-only tasks bypass AICore dispatch: they go through the + // scheduler so fanin/fanout edges stay consistent, but completion is + // signalled inline here. The ready queue is MPMC, and the fanout path + // uses per-slot locks/atomics, so multiple scheduler threads can share + // the dependency-only resolve work. + if (thread_idx < active_sched_threads_) { + constexpr int DUMMY_DRAIN_BATCH = 8; + PTO2TaskSlotState *dummy_batch[DUMMY_DRAIN_BATCH]; + int dummy_got = sched_->dummy_ready_queue.pop_batch(dummy_batch, DUMMY_DRAIN_BATCH); +#if SIMPLER_DFX + // Dummy outer phase: covers handling of all dummies popped this + // iter. Per-dummy DummyTask markers are emitted to a SEPARATE lane + // (Worker View AICPU_N) by the converter, so they do not nest + // under this bar. Resolve emits below DO land on the sched lane + // and nest under this Dummy outer by time containment. + uint64_t dummy_outer_t0 = + (dummy_got > 0 && l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) ? get_sys_cnt_aicpu() : 0; +#endif + for (int di = 0; di < dummy_got; di++) { + PTO2TaskSlotState &dummy_slot = *dummy_batch[di]; + + // ----- Resolve work: walk this dummy's consumer list. ------ + // Same 1 µs filter as the main-path Resolve emit suppresses + // dummies whose consumer release runs sub-microsecond. +#if SIMPLER_DFX + uint64_t dummy_resolve_t0 = + (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) ? get_sys_cnt_aicpu() : 0; +#endif + // [[maybe_unused]] silences -Werror=unused-but-set-variable on + // the profiling-flags-smoke build path where SIMPLER_DFX is + // OFF and the Resolve emit below is excluded. + [[maybe_unused]] uint32_t dummy_consumers = 0; +#if SIMPLER_SCHED_PROFILING + PTO2SchedulerState::TaskCompletionOutcome outcome = sched_->complete_task(dummy_slot, thread_idx); +#else + PTO2SchedulerState::TaskCompletionOutcome outcome = sched_->complete_task(dummy_slot); +#endif + dummy_consumers = outcome.fanout_edges; +#if SIMPLER_DFX + if (dummy_resolve_t0 != 0) { + uint64_t dummy_resolve_t1 = get_sys_cnt_aicpu(); + constexpr uint64_t RESOLVE_EMIT_MIN_CYCLES = PLATFORM_PROF_SYS_CNT_FREQ / 1'000'000; // 1 µs + if (dummy_resolve_t1 - dummy_resolve_t0 >= RESOLVE_EMIT_MIN_CYCLES) { + l2_swimlane_aicpu_record_sched_phase( + thread_idx, L2SwimlaneSchedPhaseKind::Resolve, dummy_resolve_t0, dummy_resolve_t1, + sched_l2_swimlane_[thread_idx].sched_loop_count, dummy_consumers + ); + } + l2_swimlane_aicpu_record_dummy_task( + thread_idx, dummy_resolve_t0, sched_l2_swimlane_[thread_idx].sched_loop_count, + dummy_slot.task->task_id.raw + ); + } +#endif + // Polling: on_task_complete already published this slot's + // completion + drained its wake list inline. There is no deferred + // producer-release phase — consumer retirement is observed via the + // per-ring completed_watermark, not by bumping producer refcounts. + int32_t prev = completed_tasks_.fetch_add(outcome.stream_tasks_completed, std::memory_order_relaxed); + last_progress_count = prev + outcome.stream_tasks_completed; + cur_thread_completed++; + } + if (dummy_got > 0) { + made_progress = true; + } +#if SIMPLER_DFX + // Emit Dummy outer over the whole dummy_drain pass. Span starts at + // dummy_outer_t0 (captured after pop_batch) and ends at "now". + // tasks_processed = dummy_got. Advancing _t0_phase here makes the + // following Dispatch / EarlyDispatch / second-Complete bars start + // at this end. + if (dummy_outer_t0 != 0) { + uint64_t dummy_outer_t1 = get_sys_cnt_aicpu(); + int16_t phase_end_shared[L2SWIMLANE_NUM_QUEUE_SHAPES]; + capture_phase_end_fresh(phase_end_shared); + l2_swimlane_aicpu_record_sched_phase( + thread_idx, L2SwimlaneSchedPhaseKind::Dummy, dummy_outer_t0, dummy_outer_t1, + l2_swimlane.sched_loop_count, static_cast(dummy_got), /*pop_hit=*/0, + /*pop_miss=*/0, phase_start_shared, phase_end_shared + ); + for (int s = 0; s < L2SWIMLANE_NUM_QUEUE_SHAPES; s++) + phase_start_shared[s] = phase_end_shared[s]; + _t0_phase = dummy_outer_t1; + // We do NOT re-sync _t0/_t1 — the dummy span will be absorbed + // into the next CYCLE_COUNT_LAP accumulator. The phase-model + // anchor (_t0_phase) is the authoritative source for bar spans + // on the swimlane; the cycle accumulators are coarse aggregates. + } +#endif + } // Phase 4: MIX-strict-priority dispatch with phase-split and // cross-thread idle gating. See dispatch_ready_tasks for the policy. @@ -1387,7 +1562,9 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ } if (idle_iterations % STALL_LOG_INTERVAL == 0) { - log_stall_diagnostics(thread_idx, total_tasks_, idle_iterations, last_progress_count); + log_stall_diagnostics( + thread_idx, total_tasks_.load(std::memory_order_relaxed), idle_iterations, last_progress_count + ); } // Wall-clock budget gate, with two fatal-latch branches: // @@ -1407,8 +1584,10 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ // observability is preserved. if (get_sys_cnt_aicpu() - last_progress_ts > scheduler_timeout_cycles) { bool self_owns = self_owns_running_task(thread_idx); - bool global_stuck = !self_owns && total_tasks_ > 0 && - completed_tasks_.load(std::memory_order_relaxed) < total_tasks_ && + const int32_t total = total_tasks_.load(std::memory_order_relaxed); + const bool host_done = header->orchestrator_done.load(std::memory_order_acquire) != 0; + bool global_stuck = host_done && !self_owns && total > 0 && + completed_tasks_.load(std::memory_order_relaxed) < total && no_thread_owns_running_task(); if (self_owns || global_stuck) { // Latch the error + emergency_shutdown, then break to the diff --git a/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp b/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp index 61e61e5152..9f5108f040 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp @@ -112,6 +112,8 @@ PTO2SchedulerLayout PTO2SchedulerState::reserve_layout(DeviceArena &arena) { layout.off_ready_sync_queue_slots[i] = ready_queue_reserve_layout(arena, PTO2_READY_QUEUE_SIZE); } layout.off_dummy_ready_queue_slots = ready_queue_reserve_layout(arena, PTO2_READY_QUEUE_SIZE); + layout.off_graph_ready_queue_slots = ready_queue_reserve_layout(arena, PTO2_READY_QUEUE_SIZE); + layout.off_graph_prepare_queue_slots = ready_queue_reserve_layout(arena, PTO2_READY_QUEUE_SIZE); for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { layout.off_early_dispatch_queue_slots[i] = ready_queue_reserve_layout(arena, PTO2_EARLY_DISPATCH_QUEUE_SIZE); } @@ -154,6 +156,14 @@ bool PTO2SchedulerState::init_data_from_layout( )) { return false; } + if (!ready_queue_init_data_from_layout( + &sched->graph_ready_queue, arena, layout.off_graph_ready_queue_slots, layout.ready_queue_capacity + ) || + !ready_queue_init_data_from_layout( + &sched->graph_prepare_queue, arena, layout.off_graph_prepare_queue_slots, layout.ready_queue_capacity + )) { + return false; + } for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { if (!ready_queue_init_data_from_layout( &sched->early_dispatch_queues[i], arena, layout.off_early_dispatch_queue_slots[i], @@ -182,6 +192,8 @@ void PTO2SchedulerState::wire_arena_pointers(const PTO2SchedulerLayout &layout, ready_queue_wire_arena_pointers(&sched->ready_sync_queues[i], arena, layout.off_ready_sync_queue_slots[i]); } ready_queue_wire_arena_pointers(&sched->dummy_ready_queue, arena, layout.off_dummy_ready_queue_slots); + ready_queue_wire_arena_pointers(&sched->graph_ready_queue, arena, layout.off_graph_ready_queue_slots); + ready_queue_wire_arena_pointers(&sched->graph_prepare_queue, arena, layout.off_graph_prepare_queue_slots); for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { ready_queue_wire_arena_pointers( &sched->early_dispatch_queues[i], arena, layout.off_early_dispatch_queue_slots[i] @@ -200,6 +212,8 @@ void PTO2SchedulerState::destroy() { ready_queue_destroy(&sched->ready_sync_queues[i]); } ready_queue_destroy(&sched->dummy_ready_queue); + ready_queue_destroy(&sched->graph_ready_queue); + ready_queue_destroy(&sched->graph_prepare_queue); for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { ready_queue_destroy(&sched->early_dispatch_queues[i]); } @@ -364,6 +378,7 @@ PTO2Runtime *runtime_init_data_from_layout( rt->gm_heap_size = total_heap_size; rt->gm_heap_owned = false; rt->total_cycles = 0; + rt->active_callable_hash = 0; if (!rt->orchestrator.init_data_from_layout( layout.orch, arena, sm_dev_base, gm_heap_dev_base, heap_sizes[0], layout.task_window_sizes[0] diff --git a/src/a2a3/runtime/host_build_graph/runtime/shared/runtime.cpp b/src/a2a3/runtime/host_build_graph/runtime/shared/runtime.cpp index 8fae1400ee..9aa891dd2a 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/shared/runtime.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/shared/runtime.cpp @@ -37,7 +37,6 @@ Runtime::Runtime() { memset(aicpu_allowed_cpus, 0, sizeof(aicpu_allowed_cpus)); aicpu_allowed_cpu_count = 0; aicpu_launch_count = 0; - host_total_tasks = 0; // Initialize shared-memory / orchestration argument plumbing gm_sm_ptr_ = nullptr; @@ -52,6 +51,7 @@ Runtime::Runtime() { dev_orch_so_size_ = 0; device_orch_func_name_[0] = '\0'; device_orch_config_name_[0] = '\0'; + deferred_host_orchestration_ = nullptr; // Initialize kernel binary tracking registered_kernel_count_ = 0; diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp index 6f5cdacb8a..b28dc7f984 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp @@ -835,7 +835,7 @@ static bool build_and_cache_prebuilt_arena( extern "C" int bind_callable_to_runtime_impl( Runtime *runtime, const HostApi *api, const ChipStorageTaskArgs *orch_args, void *host_orch_func_ptr, const ArgDirection *signature, int sig_count, const uint64_t *ring_task_window, const uint64_t *ring_heap, - const uint64_t *ring_dep_pool + const uint64_t *ring_dep_pool, int32_t /*l2_swimlane_level*/ ) { if (runtime == nullptr) { LOG_ERROR("Runtime pointer is null"); diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/runtime.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/runtime.h index 90977d3e20..bfeb133fc3 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/runtime.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/runtime.h @@ -308,6 +308,20 @@ class Runtime { // runtime_maker.cpp from orch_args at bind time, then iterated in // validate_runtime_impl. Host-only (after `dev`): never uploaded. std::vector tensor_leases_; + + // Shared platform DFX accessors. TRB orchestration runs on AICPU and has no + // host-clock envelope, so these return an empty stream. + const std::vector &get_host_orch_phase_records() const { + static const std::vector empty; + return empty; + } + uint64_t get_host_orch_start_cycles() const { return 0; } + uint64_t get_host_orch_end_cycles() const { return 0; } + uint64_t get_host_orch_first_publish_cycles() const { return 0; } + bool has_deferred_host_orchestration() const { return false; } + void notify_deferred_host_execution_started() {} + int run_deferred_host_orchestration(const HostApi *) { return 0; } + void release_deferred_host_orchestration(const HostApi *) {} }; // `dev` must be the first member so the narrowed H2D copy starts at offset 0. diff --git a/src/a5/runtime/host_build_graph/host/runtime_maker.cpp b/src/a5/runtime/host_build_graph/host/runtime_maker.cpp index ea917e6542..e296c00d3c 100644 --- a/src/a5/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a5/runtime/host_build_graph/host/runtime_maker.cpp @@ -391,7 +391,7 @@ int register_callable_impl(const ChipCallable *callable, uint64_t (*upload_fn)(c int bind_callable_to_runtime_impl( Runtime *runtime, const HostApi *api, const ChipStorageTaskArgs *orch_args, void *host_orch_func_ptr, const ArgDirection *signature, int sig_count, const uint64_t * /*ring_task_window*/, const uint64_t * /*ring_heap*/, - const uint64_t * /*ring_dep_pool*/ + const uint64_t * /*ring_dep_pool*/, int32_t /*l2_swimlane_level*/ ) { if (runtime == nullptr) { LOG_ERROR("Runtime pointer is null"); diff --git a/src/a5/runtime/host_build_graph/runtime/runtime.h b/src/a5/runtime/host_build_graph/runtime/runtime.h index 3c9655737c..f3b745e63d 100644 --- a/src/a5/runtime/host_build_graph/runtime/runtime.h +++ b/src/a5/runtime/host_build_graph/runtime/runtime.h @@ -291,6 +291,12 @@ class Runtime { int32_t *get_aicpu_allowed_cpus() { return aicpu_allowed_cpus; } size_t aicpu_allowed_cpus_capacity() const { return sizeof(aicpu_allowed_cpus) / sizeof(aicpu_allowed_cpus[0]); } + // Common platform hook; a5 host_build_graph still builds synchronously. + bool has_deferred_host_orchestration() const { return false; } + void notify_deferred_host_execution_started() {} + int run_deferred_host_orchestration(const HostApi *) { return 0; } + void release_deferred_host_orchestration(const HostApi *) {} + // ========================================================================= // Task Management // ========================================================================= diff --git a/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp b/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp index 382353333c..c9af7f0115 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp @@ -818,7 +818,7 @@ static bool build_and_cache_prebuilt_arena( extern "C" int bind_callable_to_runtime_impl( Runtime *runtime, const HostApi *api, const ChipStorageTaskArgs *orch_args, void *host_orch_func_ptr, const ArgDirection *signature, int sig_count, const uint64_t *ring_task_window, const uint64_t *ring_heap, - const uint64_t *ring_dep_pool + const uint64_t *ring_dep_pool, int32_t /*l2_swimlane_level*/ ) { if (runtime == nullptr) { LOG_ERROR("Runtime pointer is null"); diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/runtime.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/runtime.h index dfd382f9b1..c044a71eca 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/runtime.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/runtime.h @@ -264,6 +264,12 @@ class Runtime { return sizeof(dev.aicpu_allowed_cpus) / sizeof(dev.aicpu_allowed_cpus[0]); } + // Common platform hook; this runtime has device-side orchestration. + bool has_deferred_host_orchestration() const { return false; } + void notify_deferred_host_execution_started() {} + int run_deferred_host_orchestration(const HostApi *) { return 0; } + void release_deferred_host_orchestration(const HostApi *) {} + // ========================================================================= // Performance Profiling // ========================================================================= diff --git a/src/common/platform/include/aicpu/l2_swimlane_collector_aicpu.h b/src/common/platform/include/aicpu/l2_swimlane_collector_aicpu.h index 0c6dc141da..2cb2922309 100644 --- a/src/common/platform/include/aicpu/l2_swimlane_collector_aicpu.h +++ b/src/common/platform/include/aicpu/l2_swimlane_collector_aicpu.h @@ -234,6 +234,12 @@ void l2_swimlane_aicpu_record_predicated_skip( int thread_idx, uint64_t complete_time, uint32_t loop_iter, uint64_t task_id ); +/** Record one bounded Scheduler-side Graph materialization slice. */ +void l2_swimlane_aicpu_record_graph_prepare( + int thread_idx, uint64_t start_time, uint64_t end_time, uint32_t loop_iter, uint64_t task_id, + uint32_t nodes_materialized +); + /** * Set orchestrator thread index for per-task phase recording * diff --git a/src/common/platform/include/common/host_api.h b/src/common/platform/include/common/host_api.h index 9cbd54cf0a..eb7001671c 100644 --- a/src/common/platform/include/common/host_api.h +++ b/src/common/platform/include/common/host_api.h @@ -32,6 +32,11 @@ struct HostApi { void (*device_free)(void *dev_ptr); int (*copy_to_device)(void *dev_ptr, const void *host_ptr, size_t size); int (*copy_from_device)(void *host_ptr, const void *dev_ptr, size_t size); + // Publish a 32-bit host-produced control value to device memory. Unlike a + // generic byte copy, the sim implementation performs an atomic release + // store so a concurrently running simulated AICPU can acquire the value + // without a C++ data race. Onboard uses a synchronous 4-byte H2D copy. + int (*publish_i32)(void *dev_ptr, int32_t value); // Map a device buffer into host address space and return a host-readable VA // (nullptr on failure); the paired unregister releases it. The returned VA // may differ from dev_ptr, so callers must use it, not dev_ptr, for host diff --git a/src/common/platform/include/common/l2_swimlane_profiling.h b/src/common/platform/include/common/l2_swimlane_profiling.h index e1c85536f3..7fcdf24796 100644 --- a/src/common/platform/include/common/l2_swimlane_profiling.h +++ b/src/common/platform/include/common/l2_swimlane_profiling.h @@ -539,6 +539,10 @@ enum class L2SwimlaneSchedPhaseKind : uint32_t { PredicatedSkip = 12, // Per-task marker for a real task retired inline because // its dispatch predicate evaluated false. Uses the same // phase_data.dummy_task identity payload as DummyTask. + // Outer (sched lane): one bounded Graph Definition materialization slice. + // phase_data.graph_task identifies the ring-0 outer Graph task and + // tasks_processed is the number of nodes patched in this slice. + GraphPrepare = 13, }; /** Index layout of the queue-depth snapshot arrays below: AIC=0, AIV=1, MIX=2. @@ -575,6 +579,10 @@ struct L2SwimlaneAicpuSchedPhaseRecord { uint32_t local_id; // task_id bits [31:0] uint32_t ring_id; // task_id bits [63:32] } dummy_task; + struct { + uint32_t local_id; // outer Graph task_id bits [31:0] + uint32_t ring_id; // outer Graph task_id bits [63:32] + } graph_task; } phase_data; int16_t shared_depth_at_start[L2SWIMLANE_NUM_QUEUE_SHAPES]; // sched->ready_queues[shape].size() int16_t shared_depth_at_end[L2SWIMLANE_NUM_QUEUE_SHAPES]; diff --git a/src/common/platform/include/host/l2_swimlane_collector.h b/src/common/platform/include/host/l2_swimlane_collector.h index 4db5c9db54..ab22d99d2e 100644 --- a/src/common/platform/include/host/l2_swimlane_collector.h +++ b/src/common/platform/include/host/l2_swimlane_collector.h @@ -408,6 +408,12 @@ class L2SwimlaneCollector : public profiling_common::ProfilerBase &records, uint64_t host_start_cycles, + uint64_t host_end_cycles, uint64_t first_publish_cycles + ); + /** * Export collected records as a Chrome Trace Event JSON (swimlane view). * Writes /l2_swimlane_records.json — directory is captured at @@ -530,6 +536,13 @@ class L2SwimlaneCollector : public profiling_common::ProfilerBase> collected_sched_phase_records_; std::vector> collected_orch_phase_records_; + // Host and AICPU clocks do not share an epoch. These records stay in a + // separate stream and the converter aligns host-orch end to device t=0. + std::vector host_orch_phase_records_; + uint64_t host_orch_start_cycles_{0}; + uint64_t host_orch_end_cycles_{0}; + uint64_t host_orch_first_publish_cycles_{0}; + // Core-to-thread mapping (core_id → scheduler thread index, -1 = unassigned) std::vector core_to_thread_; diff --git a/src/common/platform/onboard/host/c_api_shared.cpp b/src/common/platform/onboard/host/c_api_shared.cpp index c6322d65b9..149d8570be 100644 --- a/src/common/platform/onboard/host/c_api_shared.cpp +++ b/src/common/platform/onboard/host/c_api_shared.cpp @@ -117,6 +117,11 @@ static int copy_from_device(void *host_ptr, const void *dev_ptr, size_t size) { } } +static int publish_i32(void *dev_ptr, int32_t value) { + if (dev_ptr == NULL) return -1; + return copy_to_device(dev_ptr, &value, sizeof(value)); +} + static void *register_device_memory_to_host(void *dev_ptr, size_t bytes) { try { return current_runner()->register_device_memory_to_host(dev_ptr, bytes); @@ -236,6 +241,7 @@ static const HostApi g_host_api = { .device_free = device_free, .copy_to_device = copy_to_device, .copy_from_device = copy_from_device, + .publish_i32 = publish_i32, .register_device_memory_to_host = register_device_memory_to_host, .unregister_device_memory_from_host = unregister_device_memory_from_host, .device_memset = device_memset, @@ -724,7 +730,8 @@ int simpler_prepare_run( STRACE("simpler_run.bind"); rc = runner->bind_callable_to_runtime( state->runtime, callable_id, &g_host_api, args, state->config.runtime_env.ring_task_window, - state->config.runtime_env.ring_heap, state->config.runtime_env.ring_dep_pool + state->config.runtime_env.ring_heap, state->config.runtime_env.ring_dep_pool, + state->config.enable_l2_swimlane ); } if (rc != 0) return cleanup_failed_prepare(state, rc, true); @@ -784,11 +791,41 @@ int simpler_launch_run(DeviceContextHandle ctx, RuntimeHandle runtime) { if (attach_rc == 0) { state->adopt_host_thread_state(); state->runner->activate_launch_shape(state->runtime); - { + rc = 0; + int orch_rc = 0; + std::thread orch_thread; + if (state->runtime.has_deferred_host_orchestration()) { + try { + orch_thread = state->runner->create_thread([state, ctx, &orch_rc]() { + pthread_once(&g_runner_key_once, create_runner_key); + pthread_setspecific(g_runner_key, ctx); + STRACE_CONTEXT(state->trace_inv, state->trace_hid, 1); + try { + orch_rc = state->runtime.run_deferred_host_orchestration(&g_host_api); + } catch (...) { + orch_rc = -1; + } + pthread_setspecific(g_runner_key, nullptr); + }); + } catch (...) { + rc = -1; + } + } + if (rc == 0) { STRACE("simpler_run.runner_run"); entered_run = true; - rc = state->runner->run(state->runtime, state->config); + try { + rc = state->runner->run(state->runtime, state->config); + } catch (...) { + rc = -1; + } } + // DeviceRunner opens this gate at the real kernel launch + // boundary. Also open it after an early failure so a + // deferred orchestration thread cannot remain blocked. + state->runtime.notify_deferred_host_execution_started(); + if (orch_thread.joinable()) orch_thread.join(); + if (rc == 0 && orch_rc != 0) rc = orch_rc; } else { rc = attach_rc; } diff --git a/src/common/platform/onboard/host/device_runner_base.cpp b/src/common/platform/onboard/host/device_runner_base.cpp index 16f8a6961a..ec611a61ca 100644 --- a/src/common/platform/onboard/host/device_runner_base.cpp +++ b/src/common/platform/onboard/host/device_runner_base.cpp @@ -994,12 +994,12 @@ uint64_t DeviceRunnerBase::callable_hash(int32_t callable_id) const { extern "C" int bind_callable_to_runtime_impl( Runtime *runtime, const HostApi *api, const ChipStorageTaskArgs *orch_args, void *host_orch_func_ptr, const ArgDirection *signature, int sig_count, const uint64_t *ring_task_window, const uint64_t *ring_heap, - const uint64_t *ring_dep_pool + const uint64_t *ring_dep_pool, int32_t l2_swimlane_level ); int DeviceRunnerBase::bind_callable_to_runtime( Runtime &runtime, int32_t callable_id, const HostApi *api, const void *orch_args, const uint64_t *ring_task_window, - const uint64_t *ring_heap, const uint64_t *ring_dep_pool + const uint64_t *ring_heap, const uint64_t *ring_dep_pool, int32_t l2_swimlane_level ) { auto it = callables_.find(callable_id); if (it == callables_.end()) { @@ -1030,7 +1030,7 @@ int DeviceRunnerBase::bind_callable_to_runtime( return bind_callable_to_runtime_impl( &runtime, api, reinterpret_cast(orch_args), state.host_orch_func_ptr, state.signature.empty() ? nullptr : state.signature.data(), static_cast(state.signature.size()), - ring_task_window, ring_heap, ring_dep_pool + ring_task_window, ring_heap, ring_dep_pool, l2_swimlane_level ); } diff --git a/src/common/platform/onboard/host/device_runner_base.h b/src/common/platform/onboard/host/device_runner_base.h index e5634d7dfd..70c3e43bc7 100644 --- a/src/common/platform/onboard/host/device_runner_base.h +++ b/src/common/platform/onboard/host/device_runner_base.h @@ -467,12 +467,15 @@ class DeviceRunnerBase { * @param orch_args const ChipStorageTaskArgs* for this run (void* to * keep task_interface headers out of this header). * @param ring_task_window Per-ring overrides (trb); ignored by hbg. + * @param l2_swimlane_level Per-run DFX level; hbg uses level 4 to capture + * host-side orchestration concurrent with device execution. * @return 0 on success, non-zero on failure (unregistered id, out-of-range * func_id, or the underlying bind_callable_to_runtime_impl rc). */ int bind_callable_to_runtime( Runtime &runtime, int32_t callable_id, const HostApi *api, const void *orch_args, - const uint64_t *ring_task_window, const uint64_t *ring_heap, const uint64_t *ring_dep_pool + const uint64_t *ring_task_window, const uint64_t *ring_heap, const uint64_t *ring_dep_pool, + int32_t l2_swimlane_level ); /** diff --git a/src/common/platform/shared/aicpu/l2_swimlane_collector_aicpu.cpp b/src/common/platform/shared/aicpu/l2_swimlane_collector_aicpu.cpp index 847da4cd00..230d84ee43 100644 --- a/src/common/platform/shared/aicpu/l2_swimlane_collector_aicpu.cpp +++ b/src/common/platform/shared/aicpu/l2_swimlane_collector_aicpu.cpp @@ -984,6 +984,20 @@ void l2_swimlane_aicpu_record_predicated_skip( record_aicpu_worker_task(thread_idx, L2SwimlaneSchedPhaseKind::PredicatedSkip, complete_time, loop_iter, task_id); } +void l2_swimlane_aicpu_record_graph_prepare( + int thread_idx, uint64_t start_time, uint64_t end_time, uint32_t loop_iter, uint64_t task_id, + uint32_t nodes_materialized +) { + auto *record = acquire_sched_phase_record(thread_idx); + if (record == nullptr) return; + fill_sched_phase_record( + record, L2SwimlaneSchedPhaseKind::GraphPrepare, start_time, end_time, loop_iter, nodes_materialized, + /*shared_at_start=*/nullptr, /*shared_at_end=*/nullptr + ); + record->phase_data.graph_task.local_id = static_cast(task_id); + record->phase_data.graph_task.ring_id = static_cast(task_id >> 32); +} + void l2_swimlane_aicpu_set_orch_thread_idx(int thread_idx) { s_orch_thread_idx = thread_idx; } void l2_swimlane_aicpu_record_orch_phase( diff --git a/src/common/platform/shared/host/l2_swimlane_collector.cpp b/src/common/platform/shared/host/l2_swimlane_collector.cpp index 29e1c5e183..dd9e35b32d 100644 --- a/src/common/platform/shared/host/l2_swimlane_collector.cpp +++ b/src/common/platform/shared/host/l2_swimlane_collector.cpp @@ -893,6 +893,16 @@ void L2SwimlaneCollector::set_core_types(const CoreType *types, int n) { core_types_.assign(types, types + n); } +void L2SwimlaneCollector::set_host_orch_records( + const std::vector &records, uint64_t host_start_cycles, uint64_t host_end_cycles, + uint64_t first_publish_cycles +) { + host_orch_phase_records_ = records; + host_orch_start_cycles_ = host_start_cycles; + host_orch_end_cycles_ = host_end_cycles; + host_orch_first_publish_cycles_ = first_publish_cycles; +} + // JSON v2 emit: the host now dumps raw cycle-domain per-stream records plus // metadata, and `swimlane_converter.py` performs the join (AICore↔AICPU on // reg_task_id, base_time normalization, cycles→µs conversion, sort, core_type @@ -1044,6 +1054,8 @@ int L2SwimlaneCollector::export_swimlane_json() { return "drain_publish"; case L2SwimlaneSchedPhaseKind::AsyncPoll: return "async_poll"; + case L2SwimlaneSchedPhaseKind::GraphPrepare: + return "graph_prepare"; } return "unknown"; }; @@ -1070,6 +1082,11 @@ int L2SwimlaneCollector::export_swimlane_json() { pr.phase_data.dummy_task.local_id; outfile << ", \"task_id\": " << task_id; } + if (pr.kind == L2SwimlaneSchedPhaseKind::GraphPrepare) { + uint64_t task_id = (static_cast(pr.phase_data.graph_task.ring_id) << 32) | + pr.phase_data.graph_task.local_id; + outfile << ", \"task_id\": " << task_id; + } // Queue-depth snapshots — [AIC, AIV, MIX] per L2SwimlaneAicpuSchedPhaseRecord docstring. emit_depth_array("shared_at_start", pr.shared_depth_at_start); emit_depth_array("shared_at_end", pr.shared_depth_at_end); @@ -1114,6 +1131,22 @@ int L2SwimlaneCollector::export_swimlane_json() { } outfile << " ]"; } + if (host_orch_end_cycles_ > host_orch_start_cycles_) { + outfile << ",\n \"host_orchestrator\": {\n"; + outfile << " \"start_cycles\": " << host_orch_start_cycles_ << ",\n"; + outfile << " \"end_cycles\": " << host_orch_end_cycles_ << ",\n"; + outfile << " \"first_publish_cycles\": " << host_orch_first_publish_cycles_ << ",\n"; + outfile << " \"records\": ["; + bool first = true; + for (const auto &pr : host_orch_phase_records_) { + if (!first) outfile << ","; + outfile << "\n {\"submit_idx\": " << pr.submit_idx << ", \"task_id\": " << pr.task_id + << ", \"start_cycles\": " << pr.start_time << ", \"end_cycles\": " << pr.end_time << "}"; + first = false; + } + if (!first) outfile << "\n "; + outfile << "]\n }"; + } } outfile << "\n}\n"; @@ -1241,6 +1274,10 @@ int L2SwimlaneCollector::finalize(L2SwimlaneUnregisterCallback unregister_cb, co collected_aicore_records_.clear(); collected_sched_phase_records_.clear(); collected_orch_phase_records_.clear(); + host_orch_phase_records_.clear(); + host_orch_start_cycles_ = 0; + host_orch_end_cycles_ = 0; + host_orch_first_publish_cycles_ = 0; perf_records_by_collector_.clear(); aicore_records_by_collector_.clear(); sched_phase_records_by_collector_.clear(); diff --git a/src/common/platform/sim/host/c_api_shared.cpp b/src/common/platform/sim/host/c_api_shared.cpp index d9e510b135..46b93f21c7 100644 --- a/src/common/platform/sim/host/c_api_shared.cpp +++ b/src/common/platform/sim/host/c_api_shared.cpp @@ -33,6 +33,7 @@ #include #include +#include #include #include #include @@ -110,6 +111,12 @@ static int copy_from_device(void *host_ptr, const void *dev_ptr, size_t size) { } } +static int publish_i32(void *dev_ptr, int32_t value) { + if (dev_ptr == NULL) return -1; + reinterpret_cast *>(dev_ptr)->store(value, std::memory_order_release); + return 0; +} + static void *register_device_memory_to_host(void *dev_ptr, size_t bytes) { try { return current_runner()->register_device_memory_to_host(dev_ptr, bytes); @@ -230,6 +237,7 @@ static const HostApi g_host_api = { .device_free = device_free, .copy_to_device = copy_to_device, .copy_from_device = copy_from_device, + .publish_i32 = publish_i32, .register_device_memory_to_host = register_device_memory_to_host, .unregister_device_memory_from_host = unregister_device_memory_from_host, .device_memset = device_memset, @@ -631,7 +639,8 @@ int simpler_prepare_run( STRACE("simpler_run.bind"); rc = runner->bind_callable_to_runtime( state->runtime, callable_id, &g_host_api, args, state->config.runtime_env.ring_task_window, - state->config.runtime_env.ring_heap, state->config.runtime_env.ring_dep_pool + state->config.runtime_env.ring_heap, state->config.runtime_env.ring_dep_pool, + state->config.enable_l2_swimlane ); } if (rc != 0) return cleanup_failed_prepare(state, rc, true); @@ -662,10 +671,37 @@ int simpler_launch_run(DeviceContextHandle ctx, RuntimeHandle runtime) { int attach_rc = state->runner->attach_current_thread(state->runner->device_id()); if (attach_rc == 0) { state->adopt_host_thread_state(); - { + rc = 0; + int orch_rc = 0; + std::thread orch_thread; + if (state->runtime.has_deferred_host_orchestration()) { + try { + orch_thread = state->runner->create_thread([state, ctx, &orch_rc]() { + pthread_once(&g_runner_key_once, create_runner_key); + pthread_setspecific(g_runner_key, ctx); + STRACE_CONTEXT(state->trace_inv, state->trace_hid, 1); + try { + orch_rc = state->runtime.run_deferred_host_orchestration(&g_host_api); + } catch (...) { + orch_rc = -1; + } + pthread_setspecific(g_runner_key, nullptr); + }); + } catch (...) { + rc = -1; + } + } + if (rc == 0) { STRACE("simpler_run.runner_run"); - rc = state->runner->run(state->runtime, state->config); + try { + rc = state->runner->run(state->runtime, state->config); + } catch (...) { + rc = -1; + } } + state->runtime.notify_deferred_host_execution_started(); + if (orch_thread.joinable()) orch_thread.join(); + if (rc == 0 && orch_rc != 0) rc = orch_rc; } else { rc = attach_rc; } diff --git a/src/common/platform/sim/host/device_runner_base.cpp b/src/common/platform/sim/host/device_runner_base.cpp index 2b63b13977..741a75dacd 100644 --- a/src/common/platform/sim/host/device_runner_base.cpp +++ b/src/common/platform/sim/host/device_runner_base.cpp @@ -586,12 +586,12 @@ bool SimDeviceRunnerBase::has_callable(int32_t callable_id) const { return calla extern "C" int bind_callable_to_runtime_impl( Runtime *runtime, const HostApi *api, const ChipStorageTaskArgs *orch_args, void *host_orch_func_ptr, const ArgDirection *signature, int sig_count, const uint64_t *ring_task_window, const uint64_t *ring_heap, - const uint64_t *ring_dep_pool + const uint64_t *ring_dep_pool, int32_t l2_swimlane_level ); int SimDeviceRunnerBase::bind_callable_to_runtime( Runtime &runtime, int32_t callable_id, const HostApi *api, const void *orch_args, const uint64_t *ring_task_window, - const uint64_t *ring_heap, const uint64_t *ring_dep_pool + const uint64_t *ring_heap, const uint64_t *ring_dep_pool, int32_t l2_swimlane_level ) { auto it = callables_.find(callable_id); if (it == callables_.end()) { @@ -615,7 +615,7 @@ int SimDeviceRunnerBase::bind_callable_to_runtime( return bind_callable_to_runtime_impl( &runtime, api, reinterpret_cast(orch_args), state.host_orch_func_ptr, state.signature.empty() ? nullptr : state.signature.data(), static_cast(state.signature.size()), - ring_task_window, ring_heap, ring_dep_pool + ring_task_window, ring_heap, ring_dep_pool, l2_swimlane_level ); } diff --git a/src/common/platform/sim/host/device_runner_base.h b/src/common/platform/sim/host/device_runner_base.h index 1225fd5eff..06f1b17df1 100644 --- a/src/common/platform/sim/host/device_runner_base.h +++ b/src/common/platform/sim/host/device_runner_base.h @@ -173,7 +173,8 @@ class SimDeviceRunnerBase { // header). Returns 0 on success, non-zero on failure. int bind_callable_to_runtime( Runtime &runtime, int32_t callable_id, const HostApi *api, const void *orch_args, - const uint64_t *ring_task_window, const uint64_t *ring_heap, const uint64_t *ring_dep_pool + const uint64_t *ring_task_window, const uint64_t *ring_heap, const uint64_t *ring_dep_pool, + int32_t l2_swimlane_level ); // Publish this run's core geometry onto `Runtime` before the graph is diff --git a/tests/st/a2a3/host_build_graph/graph_execution/kernels/orchestration/graph_execution_aic_aiv_orch.cpp b/tests/st/a2a3/host_build_graph/graph_execution/kernels/orchestration/graph_execution_aic_aiv_orch.cpp new file mode 100644 index 0000000000..7048ee7ff0 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/graph_execution/kernels/orchestration/graph_execution_aic_aiv_orch.cpp @@ -0,0 +1,88 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include + +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +#define FUNC_LOG_SQRT 0 +#define FUNC_MATMUL 1 +#define FUNC_ADD_EXP 2 + +namespace { + +// Qwen-style fixed decoder-layer topology: +// +// +--> matmul(weight_1) --+ +// input -> norm --| +--> add + activation +// +--> matmul(weight_2) --+ +// +// A cache miss records two AIC and two AIV tasks. Cache hits submit one outer +// Graph task and let Scheduler materialize this saved topology. +void decoder_layer(const L0TaskArgs &args) { + const Tensor &input = args.tensor(0).ref(); + const Tensor &weight_1 = args.tensor(1).ref(); + const Tensor &weight_2 = args.tensor(2).ref(); + const Tensor &output = args.tensor(3).ref(); + + const std::array shape{input.shapes[0]}; + TensorCreateInfo normalized_info(shape.data(), static_cast(shape.size()), DataType::FLOAT16); + TensorCreateInfo projected_info(shape.data(), static_cast(shape.size()), DataType::FLOAT32); + + L0TaskArgs norm_args; + norm_args.add_input(input); + norm_args.add_output(normalized_info); + TaskOutputTensors normalized_outputs = rt_submit_aiv_task(FUNC_LOG_SQRT, norm_args); + Tensor normalized = normalized_outputs.get_ref(0); + + L0TaskArgs left_args; + left_args.add_input(normalized, weight_1); + left_args.add_output(projected_info); + TaskOutputTensors left_outputs = rt_submit_aic_task(FUNC_MATMUL, left_args); + Tensor left = left_outputs.get_ref(0); + + L0TaskArgs right_args; + right_args.add_input(normalized, weight_2); + right_args.add_output(projected_info); + TaskOutputTensors right_outputs = rt_submit_aic_task(FUNC_MATMUL, right_args); + Tensor right = right_outputs.get_ref(0); + + L0TaskArgs activation_args; + activation_args.add_input(left, right); + activation_args.add_output(output); + rt_submit_aiv_task(FUNC_ADD_EXP, activation_args); +} + +void submit_layer(const L0TaskArgs &args) { rt_submit_graph(&decoder_layer, args); } + +} // namespace + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &args) { + (void)args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 6, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &args) { + for (int32_t output_index = 3; output_index < 6; ++output_index) { + L0TaskArgs layer_args; + layer_args.add_input(args.tensor(0).ref(), args.tensor(1).ref(), args.tensor(2).ref()); + layer_args.add_output(args.tensor(output_index).ref()); + submit_layer(layer_args); + } +} + +} // extern "C" diff --git a/tests/st/a2a3/host_build_graph/graph_execution/kernels/orchestration/graph_execution_mix_spmd_orch.cpp b/tests/st/a2a3/host_build_graph/graph_execution/kernels/orchestration/graph_execution_mix_spmd_orch.cpp new file mode 100644 index 0000000000..0118bbbe95 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/graph_execution/kernels/orchestration/graph_execution_mix_spmd_orch.cpp @@ -0,0 +1,57 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +#define FUNC_SPMD_MIX_AIC 0 +#define FUNC_SPMD_MIX_AIV0 1 +#define FUNC_SPMD_MIX_AIV1 2 + +namespace { + +void mix_spmd_layer(const L0TaskArgs &args) { + MixedKernels kernels; + kernels.aic_kernel_id = FUNC_SPMD_MIX_AIC; + kernels.aiv0_kernel_id = FUNC_SPMD_MIX_AIV0; + kernels.aiv1_kernel_id = FUNC_SPMD_MIX_AIV1; + + L0TaskArgs task_args; + task_args.add_inout(args.tensor(0).ref()); + task_args.add_scalar(int64_t{0}); + task_args.launch_spec.set_block_num(static_cast(rt_available_cluster_count())); + task_args.launch_spec.set_require_sync_start(true); + rt_submit_task(kernels, task_args); +} + +void submit_layer(const L0TaskArgs &args) { rt_submit_graph(&mix_spmd_layer, args); } + +} // namespace + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &args) { + (void)args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 3, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &args) { + for (int32_t output_index = 0; output_index < 3; ++output_index) { + L0TaskArgs layer_args; + layer_args.add_inout(args.tensor(output_index).ref()); + submit_layer(layer_args); + } +} + +} // extern "C" diff --git a/tests/st/a2a3/host_build_graph/graph_execution/kernels/orchestration/graph_execution_orch.cpp b/tests/st/a2a3/host_build_graph/graph_execution/kernels/orchestration/graph_execution_orch.cpp new file mode 100644 index 0000000000..6779032574 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/graph_execution/kernels/orchestration/graph_execution_orch.cpp @@ -0,0 +1,100 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include + +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +#define FUNC_ADD 0 +#define FUNC_ADD_SCALAR 1 +#define FUNC_MUL 2 + +namespace { + +void layer(const L0TaskArgs &args, float left_delta, float right_delta) { + const Tensor &a = args.tensor(0).ref(); + const Tensor &b = args.tensor(1).ref(); + const Tensor &output = args.tensor(2).ref(); + const float ndim_delta = b.ndims == 1 ? 0.0F : 2.0F; + + const std::array shape{a.shapes[0]}; + TensorCreateInfo intermediate(shape.data(), static_cast(shape.size()), DataType::FLOAT32); + + L0TaskArgs add_args; + add_args.add_input(a, b); + add_args.add_output(intermediate); + const std::array external_dep{a.owner_task_id}; + add_args.set_dependencies(external_dep.data(), static_cast(external_dep.size())); + TaskOutputTensors add_outputs = rt_submit_aiv_task(FUNC_ADD, add_args); + Tensor sum = add_outputs.get_ref(0); + + L0TaskArgs fence_args; + fence_args.add_inout(sum); + rt_submit_dummy_task(fence_args); + + L0TaskArgs left_args; + left_args.add_input(sum); + left_args.add_output(intermediate); + left_args.add_scalar(left_delta + ndim_delta); + left_args.set_allow_early_resolve(true); + TaskOutputTensors left_outputs = rt_submit_aiv_task(FUNC_ADD_SCALAR, left_args); + Tensor left = left_outputs.get_ref(0); + + L0TaskArgs right_args; + right_args.add_input(sum); + right_args.add_output(intermediate); + right_args.add_scalar(right_delta + ndim_delta); + TaskOutputTensors right_outputs = rt_submit_aiv_task(FUNC_ADD_SCALAR, right_args); + Tensor right = right_outputs.get_ref(0); + + L0TaskArgs mul_args; + mul_args.add_input(left, right); + mul_args.add_output(output); + rt_submit_aiv_task(FUNC_MUL, mul_args); +} + +void submit_layer(const L0TaskArgs &args) { rt_submit_graph(&layer, args, 1.0F, 2.0F); } + +} // namespace + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &args) { + (void)args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 5, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &args) { + const Tensor &a = args.tensor(0).ref(); + const Tensor &b = args.tensor(1).ref(); + + const std::array shape{a.shapes[0]}; + TensorCreateInfo seeded_input_info(shape.data(), static_cast(shape.size()), DataType::FLOAT32); + L0TaskArgs seed_args; + seed_args.add_input(a); + seed_args.add_output(seeded_input_info); + seed_args.add_scalar(0.0F); + TaskOutputTensors seed_outputs = rt_submit_aiv_task(FUNC_ADD_SCALAR, seed_args); + Tensor seeded_a = seed_outputs.get_ref(0); + + for (int32_t output_index = 2; output_index < 5; ++output_index) { + L0TaskArgs layer_args; + layer_args.add_input(seeded_a, b); + layer_args.add_output(args.tensor(output_index).ref()); + submit_layer(layer_args); // first call records; later calls submit one Graph task + } +} + +} // extern "C" diff --git a/tests/st/a2a3/host_build_graph/graph_execution/kernels/orchestration/qwen3_14b_3layer_graph_execution.cpp b/tests/st/a2a3/host_build_graph/graph_execution/kernels/orchestration/qwen3_14b_3layer_graph_execution.cpp new file mode 100644 index 0000000000..e3a33a9f70 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/graph_execution/kernels/orchestration/qwen3_14b_3layer_graph_execution.cpp @@ -0,0 +1,2087 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ +// Orchestration Function: decode_fwd_layers +// Generated by PyPTO IR Compiler + +#include "runtime.h" +#include + +#include +#include +#include + +#include "pto_orchestration_api.h" + +namespace { + +constexpr uint64_t kScratchAlignmentBytes = 1024; +constexpr uint32_t kBfloat16ScratchElements = 524288; +constexpr uint32_t kFloat32ScratchElements = 918016; + +class ScratchArena { +public: + explicit ScratchArena(const Tensor &storage) : + storage_(storage) {} + + Tensor allocate(const TensorCreateInfo &create_info) { + always_assert(create_info.dtype == storage_.dtype); + const uint64_t element_size = get_element_size(storage_.dtype); + const uint64_t alignment = kScratchAlignmentBytes / element_size; + cursor_ = (cursor_ + alignment - 1) / alignment * alignment; + const uint64_t tensor_elements = create_info.buffer_size_bytes() / element_size; + const uint64_t aligned_elements = (tensor_elements + alignment - 1) / alignment * alignment; + always_assert((cursor_ + aligned_elements) * element_size <= storage_.buffer.size); + + Tensor tensor; + init_tensor_from_create_info( + tensor, create_info, reinterpret_cast(static_cast(storage_.buffer.addr)), + storage_.buffer.size + ); + tensor.owner_task_id = storage_.owner_task_id; + tensor.start_offset = storage_.start_offset + cursor_; + cursor_ += aligned_elements; + return tensor; + } + +private: + const Tensor &storage_; + uint64_t cursor_{0}; +}; + +} // namespace + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 20, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + // External tensors + const Tensor &ext_hidden_states = orch_args.tensor(0).ref(); + const Tensor &ext_input_rms_weight = orch_args.tensor(1).ref(); + const Tensor &ext_wq = orch_args.tensor(2).ref(); + const Tensor &ext_wk = orch_args.tensor(3).ref(); + const Tensor &ext_wv = orch_args.tensor(4).ref(); + const Tensor &ext_q_norm_weight = orch_args.tensor(5).ref(); + const Tensor &ext_k_norm_weight = orch_args.tensor(6).ref(); + const Tensor &ext_seq_lens = orch_args.tensor(7).ref(); + const Tensor &ext_block_table = orch_args.tensor(8).ref(); + const Tensor &ext_slot_mapping = orch_args.tensor(9).ref(); + const Tensor &ext_rope_cos = orch_args.tensor(10).ref(); + const Tensor &ext_rope_sin = orch_args.tensor(11).ref(); + const Tensor &ext_k_cache = orch_args.tensor(12).ref(); + const Tensor &ext_v_cache = orch_args.tensor(13).ref(); + const Tensor &ext_wo = orch_args.tensor(14).ref(); + const Tensor &ext_w_gate = orch_args.tensor(15).ref(); + const Tensor &ext_w_up = orch_args.tensor(16).ref(); + const Tensor &ext_w_down = orch_args.tensor(17).ref(); + const Tensor &ext_post_rms_weight = orch_args.tensor(18).ref(); + const Tensor &ext_out = orch_args.tensor(19).ref(); + + // Dynamic-dim symbols (extent of the declaring argument) + int64_t BLOCK_TABLE_FLAT_DYN = (int64_t)orch_args.tensor(8).ref().shapes[0]; + int64_t KV_CACHE_ROWS_DYN = (int64_t)orch_args.tensor(12).ref().shapes[0]; + + PTO2_SCOPE() { + uint32_t pa_metadata_ci_shapes[1] = {27840}; + TensorCreateInfo pa_metadata_ci(pa_metadata_ci_shapes, 1, DataType::UINT8); + uint32_t pa_workspace_ci_shapes[1] = {66132544}; + TensorCreateInfo pa_workspace_ci(pa_workspace_ci_shapes, 1, DataType::UINT8); + uint32_t cur_ci_shapes[2] = {16, 5120}; + TensorCreateInfo cur_ci(cur_ci_shapes, 2, DataType::FLOAT32); + uint32_t normed_ci_shapes[2] = {16, 5120}; + TensorCreateInfo normed_ci(normed_ci_shapes, 2, DataType::BFLOAT16); + TaskOutputTensors alloc_0 = alloc_tensors(pa_metadata_ci, pa_workspace_ci, cur_ci, normed_ci); + const Tensor &pa_metadata = alloc_0.get_ref(0); + const Tensor &pa_workspace = alloc_0.get_ref(1); + const Tensor &cur = alloc_0.get_ref(2); + const Tensor &normed = alloc_0.get_ref(3); + int64_t pa_num_layers = 3; + int64_t pa_num_pages = (KV_CACHE_ROWS_DYN / (pa_num_layers * 1024)); + int64_t pa_max_blocks = (BLOCK_TABLE_FLAT_DYN / 16); + int32_t pa_num_pages_i32 = static_cast(pa_num_pages); + int32_t pa_max_blocks_i32 = static_cast(pa_max_blocks); + + // Spmd pa_tiling: paged_attention_tiling_cce + L0TaskArgs params_t0; + params_t0.add_input(ext_seq_lens); + params_t0.add_output(pa_metadata); + params_t0.add_scalar(pa_max_blocks_i32); + params_t0.add_scalar(pa_num_pages_i32); + params_t0.launch_spec.set_block_num(1); + params_t0.set_allow_early_resolve(true); + TaskOutputTensors task_0_outs = rt_submit_aiv_task(0, params_t0); + PTO2TaskId tiling_tid_inline0 = task_0_outs.task_id(); + PTO2TaskId pa_tiling_tid = tiling_tid_inline0; + PTO2TaskId prev_out_tid[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + prev_out_tid[__init_i] = PTO2TaskId::invalid(); + + // Phase-fence barrier 0: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_0; + TaskOutputTensors phase_fence_barrier_0_outs = rt_submit_dummy_task(params_phase_fence_barrier_0); + PTO2TaskId t = phase_fence_barrier_0_outs.task_id(); + prev_out_tid[0] = t; + for (int64_t cb0 = 0; cb0 < 16; cb0 += 16) { + PTO2_SCOPE() { + // Task 1: copy_hidden + L0TaskArgs params_t1; + params_t1.add_output(cur); + params_t1.add_input(ext_hidden_states); + params_t1.add_scalar(cb0); + TaskOutputTensors task_1_outs = rt_submit_aiv_task(1, params_t1); + PTO2TaskId ch_tid = task_1_outs.task_id(); + prev_out_tid[0] = ch_tid; + } + } + PTO2TaskId prev_normed_tid[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + prev_normed_tid[__init_i] = PTO2TaskId::invalid(); + PTO2_SCOPE(PTO2ScopeMode::MANUAL) { + PTO2TaskId _submit_deps_buf[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v3 = prev_out_tid[0]; + _submit_deps_buf[0] = t__tmp_v3; + + // Spmd x_gamma0_spmd: x_gamma0 + L0TaskArgs params_t2; + params_t2.add_output(normed); + params_t2.add_input(cur); + params_t2.add_input(ext_input_rms_weight); + params_t2.launch_spec.set_block_num(5); + params_t2.set_allow_early_resolve(true); + PTO2TaskId params_t2_deps[1]; + uint32_t params_t2_deps_count = 0; + if (_submit_deps_buf[0].is_valid()) params_t2_deps[params_t2_deps_count++] = _submit_deps_buf[0]; + params_t2.set_dependencies(params_t2_deps, params_t2_deps_count); + TaskOutputTensors task_2_outs = rt_submit_aiv_task(2, params_t2); + PTO2TaskId xgamma_tid = task_2_outs.task_id(); + prev_normed_tid[0] = xgamma_tid; + } + Tensor cur__rv_v7 = cur; + Tensor normed__rv_v5 = normed; + uint32_t bf16_scratch_shapes[1] = {kBfloat16ScratchElements}; + TensorCreateInfo bf16_scratch_ci(bf16_scratch_shapes, 1, DataType::BFLOAT16); + uint32_t fp32_scratch_shapes[1] = {kFloat32ScratchElements}; + TensorCreateInfo fp32_scratch_ci(fp32_scratch_shapes, 1, DataType::FLOAT32); + uint32_t next_hidden_ci_shapes[2] = {16, 5120}; + TensorCreateInfo next_hidden_ci(next_hidden_ci_shapes, 2, DataType::FLOAT32); + uint32_t next_normed_ci_shapes[2] = {16, 5120}; + TensorCreateInfo next_normed_ci(next_normed_ci_shapes, 2, DataType::BFLOAT16); + + // Allocate all per-layer boundary and scratch storage before recording. + // This keeps replay iterations to one Graph submit each. + TaskOutputTensors layer_storage = alloc_tensors( + next_hidden_ci, next_normed_ci, bf16_scratch_ci, fp32_scratch_ci, next_hidden_ci, next_normed_ci, + bf16_scratch_ci, fp32_scratch_ci, next_hidden_ci, next_normed_ci, bf16_scratch_ci, fp32_scratch_ci + ); + for (int64_t i = 0; i < 3; i += 1) { + PTO2_SCOPE() { + const uint32_t layer = static_cast(i); + const uint32_t storage_base = layer * 4; + const Tensor &next_hidden = layer_storage.get_ref(storage_base); + const Tensor &next_normed = layer_storage.get_ref(storage_base + 1); + const Tensor &bf16_scratch = layer_storage.get_ref(storage_base + 2); + const Tensor &fp32_scratch = layer_storage.get_ref(storage_base + 3); + auto layer_view = [](const Tensor &tensor, uint32_t layer_index, uint32_t rows) { + uint32_t offsets[2] = {layer_index * rows, 0}; + uint32_t shapes[2] = {rows, tensor.shapes[1]}; + return tensor.view(shapes, offsets); + }; + const uint32_t cache_rows = ext_k_cache.shapes[0] / 3; + Tensor next_input_rms_weight = layer_view(ext_input_rms_weight, std::min(layer + 1, 2), 1); + Tensor wq = layer_view(ext_wq, layer, 5120); + Tensor wk = layer_view(ext_wk, layer, 5120); + Tensor wv = layer_view(ext_wv, layer, 5120); + Tensor q_norm_w_inline124 = layer_view(ext_q_norm_weight, layer, 1); + Tensor k_norm_w_inline114 = layer_view(ext_k_norm_weight, layer, 1); + Tensor k_cache = layer_view(ext_k_cache, layer, cache_rows); + Tensor v_cache = layer_view(ext_v_cache, layer, cache_rows); + Tensor wo = layer_view(ext_wo, layer, 5120); + Tensor w_gate = layer_view(ext_w_gate, layer, 5120); + Tensor w_up = layer_view(ext_w_up, layer, 5120); + Tensor w_down = layer_view(ext_w_down, layer, 17408); + Tensor post_rms_weight = layer_view(ext_post_rms_weight, layer, 1); + Tensor graph_normed = normed__rv_v5; + graph_normed.owner_task_id = prev_normed_tid[0]; + + L0TaskArgs graph_args; + graph_args.add_input(cur__rv_v7); + graph_args.add_input(graph_normed); + graph_args.add_inout(next_hidden); + graph_args.add_inout(next_normed); + graph_args.add_input(next_input_rms_weight); + graph_args.add_input(wq); + graph_args.add_input(wk); + graph_args.add_input(wv); + graph_args.add_input(q_norm_w_inline124); + graph_args.add_input(k_norm_w_inline114); + graph_args.add_input(ext_seq_lens); + graph_args.add_input(ext_block_table); + graph_args.add_input(ext_slot_mapping); + graph_args.add_input(ext_rope_cos); + graph_args.add_input(ext_rope_sin); + graph_args.add_inout(k_cache); + graph_args.add_inout(v_cache); + graph_args.add_input(wo); + graph_args.add_input(w_gate); + graph_args.add_input(w_up); + graph_args.add_input(w_down); + graph_args.add_input(post_rms_weight); + graph_args.add_inout(pa_workspace); + graph_args.add_inout(pa_metadata); + graph_args.add_inout(bf16_scratch); + graph_args.add_inout(fp32_scratch); + + auto layer_definition = [](const L0TaskArgs &args) { + const Tensor &cur__rv_v7 = args.tensor(0).ref(); + const Tensor &normed__rv_v5 = args.tensor(1).ref(); + const Tensor &next_hidden = args.tensor(2).ref(); + const Tensor &next_normed = args.tensor(3).ref(); + const Tensor &ext_input_rms_weight = args.tensor(4).ref(); + const Tensor &ext_wq = args.tensor(5).ref(); + const Tensor &ext_wk = args.tensor(6).ref(); + const Tensor &ext_wv = args.tensor(7).ref(); + const Tensor &q_norm_w_inline124 = args.tensor(8).ref(); + const Tensor &k_norm_w_inline114 = args.tensor(9).ref(); + const Tensor &ext_seq_lens = args.tensor(10).ref(); + const Tensor &ext_block_table = args.tensor(11).ref(); + const Tensor &ext_slot_mapping = args.tensor(12).ref(); + const Tensor &ext_rope_cos = args.tensor(13).ref(); + const Tensor &ext_rope_sin = args.tensor(14).ref(); + const Tensor &ext_k_cache = args.tensor(15).ref(); + const Tensor &ext_v_cache = args.tensor(16).ref(); + const Tensor &ext_wo = args.tensor(17).ref(); + const Tensor &ext_w_gate = args.tensor(18).ref(); + const Tensor &ext_w_up = args.tensor(19).ref(); + const Tensor &ext_w_down = args.tensor(20).ref(); + const Tensor &ext_post_rms_weight = args.tensor(21).ref(); + const Tensor &pa_workspace = args.tensor(22).ref(); + const Tensor &pa_metadata = args.tensor(23).ref(); + const Tensor &bf16_scratch = args.tensor(24).ref(); + const Tensor &fp32_scratch = args.tensor(25).ref(); + ScratchArena bf16_arena(bf16_scratch); + ScratchArena fp32_arena(fp32_scratch); + + constexpr int64_t i = 0; + constexpr int64_t next_gamma_idx = 0; + constexpr int64_t layer_hidden_base_inline151 = 0; + constexpr int64_t layer_inter_base_inline107 = 0; + constexpr int64_t layer_cache_base_inline193 = 0; + PTO2TaskId prev_normed_tid[1] = {normed__rv_v5.owner_task_id}; + + uint32_t inv_rms_states_inline176_ci_shapes[2] = {16, 1}; + TensorCreateInfo inv_rms_states_inline176_ci( + inv_rms_states_inline176_ci_shapes, 2, DataType::FLOAT32 + ); + uint32_t q_proj_inline139_ci_shapes[2] = {16, 5120}; + TensorCreateInfo q_proj_inline139_ci(q_proj_inline139_ci_shapes, 2, DataType::FLOAT32); + uint32_t k_proj_inline135_ci_shapes[2] = {16, 1024}; + TensorCreateInfo k_proj_inline135_ci(k_proj_inline135_ci_shapes, 2, DataType::FLOAT32); + uint32_t v_proj_inline255_ci_shapes[2] = {16, 1024}; + TensorCreateInfo v_proj_inline255_ci(v_proj_inline255_ci_shapes, 2, DataType::FLOAT32); + uint32_t q_tnd_flat_inline127_ci_shapes[2] = {640, 128}; + TensorCreateInfo q_tnd_flat_inline127_ci(q_tnd_flat_inline127_ci_shapes, 2, DataType::BFLOAT16); + uint32_t attn_out_inline282_ci_shapes[2] = {16, 5120}; + TensorCreateInfo attn_out_inline282_ci(attn_out_inline282_ci_shapes, 2, DataType::BFLOAT16); + uint32_t down_acc_all_inline168_ci_shapes[2] = {16, 5120}; + TensorCreateInfo down_acc_all_inline168_ci(down_acc_all_inline168_ci_shapes, 2, DataType::FLOAT32); + uint32_t gate_acc_all_inline203_ci_shapes[2] = {16, 17408}; + TensorCreateInfo gate_acc_all_inline203_ci(gate_acc_all_inline203_ci_shapes, 2, DataType::FLOAT32); + uint32_t up_acc_all_inline303_ci_shapes[2] = {16, 17408}; + TensorCreateInfo up_acc_all_inline303_ci(up_acc_all_inline303_ci_shapes, 2, DataType::FLOAT32); + uint32_t attn_proj_fp32_inline220_ci_shapes[2] = {16, 5120}; + TensorCreateInfo attn_proj_fp32_inline220_ci( + attn_proj_fp32_inline220_ci_shapes, 2, DataType::FLOAT32 + ); + uint32_t post_norm_partial_inline118_ci_shapes[2] = {16, 5120}; + TensorCreateInfo post_norm_partial_inline118_ci( + post_norm_partial_inline118_ci_shapes, 2, DataType::FLOAT32 + ); + uint32_t mlp_norm_in_inline71_ci_shapes[2] = {16, 5120}; + TensorCreateInfo mlp_norm_in_inline71_ci(mlp_norm_in_inline71_ci_shapes, 2, DataType::BFLOAT16); + uint32_t inv_rms_tile_inline126_ci_shapes[2] = {16, 1}; + TensorCreateInfo inv_rms_tile_inline126_ci(inv_rms_tile_inline126_ci_shapes, 2, DataType::FLOAT32); + uint32_t mlp_tile_inline149_ci_shapes[2] = {16, 17408}; + TensorCreateInfo mlp_tile_inline149_ci(mlp_tile_inline149_ci_shapes, 2, DataType::BFLOAT16); + + Tensor inv_rms_states_inline176 = fp32_arena.allocate(inv_rms_states_inline176_ci); + Tensor q_proj_inline139 = fp32_arena.allocate(q_proj_inline139_ci); + Tensor k_proj_inline135 = fp32_arena.allocate(k_proj_inline135_ci); + Tensor v_proj_inline255 = fp32_arena.allocate(v_proj_inline255_ci); + Tensor q_tnd_flat_inline127 = bf16_arena.allocate(q_tnd_flat_inline127_ci); + Tensor attn_out_inline282 = bf16_arena.allocate(attn_out_inline282_ci); + Tensor down_acc_all_inline168 = fp32_arena.allocate(down_acc_all_inline168_ci); + Tensor gate_acc_all_inline203 = fp32_arena.allocate(gate_acc_all_inline203_ci); + Tensor up_acc_all_inline303 = fp32_arena.allocate(up_acc_all_inline303_ci); + Tensor attn_proj_fp32_inline220 = fp32_arena.allocate(attn_proj_fp32_inline220_ci); + Tensor post_norm_partial_inline118 = fp32_arena.allocate(post_norm_partial_inline118_ci); + Tensor mlp_norm_in_inline71 = bf16_arena.allocate(mlp_norm_in_inline71_ci); + Tensor inv_rms_tile_inline126 = fp32_arena.allocate(inv_rms_tile_inline126_ci); + Tensor mlp_tile_inline149 = bf16_arena.allocate(mlp_tile_inline149_ci); + + PTO2TaskId down_tids_inline156[85]; + for (int64_t __init_i = 0; __init_i < 85; ++__init_i) + down_tids_inline156[__init_i] = PTO2TaskId::invalid(); + PTO2_SCOPE(PTO2ScopeMode::MANUAL) { + // Phase-fence barrier 1: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_1; + TaskOutputTensors phase_fence_barrier_1_outs = + rt_submit_dummy_task(params_phase_fence_barrier_1); + PTO2TaskId seed_dummy_inline49 = phase_fence_barrier_1_outs.task_id(); + PTO2TaskId prev_normed_seed_deps_inline120[2]; + for (int64_t __init_i = 0; __init_i < 2; ++__init_i) + prev_normed_seed_deps_inline120[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v7 = prev_normed_tid[0]; + prev_normed_seed_deps_inline120[0] = t__tmp_v7; + prev_normed_seed_deps_inline120[1] = seed_dummy_inline49; + + // Task 3: attn_out_seed + L0TaskArgs params_t3; + params_t3.add_input(attn_out_inline282); + params_t3.set_allow_early_resolve(true); + TaskOutputTensors task_3_outs = rt_submit_aiv_task(3, params_t3); + PTO2TaskId attn_out_seed_tid_inline116 = task_3_outs.task_id(); + PTO2TaskId _submit_deps_buf_inline42[2]; + for (int64_t __init_i = 0; __init_i < 2; ++__init_i) + _submit_deps_buf_inline42[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v8 = prev_normed_seed_deps_inline120[0]; + _submit_deps_buf_inline42[0] = t__tmp_v8; + PTO2TaskId t__tmp_v9 = prev_normed_seed_deps_inline120[1]; + _submit_deps_buf_inline42[1] = t__tmp_v9; + + // Task 4: rms_recip + L0TaskArgs params_t4; + params_t4.add_input(cur__rv_v7); + params_t4.add_inout(inv_rms_states_inline176); + PTO2TaskId params_t4_deps[2]; + uint32_t params_t4_deps_count = 0; + if (_submit_deps_buf_inline42[0].is_valid()) + params_t4_deps[params_t4_deps_count++] = _submit_deps_buf_inline42[0]; + if (_submit_deps_buf_inline42[1].is_valid()) + params_t4_deps[params_t4_deps_count++] = _submit_deps_buf_inline42[1]; + params_t4.set_dependencies(params_t4_deps, params_t4_deps_count); + params_t4.set_allow_early_resolve(true); + TaskOutputTensors task_4_outs = rt_submit_aiv_task(4, params_t4); + PTO2TaskId rms_tid_inline148 = task_4_outs.task_id(); + + // Task 5: q_seed + L0TaskArgs params_t5; + params_t5.add_inout(q_proj_inline139); + params_t5.set_allow_early_resolve(true); + TaskOutputTensors task_5_outs = rt_submit_aiv_task(5, params_t5); + PTO2TaskId q_seed_tid_inline162 = task_5_outs.task_id(); + PTO2TaskId prev_normed_q_deps_inline105[2]; + for (int64_t __init_i = 0; __init_i < 2; ++__init_i) + prev_normed_q_deps_inline105[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v17 = prev_normed_tid[0]; + prev_normed_q_deps_inline105[0] = t__tmp_v17; + prev_normed_q_deps_inline105[1] = q_seed_tid_inline162; + PTO2TaskId _submit_deps_buf_inline182[2]; + for (int64_t __init_i = 0; __init_i < 2; ++__init_i) + _submit_deps_buf_inline182[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v18 = prev_normed_q_deps_inline105[0]; + _submit_deps_buf_inline182[0] = t__tmp_v18; + PTO2TaskId t__tmp_v19 = prev_normed_q_deps_inline105[1]; + _submit_deps_buf_inline182[1] = t__tmp_v19; + + // Spmd q_proj_spmd: q_proj + L0TaskArgs params_t6; + params_t6.add_inout(q_proj_inline139); + params_t6.add_input(normed__rv_v5); + params_t6.add_input(ext_wq); + params_t6.add_scalar(layer_hidden_base_inline151); + params_t6.launch_spec.set_block_num(50); + params_t6.set_allow_early_resolve(true); + PTO2TaskId params_t6_deps[2]; + uint32_t params_t6_deps_count = 0; + if (_submit_deps_buf_inline182[0].is_valid()) + params_t6_deps[params_t6_deps_count++] = _submit_deps_buf_inline182[0]; + if (_submit_deps_buf_inline182[1].is_valid()) + params_t6_deps[params_t6_deps_count++] = _submit_deps_buf_inline182[1]; + params_t6.set_dependencies(params_t6_deps, params_t6_deps_count); + TaskOutputTensors task_6_outs = rt_submit_aic_task(6, params_t6); + PTO2TaskId q_proj_tid_inline183 = task_6_outs.task_id(); + PTO2TaskId _submit_deps_buf_inline261[2]; + for (int64_t __init_i = 0; __init_i < 2; ++__init_i) + _submit_deps_buf_inline261[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v24 = prev_normed_seed_deps_inline120[0]; + _submit_deps_buf_inline261[0] = t__tmp_v24; + PTO2TaskId t__tmp_v25 = prev_normed_seed_deps_inline120[1]; + _submit_deps_buf_inline261[1] = t__tmp_v25; + + // Task 7: kv_seed + L0TaskArgs params_t7; + params_t7.add_inout(k_proj_inline135); + params_t7.add_inout(v_proj_inline255); + PTO2TaskId params_t7_deps[2]; + uint32_t params_t7_deps_count = 0; + if (_submit_deps_buf_inline261[0].is_valid()) + params_t7_deps[params_t7_deps_count++] = _submit_deps_buf_inline261[0]; + if (_submit_deps_buf_inline261[1].is_valid()) + params_t7_deps[params_t7_deps_count++] = _submit_deps_buf_inline261[1]; + params_t7.set_dependencies(params_t7_deps, params_t7_deps_count); + TaskOutputTensors task_7_outs = rt_submit_aiv_task(7, params_t7); + PTO2TaskId kv_seed_tid_inline238 = task_7_outs.task_id(); + PTO2TaskId _submit_deps_buf_inline267[2]; + for (int64_t __init_i = 0; __init_i < 2; ++__init_i) + _submit_deps_buf_inline267[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v28 = prev_normed_seed_deps_inline120[0]; + _submit_deps_buf_inline267[0] = t__tmp_v28; + PTO2TaskId t__tmp_v29 = prev_normed_seed_deps_inline120[1]; + _submit_deps_buf_inline267[1] = t__tmp_v29; + + // Task 8: mlp_out_seed + L0TaskArgs params_t8; + params_t8.add_inout(down_acc_all_inline168); + params_t8.add_inout(gate_acc_all_inline203); + params_t8.add_inout(up_acc_all_inline303); + params_t8.add_inout(attn_proj_fp32_inline220); + PTO2TaskId params_t8_deps[2]; + uint32_t params_t8_deps_count = 0; + if (_submit_deps_buf_inline267[0].is_valid()) + params_t8_deps[params_t8_deps_count++] = _submit_deps_buf_inline267[0]; + if (_submit_deps_buf_inline267[1].is_valid()) + params_t8_deps[params_t8_deps_count++] = _submit_deps_buf_inline267[1]; + params_t8.set_dependencies(params_t8_deps, params_t8_deps_count); + params_t8.set_allow_early_resolve(true); + TaskOutputTensors task_8_outs = rt_submit_aiv_task(8, params_t8); + PTO2TaskId mlp_out_seed_tid_inline206 = task_8_outs.task_id(); + + // Spmd k_proj_spmd: k_proj + L0TaskArgs params_t9; + params_t9.add_inout(k_proj_inline135); + params_t9.add_input(normed__rv_v5); + params_t9.add_input(ext_wk); + params_t9.add_scalar(layer_hidden_base_inline151); + params_t9.launch_spec.set_block_num(10); + params_t9.set_allow_early_resolve(true); + PTO2TaskId params_t9_deps[1]; + uint32_t params_t9_deps_count = 0; + params_t9_deps[params_t9_deps_count++] = kv_seed_tid_inline238; + params_t9.set_dependencies(params_t9_deps, params_t9_deps_count); + TaskOutputTensors task_9_outs = rt_submit_aic_task(9, params_t9); + PTO2TaskId k_proj_tid_inline136 = task_9_outs.task_id(); + + // Spmd v_proj_spmd: v_proj + L0TaskArgs params_t10; + params_t10.add_inout(v_proj_inline255); + params_t10.add_input(normed__rv_v5); + params_t10.add_input(ext_wv); + params_t10.add_scalar(layer_hidden_base_inline151); + params_t10.launch_spec.set_block_num(10); + params_t10.set_allow_early_resolve(true); + PTO2TaskId params_t10_deps[1]; + uint32_t params_t10_deps_count = 0; + params_t10_deps[params_t10_deps_count++] = kv_seed_tid_inline238; + params_t10.set_dependencies(params_t10_deps, params_t10_deps_count); + TaskOutputTensors task_10_outs = rt_submit_aic_task(10, params_t10); + PTO2TaskId v_proj_tid_inline63 = task_10_outs.task_id(); + uint32_t q_tnd_inline191_shapes[3] = {16, 40, 128}; + Tensor q_tnd_inline191 = q_tnd_flat_inline127.reshape(q_tnd_inline191_shapes, 3); + uint32_t attn_out_tnd_inline79_shapes[3] = {16, 40, 128}; + Tensor attn_out_tnd_inline79 = attn_out_inline282.reshape(attn_out_tnd_inline79_shapes, 3); + int64_t attention_core_num_inline188 = 24; + + // Group paged_attention_rope_cce: MixedKernels (AIC + AIV lanes) + L0TaskArgs params_t11; + params_t11.add_inout(attn_out_tnd_inline79); + params_t11.add_inout(q_tnd_inline191); + params_t11.add_inout(ext_k_cache); + params_t11.add_inout(ext_v_cache); + params_t11.add_input(ext_block_table); + params_t11.add_inout(pa_workspace); + params_t11.add_inout(pa_metadata); + params_t11.add_input(q_proj_inline139); + params_t11.add_input(k_proj_inline135); + params_t11.add_input(v_proj_inline255); + params_t11.add_input(q_norm_w_inline124); + params_t11.add_input(k_norm_w_inline114); + params_t11.add_input(ext_rope_cos); + params_t11.add_input(ext_rope_sin); + params_t11.add_input(inv_rms_states_inline176); + params_t11.add_input(ext_slot_mapping); + params_t11.add_input(ext_seq_lens); + params_t11.add_scalar(layer_cache_base_inline193); + MixedKernels mixed_11 = {11, 12, 12}; + params_t11.launch_spec.set_block_num(attention_core_num_inline188); + params_t11.launch_spec.set_require_sync_start(true); + params_t11.set_allow_early_resolve(true); + PTO2TaskId params_t11_deps[7]; + uint32_t params_t11_deps_count = 0; + params_t11_deps[params_t11_deps_count++] = q_proj_tid_inline183; + params_t11_deps[params_t11_deps_count++] = k_proj_tid_inline136; + params_t11_deps[params_t11_deps_count++] = v_proj_tid_inline63; + params_t11_deps[params_t11_deps_count++] = rms_tid_inline148; + params_t11_deps[params_t11_deps_count++] = attn_out_seed_tid_inline116; + params_t11_deps[params_t11_deps_count++] = mlp_out_seed_tid_inline206; + params_t11.set_dependencies(params_t11_deps, params_t11_deps_count); + TaskOutputTensors task_11_outs = rt_submit_task(mixed_11, params_t11); + const Tensor &attn_out_tnd_inline79__ssa_v1 = attn_out_tnd_inline79; + PTO2TaskId attn_done_tid_inline78 = task_11_outs.task_id(); + uint32_t attn_out_inline282__ssa_v4_shapes[2] = {16, 5120}; + Tensor attn_out_inline282__ssa_v4 = + attn_out_tnd_inline79__ssa_v1.reshape(attn_out_inline282__ssa_v4_shapes, 2); + PTO2TaskId silu_tids_inline265[17]; + for (int64_t __init_i = 0; __init_i < 17; ++__init_i) + silu_tids_inline265[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId gate_tids_inline56[85]; + for (int64_t __init_i = 0; __init_i < 85; ++__init_i) + gate_tids_inline56[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId up_tids_inline310[85]; + for (int64_t __init_i = 0; __init_i < 85; ++__init_i) + up_tids_inline310[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId cast_tids_inline88[5]; + for (int64_t __init_i = 0; __init_i < 5; ++__init_i) + cast_tids_inline88[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId gate_late_tids_inline249[5]; + for (int64_t __init_i = 0; __init_i < 5; ++__init_i) + gate_late_tids_inline249[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId up_late_tids_inline69[5]; + for (int64_t __init_i = 0; __init_i < 5; ++__init_i) + up_late_tids_inline69[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId out_tids_inline271[50]; + for (int64_t __init_i = 0; __init_i < 50; ++__init_i) + out_tids_inline271[__init_i] = PTO2TaskId::invalid(); + + // Phase-fence barrier 2: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_2; + PTO2TaskId params_phase_fence_barrier_2_deps[1]; + uint32_t params_phase_fence_barrier_2_deps_count = 0; + params_phase_fence_barrier_2_deps[params_phase_fence_barrier_2_deps_count++] = + attn_done_tid_inline78; + params_phase_fence_barrier_2.set_dependencies( + params_phase_fence_barrier_2_deps, params_phase_fence_barrier_2_deps_count + ); + PTO2TaskId out_proj_dummy_inline257 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_2_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_2_outs = + rt_submit_dummy_task(params_phase_fence_barrier_2); + out_proj_dummy_inline257 = phase_fence_barrier_2_outs.task_id(); + } + int64_t N_OUT_DIRECT_inline61 = 26; + for (int64_t out_idx_inline74 = 0; out_idx_inline74 < N_OUT_DIRECT_inline61; + out_idx_inline74 += 1) { + int64_t n_out_proj_inline185 = (out_idx_inline74 / 5); + int64_t k_split_out_inline66 = (out_idx_inline74 % 5); + int64_t n_op_inline64 = (n_out_proj_inline185 * 512); + int64_t k_op_inline266 = (k_split_out_inline66 * 1024); + + // Task 12: out_proj + L0TaskArgs params_t12; + params_t12.add_input(attn_out_inline282__ssa_v4); + params_t12.add_input(ext_wo); + params_t12.add_inout(attn_proj_fp32_inline220); + params_t12.add_scalar(k_op_inline266); + params_t12.add_scalar(layer_hidden_base_inline151); + params_t12.add_scalar(n_op_inline64); + PTO2TaskId params_t12_deps[1]; + uint32_t params_t12_deps_count = 0; + if (out_proj_dummy_inline257.is_valid()) + params_t12_deps[params_t12_deps_count++] = out_proj_dummy_inline257; + params_t12.set_dependencies(params_t12_deps, params_t12_deps_count); + TaskOutputTensors task_12_outs = rt_submit_aic_task(13, params_t12); + PTO2TaskId out_tid_inline141 = task_12_outs.task_id(); + out_tids_inline271[out_idx_inline74] = out_tid_inline141; + } + + // Spmd out_proj_spmd: out_proj_0 + L0TaskArgs params_t13; + params_t13.add_input(attn_out_inline282__ssa_v4); + params_t13.add_input(ext_wo); + params_t13.add_inout(attn_proj_fp32_inline220); + params_t13.add_scalar(N_OUT_DIRECT_inline61); + params_t13.add_scalar(layer_hidden_base_inline151); + params_t13.launch_spec.set_block_num(24); + PTO2TaskId params_t13_deps[1]; + uint32_t params_t13_deps_count = 0; + params_t13_deps[params_t13_deps_count++] = attn_done_tid_inline78; + params_t13.set_dependencies(params_t13_deps, params_t13_deps_count); + TaskOutputTensors task_13_outs = rt_submit_aic_task(14, params_t13); + PTO2TaskId out_proj_direct_tid_inline70 = task_13_outs.task_id(); + out_tids_inline271[N_OUT_DIRECT_inline61] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 1)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 2)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 3)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 4)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 5)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 6)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 7)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 8)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 9)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 10)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 11)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 12)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 13)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 14)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 15)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 16)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 17)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 18)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 19)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 20)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 21)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 22)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 23)] = out_proj_direct_tid_inline70; + int64_t k_base_inline111 = 0; + int64_t n_split_base_inline163 = 0; + PTO2TaskId _submit_deps_buf_inline165[10]; + for (int64_t __init_i = 0; __init_i < 10; ++__init_i) + _submit_deps_buf_inline165[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v39 = out_tids_inline271[(n_split_base_inline163 * 5)]; + _submit_deps_buf_inline165[0] = t__tmp_v39; + PTO2TaskId t__tmp_v40 = out_tids_inline271[((n_split_base_inline163 * 5) + 1)]; + _submit_deps_buf_inline165[1] = t__tmp_v40; + PTO2TaskId t__tmp_v41 = out_tids_inline271[((n_split_base_inline163 * 5) + 2)]; + _submit_deps_buf_inline165[2] = t__tmp_v41; + PTO2TaskId t__tmp_v42 = out_tids_inline271[((n_split_base_inline163 * 5) + 3)]; + _submit_deps_buf_inline165[3] = t__tmp_v42; + PTO2TaskId t__tmp_v43 = out_tids_inline271[((n_split_base_inline163 * 5) + 4)]; + _submit_deps_buf_inline165[4] = t__tmp_v43; + PTO2TaskId t__tmp_v44 = out_tids_inline271[((n_split_base_inline163 * 5) + 5)]; + _submit_deps_buf_inline165[5] = t__tmp_v44; + PTO2TaskId t__tmp_v45 = out_tids_inline271[((n_split_base_inline163 * 5) + 6)]; + _submit_deps_buf_inline165[6] = t__tmp_v45; + PTO2TaskId t__tmp_v46 = out_tids_inline271[((n_split_base_inline163 * 5) + 7)]; + _submit_deps_buf_inline165[7] = t__tmp_v46; + PTO2TaskId t__tmp_v47 = out_tids_inline271[((n_split_base_inline163 * 5) + 8)]; + _submit_deps_buf_inline165[8] = t__tmp_v47; + PTO2TaskId t__tmp_v48 = out_tids_inline271[((n_split_base_inline163 * 5) + 9)]; + _submit_deps_buf_inline165[9] = t__tmp_v48; + + // Task 14: residual_rms_cast + L0TaskArgs params_t14; + params_t14.add_inout(mlp_norm_in_inline71); + params_t14.add_inout(post_norm_partial_inline118); + params_t14.add_input(attn_proj_fp32_inline220); + params_t14.add_input(cur__rv_v7); + params_t14.add_input(ext_post_rms_weight); + params_t14.add_scalar(k_base_inline111); + params_t14.add_scalar(i); + PTO2TaskId params_t14_deps[10]; + uint32_t params_t14_deps_count = 0; + if (_submit_deps_buf_inline165[0].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[0]; + if (_submit_deps_buf_inline165[1].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[1]; + if (_submit_deps_buf_inline165[2].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[2]; + if (_submit_deps_buf_inline165[3].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[3]; + if (_submit_deps_buf_inline165[4].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[4]; + if (_submit_deps_buf_inline165[5].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[5]; + if (_submit_deps_buf_inline165[6].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[6]; + if (_submit_deps_buf_inline165[7].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[7]; + if (_submit_deps_buf_inline165[8].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[8]; + if (_submit_deps_buf_inline165[9].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[9]; + params_t14.set_dependencies(params_t14_deps, params_t14_deps_count); + params_t14.set_allow_early_resolve(true); + TaskOutputTensors task_14_outs = rt_submit_aiv_task(15, params_t14); + PTO2TaskId cast_tid_k_inline76 = task_14_outs.task_id(); + cast_tids_inline88[0] = cast_tid_k_inline76; + int64_t k_base_inline111__ssa_v1 = 1024; + int64_t n_split_base_inline163__ssa_v1 = 2; + PTO2TaskId _submit_deps_buf_inline165__ssa_v1[10]; + for (int64_t __init_i = 0; __init_i < 10; ++__init_i) + _submit_deps_buf_inline165__ssa_v1[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v51 = out_tids_inline271[(n_split_base_inline163__ssa_v1 * 5)]; + _submit_deps_buf_inline165__ssa_v1[0] = t__tmp_v51; + PTO2TaskId t__tmp_v52 = out_tids_inline271[((n_split_base_inline163__ssa_v1 * 5) + 1)]; + _submit_deps_buf_inline165__ssa_v1[1] = t__tmp_v52; + PTO2TaskId t__tmp_v53 = out_tids_inline271[((n_split_base_inline163__ssa_v1 * 5) + 2)]; + _submit_deps_buf_inline165__ssa_v1[2] = t__tmp_v53; + PTO2TaskId t__tmp_v54 = out_tids_inline271[((n_split_base_inline163__ssa_v1 * 5) + 3)]; + _submit_deps_buf_inline165__ssa_v1[3] = t__tmp_v54; + PTO2TaskId t__tmp_v55 = out_tids_inline271[((n_split_base_inline163__ssa_v1 * 5) + 4)]; + _submit_deps_buf_inline165__ssa_v1[4] = t__tmp_v55; + PTO2TaskId t__tmp_v56 = out_tids_inline271[((n_split_base_inline163__ssa_v1 * 5) + 5)]; + _submit_deps_buf_inline165__ssa_v1[5] = t__tmp_v56; + PTO2TaskId t__tmp_v57 = out_tids_inline271[((n_split_base_inline163__ssa_v1 * 5) + 6)]; + _submit_deps_buf_inline165__ssa_v1[6] = t__tmp_v57; + PTO2TaskId t__tmp_v58 = out_tids_inline271[((n_split_base_inline163__ssa_v1 * 5) + 7)]; + _submit_deps_buf_inline165__ssa_v1[7] = t__tmp_v58; + PTO2TaskId t__tmp_v59 = out_tids_inline271[((n_split_base_inline163__ssa_v1 * 5) + 8)]; + _submit_deps_buf_inline165__ssa_v1[8] = t__tmp_v59; + PTO2TaskId t__tmp_v60 = out_tids_inline271[((n_split_base_inline163__ssa_v1 * 5) + 9)]; + _submit_deps_buf_inline165__ssa_v1[9] = t__tmp_v60; + + // Task 15: residual_rms_cast_0 + L0TaskArgs params_t15; + params_t15.add_inout(mlp_norm_in_inline71); + params_t15.add_inout(post_norm_partial_inline118); + params_t15.add_input(attn_proj_fp32_inline220); + params_t15.add_input(cur__rv_v7); + params_t15.add_input(ext_post_rms_weight); + params_t15.add_scalar(k_base_inline111__ssa_v1); + params_t15.add_scalar(i); + PTO2TaskId params_t15_deps[10]; + uint32_t params_t15_deps_count = 0; + if (_submit_deps_buf_inline165__ssa_v1[0].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[0]; + if (_submit_deps_buf_inline165__ssa_v1[1].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[1]; + if (_submit_deps_buf_inline165__ssa_v1[2].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[2]; + if (_submit_deps_buf_inline165__ssa_v1[3].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[3]; + if (_submit_deps_buf_inline165__ssa_v1[4].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[4]; + if (_submit_deps_buf_inline165__ssa_v1[5].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[5]; + if (_submit_deps_buf_inline165__ssa_v1[6].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[6]; + if (_submit_deps_buf_inline165__ssa_v1[7].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[7]; + if (_submit_deps_buf_inline165__ssa_v1[8].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[8]; + if (_submit_deps_buf_inline165__ssa_v1[9].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[9]; + params_t15.set_dependencies(params_t15_deps, params_t15_deps_count); + params_t15.set_allow_early_resolve(true); + TaskOutputTensors task_15_outs = rt_submit_aiv_task(16, params_t15); + PTO2TaskId cast_tid_k_inline76__ssa_v1 = task_15_outs.task_id(); + cast_tids_inline88[1] = cast_tid_k_inline76__ssa_v1; + int64_t k_base_inline111__ssa_v2 = 2048; + int64_t n_split_base_inline163__ssa_v2 = 4; + PTO2TaskId _submit_deps_buf_inline165__ssa_v2[10]; + for (int64_t __init_i = 0; __init_i < 10; ++__init_i) + _submit_deps_buf_inline165__ssa_v2[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v63 = out_tids_inline271[(n_split_base_inline163__ssa_v2 * 5)]; + _submit_deps_buf_inline165__ssa_v2[0] = t__tmp_v63; + PTO2TaskId t__tmp_v64 = out_tids_inline271[((n_split_base_inline163__ssa_v2 * 5) + 1)]; + _submit_deps_buf_inline165__ssa_v2[1] = t__tmp_v64; + PTO2TaskId t__tmp_v65 = out_tids_inline271[((n_split_base_inline163__ssa_v2 * 5) + 2)]; + _submit_deps_buf_inline165__ssa_v2[2] = t__tmp_v65; + PTO2TaskId t__tmp_v66 = out_tids_inline271[((n_split_base_inline163__ssa_v2 * 5) + 3)]; + _submit_deps_buf_inline165__ssa_v2[3] = t__tmp_v66; + PTO2TaskId t__tmp_v67 = out_tids_inline271[((n_split_base_inline163__ssa_v2 * 5) + 4)]; + _submit_deps_buf_inline165__ssa_v2[4] = t__tmp_v67; + PTO2TaskId t__tmp_v68 = out_tids_inline271[((n_split_base_inline163__ssa_v2 * 5) + 5)]; + _submit_deps_buf_inline165__ssa_v2[5] = t__tmp_v68; + PTO2TaskId t__tmp_v69 = out_tids_inline271[((n_split_base_inline163__ssa_v2 * 5) + 6)]; + _submit_deps_buf_inline165__ssa_v2[6] = t__tmp_v69; + PTO2TaskId t__tmp_v70 = out_tids_inline271[((n_split_base_inline163__ssa_v2 * 5) + 7)]; + _submit_deps_buf_inline165__ssa_v2[7] = t__tmp_v70; + PTO2TaskId t__tmp_v71 = out_tids_inline271[((n_split_base_inline163__ssa_v2 * 5) + 8)]; + _submit_deps_buf_inline165__ssa_v2[8] = t__tmp_v71; + PTO2TaskId t__tmp_v72 = out_tids_inline271[((n_split_base_inline163__ssa_v2 * 5) + 9)]; + _submit_deps_buf_inline165__ssa_v2[9] = t__tmp_v72; + + // Task 16: residual_rms_cast_1 + L0TaskArgs params_t16; + params_t16.add_inout(mlp_norm_in_inline71); + params_t16.add_inout(post_norm_partial_inline118); + params_t16.add_input(attn_proj_fp32_inline220); + params_t16.add_input(cur__rv_v7); + params_t16.add_input(ext_post_rms_weight); + params_t16.add_scalar(k_base_inline111__ssa_v2); + params_t16.add_scalar(i); + PTO2TaskId params_t16_deps[10]; + uint32_t params_t16_deps_count = 0; + if (_submit_deps_buf_inline165__ssa_v2[0].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[0]; + if (_submit_deps_buf_inline165__ssa_v2[1].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[1]; + if (_submit_deps_buf_inline165__ssa_v2[2].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[2]; + if (_submit_deps_buf_inline165__ssa_v2[3].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[3]; + if (_submit_deps_buf_inline165__ssa_v2[4].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[4]; + if (_submit_deps_buf_inline165__ssa_v2[5].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[5]; + if (_submit_deps_buf_inline165__ssa_v2[6].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[6]; + if (_submit_deps_buf_inline165__ssa_v2[7].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[7]; + if (_submit_deps_buf_inline165__ssa_v2[8].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[8]; + if (_submit_deps_buf_inline165__ssa_v2[9].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[9]; + params_t16.set_dependencies(params_t16_deps, params_t16_deps_count); + params_t16.set_allow_early_resolve(true); + TaskOutputTensors task_16_outs = rt_submit_aiv_task(17, params_t16); + PTO2TaskId cast_tid_k_inline76__ssa_v2 = task_16_outs.task_id(); + cast_tids_inline88[2] = cast_tid_k_inline76__ssa_v2; + int64_t k_base_inline111__ssa_v3 = 3072; + int64_t n_split_base_inline163__ssa_v3 = 6; + PTO2TaskId _submit_deps_buf_inline165__ssa_v3[10]; + for (int64_t __init_i = 0; __init_i < 10; ++__init_i) + _submit_deps_buf_inline165__ssa_v3[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v75 = out_tids_inline271[(n_split_base_inline163__ssa_v3 * 5)]; + _submit_deps_buf_inline165__ssa_v3[0] = t__tmp_v75; + PTO2TaskId t__tmp_v76 = out_tids_inline271[((n_split_base_inline163__ssa_v3 * 5) + 1)]; + _submit_deps_buf_inline165__ssa_v3[1] = t__tmp_v76; + PTO2TaskId t__tmp_v77 = out_tids_inline271[((n_split_base_inline163__ssa_v3 * 5) + 2)]; + _submit_deps_buf_inline165__ssa_v3[2] = t__tmp_v77; + PTO2TaskId t__tmp_v78 = out_tids_inline271[((n_split_base_inline163__ssa_v3 * 5) + 3)]; + _submit_deps_buf_inline165__ssa_v3[3] = t__tmp_v78; + PTO2TaskId t__tmp_v79 = out_tids_inline271[((n_split_base_inline163__ssa_v3 * 5) + 4)]; + _submit_deps_buf_inline165__ssa_v3[4] = t__tmp_v79; + PTO2TaskId t__tmp_v80 = out_tids_inline271[((n_split_base_inline163__ssa_v3 * 5) + 5)]; + _submit_deps_buf_inline165__ssa_v3[5] = t__tmp_v80; + PTO2TaskId t__tmp_v81 = out_tids_inline271[((n_split_base_inline163__ssa_v3 * 5) + 6)]; + _submit_deps_buf_inline165__ssa_v3[6] = t__tmp_v81; + PTO2TaskId t__tmp_v82 = out_tids_inline271[((n_split_base_inline163__ssa_v3 * 5) + 7)]; + _submit_deps_buf_inline165__ssa_v3[7] = t__tmp_v82; + PTO2TaskId t__tmp_v83 = out_tids_inline271[((n_split_base_inline163__ssa_v3 * 5) + 8)]; + _submit_deps_buf_inline165__ssa_v3[8] = t__tmp_v83; + PTO2TaskId t__tmp_v84 = out_tids_inline271[((n_split_base_inline163__ssa_v3 * 5) + 9)]; + _submit_deps_buf_inline165__ssa_v3[9] = t__tmp_v84; + + // Task 17: residual_rms_cast_2 + L0TaskArgs params_t17; + params_t17.add_inout(mlp_norm_in_inline71); + params_t17.add_inout(post_norm_partial_inline118); + params_t17.add_input(attn_proj_fp32_inline220); + params_t17.add_input(cur__rv_v7); + params_t17.add_input(ext_post_rms_weight); + params_t17.add_scalar(k_base_inline111__ssa_v3); + params_t17.add_scalar(i); + PTO2TaskId params_t17_deps[10]; + uint32_t params_t17_deps_count = 0; + if (_submit_deps_buf_inline165__ssa_v3[0].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[0]; + if (_submit_deps_buf_inline165__ssa_v3[1].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[1]; + if (_submit_deps_buf_inline165__ssa_v3[2].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[2]; + if (_submit_deps_buf_inline165__ssa_v3[3].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[3]; + if (_submit_deps_buf_inline165__ssa_v3[4].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[4]; + if (_submit_deps_buf_inline165__ssa_v3[5].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[5]; + if (_submit_deps_buf_inline165__ssa_v3[6].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[6]; + if (_submit_deps_buf_inline165__ssa_v3[7].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[7]; + if (_submit_deps_buf_inline165__ssa_v3[8].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[8]; + if (_submit_deps_buf_inline165__ssa_v3[9].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[9]; + params_t17.set_dependencies(params_t17_deps, params_t17_deps_count); + params_t17.set_allow_early_resolve(true); + TaskOutputTensors task_17_outs = rt_submit_aiv_task(18, params_t17); + PTO2TaskId cast_tid_k_inline76__ssa_v3 = task_17_outs.task_id(); + cast_tids_inline88[3] = cast_tid_k_inline76__ssa_v3; + int64_t k_base_inline111__ssa_v4 = 4096; + int64_t n_split_base_inline163__ssa_v4 = 8; + PTO2TaskId _submit_deps_buf_inline165__ssa_v4[10]; + for (int64_t __init_i = 0; __init_i < 10; ++__init_i) + _submit_deps_buf_inline165__ssa_v4[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v87 = out_tids_inline271[(n_split_base_inline163__ssa_v4 * 5)]; + _submit_deps_buf_inline165__ssa_v4[0] = t__tmp_v87; + PTO2TaskId t__tmp_v88 = out_tids_inline271[((n_split_base_inline163__ssa_v4 * 5) + 1)]; + _submit_deps_buf_inline165__ssa_v4[1] = t__tmp_v88; + PTO2TaskId t__tmp_v89 = out_tids_inline271[((n_split_base_inline163__ssa_v4 * 5) + 2)]; + _submit_deps_buf_inline165__ssa_v4[2] = t__tmp_v89; + PTO2TaskId t__tmp_v90 = out_tids_inline271[((n_split_base_inline163__ssa_v4 * 5) + 3)]; + _submit_deps_buf_inline165__ssa_v4[3] = t__tmp_v90; + PTO2TaskId t__tmp_v91 = out_tids_inline271[((n_split_base_inline163__ssa_v4 * 5) + 4)]; + _submit_deps_buf_inline165__ssa_v4[4] = t__tmp_v91; + PTO2TaskId t__tmp_v92 = out_tids_inline271[((n_split_base_inline163__ssa_v4 * 5) + 5)]; + _submit_deps_buf_inline165__ssa_v4[5] = t__tmp_v92; + PTO2TaskId t__tmp_v93 = out_tids_inline271[((n_split_base_inline163__ssa_v4 * 5) + 6)]; + _submit_deps_buf_inline165__ssa_v4[6] = t__tmp_v93; + PTO2TaskId t__tmp_v94 = out_tids_inline271[((n_split_base_inline163__ssa_v4 * 5) + 7)]; + _submit_deps_buf_inline165__ssa_v4[7] = t__tmp_v94; + PTO2TaskId t__tmp_v95 = out_tids_inline271[((n_split_base_inline163__ssa_v4 * 5) + 8)]; + _submit_deps_buf_inline165__ssa_v4[8] = t__tmp_v95; + PTO2TaskId t__tmp_v96 = out_tids_inline271[((n_split_base_inline163__ssa_v4 * 5) + 9)]; + _submit_deps_buf_inline165__ssa_v4[9] = t__tmp_v96; + + // Task 18: residual_rms_cast_3 + L0TaskArgs params_t18; + params_t18.add_inout(mlp_norm_in_inline71); + params_t18.add_inout(post_norm_partial_inline118); + params_t18.add_input(attn_proj_fp32_inline220); + params_t18.add_input(cur__rv_v7); + params_t18.add_input(ext_post_rms_weight); + params_t18.add_scalar(k_base_inline111__ssa_v4); + params_t18.add_scalar(i); + PTO2TaskId params_t18_deps[10]; + uint32_t params_t18_deps_count = 0; + if (_submit_deps_buf_inline165__ssa_v4[0].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[0]; + if (_submit_deps_buf_inline165__ssa_v4[1].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[1]; + if (_submit_deps_buf_inline165__ssa_v4[2].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[2]; + if (_submit_deps_buf_inline165__ssa_v4[3].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[3]; + if (_submit_deps_buf_inline165__ssa_v4[4].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[4]; + if (_submit_deps_buf_inline165__ssa_v4[5].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[5]; + if (_submit_deps_buf_inline165__ssa_v4[6].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[6]; + if (_submit_deps_buf_inline165__ssa_v4[7].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[7]; + if (_submit_deps_buf_inline165__ssa_v4[8].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[8]; + if (_submit_deps_buf_inline165__ssa_v4[9].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[9]; + params_t18.set_dependencies(params_t18_deps, params_t18_deps_count); + params_t18.set_allow_early_resolve(true); + TaskOutputTensors task_18_outs = rt_submit_aiv_task(19, params_t18); + PTO2TaskId cast_tid_k_inline76__ssa_v4 = task_18_outs.task_id(); + cast_tids_inline88[4] = cast_tid_k_inline76__ssa_v4; + + // Task 19: post_rms_reduce + L0TaskArgs params_t19; + params_t19.add_input(attn_proj_fp32_inline220); + params_t19.add_input(cur__rv_v7); + params_t19.add_inout(inv_rms_tile_inline126); + PTO2TaskId params_t19_deps[50]; + uint32_t params_t19_deps_count = 0; + if (out_tids_inline271[0].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[0]; + if (out_tids_inline271[1].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[1]; + if (out_tids_inline271[2].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[2]; + if (out_tids_inline271[3].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[3]; + if (out_tids_inline271[4].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[4]; + if (out_tids_inline271[5].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[5]; + if (out_tids_inline271[6].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[6]; + if (out_tids_inline271[7].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[7]; + if (out_tids_inline271[8].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[8]; + if (out_tids_inline271[9].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[9]; + if (out_tids_inline271[10].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[10]; + if (out_tids_inline271[11].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[11]; + if (out_tids_inline271[12].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[12]; + if (out_tids_inline271[13].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[13]; + if (out_tids_inline271[14].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[14]; + if (out_tids_inline271[15].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[15]; + if (out_tids_inline271[16].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[16]; + if (out_tids_inline271[17].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[17]; + if (out_tids_inline271[18].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[18]; + if (out_tids_inline271[19].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[19]; + if (out_tids_inline271[20].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[20]; + if (out_tids_inline271[21].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[21]; + if (out_tids_inline271[22].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[22]; + if (out_tids_inline271[23].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[23]; + if (out_tids_inline271[24].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[24]; + if (out_tids_inline271[25].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[25]; + if (out_tids_inline271[26].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[26]; + if (out_tids_inline271[27].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[27]; + if (out_tids_inline271[28].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[28]; + if (out_tids_inline271[29].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[29]; + if (out_tids_inline271[30].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[30]; + if (out_tids_inline271[31].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[31]; + if (out_tids_inline271[32].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[32]; + if (out_tids_inline271[33].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[33]; + if (out_tids_inline271[34].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[34]; + if (out_tids_inline271[35].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[35]; + if (out_tids_inline271[36].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[36]; + if (out_tids_inline271[37].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[37]; + if (out_tids_inline271[38].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[38]; + if (out_tids_inline271[39].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[39]; + if (out_tids_inline271[40].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[40]; + if (out_tids_inline271[41].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[41]; + if (out_tids_inline271[42].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[42]; + if (out_tids_inline271[43].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[43]; + if (out_tids_inline271[44].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[44]; + if (out_tids_inline271[45].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[45]; + if (out_tids_inline271[46].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[46]; + if (out_tids_inline271[47].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[47]; + if (out_tids_inline271[48].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[48]; + if (out_tids_inline271[49].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[49]; + params_t19.set_dependencies(params_t19_deps, params_t19_deps_count); + TaskOutputTensors task_19_outs = rt_submit_aiv_task(20, params_t19); + PTO2TaskId reduce_tid_inline226 = task_19_outs.task_id(); + int64_t gu_k0_inline131 = 0; + PTO2TaskId _submit_deps_buf_inline236[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline236[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v105 = cast_tids_inline88[0]; + _submit_deps_buf_inline236[0] = t__tmp_v105; + + // Phase-fence barrier 3: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_3; + PTO2TaskId params_phase_fence_barrier_3_deps[1]; + uint32_t params_phase_fence_barrier_3_deps_count = 0; + if (_submit_deps_buf_inline236[0].is_valid()) + params_phase_fence_barrier_3_deps[params_phase_fence_barrier_3_deps_count++] = + _submit_deps_buf_inline236[0]; + params_phase_fence_barrier_3.set_dependencies( + params_phase_fence_barrier_3_deps, params_phase_fence_barrier_3_deps_count + ); + PTO2TaskId t__tmp_v106 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_3_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_3_outs = + rt_submit_dummy_task(params_phase_fence_barrier_3); + t__tmp_v106 = phase_fence_barrier_3_outs.task_id(); + } + gate_late_tids_inline249[0] = t__tmp_v106; + PTO2TaskId _submit_deps_buf_inline225[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline225[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v107 = cast_tids_inline88[0]; + _submit_deps_buf_inline225[0] = t__tmp_v107; + + // Phase-fence barrier 4: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_4; + PTO2TaskId params_phase_fence_barrier_4_deps[1]; + uint32_t params_phase_fence_barrier_4_deps_count = 0; + if (_submit_deps_buf_inline225[0].is_valid()) + params_phase_fence_barrier_4_deps[params_phase_fence_barrier_4_deps_count++] = + _submit_deps_buf_inline225[0]; + params_phase_fence_barrier_4.set_dependencies( + params_phase_fence_barrier_4_deps, params_phase_fence_barrier_4_deps_count + ); + PTO2TaskId t__tmp_v108 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_4_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_4_outs = + rt_submit_dummy_task(params_phase_fence_barrier_4); + t__tmp_v108 = phase_fence_barrier_4_outs.task_id(); + } + up_late_tids_inline69[0] = t__tmp_v108; + PTO2TaskId _submit_deps_buf_inline237[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline237[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v109 = cast_tids_inline88[0]; + _submit_deps_buf_inline237[0] = t__tmp_v109; + + // Spmd gate_proj_spmd: gate_proj + L0TaskArgs params_t20; + params_t20.add_input(mlp_norm_in_inline71); + params_t20.add_input(ext_w_gate); + params_t20.add_inout(gate_acc_all_inline203); + params_t20.add_scalar(gu_k0_inline131); + params_t20.add_scalar(layer_hidden_base_inline151); + params_t20.launch_spec.set_block_num(6); + PTO2TaskId params_t20_deps[1]; + uint32_t params_t20_deps_count = 0; + if (_submit_deps_buf_inline237[0].is_valid()) + params_t20_deps[params_t20_deps_count++] = _submit_deps_buf_inline237[0]; + params_t20.set_dependencies(params_t20_deps, params_t20_deps_count); + TaskOutputTensors task_20_outs = rt_submit_aic_task(21, params_t20); + PTO2TaskId gate_spmd_tid_inline245 = task_20_outs.task_id(); + gate_tids_inline56[0] = gate_spmd_tid_inline245; + gate_tids_inline56[5] = gate_spmd_tid_inline245; + gate_tids_inline56[10] = gate_spmd_tid_inline245; + gate_tids_inline56[15] = gate_spmd_tid_inline245; + gate_tids_inline56[20] = gate_spmd_tid_inline245; + gate_tids_inline56[25] = gate_spmd_tid_inline245; + PTO2TaskId _submit_deps_buf_inline260[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline260[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v110 = cast_tids_inline88[0]; + _submit_deps_buf_inline260[0] = t__tmp_v110; + + // Spmd up_proj_spmd: up_proj + L0TaskArgs params_t21; + params_t21.add_input(mlp_norm_in_inline71); + params_t21.add_input(ext_w_up); + params_t21.add_inout(up_acc_all_inline303); + params_t21.add_scalar(gu_k0_inline131); + params_t21.add_scalar(layer_hidden_base_inline151); + params_t21.launch_spec.set_block_num(6); + PTO2TaskId params_t21_deps[1]; + uint32_t params_t21_deps_count = 0; + if (_submit_deps_buf_inline260[0].is_valid()) + params_t21_deps[params_t21_deps_count++] = _submit_deps_buf_inline260[0]; + params_t21.set_dependencies(params_t21_deps, params_t21_deps_count); + TaskOutputTensors task_21_outs = rt_submit_aic_task(22, params_t21); + PTO2TaskId up_spmd_tid_inline264 = task_21_outs.task_id(); + up_tids_inline310[0] = up_spmd_tid_inline264; + up_tids_inline310[5] = up_spmd_tid_inline264; + up_tids_inline310[10] = up_spmd_tid_inline264; + up_tids_inline310[15] = up_spmd_tid_inline264; + up_tids_inline310[20] = up_spmd_tid_inline264; + up_tids_inline310[25] = up_spmd_tid_inline264; + int64_t gu_k0_inline131__ssa_v1 = 1024; + PTO2TaskId _submit_deps_buf_inline236__ssa_v1[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline236__ssa_v1[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v111 = cast_tids_inline88[1]; + _submit_deps_buf_inline236__ssa_v1[0] = t__tmp_v111; + + // Phase-fence barrier 5: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_5; + PTO2TaskId params_phase_fence_barrier_5_deps[1]; + uint32_t params_phase_fence_barrier_5_deps_count = 0; + if (_submit_deps_buf_inline236__ssa_v1[0].is_valid()) + params_phase_fence_barrier_5_deps[params_phase_fence_barrier_5_deps_count++] = + _submit_deps_buf_inline236__ssa_v1[0]; + params_phase_fence_barrier_5.set_dependencies( + params_phase_fence_barrier_5_deps, params_phase_fence_barrier_5_deps_count + ); + PTO2TaskId t__tmp_v112 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_5_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_5_outs = + rt_submit_dummy_task(params_phase_fence_barrier_5); + t__tmp_v112 = phase_fence_barrier_5_outs.task_id(); + } + gate_late_tids_inline249[1] = t__tmp_v112; + PTO2TaskId _submit_deps_buf_inline225__ssa_v1[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline225__ssa_v1[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v113 = cast_tids_inline88[1]; + _submit_deps_buf_inline225__ssa_v1[0] = t__tmp_v113; + + // Phase-fence barrier 6: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_6; + PTO2TaskId params_phase_fence_barrier_6_deps[1]; + uint32_t params_phase_fence_barrier_6_deps_count = 0; + if (_submit_deps_buf_inline225__ssa_v1[0].is_valid()) + params_phase_fence_barrier_6_deps[params_phase_fence_barrier_6_deps_count++] = + _submit_deps_buf_inline225__ssa_v1[0]; + params_phase_fence_barrier_6.set_dependencies( + params_phase_fence_barrier_6_deps, params_phase_fence_barrier_6_deps_count + ); + PTO2TaskId t__tmp_v114 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_6_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_6_outs = + rt_submit_dummy_task(params_phase_fence_barrier_6); + t__tmp_v114 = phase_fence_barrier_6_outs.task_id(); + } + up_late_tids_inline69[1] = t__tmp_v114; + PTO2TaskId _submit_deps_buf_inline237__ssa_v1[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline237__ssa_v1[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v115 = cast_tids_inline88[1]; + _submit_deps_buf_inline237__ssa_v1[0] = t__tmp_v115; + + // Spmd gate_proj_spmd_0: gate_proj_0 + L0TaskArgs params_t22; + params_t22.add_input(mlp_norm_in_inline71); + params_t22.add_input(ext_w_gate); + params_t22.add_inout(gate_acc_all_inline203); + params_t22.add_scalar(gu_k0_inline131__ssa_v1); + params_t22.add_scalar(layer_hidden_base_inline151); + params_t22.launch_spec.set_block_num(6); + PTO2TaskId params_t22_deps[1]; + uint32_t params_t22_deps_count = 0; + if (_submit_deps_buf_inline237__ssa_v1[0].is_valid()) + params_t22_deps[params_t22_deps_count++] = _submit_deps_buf_inline237__ssa_v1[0]; + params_t22.set_dependencies(params_t22_deps, params_t22_deps_count); + TaskOutputTensors task_22_outs = rt_submit_aic_task(23, params_t22); + PTO2TaskId gate_spmd_tid_inline245__ssa_v1 = task_22_outs.task_id(); + gate_tids_inline56[1] = gate_spmd_tid_inline245__ssa_v1; + gate_tids_inline56[6] = gate_spmd_tid_inline245__ssa_v1; + gate_tids_inline56[11] = gate_spmd_tid_inline245__ssa_v1; + gate_tids_inline56[16] = gate_spmd_tid_inline245__ssa_v1; + gate_tids_inline56[21] = gate_spmd_tid_inline245__ssa_v1; + gate_tids_inline56[26] = gate_spmd_tid_inline245__ssa_v1; + PTO2TaskId _submit_deps_buf_inline260__ssa_v1[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline260__ssa_v1[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v116 = cast_tids_inline88[1]; + _submit_deps_buf_inline260__ssa_v1[0] = t__tmp_v116; + + // Spmd up_proj_spmd_0: up_proj_0 + L0TaskArgs params_t23; + params_t23.add_input(mlp_norm_in_inline71); + params_t23.add_input(ext_w_up); + params_t23.add_inout(up_acc_all_inline303); + params_t23.add_scalar(gu_k0_inline131__ssa_v1); + params_t23.add_scalar(layer_hidden_base_inline151); + params_t23.launch_spec.set_block_num(6); + PTO2TaskId params_t23_deps[1]; + uint32_t params_t23_deps_count = 0; + if (_submit_deps_buf_inline260__ssa_v1[0].is_valid()) + params_t23_deps[params_t23_deps_count++] = _submit_deps_buf_inline260__ssa_v1[0]; + params_t23.set_dependencies(params_t23_deps, params_t23_deps_count); + TaskOutputTensors task_23_outs = rt_submit_aic_task(24, params_t23); + PTO2TaskId up_spmd_tid_inline264__ssa_v1 = task_23_outs.task_id(); + up_tids_inline310[1] = up_spmd_tid_inline264__ssa_v1; + up_tids_inline310[6] = up_spmd_tid_inline264__ssa_v1; + up_tids_inline310[11] = up_spmd_tid_inline264__ssa_v1; + up_tids_inline310[16] = up_spmd_tid_inline264__ssa_v1; + up_tids_inline310[21] = up_spmd_tid_inline264__ssa_v1; + up_tids_inline310[26] = up_spmd_tid_inline264__ssa_v1; + int64_t gu_k0_inline131__ssa_v2 = 2048; + PTO2TaskId _submit_deps_buf_inline236__ssa_v2[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline236__ssa_v2[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v117 = cast_tids_inline88[2]; + _submit_deps_buf_inline236__ssa_v2[0] = t__tmp_v117; + + // Phase-fence barrier 7: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_7; + PTO2TaskId params_phase_fence_barrier_7_deps[1]; + uint32_t params_phase_fence_barrier_7_deps_count = 0; + if (_submit_deps_buf_inline236__ssa_v2[0].is_valid()) + params_phase_fence_barrier_7_deps[params_phase_fence_barrier_7_deps_count++] = + _submit_deps_buf_inline236__ssa_v2[0]; + params_phase_fence_barrier_7.set_dependencies( + params_phase_fence_barrier_7_deps, params_phase_fence_barrier_7_deps_count + ); + PTO2TaskId t__tmp_v118 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_7_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_7_outs = + rt_submit_dummy_task(params_phase_fence_barrier_7); + t__tmp_v118 = phase_fence_barrier_7_outs.task_id(); + } + gate_late_tids_inline249[2] = t__tmp_v118; + PTO2TaskId _submit_deps_buf_inline225__ssa_v2[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline225__ssa_v2[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v119 = cast_tids_inline88[2]; + _submit_deps_buf_inline225__ssa_v2[0] = t__tmp_v119; + + // Phase-fence barrier 8: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_8; + PTO2TaskId params_phase_fence_barrier_8_deps[1]; + uint32_t params_phase_fence_barrier_8_deps_count = 0; + if (_submit_deps_buf_inline225__ssa_v2[0].is_valid()) + params_phase_fence_barrier_8_deps[params_phase_fence_barrier_8_deps_count++] = + _submit_deps_buf_inline225__ssa_v2[0]; + params_phase_fence_barrier_8.set_dependencies( + params_phase_fence_barrier_8_deps, params_phase_fence_barrier_8_deps_count + ); + PTO2TaskId t__tmp_v120 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_8_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_8_outs = + rt_submit_dummy_task(params_phase_fence_barrier_8); + t__tmp_v120 = phase_fence_barrier_8_outs.task_id(); + } + up_late_tids_inline69[2] = t__tmp_v120; + PTO2TaskId _submit_deps_buf_inline237__ssa_v2[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline237__ssa_v2[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v121 = cast_tids_inline88[2]; + _submit_deps_buf_inline237__ssa_v2[0] = t__tmp_v121; + + // Spmd gate_proj_spmd_1: gate_proj_1 + L0TaskArgs params_t24; + params_t24.add_input(mlp_norm_in_inline71); + params_t24.add_input(ext_w_gate); + params_t24.add_inout(gate_acc_all_inline203); + params_t24.add_scalar(gu_k0_inline131__ssa_v2); + params_t24.add_scalar(layer_hidden_base_inline151); + params_t24.launch_spec.set_block_num(6); + PTO2TaskId params_t24_deps[1]; + uint32_t params_t24_deps_count = 0; + if (_submit_deps_buf_inline237__ssa_v2[0].is_valid()) + params_t24_deps[params_t24_deps_count++] = _submit_deps_buf_inline237__ssa_v2[0]; + params_t24.set_dependencies(params_t24_deps, params_t24_deps_count); + TaskOutputTensors task_24_outs = rt_submit_aic_task(25, params_t24); + PTO2TaskId gate_spmd_tid_inline245__ssa_v2 = task_24_outs.task_id(); + gate_tids_inline56[2] = gate_spmd_tid_inline245__ssa_v2; + gate_tids_inline56[7] = gate_spmd_tid_inline245__ssa_v2; + gate_tids_inline56[12] = gate_spmd_tid_inline245__ssa_v2; + gate_tids_inline56[17] = gate_spmd_tid_inline245__ssa_v2; + gate_tids_inline56[22] = gate_spmd_tid_inline245__ssa_v2; + gate_tids_inline56[27] = gate_spmd_tid_inline245__ssa_v2; + PTO2TaskId _submit_deps_buf_inline260__ssa_v2[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline260__ssa_v2[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v122 = cast_tids_inline88[2]; + _submit_deps_buf_inline260__ssa_v2[0] = t__tmp_v122; + + // Spmd up_proj_spmd_1: up_proj_1 + L0TaskArgs params_t25; + params_t25.add_input(mlp_norm_in_inline71); + params_t25.add_input(ext_w_up); + params_t25.add_inout(up_acc_all_inline303); + params_t25.add_scalar(gu_k0_inline131__ssa_v2); + params_t25.add_scalar(layer_hidden_base_inline151); + params_t25.launch_spec.set_block_num(6); + PTO2TaskId params_t25_deps[1]; + uint32_t params_t25_deps_count = 0; + if (_submit_deps_buf_inline260__ssa_v2[0].is_valid()) + params_t25_deps[params_t25_deps_count++] = _submit_deps_buf_inline260__ssa_v2[0]; + params_t25.set_dependencies(params_t25_deps, params_t25_deps_count); + TaskOutputTensors task_25_outs = rt_submit_aic_task(26, params_t25); + PTO2TaskId up_spmd_tid_inline264__ssa_v2 = task_25_outs.task_id(); + up_tids_inline310[2] = up_spmd_tid_inline264__ssa_v2; + up_tids_inline310[7] = up_spmd_tid_inline264__ssa_v2; + up_tids_inline310[12] = up_spmd_tid_inline264__ssa_v2; + up_tids_inline310[17] = up_spmd_tid_inline264__ssa_v2; + up_tids_inline310[22] = up_spmd_tid_inline264__ssa_v2; + up_tids_inline310[27] = up_spmd_tid_inline264__ssa_v2; + int64_t gu_k0_inline131__ssa_v3 = 3072; + PTO2TaskId _submit_deps_buf_inline236__ssa_v3[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline236__ssa_v3[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v123 = cast_tids_inline88[3]; + _submit_deps_buf_inline236__ssa_v3[0] = t__tmp_v123; + + // Phase-fence barrier 9: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_9; + PTO2TaskId params_phase_fence_barrier_9_deps[1]; + uint32_t params_phase_fence_barrier_9_deps_count = 0; + if (_submit_deps_buf_inline236__ssa_v3[0].is_valid()) + params_phase_fence_barrier_9_deps[params_phase_fence_barrier_9_deps_count++] = + _submit_deps_buf_inline236__ssa_v3[0]; + params_phase_fence_barrier_9.set_dependencies( + params_phase_fence_barrier_9_deps, params_phase_fence_barrier_9_deps_count + ); + PTO2TaskId t__tmp_v124 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_9_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_9_outs = + rt_submit_dummy_task(params_phase_fence_barrier_9); + t__tmp_v124 = phase_fence_barrier_9_outs.task_id(); + } + gate_late_tids_inline249[3] = t__tmp_v124; + PTO2TaskId _submit_deps_buf_inline225__ssa_v3[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline225__ssa_v3[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v125 = cast_tids_inline88[3]; + _submit_deps_buf_inline225__ssa_v3[0] = t__tmp_v125; + + // Phase-fence barrier 10: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_10; + PTO2TaskId params_phase_fence_barrier_10_deps[1]; + uint32_t params_phase_fence_barrier_10_deps_count = 0; + if (_submit_deps_buf_inline225__ssa_v3[0].is_valid()) + params_phase_fence_barrier_10_deps[params_phase_fence_barrier_10_deps_count++] = + _submit_deps_buf_inline225__ssa_v3[0]; + params_phase_fence_barrier_10.set_dependencies( + params_phase_fence_barrier_10_deps, params_phase_fence_barrier_10_deps_count + ); + PTO2TaskId t__tmp_v126 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_10_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_10_outs = + rt_submit_dummy_task(params_phase_fence_barrier_10); + t__tmp_v126 = phase_fence_barrier_10_outs.task_id(); + } + up_late_tids_inline69[3] = t__tmp_v126; + PTO2TaskId _submit_deps_buf_inline237__ssa_v3[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline237__ssa_v3[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v127 = cast_tids_inline88[3]; + _submit_deps_buf_inline237__ssa_v3[0] = t__tmp_v127; + + // Spmd gate_proj_spmd_2: gate_proj_2 + L0TaskArgs params_t26; + params_t26.add_input(mlp_norm_in_inline71); + params_t26.add_input(ext_w_gate); + params_t26.add_inout(gate_acc_all_inline203); + params_t26.add_scalar(gu_k0_inline131__ssa_v3); + params_t26.add_scalar(layer_hidden_base_inline151); + params_t26.launch_spec.set_block_num(6); + PTO2TaskId params_t26_deps[1]; + uint32_t params_t26_deps_count = 0; + if (_submit_deps_buf_inline237__ssa_v3[0].is_valid()) + params_t26_deps[params_t26_deps_count++] = _submit_deps_buf_inline237__ssa_v3[0]; + params_t26.set_dependencies(params_t26_deps, params_t26_deps_count); + TaskOutputTensors task_26_outs = rt_submit_aic_task(27, params_t26); + PTO2TaskId gate_spmd_tid_inline245__ssa_v3 = task_26_outs.task_id(); + gate_tids_inline56[3] = gate_spmd_tid_inline245__ssa_v3; + gate_tids_inline56[8] = gate_spmd_tid_inline245__ssa_v3; + gate_tids_inline56[13] = gate_spmd_tid_inline245__ssa_v3; + gate_tids_inline56[18] = gate_spmd_tid_inline245__ssa_v3; + gate_tids_inline56[23] = gate_spmd_tid_inline245__ssa_v3; + gate_tids_inline56[28] = gate_spmd_tid_inline245__ssa_v3; + PTO2TaskId _submit_deps_buf_inline260__ssa_v3[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline260__ssa_v3[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v128 = cast_tids_inline88[3]; + _submit_deps_buf_inline260__ssa_v3[0] = t__tmp_v128; + + // Spmd up_proj_spmd_2: up_proj_2 + L0TaskArgs params_t27; + params_t27.add_input(mlp_norm_in_inline71); + params_t27.add_input(ext_w_up); + params_t27.add_inout(up_acc_all_inline303); + params_t27.add_scalar(gu_k0_inline131__ssa_v3); + params_t27.add_scalar(layer_hidden_base_inline151); + params_t27.launch_spec.set_block_num(6); + PTO2TaskId params_t27_deps[1]; + uint32_t params_t27_deps_count = 0; + if (_submit_deps_buf_inline260__ssa_v3[0].is_valid()) + params_t27_deps[params_t27_deps_count++] = _submit_deps_buf_inline260__ssa_v3[0]; + params_t27.set_dependencies(params_t27_deps, params_t27_deps_count); + TaskOutputTensors task_27_outs = rt_submit_aic_task(28, params_t27); + PTO2TaskId up_spmd_tid_inline264__ssa_v3 = task_27_outs.task_id(); + up_tids_inline310[3] = up_spmd_tid_inline264__ssa_v3; + up_tids_inline310[8] = up_spmd_tid_inline264__ssa_v3; + up_tids_inline310[13] = up_spmd_tid_inline264__ssa_v3; + up_tids_inline310[18] = up_spmd_tid_inline264__ssa_v3; + up_tids_inline310[23] = up_spmd_tid_inline264__ssa_v3; + up_tids_inline310[28] = up_spmd_tid_inline264__ssa_v3; + int64_t gu_k0_inline131__ssa_v4 = 4096; + PTO2TaskId _submit_deps_buf_inline236__ssa_v4[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline236__ssa_v4[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v129 = cast_tids_inline88[4]; + _submit_deps_buf_inline236__ssa_v4[0] = t__tmp_v129; + + // Phase-fence barrier 11: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_11; + PTO2TaskId params_phase_fence_barrier_11_deps[1]; + uint32_t params_phase_fence_barrier_11_deps_count = 0; + if (_submit_deps_buf_inline236__ssa_v4[0].is_valid()) + params_phase_fence_barrier_11_deps[params_phase_fence_barrier_11_deps_count++] = + _submit_deps_buf_inline236__ssa_v4[0]; + params_phase_fence_barrier_11.set_dependencies( + params_phase_fence_barrier_11_deps, params_phase_fence_barrier_11_deps_count + ); + PTO2TaskId t__tmp_v130 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_11_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_11_outs = + rt_submit_dummy_task(params_phase_fence_barrier_11); + t__tmp_v130 = phase_fence_barrier_11_outs.task_id(); + } + gate_late_tids_inline249[4] = t__tmp_v130; + PTO2TaskId _submit_deps_buf_inline225__ssa_v4[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline225__ssa_v4[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v131 = cast_tids_inline88[4]; + _submit_deps_buf_inline225__ssa_v4[0] = t__tmp_v131; + + // Phase-fence barrier 12: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_12; + PTO2TaskId params_phase_fence_barrier_12_deps[1]; + uint32_t params_phase_fence_barrier_12_deps_count = 0; + if (_submit_deps_buf_inline225__ssa_v4[0].is_valid()) + params_phase_fence_barrier_12_deps[params_phase_fence_barrier_12_deps_count++] = + _submit_deps_buf_inline225__ssa_v4[0]; + params_phase_fence_barrier_12.set_dependencies( + params_phase_fence_barrier_12_deps, params_phase_fence_barrier_12_deps_count + ); + PTO2TaskId t__tmp_v132 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_12_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_12_outs = + rt_submit_dummy_task(params_phase_fence_barrier_12); + t__tmp_v132 = phase_fence_barrier_12_outs.task_id(); + } + up_late_tids_inline69[4] = t__tmp_v132; + PTO2TaskId _submit_deps_buf_inline237__ssa_v4[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline237__ssa_v4[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v133 = cast_tids_inline88[4]; + _submit_deps_buf_inline237__ssa_v4[0] = t__tmp_v133; + + // Spmd gate_proj_spmd_3: gate_proj_3 + L0TaskArgs params_t28; + params_t28.add_input(mlp_norm_in_inline71); + params_t28.add_input(ext_w_gate); + params_t28.add_inout(gate_acc_all_inline203); + params_t28.add_scalar(gu_k0_inline131__ssa_v4); + params_t28.add_scalar(layer_hidden_base_inline151); + params_t28.launch_spec.set_block_num(6); + PTO2TaskId params_t28_deps[1]; + uint32_t params_t28_deps_count = 0; + if (_submit_deps_buf_inline237__ssa_v4[0].is_valid()) + params_t28_deps[params_t28_deps_count++] = _submit_deps_buf_inline237__ssa_v4[0]; + params_t28.set_dependencies(params_t28_deps, params_t28_deps_count); + TaskOutputTensors task_28_outs = rt_submit_aic_task(29, params_t28); + PTO2TaskId gate_spmd_tid_inline245__ssa_v4 = task_28_outs.task_id(); + gate_tids_inline56[4] = gate_spmd_tid_inline245__ssa_v4; + gate_tids_inline56[9] = gate_spmd_tid_inline245__ssa_v4; + gate_tids_inline56[14] = gate_spmd_tid_inline245__ssa_v4; + gate_tids_inline56[19] = gate_spmd_tid_inline245__ssa_v4; + gate_tids_inline56[24] = gate_spmd_tid_inline245__ssa_v4; + gate_tids_inline56[29] = gate_spmd_tid_inline245__ssa_v4; + PTO2TaskId _submit_deps_buf_inline260__ssa_v4[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline260__ssa_v4[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v134 = cast_tids_inline88[4]; + _submit_deps_buf_inline260__ssa_v4[0] = t__tmp_v134; + + // Spmd up_proj_spmd_3: up_proj_3 + L0TaskArgs params_t29; + params_t29.add_input(mlp_norm_in_inline71); + params_t29.add_input(ext_w_up); + params_t29.add_inout(up_acc_all_inline303); + params_t29.add_scalar(gu_k0_inline131__ssa_v4); + params_t29.add_scalar(layer_hidden_base_inline151); + params_t29.launch_spec.set_block_num(6); + PTO2TaskId params_t29_deps[1]; + uint32_t params_t29_deps_count = 0; + if (_submit_deps_buf_inline260__ssa_v4[0].is_valid()) + params_t29_deps[params_t29_deps_count++] = _submit_deps_buf_inline260__ssa_v4[0]; + params_t29.set_dependencies(params_t29_deps, params_t29_deps_count); + TaskOutputTensors task_29_outs = rt_submit_aic_task(30, params_t29); + PTO2TaskId up_spmd_tid_inline264__ssa_v4 = task_29_outs.task_id(); + up_tids_inline310[4] = up_spmd_tid_inline264__ssa_v4; + up_tids_inline310[9] = up_spmd_tid_inline264__ssa_v4; + up_tids_inline310[14] = up_spmd_tid_inline264__ssa_v4; + up_tids_inline310[19] = up_spmd_tid_inline264__ssa_v4; + up_tids_inline310[24] = up_spmd_tid_inline264__ssa_v4; + up_tids_inline310[29] = up_spmd_tid_inline264__ssa_v4; + for (int64_t n_out_inline275 = 6; n_out_inline275 < 17; n_out_inline275 += 1) { + int64_t n0_inline122 = (n_out_inline275 * 1024); + for (int64_t k_split_inline276 = 0; k_split_inline276 < 5; k_split_inline276 += 1) { + int64_t k0_inline113 = (k_split_inline276 * 1024); + PTO2TaskId _submit_deps_buf_inline102[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline102[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v135 = gate_late_tids_inline249[k_split_inline276]; + _submit_deps_buf_inline102[0] = t__tmp_v135; + + // Task 30: gate_proj_4 + L0TaskArgs params_t30; + params_t30.add_input(mlp_norm_in_inline71); + params_t30.add_input(ext_w_gate); + params_t30.add_inout(gate_acc_all_inline203); + params_t30.add_scalar(k0_inline113); + params_t30.add_scalar(layer_hidden_base_inline151); + params_t30.add_scalar(n0_inline122); + PTO2TaskId params_t30_deps[1]; + uint32_t params_t30_deps_count = 0; + if (_submit_deps_buf_inline102[0].is_valid()) + params_t30_deps[params_t30_deps_count++] = _submit_deps_buf_inline102[0]; + params_t30.set_dependencies(params_t30_deps, params_t30_deps_count); + TaskOutputTensors task_30_outs = rt_submit_aic_task(31, params_t30); + PTO2TaskId gate_tid_inline277 = task_30_outs.task_id(); + gate_tids_inline56[((n_out_inline275 * 5) + k_split_inline276)] = gate_tid_inline277; + PTO2TaskId _submit_deps_buf_inline246[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline246[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v136 = up_late_tids_inline69[k_split_inline276]; + _submit_deps_buf_inline246[0] = t__tmp_v136; + + // Task 31: up_proj_4 + L0TaskArgs params_t31; + params_t31.add_input(mlp_norm_in_inline71); + params_t31.add_input(ext_w_up); + params_t31.add_inout(up_acc_all_inline303); + params_t31.add_scalar(k0_inline113); + params_t31.add_scalar(layer_hidden_base_inline151); + params_t31.add_scalar(n0_inline122); + PTO2TaskId params_t31_deps[1]; + uint32_t params_t31_deps_count = 0; + if (_submit_deps_buf_inline246[0].is_valid()) + params_t31_deps[params_t31_deps_count++] = _submit_deps_buf_inline246[0]; + params_t31.set_dependencies(params_t31_deps, params_t31_deps_count); + TaskOutputTensors task_31_outs = rt_submit_aic_task(32, params_t31); + PTO2TaskId up_tid_inline290 = task_31_outs.task_id(); + up_tids_inline310[((n_out_inline275 * 5) + k_split_inline276)] = up_tid_inline290; + } + } + for (int64_t n_out_inline292 = 0; n_out_inline292 < 17; n_out_inline292 += 1) { + int64_t n0_inline122__ssa_v7 = (n_out_inline292 * 1024); + PTO2TaskId _submit_deps_buf_inline167[11]; + for (int64_t __init_i = 0; __init_i < 11; ++__init_i) + _submit_deps_buf_inline167[__init_i] = PTO2TaskId::invalid(); + _submit_deps_buf_inline167[0] = reduce_tid_inline226; + PTO2TaskId t__tmp_v137 = gate_tids_inline56[(n_out_inline292 * 5)]; + _submit_deps_buf_inline167[1] = t__tmp_v137; + PTO2TaskId t__tmp_v138 = gate_tids_inline56[((n_out_inline292 * 5) + 1)]; + _submit_deps_buf_inline167[2] = t__tmp_v138; + PTO2TaskId t__tmp_v139 = gate_tids_inline56[((n_out_inline292 * 5) + 2)]; + _submit_deps_buf_inline167[3] = t__tmp_v139; + PTO2TaskId t__tmp_v140 = gate_tids_inline56[((n_out_inline292 * 5) + 3)]; + _submit_deps_buf_inline167[4] = t__tmp_v140; + PTO2TaskId t__tmp_v141 = gate_tids_inline56[((n_out_inline292 * 5) + 4)]; + _submit_deps_buf_inline167[5] = t__tmp_v141; + PTO2TaskId t__tmp_v142 = up_tids_inline310[(n_out_inline292 * 5)]; + _submit_deps_buf_inline167[6] = t__tmp_v142; + PTO2TaskId t__tmp_v143 = up_tids_inline310[((n_out_inline292 * 5) + 1)]; + _submit_deps_buf_inline167[7] = t__tmp_v143; + PTO2TaskId t__tmp_v144 = up_tids_inline310[((n_out_inline292 * 5) + 2)]; + _submit_deps_buf_inline167[8] = t__tmp_v144; + PTO2TaskId t__tmp_v145 = up_tids_inline310[((n_out_inline292 * 5) + 3)]; + _submit_deps_buf_inline167[9] = t__tmp_v145; + PTO2TaskId t__tmp_v146 = up_tids_inline310[((n_out_inline292 * 5) + 4)]; + _submit_deps_buf_inline167[10] = t__tmp_v146; + + // Task 32: silu + L0TaskArgs params_t32; + params_t32.add_input(inv_rms_tile_inline126); + params_t32.add_inout(mlp_tile_inline149); + params_t32.add_input(gate_acc_all_inline203); + params_t32.add_input(up_acc_all_inline303); + params_t32.add_scalar(n0_inline122__ssa_v7); + PTO2TaskId params_t32_deps[11]; + uint32_t params_t32_deps_count = 0; + if (_submit_deps_buf_inline167[0].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[0]; + if (_submit_deps_buf_inline167[1].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[1]; + if (_submit_deps_buf_inline167[2].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[2]; + if (_submit_deps_buf_inline167[3].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[3]; + if (_submit_deps_buf_inline167[4].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[4]; + if (_submit_deps_buf_inline167[5].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[5]; + if (_submit_deps_buf_inline167[6].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[6]; + if (_submit_deps_buf_inline167[7].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[7]; + if (_submit_deps_buf_inline167[8].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[8]; + if (_submit_deps_buf_inline167[9].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[9]; + if (_submit_deps_buf_inline167[10].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[10]; + params_t32.set_dependencies(params_t32_deps, params_t32_deps_count); + TaskOutputTensors task_32_outs = rt_submit_aiv_task(33, params_t32); + PTO2TaskId silu_tid_inline80 = task_32_outs.task_id(); + silu_tids_inline265[n_out_inline292] = silu_tid_inline80; + } + for (int64_t n_out_inline195 = 0; n_out_inline195 < 5; n_out_inline195 += 1) { + int64_t n0_inline122__ssa_v8 = (n_out_inline195 * 1024); + for (int64_t k_split_inline302 = 0; k_split_inline302 < 17; k_split_inline302 += 1) { + int64_t k0_inline113__ssa_v8 = (k_split_inline302 * 1024); + PTO2TaskId _submit_deps_buf_inline229[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline229[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v152 = silu_tids_inline265[k_split_inline302]; + _submit_deps_buf_inline229[0] = t__tmp_v152; + + // Task 33: down_proj + L0TaskArgs params_t33; + params_t33.add_input(mlp_tile_inline149); + params_t33.add_input(ext_w_down); + params_t33.add_inout(down_acc_all_inline168); + params_t33.add_scalar(k0_inline113__ssa_v8); + params_t33.add_scalar(layer_inter_base_inline107); + params_t33.add_scalar(n0_inline122__ssa_v8); + PTO2TaskId params_t33_deps[1]; + uint32_t params_t33_deps_count = 0; + if (_submit_deps_buf_inline229[0].is_valid()) + params_t33_deps[params_t33_deps_count++] = _submit_deps_buf_inline229[0]; + params_t33.set_dependencies(params_t33_deps, params_t33_deps_count); + params_t33.set_allow_early_resolve(true); + TaskOutputTensors task_33_outs = rt_submit_aic_task(34, params_t33); + PTO2TaskId down_tid_inline210 = task_33_outs.task_id(); + down_tids_inline156[((n_out_inline195 * 17) + k_split_inline302)] = down_tid_inline210; + } + } + } + PTO2TaskId _submit_deps_buf_inline123[85]; + for (int64_t __init_i = 0; __init_i < 85; ++__init_i) + _submit_deps_buf_inline123[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v153 = down_tids_inline156[0]; + _submit_deps_buf_inline123[0] = t__tmp_v153; + PTO2TaskId t__tmp_v154 = down_tids_inline156[1]; + _submit_deps_buf_inline123[1] = t__tmp_v154; + PTO2TaskId t__tmp_v155 = down_tids_inline156[2]; + _submit_deps_buf_inline123[2] = t__tmp_v155; + PTO2TaskId t__tmp_v156 = down_tids_inline156[3]; + _submit_deps_buf_inline123[3] = t__tmp_v156; + PTO2TaskId t__tmp_v157 = down_tids_inline156[4]; + _submit_deps_buf_inline123[4] = t__tmp_v157; + PTO2TaskId t__tmp_v158 = down_tids_inline156[5]; + _submit_deps_buf_inline123[5] = t__tmp_v158; + PTO2TaskId t__tmp_v159 = down_tids_inline156[6]; + _submit_deps_buf_inline123[6] = t__tmp_v159; + PTO2TaskId t__tmp_v160 = down_tids_inline156[7]; + _submit_deps_buf_inline123[7] = t__tmp_v160; + PTO2TaskId t__tmp_v161 = down_tids_inline156[8]; + _submit_deps_buf_inline123[8] = t__tmp_v161; + PTO2TaskId t__tmp_v162 = down_tids_inline156[9]; + _submit_deps_buf_inline123[9] = t__tmp_v162; + PTO2TaskId t__tmp_v163 = down_tids_inline156[10]; + _submit_deps_buf_inline123[10] = t__tmp_v163; + PTO2TaskId t__tmp_v164 = down_tids_inline156[11]; + _submit_deps_buf_inline123[11] = t__tmp_v164; + PTO2TaskId t__tmp_v165 = down_tids_inline156[12]; + _submit_deps_buf_inline123[12] = t__tmp_v165; + PTO2TaskId t__tmp_v166 = down_tids_inline156[13]; + _submit_deps_buf_inline123[13] = t__tmp_v166; + PTO2TaskId t__tmp_v167 = down_tids_inline156[14]; + _submit_deps_buf_inline123[14] = t__tmp_v167; + PTO2TaskId t__tmp_v168 = down_tids_inline156[15]; + _submit_deps_buf_inline123[15] = t__tmp_v168; + PTO2TaskId t__tmp_v169 = down_tids_inline156[16]; + _submit_deps_buf_inline123[16] = t__tmp_v169; + PTO2TaskId t__tmp_v170 = down_tids_inline156[17]; + _submit_deps_buf_inline123[17] = t__tmp_v170; + PTO2TaskId t__tmp_v171 = down_tids_inline156[18]; + _submit_deps_buf_inline123[18] = t__tmp_v171; + PTO2TaskId t__tmp_v172 = down_tids_inline156[19]; + _submit_deps_buf_inline123[19] = t__tmp_v172; + PTO2TaskId t__tmp_v173 = down_tids_inline156[20]; + _submit_deps_buf_inline123[20] = t__tmp_v173; + PTO2TaskId t__tmp_v174 = down_tids_inline156[21]; + _submit_deps_buf_inline123[21] = t__tmp_v174; + PTO2TaskId t__tmp_v175 = down_tids_inline156[22]; + _submit_deps_buf_inline123[22] = t__tmp_v175; + PTO2TaskId t__tmp_v176 = down_tids_inline156[23]; + _submit_deps_buf_inline123[23] = t__tmp_v176; + PTO2TaskId t__tmp_v177 = down_tids_inline156[24]; + _submit_deps_buf_inline123[24] = t__tmp_v177; + PTO2TaskId t__tmp_v178 = down_tids_inline156[25]; + _submit_deps_buf_inline123[25] = t__tmp_v178; + PTO2TaskId t__tmp_v179 = down_tids_inline156[26]; + _submit_deps_buf_inline123[26] = t__tmp_v179; + PTO2TaskId t__tmp_v180 = down_tids_inline156[27]; + _submit_deps_buf_inline123[27] = t__tmp_v180; + PTO2TaskId t__tmp_v181 = down_tids_inline156[28]; + _submit_deps_buf_inline123[28] = t__tmp_v181; + PTO2TaskId t__tmp_v182 = down_tids_inline156[29]; + _submit_deps_buf_inline123[29] = t__tmp_v182; + PTO2TaskId t__tmp_v183 = down_tids_inline156[30]; + _submit_deps_buf_inline123[30] = t__tmp_v183; + PTO2TaskId t__tmp_v184 = down_tids_inline156[31]; + _submit_deps_buf_inline123[31] = t__tmp_v184; + PTO2TaskId t__tmp_v185 = down_tids_inline156[32]; + _submit_deps_buf_inline123[32] = t__tmp_v185; + PTO2TaskId t__tmp_v186 = down_tids_inline156[33]; + _submit_deps_buf_inline123[33] = t__tmp_v186; + PTO2TaskId t__tmp_v187 = down_tids_inline156[34]; + _submit_deps_buf_inline123[34] = t__tmp_v187; + PTO2TaskId t__tmp_v188 = down_tids_inline156[35]; + _submit_deps_buf_inline123[35] = t__tmp_v188; + PTO2TaskId t__tmp_v189 = down_tids_inline156[36]; + _submit_deps_buf_inline123[36] = t__tmp_v189; + PTO2TaskId t__tmp_v190 = down_tids_inline156[37]; + _submit_deps_buf_inline123[37] = t__tmp_v190; + PTO2TaskId t__tmp_v191 = down_tids_inline156[38]; + _submit_deps_buf_inline123[38] = t__tmp_v191; + PTO2TaskId t__tmp_v192 = down_tids_inline156[39]; + _submit_deps_buf_inline123[39] = t__tmp_v192; + PTO2TaskId t__tmp_v193 = down_tids_inline156[40]; + _submit_deps_buf_inline123[40] = t__tmp_v193; + PTO2TaskId t__tmp_v194 = down_tids_inline156[41]; + _submit_deps_buf_inline123[41] = t__tmp_v194; + PTO2TaskId t__tmp_v195 = down_tids_inline156[42]; + _submit_deps_buf_inline123[42] = t__tmp_v195; + PTO2TaskId t__tmp_v196 = down_tids_inline156[43]; + _submit_deps_buf_inline123[43] = t__tmp_v196; + PTO2TaskId t__tmp_v197 = down_tids_inline156[44]; + _submit_deps_buf_inline123[44] = t__tmp_v197; + PTO2TaskId t__tmp_v198 = down_tids_inline156[45]; + _submit_deps_buf_inline123[45] = t__tmp_v198; + PTO2TaskId t__tmp_v199 = down_tids_inline156[46]; + _submit_deps_buf_inline123[46] = t__tmp_v199; + PTO2TaskId t__tmp_v200 = down_tids_inline156[47]; + _submit_deps_buf_inline123[47] = t__tmp_v200; + PTO2TaskId t__tmp_v201 = down_tids_inline156[48]; + _submit_deps_buf_inline123[48] = t__tmp_v201; + PTO2TaskId t__tmp_v202 = down_tids_inline156[49]; + _submit_deps_buf_inline123[49] = t__tmp_v202; + PTO2TaskId t__tmp_v203 = down_tids_inline156[50]; + _submit_deps_buf_inline123[50] = t__tmp_v203; + PTO2TaskId t__tmp_v204 = down_tids_inline156[51]; + _submit_deps_buf_inline123[51] = t__tmp_v204; + PTO2TaskId t__tmp_v205 = down_tids_inline156[52]; + _submit_deps_buf_inline123[52] = t__tmp_v205; + PTO2TaskId t__tmp_v206 = down_tids_inline156[53]; + _submit_deps_buf_inline123[53] = t__tmp_v206; + PTO2TaskId t__tmp_v207 = down_tids_inline156[54]; + _submit_deps_buf_inline123[54] = t__tmp_v207; + PTO2TaskId t__tmp_v208 = down_tids_inline156[55]; + _submit_deps_buf_inline123[55] = t__tmp_v208; + PTO2TaskId t__tmp_v209 = down_tids_inline156[56]; + _submit_deps_buf_inline123[56] = t__tmp_v209; + PTO2TaskId t__tmp_v210 = down_tids_inline156[57]; + _submit_deps_buf_inline123[57] = t__tmp_v210; + PTO2TaskId t__tmp_v211 = down_tids_inline156[58]; + _submit_deps_buf_inline123[58] = t__tmp_v211; + PTO2TaskId t__tmp_v212 = down_tids_inline156[59]; + _submit_deps_buf_inline123[59] = t__tmp_v212; + PTO2TaskId t__tmp_v213 = down_tids_inline156[60]; + _submit_deps_buf_inline123[60] = t__tmp_v213; + PTO2TaskId t__tmp_v214 = down_tids_inline156[61]; + _submit_deps_buf_inline123[61] = t__tmp_v214; + PTO2TaskId t__tmp_v215 = down_tids_inline156[62]; + _submit_deps_buf_inline123[62] = t__tmp_v215; + PTO2TaskId t__tmp_v216 = down_tids_inline156[63]; + _submit_deps_buf_inline123[63] = t__tmp_v216; + PTO2TaskId t__tmp_v217 = down_tids_inline156[64]; + _submit_deps_buf_inline123[64] = t__tmp_v217; + PTO2TaskId t__tmp_v218 = down_tids_inline156[65]; + _submit_deps_buf_inline123[65] = t__tmp_v218; + PTO2TaskId t__tmp_v219 = down_tids_inline156[66]; + _submit_deps_buf_inline123[66] = t__tmp_v219; + PTO2TaskId t__tmp_v220 = down_tids_inline156[67]; + _submit_deps_buf_inline123[67] = t__tmp_v220; + PTO2TaskId t__tmp_v221 = down_tids_inline156[68]; + _submit_deps_buf_inline123[68] = t__tmp_v221; + PTO2TaskId t__tmp_v222 = down_tids_inline156[69]; + _submit_deps_buf_inline123[69] = t__tmp_v222; + PTO2TaskId t__tmp_v223 = down_tids_inline156[70]; + _submit_deps_buf_inline123[70] = t__tmp_v223; + PTO2TaskId t__tmp_v224 = down_tids_inline156[71]; + _submit_deps_buf_inline123[71] = t__tmp_v224; + PTO2TaskId t__tmp_v225 = down_tids_inline156[72]; + _submit_deps_buf_inline123[72] = t__tmp_v225; + PTO2TaskId t__tmp_v226 = down_tids_inline156[73]; + _submit_deps_buf_inline123[73] = t__tmp_v226; + PTO2TaskId t__tmp_v227 = down_tids_inline156[74]; + _submit_deps_buf_inline123[74] = t__tmp_v227; + PTO2TaskId t__tmp_v228 = down_tids_inline156[75]; + _submit_deps_buf_inline123[75] = t__tmp_v228; + PTO2TaskId t__tmp_v229 = down_tids_inline156[76]; + _submit_deps_buf_inline123[76] = t__tmp_v229; + PTO2TaskId t__tmp_v230 = down_tids_inline156[77]; + _submit_deps_buf_inline123[77] = t__tmp_v230; + PTO2TaskId t__tmp_v231 = down_tids_inline156[78]; + _submit_deps_buf_inline123[78] = t__tmp_v231; + PTO2TaskId t__tmp_v232 = down_tids_inline156[79]; + _submit_deps_buf_inline123[79] = t__tmp_v232; + PTO2TaskId t__tmp_v233 = down_tids_inline156[80]; + _submit_deps_buf_inline123[80] = t__tmp_v233; + PTO2TaskId t__tmp_v234 = down_tids_inline156[81]; + _submit_deps_buf_inline123[81] = t__tmp_v234; + PTO2TaskId t__tmp_v235 = down_tids_inline156[82]; + _submit_deps_buf_inline123[82] = t__tmp_v235; + PTO2TaskId t__tmp_v236 = down_tids_inline156[83]; + _submit_deps_buf_inline123[83] = t__tmp_v236; + PTO2TaskId t__tmp_v237 = down_tids_inline156[84]; + _submit_deps_buf_inline123[84] = t__tmp_v237; + + // Spmd dcr_xgamma_spmd: dcr_xgamma + L0TaskArgs params_t34; + params_t34.add_input(down_acc_all_inline168); + params_t34.add_input(post_norm_partial_inline118); + params_t34.add_inout(next_hidden); + params_t34.add_input(ext_input_rms_weight); + params_t34.add_inout(next_normed); + params_t34.add_scalar(next_gamma_idx); + params_t34.launch_spec.set_block_num(5); + params_t34.set_allow_early_resolve(true); + PTO2TaskId params_t34_deps[85]; + uint32_t params_t34_deps_count = 0; + if (_submit_deps_buf_inline123[0].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[0]; + if (_submit_deps_buf_inline123[1].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[1]; + if (_submit_deps_buf_inline123[2].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[2]; + if (_submit_deps_buf_inline123[3].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[3]; + if (_submit_deps_buf_inline123[4].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[4]; + if (_submit_deps_buf_inline123[5].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[5]; + if (_submit_deps_buf_inline123[6].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[6]; + if (_submit_deps_buf_inline123[7].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[7]; + if (_submit_deps_buf_inline123[8].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[8]; + if (_submit_deps_buf_inline123[9].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[9]; + if (_submit_deps_buf_inline123[10].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[10]; + if (_submit_deps_buf_inline123[11].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[11]; + if (_submit_deps_buf_inline123[12].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[12]; + if (_submit_deps_buf_inline123[13].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[13]; + if (_submit_deps_buf_inline123[14].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[14]; + if (_submit_deps_buf_inline123[15].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[15]; + if (_submit_deps_buf_inline123[16].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[16]; + if (_submit_deps_buf_inline123[17].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[17]; + if (_submit_deps_buf_inline123[18].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[18]; + if (_submit_deps_buf_inline123[19].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[19]; + if (_submit_deps_buf_inline123[20].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[20]; + if (_submit_deps_buf_inline123[21].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[21]; + if (_submit_deps_buf_inline123[22].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[22]; + if (_submit_deps_buf_inline123[23].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[23]; + if (_submit_deps_buf_inline123[24].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[24]; + if (_submit_deps_buf_inline123[25].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[25]; + if (_submit_deps_buf_inline123[26].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[26]; + if (_submit_deps_buf_inline123[27].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[27]; + if (_submit_deps_buf_inline123[28].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[28]; + if (_submit_deps_buf_inline123[29].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[29]; + if (_submit_deps_buf_inline123[30].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[30]; + if (_submit_deps_buf_inline123[31].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[31]; + if (_submit_deps_buf_inline123[32].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[32]; + if (_submit_deps_buf_inline123[33].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[33]; + if (_submit_deps_buf_inline123[34].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[34]; + if (_submit_deps_buf_inline123[35].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[35]; + if (_submit_deps_buf_inline123[36].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[36]; + if (_submit_deps_buf_inline123[37].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[37]; + if (_submit_deps_buf_inline123[38].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[38]; + if (_submit_deps_buf_inline123[39].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[39]; + if (_submit_deps_buf_inline123[40].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[40]; + if (_submit_deps_buf_inline123[41].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[41]; + if (_submit_deps_buf_inline123[42].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[42]; + if (_submit_deps_buf_inline123[43].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[43]; + if (_submit_deps_buf_inline123[44].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[44]; + if (_submit_deps_buf_inline123[45].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[45]; + if (_submit_deps_buf_inline123[46].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[46]; + if (_submit_deps_buf_inline123[47].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[47]; + if (_submit_deps_buf_inline123[48].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[48]; + if (_submit_deps_buf_inline123[49].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[49]; + if (_submit_deps_buf_inline123[50].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[50]; + if (_submit_deps_buf_inline123[51].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[51]; + if (_submit_deps_buf_inline123[52].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[52]; + if (_submit_deps_buf_inline123[53].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[53]; + if (_submit_deps_buf_inline123[54].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[54]; + if (_submit_deps_buf_inline123[55].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[55]; + if (_submit_deps_buf_inline123[56].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[56]; + if (_submit_deps_buf_inline123[57].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[57]; + if (_submit_deps_buf_inline123[58].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[58]; + if (_submit_deps_buf_inline123[59].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[59]; + if (_submit_deps_buf_inline123[60].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[60]; + if (_submit_deps_buf_inline123[61].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[61]; + if (_submit_deps_buf_inline123[62].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[62]; + if (_submit_deps_buf_inline123[63].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[63]; + if (_submit_deps_buf_inline123[64].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[64]; + if (_submit_deps_buf_inline123[65].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[65]; + if (_submit_deps_buf_inline123[66].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[66]; + if (_submit_deps_buf_inline123[67].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[67]; + if (_submit_deps_buf_inline123[68].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[68]; + if (_submit_deps_buf_inline123[69].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[69]; + if (_submit_deps_buf_inline123[70].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[70]; + if (_submit_deps_buf_inline123[71].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[71]; + if (_submit_deps_buf_inline123[72].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[72]; + if (_submit_deps_buf_inline123[73].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[73]; + if (_submit_deps_buf_inline123[74].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[74]; + if (_submit_deps_buf_inline123[75].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[75]; + if (_submit_deps_buf_inline123[76].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[76]; + if (_submit_deps_buf_inline123[77].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[77]; + if (_submit_deps_buf_inline123[78].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[78]; + if (_submit_deps_buf_inline123[79].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[79]; + if (_submit_deps_buf_inline123[80].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[80]; + if (_submit_deps_buf_inline123[81].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[81]; + if (_submit_deps_buf_inline123[82].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[82]; + if (_submit_deps_buf_inline123[83].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[83]; + if (_submit_deps_buf_inline123[84].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[84]; + params_t34.set_dependencies(params_t34_deps, params_t34_deps_count); + rt_submit_aiv_task(35, params_t34); + }; + rt_submit_graph(GRAPH_KEY("qwen3_14b_decoder_layer_v1"), +layer_definition, graph_args); + Tensor cur__ssa_v8 = next_hidden; + Tensor normed__ssa_v6 = next_normed; + cur__rv_v7 = cur__ssa_v8; + normed__rv_v5 = normed__ssa_v6; + } + } + for (int64_t ob0 = 0; ob0 < 16; ob0 += 16) { + PTO2_SCOPE() { + // Task 35: copy_out + L0TaskArgs params_t35; + params_t35.add_output(ext_out); + params_t35.add_input(cur__rv_v7); + params_t35.add_scalar(ob0); + rt_submit_aiv_task(36, params_t35); + } + } + } +} + +} // extern "C" diff --git a/tests/st/a2a3/host_build_graph/graph_execution/test_graph_execution.py b/tests/st/a2a3/host_build_graph/graph_execution/test_graph_execution.py new file mode 100644 index 0000000000..4876a8b7ee --- /dev/null +++ b/tests/st/a2a3/host_build_graph/graph_execution/test_graph_execution.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Graph Execution records once and replays topology with dynamic L0TaskArgs tensors.""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test + + +@scene_test(level=2, runtime="host_build_graph") +class TestGraphExecutionHostBuildGraph(SceneTestCase): + RTOL = 1e-5 + ATOL = 1e-5 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/graph_execution_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT, D.OUT, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": "../vector_example/kernels/aiv/kernel_add.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "source": "../vector_example/kernels/aiv/kernel_add_scalar.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT], + }, + { + "func_id": 2, + "source": "../vector_example/kernels/aiv/kernel_mul.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + ], + } + + CASES = [ + { + "name": "record_then_replay_1d", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4, "block_dim": 3}, + "params": {"shape": (128 * 128,)}, + }, + { + "name": "record_then_replay_2d", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4, "block_dim": 3}, + "params": {"shape": (128 * 128, 1)}, + }, + ] + + def generate_args(self, params): + shape = params["shape"] + return TaskArgsBuilder( + Tensor("a", torch.full(shape, 2.0, dtype=torch.float32)), + Tensor("b", torch.full(shape, 3.0, dtype=torch.float32)), + Tensor("output_1", torch.zeros(shape, dtype=torch.float32)), + Tensor("output_3", torch.zeros(shape, dtype=torch.float32)), + Tensor("output_5", torch.zeros(shape, dtype=torch.float32)), + ) + + def compute_golden(self, args, params): + base = args.a + args.b + ndim_delta = 0.0 if args.a.ndim == 1 else 2.0 + expected = (base + 1.0 + ndim_delta) * (base + 2.0 + ndim_delta) + args.output_1[:] = expected + args.output_3[:] = expected + args.output_5[:] = expected + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a2a3/host_build_graph/graph_execution/test_graph_execution_aic_aiv.py b/tests/st/a2a3/host_build_graph/graph_execution/test_graph_execution_aic_aiv.py new file mode 100644 index 0000000000..8a7173c083 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/graph_execution/test_graph_execution_aic_aiv.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Graph Execution covers a Qwen-style fixed AIC/AIV decoder-layer DAG.""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test + + +@scene_test(level=2, runtime="host_build_graph") +class TestGraphExecutionAicAivHostBuildGraph(SceneTestCase): + RTOL = 1e-2 + ATOL = 1e-2 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/graph_execution_aic_aiv_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.IN, D.OUT, D.OUT, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": "../matmul/kernels/aiv/kernel_log_sqrt.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT], + }, + { + "func_id": 1, + "source": "../matmul/kernels/aic/kernel_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 2, + "source": "../matmul/kernels/aiv/kernel_add_exp.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + ], + } + + CASES = [ + { + "name": "record_then_replay_aic_aiv", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4}, + "params": {}, + }, + ] + + def generate_args(self, params): + rows = 128 + columns = 128 + size = rows * columns + input_value = torch.exp(torch.tensor(4.0)).item() + weight_value = 1.0 / (2 * columns) + return TaskArgsBuilder( + Tensor("input", torch.full((size,), input_value, dtype=torch.float16)), + Tensor("weight_1", torch.full((size,), weight_value, dtype=torch.float16)), + Tensor("weight_2", torch.full((size,), weight_value, dtype=torch.float16)), + Tensor("output_1", torch.zeros(size, dtype=torch.float32)), + Tensor("output_2", torch.zeros(size, dtype=torch.float32)), + Tensor("output_3", torch.zeros(size, dtype=torch.float32)), + ) + + def compute_golden(self, args, params): + rows = 128 + columns = 128 + normalized = torch.sqrt(torch.log(args.input.reshape(rows, columns).to(torch.float32))) + left = torch.matmul(normalized, args.weight_1.reshape(rows, columns).to(torch.float32)) + right = torch.matmul(normalized, args.weight_2.reshape(rows, columns).to(torch.float32)) + expected = torch.exp(left + right).flatten().to(torch.float32) + args.output_1[:] = expected + args.output_2[:] = expected + args.output_3[:] = expected + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a2a3/host_build_graph/graph_execution/test_graph_execution_mix_spmd.py b/tests/st/a2a3/host_build_graph/graph_execution/test_graph_execution_mix_spmd.py new file mode 100644 index 0000000000..b500cf90cb --- /dev/null +++ b/tests/st/a2a3/host_build_graph/graph_execution/test_graph_execution_mix_spmd.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Graph Execution preserves MIX active slots and multi-block SPMD metadata.""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test + +FLOATS_PER_CACHE_LINE = 16 +SLOTS_PER_BLOCK = 3 +MAX_CLUSTERS = 24 +TOTAL_FLOATS = MAX_CLUSTERS * SLOTS_PER_BLOCK * FLOATS_PER_CACHE_LINE + + +@scene_test(level=2, runtime="host_build_graph") +class TestGraphExecutionMixSpmdHostBuildGraph(SceneTestCase): + RTOL = 0 + ATOL = 0 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/graph_execution_mix_spmd_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT, D.INOUT, D.INOUT], + }, + "incores": [ + { + "func_id": 0, + "source": "../../tensormap_and_ringbuffer/spmd_multiblock_mix/kernels/aic/kernel_spmd_mix.cpp", + "core_type": "aic", + "signature": [D.INOUT], + }, + { + "func_id": 1, + "source": "../../tensormap_and_ringbuffer/spmd_multiblock_mix/kernels/aiv/kernel_spmd_mix.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + { + "func_id": 2, + "source": "../../tensormap_and_ringbuffer/spmd_multiblock_mix/kernels/aiv/kernel_spmd_mix.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "record_then_replay_mix_spmd", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4}, + "params": {}, + }, + ] + + def generate_args(self, params): + return TaskArgsBuilder( + Tensor("blocks_1", torch.zeros(TOTAL_FLOATS, dtype=torch.float32)), + Tensor("blocks_2", torch.zeros(TOTAL_FLOATS, dtype=torch.float32)), + Tensor("blocks_3", torch.zeros(TOTAL_FLOATS, dtype=torch.float32)), + ) + + def compute_golden(self, args, params): + # The exact whole-device cluster count is platform-provided. Validate + # it from the first execution and require both Graph replays to match. + pass + + def compare_outputs(self, test_args, golden_args, output_names, params): + outputs = [test_args.blocks_1, test_args.blocks_2, test_args.blocks_3] + assert torch.equal(outputs[0], outputs[1]) + assert torch.equal(outputs[0], outputs[2]) + + cache_line_heads = outputs[0].reshape(-1, FLOATS_PER_CACHE_LINE)[:, 0] + nonzero = torch.nonzero(cache_line_heads, as_tuple=False) + assert nonzero.numel() > 0, "SPMD Graph did not execute any non-zero block" + cluster_count = int(cache_line_heads.max().item()) + 1 + assert 1 < cluster_count <= MAX_CLUSTERS + + expected = torch.zeros_like(cache_line_heads) + for block_idx in range(cluster_count): + begin = block_idx * SLOTS_PER_BLOCK + expected[begin : begin + SLOTS_PER_BLOCK] = float(block_idx) + assert torch.equal(cache_line_heads, expected) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a2a3/host_build_graph/graph_execution/test_qwen3_14b_3layer_graph_execution.py b/tests/st/a2a3/host_build_graph/graph_execution/test_qwen3_14b_3layer_graph_execution.py new file mode 100644 index 0000000000..c0af8248a6 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/graph_execution/test_qwen3_14b_3layer_graph_execution.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Qwen3-14B three-layer decode coverage for host-built Graph Execution. + +Layer 0 records and executes the decoder-layer task graph. Layers 1 and 2 submit +one Graph task each; the scheduler expands the saved topology and applies the +new layer's boundary tensors. +""" + +from __future__ import annotations + +import copy +import importlib.util +import sys +from pathlib import Path + +from simpler_setup import SceneTestCase, scene_test +from simpler_setup.goldens.qwen3_14b_decode import compute_golden as _decode_golden +from simpler_setup.goldens.qwen3_14b_decode import generate_inputs as _decode_generate_inputs + +N_LAYERS = 3 +HERE = Path(__file__).resolve().parent +REPO_ROOT = HERE.parents[4] +QWEN_DIR = REPO_ROOT / "examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode" + + +def _load_qwen_case(): + module_name = "_qwen3_14b_graph_execution_base" + spec = importlib.util.spec_from_file_location(module_name, QWEN_DIR / "test_qwen3_14b_decode.py") + if spec is None or spec.loader is None: + raise RuntimeError("cannot load the Qwen3-14B decode case") + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module.TestQwen314BDecode + + +def _callable(): + callable_cfg = copy.deepcopy(_load_qwen_case().CALLABLE) + callable_cfg["orchestration"]["source"] = str(HERE / "kernels/orchestration/qwen3_14b_3layer_graph_execution.cpp") + for incore in callable_cfg["incores"]: + incore["source"] = str(QWEN_DIR / incore["source"]) + return callable_cfg + + +@scene_test(level=2, runtime="host_build_graph") +class TestQwen314B3LayerGraphExecution(SceneTestCase): + RTOL = 5e-2 + ATOL = 1e-1 + CALLABLE = _callable() + + CASES = [ + { + "name": "GraphExecutionBatch16Seq3500", + "platforms": ["a2a3"], + "config": {"aicpu_thread_num": 4, "block_dim": 0}, + # The three-layer fixture is about 3 GiB and compiles the complete + # Qwen decoder kernel set, so keep it out of routine CI. + "manual": True, + "params": {"seed": 1234, "seq_len": 3500}, + }, + ] + + def generate_args(self, params): + return _decode_generate_inputs( + seed=params.get("seed", 1234), + seq_len=params.get("seq_len", 3500), + n_layers=N_LAYERS, + ) + + def compute_golden(self, args, params): + _decode_golden(args, n_layers=N_LAYERS) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/ut/py/test_swimlane_converter.py b/tests/ut/py/test_swimlane_converter.py index 0d5ff42772..b8c0d522b8 100644 --- a/tests/ut/py/test_swimlane_converter.py +++ b/tests/ut/py/test_swimlane_converter.py @@ -10,6 +10,8 @@ import json +import pytest + from simpler_setup.tools import swimlane_converter as sc @@ -163,6 +165,209 @@ def test_load_func_names_auto_discovery_and_explicit_precedence(tmp_path): assert func_names == {"0": "explicit"} +def test_graph_prepare_phases_create_graph_execution_envelopes(tmp_path): + out = tmp_path / "trace.json" + outer_a = 3 + outer_b = 7 + node_a0 = (1 << 32) | (outer_a << 10) + node_a1 = (1 << 32) | ((outer_a << 10) | 1) + node_b0 = (1 << 32) | (outer_b << 10) + scheduler_phases = [ + [ + { + "phase": "graph_prepare", + "task_id": outer_a, + "start_time_us": 1.0, + "end_time_us": 1.4, + "tasks_processed": 1, + }, + { + "phase": "graph_prepare", + "task_id": outer_a, + "start_time_us": 1.5, + "end_time_us": 1.8, + "tasks_processed": 1, + }, + { + "phase": "graph_prepare", + "task_id": outer_b, + "start_time_us": 5.0, + "end_time_us": 5.2, + "tasks_processed": 1, + }, + ] + ] + tasks = [ + _task_row(node_a0, 0, dispatch=2.0, start=2.2, end=3.0, receive=2.1), + _task_row(node_a1, 1, dispatch=3.2, start=3.4, end=4.0, receive=3.3), + _task_row(node_b0, 0, dispatch=5.3, start=5.5, end=6.0, receive=5.4), + ] + + sc.generate_chrome_trace_json(tasks, str(out), scheduler_phases=scheduler_phases, core_to_thread=[0, 0]) + + with open(out) as f: + events = json.load(f)["traceEvents"] + assert any( + event.get("ph") == "M" and event.get("pid") == 5 and event.get("args", {}).get("name") == "Graph Execution" + for event in events + ) + graph_events = [event for event in events if event.get("cat") == "graph_execution"] + assert [event["args"]["outer_task_id"] for event in graph_events] == [outer_a, outer_b] + assert graph_events[0]["args"]["visible_node_count"] == 2 + assert graph_events[0]["args"]["prepare_slice_count"] == 2 + assert graph_events[0]["ts"] == 1.0 + assert graph_events[0]["dur"] == 4.0 + assert ( + sum(event.get("cat") == "scheduler" and event.get("name", "").startswith("graph_prepare(") for event in events) + == 3 + ) + + +def test_host_orchestrator_uses_separate_clock_domain_and_lane(tmp_path): + raw = tmp_path / "l2_swimlane_records.json" + raw.write_text( + json.dumps( + { + "l2_swimlane_level": 4, + "metadata": {"clock_freq_hz": 50_000_000, "num_cores": 0, "core_types": []}, + "aicore_tasks": [], + "aicpu_tasks": [], + "aicpu_scheduler_phases": [], + "host_orchestrator": { + "start_cycles": 100, + "end_cycles": 200, + "records": [{"submit_idx": 0, "task_id": 3, "start_cycles": 120, "end_cycles": 140}], + }, + } + ) + ) + parsed = sc.read_perf_data(raw) + host = parsed["host_orchestrator"] + assert host["start_time_us"] == -2.0 + assert host["end_time_us"] == 0.0 + assert host["records"][0]["start_time_us"] == -1.6 + assert host["records"][0]["end_time_us"] == -1.2 + + out = tmp_path / "trace.json" + tasks = [_task_row(1, 0, dispatch=0.5, start=1.0, end=1.5, receive=0.75)] + sc.generate_chrome_trace_json(tasks, str(out), host_orchestrator=host) + with open(out) as f: + events = json.load(f)["traceEvents"] + assert any( + event.get("ph") == "M" and event.get("pid") == 1 and event.get("args", {}).get("name") == "Host Orchestrator" + for event in events + ) + assert not any( + event.get("ph") == "M" and event.get("args", {}).get("name") == "AICPU Orchestrator" for event in events + ) + envelope = next(event for event in events if event.get("name") == "host_orchestration") + assert envelope["ts"] == 0.0 + assert envelope["dur"] == 2.0 + host_submit = next(event for event in events if event.get("name") == "task_submit(t3)") + assert host_submit["ts"] == pytest.approx(0.4) + worker = next(event for event in events if event.get("pid") == 4 and event.get("ph") == "X") + assert worker["ts"] == 2.75 + + +def test_host_orchestrator_envelope_does_not_require_submit_records(tmp_path): + raw = tmp_path / "l2_swimlane_records.json" + raw.write_text( + json.dumps( + { + "l2_swimlane_level": 4, + "metadata": {"clock_freq_hz": 50_000_000, "num_cores": 0, "core_types": []}, + "aicore_tasks": [], + "aicpu_tasks": [], + "aicpu_scheduler_phases": [], + "host_orchestrator": {"start_cycles": 100, "end_cycles": 200, "records": []}, + } + ) + ) + + host = sc.read_perf_data(raw)["host_orchestrator"] + out = tmp_path / "trace.json" + sc.generate_chrome_trace_json([], str(out), host_orchestrator=host) + + with open(out) as f: + events = json.load(f)["traceEvents"] + envelopes = [event for event in events if event.get("name") == "host_orchestration"] + assert len(envelopes) == 1 + assert envelopes[0]["ts"] == 0.0 + assert envelopes[0]["dur"] == 2.0 + + +def test_streaming_host_orchestrator_aligns_first_publish_to_device_zero(tmp_path): + raw = tmp_path / "l2_swimlane_records.json" + raw.write_text( + json.dumps( + { + "l2_swimlane_level": 4, + "metadata": {"clock_freq_hz": 50_000_000, "num_cores": 0, "core_types": []}, + "aicore_tasks": [], + "aicpu_tasks": [], + "aicpu_scheduler_phases": [], + "host_orchestrator": { + "start_cycles": 100, + "end_cycles": 300, + "first_publish_cycles": 200, + "records": [ + {"submit_idx": 0, "task_id": 1, "start_cycles": 140, "end_cycles": 160}, + {"submit_idx": 1, "task_id": 2, "start_cycles": 240, "end_cycles": 260}, + ], + }, + } + ) + ) + + host = sc.read_perf_data(raw)["host_orchestrator"] + assert host["clock_alignment"] == "host_first_publish_aligned_to_first_device_event" + assert host["start_time_us"] == -2.0 + assert host["end_time_us"] == 2.0 + assert host["records"][0]["start_time_us"] == -1.2 + assert host["records"][1]["start_time_us"] == 0.8 + + +def test_host_submit_records_distinguish_graph_outer_task(tmp_path): + out = tmp_path / "trace.json" + outer_task_id = 3 + node_task_id = (1 << 32) | (outer_task_id << 10) + tasks = [_task_row(node_task_id, 0, dispatch=2.0, start=2.2, end=3.0, receive=2.1)] + scheduler_phases = [ + [ + { + "phase": "graph_prepare", + "task_id": outer_task_id, + "start_time_us": 1.0, + "end_time_us": 1.5, + "tasks_processed": 1, + } + ] + ] + host_orchestrator = { + "start_time_us": -1.0, + "end_time_us": 0.0, + "clock_alignment": "host_orch_end_aligned_to_device_zero", + "records": [ + {"submit_idx": 0, "task_id": 2, "start_time_us": -0.9, "end_time_us": -0.8}, + {"submit_idx": 1, "task_id": outer_task_id, "start_time_us": -0.7, "end_time_us": -0.6}, + ], + } + + sc.generate_chrome_trace_json( + tasks, + str(out), + scheduler_phases=scheduler_phases, + core_to_thread=[0], + host_orchestrator=host_orchestrator, + ) + + with open(out) as f: + events = json.load(f)["traceEvents"] + host_submits = [event for event in events if event.get("cat") == "host_orchestrator" and "submit(" in event["name"]] + assert [event["name"] for event in host_submits] == ["task_submit(t2)", "graph_submit(t3)"] + assert [event["args"]["submit_kind"] for event in host_submits] == ["task_submit", "graph_submit"] + + def test_spmd_pred_routes_dependency_to_earliest_slice(tmp_path): pred_id = 100 succ_id = 200