Fix: kernel-mode capacity refusals no longer free what they protect (⑤K3)

[sunkaixuan2018]

Depends on #2064 (①K1) @ 86dd62b4912866a6e96d63891b4c2838a2635b5a

This branch is built directly on that commit, which is not yet merged. The
main base above is a GitHub constraint — a branch on a fork cannot be a PR
base — so the diff shown against main includes all of K1. Exactly one
commit is mine:

• b8e739d9 — kernel-mode capacity refusals no longer free what they protect

Incremental diff (this PR only):
sunkaixuan2018/simpler-PTO@86dd62b...skx/kernel-capacity-freeze-k3

Draft until #2064 merges, then rebased and marked ready.

What this is

A kernel-mode context's device buffers must keep their addresses for the life of
the context, because an ACLGraph replays the addresses that were captured. Every
runtime-time growth is release + reserve or free + malloc, which re-bases; the
failure mode is silently wrong data on replay, not a crash. That makes this a
correctness change, not a sizing one.

K1 landed the first guard on that invariant, in setup_static_arena. Its
predicate is right and is unchanged here. What this PR fixes is that its
refusal, and the refusal on the second growth point, were both destructive.

Two refusals that destroyed what they protected

1. setup_static_arena — the guard freed the bases it had just declined to
move. The refusal returned PTO_RUNTIME_ERR_INTERNAL, the caller collapsed that
to ok = false, and the !ok block released all three regions unconditionally.
DeviceArena::release() frees the backing buffer, so a fired guard dropped
exactly the base addresses it names in its own error message. K1's comment above
the guard said so; nothing acted on it.

2. RetainedTempBump::begin — the free came before anything could refuse. The
grow path is device_free(old) then device_malloc(bigger). A refusal placed
anywhere downstream of that — including at the platform's device_malloc —
arrives after the address is already gone, and the existing failure path then
clears the slot to {nullptr, 0}. The guard has to sit ahead of the free.

The invariant both fixes share, and the one every later growth-point guard will
need: a refusal must have no side effects.

What changed

host/static_arena_bank.h (new)	the bank's commit rule, shared by the onboard and simulation runners instead of hand-duplicated in both
the rule's failure handling	an allocation failure still rolls the whole bank back; a capacity refusal leaves every region committed and every cached size intact
both setup_static_arena bodies	~55 lines each become a request literal plus a call; each keeps its own prebuilt-arena cache invalidation, now keyed on bases_changed
RetainedTempBump::begin (a2a3 + a5)	refuses a kernel-mode re-base ahead of the device_free, leaving the slot as the previous run left it
HostApiOps::is_kernel_mode (new)	how the context's identity reaches runtime code; a table that does not supply it reports program mode
six contract sites	the docs that described the behavior these guards make conditional, listed below

is_kernel_mode is a query of the existing ExecutionModeLatch, not a new
switch — no environment variable and no macro, per
env-macro-gating.md §1. Its producer
(c_api_shared.cpp) and its consumer (runtime_maker.cpp) are compiled into the
same libhost_runtime.so, so appending to the ops table creates no ABI skew.

The retained-buffer guard refuses a re-base, not an allocation

This is the part worth a reviewer's attention, because I got it wrong first.

The guard fires on required > size && addr != nullptr. The addr != nullptr
half matters: an empty slot holds no address a captured graph could reference, so
a context's first allocation moves nothing and is taken normally. Without that
clause the guard refuses the first bind too, and since the grow path is the only
writer of the slot, the slot would stay {nullptr, 0} forever — a kernel context
could never reach its frozen size at all.

My first version omitted it, on the premise that a kernel run never stages host
tensors anyway. That premise is not enforced anywhere in the tree, and I could
not support it on a re-read:

• validate_kernel_launch_args checks null-ness and the callable-id range. It
  never inspects tensors.
• ChipTensor::address_space defaults to HOST, so the default goes the
  opposite way.
• SimplerKernelInvocationHeader::host_copy_tensor_count exists precisely to
  carry host-memory args in kernel mode, documented as "zero until the host-only
  copy contract lands". Host tensors there are a planned contract, not an excluded
  one.

So the narrow guard is the correct one: it still refuses every real re-base, and
it refuses nothing else.

Doc updates in the same commit

The behavior these guards make conditional was described in six places, all of
which now read false for a kernel context and move with the code:
setup_static_arena and the retained-temp-buffer entries in host_api.h,
DeviceRunnerBase::setup_static_arena's rollback promise (and its @return,
which named -1 for a function that returns PTO_RUNTIME_ERR_INTERNAL), the
kernel-mode capacity paragraph in runtime_c_api.h, both runners' latch comments,
docs/task-flow.md, and RUNTIME_LOGIC.md in both architecture trees.

The extraction is a flagged deviation from "program path byte-identical"

Behavior in program mode is unchanged — kernel_mode is false there, so the new
branch is unreachable and growth, release and rollback all take the paths they
took before. But the code moved, so this is not literally untouched, and it is
deliberate for two reasons. The rule was two hand-maintained copies that
codestyle.md §10 warns about, and onboard/sim
symmetry is now structural rather than a review obligation. And it is what makes
the fix testable without a device: DeviceRunnerBase needs CANN, the rule needs
nothing.

The allocate_tensor design question, answered

