Add: rt_available_cluster_count / rt_available_aiv_count orch helpers

[yanghaoran29]

What

Orchestration had no way to ask how many clusters the current run actually has,
so MIX / AIV cohorts hardcoded 24 or 36. This exposes the counts the runtime
already latches in runtime_finalize_after_wire through the ops table:

int32_t clusters = rt_available_cluster_count();  // MIX clusters (one AIC each)
int32_t aivs     = rt_available_aiv_count();      // standalone AIV cores

Available on tensormap_and_ringbuffer (a2a3, a5) and host_build_graph (a2a3).

Why the host path had to change first

host_build_graph's orchestrator runs on the host, inside
bind_callable_to_runtime. But worker_count was not published until
prepare_runtime_for_launch, one step later in run():

new (runtime) Runtime()      → worker_count = 0
bind_callable_to_runtime()   → host orchestration runs here, reads worker_count ← 0
run()                        → resolve_block_dim() → set_worker_count()

So host orchestration read 0. That silently disabled the require_sync_start
block_num deadlock guard, whose only consumer treats a limit of 0 as
"unknown, skip" — no in-repo hbg case uses require_sync_start today, so CI
stayed green while the check was off.

The fix is to resolve the geometry before the graph is built. The expensive
part is the device query, so that gets cached once:

• DeviceRunner latches the AICore stream's block_dim ceiling in
  ensure_device_initialized(). Nothing in the codebase calls
  aclrtSetStreamResLimit, so the limits hold for that stream's lifetime;
  finalize_common() clears them alongside the stream.
• resolve_block_dim() becomes pure arithmetic on that ceiling, so the auto
  and explicit branches share one range check instead of the auto branch
  skipping validation to save an ACL round trip.
• The core-geometry half of prepare_runtime_for_launch splits out as
  prepare_launch_shape() and runs before the bind; the task-address half
  stays in run() as resolve_task_binary_addrs(), where the task table
  exists.

The move is a no-op for tensormap_and_ringbuffer:
src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp is the only
bind-time reader of worker_count in the tree.

Also in this change

• Sim auto width. Sim resolves the block_dim == 0 sentinel to
  SIM_AUTO_BLOCKDIM (8) instead of the modelled chip's full width. Sim runs
  one OS thread per AICore, so 24 / 36 clusters is 72-108 threads per case for
  no added coverage. It is not a ceiling — an explicit block_dim is still
  honoured up to PLATFORM_MAX_BLOCKDIM, so the 178 cases that pin a value
  today are unaffected.
• No more /3. Cluster counts derive from
  PLATFORM_{AIC,AIV}_CORES_PER_BLOCKDIM instead of open-coded / 3 and
  * 2 / 3 in the two scheduler cold paths and on the host.

Test

The scene test spends the counts rather than pinning them — there is no
constant to pin, since sim and onboard resolve different values and onboard's
comes from the driver. The orchestration reports what it saw in one tensor and
launches a MIX cohort of exactly cluster_count blocks, each writing its block
index into another; the host reads the width back out and checks that many
slots (via the compare_outputs hook from #1470).

What would fail:

Defect	Caught by
count is 0 or exceeds the platform ceiling	host range check
count larger than the run really has	cohort's require_sync_start needs every block co-resident → deadlock guard fires on device
count smaller than it should be	untouched tail slots stay zero

Verification

Suite	Result
pytest examples tests/st --platform a2a3sim	58/58 cases
pytest examples tests/st --platform a5sim	43/43 cases
pytest examples tests/st --platform a2a3 (onboard, task-submit)	54/54 cases
pytest tests/ut -m "not requires_hardware"	803 passed
ctest C++ unit tests	60/60
pre-commit (all hooks)	passed

Confirmed on a2a3sim that the sim auto width takes effect and is really spent:

prepare_launch_shape: block_dim auto-resolved to 8
run: Launching 24 AICore thread(s)          # was 72
aicpu_orchestration_entry: [available_aicore_counts] clusters=8 aiv=16

One earlier onboard full-suite run had two allreduce L3 cases fail with an
AICore ub address out of bounds; it did not reproduce (10/10 in isolation,
54/54 on an immediate re-run) and cannot come from this change, since those
cases run on tensormap_and_ringbuffer whose bind path never reads
worker_count. Logged locally rather than dismissed.

Follow-up

This is the groundwork for removing the user-specified block_dim entirely, so
runs always take the whole device and SPMD cohorts size themselves from
rt_available_cluster_count(). That flips the 178 pinned cases and reworks the
require_sync_start cases — see #1309.

[coderabbitai]

📝 Walkthrough

Walkthrough

The runtime ops tables now expose available cluster counts in both runtime variants. Orchestration adds helpers for AIC counts and derives AIV counts by doubling the AIC value. Implementations return total_cluster_count.

Changes

Runtime cluster count query

Layer / File(s)	Summary
Count contracts and orchestration helpers src/a2a3/runtime/.../pto_orchestration_api.h, src/a2a3/runtime/.../pto_runtime2.h	Both ops-table declarations add available_cluster_count; orchestration exposes AIC and derived AIV count helpers.
Runtime count implementation and dispatch src/a2a3/runtime/.../pto_runtime2.cpp	Both runtime implementations return orchestrator.total_cluster_count and register the callback in s_runtime_ops.

Estimated code review effort: 2 (Simple) | ~10 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Orchestration
  participant RuntimeOps
  participant RuntimeImplementation
  participant PTO2Runtime
  Orchestration->>RuntimeOps: query available_cluster_count
  RuntimeOps->>RuntimeImplementation: dispatch callback
  RuntimeImplementation->>PTO2Runtime: read orchestrator.total_cluster_count
  RuntimeImplementation-->>Orchestration: return AIC count
  Orchestration->>Orchestration: derive AIV count as AIC count * 2

Loading

Poem

A rabbit hops through ops so neat,
Counts of clusters land at its feet.
AICs are read, AIVs twice,
Runtime tables make it precise.
“Hop!” says the hare, “the numbers align!”

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly matches the main change: adding orchestration helpers for available cluster/AIV counts.
Description check	✅ Passed	The description is directly about exposing runtime cluster/AIV counts through orchestration helpers, matching the changeset.

---

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