Add: TMR kernel invocation snapshots and native transport adapter（K4）

[Leaf-Salix]

Depends on #2064, #2177 and #2176 — keep this PR in Draft. Align and revalidate against the merged prerequisites before merging. See #2064, #2177, and #2176.

Based on #2177 at 2153406a0420e8de17cf7397d30864939b1bca0c, whose K1 baseline is dc1268cd55405fc56137eb16d04227ea797a419f. The original K4 foundation is 3e822453. The follow-up b0943525 integrates K2's 48b120b15cbcc7ea1f28ae36ec3403e17cfe48b1 with K4 admission/transport/consumption interfaces. K2 is included as source changes, not as a merge parent; the combined diff must not be attributed entirely to K4. The HBG/2b stack is not included.

Motivation

TMR (tensormap_and_ringbuffer) kernel mode needs a per-invocation argument snapshot that remains independent of mutable Host inputs and later calls. Reusing program bind/staging or overwriting a shared Runtime to pass each invocation would couple parameter lifetime to execution resources and make asynchronous submission difficult to reason about.

This PR provides shared invocation validation → TMR snapshot encoding/decoding → a native CPU transport adapter, with explicit interfaces for trusted resource metadata:

Prepared callable metadata + stable execution-binding view
  + current ChipStorageTaskArgs
  → validate and encode an independently owned candidate packet
  → native CPU submission adapter
  → CANN-owned argument copy
  → bounded TMR decoding and caller-owned EntryArgsStorage

Host template cache:
  updated only after the owner reports successful complete enqueue

This is K4 foundation plus K2 integration, not the complete kernel execution path. Onboard TMR init/prepare/close can exercise K2 resources, and real prepare consumes the shared callable admission checks. Sim init retains UNSUPPORTED after contract validation; HBG kernel init remains unsupported without 2b. simpler_kernel_mode_supported() remains 0, and public launch remains fail-closed rather than reaching the adapter. Stable execution-binding and residency-generation providers, the production device dispatch, and the complete submission binder remain dependencies. The included transport probe uses explicitly test-only metadata.

Invocation contract and adaptation

Shared validation without a shared runtime payload

src/common/task_interface/kernel_invocation_validation.h provides:

• PreparedInvocationView: trusted callable ID, effective tensor/scalar counts, and callable residency generation.
• derive_invocation_counts(): validates the signature and derives effective counts. A nonzero cached scalar count must agree with the signature; zero uses the signature-derived value.
• validate_invocation_header(): checks the formal K9 framing, kernel mode, callable ID range, count limits, payload length against the supplied readable span, and identity/count agreement with the trusted callable.
• host_copy_tensor_count != 0 is rejected.

These checks do not interpret TMR payload contents. HBG can reuse the common admission rules while retaining its own graph/blob validation. A trusted callable view must come from an independently validated, live registration; constructing it from the packet being checked would not establish trust.

No public C query or lifecycle API is added. The combined branch adopts K2's four-argument simpler_kernel_mode_prepare_callable(ctx, cid, callable, callable_size) signature; prepare no longer receives a caller stream. This is an intentional prerequisite ABI alignment from the original K4 baseline, not a claim that all old signatures remain unchanged. Execution-mode definitions and formal K9 layout remain unchanged.

src/common/platform/include/host/kernel_entry_validation.h consumes the shared effective-count rules in real onboard/sim prepare admission. It also checks canonical callable names, child counts, alignment, offsets and complete variable-length storage bounds before hashing, registration or state mutation. These checks are kernel-prepare-only; program registration does not acquire this new validation path. Combined onboard init preserves 2a's builder/contract admission before mode latch, executor replacement or context resource creation.

TMR-specific wire and consumption

src/common/tensormap_and_ringbuffer/kernel_invocation.h defines the runtime-specific binding reference and bounded decoder. The packet layout is:

SimplerKernelInvocationHeader
TmrBindingRef { device_binding_addr, context_generation }
ChipTensor[tensor_count]
uint64_t[scalar_count]

The packet length is derived from the formal types and trusted effective counts:

sizeof(SimplerKernelInvocationHeader) + sizeof(TmrBindingRef)
    + tensor_count * sizeof(ChipTensor)
    + scalar_count * sizeof(uint64_t)

The inherited K9 is currently 40 bytes; TmrBindingRef is 16 bytes, with size/offset assertions. Packets contain descriptors and scalar bits, not Tensor data, callable binaries, Host pointers, STL objects, or an entire Runtime. No global maximum-sized Host packet is allocated.

The two version fields protect different resources:

Field	Meaning and provider
K9 generation / PreparedInvocationView::slot_generation	Callable residency version, supplied by the prepared-callable owner
TmrBindingRef::context_generation	Stable execution-binding/context version, supplied by the execution-resource owner

