diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index 43c2e6e9c..5b4e8b870 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -873,6 +873,17 @@ NB_MODULE(_task_interface, m) { "Launch a callable_id from a TaskArgs (used for in-process callers). " "Returns None; timing is emitted as `[STRACE]` log markers." ) + .def( + "prewarm_config", + [](ChipWorker &self, const CallConfig &config) { + self.prewarm_config(config); + }, + nb::arg("config"), + "Eagerly build + cache the prebuilt runtime-arena image for config's ring " + "sizing so the first run() with the same sizing skips the (~800ms) cold " + "build. Config-only (no callable_id, no args); requires init() first. " + "A no-op for runtimes without a prebuilt arena (host_build_graph)." + ) .def( "run_from_blob", [](ChipWorker &self, int32_t callable_id, uint64_t args_blob_ptr, size_t blob_capacity, diff --git a/python/bindings/worker_bind.h b/python/bindings/worker_bind.h index 5b4893b63..cbba67f6a 100644 --- a/python/bindings/worker_bind.h +++ b/python/bindings/worker_bind.h @@ -435,6 +435,16 @@ inline void bind_worker(nb::module_ &m) { }, nb::arg("worker_id"), nb::arg("digest"), "Prewarm a NEXT_LEVEL child by callable digest." ) + .def( + "control_prewarm_config", + [](Worker &self, int worker_id, const CallConfig &config) { + nb::gil_scoped_release release; + self.control_prewarm_config(worker_id, &config, sizeof(CallConfig)); + }, + nb::arg("worker_id"), nb::arg("config"), + "Prewarm a NEXT_LEVEL child's prebuilt runtime-arena cache from a " + "CallConfig so its first run skips the cold arena build." + ) .def( "broadcast_register_all", [](Worker &self, uint64_t blob_ptr, uint64_t blob_size, nb::object digest) { diff --git a/python/simpler/task_interface.py b/python/simpler/task_interface.py index b29e24911..9f69b7f31 100644 --- a/python/simpler/task_interface.py +++ b/python/simpler/task_interface.py @@ -1153,6 +1153,24 @@ def run(self, handle, args, config=None, **kwargs): state = self._resolve_handle(handle) self._run_slot(state.slot_id, args, config, **kwargs) + def prewarm_config(self, config=None, **kwargs) -> None: + """Eagerly build + cache the prebuilt runtime-arena image for ``config``. + + Config-only warm-up: only the ``runtime_env`` ring sizing is read, so no + callable handle or args are needed. The first ``run`` with the same ring + sizing then hits the cache and skips the (~800ms) cold build. A no-op for + runtimes without a prebuilt arena (host_build_graph). Requires ``init``. + + Args: + config: Optional CallConfig; if None a default is created. + **kwargs: Overrides applied to config (e.g. ``ring_heap=...``). + """ + if config is None: + config = CallConfig() + for k, v in kwargs.items(): + setattr(config, k, v) + self._impl.prewarm_config(config) + def unregister_callable(self, handle) -> None: """Drop one live callable handle and release its private resources when final.""" with self._registry_lock: diff --git a/python/simpler/worker.py b/python/simpler/worker.py index c0ebba1d0..b01ed21ad 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -219,6 +219,12 @@ def my_l4_orch(orch, args, config): _CTRL_MAP_HOST = 14 _CTRL_UNMAP_HOST = 15 +# Eagerly build + cache the prebuilt runtime-arena image for a run config's ring +# sizing so the first run() with the same sizing skips the (~800ms) cold build. +# Sent from the parent at end of init() alongside CTRL_PREPARE. Carries the +# packed CallConfig at _OFF_CONFIG (same slot a task dispatch uses). +_CTRL_PREWARM_CONFIG = 16 + # MAP_HOST payload: token (u64), parent_va (u64), nbytes (u64), then the # NUL-free host-buffer shm name as the trailing bytes. UNMAP_HOST payload is the # token alone. @@ -1226,6 +1232,12 @@ def _run_chip_main_loop( # noqa: PLR0912, PLR0913, PLR0915 -- unified TASK_READ f"prepare chip={device_id}: callable hash {_format_digest(digest)} not registered" ) _ensure_prepared(cw, registry, prepared, int(cid), device_id=device_id) + elif sub_cmd == _CTRL_PREWARM_CONFIG: + # Config-only: build + cache the prebuilt runtime arena for + # this run config's ring sizing so the first run() hits the + # cache. No callable digest — the arena is callable-agnostic. + cfg = _read_config_from_mailbox(buf) + cw.prewarm_config(cfg) elif sub_cmd == _CTRL_REGISTER: digest = _read_control_digest(buf) payload_size = struct.unpack_from("Q", buf, _CTRL_OFF_ARG0)[0] @@ -1566,6 +1578,11 @@ def __init__( self._hierarchical_start_mu = threading.Lock() self._hierarchical_start_cv = threading.Condition(self._hierarchical_start_mu) + # Optional CallConfig whose ring sizing is pre-warmed at init() so the + # first run() with the same sizing skips the (~800ms) cold prebuilt + # runtime-arena build. Set by init(prewarm_config=...); None = disabled. + self._prewarm_config: Any | None = None + # Level-2 internals self._chip_worker: ChipWorker | None = None @@ -3046,9 +3063,20 @@ def add_worker(self, worker: Worker) -> int: # init — auto-discovery # ------------------------------------------------------------------ - def init(self) -> None: + def init(self, prewarm_config=None) -> None: + """Initialize the worker. + + Args: + prewarm_config: Optional CallConfig. When given, its ring sizing + (``runtime_env.ring_task_window`` / ``ring_heap`` / + ``ring_dep_pool``) is eagerly built + cached at init so the first + ``run`` with the same sizing skips the (~800ms) cold prebuilt + runtime-arena build. A no-op for runtimes without a prebuilt arena + (host_build_graph). ``None`` (default) disables prewarm. + """ if self._initialized: raise RuntimeError("Worker already initialized") + self._prewarm_config = prewarm_config try: if self.level == 2: @@ -3083,6 +3111,11 @@ def _init_level2(self) -> None: if isinstance(target, ChipCallable): self._chip_worker._register_callable_at_slot(cid, target) + # Pre-warm the prebuilt runtime-arena cache for the declared config so the + # first run() with matching ring sizing skips the cold arena build. + if self._prewarm_config is not None: + self._chip_worker.prewarm_config(self._prewarm_config) + def _init_hierarchical(self) -> None: device_ids = self._config.get("device_ids", []) n_sub = self._config.get("num_sub_workers", 0) @@ -3311,6 +3344,14 @@ def _start_hierarchical(self) -> None: # noqa: PLR0912 -- three parallel fork l for worker_id in range(len(self._chip_shms)): dw.control_prepare(worker_id, digest) + # Pre-warm each chip child's prebuilt runtime-arena cache for the + # declared config so the first dispatch with matching ring sizing + # skips the (~800ms) cold arena build. No-op when no prewarm config + # was passed to init() or for runtimes without a prebuilt arena. + if device_ids and self._prewarm_config is not None: + for worker_id in range(len(self._chip_shms)): + dw.control_prewarm_config(worker_id, self._prewarm_config) + self._hierarchical_started = True with self._hierarchical_start_cv: self._hierarchical_start_state = "started" diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp index 1fef29985..7cd5bfada 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp @@ -622,8 +622,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 = @@ -651,7 +650,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) { @@ -678,8 +676,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 @@ -705,24 +703,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 @@ -754,7 +734,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) { @@ -766,6 +746,40 @@ static void store_prebuilt_runtime_image( ); } +// Populate the DeviceRunnerBase prebuilt-arena cache for `sizing`: reserve the +// pooled arenas, build the host image, rtMemcpy it to the pooled runtime-arena +// region, and record it in the cache. Needs no Runtime and no per-run args — +// the image is arg-independent (args are wired only on the cache-hit path via +// bind_cached_runtime_image). Shared by the lazy first-run miss path and the +// eager prewarm_config_impl entry, so both populate the cache identically. +static bool build_and_cache_prebuilt_arena(const HostApi *api, const ArenaSizingConfig &sizing) { + 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); + 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 @@ -777,9 +791,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 @@ -856,26 +870,18 @@ extern "C" int bind_callable_to_runtime_impl( return -1; } if (cache_rc != 0) { - ArenaStaticSizes sizes; - if (!derive_arena_static_sizes(sizing, &sizes)) { + // Miss: build + upload + cache the arena image, then re-run the + // cache lookup — which now hits and wires the runtime (args, gm/sm + // pointer, prebuilt-arena base). This is the same code the eager + // prewarm_config_impl path runs, so the cache contract is identical. + if (!build_and_cache_prebuilt_arena(api, sizing)) { return -1; } - - StaticArenaPtrs ptrs; - if (!ensure_static_arenas(runtime, api, sizing, sizes, &ptrs)) { + int rebind_rc = bind_cached_runtime_image(runtime, api, cache_probe, device_args); + if (rebind_rc != 0) { + LOG_ERROR("prebuilt runtime arena cache miss after eager build (rc=%d)", rebind_rc); 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)) { - return -1; - } - store_prebuilt_runtime_image(runtime, api, cache_probe, ptrs, layout, host_arena); } } int64_t t_prebuilt_end = _now_ms(); @@ -890,6 +896,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/src/common/hierarchical/worker.h b/src/common/hierarchical/worker.h index e915b428f..2e87732ad 100644 --- a/src/common/hierarchical/worker.h +++ b/src/common/hierarchical/worker.h @@ -100,6 +100,12 @@ class Worker { // used by the Python facade at end of _start_hierarchical). void control_prepare(int worker_id, const uint8_t *digest) { manager_.control_prepare(worker_id, digest); } + // Forward CTRL_PREWARM_CONFIG to a specific NEXT_LEVEL worker (config-arena + // prewarm path used by the Python facade at end of _start_hierarchical). + void control_prewarm_config(int worker_id, const void *config, size_t config_size) { + manager_.control_prewarm_config(worker_id, config, config_size); + } + // Drive a single chip child through one CommDomain alloc / release. The // Python orch facade is expected to call this on every participating chip // in parallel (one thread per chip) so the child-side `file_barrier` can diff --git a/src/common/hierarchical/worker_manager.cpp b/src/common/hierarchical/worker_manager.cpp index ccc526682..729c49f7a 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -563,6 +563,17 @@ void LocalMailboxEndpoint::control_prepare(const uint8_t *digest) { run_control_command("control_prepare"); } +void LocalMailboxEndpoint::control_prewarm_config(const void *config, size_t config_size) { + std::lock_guard lk(mailbox_mu_); + write_control_args(mbox(), CTRL_PREWARM_CONFIG); + // Stage the packed CallConfig at MAILBOX_OFF_CONFIG (same slot a task dispatch + // uses) so the child reads it with its normal config decode. Clamp to the + // CallConfig region so a mis-sized caller cannot overrun the args area. + size_t n = config_size < sizeof(CallConfig) ? config_size : sizeof(CallConfig); + std::memcpy(mbox() + MAILBOX_OFF_CONFIG, config, n); + run_control_command("control_prewarm_config"); +} + void LocalMailboxEndpoint::control_register(const char *shm_name, size_t blob_size, const uint8_t *digest) { std::lock_guard lk(mailbox_mu_); // OFF_ERROR / OFF_ERROR_MSG are cleared by run_control_command — no @@ -749,6 +760,11 @@ void WorkerThread::control_prepare(const uint8_t *digest) { endpoint_->control_prepare(digest); } +void WorkerThread::control_prewarm_config(const void *config, size_t config_size) { + if (!endpoint_) throw std::runtime_error("control_prewarm_config: null endpoint"); + endpoint_->control_prewarm_config(config, config_size); +} + void WorkerThread::control_register(const char *shm_name, size_t blob_size, const uint8_t *digest) { if (!endpoint_) throw std::runtime_error("control_register: null endpoint"); endpoint_->control_register(shm_name, blob_size, digest); @@ -977,6 +993,14 @@ void WorkerManager::control_prepare(int worker_id, const uint8_t *digest) { wt->control_prepare(digest); } +void WorkerManager::control_prewarm_config(int worker_id, const void *config, size_t config_size) { + auto *wt = get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); + if (wt == nullptr) { + throw std::runtime_error("control_prewarm_config: invalid worker_id " + std::to_string(worker_id)); + } + wt->control_prewarm_config(config, config_size); +} + void WorkerManager::control_alloc_domain(int worker_id, const char *request_shm_name, const char *reply_shm_name) { auto *wt = get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); if (wt == nullptr) { diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index f1b0614ba..6d93d496d 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -152,6 +152,11 @@ static constexpr uint64_t CTRL_COMM_INIT = 9; static constexpr uint64_t CTRL_PY_REGISTER = 10; static constexpr uint64_t CTRL_PY_UNREGISTER = 11; static constexpr uint64_t CTRL_L3_L2_ORCH_COMM_INIT = 13; +// Eagerly build + cache the prebuilt runtime-arena image for a run config's ring +// sizing, so the first simpler_run with the same sizing skips the (~800ms) cold +// build. Issued at end of init() alongside CTRL_PREPARE. Carries the packed +// CallConfig bytes at MAILBOX_OFF_CONFIG (same layout as a task dispatch). +static constexpr uint64_t CTRL_PREWARM_CONFIG = 16; // Control args reuse the task mailbox region (mutually exclusive with task dispatch): // offset 16: uint64 arg0 (size for malloc/register; ptr for free; dst for copy) @@ -205,6 +210,9 @@ class WorkerEndpoint { virtual void control_copy_to(uint64_t dst, uint64_t src, size_t size); virtual void control_copy_from(uint64_t dst, uint64_t src, size_t size); virtual void control_prepare(const uint8_t *digest); + // Prewarm is optional: the base default is a no-op so endpoints that do not + // support it (remote) simply skip it. LocalMailboxEndpoint overrides. + virtual void control_prewarm_config(const void * /*config*/, size_t /*config_size*/) {} virtual void control_register(const char *shm_name, size_t blob_size, const uint8_t *digest); virtual void control_unregister(const uint8_t *digest); virtual void control_remote_prepare_register( @@ -255,6 +263,7 @@ class LocalMailboxEndpoint : public WorkerEndpoint { void control_copy_to(uint64_t dst, uint64_t src, size_t size) override; void control_copy_from(uint64_t dst, uint64_t src, size_t size) override; void control_prepare(const uint8_t *digest) override; + void control_prewarm_config(const void *config, size_t config_size) override; void control_register(const char *shm_name, size_t blob_size, const uint8_t *digest) override; void control_unregister(const uint8_t *digest) override; void control_remote_prepare_register( @@ -377,6 +386,10 @@ class WorkerThread { // target-local slot via CTRL_PREPARE. void control_prepare(const uint8_t *digest); + // Pre-warm a chip child's prebuilt runtime-arena cache for a run config via + // CTRL_PREWARM_CONFIG. `config` points to the packed CallConfig bytes. + void control_prewarm_config(const void *config, size_t config_size); + // Dynamic post-init register/unregister of a ChipCallable identity. // `shm_name` is the (NUL-terminated, ≤ CTRL_SHM_NAME_BYTES-1) POSIX shm // name where the ChipCallable bytes are staged; `blob_size` is the exact @@ -481,6 +494,10 @@ class WorkerManager { // Python facade can prewarm without reaching into individual WorkerThreads. void control_prepare(int worker_id, const uint8_t *digest); + // Forward CTRL_PREWARM_CONFIG to a specific NEXT_LEVEL worker (config-arena + // prewarm). `config` points to the packed CallConfig bytes. + void control_prewarm_config(int worker_id, const void *config, size_t config_size); + // Forward CTRL_ALLOC_DOMAIN / CTRL_RELEASE_DOMAIN to a specific NEXT_LEVEL // worker. Used by the Python orch facade to drive collective domain // allocation across a subset of chips — caller dispatches to each diff --git a/src/common/platform/onboard/host/c_api_shared.cpp b/src/common/platform/onboard/host/c_api_shared.cpp index 42f6c86f3..8111ccb69 100644 --- a/src/common/platform/onboard/host/c_api_shared.cpp +++ b/src/common/platform/onboard/host/c_api_shared.cpp @@ -571,6 +571,31 @@ int simpler_run( } } +int simpler_prewarm_config(DeviceContextHandle ctx, const CallConfig *config) { + if (ctx == NULL || config == NULL) return -1; + DeviceRunnerBase *runner = static_cast(ctx); + + pthread_once(&g_runner_key_once, create_runner_key); + pthread_setspecific(g_runner_key, ctx); + auto tsd_guard = RAIIScopeGuard([]() { + pthread_setspecific(g_runner_key, nullptr); + }); + + STRACE_NEW_INV(); + STRACE("simpler_prewarm"); + + try { + int rc = runner->attach_current_thread(runner->device_id()); + if (rc != 0) return rc; + return runner->prewarm_config( + &g_host_api, config->runtime_env.ring_task_window, config->runtime_env.ring_heap, + config->runtime_env.ring_dep_pool + ); + } catch (...) { + return -1; + } +} + int simpler_unregister_callable(DeviceContextHandle ctx, int32_t callable_id) { if (ctx == NULL) return -1; try { diff --git a/src/common/platform/onboard/host/device_runner_base.cpp b/src/common/platform/onboard/host/device_runner_base.cpp index 79378df63..d28e4cfa3 100644 --- a/src/common/platform/onboard/host/device_runner_base.cpp +++ b/src/common/platform/onboard/host/device_runner_base.cpp @@ -955,6 +955,25 @@ int DeviceRunnerBase::bind_callable_to_runtime( ); } +// Eager prebuilt-arena warm-up half. A runtime that has a prebuilt runtime arena +// (tensormap_and_ringbuffer) provides a strong prewarm_config_impl in its +// runtime_maker.cpp that overrides this weak no-op default. Runtimes without one +// (host_build_graph, or an arch that has not implemented it yet) link this weak +// default and treat prewarm as a no-op — so DeviceRunner::prewarm_config always +// resolves without every arch/runtime having to define the symbol. +extern "C" __attribute__((weak)) int prewarm_config_impl( + const HostApi * /*api*/, const uint64_t * /*ring_task_window*/, const uint64_t * /*ring_heap*/, + const uint64_t * /*ring_dep_pool*/ +) { + return 0; +} + +int DeviceRunnerBase::prewarm_config( + const HostApi *api, const uint64_t *ring_task_window, const uint64_t *ring_heap, const uint64_t *ring_dep_pool +) { + return prewarm_config_impl(api, ring_task_window, ring_heap, ring_dep_pool); +} + void DeviceRunnerBase::apply_call_config(const CallConfig &config) { set_l2_swimlane_enabled(config.enable_l2_swimlane); set_dump_tensor_enabled(config.enable_dump_tensor); diff --git a/src/common/platform/onboard/host/device_runner_base.h b/src/common/platform/onboard/host/device_runner_base.h index 9f7616f3f..6a74a703a 100644 --- a/src/common/platform/onboard/host/device_runner_base.h +++ b/src/common/platform/onboard/host/device_runner_base.h @@ -352,6 +352,22 @@ class DeviceRunnerBase : public L3L2OrchCommBackend { const uint64_t *ring_task_window, const uint64_t *ring_heap, const uint64_t *ring_dep_pool ); + /** + * Eagerly build + upload + cache the prebuilt runtime-arena image for a run + * config, so the first bind_callable_to_runtime with the same ring sizing + * hits the cache and skips the build (~800ms on the first cold run). + * Config-only: needs no callable_id and no per-run args — the image depends + * solely on the ring sizing. Forwards to the runtime's prewarm_config_impl + * (trb builds + caches; hbg is a no-op). The device must already be + * initialized so the pooled-arena device_malloc + rtMemcpy have a context. + * + * @param ring_task_window Per-ring overrides (trb); ignored by hbg. + * @return 0 on success, non-zero on failure. + */ + int prewarm_config( + const HostApi *api, const uint64_t *ring_task_window, const uint64_t *ring_heap, const uint64_t *ring_dep_pool + ); + /** * Number of distinct callable_ids the AICPU has been asked to * dlopen for. Monotonically increases when an AICPU load succeeds diff --git a/src/common/platform/sim/host/c_api_shared.cpp b/src/common/platform/sim/host/c_api_shared.cpp index 760c9245a..0fb9f1781 100644 --- a/src/common/platform/sim/host/c_api_shared.cpp +++ b/src/common/platform/sim/host/c_api_shared.cpp @@ -521,6 +521,29 @@ int simpler_run( } } +int simpler_prewarm_config(DeviceContextHandle ctx, const CallConfig *config) { + if (ctx == NULL || config == NULL) return -1; + SimDeviceRunnerBase *runner = static_cast(ctx); + + pthread_once(&g_runner_key_once, create_runner_key); + pthread_setspecific(g_runner_key, ctx); + + STRACE_NEW_INV(); + STRACE("simpler_prewarm"); + + try { + int rc = runner->prewarm_config( + &g_host_api, config->runtime_env.ring_task_window, config->runtime_env.ring_heap, + config->runtime_env.ring_dep_pool + ); + pthread_setspecific(g_runner_key, nullptr); + return rc; + } catch (...) { + pthread_setspecific(g_runner_key, nullptr); + return -1; + } +} + int simpler_unregister_callable(DeviceContextHandle ctx, int32_t callable_id) { if (ctx == NULL) return -1; try { diff --git a/src/common/platform/sim/host/device_runner_base.cpp b/src/common/platform/sim/host/device_runner_base.cpp index d69a25803..7abf32635 100644 --- a/src/common/platform/sim/host/device_runner_base.cpp +++ b/src/common/platform/sim/host/device_runner_base.cpp @@ -574,6 +574,24 @@ int SimDeviceRunnerBase::bind_callable_to_runtime( ); } +// Eager prebuilt-arena warm-up half. A runtime with a prebuilt runtime arena +// (tensormap_and_ringbuffer) provides a strong prewarm_config_impl in its +// runtime_maker.cpp that overrides this weak no-op default; runtimes without one +// link the weak default and treat prewarm as a no-op, so SimDeviceRunner:: +// prewarm_config always resolves without every arch/runtime defining the symbol. +extern "C" __attribute__((weak)) int prewarm_config_impl( + const HostApi * /*api*/, const uint64_t * /*ring_task_window*/, const uint64_t * /*ring_heap*/, + const uint64_t * /*ring_dep_pool*/ +) { + return 0; +} + +int SimDeviceRunnerBase::prewarm_config( + const HostApi *api, const uint64_t *ring_task_window, const uint64_t *ring_heap, const uint64_t *ring_dep_pool +) { + return prewarm_config_impl(api, ring_task_window, ring_heap, ring_dep_pool); +} + void SimDeviceRunnerBase::apply_call_config(const CallConfig &config) { set_l2_swimlane_enabled(config.enable_l2_swimlane); set_dump_tensor_enabled(config.enable_dump_tensor); diff --git a/src/common/platform/sim/host/device_runner_base.h b/src/common/platform/sim/host/device_runner_base.h index 911ad53e2..2c49ea6e1 100644 --- a/src/common/platform/sim/host/device_runner_base.h +++ b/src/common/platform/sim/host/device_runner_base.h @@ -123,6 +123,13 @@ class SimDeviceRunnerBase : public L3L2OrchCommBackend { Runtime &runtime, int32_t callable_id, const HostApi *api, const void *orch_args, const uint64_t *ring_task_window, const uint64_t *ring_heap, const uint64_t *ring_dep_pool ); + // Eagerly build + upload + cache the prebuilt runtime-arena image for a run + // config so the first bind with the same ring sizing hits the cache and + // skips the build. Config-only (no callable_id, no args). Forwards to the + // runtime's prewarm_config_impl (trb builds + caches; hbg no-ops). + int prewarm_config( + const HostApi *api, const uint64_t *ring_task_window, const uint64_t *ring_heap, const uint64_t *ring_dep_pool + ); uint64_t upload_chip_callable_buffer(const ChipCallable *callable); int launch_device_register(int32_t callable_id); int commit_device_register(int32_t callable_id); diff --git a/src/common/worker/chip_worker.cpp b/src/common/worker/chip_worker.cpp index a809c9d6e..571216c1b 100644 --- a/src/common/worker/chip_worker.cpp +++ b/src/common/worker/chip_worker.cpp @@ -105,6 +105,7 @@ void ChipWorker::init( simpler_init_fn_ = load_symbol(handle, "simpler_init"); register_callable_fn_ = load_symbol(handle, "simpler_register_callable"); run_fn_ = load_symbol(handle, "simpler_run"); + prewarm_config_fn_ = load_symbol(handle, "simpler_prewarm_config"); unregister_callable_fn_ = load_symbol(handle, "simpler_unregister_callable"); get_aicpu_dlopen_count_fn_ = load_symbol(handle, "get_aicpu_dlopen_count"); get_host_dlopen_count_fn_ = load_symbol(handle, "get_host_dlopen_count"); @@ -188,6 +189,7 @@ void ChipWorker::init( simpler_init_fn_ = nullptr; register_callable_fn_ = nullptr; run_fn_ = nullptr; + prewarm_config_fn_ = nullptr; unregister_callable_fn_ = nullptr; get_aicpu_dlopen_count_fn_ = nullptr; get_host_dlopen_count_fn_ = nullptr; @@ -228,6 +230,7 @@ void ChipWorker::init( simpler_init_fn_ = nullptr; register_callable_fn_ = nullptr; run_fn_ = nullptr; + prewarm_config_fn_ = nullptr; unregister_callable_fn_ = nullptr; get_aicpu_dlopen_count_fn_ = nullptr; get_host_dlopen_count_fn_ = nullptr; @@ -282,6 +285,7 @@ void ChipWorker::finalize() { get_runtime_size_fn_ = nullptr; register_callable_fn_ = nullptr; run_fn_ = nullptr; + prewarm_config_fn_ = nullptr; unregister_callable_fn_ = nullptr; get_aicpu_dlopen_count_fn_ = nullptr; get_host_dlopen_count_fn_ = nullptr; @@ -340,6 +344,19 @@ void ChipWorker::run(int32_t callable_id, const ChipStorageTaskArgs *args, const } } +void ChipWorker::prewarm_config(const CallConfig &config) { + config.validate(); + if (!initialized_) { + throw std::runtime_error("ChipWorker not initialized; call init() first"); + } + // Config-only warm-up: only runtime_env ring sizing is read; no callable_id, + // no args, no Runtime buffer. A no-op for runtimes without a prebuilt arena. + int rc = prewarm_config_fn_(device_ctx_, &config); + if (rc != 0) { + throw std::runtime_error("prewarm_config failed with code " + std::to_string(rc)); + } +} + void ChipWorker::unregister_callable(int32_t callable_id) { if (!initialized_) { throw std::runtime_error("ChipWorker not initialized; call init() first"); diff --git a/src/common/worker/chip_worker.h b/src/common/worker/chip_worker.h index 6dc7cd214..f77b485eb 100644 --- a/src/common/worker/chip_worker.h +++ b/src/common/worker/chip_worker.h @@ -67,6 +67,12 @@ class ChipWorker { void register_callable(int32_t callable_id, const void *callable); void unregister_callable(int32_t callable_id); + // Eagerly build + cache the prebuilt runtime-arena image for `config`'s ring + // sizing so the first run() with the same sizing skips the (~800ms) cold + // build. Config-only: needs no callable_id and no args. Requires init() + // first. A no-op for runtimes without a prebuilt arena (host_build_graph). + void prewarm_config(const CallConfig &config); + /// Number of distinct callable_ids the AICPU has been asked to dlopen for /// on the bound device. Returns 0 when not initialized or the runtime /// variant has no per-cid registration support. Used by tests to assert @@ -146,6 +152,7 @@ class ChipWorker { int (*)(void *, int, const uint8_t *, size_t, const uint8_t *, size_t, const uint8_t *, size_t); using SimplerRegisterCallableFn = int (*)(void *, int32_t, const void *); using SimplerRunFn = int (*)(void *, void *, int32_t, const void *, const CallConfig *); + using SimplerPrewarmConfigFn = int (*)(void *, const CallConfig *); using SimplerUnregisterCallableFn = int (*)(void *, int32_t); using GetAicpuDlopenCountFn = size_t (*)(void *); using FinalizeDeviceFn = int (*)(void *); @@ -193,6 +200,7 @@ class ChipWorker { SimplerInitFn simpler_init_fn_ = nullptr; SimplerRegisterCallableFn register_callable_fn_ = nullptr; SimplerRunFn run_fn_ = nullptr; + SimplerPrewarmConfigFn prewarm_config_fn_ = nullptr; SimplerUnregisterCallableFn unregister_callable_fn_ = nullptr; GetAicpuDlopenCountFn get_aicpu_dlopen_count_fn_ = nullptr; GetAicpuDlopenCountFn get_host_dlopen_count_fn_ = nullptr; diff --git a/src/common/worker/pto_runtime_c_api.h b/src/common/worker/pto_runtime_c_api.h index ba6ca733e..0f4f28742 100644 --- a/src/common/worker/pto_runtime_c_api.h +++ b/src/common/worker/pto_runtime_c_api.h @@ -208,6 +208,22 @@ int simpler_run( DeviceContextHandle ctx, RuntimeHandle runtime, int32_t callable_id, const void *args, const CallConfig *config ); +/** + * Eagerly build + upload + cache the prebuilt runtime-arena image for the ring + * sizing in `config->runtime_env`, so the first `simpler_run` with the same + * sizing hits the cache and skips the (~800ms) cold build. Config-only: needs + * no callable_id and no per-run args — the arena image depends solely on the + * ring sizing, so this may be called at init time before any run. + * + * Ring overrides are consumed by tensormap_and_ringbuffer only; other runtime + * variants treat this as a no-op and return 0. Requires the device to be + * initialized (`simpler_init` done) so the pooled-arena device_malloc + rtMemcpy + * have a live context. Only `config->runtime_env` is read. + * + * @return 0 on success (including the no-op variants), negative on error. + */ +int simpler_prewarm_config(DeviceContextHandle ctx, const CallConfig *config); + /** * Drop the prepared state for `callable_id` and release the per-id share of * the device orch SO buffer. The buffer itself is freed only when its