Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
131 changes: 88 additions & 43 deletions src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -620,8 +620,7 @@ static void apply_orch_sched_env_flags(Runtime *runtime) {
// reserve sequence on a throwaway host arena. Idempotent across runs — the
// pools are owned by DeviceRunner and freed in DeviceRunner::finalize().
static bool ensure_static_arenas(
Runtime *runtime, const HostApi *api, const ArenaSizingConfig &sizing, const ArenaStaticSizes &sizes,
StaticArenaPtrs *out
const HostApi *api, const ArenaSizingConfig &sizing, const ArenaStaticSizes &sizes, StaticArenaPtrs *out
) {
DeviceArena sizing_arena; // discarded; only its computed arena_size is read
PTO2RuntimeArenaLayout layout =
Expand Down Expand Up @@ -649,7 +648,6 @@ static bool ensure_static_arenas(
LOG_ERROR("Failed to acquire pooled PTO2 shared memory");
return false;
}
runtime->set_gm_sm_ptr(out->gm_sm);

out->runtime_arena_dev = api->acquire_pooled_runtime_arena();
if (out->runtime_arena_dev == nullptr) {
Expand All @@ -676,8 +674,8 @@ static bool ensure_static_arenas(
// The layout is stashed inside the image (rt->prebuilt_layout) so the AICPU can
// recover every arena-internal offset after the rtMemcpy. Returns the layout
// via `out_layout`; the runtime-arena device base travels separately on the
// host Runtime (bind_launch_state), since the AICPU needs that pointer *before*
// it can dereference the image.
// host Runtime (set on the cache-hit path), since the AICPU needs that pointer
// *before* it can dereference the image.
static bool build_runtime_image(
const ArenaSizingConfig &sizing, const ArenaStaticSizes &sizes, const StaticArenaPtrs &ptrs,
DeviceArena *host_arena, PTO2RuntimeArenaLayout *out_layout
Expand All @@ -703,24 +701,6 @@ static bool build_runtime_image(
return true;
}

// per-run: publish the launch state. Copy the staged args onto the runtime,
// rtMemcpy the host image into the pooled runtime-arena region, and record the
// device base + runtime offset the AICPU reads before dereferencing the image.
static bool bind_launch_state(
Runtime *runtime, const HostApi *api, const StaticArenaPtrs &ptrs, const DeviceArena &host_arena,
const PTO2RuntimeArenaLayout &layout, const ChipStorageTaskArgs &device_args
) {
runtime->set_orch_args(device_args);

int rc_upload = api->copy_to_device(ptrs.runtime_arena_dev, host_arena.base(), layout.offsets.arena_size);
if (rc_upload != 0) {
LOG_ERROR("Failed to rtMemcpy prebuilt runtime arena to device (rc=%d)", rc_upload);
return false;
}
runtime->set_prebuilt_arena(ptrs.runtime_arena_dev, layout.offsets.off_runtime);
return true;
}

static int bind_cached_runtime_image(
Runtime *runtime, const HostApi *api, const PrebuiltRuntimeArenaCacheProbe &probe,
const ChipStorageTaskArgs &device_args
Expand Down Expand Up @@ -752,7 +732,7 @@ static int bind_cached_runtime_image(
}

static void store_prebuilt_runtime_image(
Runtime *runtime, const HostApi *api, const PrebuiltRuntimeArenaCacheProbe &probe, const StaticArenaPtrs &ptrs,
const HostApi *api, const PrebuiltRuntimeArenaCacheProbe &probe, const StaticArenaPtrs &ptrs,
const PTO2RuntimeArenaLayout &layout, const DeviceArena &host_arena
) {
if (api->mark_prebuilt_runtime_arena_cached == nullptr) {
Expand All @@ -764,6 +744,51 @@ static void store_prebuilt_runtime_image(
);
}

// Reserve the pooled arenas, build the host image, rtMemcpy it to the pooled
// runtime-arena region, and record it in the DeviceRunnerBase prebuilt-arena
// cache for `sizing`. Needs no Runtime and no per-run args — the image is
// arg-independent. The cache store is best-effort (a no-op on backends without
// cache callbacks); `out_ptrs`/`out_layout` return the freshly built arena so
// the run path can wire the runtime directly instead of depending on a cache
// round-trip. Shared by the lazy first-run miss path and the eager
// prewarm_config_impl entry, so both build the arena identically.
static bool build_and_cache_prebuilt_arena(
const HostApi *api, const ArenaSizingConfig &sizing, StaticArenaPtrs *out_ptrs = nullptr,
PTO2RuntimeArenaLayout *out_layout = nullptr
) {
ArenaStaticSizes sizes;
if (!derive_arena_static_sizes(sizing, &sizes)) {
return false;
}

StaticArenaPtrs ptrs;
if (!ensure_static_arenas(api, sizing, sizes, &ptrs)) {
return false;
}

DeviceArena host_arena; // libc malloc backend; owns the image until upload
PTO2RuntimeArenaLayout layout;
if (!build_runtime_image(sizing, sizes, ptrs, &host_arena, &layout)) {
return false;
}

int rc_upload = api->copy_to_device(ptrs.runtime_arena_dev, host_arena.base(), layout.offsets.arena_size);
if (rc_upload != 0) {
LOG_ERROR("Failed to rtMemcpy prebuilt runtime arena to device (rc=%d)", rc_upload);
return false;
}

PrebuiltRuntimeArenaCacheProbe probe = make_prebuilt_runtime_arena_cache_probe(sizing);
store_prebuilt_runtime_image(api, probe, ptrs, layout, host_arena);
if (out_ptrs != nullptr) {
*out_ptrs = ptrs;
}
if (out_layout != nullptr) {
*out_layout = layout;
}
return true;
}

/**
* Per-run binding: build device-side argument storage (tensor copy-out, GM
* heap, PTO2 shared memory) and publish it to the runtime. Assumes the
Expand All @@ -775,9 +800,9 @@ static void store_prebuilt_runtime_image(
* half runs only once per callable_id.
*
* Orchestrates the three lifecycles behind the bind: per-config arena sizing
* (resolve_arena_sizing) + static pools (ensure_static_arenas) + host image
* (build_runtime_image), and per-run args (stage_device_args) + launch publish
* (bind_launch_state).
* (resolve_arena_sizing) + per-run args (stage_device_args) + the prebuilt
* runtime-arena image (build_and_cache_prebuilt_arena on a cache miss, then
* bind_cached_runtime_image wires the pointers onto the runtime).
*
* @param runtime Pointer to pre-constructed Runtime
* @param orch_args Separated tensor/scalar arguments for this run
Expand Down Expand Up @@ -854,26 +879,20 @@ extern "C" int bind_callable_to_runtime_impl(
return -1;
}
if (cache_rc != 0) {
ArenaStaticSizes sizes;
if (!derive_arena_static_sizes(sizing, &sizes)) {
return -1;
}

// Miss: build + upload the arena image, then wire the runtime
// directly from the freshly built arena (same three fields the
// cache-hit path sets). The store inside build_and_cache is
// best-effort for the NEXT bind — this bind must not depend on the
// cache round-trip, so a backend with no-op cache callbacks still
// binds successfully.
StaticArenaPtrs ptrs;
if (!ensure_static_arenas(runtime, api, sizing, sizes, &ptrs)) {
return -1;
}

DeviceArena host_arena; // libc malloc backend; owns the image until upload
PTO2RuntimeArenaLayout layout;
if (!build_runtime_image(sizing, sizes, ptrs, &host_arena, &layout)) {
return -1;
}

if (!bind_launch_state(runtime, api, ptrs, host_arena, layout, device_args)) {
if (!build_and_cache_prebuilt_arena(api, sizing, &ptrs, &layout)) {
return -1;
}
store_prebuilt_runtime_image(runtime, api, cache_probe, ptrs, layout, host_arena);
runtime->set_orch_args(device_args);
runtime->set_gm_sm_ptr(ptrs.gm_sm);
runtime->set_prebuilt_arena(ptrs.runtime_arena_dev, layout.offsets.off_runtime);
}
}
int64_t t_prebuilt_end = _now_ms();
Expand All @@ -888,6 +907,32 @@ extern "C" int bind_callable_to_runtime_impl(
return 0;
}

/**
* Eagerly populate the prebuilt runtime-arena cache for a run config, so the
* first bind_callable_to_runtime_impl with the same sizing hits the cache and
* skips the (~800ms) build + upload. Config-only: no callable, no per-run args
* — the arena image depends solely on the ring sizing. Requires the device to
* be initialized (pooled-arena device_malloc + rtMemcpy need a live context).
*
* @return 0 on success, -1 on failure
*/
extern "C" int prewarm_config_impl(
const HostApi *api, const uint64_t *ring_task_window, const uint64_t *ring_heap, const uint64_t *ring_dep_pool
) {
if (api == nullptr) {
LOG_ERROR("HostApi pointer is null");
return -1;
}

ArenaSizingConfig sizing;
if (!resolve_arena_sizing(ring_task_window, ring_heap, ring_dep_pool, &sizing)) {
return -1;
}

STRACE("simpler_prewarm.build");
return build_and_cache_prebuilt_arena(api, sizing) ? 0 : -1;
}

/**
* Validate runtime results and cleanup.
*
Expand Down
188 changes: 188 additions & 0 deletions tests/st/a5/tensormap_and_ringbuffer/bench_prewarm_timing.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,188 @@
#!/usr/bin/env python3
# Copyright (c) PyPTO Contributors.
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
# CANN Open Software License Agreement Version 2.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# -----------------------------------------------------------------------------------------------------------
"""Onboard benchmark: cold (no prewarm) vs hot (init prewarm) prebuilt_runtime_arena TIMING.

Each scenario runs in a fresh process so the DeviceRunner prebuilt-arena cache does
not carry over between cold and hot measurements.
"""

from __future__ import annotations

import argparse
import logging
import re
import subprocess
import sys
from pathlib import Path

from simpler import _log
from simpler.task_interface import CallConfig
from simpler.worker import Worker

from simpler_setup.scene_test import _build_chip_task_args

_HERE = Path(__file__).resolve().parent
if str(_HERE / "dummy_task") not in sys.path:
sys.path.insert(0, str(_HERE / "dummy_task"))

from test_dummy_task import TestDummyTask # noqa: E402

_RUNTIME = "tensormap_and_ringbuffer"
_CASE_NAME = "SingleDummyAutoDep"
_RING = 64

_TIMING_RE = re.compile(r"TIMING: prebuilt_runtime_arena = (\d+)ms")
_STRACE_BIND_RE = re.compile(
r"\[STRACE\].*name=simpler_run\.bind\.prebuilt\s+ts=\d+\s+dur=(\d+)"
)
_STRACE_PREWARM_RE = re.compile(
r"\[STRACE\].*name=simpler_prewarm\.build\s+ts=\d+\s+dur=(\d+)"
)


def _last_ms_from_ns(values: list[int]) -> float | None:
if not values:
return None
return values[-1] / 1_000_000.0


def _parse_log(text: str) -> dict[str, float | None]:
timing_ms = [int(m) for m in _TIMING_RE.findall(text)]
bind_ns = [int(m) for m in _STRACE_BIND_RE.findall(text)]
prewarm_ns = [int(m) for m in _STRACE_PREWARM_RE.findall(text)]
bind_ms = _last_ms_from_ns(bind_ns)
return {
"prebuilt_runtime_arena_ms": float(timing_ms[-1]) if timing_ms else bind_ms,
"simpler_prewarm_build_ms": _last_ms_from_ns(prewarm_ns),
"bind_prebuilt_strace_ms": bind_ms,
}


def _run_scenario(*, device_id: int, platform: str, prewarm: bool) -> None:
logging.getLogger("simpler").setLevel(_log.V9)

case = next(c for c in TestDummyTask.CASES if c["name"] == _CASE_NAME)
callable_obj = TestDummyTask.compile_chip_callable(platform)

worker = Worker(
level=2,
device_id=device_id,
platform=platform,
runtime=_RUNTIME,
)
handle = worker.register(callable_obj)

ring_cfg = CallConfig()
ring_cfg.runtime_env.ring_task_window = _RING

try:
if prewarm:
worker.init(prewarm_config=ring_cfg)
else:
worker.init()

params = case.get("params", {})
config_dict = case.get("config", {})
orch_sig = TestDummyTask.CALLABLE.get("orchestration", {}).get("signature", [])

test_args = TestDummyTask().generate_args(params)
chip_args, _output_names = _build_chip_task_args(test_args, orch_sig)

run_cfg = CallConfig()
run_cfg.block_dim = config_dict.get("block_dim", 1)
run_cfg.aicpu_thread_num = config_dict.get("aicpu_thread_num", 2)
run_cfg.runtime_env.ring_task_window = _RING

worker.run(handle, chip_args, config=run_cfg)
finally:
worker.close()


def _child_main(mode: str, device_id: int, platform: str) -> int:
_run_scenario(device_id=device_id, platform=platform, prewarm=(mode == "hot"))
return 0


def _spawn(mode: str, device_id: int, platform: str) -> tuple[str, dict[str, float | None]]:
proc = subprocess.run(
[
sys.executable,
str(Path(__file__).resolve()),
"--child",
mode,
"--device",
str(device_id),
"--platform",
platform,
],
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
check=False,
)
log = proc.stdout or ""
if proc.returncode != 0:
print(log, file=sys.stderr)
raise SystemExit(proc.returncode)
return log, _parse_log(log)


def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--platform", default="a5", choices=["a5", "a5sim"])
parser.add_argument("--device", type=int, required=True)
parser.add_argument("--mode", choices=["cold", "hot", "both"], default="both")
parser.add_argument("--child", choices=["cold", "hot"], help=argparse.SUPPRESS)
args = parser.parse_args()

if args.child:
raise SystemExit(_child_main(args.child, args.device, args.platform))

modes = ["cold", "hot"] if args.mode == "both" else [args.mode]
results: dict[str, dict[str, float | None]] = {}

for mode in modes:
print(f"=== {mode.upper()} (prewarm={'yes' if mode == 'hot' else 'no'}) ===", flush=True)
log, parsed = _spawn(mode, args.device, args.platform)
results[mode] = parsed
for line in log.splitlines():
if any(
token in line
for token in (
"TIMING: prebuilt_runtime_arena",
"name=simpler_prewarm.build",
"name=simpler_run.bind.prebuilt",
)
):
print(line)
prewarm_val = parsed["simpler_prewarm_build_ms"]
prewarm_text = f"{prewarm_val:.2f}ms" if prewarm_val is not None else "None"
print(
f" -> bind.prebuilt(STRACE)={parsed['bind_prebuilt_strace_ms']:.2f}ms"
f", prewarm.build={prewarm_text}",
flush=True,
)

if args.mode == "both":
cold = results["cold"]["bind_prebuilt_strace_ms"]
hot = results["hot"]["bind_prebuilt_strace_ms"]
prewarm_ms = results["hot"]["simpler_prewarm_build_ms"]
print("\n=== SUMMARY ===")
print(f"cold first-run bind.prebuilt (STRACE) : {cold:.2f} ms")
print(f"hot first-run bind.prebuilt (STRACE) : {hot:.2f} ms")
print(f"hot init simpler_prewarm.build : {prewarm_ms:.2f} ms")
if cold is not None and hot is not None:
print(f"first-run savings (cold - hot) : {cold - hot:.2f} ms")
if prewarm_ms is not None:
print(f"(build cost shifted to init prewarm) : ~{prewarm_ms:.2f} ms")
Comment on lines +165 to +184

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Unguarded .2f formatting on possibly-None timing values.

_parse_log types bind_prebuilt_strace_ms/simpler_prewarm_build_ms as float | None, and line 166 already guards against None for prewarm_text, but lines 168, 178, 179, and 180 format these values directly with :.2f with no None check. If a TIMING/STRACE log line fails to match (e.g. log-level or format drift), this raises TypeError and the benchmark crashes without reporting partial results, unlike the guarded delta computation at 181-184.

🐛 Proposed fix for None-safe formatting
+    def _fmt(v: float | None) -> str:
+        return f"{v:.2f}ms" if v is not None else "None"
+
         prewarm_val = parsed["simpler_prewarm_build_ms"]
         prewarm_text = f"{prewarm_val:.2f}ms" if prewarm_val is not None else "None"
         print(
-            f"  -> bind.prebuilt(STRACE)={parsed['bind_prebuilt_strace_ms']:.2f}ms"
+            f"  -> bind.prebuilt(STRACE)={_fmt(parsed['bind_prebuilt_strace_ms'])}"
             f", prewarm.build={prewarm_text}",
             flush=True,
         )
 
     if args.mode == "both":
         cold = results["cold"]["bind_prebuilt_strace_ms"]
         hot = results["hot"]["bind_prebuilt_strace_ms"]
         prewarm_ms = results["hot"]["simpler_prewarm_build_ms"]
         print("\n=== SUMMARY ===")
-        print(f"cold first-run bind.prebuilt (STRACE) : {cold:.2f} ms")
-        print(f"hot  first-run bind.prebuilt (STRACE) : {hot:.2f} ms")
-        print(f"hot  init simpler_prewarm.build        : {prewarm_ms:.2f} ms")
+        print(f"cold first-run bind.prebuilt (STRACE) : {_fmt(cold)}")
+        print(f"hot  first-run bind.prebuilt (STRACE) : {_fmt(hot)}")
+        print(f"hot  init simpler_prewarm.build        : {_fmt(prewarm_ms)}")
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
prewarm_val = parsed["simpler_prewarm_build_ms"]
prewarm_text = f"{prewarm_val:.2f}ms" if prewarm_val is not None else "None"
print(
f" -> bind.prebuilt(STRACE)={parsed['bind_prebuilt_strace_ms']:.2f}ms"
f", prewarm.build={prewarm_text}",
flush=True,
)
if args.mode == "both":
cold = results["cold"]["bind_prebuilt_strace_ms"]
hot = results["hot"]["bind_prebuilt_strace_ms"]
prewarm_ms = results["hot"]["simpler_prewarm_build_ms"]
print("\n=== SUMMARY ===")
print(f"cold first-run bind.prebuilt (STRACE) : {cold:.2f} ms")
print(f"hot first-run bind.prebuilt (STRACE) : {hot:.2f} ms")
print(f"hot init simpler_prewarm.build : {prewarm_ms:.2f} ms")
if cold is not None and hot is not None:
print(f"first-run savings (cold - hot) : {cold - hot:.2f} ms")
if prewarm_ms is not None:
print(f"(build cost shifted to init prewarm) : ~{prewarm_ms:.2f} ms")
def _fmt(v: float | None) -> str:
return f"{v:.2f}ms" if v is not None else "None"
prewarm_val = parsed["simpler_prewarm_build_ms"]
prewarm_text = f"{prewarm_val:.2f}ms" if prewarm_val is not None else "None"
print(
f" -> bind.prebuilt(STRACE)={_fmt(parsed['bind_prebuilt_strace_ms'])}"
f", prewarm.build={prewarm_text}",
flush=True,
)
if args.mode == "both":
cold = results["cold"]["bind_prebuilt_strace_ms"]
hot = results["hot"]["bind_prebuilt_strace_ms"]
prewarm_ms = results["hot"]["simpler_prewarm_build_ms"]
print("\n=== SUMMARY ===")
print(f"cold first-run bind.prebuilt (STRACE) : {_fmt(cold)}")
print(f"hot first-run bind.prebuilt (STRACE) : {_fmt(hot)}")
print(f"hot init simpler_prewarm.build : {_fmt(prewarm_ms)}")
if cold is not None and hot is not None:
print(f"first-run savings (cold - hot) : {cold - hot:.2f} ms")
if prewarm_ms is not None:
print(f"(build cost shifted to init prewarm) : ~{prewarm_ms:.2f} ms")
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/st/a5/tensormap_and_ringbuffer/bench_prewarm_timing.py` around lines
165 - 184, Make all timing output in the displayed per-run and “both” summary
paths None-safe. Reuse the existing formatting approach from prewarm_text for
parsed bind_prebuilt_strace_ms, cold, hot, and prewarm_ms values, ensuring
missing timings print a fallback such as “None” instead of applying :.2f;
preserve the existing guarded savings calculations.



if __name__ == "__main__":
main()
Loading
Loading