device_binding_addr is opaque to the codec. The provider owns the referenced binding's format, capacity validation, publication, and lifetime. A nonzero address or a matching generation does not independently prove that resources exist or remain alive.

Encoding normalizes Tensor fields, excludes padding and unused dimensions from identity, and accepts only Device Tensor descriptors. Nonempty tensors undergo extent/backing and arithmetic checks. Empty tensors follow the existing boundary-conversion behavior; this is not a promise that arbitrary operators support empty inputs.

Decoding uses a pointer-plus-length ByteSpan, validates exact TMR size and trusted bindings, and reads unaligned input through bounded copies. Failure leaves the previous output unchanged. The resulting view borrows immutable packet storage.

kernel_invocation_args.h provides consume_tmr_invocation(packet, trusted_callable, trusted_binding, out): decode and validate before materializing into consumer-owned EntryArgsStorage, preserving out on rejection. It is an explicit consumption interface, not a registered production AICPU dispatch. It does not call program Runtime::set_orch_args() or overwrite dynamic arguments in a persistent Runtime. Storage must remain alive through the last config/orchestration access and every derived Tensor reference.

Host templates and native transport

src/common/worker/tmr_kernel_invocation.h provides one Host encoding-cache object per prepared callable and independently owned candidates:

• The structural template excludes Tensor addresses and scalar values, while retaining scope, generations, geometry, offsets, and backing sizes.
• A cache hit still validates current inputs and fills current addresses/scalars.
• Encoding or submission failure must leave the previous cache in place. Cache commit is explicit and non-throwing; the owner calls it only after the complete enqueue sequence succeeds.
• Cache storage does not grow with invocation count. In-flight task/graph copies have a separate owner and lifetime.

src/common/platform/onboard/host/tmr_kernel_invocation.{h,cpp} provides enqueue_tmr_invocation_aicpu(DeviceRunnerBase&, ...). It revalidates the candidate's owned length and trusted metadata, then calls the initialized runner's launch_aicpu_payload() CPU transport. The implementation is compiled into both architecture-specific onboard targets. The direct LoadAicpuOp overload remains available for the isolated transport probe. The owner supplies the dedicated AICPU stream and establishes the complete event ordering outside this adapter; linking the adapter does not wire public launch or install a device consumer.

The adapter targets the kernel-specific simpler_aicpu_kernel_exec symbol, not the existing program KernelArgs* entry. The production symbol/consumer is not registered by this PR; it exists in the isolated probe for transport validation.

This does not require TMR and HBG to use the same SDK launch API. The common requirement is task-owned argument-copy semantics with a valid Host buffer until the native copy returns. The TMR probe exercises the existing rtsLaunchCpuKernel path, not an assumed WithHostArgs replacement.

Lifecycle and program-path impact

Submission transaction: The owner must protect the entire sequence, not just individual accessor calls:

enter owner submission protection
  → acquire trusted callable and binding views
  → validate and encode a private candidate
  → perform the complete native enqueue sequence
  → commit the template on success
  → release Host packet
leave owner submission protection

The same protection must coordinate with prepare replacement and close. K4 introduces no independent cache mutex or parallel resource registry. Independent callers may use independent caches/outputs; shared-cache access requires owner serialization.

Ownership: K4 owns the Host packet until the native API completes its copy. Native success means enqueue, not device completion. CANN owns the task/graph argument copy; the resource providers keep callable and stable execution resources alive. Generation validation does not replace lifetime management. Partial enqueue failure and already-submitted work require binder error handling; leaving the Host cache unchanged does not cancel device work.

Alignment with other modules: The K2 stream/persistent-argument owner is included, not duplicated. Its arbitrary-payload transport is the runner adapter's integration point. The included K2 snapshot owns a private AICPU stream, a hidden AICore stream and five context events; this must not be described as proof of two-hidden-stream capture. K3/K5/K10 must provide actual trusted views; K4 does not substitute fixed generations, fake addresses, or the program Runtime. No cache sidecar or second registry is attached to context/callable state before a real owner consumes it.

Program-path impact: K4 does not redirect program bind, staging, Runtime upload or per-run cleanup through invocation snapshots. However, the combined diff includes K2 changes in shared platform, device-runner and worker code, so it is not accurate to say no existing production source is modified. It also includes two shared buffer-pool lint fixes: remove unspecified pointer-address ordering while retaining collect/deduplicate-before-release behavior, and use a single-mutex scoped_lock with the same notify lifetime. A regression test covers release-once and queue clearing; no warning suppression is added.

