dfx: capture outer Graph tasks in host dep-gen

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<int32_t, PTO2_SUBTASK_SLOT_COUNT> 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;
 

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<int32_t, PTO2_SUBTASK_SLOT_COUNT> 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;
 

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__)

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__)