The handover asked whether kernel mode should refuse allocate_tensor
unconditionally, or only outside a prepare window borrowed from ④K2. The
unconditional form holds: #2176's prepare-once allocator reaches mem_alloc_
directly through its own context ops, never through HostApi::device_malloc, so
refusing the runtime-facing surface would not block a legitimate kernel-time
allocation. This PR does not install that refusal, because the growth point it is
in scope for needs its guard higher up anyway — but the question is settled and the
next guard does not have to re-derive it.

Out of scope, and one question for a reviewer

Growth points #1 acquire_graph_definition_block and #2 HBG host-tensor staging
are not closed here. Both are HBG-only and both live in files #2173 rewrites;
#2173 already carries a disposition for the definition block. @TaoZQY — do you want
to keep #1 (and #2) on your side, or should a follow-up here close both against the
is_kernel_mode query this PR adds? Splitting one growth point across two PRs is
the outcome worth avoiding. Note that acquire_graph_definition_block is
grow-by-replacement of a device block, so it is the same silent-replay hazard,
not a lesser one.

DFX is outside the frozen capacity, but it is not address-stable. This is the
question the handover asked, so here is the full answer rather than half of it. The
DFX device buffers — the device-wall buffer and the collectors' workspaces — are
separate mem_alloc_ allocations, not regions of an arena bank, so nothing this PR
freezes covers them and it correctly writes no DFX teardown logic. The device-wall
buffer is allocated lazily once and never grown, so it does not re-base. The
collector pools do: prepare_execution calls finalize_collectors() whenever
collector_shape_is_stale(...), so a run whose core or AICPU-thread counts differ
from the previous one tears their device memory down and rebuilds it at a new
address. A kernel context's shape is context-static, so that should not fire — but
nothing enforces it, and if the DFX buffers are ever reachable from a captured
graph this is a growth point in its own right. Flagging it rather than guarding it
here: it belongs with whoever owns DFX under the execution claim (#2163).

The sim-side guard stays. Simulation never latches KERNEL, so that branch is
never true, but removing it would break the onboard/sim symmetry the stub-parity
argument rests on.

Verification

All on myserver (aarch64, CANN 9.0.0). No device is needed for any of it.

Lane	Result
cpput, no hardware	144/144 passed — base 86dd62b4 is 143/143, built and run the same way
new test_static_arena_bank	11/11, no_hardware label
test_trb_runtime_temp_buffer	16/16, was 10 — and identically for test_a5_trb_runtime_temp_buffer
editable install	all 8 libhost_runtime.so built (a2a3 + a5 × onboard + sim × hbg + trb); no warning names a changed file
tests/ut/py/test_host_runtime_abi.py	4 passed, 4 skipped — the exported symbol set is unchanged
pre-commit	all hooks pass, clang-tidy included, in CI

The five acceptance criteria are covered twice, once per slot kind:

	arena bank	retained temp buffer
base and capacity constant across N calls	RepeatedCommitsMoveNoBaseAndCallNoAllocator	KernelModeHoldsOneAddressAcrossRepeatedRuns
zero underlying alloc/free	same, via DeviceArena::alloc_count()	same, via the fake's malloc/free counters
exact capacity accepted	ExactAndSmallerRequestsAreServedInPlace	same test, a run that exactly fills the buffer
over capacity refused with INTERNAL	OneByteOverCapacityIsRefusedAndChangesNothing	KernelModeRefusesToGrowTheRetainedBuffer
refusal does not break the held plan	same test's second half	KernelModeRefusalLeavesTheRetainedBufferUsable

KernelModeAllocatesOnceThenRefusesToGrow is the production shape: the latch is
write-once, so a real context is in kernel mode from its very first bind. It
asserts one allocation, then a refusal, then the sized run still binding from the
same address.

The other exit is pinned too. AllocationFailureRollsTheWholeBankBack and
AllocationFailureOnTheFirstCommitStillRollsBack drive a failing backing
allocator through DeviceArena's injectable alloc/free, so the rollback path is
reachable with no device: an allocation failure releases every region and zeroes
every remembered size, in kernel mode exactly as in program mode, because a setup
that never completed published no address to protect.

Negative controls

A test that passes against the fix is not evidence it would have caught the
defect, so each defect was re-introduced and the suites re-run.

• Routing the capacity refusal back through the rollback (dropping the
  capacity_refused branch): 3 kernel-mode arena tests fail, and both
  program-mode tests still pass.
• Removing the retained-buffer guard: 2 tests fail, identically on a2a3 and
  a5; the rest pass.
• Dropping the addr != nullptr clause: exactly
  KernelModeAllocatesOnceThenRefusesToGrow fails, on both architectures — the
  first-bind regression has a precise barrier.
• Removing the rollback entirely: exactly the two allocation-failure tests
  fail; the nine capacity tests still pass, so the two failure classes are
  pinned independently of each other.

How the findings above were found

The commit was put through a self-review before this update: eight independent
reviewers over separate dimensions of the diff, then three adversarial refutation
lenses per finding. It raised 22 findings and the skeptic panel refuted all 22 —
a verdict I do not report as a clean bill of health, because the refutation lens
set included reachability, and since nothing latches KERNEL yet, every
kernel-mode finding is trivially unreachable today. The panel's value was in
surfacing candidates, not in adjudicating them.

Four of those candidates were textually verifiable and are fixed here regardless
of the vote: the first-allocation refusal, the is_kernel_mode comment claiming
table-wide address stability that acquire_graph_definition_block contradicts,
the stale contracts listed earlier, and the two test defects above — the
"smaller" run that packed to exactly the retained capacity, and the missing
rollback coverage.

🤖 Generated with Claude Code

[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.}