Left to follow-up PRs: Trusted Host/device provider publication; capacity establishment and callable residency; complete stream/event submission and cancellation; production device dispatch; Torch/taskQueue/storage retention; ACLGraph capture/replay; and close/graph quiescence. Host concurrency tests do not establish safety for device execution, replay, or close races.

Tests

Evidence is separated below between the original 3e822453 foundation and the combined K2/K4 source snapshot. Neither is a claim that the full CI matrix or the complete kernel pipeline has passed.

Coverage added in this PR

• tests/ut/cpp/types/test_kernel_invocation_validation.cpp: 3 cases covering effective scalar rules, invalid signatures/counts/modes/IDs/generations, unaligned framing, output preservation, and acceptance of opaque non-TMR payloads with independently trusted counts.
• tests/ut/cpp/common/test_tmr_kernel_invocation.cpp: 11 cases covering round-trip conversion, cache identity and dynamic patching, unchanged cache/output on failure, canonical padding, stale bindings, short/corrupt packets, self-consistent K9 lengths that still violate exact TMR size, tensor arithmetic/Device-only rules, empty boundary conversion, minimum/maximum/scalar-only packets, and independent/owner-serialized concurrent calls.
• tools/cann-examples/tmr-invocation-snapshot/: A committed, standalone Host/device probe using the actual CPU loader and packet layout. It submits 24 asynchronous packets, overwrites and releases Host packet storage, and verifies copied values only after a test-only device-read gate is released. Result buffers start with failure sentinels. The gate has a timeout and is not part of production execution.

The combined update adds prepare signature/FAM rejection cases, real sim prepare entry coverage, trusted-view consumption/output-preservation cases, and 10,000 repeated encodes showing bounded logical template bytes. The latter is not an RSS, peak-heap or CANN task-memory measurement. It also includes K2's lifecycle tests and a buffer-pool release regression.

Original foundation results (3e822453)

Validation	Result and scope
New K4 C++ UT	14 cases passed in 2 targets, on the local Host and hwServer
Targeted C++ regression	9/9 CTest targets passed locally and on hwServer, including the two K4 targets plus inherited K1/2a contract, loader, entry, state, wire, and both architecture-specific TMR maker targets
ASan/UBSan	2/2 K4 targets passed locally; not counted as a remote sanitizer run
Python sim ABI	13 passed, 4 onboard parameterizations deselected, locally and on hwServer; deselected cases are not passes
Program sim smoke	TMR/HBG vector on both a2a3sim and a5sim: 4 passed on hwServer
Native CPU snapshot probe	A2/A3 hardware, CANN 9.0.0, exclusive device 3: 24 gated packets passed, process exit 0
Build and lint	Isolated rebuilds with explicit ccache and four-way parallelism; all applicable incremental pre-commit hooks passed for the committed files

The hwServer run used a fixed archive of the pushed commit in a new isolated directory. Archive hashes matched; post-test comparison found no source-content differences. Python/native imports were verified against this snapshot. Existing Torch dependencies were reused read-only; no Torch installation or shared-environment modification was performed.

Combined K2/K4 results (2026-09-10)

Validation	Result and scope
Targeted C++ regression	11/11 targets passed, locally and on hwServer, including K1/2a, K2 persistent args, K4 validation/codec and both TMR makers
ASan/UBSan	3 targeted admission/K4 targets passed locally
Python sim ABI and prepare admission	25 passed, 36 deselected, locally and on hwServer
Program sim smoke	Both architectures, TMR/HBG vector: 4 passed on hwServer
K2 onboard lifecycle	A2/A3 TMR borrowed-device init/close and prepare/reuse/close: 2 passed, exclusive device 3
Runtime build	All eight hwServer components rebuilt; both onboard TMR libraries link the runner adapter
Subsequent shared-header lint fix	Buffer-pool/profiler tests passed locally before and after the fix; both targets also passed ASan/UBSan
Incremental pre-commit	All applicable checks passed for the complete modified set, including the new adapter .cpp; no bypass or suppression

The remote combined source archive SHA-256 is e9b0fe33f0ae5fc4bef96124384ffdb6ca3959289f298056d790d9333812c1b2. That archive predates the final buffer-pool lint-only fix, which has separate local rebuild/test/sanitizer evidence. The original 24-packet probe was not rerun on this combined snapshot. K2 lifecycle success is not K4 invocation execution or capture/replay evidence. No Torch installation was performed.

The old #2180 network1 failure at 3e822453 was investigated separately: the CI daemon could not bind its occupied port, and tests reached a different daemon. This was not a K4 regression; no CI/daemon changes or rerun were made for that incident. The combined branch's future CI result is not inferred from that diagnosis.

