[Feature] Introduce reusable PreparedCall with production TMR program-path integration

[ChaoWao]

Summary

Introduce an immutable, reusable PreparedCall containing the resolved host-side result of a TMR invocation, and make the existing program execution path actually consume it.

This is the next bounded milestone after merged #2176, derived from the shared program/kernel execution design. Work from current main and selectively reuse existing implementations; existing PR numbering or dependency stacks do not dictate the design or merge order.

The deliverable is a production-used preparation result with demonstrable isolation and reuse. A types-only implementation or standalone codec with fake consumers does not complete this issue.

Motivation / Use Case

The target composition is:

program: device ownership + tensor bindings/transfers ─┐
                                                     ├→ prepare → PreparedCall → submit → Submission
kernel: caller-owned device arguments ────────────────┘

Kernel execution may own resident runtime metadata and internal workspaces while borrowing the device and tensor storage. Program mode adds device ownership and tensor parameter management around that common execution capability.

At the inspected main baseline fd5f2d93fc55c784a002049577b0015ebfa4194e:

• Add: persistent kernel-context execution resources #2176 establishes cold-path resources. prepare_kernel_callable() registers the callable and establishes persistent arguments; it does not produce a per-invocation prepared result. Public kernel launch remains unsupported.
• PreparedExecution combines a borrowed Runtime*, configuration, launch fields, native run identity, pipeline slot, and retirement state. ActiveExecution takes ownership of that object rather than allowing the preparation result to be reused.
• launch_run() still performs DFX arming and other host work; Runtime publication is synchronous and includes handshake state.

We need the explicit preparation result before introducing a common stream-ordered publication/reset/launch closure. Removing admission guards or enabling the launch stub first would preserve the coupling that the design is intended to remove.

Proposed API / Behavior

1. Three distinct lifetimes

Object	Responsibility
PreparedCall	Immutable resolved invocation description and references to everything needed to consume it; reusable across executions
Submission (or the existing per-run state adapted to this role)	Identity, ordering/completion state, status, pipeline lease during migration, and references for one execution
WorkingSet	Mutable device execution storage; remains under existing serial execution ownership in this milestone

One PreparedCall may be used by multiple submissions. Submission does not consume or mutate the PreparedCall. This milestone validates sequential reuse and does not promise simultaneous device execution or multiple queued launches.

PreparedCall is not an ACLGraph: TMR may still build its internal task graph on the device. Do not represent the preparation result as a log of host function calls.

2. Extract a real, owned preparation result

Start with the TMR path:

• Snapshot normalized device tensor descriptors and scalar values, resolved launch shape/thread/affinity settings, callable/function-table identity, and the effective per-call execution configuration.
• Give image/argument storage and callable/code/resource references explicit lifetimes. No borrowing from a temporary Runtime, input descriptor array, mutable caller configuration, or temporary host orchestration object.
• Retain stable resource bindings under their existing owner and identity rules. Keep references alive while PreparedCall or an execution still uses them; an unsupported invalidation must be rejected explicitly rather than yielding a dangling reference.
• Move native run identity, pipeline slot and completion/retirement flags out of the immutable result. Existing per-run machinery can carry them during migration; a new general-purpose framework is not required.
• A new or failed preparation must not overwrite a prior PreparedCall. A bank/slot being selected for execution is not the identity of the prepared result.
• Keep host-only diagnostics/export metadata separate from device execution data.

Descriptor/scalar snapshots do not snapshot device tensor contents. Preserve existing tensor data semantics and document borrowed-storage requirements where ownership cannot be retained automatically. Program tensor allocations/slices and copy-back leases must not be recycled while a prepared result or execution still references them.

For DFX, distinguish resolved host configuration/resource requirements from per-execution reset/arming effects. Carry the preparation data explicitly; the serial execution adapter may still apply effects at the existing safe point in this milestone. Do not claim all launch-time host work or capture requirements have been eliminated. Do not regress existing program DFX support to make the new path pass.

3. Connect the existing program path

Adapt the current TMR program preparation/execution path to create and consume PreparedCall, while preserving the current execution claim, serial dispatch and tensor transfer behavior.

A small internal seam is sufficient. No new public Python API, duplicate kernel-suffixed C API family, or generic runner hierarchy is required here. Compatibility adapters may keep existing entry points working during migration.

The result must be reusable through the production execution path, not just encode/decode unit tests. Submission-specific data and mutable execution storage must be rebuilt/reset independently of the retained prepared description.

