From 66f183e8f05ea28d0f4525633c4840821d3fe00e Mon Sep 17 00:00:00 2001 From: TaoZQY Date: Wed, 22 Jul 2026 01:53:36 -0700 Subject: [PATCH] Add: graph execution to host_build_graph - Cache repeated task DAGs with saved topology and dynamic TaskArgs sources - Submit one Graph task on cache hits and expand nodes in the device Scheduler - Reuse compact host submissions and graph-affine AICPU execution blocks - Add Graph/Host-Orchestrator DFX lanes, simulation coverage, and Qwen API docs - Fix predicate-field compatibility in a2a3/a5 runtime builds Co-authored-by: Crane-Liu --- simpler_setup/tools/swimlane_converter.py | 244 +++- .../platform/onboard/host/device_runner.cpp | 6 + src/a2a3/platform/sim/host/device_runner.cpp | 6 + .../host_build_graph/aicpu/aicpu_executor.cpp | 5 + .../host_build_graph/docs/GRAPH_EXECUTION.md | 171 +++ .../host_build_graph/host/runtime_maker.cpp | 120 +- .../orchestration/pto_orchestration_api.h | 52 + .../orchestrator_core/pto_graph_execution.cpp | 338 +++++ .../orchestrator_core/pto_orchestrator.cpp | 1238 ++++++++++++++++- .../orchestrator_core/pto_runtime2.cpp | 11 + .../runtime/pto_graph_cache.h | 87 ++ .../runtime/pto_graph_execution.h | 158 +++ .../runtime/pto_orchestrator.h | 19 + .../host_build_graph/runtime/pto_runtime2.h | 6 + .../runtime/pto_runtime2_types.h | 21 +- .../host_build_graph/runtime/pto_types.h | 79 +- .../host_build_graph/runtime/runtime.h | 17 + .../runtime/scheduler/pto_scheduler.h | 172 ++- .../scheduler/scheduler_completion.cpp | 7 +- .../runtime/scheduler/scheduler_dispatch.cpp | 56 +- .../runtime/shared/pto_runtime2_init.cpp | 15 + .../host/runtime_maker.cpp | 2 +- .../runtime/runtime.h | 9 + .../runtime/scheduler/scheduler_dispatch.cpp | 2 +- .../host_build_graph/host/runtime_maker.cpp | 2 +- .../host/runtime_maker.cpp | 2 +- .../runtime/scheduler/scheduler_dispatch.cpp | 2 +- .../aicpu/l2_swimlane_collector_aicpu.h | 6 + .../include/common/l2_swimlane_profiling.h | 8 + .../include/host/l2_swimlane_collector.h | 11 + .../platform/onboard/host/c_api_shared.cpp | 2 +- .../onboard/host/device_runner_base.cpp | 6 +- .../onboard/host/device_runner_base.h | 5 +- .../aicpu/l2_swimlane_collector_aicpu.cpp | 14 + .../shared/host/l2_swimlane_collector.cpp | 33 + src/common/platform/sim/host/c_api_shared.cpp | 2 +- .../platform/sim/host/device_runner_base.cpp | 6 +- .../platform/sim/host/device_runner_base.h | 3 +- .../orchestration/graph_execution_orch.cpp | 85 ++ .../graph_execution/test_graph_execution.py | 78 ++ tests/ut/py/test_swimlane_converter.py | 174 +++ 41 files changed, 3215 insertions(+), 65 deletions(-) create mode 100644 src/a2a3/runtime/host_build_graph/docs/GRAPH_EXECUTION.md create mode 100644 src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_graph_execution.cpp create mode 100644 src/a2a3/runtime/host_build_graph/runtime/pto_graph_cache.h create mode 100644 src/a2a3/runtime/host_build_graph/runtime/pto_graph_execution.h create mode 100644 tests/st/a2a3/host_build_graph/graph_execution/kernels/orchestration/graph_execution_orch.cpp create mode 100644 tests/st/a2a3/host_build_graph/graph_execution/test_graph_execution.py diff --git a/simpler_setup/tools/swimlane_converter.py b/simpler_setup/tools/swimlane_converter.py index 08832c7d6b..681417db68 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,7 @@ 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, records: [...]} } aicore_tasks columns (v3 schema): the trailing receive_to_start_cycles @@ -191,6 +285,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 +465,29 @@ 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)) + if host_end_cycles > host_start_cycles: + # Host CLOCK_MONOTONIC and AICPU syscnt do not share an epoch. Preserve + # measured host durations and first align Host Orch end to device t=0. + # generate_chrome_trace_json() then rebases the complete trace so the + # earliest Host Orch timestamp is zero; Perfetto clips negative-time + # slices from its default viewport. + host_orchestrator = { + "start_time_us": (host_start_cycles - host_end_cycles) * cycles_to_us_factor, + "end_time_us": 0.0, + "clock_alignment": "host_orch_end_aligned_to_device_zero", + "records": [], + } + for record in host_orchestrator_raw.get("records", []): + converted = dict(record) + converted["start_time_us"] = (int(record["start_cycles"]) - host_end_cycles) * cycles_to_us_factor + converted["end_time_us"] = (int(record["end_cycles"]) - host_end_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 +496,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 +1216,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 +1234,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 +1272,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 +1291,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 +1656,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 +1818,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 +1938,7 @@ def _find_containing_complete(thread_idx: int, finish_us: float): # orchestrator timing. There is no separate aggregate summary — the # device-side LOG_INFO_V9 "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 +2663,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 +2903,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 90b8ad1a20..4cd980f95e 100644 --- a/src/a2a3/platform/onboard/host/device_runner.cpp +++ b/src/a2a3/platform/onboard/host/device_runner.cpp @@ -333,6 +333,12 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { LOG_ERROR("init_l2_swimlane failed: %d", rc); return rc; } + if (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() + ); + } } if (enable_dump_args_) { diff --git a/src/a2a3/platform/sim/host/device_runner.cpp b/src/a2a3/platform/sim/host/device_runner.cpp index 614c714b77..341665e941 100644 --- a/src/a2a3/platform/sim/host/device_runner.cpp +++ b/src/a2a3/platform/sim/host/device_runner.cpp @@ -317,6 +317,12 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { core_types[i] = runtime.get_workers()[i].core_type; } l2_swimlane_collector_.set_core_types(core_types.data(), num_aicore); + if (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() + ); + } } if (enable_dump_args_) { 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 b2f7df3572..11a2498e7f 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 "pto_graph_execution.h" #include "pto_shared_memory.h" // Performance profiling headers @@ -321,6 +322,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. + pto2_graph_execution_collect_retired(); runtime_destroy(rt, runtime_arena_); rt = nullptr; } 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..3ff277f60f --- /dev/null +++ b/src/a2a3/runtime/host_build_graph/docs/GRAPH_EXECUTION.md @@ -0,0 +1,171 @@ +# Graph Execution + +Graph Execution is available only in the `host_build_graph` runtime. It caches a repeated task DAG and replaces each +later invocation with one Graph task in the host-built task window. The device Scheduler expands that task and submits +the cached nodes to AICore; the host Orchestrator does not replay the nodes one by one. + +## API + +Define a layer as a function that accepts the existing `L2TaskArgs` type: + +```cpp +void layer(const L2TaskArgs &args) { + const Tensor &input = args.tensor(0).ref(); + const Tensor &weight = args.tensor(1).ref(); + const Tensor &output = args.tensor(2).ref(); + + L0TaskArgs task_args; + task_args.add_input(input, weight); + task_args.add_output(output); + task_args.add_scalar(args.scalar(0)); // dynamic: read from this invocation + task_args.add_scalar(uint32_t{16}); // static: stored in the definition + rt_submit_aic_task(FUNC_MATMUL, task_args); +} + +void submit_layer(const L2TaskArgs &args) { + rt_submit_graph(&layer, args); +} +``` + +The upper layer constructs `L2TaskArgs` before calling `submit_layer`; the +wrapper does not repeat tensors or configuration values in its own signature. +The function pointer is the default Graph ID. An explicit stable ID can be +supplied when needed: + +```cpp +rt_submit_graph(PTO2_GRAPH_KEY("decoder_layer_v1"), &layer, args); +``` + +There are no public `GraphBindings`, `Patch`, or `GraphArgs` types. Tensors, configuration values, and dynamic scalars +all enter through `L2TaskArgs`. + +### Qwen decoder-layer integration + +The upper layer packages every value that changes between decoder layers or +decode rounds into `L2TaskArgs`. The Graph function reads those values and +submits the layer's ordinary AIC/AIV tasks. Kernel constants remain literals in +the Graph function and are stored once in the Graph definition. + +The following abbreviated Qwen3-14B layer illustrates the calling pattern. The +real layer may submit attention, normalization, matrix multiplication, and MLP +tasks in any DAG supported by the ordinary task API. + +```cpp +enum QwenArg : int32_t { + HIDDEN_STATES, + ATTN_WEIGHT, + MLP_WEIGHT, + OUTPUT, +}; + +enum QwenScalar : int32_t { + LAYER_ID, + TOKEN_POSITION, +}; + +void qwen_decoder_layer(const L2TaskArgs &args) { + const Tensor &hidden = args.tensor(HIDDEN_STATES).ref(); + const Tensor &attn_weight = args.tensor(ATTN_WEIGHT).ref(); + const Tensor &mlp_weight = args.tensor(MLP_WEIGHT).ref(); + const Tensor &output = args.tensor(OUTPUT).ref(); + + uint32_t intermediate_shape[] = {hidden.shapes[0]}; + TensorCreateInfo intermediate( + intermediate_shape, 1, hidden.dtype + ); + + L0TaskArgs attention_args; + attention_args.add_input(hidden, attn_weight); + attention_args.add_output(intermediate); + attention_args.add_scalar(args.scalar(LAYER_ID)); + attention_args.add_scalar(args.scalar(TOKEN_POSITION)); + attention_args.add_scalar(uint32_t{16}); // fixed head-group size + Tensor attention = + rt_submit_aic_task(FUNC_ATTENTION, attention_args).get_ref(0); + + MixedKernels mlp_kernels; + mlp_kernels.aic_kernel_id = FUNC_MLP_AIC; + mlp_kernels.aiv0_kernel_id = FUNC_MLP_AIV; + + L0TaskArgs mlp_args; + mlp_args.add_input(attention, mlp_weight); + mlp_args.add_output(output); + rt_submit_task(mlp_kernels, mlp_args); +} + +void submit_qwen_decoder_layer(const L2TaskArgs &args) { + rt_submit_graph(&qwen_decoder_layer, args); +} + +void decode_three_layers(const L2TaskArgs layer_args[3]) { + for (uint32_t layer_id = 0; layer_id < 3; ++layer_id) { + submit_qwen_decoder_layer(layer_args[layer_id]); + } +} +``` + +Each `layer_args[i]` contains that layer's hidden state, weights, output, +`layer_id`, and `token_position`. The common function pointer supplies the same +Graph ID for all structurally identical decoder layers. The first layer records +the task DAG. Later layers submit one Graph task because tensor addresses and +scalar values are dynamic. Tensor shape, dtype, stride, size, and direction +must remain structurally compatible; changing them selects a different cache +entry. If a model has decoder-layer variants with different task topology, use +the explicit-key overload with a distinct key for each variant. + +## Record and execute + +On the first call for a structural key, `rt_submit_graph` executes the function normally. The Orchestrator records each +C, V, MIX, or Dummy task and builds a compact definition containing: + +- topological node order and kernel/launch metadata; +- internal fanin counts, fanout adjacency, and root nodes; +- tensor argument sources and per-node packed-output offsets; +- static scalar values and dynamic scalar source indices; +- the external tensor boundary and its input/output directions. + +On a cache hit, the Orchestrator allocates the Graph's combined intermediate heap range, computes only its external +dependencies, and places one `GRAPH` descriptor in the task window. The uploaded submission contains only the compact +definition and current `L2TaskArgs`; it does not contain the expanded node array. The host uploads exactly the used bytes +and immediately returns its temporary block to a bounded pool. This avoids copying the former fixed 6.14 MiB execution +object for every invocation. + +On first observation of the outer task, the Scheduler acquires an AICPU-local execution block and copies the compact +definition into it. It materializes four nodes per scheduling slice while the outer Graph waits for external fanin. Once +both preparation and external readiness are satisfied, it activates the saved roots. Node completion directly releases +the saved fanout list. The outer Graph completes only after all internal nodes complete, so downstream tasks observe it +as one normal dependency producer. + +Host submission blocks and AICPU-local execution blocks use separate bounded pools (16 MiB and 64 blocks each). The +execution pool prefers a block previously used by the same Graph key. That graph-affine path keeps the local definition +and static node fields, then refreshes only task IDs, dynamic tensors/scalars, packed-buffer bases, and scheduling state. +Completed executions are reclaimed after their internal nodes retire; the final Scheduler thread performs a last +collection before runtime shutdown. + +## DFX lanes + +`--enable-l2-swimlane 4` produces five logical lanes in the converted Perfetto JSON: + +- `Host Orchestrator`: one envelope for the host build/submit interval, with + `task_submit(tN)` children for cache-miss nodes and one `graph_submit(tN)` child per cache-hit Graph invocation; +- `Graph Execution`: one envelope per outer Graph task, from its first prepare slice through its last visible node; +- `AICPU Scheduler`: includes `graph_prepare` separately from normal `dispatch`; +- `Scheduler View` and `Worker View`: the existing task scheduling and AICore execution lanes. + +Host and device counters are separate clock domains. The converter logically places Host Orchestrator immediately before +device time zero, then rebases the complete trace so Host Orch starts at trace time zero and every device event follows +it. This preserves ordering without implying clock synchronization or host/device overlap, and avoids Perfetto clipping +the host slice as a negative-time event. Keeping `graph_prepare` out of `dispatch` makes Graph expansion cost visible +without inflating ordinary Scheduler dispatch time. + +## Dynamic arguments and cache keys + +Tensor addresses and scalar values may change between invocations. Tensor shape, dtype, stride, size, direction, and +other structural metadata participate in the cache key; addresses and scalar values do not. A scalar copied from +`args.scalar(i)` records source index `i` and is read from the current invocation during materialization. A literal or +locally computed scalar added directly to `L0TaskArgs` is stored in the definition. + +The current implementation caches up to 16 definitions, 1024 nodes per definition, and 32 boundary tensors/scalars. +Explicit dependencies between nodes inside one Graph are saved in the topology. Nested Graph scopes, explicit +dependencies that cross the Graph boundary, and dispatch predicates are not cached yet; those calls still execute +their function body on the ordinary task-submission path. 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 4d4a59dadc..78450c3423 100644 --- a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp @@ -32,6 +32,7 @@ #include #include #include +#include #include #include @@ -50,6 +51,7 @@ #include "../common/pto_runtime_status.h" #include "../runtime/common.h" #include "../runtime/pto_orchestrator.h" +#include "../runtime/pto_graph_execution.h" #include "../runtime/pto_runtime2.h" #include "../runtime/pto_shared_memory.h" #include "../runtime/pto_types.h" @@ -75,6 +77,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 @@ -478,17 +487,78 @@ static bool relocate_host_orch_image( reloc_ready(rt->scheduler.ready_sync_queues[i]); } reloc_ready(rt->scheduler.dummy_ready_queue); + reloc_ready(rt->scheduler.graph_ready_queue); + reloc_ready(rt->scheduler.graph_prepare_queue); for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { reloc_ready(rt->scheduler.early_dispatch_queues[i]); } return ok; } +static void discard_pending_graph_submissions(PTO2GraphSubmission *submission) { + while (submission != nullptr) { + PTO2GraphSubmission *next = submission->upload_next; + pto2_graph_submission_discard(submission); + submission = next; + } +} + +// Graph definitions live in a process-local cache, while each invocation's +// compact execution image must be device-visible. Upload only the bytes used +// by that invocation, rewrite the single outer Graph descriptor to the device +// image, and return the host block to its bounded pool immediately. +static bool upload_pending_graph_executions( + Runtime *runtime, const HostApi *api, PTO2Runtime *rt, uintptr_t host_sm_begin, size_t host_sm_size, + intptr_t sm_delta +) { + PTO2GraphSubmission *submission = rt->orchestrator.pending_graph_uploads; + rt->orchestrator.pending_graph_uploads = nullptr; + + while (submission != nullptr) { + PTO2GraphSubmission *next = submission->upload_next; + PTO2TaskSlotState *outer_slot = submission->outer_slot; + PTO2TaskDescriptor *outer_task = outer_slot == nullptr ? nullptr : outer_slot->task; + const size_t bytes = pto2_graph_submission_allocation_bytes(submission); + if (outer_task == nullptr || outer_task->kind != PTO2TaskKind::GRAPH || bytes == 0) { + LOG_ERROR("host-orch: invalid pending Graph submission image"); + pto2_graph_submission_discard(submission); + discard_pending_graph_submissions(next); + return false; + } + + void *device_submission = api->device_malloc(bytes); + if (device_submission == nullptr) { + LOG_ERROR("host-orch: failed to allocate %zu bytes for Graph submission", bytes); + pto2_graph_submission_discard(submission); + discard_pending_graph_submissions(next); + return false; + } + + if (!pto2_graph_definition_relocate_for_upload(submission, device_submission) || + !pto2_graph_submission_relocate_for_upload( + submission, device_submission, host_sm_begin, host_sm_size, sm_delta + ) || + api->copy_to_device(device_submission, submission, bytes) != 0) { + LOG_ERROR("host-orch: failed to relocate or upload Graph submission image"); + api->device_free(device_submission); + pto2_graph_submission_discard(submission); + discard_pending_graph_submissions(next); + return false; + } + + outer_task->graph_execution = device_submission; + runtime->tensor_pairs_.push_back({nullptr, device_submission, bytes, false}); + pto2_graph_submission_release_uploaded(submission); + submission = next; + } + 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 + void *host_orch_func_ptr, const L2TaskArgs &orch_l2, bool capture_host_orch ) { std::vector host_sm_buf(sm_size, 0); void *host_sm = host_sm_buf.data(); @@ -526,6 +596,7 @@ int32_t run_host_orchestration( // 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); + rt->active_callable_hash = reinterpret_cast(eps->entry); framework_bind_runtime(rt); if (eps->bind != nullptr) { eps->bind(rt); @@ -534,10 +605,45 @@ int32_t run_host_orchestration( return -1; } + // Host orchestration precedes collector initialization. At level 4 keep + // both the whole-orchestration envelope and one record per task/Graph + // submit. Point the current orchestrator state at the per-run capture + // buffer before entering the user function and retain the populated prefix + // afterwards. + runtime->host_orch_phase_records_.clear(); + runtime->host_orch_start_cycles_ = 0; + runtime->host_orch_end_cycles_ = 0; +#if SIMPLER_DFX + if (capture_host_orch) { + rt->orchestrator.l2_swimlane_level = L2SwimlaneLevel::ORCH_PHASES; + size_t record_capacity = static_cast(eff_task_window_sizes[0]); + if (record_capacity > PLATFORM_PHASE_RECORDS_PER_THREAD) { + record_capacity = PLATFORM_PHASE_RECORDS_PER_THREAD; + } + runtime->host_orch_phase_records_.resize(record_capacity); + rt->orchestrator.host_orch_phase_records = runtime->host_orch_phase_records_.data(); + rt->orchestrator.host_orch_phase_capacity = record_capacity; + rt->orchestrator.host_orch_phase_count = 0; + runtime->host_orch_start_cycles_ = host_prof_cycles(); + } +#endif rt_scope_begin(rt); eps->entry(orch_l2); rt_scope_end(rt); rt_orchestration_done(rt); +#if SIMPLER_DFX + if (capture_host_orch) { + runtime->host_orch_end_cycles_ = host_prof_cycles(); + size_t record_count = rt->orchestrator.host_orch_phase_count; + if (record_count > runtime->host_orch_phase_records_.size()) { + record_count = runtime->host_orch_phase_records_.size(); + } + runtime->host_orch_phase_records_.resize(record_count); + rt->orchestrator.host_orch_phase_records = nullptr; + rt->orchestrator.host_orch_phase_capacity = 0; + rt->orchestrator.host_orch_phase_count = 0; + } +#endif int32_t total_tasks = pto2_sm_layout::ring_current_task_index_addr(host_sm)->load(std::memory_order_acquire); @@ -550,6 +656,9 @@ int32_t run_host_orchestration( 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 (!upload_pending_graph_executions(runtime, api, rt, reinterpret_cast(host_sm), sm_size, sm_delta)) { + return -1; + } 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 @@ -685,7 +794,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, - 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"); @@ -900,7 +1009,8 @@ extern "C" int bind_callable_to_runtime_impl( 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 + eff_task_window_sizes, host_orch_func_ptr, orch_l2, + l2_swimlane_level >= static_cast(L2SwimlaneLevel::ORCH_PHASES) ); if (total_tasks < 0) { LOG_ERROR("host-orch: orchestration run failed"); @@ -1027,7 +1137,9 @@ 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); } } 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 f8321ece0f..7395375f22 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 @@ -34,6 +34,7 @@ // Type headers needed by orchestration #include "common.h" // framework_bind_runtime / framework_current_runtime +#include "pto_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 @@ -86,6 +87,8 @@ typedef struct PTO2RuntimeOps { ); TaskOutputTensors (*alloc_tensors)(PTO2Runtime *rt, const L0TaskArgs &args); TaskOutputTensors (*submit_dummy_task)(PTO2Runtime *rt, const L0TaskArgs &args); + PTO2GraphScopeResult (*graph_begin)(PTO2Runtime *rt, uint64_t graph_key, const L2TaskArgs &args); + void (*graph_end)(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 @@ -203,6 +206,22 @@ static inline TaskOutputTensors rt_submit_dummy_task(const L0TaskArgs &args) { return rt->ops->submit_dummy_task(rt, args); } +static inline PTO2GraphScopeResult rt_graph_begin(uint64_t graph_key, const L2TaskArgs &args) { + PTO2Runtime *rt = current_runtime(); + if (rt->ops->is_fatal(rt) || rt->ops->graph_begin == nullptr) { + return PTO2GraphScopeResult{}; + } + 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_scope_begin(PTO2ScopeMode mode = PTO2ScopeMode::AUTO) { PTO2Runtime *rt = current_runtime(); if (rt->ops->is_fatal(rt)) { @@ -347,6 +366,39 @@ 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 L2TaskArgs. +using PTO2GraphFunction = void (*)(const L2TaskArgs &); + +static inline uint64_t rt_graph_function_id(PTO2GraphFunction function) { + static_assert(sizeof(function) <= sizeof(uint64_t), "Graph function pointer must fit in a 64-bit identity"); + uint64_t function_id = 0; + memcpy(&function_id, &function, sizeof(function)); + return function_id; +} + +static inline PTO2GraphSubmitResult +rt_submit_graph(uint64_t graph_id, PTO2GraphFunction function, const L2TaskArgs &args) { + if (!rt_graph_args_cacheable(args)) { + if (function != nullptr) function(args); + return PTO2GraphSubmitResult{}; + } + PTO2GraphScopeResult result = rt_graph_begin(rt_graph_make_key(graph_id, args), args); + if (result.execute_block && function != nullptr) { + function(args); + } + if (result.recording) { + rt_graph_end(); + } + return result; +} + +static inline PTO2GraphSubmitResult rt_submit_graph(PTO2GraphFunction function, const L2TaskArgs &args) { + return rt_submit_graph(rt_graph_function_id(function), function, args); +} + #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/orchestrator_core/pto_graph_execution.cpp b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_graph_execution.cpp new file mode 100644 index 0000000000..fa921d8527 --- /dev/null +++ b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_graph_execution.cpp @@ -0,0 +1,338 @@ +/* + * 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 "pto_graph_execution.h" + +#include +#include +#include + +namespace { + +PTO2GraphExecution *g_graph_executions = nullptr; +PTO2GraphExecution *g_graph_execution_pool = nullptr; +size_t g_graph_execution_pool_bytes = 0; +int32_t g_graph_execution_pool_blocks = 0; +std::atomic_flag g_graph_execution_lock = ATOMIC_FLAG_INIT; + +PTO2GraphSubmission *g_graph_submission_pool = nullptr; +size_t g_graph_submission_pool_bytes = 0; +int32_t g_graph_submission_pool_blocks = 0; + +constexpr size_t PTO2_GRAPH_EXECUTION_POOL_MAX_BYTES = 16ULL * 1024 * 1024; +constexpr int32_t PTO2_GRAPH_EXECUTION_POOL_MAX_BLOCKS = 64; +constexpr size_t PTO2_GRAPH_SUBMISSION_POOL_MAX_BYTES = 16ULL * 1024 * 1024; +constexpr int32_t PTO2_GRAPH_SUBMISSION_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) return false; + if (static_cast(node_capacity) > std::numeric_limits::max() / sizeof(PTO2GraphNodeStorage)) { + return false; + } + const size_t nodes_bytes = static_cast(node_capacity) * sizeof(PTO2GraphNodeStorage); + if (!checked_align_up(sizeof(PTO2GraphExecution), alignof(PTO2GraphNodeStorage), nodes_offset) || + *nodes_offset > std::numeric_limits::max() - nodes_bytes || + !checked_align_up(*nodes_offset + nodes_bytes, alignof(Tensor), definition_offset) || + *definition_offset > std::numeric_limits::max() - definition_capacity) { + return false; + } + return checked_align_up(*definition_offset + definition_capacity, alignof(PTO2GraphNodeStorage), allocation_bytes); +} + +bool graph_submission_layout(size_t definition_capacity, size_t *definition_offset, size_t *allocation_bytes) { + if (definition_capacity == 0 || + !checked_align_up(sizeof(PTO2GraphSubmission), alignof(Tensor), definition_offset) || + *definition_offset > std::numeric_limits::max() - definition_capacity) { + return false; + } + return checked_align_up(*definition_offset + definition_capacity, alignof(Tensor), allocation_bytes); +} + +void destroy_execution_nodes(PTO2GraphExecution *execution) { + for (int32_t i = 0; i < execution->constructed_nodes; ++i) { + execution->node_storage[i].~PTO2GraphNodeStorage(); + } + execution->constructed_nodes = 0; +} + +void reset_execution(PTO2GraphExecution *execution, int32_t node_count, uint64_t graph_key) { + 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_node_count == node_count && + execution->constructed_nodes >= node_count; + execution->state.store(PTO2GraphExecutionState::SUBMITTED, std::memory_order_relaxed); + execution->activation_requested.store(0, 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->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->graph_definition = nullptr; + execution->topology = PTO2GraphTopologyView{}; + } + execution->args.clear(); + execution->next = nullptr; +} + +void destroy_execution(PTO2GraphExecution *execution) { + if (execution == nullptr) return; + destroy_execution_nodes(execution); + execution->~PTO2GraphExecution(); + std::free(execution); +} + +void recycle_execution_locked(PTO2GraphExecution *execution) { + if (execution == nullptr) return; + const size_t bytes = execution->allocation_bytes; + if (bytes == 0 || bytes > PTO2_GRAPH_EXECUTION_POOL_MAX_BYTES || + g_graph_execution_pool_blocks >= PTO2_GRAPH_EXECUTION_POOL_MAX_BLOCKS || + g_graph_execution_pool_bytes > PTO2_GRAPH_EXECUTION_POOL_MAX_BYTES - bytes) { + destroy_execution(execution); + return; + } + execution->outer_slot = nullptr; + execution->nodes = nullptr; + execution->next = g_graph_execution_pool; + g_graph_execution_pool = execution; + g_graph_execution_pool_bytes += bytes; + g_graph_execution_pool_blocks++; +} + +PTO2GraphExecution *take_pooled_execution_locked(int32_t node_count, uint64_t graph_key, size_t definition_bytes) { + PTO2GraphExecution **best_link = nullptr; + PTO2GraphExecution *best = nullptr; + bool best_is_affine = false; + PTO2GraphExecution **link = &g_graph_execution_pool; + while (*link != nullptr) { + PTO2GraphExecution *candidate = *link; + if (candidate->node_capacity >= node_count && candidate->definition_capacity >= definition_bytes) { + const bool candidate_is_affine = graph_key != 0 && candidate->materialized_graph_key == graph_key && + candidate->materialized_node_count == node_count && + candidate->constructed_nodes >= node_count; + if (best == nullptr || (candidate_is_affine && !best_is_affine) || + (candidate_is_affine == best_is_affine && candidate->allocation_bytes < best->allocation_bytes)) { + best = candidate; + best_link = link; + best_is_affine = candidate_is_affine; + } + } + link = &candidate->next; + } + 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); + return best; +} + +void reset_submission(PTO2GraphSubmission *submission, int32_t node_count, uint64_t graph_key) { + size_t definition_offset = 0; + size_t bytes = 0; + if (!graph_submission_layout(submission->definition_capacity, &definition_offset, &bytes)) return; + submission->local_execution.store(nullptr, std::memory_order_relaxed); + submission->activation_requested.store(0, std::memory_order_relaxed); + submission->node_count = node_count; + submission->allocation_bytes = bytes; + submission->graph_key = graph_key; + submission->outer_slot = nullptr; + submission->definition_storage = reinterpret_cast(submission) + definition_offset; + submission->graph_definition = nullptr; + submission->args.clear(); + submission->next = nullptr; + submission->upload_next = nullptr; +} + +void recycle_submission(PTO2GraphSubmission *submission) { + if (submission == nullptr) return; + const size_t bytes = submission->allocation_bytes; + if (bytes == 0 || bytes > PTO2_GRAPH_SUBMISSION_POOL_MAX_BYTES || + g_graph_submission_pool_blocks >= PTO2_GRAPH_SUBMISSION_POOL_MAX_BLOCKS || + g_graph_submission_pool_bytes > PTO2_GRAPH_SUBMISSION_POOL_MAX_BYTES - bytes) { + submission->~PTO2GraphSubmission(); + std::free(submission); + return; + } + submission->next = g_graph_submission_pool; + submission->upload_next = nullptr; + g_graph_submission_pool = submission; + g_graph_submission_pool_bytes += bytes; + g_graph_submission_pool_blocks++; +} + +PTO2GraphSubmission *take_pooled_submission(int32_t node_count, uint64_t graph_key, size_t definition_bytes) { + PTO2GraphSubmission **best_link = nullptr; + PTO2GraphSubmission *best = nullptr; + PTO2GraphSubmission **link = &g_graph_submission_pool; + while (*link != nullptr) { + PTO2GraphSubmission *candidate = *link; + if (candidate->definition_capacity >= definition_bytes && + (best == nullptr || candidate->allocation_bytes < best->allocation_bytes)) { + best = candidate; + best_link = link; + } + link = &candidate->next; + } + if (best == nullptr) return nullptr; + *best_link = best->next; + g_graph_submission_pool_bytes -= best->allocation_bytes; + g_graph_submission_pool_blocks--; + reset_submission(best, node_count, graph_key); + return best; +} + +template +bool relocate_graph_pointer(T *&ptr, uintptr_t host_begin, size_t bytes, intptr_t delta) { + if (ptr == nullptr) return true; + uintptr_t value = reinterpret_cast(ptr); + if (value < host_begin || value >= host_begin + bytes) return false; + ptr = reinterpret_cast(static_cast(value) + delta); + return true; +} + +} // namespace + +PTO2GraphExecution *pto2_graph_execution_create(int32_t node_count, uint64_t graph_key, size_t definition_bytes) { + if (node_count <= 0 || definition_bytes == 0) return nullptr; + { + GraphExecutionLockGuard guard; + if (PTO2GraphExecution *pooled = take_pooled_execution_locked(node_count, graph_key, 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; + } + void *storage = nullptr; + if (::posix_memalign(&storage, alignof(PTO2GraphNodeStorage), allocation_bytes) != 0) return nullptr; + auto *execution = new (storage) PTO2GraphExecution{}; + execution->node_capacity = node_count; + execution->definition_capacity = definition_bytes; + execution->allocation_bytes = allocation_bytes; + reset_execution(execution, node_count, graph_key); + return execution; +} + +void pto2_graph_execution_discard(PTO2GraphExecution *execution) { + GraphExecutionLockGuard guard; + recycle_execution_locked(execution); +} + +void pto2_graph_execution_publish(PTO2GraphExecution *execution) { + if (execution == nullptr) return; + GraphExecutionLockGuard guard; + execution->next = g_graph_executions; + g_graph_executions = execution; +} + +void pto2_graph_execution_collect_retired() { + GraphExecutionLockGuard guard; + PTO2GraphExecution **link = &g_graph_executions; + while (*link != nullptr) { + PTO2GraphExecution *execution = *link; + const bool completed = execution->state.load(std::memory_order_acquire) == PTO2GraphExecutionState::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); + } +} + +PTO2GraphSubmission *pto2_graph_submission_create(int32_t node_count, uint64_t graph_key, size_t definition_bytes) { + if (node_count <= 0 || definition_bytes == 0) return nullptr; + if (PTO2GraphSubmission *pooled = take_pooled_submission(node_count, graph_key, definition_bytes)) return pooled; + + size_t definition_offset = 0; + size_t allocation_bytes = 0; + if (!graph_submission_layout(definition_bytes, &definition_offset, &allocation_bytes)) return nullptr; + void *storage = nullptr; + if (::posix_memalign(&storage, alignof(Tensor), allocation_bytes) != 0) return nullptr; + auto *submission = new (storage) PTO2GraphSubmission{}; + submission->definition_capacity = definition_bytes; + submission->allocation_bytes = allocation_bytes; + reset_submission(submission, node_count, graph_key); + return submission; +} + +void pto2_graph_submission_discard(PTO2GraphSubmission *submission) { recycle_submission(submission); } + +void pto2_graph_submission_release_uploaded(PTO2GraphSubmission *submission) { recycle_submission(submission); } + +size_t pto2_graph_submission_allocation_bytes(const PTO2GraphSubmission *submission) { + return submission == nullptr ? 0 : submission->allocation_bytes; +} + +bool pto2_graph_submission_relocate_for_upload( + PTO2GraphSubmission *submission, void *device_base, uintptr_t host_sm_begin, size_t host_sm_size, intptr_t sm_delta +) { + if (submission == nullptr || device_base == nullptr || submission->allocation_bytes == 0) return false; + const uintptr_t host_begin = reinterpret_cast(submission); + const intptr_t graph_delta = + static_cast(reinterpret_cast(device_base)) - static_cast(host_begin); + + if (submission->outer_slot != nullptr) { + const uintptr_t outer = reinterpret_cast(submission->outer_slot); + if (outer < host_sm_begin || outer >= host_sm_begin + host_sm_size) return false; + submission->outer_slot = reinterpret_cast(static_cast(outer) + sm_delta); + } + if (!relocate_graph_pointer( + submission->definition_storage, host_begin, submission->allocation_bytes, graph_delta + )) { + return false; + } + if (submission->graph_definition != nullptr) { + auto *definition = const_cast(submission->graph_definition); + if (!relocate_graph_pointer(definition, host_begin, submission->allocation_bytes, graph_delta)) return false; + submission->graph_definition = definition; + } + submission->local_execution.store(nullptr, std::memory_order_relaxed); + submission->next = nullptr; + submission->upload_next = nullptr; + return true; +} 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 235b7e9cbd..39139e1a59 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,11 +27,17 @@ #include #include +#include +#include +#include +#include + #include "aicpu/dep_gen_collector_aicpu.h" #include "common/dep_gen.h" #include "common/platform_config.h" #include "common/unified_log.h" #include "pto_dep_compute.h" +#include "pto_graph_execution.h" #include "pto_runtime2_types.h" #include "pto_shared_memory.h" #include "pto_tensormap.h" @@ -88,6 +94,26 @@ __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; + } + auto &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" @@ -118,7 +144,12 @@ __attribute__((weak, visibility("hidden"))) uint64_t get_sys_cnt_aicpu() { // The strong symbol from the AICPU build wins when profiling is available. // Also hidden to prevent HOST .so from polluting the global symbol table. __attribute__((weak, visibility("hidden"))) void -l2_swimlane_aicpu_record_orch_phase(uint64_t, uint64_t, uint64_t, uint32_t) {} +l2_swimlane_aicpu_record_orch_phase(uint64_t start, uint64_t end, uint64_t task_id, uint32_t submit_idx) { + (void)start; + (void)end; + (void)task_id; + (void)submit_idx; +} // Accumulated cycles per sub-step (only needed for ORCH_PROFILING export) static uint64_t g_orch_sync_cycle = 0; // tensormap sync static uint64_t g_orch_alloc_cycle = 0; // unified task+heap alloc @@ -154,11 +185,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" @@ -176,7 +207,12 @@ __attribute__((weak, visibility("hidden"))) uint64_t get_sys_cnt_aicpu() { static_cast(ts.tv_nsec) * PLATFORM_PROF_SYS_CNT_FREQ / 1000000000ull; } __attribute__((weak, visibility("hidden"))) void -l2_swimlane_aicpu_record_orch_phase(uint64_t, uint64_t, uint64_t, uint32_t) {} +l2_swimlane_aicpu_record_orch_phase(uint64_t start, uint64_t end, uint64_t task_id, uint32_t submit_idx) { + (void)start; + (void)end; + (void)task_id; + (void)submit_idx; +} // submit_idx needed for swimlane task_id tagging (no cycle accumulation at this level) static uint32_t g_orch_submit_idx = 0; #define CYCLE_COUNT_START() \ @@ -186,12 +222,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() @@ -252,6 +288,437 @@ void PTO2OrchestratorState::report_fatal(int32_t error_code, const char *func, c va_end(args); } +namespace { + +constexpr int32_t PTO2_GRAPH_DEFINITION_CAP = 16; +constexpr int32_t PTO2_GRAPH_MAX_TASKS = 1024; +constexpr int32_t PTO2_GRAPH_MAX_FANIN_PER_TASK = 128; +constexpr int32_t PTO2_GRAPH_MAX_INTERNAL_EDGES = PTO2_GRAPH_MAX_TASKS * PTO2_GRAPH_MAX_FANIN_PER_TASK; + +enum class PTO2GraphTensorSource : uint8_t { + BOUNDARY_EXACT = 0, + BOUNDARY_VIEW = 1, + INTERNAL = 2, + OWN_OUTPUT = 3, +}; + +enum class PTO2GraphFaninSource : uint8_t { + INTERNAL = 0, + EXTERNAL_LOCAL_DELTA = 1, +}; + +struct PTO2GraphTensorSourceRef { + PTO2GraphTensorSource source{PTO2GraphTensorSource::BOUNDARY_EXACT}; + uint16_t source_index{0}; + uint64_t packed_offset{0}; +}; + +struct PTO2GraphFaninRef { + PTO2GraphFaninSource source{PTO2GraphFaninSource::INTERNAL}; + int32_t value{0}; +}; + +// Capture-only representation. Its fixed arrays make recording allocation-free +// but are never copied wholesale into the process-local Definition cache. +struct PTO2GraphRecordedNode { + int32_t kernel_id[PTO2_SUBTASK_SLOT_COUNT]{INVALID_KERNEL_ID, INVALID_KERNEL_ID, INVALID_KERNEL_ID}; + ActiveMask active_mask{}; + int16_t logical_block_num{1}; + int16_t total_required_subtasks{0}; + int32_t task_timing_slot{TASK_TIMING_SLOT_NONE}; + int32_t tensor_count{0}; + int32_t scalar_count{0}; + int32_t total_output_size{0}; + uint64_t record_packed_base{0}; + Tensor tensors[MAX_TENSOR_ARGS]; + PTO2GraphTensorSourceRef tensor_sources[MAX_TENSOR_ARGS]; + TensorArgType tensor_arg_types[MAX_TENSOR_ARGS]; + uint64_t scalars[MAX_SCALAR_ARGS]; + uint16_t scalar_source_indices[MAX_SCALAR_ARGS]; + int32_t fanin_count{0}; + PTO2GraphFaninRef fanins[PTO2_GRAPH_MAX_FANIN_PER_TASK]; +}; + +static_assert(std::is_trivially_copyable_v, "Graph capture nodes must remain byte-copyable"); + +// Cache/runtime representation. Variable-length tensors, sources, scalars, +// and scalar-source indices live in Definition-wide packed arrays. +struct PTO2GraphNodeDefinition { + int32_t kernel_id[PTO2_SUBTASK_SLOT_COUNT]; + ActiveMask active_mask; + int16_t logical_block_num; + int16_t total_required_subtasks; + int32_t task_timing_slot; + int32_t tensor_count; + int32_t scalar_count; + int32_t total_output_size; + uint32_t tensor_offset; + uint32_t scalar_offset; +}; + +static_assert( + std::is_trivially_copyable_v, "Compact Graph Definition nodes must remain byte-copyable" +); + +struct PTO2GraphReplayPlan { + uint64_t required_heap{0}; + int32_t boundary_count{0}; + uint32_t required_tensor_count{0}; + uint32_t required_scalar_count{0}; + uint16_t boundary_indices[PTO2_GRAPH_MAX_TENSOR_ARGS]{}; + TensorArgType boundary_types[PTO2_GRAPH_MAX_TENSOR_ARGS]{}; +}; + +struct PTO2GraphRecordingDefinition { + bool in_use{false}; + uint64_t full_key{0}; + int32_t task_count{0}; + int32_t edge_count{0}; + int32_t root_count{0}; + PTO2GraphReplayPlan replay_plan; + uint32_t fanout_offsets[PTO2_GRAPH_MAX_TASKS + 1]{}; + uint16_t fanout_indices[PTO2_GRAPH_MAX_INTERNAL_EDGES]{}; + uint16_t fanin_counts[PTO2_GRAPH_MAX_TASKS]{}; + uint16_t root_indices[PTO2_GRAPH_MAX_TASKS]{}; + uint64_t node_offsets[PTO2_GRAPH_MAX_TASKS]{}; + PTO2GraphRecordedNode tasks[PTO2_GRAPH_MAX_TASKS]; +}; + +struct PTO2GraphDefinition { + bool in_use{false}; + uint64_t full_key{0}; + int32_t task_count{0}; + int32_t edge_count{0}; + int32_t root_count{0}; + PTO2GraphReplayPlan replay_plan; + size_t storage_bytes{0}; + void *storage{nullptr}; + uint32_t *fanout_offsets{nullptr}; + uint16_t *fanout_indices{nullptr}; + uint16_t *fanin_counts{nullptr}; + uint16_t *root_indices{nullptr}; + uint64_t *node_offsets{nullptr}; + PTO2GraphNodeDefinition *tasks{nullptr}; + Tensor *tensors{nullptr}; + PTO2GraphTensorSourceRef *tensor_sources{nullptr}; + uint64_t *scalars{nullptr}; + uint16_t *scalar_source_indices{nullptr}; +}; + +struct PTO2GraphRecordingState { + bool active{false}; + bool pending_finalize{false}; + bool unsupported{false}; + int32_t unsupported_reason{0}; + int32_t unsupported_task_index{-1}; + int32_t unsupported_tensor_index{-1}; + uint64_t full_key{0}; + int32_t start_local_task_id{0}; + int32_t current_task_index{-1}; + int32_t current_fanin_count{0}; + PTO2GraphFaninRef current_fanins[PTO2_GRAPH_MAX_FANIN_PER_TASK]; + ChipStorageTaskArgs args; + PTO2GraphRecordingDefinition temp; +}; + +// Definitions survive repeated simpler_run calls while the AICPU runtime DSO +// remains loaded for the DeviceRunner lifetime. +PTO2GraphDefinition g_graph_definitions[PTO2_GRAPH_DEFINITION_CAP]; +int32_t g_graph_next_replace = 0; +PTO2GraphRecordingState g_graph_recording; + +enum PTO2GraphUnsupportedReason : int32_t { + PTO2_GRAPH_UNSUPPORTED_NONE = 0, + PTO2_GRAPH_UNSUPPORTED_TASK_WINDOW = 1, + PTO2_GRAPH_UNSUPPORTED_NULL_PRODUCER = 2, + PTO2_GRAPH_UNSUPPORTED_EXTERNAL_PRODUCER = 3, + PTO2_GRAPH_UNSUPPORTED_FANIN_OVERFLOW = 4, + PTO2_GRAPH_UNSUPPORTED_RECORD_TASK_ORDER = 5, + PTO2_GRAPH_UNSUPPORTED_RECORD_TASK_NULL = 6, + PTO2_GRAPH_UNSUPPORTED_ARG_OVERFLOW = 7, + PTO2_GRAPH_UNSUPPORTED_TENSOR_SOURCE = 8, + PTO2_GRAPH_UNSUPPORTED_NESTED_SCOPE = 9, + PTO2_GRAPH_UNSUPPORTED_EXTERNAL_EXPLICIT_DEP = 10, + PTO2_GRAPH_UNSUPPORTED_PREDICATE = 11, +}; + +void graph_mark_unsupported(PTO2GraphUnsupportedReason reason, int32_t task_index = -1, int32_t tensor_index = -1) { + if (!g_graph_recording.unsupported) { + g_graph_recording.unsupported_reason = static_cast(reason); + g_graph_recording.unsupported_task_index = task_index; + g_graph_recording.unsupported_tensor_index = tensor_index; + } + g_graph_recording.unsupported = true; +} + +void reset_graph_definition_header(PTO2GraphRecordingDefinition *templ) { + templ->in_use = false; + templ->full_key = 0; + templ->task_count = 0; + templ->edge_count = 0; + templ->root_count = 0; + templ->replay_plan = PTO2GraphReplayPlan{}; +} + +void reset_graph_recording() { + g_graph_recording.active = false; + g_graph_recording.pending_finalize = false; + g_graph_recording.unsupported = false; + g_graph_recording.unsupported_reason = PTO2_GRAPH_UNSUPPORTED_NONE; + g_graph_recording.unsupported_task_index = -1; + g_graph_recording.unsupported_tensor_index = -1; + g_graph_recording.full_key = 0; + g_graph_recording.start_local_task_id = 0; + g_graph_recording.current_task_index = -1; + g_graph_recording.current_fanin_count = 0; + g_graph_recording.args.clear(); + reset_graph_definition_header(&g_graph_recording.temp); +} + +uint64_t graph_full_key(uint64_t callable_hash, uint64_t graph_key) { + uint64_t h = 1469598103934665603ULL; + h = pto2_graph_hash_bytes(h, &PTO2_GRAPH_CACHE_SCHEMA_VERSION, sizeof(PTO2_GRAPH_CACHE_SCHEMA_VERSION)); + h = pto2_graph_hash_bytes(h, &callable_hash, sizeof(callable_hash)); + return pto2_graph_hash_bytes(h, &graph_key, sizeof(graph_key)); +} + +bool graph_snapshot_args(const L2TaskArgs &source, ChipStorageTaskArgs *snapshot) { + if (snapshot == nullptr || source.has_error || + source.tensor_count() > static_cast(PTO2_GRAPH_MAX_TENSOR_ARGS) || + source.scalar_count() > static_cast(PTO2_GRAPH_MAX_SCALAR_ARGS)) { + return false; + } + snapshot->clear(); + for (int32_t i = 0; i < source.tensor_count(); ++i) { + if (source.tag(i) == TensorArgType::OUTPUT) return false; + snapshot->add_tensor(source.tensor(i).ref()); + } + for (int32_t i = 0; i < source.scalar_count(); ++i) { + snapshot->add_scalar(source.scalar(i)); + } + return true; +} + +PTO2GraphDefinition *find_graph_definition(uint64_t full_key) { + for (int32_t i = 0; i < PTO2_GRAPH_DEFINITION_CAP; ++i) { + if (g_graph_definitions[i].in_use && g_graph_definitions[i].full_key == full_key) { + return &g_graph_definitions[i]; + } + } + return nullptr; +} + +bool graph_definition_append_storage( + size_t *cursor, size_t count, size_t element_size, size_t alignment, size_t *offset +) { + if (cursor == nullptr || offset == nullptr || alignment == 0 || (alignment & (alignment - 1)) != 0) return false; + if (*cursor > SIZE_MAX - (alignment - 1)) return false; + size_t aligned = (*cursor + alignment - 1) & ~(alignment - 1); + if (count > 0 && element_size > SIZE_MAX / count) return false; + size_t bytes = count * element_size; + if (aligned > SIZE_MAX - bytes) return false; + *offset = aligned; + *cursor = aligned + bytes; + return true; +} + +void release_graph_definition(PTO2GraphDefinition *definition) { + if (definition == nullptr) return; + free(definition->storage); + *definition = PTO2GraphDefinition{}; +} + +bool compact_graph_definition(PTO2GraphDefinition *definition, const PTO2GraphRecordingDefinition &recorded) { + if (definition == nullptr || recorded.task_count <= 0 || recorded.edge_count < 0 || recorded.root_count < 0) { + return false; + } + + size_t tensor_arg_count = 0; + size_t scalar_arg_count = 0; + for (int32_t i = 0; i < recorded.task_count; ++i) { + const PTO2GraphRecordedNode &node = recorded.tasks[i]; + if (node.tensor_count < 0 || node.tensor_count > MAX_TENSOR_ARGS || node.scalar_count < 0 || + node.scalar_count > MAX_SCALAR_ARGS) { + return false; + } + tensor_arg_count += static_cast(node.tensor_count); + scalar_arg_count += static_cast(node.scalar_count); + } + if (tensor_arg_count > UINT32_MAX || scalar_arg_count > UINT32_MAX) return false; + + size_t cursor = 0; + size_t fanout_offsets_offset = 0; + size_t fanout_indices_offset = 0; + size_t fanin_counts_offset = 0; + size_t root_indices_offset = 0; + size_t node_offsets_offset = 0; + size_t tasks_offset = 0; + size_t tensors_offset = 0; + size_t tensor_sources_offset = 0; + size_t scalars_offset = 0; + size_t scalar_source_indices_offset = 0; + if (!graph_definition_append_storage( + &cursor, static_cast(recorded.task_count) + 1, sizeof(uint32_t), alignof(uint32_t), + &fanout_offsets_offset + ) || + !graph_definition_append_storage( + &cursor, static_cast(recorded.edge_count), sizeof(uint16_t), alignof(uint16_t), + &fanout_indices_offset + ) || + !graph_definition_append_storage( + &cursor, static_cast(recorded.task_count), sizeof(uint16_t), alignof(uint16_t), &fanin_counts_offset + ) || + !graph_definition_append_storage( + &cursor, static_cast(recorded.root_count), sizeof(uint16_t), alignof(uint16_t), &root_indices_offset + ) || + !graph_definition_append_storage( + &cursor, static_cast(recorded.task_count), sizeof(uint64_t), alignof(uint64_t), &node_offsets_offset + ) || + !graph_definition_append_storage( + &cursor, static_cast(recorded.task_count), sizeof(PTO2GraphNodeDefinition), + alignof(PTO2GraphNodeDefinition), &tasks_offset + ) || + !graph_definition_append_storage(&cursor, tensor_arg_count, sizeof(Tensor), alignof(Tensor), &tensors_offset) || + !graph_definition_append_storage( + &cursor, tensor_arg_count, sizeof(PTO2GraphTensorSourceRef), alignof(PTO2GraphTensorSourceRef), + &tensor_sources_offset + ) || + !graph_definition_append_storage( + &cursor, scalar_arg_count, sizeof(uint64_t), alignof(uint64_t), &scalars_offset + ) || + !graph_definition_append_storage( + &cursor, scalar_arg_count, sizeof(uint16_t), alignof(uint16_t), &scalar_source_indices_offset + )) { + return false; + } + + void *storage = nullptr; + if (posix_memalign(&storage, alignof(Tensor), cursor) != 0) return false; + auto *base = static_cast(storage); + + PTO2GraphDefinition compact; + compact.in_use = true; + compact.full_key = recorded.full_key; + compact.task_count = recorded.task_count; + compact.edge_count = recorded.edge_count; + compact.root_count = recorded.root_count; + compact.replay_plan = recorded.replay_plan; + compact.storage_bytes = cursor; + compact.storage = storage; + compact.fanout_offsets = reinterpret_cast(base + fanout_offsets_offset); + compact.fanout_indices = reinterpret_cast(base + fanout_indices_offset); + compact.fanin_counts = reinterpret_cast(base + fanin_counts_offset); + compact.root_indices = reinterpret_cast(base + root_indices_offset); + compact.node_offsets = reinterpret_cast(base + node_offsets_offset); + compact.tasks = reinterpret_cast(base + tasks_offset); + compact.tensors = reinterpret_cast(base + tensors_offset); + compact.tensor_sources = reinterpret_cast(base + tensor_sources_offset); + compact.scalars = reinterpret_cast(base + scalars_offset); + compact.scalar_source_indices = reinterpret_cast(base + scalar_source_indices_offset); + + memcpy( + compact.fanout_offsets, recorded.fanout_offsets, + sizeof(uint32_t) * (static_cast(recorded.task_count) + 1) + ); + memcpy( + compact.fanout_indices, recorded.fanout_indices, sizeof(uint16_t) * static_cast(recorded.edge_count) + ); + memcpy(compact.fanin_counts, recorded.fanin_counts, sizeof(uint16_t) * static_cast(recorded.task_count)); + memcpy(compact.root_indices, recorded.root_indices, sizeof(uint16_t) * static_cast(recorded.root_count)); + memcpy(compact.node_offsets, recorded.node_offsets, sizeof(uint64_t) * static_cast(recorded.task_count)); + + uint32_t tensor_cursor = 0; + uint32_t scalar_cursor = 0; + for (int32_t i = 0; i < recorded.task_count; ++i) { + const PTO2GraphRecordedNode &src = recorded.tasks[i]; + PTO2GraphNodeDefinition &dst = compact.tasks[i]; + for (int32_t k = 0; k < PTO2_SUBTASK_SLOT_COUNT; ++k) + dst.kernel_id[k] = src.kernel_id[k]; + dst.active_mask = src.active_mask; + dst.logical_block_num = src.logical_block_num; + dst.total_required_subtasks = src.total_required_subtasks; + dst.task_timing_slot = src.task_timing_slot; + dst.tensor_count = src.tensor_count; + dst.scalar_count = src.scalar_count; + dst.total_output_size = src.total_output_size; + dst.tensor_offset = tensor_cursor; + dst.scalar_offset = scalar_cursor; + + memcpy(compact.tensors + tensor_cursor, src.tensors, sizeof(Tensor) * static_cast(src.tensor_count)); + memcpy( + compact.tensor_sources + tensor_cursor, src.tensor_sources, + sizeof(PTO2GraphTensorSourceRef) * static_cast(src.tensor_count) + ); + memcpy(compact.scalars + scalar_cursor, src.scalars, sizeof(uint64_t) * static_cast(src.scalar_count)); + memcpy( + compact.scalar_source_indices + scalar_cursor, src.scalar_source_indices, + sizeof(uint16_t) * static_cast(src.scalar_count) + ); + tensor_cursor += static_cast(src.tensor_count); + scalar_cursor += static_cast(src.scalar_count); + } + + release_graph_definition(definition); + *definition = compact; + return true; +} + +bool store_graph_definition(const PTO2GraphRecordingDefinition &templ) { + int32_t target = -1; + for (int32_t i = 0; i < PTO2_GRAPH_DEFINITION_CAP; ++i) { + if (!g_graph_definitions[i].in_use) { + target = i; + break; + } + } + if (target < 0) { + target = g_graph_next_replace; + g_graph_next_replace = (g_graph_next_replace + 1) % PTO2_GRAPH_DEFINITION_CAP; + } + if (target < 0) return false; + if (!compact_graph_definition(&g_graph_definitions[target], templ)) return false; + return true; +} + +void graph_record_begin_task(PTO2TaskId task_id) { + if (!g_graph_recording.active || g_graph_recording.unsupported) return; + int32_t idx = static_cast(task_id.local()) - g_graph_recording.start_local_task_id; + if (idx < 0 || idx >= PTO2_GRAPH_MAX_TASKS || idx != g_graph_recording.temp.task_count) { + graph_mark_unsupported(PTO2_GRAPH_UNSUPPORTED_TASK_WINDOW, idx); + return; + } + g_graph_recording.current_task_index = idx; + g_graph_recording.current_fanin_count = 0; +} + +void graph_record_note_fanin(PTO2TaskSlotState *producer) { + if (!g_graph_recording.active || g_graph_recording.unsupported) return; + if (producer == nullptr || producer->task == nullptr) { + graph_mark_unsupported(PTO2_GRAPH_UNSUPPORTED_NULL_PRODUCER, g_graph_recording.current_task_index); + return; + } + int32_t producer_local = static_cast(producer->task->task_id.local()); + int32_t producer_index = producer_local - g_graph_recording.start_local_task_id; + if (producer_index >= g_graph_recording.current_task_index) { + graph_mark_unsupported(PTO2_GRAPH_UNSUPPORTED_EXTERNAL_PRODUCER, g_graph_recording.current_task_index); + return; + } + if (g_graph_recording.current_fanin_count >= PTO2_GRAPH_MAX_FANIN_PER_TASK) { + graph_mark_unsupported(PTO2_GRAPH_UNSUPPORTED_FANIN_OVERFLOW, g_graph_recording.current_task_index); + return; + } + PTO2GraphFaninRef &fanin = g_graph_recording.current_fanins[g_graph_recording.current_fanin_count++]; + if (producer_index >= 0) { + fanin.source = PTO2GraphFaninSource::INTERNAL; + fanin.value = producer_index; + } else { + fanin.source = PTO2GraphFaninSource::EXTERNAL_LOCAL_DELTA; + fanin.value = producer_local - g_graph_recording.start_local_task_id; + } +} + +} // namespace + static uint32_t next_fanin_seen_epoch(PTO2OrchestratorState *orch) { uint32_t next = orch->fanin_seen_current_epoch + 1; if (next == 0) { @@ -353,6 +820,8 @@ static bool append_fanin_or_fail( return true; } + graph_record_note_fanin(prod_state); + if (fanin_builder->count < PTO2_FANIN_INLINE_CAP) { fanin_builder->inline_slots[fanin_builder->count++] = prod_state; return true; @@ -472,6 +941,670 @@ struct PTO2PreparedTask { PTO2TaskSlotState *slot_state = nullptr; }; +namespace { + +static bool graph_tensor_metadata_equal(const Tensor &lhs, const Tensor &rhs) { + if (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 memcmp(lhs.shapes, rhs.shapes, sizeof(uint32_t) * lhs.ndims) == 0 && + memcmp(lhs.strides, rhs.strides, sizeof(uint32_t) * lhs.ndims) == 0; +} + +static bool +graph_tensor_from_args(const Tensor &tensor, const ChipStorageTaskArgs &args, PTO2GraphTensorSourceRef *source_ref) { + for (int32_t i = 0; i < args.tensor_count(); ++i) { + if (graph_tensor_metadata_equal(tensor, args.tensor(i))) { + source_ref->source = PTO2GraphTensorSource::BOUNDARY_EXACT; + source_ref->source_index = static_cast(i); + source_ref->packed_offset = 0; + return true; + } + } + uint16_t best_index = UINT16_MAX; + uint64_t best_offset = UINT64_MAX; + for (int32_t i = 0; i < args.tensor_count(); ++i) { + const Tensor &boundary = args.tensor(i); + if (tensor.buffer.addr == boundary.buffer.addr && tensor.buffer.size == boundary.buffer.size && + tensor.start_offset >= boundary.start_offset) { + uint64_t offset = tensor.start_offset - boundary.start_offset; + if (offset < best_offset) { + best_index = static_cast(i); + best_offset = offset; + } + } + } + if (best_index == UINT16_MAX) return false; + source_ref->source = PTO2GraphTensorSource::BOUNDARY_VIEW; + source_ref->source_index = best_index; + source_ref->packed_offset = best_offset; + return true; +} + +static bool +graph_classify_tensor(PTO2GraphRecordedNode *task, int32_t tensor_index, const Tensor &tensor, int32_t task_index) { + PTO2GraphTensorSourceRef &source_ref = task->tensor_sources[tensor_index]; + if (graph_tensor_from_args(tensor, g_graph_recording.args, &source_ref)) return true; + + uint64_t tensor_addr = tensor.buffer.addr; + for (int32_t producer_index = task_index; producer_index >= 0; --producer_index) { + const PTO2GraphRecordedNode &producer = + (producer_index == task_index) ? *task : g_graph_recording.temp.tasks[producer_index]; + if (producer.record_packed_base == 0 || producer.total_output_size <= 0) continue; + uint64_t producer_begin = producer.record_packed_base; + uint64_t producer_end = producer_begin + static_cast(producer.total_output_size); + if (tensor_addr < producer_begin || tensor_addr >= producer_end) continue; + source_ref.source = + producer_index == task_index ? PTO2GraphTensorSource::OWN_OUTPUT : PTO2GraphTensorSource::INTERNAL; + source_ref.source_index = static_cast(producer_index); + source_ref.packed_offset = tensor_addr - producer_begin; + return true; + } + return false; +} + +static void graph_record_task(const PTO2PreparedTask &prepared, const L0TaskArgs &args) { + if (!g_graph_recording.active || g_graph_recording.unsupported) return; + int32_t task_index = static_cast(prepared.task_id.local()) - g_graph_recording.start_local_task_id; + if (task_index < 0 || task_index >= PTO2_GRAPH_MAX_TASKS || task_index != g_graph_recording.temp.task_count) { + graph_mark_unsupported(PTO2_GRAPH_UNSUPPORTED_RECORD_TASK_ORDER, task_index); + return; + } + if (prepared.task == nullptr || prepared.payload == nullptr || prepared.slot_state == nullptr) { + graph_mark_unsupported(PTO2_GRAPH_UNSUPPORTED_RECORD_TASK_NULL, task_index); + return; + } + + PTO2GraphRecordedNode &task = g_graph_recording.temp.tasks[task_index]; + task = PTO2GraphRecordedNode{}; + for (int i = 0; i < PTO2_SUBTASK_SLOT_COUNT; ++i) + task.kernel_id[i] = prepared.task->kernel_id[i]; + task.active_mask = prepared.slot_state->active_mask; + task.logical_block_num = prepared.slot_state->logical_block_num; + task.total_required_subtasks = prepared.slot_state->total_required_subtasks; + task.task_timing_slot = prepared.task->task_timing_slot; + task.tensor_count = prepared.payload->tensor_count; + task.scalar_count = prepared.payload->scalar_count; + task.total_output_size = static_cast( + reinterpret_cast(prepared.task->packed_buffer_end) - + reinterpret_cast(prepared.task->packed_buffer_base) + ); + task.record_packed_base = reinterpret_cast(prepared.task->packed_buffer_base); + if (task.tensor_count < 0 || task.tensor_count > MAX_TENSOR_ARGS || task.scalar_count < 0 || + task.scalar_count > MAX_SCALAR_ARGS) { + graph_mark_unsupported(PTO2_GRAPH_UNSUPPORTED_ARG_OVERFLOW, task_index); + return; + } + + for (int32_t i = 0; i < task.tensor_count; ++i) { + task.tensors[i].copy(prepared.payload->tensors[i]); + task.tensor_arg_types[i] = args.tag(i); + if (!graph_classify_tensor(&task, i, task.tensors[i], task_index)) { + graph_mark_unsupported(PTO2_GRAPH_UNSUPPORTED_TENSOR_SOURCE, task_index, i); + return; + } + } + memcpy(task.scalars, prepared.payload->scalars, sizeof(uint64_t) * static_cast(task.scalar_count)); + for (int32_t i = 0; i < task.scalar_count; ++i) + task.scalar_source_indices[i] = args.scalar_source_index(i); + for (uint32_t i = 0; i < args.explicit_dep_count(); ++i) { + PTO2TaskId dependency = args.explicit_dep(i); + int32_t dependency_index = static_cast(dependency.local()) - g_graph_recording.start_local_task_id; + if (!dependency.is_valid() || dependency.ring() != 0 || dependency_index < 0 || + dependency_index >= task_index) { + graph_mark_unsupported(PTO2_GRAPH_UNSUPPORTED_EXTERNAL_EXPLICIT_DEP, task_index); + return; + } + } + if (args.predicate().op != PredicateOp::NONE) { + graph_mark_unsupported(PTO2_GRAPH_UNSUPPORTED_PREDICATE, task_index); + return; + } + task.fanin_count = g_graph_recording.current_fanin_count; + for (int32_t i = 0; i < task.fanin_count; ++i) + task.fanins[i] = g_graph_recording.current_fanins[i]; + + g_graph_recording.temp.task_count++; + g_graph_recording.current_task_index = -1; +} + +static void graph_reset_payload(PTO2TaskPayload *payload) { + 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); +} + +static TensorArgType graph_boundary_type(bool reads, bool writes, bool no_dep) { + if (reads && writes) return TensorArgType::INOUT; + if (writes) return TensorArgType::OUTPUT_EXISTING; + if (reads) return TensorArgType::INPUT; + return no_dep ? TensorArgType::NO_DEP : TensorArgType::INPUT; +} + +static bool graph_build_definition(PTO2GraphRecordingDefinition *templ, const ChipStorageTaskArgs &args) { + if (templ == nullptr || templ->task_count <= 0 || templ->task_count > PTO2_GRAPH_MAX_TASKS) return false; + PTO2GraphReplayPlan plan; + bool boundary_seen[PTO2_GRAPH_MAX_TENSOR_ARGS]{}; + bool boundary_reads[PTO2_GRAPH_MAX_TENSOR_ARGS]{}; + bool boundary_writes[PTO2_GRAPH_MAX_TENSOR_ARGS]{}; + bool boundary_no_dep[PTO2_GRAPH_MAX_TENSOR_ARGS]{}; + uint32_t fanout_counts[PTO2_GRAPH_MAX_TASKS]{}; + + for (int32_t i = 0; i < templ->task_count; ++i) { + const PTO2GraphRecordedNode &task = templ->tasks[i]; + if (task.tensor_count < 0 || task.tensor_count > MAX_TENSOR_ARGS || task.scalar_count < 0 || + task.scalar_count > MAX_SCALAR_ARGS || task.fanin_count < 0 || + task.fanin_count > PTO2_GRAPH_MAX_FANIN_PER_TASK || task.total_output_size < 0) { + return false; + } + templ->node_offsets[i] = plan.required_heap; + uint64_t aligned_output = PTO2_ALIGN_UP(static_cast(task.total_output_size), PTO2_ALIGN_SIZE); + if (plan.required_heap > UINT64_MAX - aligned_output) return false; + plan.required_heap += aligned_output; + + int32_t internal_fanin_count = 0; + for (int32_t e = 0; e < task.fanin_count; ++e) { + const PTO2GraphFaninRef &fanin = task.fanins[e]; + if (fanin.source != PTO2GraphFaninSource::INTERNAL) continue; + if (fanin.value < 0 || fanin.value >= i) return false; + if (templ->edge_count >= PTO2_GRAPH_MAX_INTERNAL_EDGES) return false; + fanout_counts[fanin.value]++; + templ->edge_count++; + internal_fanin_count++; + } + templ->fanin_counts[i] = static_cast(internal_fanin_count); + if (internal_fanin_count == 0) { + templ->root_indices[templ->root_count++] = static_cast(i); + } + + for (int32_t j = 0; j < task.tensor_count; ++j) { + const PTO2GraphTensorSourceRef &source_ref = task.tensor_sources[j]; + bool is_boundary = source_ref.source == PTO2GraphTensorSource::BOUNDARY_EXACT || + source_ref.source == PTO2GraphTensorSource::BOUNDARY_VIEW; + if (is_boundary && source_ref.source_index >= args.tensor_count()) return false; + if (source_ref.source == PTO2GraphTensorSource::INTERNAL && source_ref.source_index >= i) return false; + if (is_boundary) { + uint16_t index = source_ref.source_index; + uint32_t required = static_cast(index) + 1; + if (required > plan.required_tensor_count) plan.required_tensor_count = required; + TensorArgType type = task.tensor_arg_types[j]; + boundary_seen[index] = true; + boundary_reads[index] |= type == TensorArgType::INPUT || type == TensorArgType::INOUT; + boundary_writes[index] |= type == TensorArgType::INOUT || type == TensorArgType::OUTPUT_EXISTING; + boundary_no_dep[index] |= type == TensorArgType::NO_DEP; + } + } + for (int32_t j = 0; j < task.scalar_count; ++j) { + uint16_t source_index = task.scalar_source_indices[j]; + if (source_index != PTO2_TASK_ARG_STATIC && source_index >= args.scalar_count()) return false; + if (source_index != PTO2_TASK_ARG_STATIC) { + uint32_t required = static_cast(source_index) + 1; + if (required > plan.required_scalar_count) plan.required_scalar_count = required; + } + } + } + + for (int32_t i = 0; i < args.tensor_count(); ++i) { + if (!boundary_seen[i]) continue; + TensorArgType type = graph_boundary_type(boundary_reads[i], boundary_writes[i], boundary_no_dep[i]); + plan.boundary_indices[plan.boundary_count] = static_cast(i); + plan.boundary_types[plan.boundary_count] = type; + plan.boundary_count++; + } + + templ->fanout_offsets[0] = 0; + for (int32_t i = 0; i < templ->task_count; ++i) { + templ->fanout_offsets[i + 1] = templ->fanout_offsets[i] + fanout_counts[i]; + } + uint32_t fanout_cursors[PTO2_GRAPH_MAX_TASKS]; + memcpy(fanout_cursors, templ->fanout_offsets, sizeof(uint32_t) * static_cast(templ->task_count)); + for (int32_t consumer_index = 0; consumer_index < templ->task_count; ++consumer_index) { + const PTO2GraphRecordedNode &task = templ->tasks[consumer_index]; + for (int32_t e = 0; e < task.fanin_count; ++e) { + const PTO2GraphFaninRef &fanin = task.fanins[e]; + if (fanin.source != PTO2GraphFaninSource::INTERNAL) continue; + templ->fanout_indices[fanout_cursors[fanin.value]++] = static_cast(consumer_index); + } + } + + templ->replay_plan = plan; + return templ->fanout_offsets[templ->task_count] == static_cast(templ->edge_count); +} + +static bool graph_submission_preflight( + PTO2OrchestratorState *orch, const PTO2GraphDefinition &templ, const ChipStorageTaskArgs &args +) { + const PTO2GraphReplayPlan &plan = templ.replay_plan; + if (args.tensor_count() < static_cast(plan.required_tensor_count) || + args.scalar_count() < static_cast(plan.required_scalar_count)) { + return false; + } + auto &allocator = orch->ring.task_allocator; + if (allocator.active_count() + 1 >= allocator.window_size()) return false; + if (plan.required_heap > allocator.heap_available() || plan.required_heap > INT32_MAX) return false; + int32_t tensormap_entries = 0; + for (int32_t i = 0; i < plan.boundary_count; ++i) { + uint16_t arg_index = plan.boundary_indices[i]; + if (arg_index >= args.tensor_count()) return false; + TensorArgType type = plan.boundary_types[i]; + if ((type == TensorArgType::INOUT || type == TensorArgType::OUTPUT_EXISTING) && + !args.tensor(arg_index).manual_dep) { + tensormap_entries++; + } + } + if (tensormap_entries > orch->tensor_map.free_entries()) return false; + return true; +} + +static size_t graph_definition_image_bytes(const PTO2GraphDefinition &templ) { + size_t header_bytes = PTO2_ALIGN_UP(sizeof(PTO2GraphDefinition), alignof(Tensor)); + if (templ.storage_bytes > SIZE_MAX - header_bytes) return 0; + return header_bytes + templ.storage_bytes; +} + +template +static bool graph_rebase_cloned_pointer(T *&ptr, uintptr_t source_begin, size_t source_bytes, intptr_t delta) { + if (ptr == nullptr) return true; + uintptr_t value = reinterpret_cast(ptr); + if (value < source_begin || value >= source_begin + source_bytes) return false; + ptr = reinterpret_cast(static_cast(value) + delta); + return true; +} + +static PTO2GraphDefinition * +graph_clone_definition_image(void *definition_storage, size_t definition_capacity, const PTO2GraphDefinition &templ) { + if (definition_storage == nullptr || templ.storage == nullptr) return nullptr; + const size_t header_bytes = PTO2_ALIGN_UP(sizeof(PTO2GraphDefinition), alignof(Tensor)); + if (definition_capacity < header_bytes + templ.storage_bytes) return nullptr; + + auto *copy = new (definition_storage) PTO2GraphDefinition(templ); + void *copy_storage = static_cast(definition_storage) + header_bytes; + memcpy(copy_storage, templ.storage, templ.storage_bytes); + const uintptr_t source_begin = reinterpret_cast(templ.storage); + const intptr_t delta = + static_cast(reinterpret_cast(copy_storage)) - static_cast(source_begin); + copy->storage = copy_storage; + if (!graph_rebase_cloned_pointer(copy->fanout_offsets, source_begin, templ.storage_bytes, delta) || + !graph_rebase_cloned_pointer(copy->fanout_indices, source_begin, templ.storage_bytes, delta) || + !graph_rebase_cloned_pointer(copy->fanin_counts, source_begin, templ.storage_bytes, delta) || + !graph_rebase_cloned_pointer(copy->root_indices, source_begin, templ.storage_bytes, delta) || + !graph_rebase_cloned_pointer(copy->node_offsets, source_begin, templ.storage_bytes, delta) || + !graph_rebase_cloned_pointer(copy->tasks, source_begin, templ.storage_bytes, delta) || + !graph_rebase_cloned_pointer(copy->tensors, source_begin, templ.storage_bytes, delta) || + !graph_rebase_cloned_pointer(copy->tensor_sources, source_begin, templ.storage_bytes, delta) || + !graph_rebase_cloned_pointer(copy->scalars, source_begin, templ.storage_bytes, delta) || + !graph_rebase_cloned_pointer(copy->scalar_source_indices, source_begin, templ.storage_bytes, delta)) { + return nullptr; + } + return copy; +} + +static bool graph_clone_definition(PTO2GraphSubmission *submission, const PTO2GraphDefinition &templ) { + if (submission == nullptr) return false; + submission->graph_definition = + graph_clone_definition_image(submission->definition_storage, submission->definition_capacity, templ); + return submission->graph_definition != nullptr; +} + +static bool graph_clone_definition(PTO2GraphExecution *execution, const PTO2GraphDefinition &templ) { + if (execution == nullptr) return false; + auto *copy = graph_clone_definition_image(execution->definition_storage, execution->definition_capacity, templ); + if (copy == nullptr) return false; + execution->graph_definition = copy; + execution->topology.edge_count = copy->edge_count; + execution->topology.root_count = copy->root_count; + execution->topology.fanout_offsets = copy->fanout_offsets; + execution->topology.fanout_indices = copy->fanout_indices; + execution->topology.fanin_counts = copy->fanin_counts; + execution->topology.root_indices = copy->root_indices; + return true; +} + +static bool graph_submit_definition( + PTO2OrchestratorState *orch, const PTO2GraphDefinition &templ, const ChipStorageTaskArgs &args, + uint64_t *orch_record_task_id +) { + if (orch_record_task_id != nullptr) *orch_record_task_id = PTO2TaskId::invalid().raw; + if (!graph_submission_preflight(orch, templ, args)) return false; + + const size_t definition_bytes = graph_definition_image_bytes(templ); + if (definition_bytes == 0) return false; + PTO2GraphSubmission *submission = pto2_graph_submission_create(templ.task_count, templ.full_key, definition_bytes); + if (submission == nullptr || !graph_clone_definition(submission, templ)) { + pto2_graph_submission_discard(submission); + return false; + } + submission->args = args; + + const PTO2GraphReplayPlan &plan = templ.replay_plan; + PTO2TaskAllocResult outer_alloc = orch->ring.task_allocator.alloc(static_cast(plan.required_heap)); + if (outer_alloc.failed()) { + pto2_graph_submission_discard(submission); + orch_mark_fatal(orch, PTO2_ERROR_HEAP_RING_DEADLOCK); + return false; + } + + PTO2TaskId outer_task_id = PTO2TaskId::make(0, static_cast(outer_alloc.task_id)); + PTO2SharedMemoryRingHeader &ring = orch->sm_header->ring; + PTO2TaskDescriptor &outer_task = ring.task_descriptors[outer_alloc.slot]; + PTO2TaskPayload &outer_payload = ring.task_payloads[outer_alloc.slot]; + PTO2TaskSlotState &outer_slot = ring.get_slot_state_by_slot(outer_alloc.slot); + outer_task.task_id = outer_task_id; + for (int k = 0; k < PTO2_SUBTASK_SLOT_COUNT; ++k) + outer_task.kernel_id[k] = INVALID_KERNEL_ID; + outer_task.task_timing_slot = TASK_TIMING_SLOT_NONE; + outer_task.packed_buffer_base = outer_alloc.packed_base; + outer_task.packed_buffer_end = outer_alloc.packed_end; + outer_task.graph_execution = submission; + outer_task.graph_node_index = -1; + outer_task.kind = PTO2TaskKind::GRAPH; + outer_payload.tensor_count = 0; + outer_payload.scalar_count = 0; + outer_payload.fanin_actual_count = 0; + outer_payload.fanin_spill_start = 0; + outer_payload.fanin_spill_pool = &orch->ring.fanin_pool; + outer_payload.predicate = DispatchPredicate{}; + graph_reset_payload(&outer_payload); + outer_slot.bind_buffers(&outer_payload, &outer_task); + outer_slot.task_state.store(PTO2_TASK_PENDING, std::memory_order_relaxed); + outer_slot.active_mask = ActiveMask{}; + outer_slot.set_allow_early_resolve(false); + outer_slot.total_required_subtasks = 0; + outer_slot.logical_block_num = 1; + scope_tasks_push(orch, &outer_slot); + + submission->outer_slot = &outer_slot; + + TensorRef boundary_tensors[PTO2_GRAPH_MAX_TENSOR_ARGS]; + TensorArgType boundary_types[PTO2_GRAPH_MAX_TENSOR_ARGS]; + for (int32_t i = 0; i < plan.boundary_count; ++i) { + boundary_tensors[i] = &args.tensor(plan.boundary_indices[i]); + boundary_types[i] = plan.boundary_types[i]; + } + DepInputs boundary_inputs{plan.boundary_count, boundary_tensors, boundary_types, 0, nullptr}; + PTO2FaninBuilder fanin_builder(orch, orch->ring.fanin_pool, next_fanin_seen_epoch(orch)); + orch->tensor_map.sync_tensormap(outer_task_id, ring.fc.last_task_alive.load(std::memory_order_acquire)); + auto boundary_emit = [&](PTO2TaskId producer_task_id) -> bool { + uint8_t producer_ring_id = producer_task_id.ring(); + int32_t producer_local = static_cast(producer_task_id.local()); + 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_ring_id, producer_slot, producer, producer_task_id, &fanin_builder); + }; + if (!compute_task_fanin(boundary_inputs, orch->tensor_map, orch->in_manual_scope(), boundary_emit)) { + pto2_graph_submission_discard(submission); + return false; + } + register_task_outputs(boundary_inputs, outer_task_id, orch->tensor_map, orch->in_manual_scope()); + + outer_payload.fanin_actual_count = fanin_builder.count; + outer_payload.fanin_spill_start = fanin_builder.spill_start; + outer_payload.fanin_spill_pool = &fanin_builder.spill_pool; + const int32_t inline_count = std::min(fanin_builder.count, PTO2_FANIN_INLINE_CAP); + for (int32_t i = 0; i < inline_count; ++i) + outer_payload.fanin_inline_slot_states[i] = fanin_builder.inline_slots[i]; + + if (fanin_builder.count == 0) { + outer_slot.fanin_count = 1; + outer_slot.fanin_refcount.store(1, std::memory_order_release); + orch_mark_dep_pool_position(orch, outer_slot); + orch->scheduler->push_ready_routed(&outer_slot); + } else if (all_claimed_fanin_completed(fanin_builder)) { + int32_t ready_seed = fanin_builder.count + 1; + outer_slot.fanin_count = ready_seed; + outer_slot.fanin_refcount.store(ready_seed, std::memory_order_release); + orch_mark_dep_pool_position(orch, outer_slot); + orch->scheduler->push_ready_routed(&outer_slot); + } else if (!orch_wire_live_fanin_task(orch, outer_slot, fanin_builder.count)) { + pto2_graph_submission_discard(submission); + return false; + } + + if (orch_record_task_id != nullptr) *orch_record_task_id = outer_task_id.raw; + if (!orch->scheduler->graph_prepare_queue.push_tagged(&outer_slot, outer_task_id.raw)) { + orch_mark_fatal(orch, PTO2_ERROR_HEAP_RING_DEADLOCK); + pto2_graph_submission_discard(submission); + return false; + } + submission->upload_next = orch->pending_graph_uploads; + orch->pending_graph_uploads = submission; +#if SIMPLER_DFX + orch->tasks_submitted++; +#endif + return true; +} + +} // namespace + +bool pto2_graph_definition_relocate_for_upload(PTO2GraphSubmission *submission, void *device_base) { + if (submission == nullptr || device_base == nullptr || submission->graph_definition == nullptr) return false; + auto *definition = + const_cast(static_cast(submission->graph_definition)); + const uintptr_t host_begin = reinterpret_cast(submission); + const intptr_t delta = + static_cast(reinterpret_cast(device_base)) - static_cast(host_begin); + const size_t bytes = submission->allocation_bytes; + return graph_rebase_cloned_pointer(definition->storage, host_begin, bytes, delta) && + graph_rebase_cloned_pointer(definition->fanout_offsets, host_begin, bytes, delta) && + graph_rebase_cloned_pointer(definition->fanout_indices, host_begin, bytes, delta) && + graph_rebase_cloned_pointer(definition->fanin_counts, host_begin, bytes, delta) && + graph_rebase_cloned_pointer(definition->root_indices, host_begin, bytes, delta) && + graph_rebase_cloned_pointer(definition->node_offsets, host_begin, bytes, delta) && + graph_rebase_cloned_pointer(definition->tasks, host_begin, bytes, delta) && + graph_rebase_cloned_pointer(definition->tensors, host_begin, bytes, delta) && + graph_rebase_cloned_pointer(definition->tensor_sources, host_begin, bytes, delta) && + graph_rebase_cloned_pointer(definition->scalars, host_begin, bytes, delta) && + graph_rebase_cloned_pointer(definition->scalar_source_indices, host_begin, bytes, delta); +} + +PTO2GraphExecution *pto2_graph_execution_localize(PTO2TaskSlotState &outer_slot) { + if (outer_slot.task == nullptr || outer_slot.task->kind != PTO2TaskKind::GRAPH) return nullptr; + PTO2GraphSubmission *submission = pto2_graph_submission_from_task(*outer_slot.task); + if (submission == nullptr || submission->graph_definition == nullptr) return nullptr; + + PTO2GraphExecution *execution = submission->local_execution.load(std::memory_order_acquire); + if (execution != nullptr) return execution; + + // Reclaim completed local blocks before allocating. This pool is owned by + // the Scheduler DSO and survives run boundaries, so later invocations can + // reuse both the allocation and graph-affine static node contents. + pto2_graph_execution_collect_retired(); + const auto &templ = *static_cast(submission->graph_definition); + const size_t definition_bytes = graph_definition_image_bytes(templ); + execution = pto2_graph_execution_create(submission->node_count, submission->graph_key, definition_bytes); + if (execution == nullptr) return nullptr; + + if (!execution->definition_affine_reuse && !graph_clone_definition(execution, templ)) { + pto2_graph_execution_discard(execution); + return nullptr; + } + execution->args = submission->args; + execution->outer_slot = &outer_slot; + + PTO2GraphExecution *expected = nullptr; + if (!submission->local_execution.compare_exchange_strong( + expected, execution, std::memory_order_release, std::memory_order_acquire + )) { + pto2_graph_execution_discard(execution); + return expected; + } + pto2_graph_execution_publish(execution); + return execution; +} + +bool PTO2OrchestratorState::finalize_pending_graph_definition() { + if (!g_graph_recording.pending_finalize) return true; + + bool built = graph_build_definition(&g_graph_recording.temp, g_graph_recording.args); + bool stored = built && store_graph_definition(g_graph_recording.temp); + if (stored) { + PTO2GraphDefinition *definition = find_graph_definition(g_graph_recording.full_key); + LOG_INFO_V0( + "[GraphExecution] define key=0x%llx nodes=%d bytes=%zu", + static_cast(g_graph_recording.full_key), g_graph_recording.temp.task_count, + definition == nullptr ? 0 : definition->storage_bytes + ); + } + reset_graph_recording(); + return stored; +} + +PTO2GraphMaterializeResult pto2_graph_execution_materialize_slice( + PTO2TaskSlotState &outer_slot, PTO2GraphExecution &execution_ref, int32_t max_nodes, int32_t *nodes_materialized +) { + if (nodes_materialized != nullptr) *nodes_materialized = 0; + if (outer_slot.task == nullptr || outer_slot.task->kind != PTO2TaskKind::GRAPH || max_nodes <= 0) { + return PTO2GraphMaterializeResult::INVALID; + } + PTO2GraphExecution *execution = &execution_ref; + if (execution == nullptr || execution->graph_definition == nullptr || execution->node_storage == nullptr) { + return PTO2GraphMaterializeResult::INVALID; + } + + PTO2GraphExecutionState state = execution->state.load(std::memory_order_acquire); + if (state >= PTO2GraphExecutionState::PREPARED) return PTO2GraphMaterializeResult::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 PTO2GraphMaterializeResult::BUSY; + } + + state = execution->state.load(std::memory_order_acquire); + if (state == PTO2GraphExecutionState::SUBMITTED) { + if (!pto2_graph_execution_begin_materialize(*execution)) { + execution->materialize_busy.store(0, std::memory_order_release); + state = execution->state.load(std::memory_order_acquire); + return state >= PTO2GraphExecutionState::PREPARED ? PTO2GraphMaterializeResult::PREPARED : + PTO2GraphMaterializeResult::BUSY; + } + } else if (state != PTO2GraphExecutionState::MATERIALIZING) { + execution->materialize_busy.store(0, std::memory_order_release); + return state >= PTO2GraphExecutionState::PREPARED ? PTO2GraphMaterializeResult::PREPARED : + PTO2GraphMaterializeResult::INVALID; + } + + const auto &templ = *static_cast(execution->graph_definition); + const ChipStorageTaskArgs &args = execution->args; + const bool reuse_static = execution->definition_affine_reuse; + const int32_t first_node = execution->materialized_nodes; + if (first_node == 0 && !reuse_static) { + execution->materialized_graph_key = 0; + execution->materialized_node_count = 0; + } + const int32_t last_node = std::min(templ.task_count, first_node + max_nodes); + for (int32_t i = first_node; i < last_node; ++i) { + const PTO2GraphNodeDefinition &src = templ.tasks[i]; + PTO2GraphNodeStorage *node = &execution->node_storage[i]; + if (i >= execution->constructed_nodes) { + // Default-initialize object lifetimes without value-initializing + // the 4 KiB payload arrays. Every live descriptor/slot field and + // every tensor/scalar below the published counts is overwritten + // before PREPARED is release-published. + node = new (node) PTO2GraphNodeStorage; + execution->constructed_nodes++; + } + execution->materialized_nodes++; + PTO2TaskDescriptor &task = node->task; + PTO2TaskPayload &payload = node->payload; + PTO2TaskSlotState &slot = node->slot; + + uint32_t synthetic_local = + (static_cast(outer_slot.task->task_id.local()) << 10) | static_cast(i); + task.task_id = PTO2TaskId::make(1, synthetic_local); + task.packed_buffer_base = static_cast(outer_slot.task->packed_buffer_base) + templ.node_offsets[i]; + task.packed_buffer_end = static_cast(task.packed_buffer_base) + + PTO2_ALIGN_UP(static_cast(src.total_output_size), PTO2_ALIGN_SIZE); + if (!reuse_static) { + for (int k = 0; k < PTO2_SUBTASK_SLOT_COUNT; ++k) + task.kernel_id[k] = src.kernel_id[k]; + task.task_timing_slot = src.task_timing_slot; + task.graph_execution = execution; + task.graph_node_index = i; + task.kind = PTO2TaskKind::GRAPH_NODE; + } + + slot.reset_for_reuse(); + if (!reuse_static) slot.bind_buffers(&payload, &task); + slot.task_state.store(PTO2_TASK_PENDING, std::memory_order_relaxed); + slot.fanin_refcount.store(1, std::memory_order_relaxed); + if (!reuse_static) { + slot.fanin_count = static_cast(templ.fanin_counts[i]) + 1; + slot.active_mask = src.active_mask; + // Graph topology is released directly, so ordinary fanout-based + // early propagation is intentionally disabled for Graph nodes. + slot.set_allow_early_resolve(false); + slot.total_required_subtasks = src.total_required_subtasks; + slot.logical_block_num = src.logical_block_num; + } + + if (!reuse_static) { + payload.tensor_count = src.tensor_count; + payload.scalar_count = src.scalar_count; + } + payload.fanin_actual_count = 0; + payload.fanin_spill_start = 0; + payload.fanin_spill_pool = nullptr; + payload.predicate = DispatchPredicate{}; + for (int32_t j = 0; j < src.tensor_count; ++j) { + const uint32_t tensor_index = src.tensor_offset + static_cast(j); + Tensor &tensor = payload.tensors[j]; + if (!reuse_static) tensor.copy(templ.tensors[tensor_index]); + const PTO2GraphTensorSourceRef &source_ref = templ.tensor_sources[tensor_index]; + if (source_ref.source == PTO2GraphTensorSource::BOUNDARY_EXACT) { + tensor.copy(args.tensor(source_ref.source_index)); + } else if (source_ref.source == PTO2GraphTensorSource::BOUNDARY_VIEW) { + const Tensor &boundary = args.tensor(source_ref.source_index); + tensor.buffer = boundary.buffer; + tensor.owner_task_id = boundary.owner_task_id; + tensor.start_offset = boundary.start_offset + source_ref.packed_offset; + tensor.version = boundary.version; + tensor.child_memory = boundary.child_memory; + } else { + int32_t producer_index = source_ref.source == PTO2GraphTensorSource::OWN_OUTPUT ? + i : + static_cast(source_ref.source_index); + PTO2TaskDescriptor &producer = execution->node_storage[producer_index].task; + tensor.buffer.addr = reinterpret_cast(producer.packed_buffer_base) + source_ref.packed_offset; + tensor.owner_task_id = producer.task_id; + } + } + for (int32_t j = 0; j < src.scalar_count; ++j) { + const uint32_t scalar_index = src.scalar_offset + static_cast(j); + uint16_t source_index = templ.scalar_source_indices[scalar_index]; + if (!reuse_static || source_index != PTO2_TASK_ARG_STATIC) { + payload.scalars[j] = + source_index == PTO2_TASK_ARG_STATIC ? templ.scalars[scalar_index] : args.scalar(source_index); + } + } + graph_reset_payload(&payload); + } + if (nodes_materialized != nullptr) *nodes_materialized = last_node - first_node; + + if (execution->materialized_nodes < templ.task_count) { + execution->materialize_busy.store(0, std::memory_order_release); + return PTO2GraphMaterializeResult::PENDING; + } + + execution->nodes = execution->node_storage; + execution->materialized_graph_key = execution->graph_key; + execution->materialized_node_count = execution->node_count; + pto2_graph_execution_publish_materialized(*execution); + execution->materialize_busy.store(0, std::memory_order_release); + return PTO2GraphMaterializeResult::PREPARED; +} + static PTO2OutputLayout calculate_output_layout(const L0TaskArgs &args) { PTO2OutputLayout layout; for (int32_t i = 0; i < args.tensor_count(); i++) { @@ -541,6 +1674,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(out->task_id); + out->payload->prefetch(args.tensor_count(), args.scalar_count()); // Re-bind payload/task pointers each submit. Value is per-slot constant @@ -927,6 +2062,9 @@ static TaskOutputTensors submit_task_common( task.task_timing_slot = args.task_timing_slot(); task.packed_buffer_base = prepared.alloc_result.packed_base; task.packed_buffer_end = prepared.alloc_result.packed_end; + task.graph_execution = nullptr; + task.graph_node_index = -1; + task.kind = static_cast(active_mask) ? PTO2TaskKind::KERNEL : PTO2TaskKind::DUMMY; // fanout_count was already incremented per live producer inside // append_fanin_or_fail, atomically with the consumed/generation check under @@ -963,6 +2101,7 @@ static TaskOutputTensors submit_task_common( payload.predicate.op = PredicateOp::NONE; } } + graph_record_task(prepared, args); #if SIMPLER_DFX if (is_dump_args_enabled()) { if (args.scalar_count() > 0) { @@ -1188,6 +2327,9 @@ TaskOutputTensors PTO2OrchestratorState::alloc_tensors(const L0TaskArgs &args) { task.task_timing_slot = TASK_TIMING_SLOT_NONE; task.packed_buffer_base = prepared.alloc_result.packed_base; task.packed_buffer_end = prepared.alloc_result.packed_end; + task.graph_execution = nullptr; + task.graph_node_index = -1; + task.kind = PTO2TaskKind::DUMMY; TaskOutputTensors outputs; outputs.set_task_id(prepared.task_id); @@ -1217,6 +2359,7 @@ TaskOutputTensors PTO2OrchestratorState::alloc_tensors(const L0TaskArgs &args) { prepared.slot_state->allow_early_resolve = true; prepared.slot_state->mark_completed(); } + graph_record_task(prepared, args); orch->inline_completed_tasks++; CYCLE_COUNT_LAP(g_orch_fanin_cycle); @@ -1233,12 +2376,81 @@ TaskOutputTensors PTO2OrchestratorState::alloc_tensors(const L0TaskArgs &args) { return outputs; } +PTO2GraphScopeResult +PTO2OrchestratorState::graph_begin(uint64_t graph_key, const L2TaskArgs &args, uint64_t callable_hash) { + PTO2GraphScopeResult result; + auto *orch = this; + if (g_graph_recording.pending_finalize) finalize_pending_graph_definition(); + + ChipStorageTaskArgs args_snapshot; + if (orch->fatal || !graph_snapshot_args(args, &args_snapshot)) return result; + if (g_graph_recording.active) { + graph_mark_unsupported(PTO2_GRAPH_UNSUPPORTED_NESTED_SCOPE); + return result; + } + + uint64_t full_key = graph_full_key(callable_hash, graph_key); + PTO2GraphDefinition *cached = find_graph_definition(full_key); + if (cached != nullptr) { +#if SIMPLER_DFX + bool record_replay_orch = orch->l2_swimlane_level >= L2SwimlaneLevel::ORCH_PHASES; + uint64_t replay_start_ts = record_replay_orch ? get_sys_cnt_aicpu() : 0; +#endif + uint64_t replay_task_id = PTO2TaskId::invalid().raw; + if (graph_submit_definition(orch, *cached, args_snapshot, &replay_task_id)) { +#if SIMPLER_DFX + if (record_replay_orch) { + record_host_orch_phase(orch, replay_start_ts, get_sys_cnt_aicpu(), replay_task_id, g_orch_submit_idx++); + } +#endif + LOG_INFO_V0( + "[GraphExecution] submit key=0x%llx nodes=%d", static_cast(full_key), + cached->task_count + ); + result.execute_block = false; + result.recording = false; + result.task_id = PTO2TaskId{replay_task_id}; + return result; + } + if (orch->fatal) return result; + } + + reset_graph_recording(); + g_graph_recording.active = true; + g_graph_recording.full_key = full_key; + g_graph_recording.start_local_task_id = orch->ring.task_allocator.task_head(); + g_graph_recording.args = args_snapshot; + reset_graph_definition_header(&g_graph_recording.temp); + g_graph_recording.temp.full_key = full_key; + result.execute_block = true; + result.recording = true; + return result; +} + +void PTO2OrchestratorState::graph_end() { + if (!g_graph_recording.active) return; + if (g_graph_recording.unsupported || g_graph_recording.temp.task_count <= 0) { + if (g_graph_recording.unsupported) { + LOG_WARN( + "Graph Execution definition skipped: reason=%d task_index=%d tensor_index=%d recorded_tasks=%d", + g_graph_recording.unsupported_reason, g_graph_recording.unsupported_task_index, + g_graph_recording.unsupported_tensor_index, g_graph_recording.temp.task_count + ); + } + reset_graph_recording(); + return; + } + g_graph_recording.active = false; + g_graph_recording.pending_finalize = true; +} + // ============================================================================= // Flow Control // ============================================================================= void PTO2OrchestratorState::mark_done() { auto *orch = this; + if (!orch->fatal) finalize_pending_graph_definition(); for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { int32_t total_tasks = orch->ring.task_allocator.active_count(); if (total_tasks > 0) { 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 2f1bcae127..5198139574 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,15 @@ static TaskOutputTensors submit_dummy_task_impl(PTO2Runtime *rt, const L0TaskArg return rt->orchestrator.submit_dummy_task(args); } +static PTO2GraphScopeResult graph_begin_impl(PTO2Runtime *rt, uint64_t graph_key, const L2TaskArgs &args) { + if (rt == nullptr) return PTO2GraphScopeResult{}; + 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(); +} + void rt_scope_begin(PTO2Runtime *rt) { PTO2ScopeMode mode = rt->pending_scope_mode; rt->pending_scope_mode = PTO2ScopeMode::AUTO; @@ -295,6 +304,8 @@ static const PTO2RuntimeOps s_runtime_ops = { .set_tensor_data = set_tensor_data, .alloc_tensors = alloc_tensors_impl, .submit_dummy_task = submit_dummy_task_impl, + .graph_begin = graph_begin_impl, + .graph_end = graph_end_impl, #if SIMPLER_DFX .scope_set_site = scope_set_site_impl, #else diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_graph_cache.h b/src/a2a3/runtime/host_build_graph/runtime/pto_graph_cache.h new file mode 100644 index 0000000000..7b1d58476a --- /dev/null +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_graph_cache.h @@ -0,0 +1,87 @@ +/* + * 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" + +// Versioned independently from the runtime ABI so cached definitions are +// invalidated when Graph Execution materialization semantics change. +inline constexpr uint64_t PTO2_GRAPH_CACHE_SCHEMA_VERSION = 6; +inline constexpr uint32_t PTO2_GRAPH_MAX_TENSOR_ARGS = 32; +inline constexpr uint32_t PTO2_GRAPH_MAX_SCALAR_ARGS = 32; + +struct PTO2GraphScopeResult { + bool execute_block{true}; + bool recording{false}; + PTO2TaskId task_id{PTO2TaskId::invalid()}; +}; + +using PTO2GraphSubmitResult = PTO2GraphScopeResult; + +constexpr uint64_t pto2_graph_hash_byte(uint64_t h, uint8_t b) { + return (h ^ static_cast(b)) * 1099511628211ULL; +} + +inline uint64_t pto2_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 = pto2_graph_hash_byte(h, p[i]); + } + return h; +} + +constexpr uint64_t pto2_graph_const_hash_impl(const char *s, uint64_t h) { + return (*s == '\0') ? h : pto2_graph_const_hash_impl(s + 1, pto2_graph_hash_byte(h, static_cast(*s))); +} + +constexpr uint64_t PTO2_GRAPH_KEY(const char *s) { return pto2_graph_const_hash_impl(s, 1469598103934665603ULL); } + +inline bool rt_graph_args_cacheable(const L2TaskArgs &args) { + if (args.has_error || args.tensor_count() > static_cast(PTO2_GRAPH_MAX_TENSOR_ARGS) || + args.scalar_count() > static_cast(PTO2_GRAPH_MAX_SCALAR_ARGS)) { + 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, const L2TaskArgs &args) { + uint64_t h = 1469598103934665603ULL; + h = pto2_graph_hash_bytes(h, &PTO2_GRAPH_CACHE_SCHEMA_VERSION, sizeof(PTO2_GRAPH_CACHE_SCHEMA_VERSION)); + h = pto2_graph_hash_bytes(h, &graph_id, sizeof(graph_id)); + int32_t tensor_count = args.tensor_count(); + int32_t scalar_count = args.scalar_count(); + h = pto2_graph_hash_bytes(h, &tensor_count, sizeof(tensor_count)); + h = pto2_graph_hash_bytes(h, &scalar_count, sizeof(scalar_count)); + for (int32_t i = 0; i < tensor_count; ++i) { + const Tensor &tensor = args.tensor(i).ref(); + TensorArgType type = args.tag(i); + h = pto2_graph_hash_bytes(h, &tensor.buffer.size, sizeof(tensor.buffer.size)); + h = pto2_graph_hash_bytes(h, &tensor.ndims, sizeof(tensor.ndims)); + h = pto2_graph_hash_bytes(h, &tensor.dtype, sizeof(tensor.dtype)); + h = pto2_graph_hash_bytes(h, &tensor.manual_dep, sizeof(tensor.manual_dep)); + h = pto2_graph_hash_bytes(h, &tensor.is_contiguous, sizeof(tensor.is_contiguous)); + h = pto2_graph_hash_bytes(h, &type, sizeof(type)); + h = pto2_graph_hash_bytes(h, tensor.shapes, sizeof(uint32_t) * tensor.ndims); + h = pto2_graph_hash_bytes(h, tensor.strides, sizeof(uint32_t) * tensor.ndims); + } + return h; +} diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_graph_execution.h b/src/a2a3/runtime/host_build_graph/runtime/pto_graph_execution.h new file mode 100644 index 0000000000..2309883e4c --- /dev/null +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_graph_execution.h @@ -0,0 +1,158 @@ +/* + * 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 "pto_graph_cache.h" +#include "pto_runtime2_types.h" +#include "task_args.h" + +enum class PTO2GraphExecutionState : uint8_t { + SUBMITTED = 0, + MATERIALIZING = 1, + PREPARED = 2, + ACTIVE = 3, + COMPLETED = 4, +}; + +enum class PTO2GraphMaterializeResult : uint8_t { + INVALID = 0, + BUSY = 1, + PENDING = 2, + PREPARED = 3, +}; + +constexpr int32_t PTO2_GRAPH_MATERIALIZE_SLICE_NODES = 4; + +struct PTO2GraphTopologyView { + int32_t edge_count{0}; + int32_t root_count{0}; + const uint32_t *fanout_offsets{nullptr}; + const uint16_t *fanout_indices{nullptr}; + const uint16_t *fanin_counts{nullptr}; + const uint16_t *root_indices{nullptr}; +}; + +struct alignas(64) PTO2GraphNodeStorage { + PTO2TaskDescriptor task; + PTO2TaskPayload payload; + PTO2TaskSlotState slot; +}; + +struct PTO2GraphExecution { + std::atomic state{PTO2GraphExecutionState::SUBMITTED}; + std::atomic activation_requested{0}; + // A prepare token owns at most one bounded slice; this claim protects the + // cursor if another scheduler observes a requeued token concurrently. + 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}; + // Static node contents survive in the bounded AICPU-local pool. A block + // with the same graph key can therefore patch only per-invocation fields. + 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 materialized_graph_key{0}; + bool definition_affine_reuse{false}; + PTO2TaskSlotState *outer_slot{nullptr}; + PTO2GraphNodeStorage *nodes{nullptr}; + PTO2GraphNodeStorage *node_storage{nullptr}; + void *definition_storage{nullptr}; + const void *graph_definition{nullptr}; + PTO2GraphTopologyView topology; + ChipStorageTaskArgs args; + PTO2GraphExecution *next{nullptr}; +}; + +// Compact host-built object uploaded for one Graph task. It contains the +// immutable Definition plus this invocation's TaskArgs, but deliberately no +// expanded node storage. Scheduler workers localize it into an AICPU pool. +struct PTO2GraphSubmission { + std::atomic local_execution{nullptr}; + std::atomic activation_requested{0}; + int32_t node_count{0}; + size_t allocation_bytes{0}; + size_t definition_capacity{0}; + uint64_t graph_key{0}; + PTO2TaskSlotState *outer_slot{nullptr}; + void *definition_storage{nullptr}; + const void *graph_definition{nullptr}; + ChipStorageTaskArgs args; + PTO2GraphSubmission *next{nullptr}; + PTO2GraphSubmission *upload_next{nullptr}; +}; + +PTO2GraphExecution * +pto2_graph_execution_create(int32_t node_count, uint64_t graph_key = 0, size_t definition_bytes = 0); +void pto2_graph_execution_discard(PTO2GraphExecution *execution); +void pto2_graph_execution_publish(PTO2GraphExecution *execution); +void pto2_graph_execution_collect_retired(); + +PTO2GraphSubmission * +pto2_graph_submission_create(int32_t node_count, uint64_t graph_key = 0, size_t definition_bytes = 0); +void pto2_graph_submission_discard(PTO2GraphSubmission *submission); +void pto2_graph_submission_release_uploaded(PTO2GraphSubmission *submission); +size_t pto2_graph_submission_allocation_bytes(const PTO2GraphSubmission *submission); +bool pto2_graph_submission_relocate_for_upload( + PTO2GraphSubmission *submission, void *device_base, uintptr_t host_sm_begin, size_t host_sm_size, intptr_t sm_delta +); +bool pto2_graph_definition_relocate_for_upload(PTO2GraphSubmission *submission, void *device_base); +PTO2GraphExecution *pto2_graph_execution_localize(PTO2TaskSlotState &outer_slot); +PTO2GraphMaterializeResult pto2_graph_execution_materialize_slice( + PTO2TaskSlotState &outer_slot, PTO2GraphExecution &execution, int32_t max_nodes, + int32_t *nodes_materialized = nullptr +); + +inline PTO2GraphExecution *pto2_graph_execution_from_task(const PTO2TaskDescriptor &task) { + return static_cast(task.graph_execution); +} + +inline PTO2GraphSubmission *pto2_graph_submission_from_task(const PTO2TaskDescriptor &task) { + return static_cast(task.graph_execution); +} + +inline bool pto2_graph_execution_begin_materialize(PTO2GraphExecution &execution) { + PTO2GraphExecutionState expected = PTO2GraphExecutionState::SUBMITTED; + return execution.state.compare_exchange_strong( + expected, PTO2GraphExecutionState::MATERIALIZING, std::memory_order_acq_rel, std::memory_order_acquire + ); +} + +inline void pto2_graph_execution_publish_materialized(PTO2GraphExecution &execution) { + execution.state.store(PTO2GraphExecutionState::PREPARED, std::memory_order_release); +} + +inline bool pto2_graph_execution_activate(PTO2GraphExecution &execution) { + PTO2GraphExecutionState expected = PTO2GraphExecutionState::PREPARED; + return execution.state.compare_exchange_strong( + expected, PTO2GraphExecutionState::ACTIVE, std::memory_order_acq_rel, std::memory_order_acquire + ); +} + +inline bool pto2_graph_execution_complete_node(PTO2GraphExecution &execution) { + return execution.remaining_nodes.fetch_sub(1, std::memory_order_acq_rel) == 1; +} + +inline void pto2_graph_execution_mark_completed(PTO2GraphExecution &execution) { + execution.state.store(PTO2GraphExecutionState::COMPLETED, std::memory_order_release); +} + +inline void pto2_graph_execution_retire_node(PTO2GraphExecution &execution) { + execution.retired_nodes.fetch_add(1, std::memory_order_release); +} 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 b252f3a19a..8265b15305 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 "pto_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 PTO2GraphSubmission; + /** * Layout descriptor produced by PTO2OrchestratorState::reserve_layout(). Holds * arena offsets for every sub-region the orchestrator owns (per-ring fanin @@ -115,11 +118,24 @@ struct PTO2OrchestratorState { // after orchestration finishes so shutdown/profiling totals remain closed. int64_t inline_completed_tasks{0}; + // Host-side execution images waiting for runtime_maker to relocate and + // upload. Each replayed Graph contributes one compact block; ordinary + // task-window storage still contains only its outer Graph descriptor. + PTO2GraphSubmission *pending_graph_uploads{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; } @@ -157,6 +173,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); + PTO2GraphScopeResult graph_begin(uint64_t graph_key, const L2TaskArgs &args, uint64_t callable_hash); + void graph_end(); + bool finalize_pending_graph_definition(); 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 4ecd2a14ac..39ab8dfeb8 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 "pto_graph_cache.h" #include "pto_submit_types.h" #include "pto_shared_memory.h" #include "pto_ring_buffer.h" @@ -89,6 +90,8 @@ struct PTO2RuntimeOps { ); TaskOutputTensors (*alloc_tensors)(PTO2Runtime *rt, const L0TaskArgs &args); TaskOutputTensors (*submit_dummy_task)(PTO2Runtime *rt, const L0TaskArgs &args); + PTO2GraphScopeResult (*graph_begin)(PTO2Runtime *rt, uint64_t graph_key, const L2TaskArgs &args); + void (*graph_end)(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. @@ -146,6 +149,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 0548bd739e..71892d07e7 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 @@ -160,6 +160,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] = {}; @@ -219,12 +226,18 @@ struct PTO2TaskDescriptor { // Packed output buffer (all outputs packed into single contiguous buffer) void *packed_buffer_base; // Start of packed buffer in GM Heap void *packed_buffer_end; // End of packed buffer (for heap reclamation) + + // Ordinary tasks leave these fields null/-1. A Graph task points at the + // separately uploaded GraphExecution block; scheduler-owned Graph nodes + // point at the same block and carry their immutable topological index. + void *graph_execution{nullptr}; + int32_t graph_node_index{-1}; + PTO2TaskKind kind{PTO2TaskKind::KERNEL}; }; -// task_timing_slot must occupy the former padding after kernel_id[3] without -// growing the descriptor or shifting packed_buffer_base — the scheduler and -// shared-memory ABI depend on the existing size and pointer offset. -static_assert(sizeof(PTO2TaskDescriptor) == 40, "PTO2TaskDescriptor must not grow: slot uses the existing pad"); +// task_timing_slot retains its established location. Graph metadata follows +// the original descriptor prefix so AICore payload consumers keep the same +// offsets for task identity, kernels, timing and packed buffers. static_assert( offsetof(PTO2TaskDescriptor, task_timing_slot) == offsetof(PTO2TaskDescriptor, kernel_id) + sizeof(int32_t) * PTO2_SUBTASK_SLOT_COUNT, diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_types.h b/src/a2a3/runtime/host_build_graph/runtime/pto_types.h index ad0109c9be..1324ad700c 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_types.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_types.h @@ -66,6 +66,11 @@ enum class PTO2ScopeMode : uint8_t { MANUAL = 1, }; +// Scalar values copied from L2TaskArgs retain their invocation argument index. +// Scalars constructed directly inside a graph function are part of the static +// Graph Definition. +inline constexpr uint16_t PTO2_TASK_ARG_STATIC = UINT16_MAX; + /** * TaskOutputTensors — returned by submit, holds materialized output Tensors. * @@ -176,6 +181,25 @@ class TensorRef { bool refers_to(const TensorCreateInfo *ci) const { return create_info_ == ci; } }; +class ScalarRef { +public: + ScalarRef(uint64_t value, uint16_t source_index) : + value_(value), + source_index_(source_index) {} + + operator uint64_t() const { return value_; } + uint64_t value() const { return value_; } + uint16_t source_index() const { return source_index_; } + +private: + uint64_t value_; + uint16_t source_index_; +}; + +template +inline constexpr bool is_task_scalar_arg_v = + is_supported_scalar_arg_v || std::is_same_v>, ScalarRef>; + /** * Aggregated argument container for pto_submit_task * @@ -253,6 +277,21 @@ struct Arg : TaskArgsTpl { void set_allow_early_resolve(bool v = true) { allow_early_resolve_ = v; } bool allow_early_resolve() const { return allow_early_resolve_; } + ScalarRef scalar(int32_t index) const { + always_assert(index >= 0 && index < scalar_count_); + return ScalarRef{scalars_[index], scalar_source_indices_[index]}; + } + + uint64_t &scalar(int32_t index) { + always_assert(index >= 0 && index < scalar_count_); + return scalars_[index]; + } + + uint16_t scalar_source_index(int32_t index) const { + always_assert(index >= 0 && index < scalar_count_); + return scalar_source_indices_[index]; + } + // Dispatch predicate (codegen-author set; default op == NONE = always // dispatch). A FALSE result at the dispatch point retires the task inline // through the dep-only path — never dispatched to an AICore — while still @@ -427,7 +466,10 @@ struct Arg : TaskArgsTpl { template void add_scalar(Args &&...args) { static_assert(sizeof...(Args) >= 1, "add_scalar: at least one argument required"); - static_assert((is_supported_scalar_arg_v && ...), "add_scalar: all types must be arithmetic or enum"); + static_assert( + (is_task_scalar_arg_v && ...), + "add_scalar: all types must be arithmetic, enum, or scalar references from another Args object" + ); if (scalar_count_ + sizeof...(Args) > MaxS) { set_error(scalar_cap_msg()); return; @@ -441,6 +483,8 @@ struct Arg : TaskArgsTpl { return; } memcpy(&scalars_[scalar_count_], values, count * sizeof(uint64_t)); + for (int i = 0; i < count; ++i) + scalar_source_indices_[scalar_count_ + i] = PTO2_TASK_ARG_STATIC; #if SIMPLER_DFX dump_arg_selection_.clear_scalar_metadata(scalar_count_, count); #endif @@ -479,6 +523,8 @@ struct Arg : TaskArgsTpl { #if SIMPLER_DFX dump_arg_selection_.clear_scalar_metadata(scalar_count_, count); #endif + for (int i = 0; i < count; ++i) + scalar_source_indices_[scalar_count_ + i] = PTO2_TASK_ARG_STATIC; scalar_count_ += count; } @@ -496,6 +542,9 @@ struct Arg : TaskArgsTpl { return; } memcpy(&scalars_[scalar_count_], &src.scalars_[src_offset], count * sizeof(uint64_t)); + memcpy( + &scalar_source_indices_[scalar_count_], &src.scalar_source_indices_[src_offset], count * sizeof(uint16_t) + ); #if SIMPLER_DFX dump_arg_selection_.copy_scalar_dtypes_from(src.dump_arg_selection_, scalar_count_, src_offset, count); #endif @@ -515,6 +564,7 @@ struct Arg : TaskArgsTpl { #endif const PTO2TaskId *explicit_deps_{nullptr}; uint32_t explicit_dep_count_{0}; + uint16_t scalar_source_indices_[MaxS]{}; #if SIMPLER_DFX template static constexpr bool is_supported_dump_arg_v = @@ -535,16 +585,24 @@ struct Arg : TaskArgsTpl { } template - void add_scalar_one(T &&value) { - scalars_[scalar_count_] = to_u64(value); + void add_scalar_one(T &&value, uint16_t source_index = PTO2_TASK_ARG_STATIC) { + using ValueT = std::remove_cv_t>; + if constexpr (std::is_same_v) { + scalars_[scalar_count_] = value.value(); + scalar_source_indices_[scalar_count_] = + source_index == PTO2_TASK_ARG_STATIC ? value.source_index() : source_index; + } else { + scalars_[scalar_count_] = to_u64(value); + scalar_source_indices_[scalar_count_] = source_index; + } #if SIMPLER_DFX uintptr_t scalar_source_ptr = 0; if constexpr (std::is_lvalue_reference_v) { scalar_source_ptr = reinterpret_cast(&value); } - dump_arg_selection_.record_scalar_source( - scalar_count_, scalar_source_ptr, dtype_of>>() - ); + uint8_t scalar_dtype = static_cast(DataType::UINT64); + if constexpr (!std::is_same_v) scalar_dtype = dtype_of(); + dump_arg_selection_.record_scalar_source(scalar_count_, scalar_source_ptr, scalar_dtype); #endif scalar_count_++; } @@ -640,6 +698,15 @@ using L0TaskArgs = Arg; // already-allocated inputs (capacity matches ChipStorageTaskArgs). // aicpu_orchestration_entry/config receive a const L2TaskArgs&. struct L2TaskArgs : Arg { + using Base = Arg; + + ScalarRef scalar(int32_t index) const { + always_assert(index >= 0 && index < this->scalar_count_); + return ScalarRef{this->scalars_[index], static_cast(index)}; + } + + uint64_t &scalar(int32_t index) { return Base::scalar(index); } + // Build from the executor's ChipStorageTaskArgs: each input becomes a // TensorRef pointing at src's Tensor, so `src` must outlive this (on the // executor path src is runtime->orch_args_storage_, alive for the whole run). diff --git a/src/a2a3/runtime/host_build_graph/runtime/runtime.h b/src/a2a3/runtime/host_build_graph/runtime/runtime.h index d49e6d7242..8746820087 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/runtime.h +++ b/src/a2a3/runtime/host_build_graph/runtime/runtime.h @@ -314,6 +314,23 @@ 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 { + return host_orch_phase_records_; + } + uint64_t get_host_orch_start_cycles() const { return host_orch_start_cycles_; } + uint64_t get_host_orch_end_cycles() const { return host_orch_end_cycles_; } }; // 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/pto_scheduler.h b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h index a46f484df3..af2a57dba3 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 @@ -36,6 +36,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 "pto_graph_execution.h" #include "pto_ring_buffer.h" #include "pto_runtime2_types.h" #include "pto_shared_memory.h" @@ -417,6 +418,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; size_t off_dep_pool_entries; @@ -486,6 +489,12 @@ struct PTO2SchedulerState { // the dispatch loop and completed inline -- never goes to AICore. PTO2ReadyQueue dummy_ready_queue; + // Graph outer tasks are scheduler control work. Preparation expands the + // immutable definition in bounded slices; activation remains gated by the + // outer task's external dependencies. + PTO2ReadyQueue graph_ready_queue; + PTO2ReadyQueue graph_prepare_queue; + alignas(64) AsyncWaitList async_wait_list; // Statistics (cold path, isolated from hot-path fields) @@ -503,6 +512,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 != nullptr && 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->active_mask.has_predicate() && !slot_state->payload->predicate.pass())) { @@ -907,6 +920,11 @@ struct PTO2SchedulerState { bool route_ready_once(PTO2TaskSlotState &slot_state, EarlyDispatchReleaseSink *sink = nullptr) { if (!try_claim_ready_once(slot_state)) return false; + if (slot_state.task != nullptr && slot_state.task->kind == PTO2TaskKind::GRAPH) { + graph_ready_queue.push(&slot_state); + return true; + } + // Early-dispatch: pre-staged tasks are released by doorbell // here, skipping the ready-queue round-trip entirely. bool early_handled = try_early_dispatch_release(slot_state, sink); @@ -948,6 +966,11 @@ struct PTO2SchedulerState { atomic_count += 1; // failed lifecycle_flags CAS } + if (slot_state.task != nullptr && slot_state.task->kind == PTO2TaskKind::GRAPH) { + graph_ready_queue.push(&slot_state, atomic_count, push_wait); + return true; + } + // Early-dispatch: pre-staged tasks are released by doorbell // here, skipping the ready-queue round-trip entirely. bool early_handled = try_early_dispatch_release(slot_state, sink); @@ -1028,6 +1051,49 @@ struct PTO2SchedulerState { #endif } + 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 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 + if (slot_state.task == nullptr || slot_state.task->kind != PTO2TaskKind::GRAPH_NODE) { + outcome.stream_tasks_completed = 1; + return outcome; + } + + PTO2GraphExecution *execution = pto2_graph_execution_from_task(*slot_state.task); + if (execution == nullptr || !pto2_graph_execution_complete_node(*execution)) return outcome; + + PTO2TaskSlotState *outer = execution->outer_slot; + if (outer != nullptr) { +#if SIMPLER_SCHED_PROFILING + CompletionStats outer_stats = on_task_complete(*outer, thread_idx); + outcome.fanout_edges += static_cast(outer_stats.fanout_edges); + (void)on_task_release(*outer, thread_idx); +#else + outcome.fanout_edges += on_task_complete(*outer); + (void)on_task_release(*outer); +#endif + outcome.stream_tasks_completed = 1; + } + pto2_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. @@ -1041,6 +1107,52 @@ struct PTO2SchedulerState { return (prev + 1) == slot_state.total_required_subtasks; } + int32_t activate_prepared_graph_task(PTO2GraphExecution &execution) { + if (!pto2_graph_execution_activate(execution)) return 0; + int32_t roots = 0; + for (int32_t i = 0; i < execution.topology.root_count; ++i) { + uint16_t node_index = execution.topology.root_indices[i]; + if (node_index >= execution.node_count) continue; + route_ready_once(execution.nodes[node_index].slot); + roots++; + } + return roots; + } + + PTO2GraphMaterializeResult prepare_graph_task( + PTO2TaskSlotState &outer_slot, int32_t max_nodes = PTO2_GRAPH_MATERIALIZE_SLICE_NODES, + int32_t *nodes_materialized = nullptr + ) { + if (outer_slot.task == nullptr || outer_slot.task->kind != PTO2TaskKind::GRAPH) { + return PTO2GraphMaterializeResult::INVALID; + } + PTO2GraphSubmission *submission = pto2_graph_submission_from_task(*outer_slot.task); + if (submission == nullptr) return PTO2GraphMaterializeResult::INVALID; + PTO2GraphExecution *execution = pto2_graph_execution_localize(outer_slot); + if (execution == nullptr) return PTO2GraphMaterializeResult::INVALID; + + PTO2GraphMaterializeResult result = + pto2_graph_execution_materialize_slice(outer_slot, *execution, max_nodes, nodes_materialized); + PTO2GraphExecutionState state = execution->state.load(std::memory_order_acquire); + if (result == PTO2GraphMaterializeResult::PREPARED && state == PTO2GraphExecutionState::PREPARED && + submission->activation_requested.load(std::memory_order_acquire) != 0) { + activate_prepared_graph_task(*execution); + } + return result; + } + + int32_t activate_graph_task(PTO2TaskSlotState &outer_slot) { + if (outer_slot.task == nullptr || outer_slot.task->kind != PTO2TaskKind::GRAPH) return 0; + PTO2GraphSubmission *submission = pto2_graph_submission_from_task(*outer_slot.task); + if (submission == nullptr) return 0; + + submission->activation_requested.store(1, std::memory_order_release); + PTO2GraphExecution *execution = submission->local_execution.load(std::memory_order_acquire); + if (execution == nullptr) return 0; + if (execution->state.load(std::memory_order_acquire) != PTO2GraphExecutionState::PREPARED) return 0; + return activate_prepared_graph_task(*execution); + } + /** * Two-stage completion: second stage. * Called exactly once when all subtasks of a task are done (i.e., @@ -1078,14 +1190,20 @@ struct PTO2SchedulerState { PTO2_SCHED_CYCLE_START(); #endif + const bool graph_node = slot_state.task != nullptr && slot_state.task->kind == PTO2TaskKind::GRAPH_NODE; + PTO2DepListEntry *current = nullptr; + if (graph_node) { + slot_state.mark_completed(); + } else { #if SIMPLER_SCHED_PROFILING - slot_state.lock_fanout(lock_atomics, lock_wait); + slot_state.lock_fanout(lock_atomics, lock_wait); #else - slot_state.lock_fanout(); + slot_state.lock_fanout(); #endif - slot_state.mark_completed(); - PTO2DepListEntry *current = slot_state.fanout_head; // Protected by fanout_lock - slot_state.unlock_fanout(); + slot_state.mark_completed(); + current = slot_state.fanout_head; // Protected by fanout_lock + slot_state.unlock_fanout(); + } #if SIMPLER_SCHED_PROFILING lock_atomics += 3; // task_state.store + lifecycle_flags.fetch_or + unlock.store @@ -1114,8 +1232,7 @@ struct PTO2SchedulerState { // below; their dispatch_fanin propagation is collected here and replayed // after the walk, so no consumer's doorbell waits on a sibling's propagate. EarlyDispatchReleaseSink rel_sink; - while (current != nullptr) { - PTO2TaskSlotState &consumer_slot = *current->slot_state; + auto release_consumer = [&](PTO2TaskSlotState &consumer_slot) { #if SIMPLER_SCHED_PROFILING stats.fanout_edges++; if (release_fanin_and_check_ready(consumer_slot, fanout_atomics, push_wait, &rel_sink)) { @@ -1125,7 +1242,26 @@ struct PTO2SchedulerState { consumer_walk_count++; release_fanin_and_check_ready(consumer_slot, &rel_sink); #endif - current = current->next; + }; + if (graph_node) { + PTO2GraphExecution *execution = pto2_graph_execution_from_task(*slot_state.task); + int32_t node_index = slot_state.task->graph_node_index; + if (execution != nullptr && execution->nodes != nullptr && node_index >= 0 && + node_index < execution->node_count) { + uint32_t begin = execution->topology.fanout_offsets[node_index]; + uint32_t end = execution->topology.fanout_offsets[node_index + 1]; + for (uint32_t edge = begin; edge < end; ++edge) { + uint16_t consumer_index = execution->topology.fanout_indices[edge]; + if (consumer_index < execution->node_count) { + release_consumer(execution->nodes[consumer_index].slot); + } + } + } + } else { + while (current != nullptr) { + release_consumer(*current->slot_state); + current = current->next; + } } for (int i = 0; i < rel_sink.n; i++) { propagate_dispatch_fanin(*rel_sink.items[i]); @@ -1155,6 +1291,14 @@ struct PTO2SchedulerState { #else int32_t on_task_release(PTO2TaskSlotState &slot_state) { #endif + if (slot_state.task != nullptr && slot_state.task->kind == PTO2TaskKind::GRAPH_NODE) { + PTO2GraphExecution *execution = pto2_graph_execution_from_task(*slot_state.task); + if (execution != nullptr) pto2_graph_execution_retire_node(*execution); +#if SIMPLER_SCHED_PROFILING + g_sched_complete_count[thread_idx]++; +#endif + return 0; + } PTO2TaskPayload *payload = slot_state.payload; for_each_fanin_slot_state(*payload, [&](PTO2TaskSlotState *producer_slot_state) { #if SIMPLER_SCHED_PROFILING @@ -1222,9 +1366,9 @@ 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 if (*sink.deferred_release_count >= sink.deferred_release_capacity) { while (*sink.deferred_release_count > 0) { @@ -1238,7 +1382,7 @@ AsyncWaitList::try_inline_complete_locked(AsyncWaitList::DrainCompletionSink &si } } sink.deferred_release_slot_states[(*sink.deferred_release_count)++] = &slot_state; - sink.inline_completed++; + sink.inline_completed += outcome.stream_tasks_completed; return true; } @@ -1303,9 +1447,9 @@ 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 // Drain deferred_release in place when the buffer fills — same // overflow-drain idiom used by complete_slot_task's inline path @@ -1322,7 +1466,7 @@ inline AsyncPollResult AsyncWaitList::poll_and_complete( } } deferred_release_slot_states[deferred_release_count++] = entry.slot_state; - 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_completion.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp index dc39400a0a..acbb631c53 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp @@ -204,10 +204,11 @@ void SchedulerContext::complete_slot_task( // counter side-effects (g_sched_*_atomic_count[thread_idx], consumed // by the otc_* log lines). It returns CompletionStats whose // `fanout_edges` is the consumer-walk count. - consumers_resolved = sched_->on_task_complete(slot_state, thread_idx).fanout_edges; + PTO2SchedulerState::TaskCompletionOutcome outcome = sched_->complete_task(slot_state, thread_idx); #else - consumers_resolved = sched_->on_task_complete(slot_state); + PTO2SchedulerState::TaskCompletionOutcome outcome = sched_->complete_task(slot_state); #endif + consumers_resolved = outcome.fanout_edges; #if SIMPLER_DFX if (resolve_t0 != 0) { uint64_t resolve_t1 = get_sys_cnt_aicpu(); @@ -240,7 +241,7 @@ void SchedulerContext::complete_slot_task( } deferred_release_slot_states[deferred_release_count++] = &slot_state; } - completed_this_turn++; + completed_this_turn += outcome.stream_tasks_completed; } #if SIMPLER_DFX 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 a3353a1ab5..5fb2340341 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 @@ -1133,6 +1133,53 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ continue; } + // Graph control work never consumes an AICore. An outer Graph becomes + // activatable only after its external fanin is ready; materialization + // can run independently in bounded slices and therefore overlaps other + // scheduler work instead of creating one long expansion bar. + if (thread_idx < 3) { + // host_build_graph never reuses task-window slots during a run, so + // the Graph ready queue does not need a generation tag. The common + // ready routing path deliberately uses push(), whose tag is zero. + PTO2TaskSlotState *graph_slot = sched_->graph_ready_queue.pop(); + if (graph_slot != nullptr && graph_slot->task != nullptr && graph_slot->task->kind == PTO2TaskKind::GRAPH) { + (void)sched_->activate_graph_task(*graph_slot); + made_progress = true; + } + + uint64_t prepare_task_id = 0; + PTO2TaskSlotState *prepare_slot = sched_->graph_prepare_queue.pop_tagged(&prepare_task_id); + if (prepare_slot != nullptr && prepare_slot->task != nullptr && + prepare_slot->task->kind == PTO2TaskKind::GRAPH && prepare_slot->task->task_id.raw == prepare_task_id) { +#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; + PTO2GraphMaterializeResult result = + sched_->prepare_graph_task(*prepare_slot, PTO2_GRAPH_MATERIALIZE_SLICE_NODES, &nodes_materialized); + if (result == PTO2GraphMaterializeResult::PENDING || result == PTO2GraphMaterializeResult::BUSY) { + while (!sched_->graph_prepare_queue.push_tagged(prepare_slot, prepare_task_id)) { + SPIN_WAIT_HINT(); + } + } + if (nodes_materialized > 0 || result == PTO2GraphMaterializeResult::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) + ); + // Keep Graph expansion out of the next Dispatch envelope. + _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 @@ -1168,10 +1215,11 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ // OFF and the Resolve emit below is excluded. [[maybe_unused]] uint32_t dummy_consumers = 0; #if SIMPLER_SCHED_PROFILING - dummy_consumers = sched_->on_task_complete(dummy_slot, thread_idx).fanout_edges; + PTO2SchedulerState::TaskCompletionOutcome outcome = sched_->complete_task(dummy_slot, thread_idx); #else - dummy_consumers = sched_->on_task_complete(dummy_slot); + 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(); @@ -1203,8 +1251,8 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ #endif } } - int32_t prev = completed_tasks_.fetch_add(1, std::memory_order_relaxed); - last_progress_count = prev + 1; + 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) { 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 9983b2bab8..c4a7a97e92 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 @@ -117,6 +117,8 @@ PTO2SchedulerLayout PTO2SchedulerState::reserve_layout(DeviceArena &arena, int32 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); } @@ -162,6 +164,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], @@ -194,6 +204,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] @@ -215,6 +227,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]); } @@ -398,6 +412,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/tensormap_and_ringbuffer/host/runtime_maker.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp index c3f0092781..513c3facec 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp @@ -811,7 +811,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..a93a92e946 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/runtime.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/runtime.h @@ -308,6 +308,15 @@ 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; } }; // `dev` must be the first member so the narrowed H2D copy starts at offset 0. diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp index ab5a0b34f9..c090f94cc9 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp @@ -1219,7 +1219,7 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ sched_l2_swimlane_[thread_idx].sched_loop_count, dummy_consumers ); } - if (dummy_slot.active_mask.has_predicate()) { + if (dummy_slot.task_attrs.has_predicate()) { l2_swimlane_aicpu_record_predicated_skip( thread_idx, dummy_resolve_t0, sched_l2_swimlane_[thread_idx].sched_loop_count, dummy_slot.task->task_id.raw 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 e3b55486aa..6b6392b1f7 100644 --- a/src/a5/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a5/runtime/host_build_graph/host/runtime_maker.cpp @@ -384,7 +384,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/tensormap_and_ringbuffer/host/runtime_maker.cpp b/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp index c355f2e11c..507a1b3b61 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp @@ -786,7 +786,7 @@ static void store_prebuilt_runtime_image( 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/scheduler/scheduler_dispatch.cpp b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp index 2651d15323..ca23c9cb1b 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp @@ -1028,7 +1028,7 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ #endif #if SIMPLER_DFX if (dummy_resolve_t0 != 0) { - if (dummy_slot.active_mask.has_predicate()) { + if (dummy_slot.task_attrs.has_predicate()) { l2_swimlane_aicpu_record_predicated_skip( thread_idx, dummy_resolve_t0, sched_l2_swimlane_[thread_idx].sched_loop_count, dummy_slot.task->task_id.raw 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/l2_swimlane_profiling.h b/src/common/platform/include/common/l2_swimlane_profiling.h index 53344d93a4..60bbed19fc 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..d9cfa24681 100644 --- a/src/common/platform/include/host/l2_swimlane_collector.h +++ b/src/common/platform/include/host/l2_swimlane_collector.h @@ -408,6 +408,11 @@ class L2SwimlaneCollector : public profiling_common::ProfilerBase &records, uint64_t host_start_cycles, uint64_t host_end_cycles + ); + /** * Export collected records as a Chrome Trace Event JSON (swimlane view). * Writes /l2_swimlane_records.json — directory is captured at @@ -530,6 +535,12 @@ 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}; + // 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 69f91832db..03befb1315 100644 --- a/src/common/platform/onboard/host/c_api_shared.cpp +++ b/src/common/platform/onboard/host/c_api_shared.cpp @@ -589,7 +589,7 @@ int simpler_run( // the runner — no longer returned across this boundary. rc = runner->bind_callable_to_runtime( *r, callable_id, &g_host_api, args, config->runtime_env.ring_task_window, config->runtime_env.ring_heap, - config->runtime_env.ring_dep_pool + config->runtime_env.ring_dep_pool, config->enable_l2_swimlane ); } if (rc != 0) { diff --git a/src/common/platform/onboard/host/device_runner_base.cpp b/src/common/platform/onboard/host/device_runner_base.cpp index 7a47b011bc..7c41e21364 100644 --- a/src/common/platform/onboard/host/device_runner_base.cpp +++ b/src/common/platform/onboard/host/device_runner_base.cpp @@ -860,12 +860,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()) { @@ -896,7 +896,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 5720c7f77d..494038af32 100644 --- a/src/common/platform/onboard/host/device_runner_base.h +++ b/src/common/platform/onboard/host/device_runner_base.h @@ -362,12 +362,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 before device launch. * @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 3f716408fa..5e167dae0d 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 10c4519fd5..6a9d0e9032 100644 --- a/src/common/platform/shared/host/l2_swimlane_collector.cpp +++ b/src/common/platform/shared/host/l2_swimlane_collector.cpp @@ -895,6 +895,14 @@ 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 +) { + host_orch_phase_records_ = records; + host_orch_start_cycles_ = host_start_cycles; + host_orch_end_cycles_ = host_end_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 @@ -1046,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"; }; @@ -1072,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); @@ -1116,6 +1131,21 @@ 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 << " \"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"; @@ -1243,6 +1273,9 @@ 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; 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 963891ef4c..459601c382 100644 --- a/src/common/platform/sim/host/c_api_shared.cpp +++ b/src/common/platform/sim/host/c_api_shared.cpp @@ -545,7 +545,7 @@ int simpler_run( // host_orch_func_ptr + signature stay inside the runner. rc = runner->bind_callable_to_runtime( *r, callable_id, &g_host_api, args, config->runtime_env.ring_task_window, config->runtime_env.ring_heap, - config->runtime_env.ring_dep_pool + config->runtime_env.ring_dep_pool, config->enable_l2_swimlane ); } if (rc != 0) { diff --git a/src/common/platform/sim/host/device_runner_base.cpp b/src/common/platform/sim/host/device_runner_base.cpp index 023afc0113..f2df901ec2 100644 --- a/src/common/platform/sim/host/device_runner_base.cpp +++ b/src/common/platform/sim/host/device_runner_base.cpp @@ -474,12 +474,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()) { @@ -503,7 +503,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 b1ecca0ad5..fe7ba4838e 100644 --- a/src/common/platform/sim/host/device_runner_base.h +++ b/src/common/platform/sim/host/device_runner_base.h @@ -127,7 +127,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 ); uint64_t upload_chip_callable_buffer(const ChipCallable *callable); int release_chip_callable_buffer(uint64_t hash); 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..070cd44bc1 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/graph_execution/kernels/orchestration/graph_execution_orch.cpp @@ -0,0 +1,85 @@ +/* + * 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_ADD 0 +#define FUNC_ADD_SCALAR 1 +#define FUNC_MUL 2 + +namespace { + +void layer(const L2TaskArgs &args) { + const Tensor &a = args.tensor(0).ref(); + const Tensor &b = args.tensor(1).ref(); + const Tensor &output = args.tensor(2).ref(); + + uint32_t shape[1] = {a.shapes[0]}; + TensorCreateInfo intermediate(shape, 1, DataType::FLOAT32); + + L0TaskArgs add_args; + add_args.add_input(a, b); + add_args.add_output(intermediate); + TaskOutputTensors add_outputs = rt_submit_aiv_task(FUNC_ADD, add_args); + Tensor sum = add_outputs.get_ref(0); + + L0TaskArgs dynamic_args; + dynamic_args.add_input(sum); + dynamic_args.add_output(intermediate); + dynamic_args.add_scalar(args.scalar(0)); + TaskOutputTensors dynamic_outputs = rt_submit_aiv_task(FUNC_ADD_SCALAR, dynamic_args); + Tensor dynamic_sum = dynamic_outputs.get_ref(0); + + L0TaskArgs static_args; + static_args.add_input(sum); + static_args.add_output(intermediate); + static_args.add_scalar(2.0F); + PTO2TaskId explicit_dep = dynamic_outputs.task_id(); + static_args.set_dependencies(&explicit_dep, 1); + TaskOutputTensors static_outputs = rt_submit_aiv_task(FUNC_ADD_SCALAR, static_args); + Tensor static_sum = static_outputs.get_ref(0); + + L0TaskArgs mul_args; + mul_args.add_input(dynamic_sum, static_sum); + mul_args.add_output(output); + rt_submit_aiv_task(FUNC_MUL, mul_args); +} + +void submit_layer(const Tensor &a, const Tensor &b, const Tensor &output, float dynamic_scalar) { + L2TaskArgs args; + args.add_input(a, b); + args.add_output(output); + args.add_scalar(dynamic_scalar); + rt_submit_graph(&layer, args); +} + +} // 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(); + submit_layer(a, b, args.tensor(2).ref(), 1.0F); // cache miss: record four tasks + submit_layer(a, b, args.tensor(3).ref(), 3.0F); // cache hit: one Graph task + submit_layer(a, b, args.tensor(4).ref(), 5.0F); // cache hit: one Graph task +} + +} // 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..62d0b3c8dc --- /dev/null +++ b/tests/st/a2a3/host_build_graph/graph_execution/test_graph_execution.py @@ -0,0 +1,78 @@ +#!/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 TaskArgs.""" + +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", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4, "block_dim": 3}, + "params": {}, + }, + ] + + def generate_args(self, params): + size = 128 * 128 + return TaskArgsBuilder( + Tensor("a", torch.full((size,), 2.0, dtype=torch.float32)), + Tensor("b", torch.full((size,), 3.0, dtype=torch.float32)), + Tensor("output_1", torch.zeros(size, dtype=torch.float32)), + Tensor("output_3", torch.zeros(size, dtype=torch.float32)), + Tensor("output_5", torch.zeros(size, dtype=torch.float32)), + ) + + def compute_golden(self, args, params): + base = args.a + args.b + args.output_1[:] = (base + 1.0) * (base + 2.0) + args.output_3[:] = (base + 3.0) * (base + 2.0) + args.output_5[:] = (base + 5.0) * (base + 2.0) + + +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..7bfe3b13ef 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,178 @@ 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_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