The native probe is eager CPU transport evidence, not formal TMR executor, A5 hardware, or ACLGraph capture/replay evidence. Test-only callable/binding metadata does not establish production provider correctness. Program smoke does not establish kernel-mode execution support. Local and remote runs of the same tests are not summed as unique cases.

Main reproduction commands, after isolated installation and building the targets according to the repository testing guide:

ctest --test-dir tests/ut/cpp/build --output-on-failure \
  -R '^(test_kernel_invocation_validation|test_tmr_kernel_invocation|test_pipeline_contract|test_pipeline_contract_loader|test_kernel_entry_validation|test_kernel_execution_state|test_kernel_persistent_args|test_hbg_kernel_persistent_args|test_kernel_invocation_header|test_trb_runtime_temp_buffer|test_a5_trb_runtime_temp_buffer)$'

python -m pytest tests/ut/py/test_host_runtime_abi.py tests/ut/py/test_kernel_mode_c_api.py \
  -k 'not onboard' -v --forked

for arch in a2a3 a5; do
  python -m pytest "examples/$arch/tensormap_and_ringbuffer/vector_example" \
    --platform "${arch}sim" --manual include --pto-session-timeout 600 -v --forked
  python -m pytest "tests/st/$arch/host_build_graph/vector_example" \
    --platform "${arch}sim" --manual include --pto-session-timeout 600 -v --forked
done

# Build the probe's host/device CMake projects with the CANN toolchain.
# Run under an exclusive device lock; TASK_DEVICE is assigned by the queue.
<host-build>/snapshot_probe "$TASK_DEVICE" \
  <a2a3-dispatcher>/libsimpler_aicpu_dispatcher.so \
  <device-build>/libsnapshot_probe_device.so

Not validated or not implemented: Full UT/ST suites, A5 hardware execution, empty-Tensor program execution against the baseline, production metadata/binding providers and device consumption, capture/replay, alternating captured graphs, asynchronous cancellation, and close/resource-reclamation races.

Commits

The branch retains the original foundation and adds the combined integration on top of #2177:

1. 3e822453 — Add: TMR kernel invocation snapshots and transport adapter: Shared validation, TMR wire/codec, transactional Host templates, native adapter, argument conversion, unit tests, and the standalone transport probe.
2. b0943525 — Add: integrate TMR snapshots with kernel context resources: K2 48b120b15 resource lifecycle, 2a admission-preserving init integration, shared kernel prepare validation, runner transport adapter, explicit TMR consumption, regression tests and shared-header lint fixes. K2 human authorship is preserved in the commit trailer.

Prerequisite order: #2064 → {#2177 (2a), #2176 (K2)} → this K4 PR. This is a dependency set, not a requirement that 2a and K2 merge in a specific order relative to each other. Later execution integration must consume the remaining resource/provider/binder modules; this PR neither replaces them nor declares their work complete.

中文总结

• 合并依赖： 基于 Add: TMR kernel-mode resource contracts and init admission（2a） #2177 的 2153406a，保留 3e822453，叠加 K2 48b120b15 并完成组合接线；等待 Add: kernel-mode C ABI skeleton and wire headers (K1) #2064、Add: TMR kernel-mode resource contracts and init admission（2a） #2177、Add: persistent kernel-context execution resources (K2) #2176 合并后对齐复验，保持 Draft。
• 交付内容： 公共 K9/callable 校验、TMR 定长快照、Host 模板缓存、真实 prepare 准入、runner CPU adapter 和显式消费接口；不是完整 kernel executor，公开 launch 未接通，supported 仍为 0。onboard TMR init/prepare/close 可测试，不再笼统描述所有 init 都 unsupported。
• ABI 与解耦： prepare 对齐 K2 四参数 ABI；共用外层准入，不统一 HBG/TMR payload，只消费可信只读 view，不复制 K2/K3/K5/K10 owner。K4 不改 program 参数生命周期，但组合 diff 确实包含 K2 共享平台改动及公共 lint 修复。
• 生命周期： 每次独立 packet；owner 保护从获取视图到完整 enqueue/缓存提交，并与 close 协调。callable 驻留 generation 与 context generation 分开，版本检查不能替代保活。
• 组合验证： 本地/远端 11 组 C++ target、25 项 sim ABI/准入通过；本地 3 组准入/K4 sanitizer 通过；远端四项 program smoke、两项 A2/A3 K2 生命周期通过。后续 buffer-pool lint 修复另有本地两组普通/sanitizer 及完整增量 pre-commit 证据。原版 24 包 probe 未在组合版本重跑，不混算。未安装 Torch。
• 未完成范围： 真实 provider/设备消费、ACLGraph capture/replay、两图交替、异步取消及 close/graph 生命周期仍需后续集成，不以 UT 或测试 fixture 代替。

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}