Attach the device before drain and retire the dead dep_gen capture API

docs/dfx/dep-gen.md

@@ -91,8 +91,8 @@ nothing to capture-then-reconstruct.
   (`dep_gen_host_graph_active()` tells the runner). Nothing is dropped under
   back-pressure because nothing is streamed.
 - **Output.** The same `deps.json`, written during the device-runner drain.
-  After prepare's host orchestration builds the graph, the phased runtime moves
-  it into run-owned storage and the executor adopts it before enqueue.
+  The graph is thread-local, and prepare's host orchestration and the drain that
+  emits it both run on the child progress loop's single thread.
 
 ---
 
@@ -395,7 +395,7 @@ list; only the dep_gen replay graph loses the tail.
 | Capture call site (device-orch) | `src/{a2a3,a5}/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp` `submit_task_common` | One conditional block that snapshots inputs into the ring when `is_dep_gen_enabled()`; fires for both `submit_task` and `submit_dummy_task`. The schema carries `kernel_ids[3] = {aic, aiv0, aiv1}` so the swimlane post-processor can resolve `task_id → kernel` from `deps.json` at level=1 where the AICore record is the sole device-side identity source. Inactive subslots stay at `INVALID_KERNEL_ID = -1`. It also carries the SPMD logical block num (`block_num` on a2a3, `core_num` on a5's launch spec) as `tasks[].block_num`. |
 | Replay | `src/{a2a3,a5}/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.{h,cpp}` | Pure CPU; runs dual-pass differential replay — `compute_task_fanin` (oracle) + inlined STEP A/B mirror (annotated) against two `PTO2TensorMap` instances. Emits `deps.json` when both passes agree per record. Platform-agnostic — a5 reuses the a2a3 source verbatim. |
 | Host-direct capture (host-orch) | `src/a2a3/runtime/host_build_graph/runtime/dep_gen_host_graph.h`, `src/a2a3/runtime/host_build_graph/host/dep_gen_host_graph.cpp` | Task / tensor / edge tables filled from `submit_task_common` + `compute_task_fanin`'s `Annotate` hooks (`src/a2a3/runtime/host_build_graph/runtime/pto_dep_compute.h`), reset per orchestration by `run_host_orchestration`, serialized by the same `deps.json` writer. The runtime translation unit carries weak no-op fallbacks so the AICPU build links without it. |
-| Device-runner hookup | `src/{a2a3,a5}/platform/{onboard,sim}/host/device_runner.cpp` | `dep_gen_host_graph_active()` picks the shape: host-orch moves the prepare thread's completed capture into run-owned storage, adopts it on the executor, calls `dep_gen_host_graph_emit(deps_path)` at teardown, and skips collector init/start/reconcile entirely; device-orch calls `dep_gen_replay_emit_deps_json(records.data(), records.size(), deps_path)` post-`reconcile_counters`. The c_api latches the CallConfig before the bind so host capture is armed before the orchestration it records. |
+| Device-runner hookup | `src/{a2a3,a5}/platform/{onboard,sim}/host/device_runner.cpp` | `dep_gen_host_graph_active()` picks the shape: host-orch calls `dep_gen_host_graph_emit(deps_path)` at teardown — reading the thread-local graph its own orchestration built on the same child progress thread — and skips collector init/start/reconcile entirely; device-orch calls `dep_gen_replay_emit_deps_json(records.data(), records.size(), deps_path)` post-`reconcile_counters`. The c_api latches the CallConfig before the bind so host capture is armed before the orchestration it records. |
 | Viewer | `simpler_setup/tools/deps_viewer.py` | `deps.json` → text (default) or pan/zoom HTML |
 | Test | `tests/st/{a2a3,a5}/tensormap_and_ringbuffer/dfx/dep_gen/test_dep_gen.py` + `test_dep_gen_chain.py`, `tests/st/a2a3/host_build_graph/dfx/dep_gen/test_dep_gen.py` | Smoke test + 6-edge validation against `vector_example` orchestration (both platforms share byte-identical orchestration code). The host_build_graph case runs the *same* orchestration through host-direct capture and asserts the same 6 edges, so a divergence between the two shapes fails a test. |
 

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

@@ -717,8 +717,8 @@ int DeviceRunner::reap_run(unsigned slot) {
     // order (mgmt's final-drain pass into L2 has poll as its consumer).
     teardown_shared_collectors_after_run();
 
-    // a2a3-only dep_gen teardown: host-orch emits the graph snapshot adopted
-    // from the prepare thread; device-orch stops the collector, reconciles the
+    // a2a3-only dep_gen teardown: host-orch emits the graph its orchestration
+    // built on this same thread; device-orch stops the collector, reconciles the
     // ring, and replays the records.
     if (enable_dep_gen_) {
         const std::string deps = make_deps_json_path(output_prefix_);

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

@@ -671,7 +671,7 @@ int DeviceRunner::drain_execution(ActiveExecution &) {
         pmu_collector_.reconcile_counters();
     }
 
-    // Host-orch emits the graph snapshot adopted from the prepare thread;
+    // Host-orch emits the graph its orchestration built on this same thread;
     // device-orch stops the collector, reconciles the ring, and replays.
     if (enable_dep_gen_) {
         const std::string deps = make_deps_json_path(output_prefix_);

src/a2a3/runtime/host_build_graph/host/dep_gen_host_graph.cpp

@@ -195,7 +195,7 @@ void fill_producer(EdgeAnnot &e, const PTO2TensorMapEntry &entry) {
 }
 
 // ---------------------------------------------------------------------------
-// Capture state — thread-local while built, then moved with its native run
+// Capture state — thread-local, built and emitted on the same thread
 // ---------------------------------------------------------------------------
 
 struct HostGraphState {
@@ -495,25 +495,6 @@ extern "C" void dep_gen_host_graph_set_enabled(bool enable) { state().enabled =
 
 extern "C" bool dep_gen_host_graph_active() { return true; }
 
-extern "C" void *dep_gen_host_graph_take_capture() {
-    HostGraphState &current = state();
-    if (!current.enabled) return nullptr;
-    auto *capture = new HostGraphState(std::move(current));
-    current = HostGraphState{};
-    return capture;
-}
-
-extern "C" void dep_gen_host_graph_adopt_capture(void *capture) noexcept {
-    if (capture == nullptr) return;
-    auto *captured_state = static_cast<HostGraphState *>(capture);
-    state() = std::move(*captured_state);
-    delete captured_state;
-}
-
-extern "C" void dep_gen_host_graph_destroy_capture(void *capture) noexcept {
-    delete static_cast<HostGraphState *>(capture);
-}
-
 extern "C" int dep_gen_host_graph_emit(const char *deps_json_path) {
     if (deps_json_path == nullptr) {
         LOG_ERROR("dep_gen host graph: null deps_json_path");
@@ -523,11 +504,9 @@ extern "C" int dep_gen_host_graph_emit(const char *deps_json_path) {
     if (!s.captured) {
         // An empty graph here is not "the orchestration submitted nothing" —
         // begin_task() would have set captured even for a graph of one task.
-        // It means capture was never armed or the run-owned snapshot was not
-        // adopted onto this progress thread before teardown.
-        LOG_ERROR(
-            "dep_gen host graph: no capture was adopted on this thread — deps.json not written to %s", deps_json_path
-        );
+        // The graph is thread-local, so it means either capture was never armed
+        // or this run's orchestration ran on a different thread than this emit.
+        LOG_ERROR("dep_gen host graph: no capture on this thread — deps.json not written to %s", deps_json_path);
         return -3;
     }
     if (!write_deps_json(deps_json_path, s.tasks, s.tensors, s.edges)) {

src/a2a3/runtime/host_build_graph/runtime/dep_gen_host_graph.h

@@ -31,17 +31,17 @@
  *   end_task()              — closes the task, after its last dependency step
  *
  * Control surface, called from the device runner (same host_runtime.so):
- *   set_enabled() / active() / take_capture() / adopt_capture() / emit()
+ *   set_enabled() / active() / emit()
  *
  * The runtime translation unit links weak no-op fallbacks (pto_orchestrator.cpp)
  * so the AICPU build, which has no host graph, resolves without this .cpp.
  *
- * The graph is per-thread state while it is being built. After bind, prepare
- * moves the completed graph into run-owned storage; launch adopts that snapshot
- * into the progress thread's local state before enqueue, and drain emits it.
- * This keeps capture lock-free while allowing serialized lifecycle calls to
- * use different host threads and preventing two prepared contexts on one
- * thread from overwriting one another.
+ * The graph lives in thread-local state, so capture is lock-free and two
+ * prepared contexts on different threads cannot overwrite one another. Emit
+ * reads the calling thread's state, so a run's orchestration (which builds the
+ * graph) and its drain (which emits it) must land on the same thread. The child
+ * progress loop satisfies this by being single-threaded; emit returns -3 if the
+ * invariant is ever broken.
  *
  * Per-task producer dedup mirrors PTO2FaninBuilder, which keys on (ring, slot);
  * this keys on producer task id. The two agree only because host_build_graph is
@@ -124,15 +124,6 @@ void dep_gen_host_graph_set_enabled(bool enable);
  */
 bool dep_gen_host_graph_active();
 
-/** Move the current thread's capture into an opaque, caller-owned snapshot. */
-void *dep_gen_host_graph_take_capture();
-
-/** Adopt and consume a snapshot on the current execution thread. */
-void dep_gen_host_graph_adopt_capture(void *capture) noexcept;
-
-/** Destroy a snapshot that will not be launched. */
-void dep_gen_host_graph_destroy_capture(void *capture) noexcept;
-
 /**
  * Write the captured graph to `deps_json_path`. Returns 0 on success, non-zero
  * if capture was off/empty or the file could not be written.

src/a5/runtime/host_build_graph/host/dep_gen_host_graph.cpp

@@ -195,7 +195,7 @@ void fill_producer(EdgeAnnot &e, const PTO2TensorMapEntry &entry) {
 }
 
 // ---------------------------------------------------------------------------
-// Capture state — thread-local while built, then moved with its native run
+// Capture state — thread-local, built and emitted on the same thread
 // ---------------------------------------------------------------------------
 
 struct HostGraphState {
@@ -495,25 +495,6 @@ extern "C" void dep_gen_host_graph_set_enabled(bool enable) { state().enabled =
 
 extern "C" bool dep_gen_host_graph_active() { return true; }
 
-extern "C" void *dep_gen_host_graph_take_capture() {
-    HostGraphState &current = state();
-    if (!current.enabled) return nullptr;
-    auto *capture = new HostGraphState(std::move(current));
-    current = HostGraphState{};
-    return capture;
-}
-
-extern "C" void dep_gen_host_graph_adopt_capture(void *capture) noexcept {
-    if (capture == nullptr) return;
-    auto *captured_state = static_cast<HostGraphState *>(capture);
-    state() = std::move(*captured_state);
-    delete captured_state;
-}
-
-extern "C" void dep_gen_host_graph_destroy_capture(void *capture) noexcept {
-    delete static_cast<HostGraphState *>(capture);
-}
-
 extern "C" int dep_gen_host_graph_emit(const char *deps_json_path) {
     if (deps_json_path == nullptr) {
         LOG_ERROR("dep_gen host graph: null deps_json_path");
@@ -523,11 +504,9 @@ extern "C" int dep_gen_host_graph_emit(const char *deps_json_path) {
     if (!s.captured) {
         // An empty graph here is not "the orchestration submitted nothing" —
         // begin_task() would have set captured even for a graph of one task.
-        // It means capture was never armed or the run-owned snapshot was not
-        // adopted onto this progress thread before teardown.
-        LOG_ERROR(
-            "dep_gen host graph: no capture was adopted on this thread — deps.json not written to %s", deps_json_path
-        );
+        // The graph is thread-local, so it means either capture was never armed
+        // or this run's orchestration ran on a different thread than this emit.
+        LOG_ERROR("dep_gen host graph: no capture on this thread — deps.json not written to %s", deps_json_path);
         return -3;
     }
     if (!write_deps_json(deps_json_path, s.tasks, s.tensors, s.edges)) {

src/a5/runtime/host_build_graph/runtime/dep_gen_host_graph.h

@@ -31,17 +31,17 @@
  *   end_task()              — closes the task, after its last dependency step
  *
  * Control surface, called from the device runner (same host_runtime.so):
- *   set_enabled() / active() / take_capture() / adopt_capture() / emit()
+ *   set_enabled() / active() / emit()
  *
  * The runtime translation unit links weak no-op fallbacks (pto_orchestrator.cpp)
  * so the AICPU build, which has no host graph, resolves without this .cpp.
  *
- * The graph is per-thread state while it is being built. After bind, prepare
- * moves the completed graph into run-owned storage; launch adopts that snapshot
- * into the progress thread's local state before enqueue, and drain emits it.
- * This keeps capture lock-free while allowing serialized lifecycle calls to
- * use different host threads and preventing two prepared contexts on one
- * thread from overwriting one another.
+ * The graph lives in thread-local state, so capture is lock-free and two
+ * prepared contexts on different threads cannot overwrite one another. Emit
+ * reads the calling thread's state, so a run's orchestration (which builds the
+ * graph) and its drain (which emits it) must land on the same thread. The child
+ * progress loop satisfies this by being single-threaded; emit returns -3 if the
+ * invariant is ever broken.
  *
  * Per-task producer dedup mirrors PTO2FaninBuilder, which keys on (ring, slot);
  * this keys on producer task id. The two agree only because host_build_graph is
@@ -123,15 +123,6 @@ void dep_gen_host_graph_set_enabled(bool enable);
  */
 bool dep_gen_host_graph_active();
 
-/** Move the current thread's capture into an opaque, caller-owned snapshot. */
-void *dep_gen_host_graph_take_capture();
-
-/** Adopt and consume a snapshot on the current execution thread. */
-void dep_gen_host_graph_adopt_capture(void *capture) noexcept;
-
-/** Destroy a snapshot that will not be launched. */
-void dep_gen_host_graph_destroy_capture(void *capture) noexcept;
-
 /**
  * Write the captured graph to `deps_json_path`. Returns 0 on success, non-zero
  * if capture was off/empty or the file could not be written.

src/common/platform/onboard/host/c_api_shared.cpp

@@ -816,12 +816,24 @@ int simpler_wait_run(DeviceContextHandle ctx, RuntimeHandle runtime) {
     NativeRunPhase phase = state->phase.load(std::memory_order_acquire);
     if (phase == NativeRunPhase::Prepared) return -1;
     if (phase == NativeRunPhase::Complete) return state->completion_rc;
+    // drain_execution() synchronizes and destroys streams, reads device memory
+    // and frees device allocations, all of which need this thread's CANN
+    // device context. rtSetDevice is idempotent on an already-attached thread.
     int drain_rc = -1;
     try {
-        if (state->active_execution != nullptr) {
-            drain_rc = state->runner->drain_execution(*state->active_execution);
+        drain_rc = state->runner->attach_current_thread(state->runner->device_id());
+        if (drain_rc != 0) {
+            LOG_ERROR("simpler_wait_run: attach_current_thread failed: %d (%s)", drain_rc, state->trace_attrs);
+        } else {
+            drain_rc = -1;
+            if (state->active_execution != nullptr) {
+                drain_rc = state->runner->drain_execution(*state->active_execution);
+            }
         }
-    } catch (...) {}
+    } catch (...) {
+        drain_rc = -1;
+        LOG_ERROR("simpler_wait_run: drain threw (%s)", state->trace_attrs);
+    }
     if (state->completion_rc == 0) state->completion_rc = drain_rc;
     state->phase.store(NativeRunPhase::Complete, std::memory_order_release);
     emit_native_run_runner_wall(state);
@@ -842,11 +854,28 @@ int simpler_finalize_run(DeviceContextHandle ctx, RuntimeHandle runtime) {
 
     int execution_rc = state->completion_rc;
     const bool launched = state->active_execution != nullptr;
+    // Both drain_execution() and validate_runtime_impl() touch the device, so
+    // the attach covers each of them. rtSetDevice is idempotent on an
+    // already-attached thread.
+    int attach_rc = -1;
+    try {
+        attach_rc = state->runner->attach_current_thread(state->runner->device_id());
+    } catch (...) {
+        attach_rc = -1;
+    }
+    if (attach_rc != 0) {
+        LOG_ERROR("simpler_finalize_run: attach_current_thread failed: %d (%s)", attach_rc, state->trace_attrs);
+    }
     if (phase == NativeRunPhase::Running && launched) {
-        int drain_rc = -1;
-        try {
-            drain_rc = state->runner->drain_execution(*state->active_execution);
-        } catch (...) {}
+        int drain_rc = attach_rc;
+        if (attach_rc == 0) {
+            drain_rc = -1;
+            try {
+                drain_rc = state->runner->drain_execution(*state->active_execution);
+            } catch (...) {
+                LOG_ERROR("simpler_finalize_run: drain_execution threw (%s)", state->trace_attrs);
+            }
+        }
         if (execution_rc == 0) execution_rc = drain_rc;
         state->completion_rc = execution_rc;
         state->phase.store(NativeRunPhase::Complete, std::memory_order_release);
@@ -856,7 +885,6 @@ int simpler_finalize_run(DeviceContextHandle ctx, RuntimeHandle runtime) {
     int validation_rc = -1;
     try {
         if (!launched) state->runtime.set_gm_sm_ptr(nullptr);
-        int attach_rc = state->runner->attach_current_thread(state->runner->device_id());
         if (attach_rc == 0) {
             {
                 STRACE("simpler_run.validate");