4. Keep the interface compatible with the eventual HBG adapter

Review the representation against HBG's known needs: pristine image sections, relocation/bank-dependent bindings, definition/code references and tensor backing ownership. Keep runtime-specific payloads explicit; do not make a TMR packet layout the universal representation.

Implementing HBG pristine restoration or enabling HBG kernel execution is not required to close this issue.

Acceptance criteria

• Prepare A, then B with different invocation values; A's stored descriptors, scalars and effective configuration remain unchanged.
• Mutate/release the original host argument/configuration containers after preparation; the owned preparation result remains valid. Device tensor contents retain their documented external semantics.
• Execute A → B → A through the production TMR program path with independent expected outputs. The second A uses the same PreparedCall without rebuilding its resolved host preparation result. Per-submission reset, IO and bookkeeping are allowed.
• Reject an invalid or over-capacity B without damaging a previously valid A or its owned resources.
• Each execution gets independent run/completion/retirement state; prepared-result reuse does not reuse an old status or consume the object.
• Owner/lifetime tests cover temporary input containers, prepared-result lifetime, and program tensor backing through execution/copy-back; no dangling Runtime or function-table pointer is hidden behind the new type.
• Preserve existing program outputs, IN/OUT/INOUT transfer semantics and supported DFX behavior. Run relevant program regression checks and report architecture/runtime coverage honestly.
• Document the remaining launch-time host effects and the migration seam for the next submission-closure milestone. Keep kernel capability reporting honest while launch remains unsupported.

Scope boundary / stopping rule

Close this issue when the production-used immutable preparation result and the acceptance criteria above are demonstrated under current serial scheduling.

The following belong to later milestones and must not grow this PR:

• Publication/handshake device ABI split, asynchronous publication banks, and the full shared reset/launch/completion closure.
• Removing Orchestrator/lane admission restrictions or enabling arbitrary host run-ahead.
• Full kernel launch enablement, ACLGraph capture/replay, multi-graph lifetime handling or graph parameter updates.
• Complete resident DFX redesign, HBG replay restoration, or a general runner/resource-management refactor.
• The unload/reset double-failure follow-up in [Bug] Program loader loses explicit cleanup path when unload and device reset both fail #2269.

Alternatives Considered

Candidate	Reusable contribution	Required adaptation / reason not to adopt wholesale
#2180	Invocation validation, normalization, independently owned snapshots and bounded decoding	Its packet-to-AICPU transport is narrower than PreparedCall. It does not provide the complete resolved host result, publication or owner contract.
#2189	Shared ExecutionInputs and executor/scheduler consumption	Reconcile context-fixed configuration and unsupported kernel DFX with the eventual per-call/resident split. Reuse useful consumers without importing the prerequisite stack.
#2187	AICore-first submission, two-sided completion and partial-enqueue failure handling	Later closure work. Its PrepareTail wait and lack of the target publication-bank/Root protocol require redesign before capture integration.
#2028	Per-run completion events, status retirement and FIFO ownership	Later queue work. Its two-run scope does not replace the full execution-domain migration.
#2193	Non-destructive capacity refusal	Potential supporting fix after checking the current freeze predicate and rebasing; not a substitute for immutable preparation.

These are sources of implementation ideas, not merge approvals or prerequisites to merge unchanged. Rewriting the integration on current main is acceptable when it is smaller and clearer than preserving the old stacks.

Merely renaming PreparedExecution leaves its mixed lifetimes intact. Adding a second kernel-only executor duplicates scheduling and keeps program mode outside the design. Removing pipeline guards first exposes shared-state writes before their ordering is made safe.

Additional Context

Related: #2162 (runner lifecycle contracts), #1853 (queued execution), #2254 (launch-state publication). This issue establishes the preparation result before those execution/publication changes. In particular, it does not adopt an all-by-value launch payload as the final replay-update mechanism; capture freezes inline arguments, while the agreed design uses explicitly owned publication data for mutable replay bindings.

Inspected baseline and source anchors

Main after #2176: fd5f2d93fc55c784a002049577b0015ebfa4194e.

• PreparedExecution and ActiveExecution
• Current prepare_kernel_callable
• Launch-time host work
• Current synchronous Runtime publication

Design context: the agreed prepare/submit design and unified kernel-mode HTML. The operative contract is restated above so the issue does not depend on access to local design files. Existing primitive probes support the eventual queue/capture design; they do not establish that this production refactor has passed.

