diff --git a/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp b/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp index c355f2e11..c3f009278 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp @@ -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 = @@ -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) { @@ -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 @@ -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 @@ -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) { @@ -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 @@ -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 @@ -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(); @@ -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. * diff --git a/tests/st/a5/tensormap_and_ringbuffer/bench_prewarm_timing.py b/tests/st/a5/tensormap_and_ringbuffer/bench_prewarm_timing.py new file mode 100644 index 000000000..964f43d7e --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/bench_prewarm_timing.py @@ -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") + + +if __name__ == "__main__": + main() diff --git a/tests/st/a5/tensormap_and_ringbuffer/test_prewarm_config.py b/tests/st/a5/tensormap_and_ringbuffer/test_prewarm_config.py new file mode 100644 index 000000000..9ef4b64cd --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/test_prewarm_config.py @@ -0,0 +1,35 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Regression test for runtime-arena prewarm during L2 worker initialization.""" + +import pytest +from simpler.task_interface import CallConfig +from simpler.worker import Worker + +_RUNTIME = "tensormap_and_ringbuffer" + + +@pytest.mark.platforms(["a5sim", "a5"]) +@pytest.mark.device_count(1) +@pytest.mark.runtime(_RUNTIME) +def test_l2_init_with_prewarm_config(st_platform, st_device_ids): + config = CallConfig() + config.runtime_env.ring_task_window = 64 + + worker = Worker( + level=2, + device_id=int(st_device_ids[0]), + platform=st_platform, + runtime=_RUNTIME, + ) + try: + worker.init(prewarm_config=config) + finally: + worker.close()