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 a71f57c332..7c1b5ae9f7 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 @@ -72,6 +72,20 @@ __attribute__((weak, visibility("hidden"))) void dep_gen_host_graph_begin_task( uint64_t, bool, bool, const int32_t[3], int32_t, int32_t, const TensorRef *, const TensorArgType * ) {} __attribute__((weak, visibility("hidden"))) void dep_gen_host_graph_end_task() {} + +// Raises the two edge kinds compute_task_fanin can discover, for the capture +// instantiation. Shared by the ordinary submit path and the outer GRAPH task so +// both describe an edge the same way. +struct DepGraphAnnotate { + void creator(int32_t arg_idx, const ChipTensor &consumer, PTO2TaskId producer) const { + dep_gen_host_graph_add_creator_edge(producer.raw, arg_idx, consumer); + } + void tensormap( + int32_t arg_idx, const ChipTensor &consumer, const PTO2TensorMapEntry &entry, OverlapStatus overlap + ) const { + dep_gen_host_graph_add_tensormap_edge(entry.producer_task_id.raw, arg_idx, consumer, entry, overlap); + } +}; __attribute__((weak, visibility("hidden"))) void dep_gen_host_graph_add_explicit_edge(uint64_t) {} __attribute__((weak, visibility("hidden"))) void dep_gen_host_graph_add_creator_edge(uint64_t, int32_t, const ChipTensor &) {} @@ -1229,16 +1243,6 @@ static TaskOutputTensors submit_task_common( // The capture branch instantiates compute_task_fanin with a live Annotate; // the plain branch keeps the un-annotated instantiation the hot path had. if (capture_dep_graph) { - struct DepGraphAnnotate { - void creator(int32_t arg_idx, const ChipTensor &consumer, PTO2TaskId producer) const { - dep_gen_host_graph_add_creator_edge(producer.raw, arg_idx, consumer); - } - void tensormap( - int32_t arg_idx, const ChipTensor &consumer, const PTO2TensorMapEntry &entry, OverlapStatus overlap - ) const { - dep_gen_host_graph_add_tensormap_edge(entry.producer_task_id.raw, arg_idx, consumer, entry, overlap); - } - }; const bool ok = compute_task_fanin(dep_inputs, orch->tensor_map, orch->in_manual_scope(), runtime_emit, DepGraphAnnotate{}); // STEP 3 is this task's last capture point, so the entry closes here @@ -1546,7 +1550,32 @@ bool graph_submit_outer( PTO2TaskSlotState *producer = &ring.get_slot_state_by_slot(producer_slot); return append_fanin_or_fail(orch, producer_id.ring(), producer_slot, producer, producer_id, &fanin_builder); }; - if (!compute_task_fanin(boundary_inputs, orch->tensor_map, orch->in_manual_scope(), emit)) return false; + // An outer GRAPH task is a ring task like any other, so the dependency graph + // has to carry it: without this the whole Graph — and every edge into it — + // is absent from deps.json, leaving a run of 40 replays described by only its + // handful of non-Graph tasks. It dispatches no kernel of its own and the + // sub-DAG it replays owns no ring slots, so what is captured is its boundary: + // the args it consumes and the edges those produce. + const bool capture_dep_graph = dep_gen_host_graph_enabled(); + if (capture_dep_graph) { + const std::array kernel_ids_capture{ + INVALID_KERNEL_ID, + INVALID_KERNEL_ID, + INVALID_KERNEL_ID, + }; + dep_gen_host_graph_begin_task( + task_id.raw, orch->in_manual_scope(), /*early_dispatch=*/false, kernel_ids_capture.data(), + slot.logical_block_num, args.tensor_count(), args.tensor_data(), args.tag_data() + ); + const bool ok = + compute_task_fanin(boundary_inputs, orch->tensor_map, orch->in_manual_scope(), emit, DepGraphAnnotate{}); + // The task's last capture point, so the entry closes whether or not the + // fanin computation succeeded. + dep_gen_host_graph_end_task(); + if (!ok) return false; + } else if (!compute_task_fanin(boundary_inputs, orch->tensor_map, orch->in_manual_scope(), emit)) { + return false; + } register_task_outputs(boundary_inputs, task_id, orch->tensor_map, orch->in_manual_scope()); payload.fanin_count = fanin_builder.count; diff --git a/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp b/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp index a71f57c332..7c1b5ae9f7 100644 --- a/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp +++ b/src/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp @@ -72,6 +72,20 @@ __attribute__((weak, visibility("hidden"))) void dep_gen_host_graph_begin_task( uint64_t, bool, bool, const int32_t[3], int32_t, int32_t, const TensorRef *, const TensorArgType * ) {} __attribute__((weak, visibility("hidden"))) void dep_gen_host_graph_end_task() {} + +// Raises the two edge kinds compute_task_fanin can discover, for the capture +// instantiation. Shared by the ordinary submit path and the outer GRAPH task so +// both describe an edge the same way. +struct DepGraphAnnotate { + void creator(int32_t arg_idx, const ChipTensor &consumer, PTO2TaskId producer) const { + dep_gen_host_graph_add_creator_edge(producer.raw, arg_idx, consumer); + } + void tensormap( + int32_t arg_idx, const ChipTensor &consumer, const PTO2TensorMapEntry &entry, OverlapStatus overlap + ) const { + dep_gen_host_graph_add_tensormap_edge(entry.producer_task_id.raw, arg_idx, consumer, entry, overlap); + } +}; __attribute__((weak, visibility("hidden"))) void dep_gen_host_graph_add_explicit_edge(uint64_t) {} __attribute__((weak, visibility("hidden"))) void dep_gen_host_graph_add_creator_edge(uint64_t, int32_t, const ChipTensor &) {} @@ -1229,16 +1243,6 @@ static TaskOutputTensors submit_task_common( // The capture branch instantiates compute_task_fanin with a live Annotate; // the plain branch keeps the un-annotated instantiation the hot path had. if (capture_dep_graph) { - struct DepGraphAnnotate { - void creator(int32_t arg_idx, const ChipTensor &consumer, PTO2TaskId producer) const { - dep_gen_host_graph_add_creator_edge(producer.raw, arg_idx, consumer); - } - void tensormap( - int32_t arg_idx, const ChipTensor &consumer, const PTO2TensorMapEntry &entry, OverlapStatus overlap - ) const { - dep_gen_host_graph_add_tensormap_edge(entry.producer_task_id.raw, arg_idx, consumer, entry, overlap); - } - }; const bool ok = compute_task_fanin(dep_inputs, orch->tensor_map, orch->in_manual_scope(), runtime_emit, DepGraphAnnotate{}); // STEP 3 is this task's last capture point, so the entry closes here @@ -1546,7 +1550,32 @@ bool graph_submit_outer( PTO2TaskSlotState *producer = &ring.get_slot_state_by_slot(producer_slot); return append_fanin_or_fail(orch, producer_id.ring(), producer_slot, producer, producer_id, &fanin_builder); }; - if (!compute_task_fanin(boundary_inputs, orch->tensor_map, orch->in_manual_scope(), emit)) return false; + // An outer GRAPH task is a ring task like any other, so the dependency graph + // has to carry it: without this the whole Graph — and every edge into it — + // is absent from deps.json, leaving a run of 40 replays described by only its + // handful of non-Graph tasks. It dispatches no kernel of its own and the + // sub-DAG it replays owns no ring slots, so what is captured is its boundary: + // the args it consumes and the edges those produce. + const bool capture_dep_graph = dep_gen_host_graph_enabled(); + if (capture_dep_graph) { + const std::array kernel_ids_capture{ + INVALID_KERNEL_ID, + INVALID_KERNEL_ID, + INVALID_KERNEL_ID, + }; + dep_gen_host_graph_begin_task( + task_id.raw, orch->in_manual_scope(), /*early_dispatch=*/false, kernel_ids_capture.data(), + slot.logical_block_num, args.tensor_count(), args.tensor_data(), args.tag_data() + ); + const bool ok = + compute_task_fanin(boundary_inputs, orch->tensor_map, orch->in_manual_scope(), emit, DepGraphAnnotate{}); + // The task's last capture point, so the entry closes whether or not the + // fanin computation succeeded. + dep_gen_host_graph_end_task(); + if (!ok) return false; + } else if (!compute_task_fanin(boundary_inputs, orch->tensor_map, orch->in_manual_scope(), emit)) { + return false; + } register_task_outputs(boundary_inputs, task_id, orch->tensor_map, orch->in_manual_scope()); payload.fanin_count = fanin_builder.count; 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 index b578eed10c..88d466d041 100644 --- 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 @@ -9,10 +9,13 @@ # ----------------------------------------------------------------------------------------------------------- """Graph Execution replays dynamic tensor and scalar bindings across config-keyed definitions.""" +import json + import torch from simpler.task_interface import ArgDirection as D from simpler_setup import SceneTestCase, TaskArgsBuilder, TensorArg, scene_test +from simpler_setup.scene_test import _outputs_dir, _sanitize_for_filename @scene_test(level=2, runtime="host_build_graph") @@ -80,6 +83,53 @@ def compute_golden(self, args, params): args.output_3[:] = (base + 100.0) * (base + 4.0) args.output_5[:] = (base + 5.0) * (base + 6.0) + def test_run(self, st_platform, st_worker, request): + matched_cases = self._matching_cases(st_platform, request) + prior_mtimes = {case["name"]: self._matching_output_mtimes(case) for case in matched_cases} + super().test_run(st_platform, st_worker, request) + if not self._effective_enable_dep_gen(request): + return + for case in matched_cases: + self._validate_outer_graph_capture(case, prior_mtimes[case["name"]]) + + @staticmethod + def _matching_output_mtimes(case): + safe_label = _sanitize_for_filename(f"TestGraphExecutionHostBuildGraph_{case['name']}") + return {path: path.stat().st_mtime_ns for path in _outputs_dir().glob(f"{safe_label}_*")} + + @classmethod + def _validate_outer_graph_capture(cls, case, prior_mtimes): + matches = [ + path + for path, mtime_ns in cls._matching_output_mtimes(case).items() + if path not in prior_mtimes or mtime_ns > prior_mtimes[path] + ] + assert matches, f"no dep-gen output directory created for {case['name']}" + deps_path = max(matches, key=lambda path: path.stat().st_mtime_ns) / "deps.json" + assert deps_path.exists(), f"deps.json missing for {case['name']}" + with deps_path.open() as f: + deps = json.load(f) + + tasks = deps.get("tasks", []) + assert [int(task["task_id"]) for task in tasks] == list(range(8)), ( + "dep-gen must capture the seed, two outer GRAPH tasks, and the five-task fallback Graph body" + ) + outer_graph_ids = { + int(task["task_id"]) + for task in tasks + if task.get("kernel_ids") == [-1, -1, -1] and len(task.get("args", [])) == 3 + } + assert outer_graph_ids == {1, 7}, f"outer GRAPH tasks missing from deps.json: {outer_graph_ids}" + + outer_edges = { + (int(edge["pred"]), int(edge["succ"]), edge.get("source")) + for edge in deps.get("edges", []) + if int(edge["succ"]) in outer_graph_ids + } + assert outer_edges == {(0, 1, "creator"), (0, 7, "creator")}, ( + f"outer GRAPH boundary edges differ from runtime dependencies: {outer_edges}" + ) + if __name__ == "__main__": SceneTestCase.run_module(__name__) diff --git a/tests/st/a2a3/host_build_graph/graph_execution/test_graph_execution_mix_spmd.py b/tests/st/a2a3/host_build_graph/graph_execution/test_graph_execution_mix_spmd.py index 3e5c490201..fe7d61c693 100644 --- a/tests/st/a2a3/host_build_graph/graph_execution/test_graph_execution_mix_spmd.py +++ b/tests/st/a2a3/host_build_graph/graph_execution/test_graph_execution_mix_spmd.py @@ -9,10 +9,13 @@ # ----------------------------------------------------------------------------------------------------------- """Graph Execution preserves MIX active slots and multi-block SPMD metadata.""" +import json + import torch from simpler.task_interface import ArgDirection as D from simpler_setup import SceneTestCase, TaskArgsBuilder, TensorArg, scene_test +from simpler_setup.scene_test import _outputs_dir, _sanitize_for_filename FLOATS_PER_CACHE_LINE = 16 SLOTS_PER_BLOCK = 3 @@ -91,6 +94,43 @@ def compare_outputs(self, test_args, golden_args, output_names, params): expected[begin : begin + SLOTS_PER_BLOCK] = float(block_idx) assert torch.equal(cache_line_heads, expected) + def test_run(self, st_platform, st_worker, request): + matched_cases = self._matching_cases(st_platform, request) + prior_mtimes = {case["name"]: self._matching_output_mtimes(case) for case in matched_cases} + super().test_run(st_platform, st_worker, request) + if not self._effective_enable_dep_gen(request): + return + for case in matched_cases: + self._validate_outer_graph_capture(case, prior_mtimes[case["name"]]) + + @staticmethod + def _matching_output_mtimes(case): + safe_label = _sanitize_for_filename(f"TestGraphExecutionMixSpmdHostBuildGraph_{case['name']}") + return {path: path.stat().st_mtime_ns for path in _outputs_dir().glob(f"{safe_label}_*")} + + @classmethod + def _validate_outer_graph_capture(cls, case, prior_mtimes): + matches = [ + path + for path, mtime_ns in cls._matching_output_mtimes(case).items() + if path not in prior_mtimes or mtime_ns > prior_mtimes[path] + ] + assert matches, f"no dep-gen output directory created for {case['name']}" + deps_path = max(matches, key=lambda path: path.stat().st_mtime_ns) / "deps.json" + assert deps_path.exists(), f"deps.json missing for {case['name']}" + with deps_path.open() as f: + deps = json.load(f) + + tasks = deps.get("tasks", []) + assert [int(task["task_id"]) for task in tasks] == [0, 1, 2] + outer_graph_ids = { + int(task["task_id"]) + for task in tasks + if task.get("kernel_ids") == [-1, -1, -1] and len(task.get("args", [])) == 1 + } + assert outer_graph_ids == {0, 1, 2} + assert deps.get("edges", []) == [] + if __name__ == "__main__": SceneTestCase.run_module(__name__)