Subsequent milestones (context only)

1. This issue: owned preparation result and production program consumption.
2. Shared publication/reset/AICore+AICPU/completion closure, consumed by program and kernel; validate real TMR eager and fixed-binding replay.
3. Bank lifetimes and per-submission completion, followed by coordinated Orchestrator/lane/execution-claim migration; prove third and later submissions can enqueue before predecessor completion.
4. HBG restoration, complete DFX coverage, multi-graph use and updateable bindings.

[ChaoWao]

Seven suggestions. The first two are worth settling before implementation starts, because one is about the acceptance criteria themselves and the other about the interface shape. The rest are anchor corrections.

Anchors below are on f69c77491 (main after #2271), not the fd5f2d93 baseline the issue inspected, so a few line numbers have moved.

---

S1 — acceptance criterion 3 has no discriminator, and the repo just produced the right pattern

Execute A → B → A through the production TMR program path with independent expected outputs. The second A uses the same PreparedCall without rebuilding its resolved host preparation result.

The first clause is checkable from outputs. The second is not: A → B → A produces the same outputs whether the second A reused the prepared result or silently rebuilt it. As written the criterion passes in both worlds, so it cannot fail for the reason it exists.

This is the failure mode I have hit four times in my own probe work — an assertion that cannot fail is worse than a missing one, because nothing in hand can falsify it. In my case: three probes returned success while their positive arms failed; a two-graph arm compared only final values, so deleting the first replay entirely still passed; and a fence arm used the wrong event-status polarity, fenced on nothing, and reported a result that read as a significant finding.

#2271 shipped the fix pattern two days ago. Its double-failure test does not assert "the second lifecycle worked" — it counts:

assert lib.simpler_init(*init_args) == 0
assert lib.simpler_register_callable(ctx, 0, image) == 0
assert lib.finalize_device(ctx) == 0
# Two: the double failure's own attempt, and this lifecycle's success.
# Three would mean the abandoned handle was resurrected into it.
assert faults.unload_call_count() == 2

The comment is the criterion: a specific count, with a stated meaning for the wrong value. There is also precedent for exactly the quantity this issue needs — my Probe H gated its replay arm on a host-build counter being unchanged across 32 replays, which is what distinguishes "replayed" from "rebuilt and replayed".

So criterion 3 should name a counter, not a behaviour: "a host-preparation counter is unchanged across the second A", plus what a changed value would mean. record_bind_phase(HostPhaseKind::BindArenaH2d, …) already exists as instrumentation of this shape on the HBG path, so the mechanism is not new.

Related, same reason: criterion 5 ("Each execution gets independent run/completion/retirement state") also needs the negative half stated — what observation would show the second submission inherited the first's state.

S2 — HBG's preparation is materialization into a bank, which contradicts the central reuse claim

§4 asks to keep the representation compatible with an eventual HBG adapter. This is the specific incompatibility, and it is visible in HBG's own contract rather than inferable only from design discussion.

src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp:107-125:

// Host orchestration materializes this run's own graph into the image it
// uploads, so every device-resident region carries per-run content.
...
{PTO_PIPELINE_GM_HEAP,      PTO_PIPELINE_HOST_PER_RUN, 0},
{PTO_PIPELINE_RUNTIME_IMAGE, PTO_PIPELINE_HOST_PER_RUN, 0},

and :128-132:

extern "C" int concurrent_native_prepare_supported_impl(void) {
    // HBG can materialize a complete graph into the lease-selected unpublished
    // arena bank.
    return 1;
}

So HBG declares both device-resident regions HOST_PER_RUN at depth 2, and its preparation step writes the image into whichever bank the lease selected. Two of the issue's stated principles are TMR-shaped against that:

• "One PreparedCall may be used by multiple submissions" — HBG's prepared artifact contains bank-dependent addresses, so reuse across submissions requires either relocation on publish or one prepared result per bank.
• "A bank/slot being selected for execution is not the identity of the prepared result" — right as a design goal, and not currently achievable for HBG, where bank selection happens during preparation, not at execution.

Implementing HBG is explicitly out of scope and should stay out. What is cheap now is naming this in §4 as the review item: the representation needs a place for a runtime whose prepared payload is bank-relative, or an explicit statement that relocation-on-publish is the intended answer. Deciding it after a TMR-shaped interface exists is the expensive order.

S3 — the DFX paragraph should name begin_run's claim dependency

"Distinguish resolved host configuration/resource requirements from per-execution reset/arming effects" is the right split. The specific effect whose safety argument is about to change:

device_runner_base.cpp:2531 — begin_run() "drops the previous run's records and republishes the device level, which is only safe while this run holds the execution claim".

Under this milestone's serial scheduling that holds. But criterion 8 asks to document the remaining launch-time host effects, and this is the one whose justification evaporates in milestone 3, when submissions overlap. Naming it here makes it an inherited, known debt rather than a rediscovery.

S4 — there are two independent synchronous H2Ds in the launch path, not one

Motivation says "Runtime publication is synchronous and includes handshake state" and anchors device_runner_helpers.cpp. The second one:

• arm_device_wall_buffer (device_runner_base.cpp:2176) calls copy_to_device(slot_ptr, &init, sizeof(init)) at :2197, resetting the whole phase-buffer image;
• copy_to_device (:262) is a synchronous rtMemcpy(..., RT_MEMCPY_HOST_TO_DEVICE).

It matters for criterion 8 specifically because it is DFX-gated (device_wall_data_base == 0 returns early), so validating with DFX off will not surface it. For the later capture milestone it matters more: a synchronous copy under capture is rejected with 107030 while the capture stays ACTIVE and CaptureEnd still returns 0 — the graph is produced without that operation, silently. That is measured, not inferred: in one capture containing a rejected synchronous copy and a legal asynchronous one, the replayed graph carried only the asynchronous one's effect.

S5 — add a criterion for preparing the same invocation twice

Criteria cover A-versus-B independence and rejecting an invalid B. Nothing covers two prepared results with identical resolved content. Given that callable images on this path are already content-hash deduplicated, and that getting the retained-owner protocol right for that map took three review rounds in #2176 plus the follow-up in #2271, "do two identical PreparedCalls share storage, and does releasing one damage the other" belongs in the acceptance criteria rather than in the implementation's discovery log.

S6 — the #2187 row can be stated harder

The table says its PrepareTail wait "require[s] redesign before capture integration". Since #2176 merged, main's own event contract settles it — there is no PrepareTail among the five event kinds, and kernel_execution_state.h:88 states the rule directly:

Launch must therefore enqueue every record and its matching wait inside the same capture; prepare must not publish a tail for a later captured launch to consume.

So it is not pending redesign, it is contradicted by merged main. (The underlying constraint is CANN's: waiting during capture on an event recorded before capture began is rejected with 107024.) Same header also records the reason the topology is a chain rather than a fork — AicoreStart must be recorded on the aicpu stream before the AICPU launch, because the AICPU orchestrator spins on AICore's handshake report, so recording it after the launch could only fire once the AICPU task completed. Worth knowing for milestone 2, since it constrains the shared closure's shape.

S7 — ensure_run_streams() is structural launch-time host work

Criterion 8 asks which host effects remain at launch. This one cannot move to prepare, and launch_run states why (a2a3/platform/onboard/host/device_runner.cpp:591, called at :594):

The pair is readied here rather than at prepare because this is the first point the caller holds the execution claim: a prepared successor overlaps its predecessor's execution, so replacing a stale AICore stream during preparation would destroy a stream the predecessor is still running on.

So it is not incidental launch-time work that a cleaner split would remove — it is load-bearing on the claim, like begin_run in S3. Both belong in the criterion-8 documentation with that reason attached.

---

One thing in the issue I want to endorse explicitly, since it is a line I have been sloppier about than this text is:

Existing primitive probes support the eventual queue/capture design; they do not establish that this production refactor has passed.

That is the correct reading of what those measurements are worth.

[ChaoWao]

Design change: withdraw the reusable PreparedCall milestone

The agreed implementation route is now:

Pipeline A: program queueing, per-run results/errors and safe retirement
    → kernel eager integration using the same execution machinery
    → capture/replay with explicit publication, state restoration and graph ownership

Each ordinary invocation prepares its own inputs and submits its execution context once. Independent invocations can queue ahead. We will not introduce an independently reusable PreparedCall, an interface for resubmitting a prepared context, or an automatic cross-invocation preparation cache. These are withdrawn requirements, not optional future branches of this design. Existing callable, allocation and runtime-internal caches remain governed by their existing contracts.

Consequently, this issue's central requirement—production A → B → A using the same PreparedCall without preparing the second A—is withdrawn. The previous issue text and guidance requiring that behavior are superseded by this decision.

Retain already-validated argument snapshots, ownership fixes and useful preparation/execution separation only where they serve Pipeline A's execution context or the subsequent kernel integration. Do not retain a separate PreparedCall store, reuse controls or reuse-specific interfaces merely to preserve the old milestone. The code retained for the new scope must pass its production regressions; disabling reuse does not, by itself, establish that the remaining change is necessary or ready to merge. Preserve the reported reuse-failure evidence rather than calling those failures acceptable.

Disposition: this original feature request should be closed as superseded / not planned, not completed. Necessary remaining work belongs to the Pipeline A or kernel/capture scope with corresponding acceptance criteria; do not invert reuse assertions and claim the original criterion passed.

This is a change to the requested design, not an implementation failure to follow the original issue. The same scope correction applies to #2281 and its implementation PR #2285. This comment does not itself change issue or PR state.

[ChaoWao]

Closing out my side of the withdrawn milestone, and recording the reuse-failure evidence so it is not lost with the requirement.

The reuse regression had a root cause, and it was mine

I reported 7 onboard failures under reuse. They were not timing sensitivity and not a pre-existing flake. Cause: PreparedCallStore and PreparedCall::use_count_ were unsynchronized, while pipeline depth two reaches the store from two threads — the prepare of one run and the finalize of its predecessor.

The fatal interleaving is match-then-reference as two steps: a concurrent resolve is free to evict the matched entry between them, so the execution that matched it launches against bindings another thread has already released. On the device that is a run which never progresses, which is why every failure read as a scheduler timeout rather than an error:

finalize_native_run failed with code -100
sched_error_code=100 SCHEDULER_TIMEOUT
sub_class=S1:running-stalled completed=7/513 running=1 ready=0 waiting=505
orch_done=1 stuck_task_id=7 stuck_core=24

always on run_id=1, the second run. host_build_graph was hit too even though it opts out of matching, because it still stages, commits and evicts through the store, and setup_static_arena's bank retirement walks every slot.

The controlled A/B

Narrowed to the four failing directories in the sweep's own shape (xdist, 4 devices — the box's single 8-card lane was head-of-line blocked by a pending 16-card job for over an hour, so both arms ran at 4):

Arm	Result
store unsynchronized	7 FAIL / 2 PASS, signature above
every store entry point under one lock, match+reference made a single acquire_matching	9 PASS / 0 FAIL

Both arms are the same tree and the same shape; only the synchronization differs. Off-device gates on the fixed arm: cpput 152/152, pyut 2401 passed / 46 skipped, a2a3sim 40/40, a5sim 36/36.

What this does and does not establish

• It establishes that the reported failures were a real defect in the change and that they are diagnosed rather than tolerated.
• It does not establish the acceptance criterion. The closing gate was to be the full 8-card sweep, which never ran, and the 4-device shape is half the sweep's concurrency. Criterion 3 is not met and I am not claiming it. No assertion was inverted; the two scene-test methods still assert a resolve delta of 2 and a reuse delta of 1 across A → B → A.
• The commit was never pushed and no PR was opened.

Disposition on my side

Per the design change, I am not retaining the store, the reuse controls, the reuse counters or the reuse-specific tests and documentation. I am keeping only what serves Pipeline A's execution context, under that scope's own acceptance criteria: the owned argument snapshot and the ownership fixes behind it, the split of validation into input staging / output copy-back / binding release, and the separation of the resolve step from the per-submission step.

Two defects the onboard lane surfaced along the way are worth naming because they are independent of reuse and any prepare/submit split will meet them again:

1. A stack-local HostApi passed into bind_callable_to_runtime_impl is retained by a5 host_build_graph in SchedulerStateOwner::api and dereferenced at drain — a bus error at _wait_run_handle. Whatever owns a run's preparation has to own the HostApi it binds with.
2. setup_static_arena sets arena_changed on the first commit of a new bank, so any store-wide invalidation there releases the bindings of a concurrently executing run on another slot. host_build_graph gives each pipeline slot its own bank, so slot 1's first bind killed slot 0's in-flight run. Invalidation at that point has to be bank-scoped and must refuse to take bindings an execution still references.

One correction to my own record, for S1's sake: the reuse counters I added are the right kind of discriminator, but they were never the thing that failed — they were green in isolation throughout. The signal that mattered came from the sweep, and I read it late.

🤖 Generated with Claude Code