From a7a509b3f3186119b0025f4d71906edb08f710bc Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Tue, 18 Aug 2026 17:33:31 -0700 Subject: [PATCH 1/3] perf(executorch): opt-in shared per-device activation-scratch pool A multi-layer model lowered to the ExecuTorch TensorRT delegate becomes N separate single-layer engines, each with its own execution context. By default every context allocates its own activation scratch (`getDeviceMemorySizeV2` bytes) and holds it for as long as the context lives, so device memory scales with the layer count and multi-layer models OOM at runtime. The scratch of one engine need not sit alongside the scratch of the next, because the delegates of a Method are submitted one at a time: `Method::execute()` advances `step_state_` through the instruction stream one instruction at a time on the calling thread, and a `DelegateCall` is one instruction. Their enqueues can still overlap on the device, but that is orderable, whereas N resident copies are not. Submission order is not stream order, though. That two consecutive delegates land on the same stream is not a property of `Method`; it holds only because both read the same thread-local caller stream, and a caller that runs two Methods under two different `CallerStreamGuard` streams breaks it. So the pool orders the handoff itself rather than relying on the stream. Add an opt-in shared pool that backs all contexts on a device from one buffer: - the `use_shared_activation_scratch` runtime backend option enables it. It is a boolean, defaults to false, and is delivered with `executorch::runtime::set_option("TensorRTBackend", options.view())`. With it unset each context owns its private `kSTATIC` scratch, the delegate emits no extra log output, and the only work it adds is one relaxed atomic load per engine init and a bool test or two per `execute()`; - when enabled, create each execution context with `kUSER_MANAGED` so it allocates no scratch of its own (`initialize_engine_io`); - in `execute()`, once the input shapes are bound, query the exact requirement with `updateDeviceMemorySizeForShapes()`, grow a per-device pool to it, and point the context at the current buffer via `setDeviceMemoryV2`. The pool grows monotonically to the largest engine's need and syncs the device before freeing a replaced buffer. N per-layer scratch copies collapse to one (the (N-1)x duplication is reclaimed). Measured with TensorRT 11.2.1.2 and CUDA 13 on one 80GB NVIDIA PG509-210, in a CMake reference runner that also loads the ExecuTorch CUDA/AOTI backend, reading `cudaMemGetInfo` after a `cudaFree(0)` baseline, on a deterministic non-uniform fp32 input: - four execution contexts of one engine holding two fp32 8-head attention blocks over `[1,2048,512]` (285,212,672 B of scratch) go from 1188MB to 372MB, 3 x 272MB reclaimed; - a single Method holding six one-block engines of the same shape, interleaved with six CUDA delegates (281,018,368 B each), goes from 1656MB to 316MB, 5 x 268MB reclaimed. Outputs are identical between the two modes in both cases. Two consequences a caller feels once the option is on. The pool is never freed, so a device keeps the largest scratch it was ever asked for until the process exits, where per-context `kSTATIC` scratch is released with its context. And a growth allocates the new buffer before releasing the old one, so both are resident for that moment -- that ordering is what leaves the existing buffer usable when an allocation fails. An execution context's allocation strategy is fixed when the context is created, so each engine captures the setting in effect at its own init and keeps it. A later `set_option` decides what the engines loaded after it are built with and changes nothing about the ones already running, so a `kSTATIC` context and a `kUSER_MANAGED` context coexist in one process. Why opt-in, not default-on: one buffer serves every context on a device, and a context holds its scratch for the whole enqueue -- which under a `CallerStreamGuard` can still be in flight when `execute()` returns -- so two enqueues must never hold it at once. The pool records each enqueue on a per-device event and makes the next one wait on it. An event, not the previous stream: synchronizing on a destroyed stream handle crashes rather than returning an error; CUDA recycles handle values, so two distinct streams can compare equal; and the NULL stream is a legal caller stream that no stream-handle sentinel can tell from "no previous user". Waiting from the stream that recorded the event is already satisfied, so the single-stream case pays a host call and no device stall. Not covered: concurrent same-device `execute()` on several threads, because the pool mutex is released before either enqueue is submitted. A default-on version needs the scratch keyed per stream instead of one buffer per device. This is orthogonal to weight streaming (#4336), which targets engine *weight* memory rather than activation scratch, and to export-time OOM. The grow/reuse/per-device policy and the handoff rule are factored into a header-only helper (`SharedScratchPool.h`) so they are unit-tested without a device (fake allocator, fake event factory). The CUDA path supplies the three callables it takes -- `cudaMalloc`, `cudaFree` and `cudaEventCreateWithFlags`; the `cudaStreamWaitEvent` and `cudaEventRecord` half of the handoff is the caller's, issued in response to what the helper returns. The helper needs the `cudaEvent_t` typedef, which is a compile-time dependency and not a runtime one; the repo had no headers-only CUDA target, so `third_party/cuda/BUILD` gains one rather than the helper depending on `cudart` and putting `libcudart` in the DT_NEEDED of a test that makes no CUDA call. `set_option` itself is not unit-tested, because no target in `tests/cpp/executorch/` links the backend. Three behaviours live only there: skipping a key this backend does not read, storing a valid boolean, and rejecting a non-boolean with `Error::InvalidArgument` instead of dropping it silently. The store is exercised by the measurement above, which reaches the pool through `set_option`; the key skip and the wrong-type rejection are covered nowhere, as `CudaBackend`'s equivalents also are. --- cpp/BUILD | 24 ++ .../executorch/SharedScratchPool.h | 123 ++++++++ .../executorch/TensorRTBackend.h | 11 + cpp/src/torch_tensorrt/executorch/README.md | 36 +++ .../executorch/TensorRTBackend.cpp | 200 +++++++++++- tests/cpp/executorch/BUILD | 10 + .../executorch/test_shared_scratch_pool.cpp | 286 ++++++++++++++++++ third_party/cuda/BUILD | 11 + 8 files changed, 699 insertions(+), 2 deletions(-) create mode 100644 cpp/include/torch_tensorrt/executorch/SharedScratchPool.h create mode 100644 tests/cpp/executorch/test_shared_scratch_pool.cpp diff --git a/cpp/BUILD b/cpp/BUILD index 30619cda92..5c23701060 100644 --- a/cpp/BUILD +++ b/cpp/BUILD @@ -191,6 +191,28 @@ cc_library( ], ) +cc_library( + name = "tensorrt_executorch_shared_scratch_pool", + hdrs = [ + "include/torch_tensorrt/executorch/SharedScratchPool.h", + ], + strip_include_prefix = "include", + target_compatible_with = select({ + ":linux_x86_64": [], + ":sbsa": [], + "//conditions:default": ["@platforms//:incompatible"], + }), + deps = select({ + ":linux_x86_64": [ + "@cuda//:cuda_headers", + ], + ":sbsa": [ + "@cuda//:cuda_headers", + ], + "//conditions:default": [], + }), +) + cc_library( name = "tensorrt_executorch_backend", srcs = [ @@ -211,6 +233,7 @@ cc_library( deps = [ ":tensorrt_executorch_binding_names", ":tensorrt_executorch_blob_header", + ":tensorrt_executorch_shared_scratch_pool", ":tensorrt_executorch_weight_streaming_budget", ] + select({ ":linux_x86_64": [ @@ -254,6 +277,7 @@ filegroup( filegroup( name = "executorch_api_headers", srcs = [ + "include/torch_tensorrt/executorch/SharedScratchPool.h", "include/torch_tensorrt/executorch/TensorRTBackend.h", "include/torch_tensorrt/executorch/TensorRTBindingNames.h", "include/torch_tensorrt/executorch/TensorRTBlobHeader.h", diff --git a/cpp/include/torch_tensorrt/executorch/SharedScratchPool.h b/cpp/include/torch_tensorrt/executorch/SharedScratchPool.h new file mode 100644 index 0000000000..62464e3aaf --- /dev/null +++ b/cpp/include/torch_tensorrt/executorch/SharedScratchPool.h @@ -0,0 +1,123 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +// Bookkeeping for the TensorRT backend's shared per-device activation-scratch +// pool: the grow/reuse/per-device policy and the enqueue-handoff rule. +// Allocation and event creation arrive as callables rather than being made here. + +#include + +#include +#include +#include + +namespace torch_tensorrt { +namespace executorch_backend { + +// Runtime backend option that backs execution-context activation scratch with a +// shared per-device pool instead of giving every context its own. Boolean, +// default false. Delivered as +// executorch::runtime::set_option("TensorRTBackend", options.view()) +// A context's allocation strategy is fixed when the context is created, so a +// later call governs only the engines loaded after it, and a pooled context and +// a private-scratch one coexist in one process. +inline constexpr char kSharedActivationScratchKey[] = "use_shared_activation_scratch"; + +// Per-device handoff marker for the shared scratch buffer: the pool-owned CUDA +// event that the last enqueue against the buffer was recorded on. +struct SharedScratchMarker { + cudaEvent_t event = nullptr; // never destroyed; one event serves the slot for the process lifetime + bool pending = false; // an enqueue against the buffer has been recorded on `event` +}; + +// What a caller about to enqueue against a device's shared scratch has to do: +// when `needs_wait`, make its stream wait on `event` first; once the enqueue is +// submitted, record it on `event`. `event` is null only when the slot has no +// event and one could not be created. +struct SharedScratchHandoff { + cudaEvent_t event = nullptr; + bool needs_wait = false; +}; + +// Claims a device's handoff for a caller about to enqueue against its shared +// scratch, creating the marker's event on first use. +// +// `create_event` returns a CUDA event, or nullptr if one could not be created, +// in which case the slot stays empty and the next call retries. +// +// The ordering between one enqueue and the next is carried by an event rather +// than by the stream the previous enqueue used, because a stream handle cannot +// carry it: synchronizing on a handle whose stream the caller has since +// destroyed is a crash rather than an error return, CUDA recycles handle values +// so a genuinely different stream can compare equal to the recorded one, and the +// NULL stream is both a legal stream a caller can select and the only available +// "no previous user" sentinel. An event names the work instead of the queue -- +// it stays valid after the stream that recorded it is destroyed, and waiting on +// it from the stream that recorded it is already satisfied, so the common +// single-stream case costs a host call and no device stall. +template +SharedScratchHandoff shared_scratch_claim_event( + std::unordered_map& markers, + int device_id, + CreateEvent create_event) { + SharedScratchMarker& marker = markers[device_id]; + if (marker.event == nullptr) { + marker.event = create_event(); + } + // A slot with no event is never marked, so a failed creation reports nothing to + // wait for rather than a wait the caller has no event to perform. + return {marker.event, marker.pending}; +} + +// The mark precedes the record, so a failed record leaves the slot claiming an +// enqueue the event does not cover -- the caller must then synchronize the stream +// itself before returning the error. +inline cudaEvent_t shared_scratch_mark_in_flight(std::unordered_map& markers, int device_id) { + SharedScratchMarker& marker = markers[device_id]; + if (marker.event != nullptr) { + marker.pending = true; + } + return marker.event; +} + +// Bookkeeping for a per-device pool of device-memory buffers that grows +// monotonically to the largest requested size. +// +// `alloc` returns nullptr on failure; the slot is then left untouched. +// Allocating before releasing is what makes that true, and it costs peak +// residency: while a slot grows, the old and the new buffer are both resident. +// `release` must leave no in-flight enqueue pointing at the buffer it frees -- +// the CUDA caller syncs the device first. +template +void* shared_scratch_get_or_grow( + std::unordered_map>& pool, + int device_id, + std::size_t need, + std::size_t& out_size, + Alloc alloc, + Release release) { + auto& slot = pool[device_id]; + if (slot.first != nullptr && slot.second >= need) { + out_size = slot.second; + return slot.first; + } + void* p = alloc(need); + if (p == nullptr) { + return nullptr; + } + if (slot.first != nullptr) { + release(slot.first); + } + slot = {p, need}; + out_size = need; + return p; +} + +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index b33d712d40..6713a950a2 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -75,6 +75,10 @@ struct EngineHandle { size_t num_aliased_outputs = 0; int device_id = 0; bool unified_memory = false; + // Whether exec_ctx was created kUSER_MANAGED and draws its activation scratch + // from the shared per-device pool (kSharedActivationScratchKey, + // SharedScratchPool.h). + bool shared_scratch = false; std::mutex mu; // Makes the skip-sync fast path safe to reuse: TensorRT forbids reconfiguring or // destroying an execution context while one of its enqueues is in flight, so when @@ -109,6 +113,13 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { ::executorch::runtime::DelegateHandle* handle, ::executorch::runtime::Span<::executorch::runtime::EValue*> args) const override; + // Applies the runtime backend options a caller passes to + // executorch::runtime::set_option("TensorRTBackend", ...). The only key read is + // kSharedActivationScratchKey (SharedScratchPool.h), a boolean. + ::executorch::runtime::Error set_option( + ET_UNUSED ::executorch::runtime::BackendOptionContext& context, + const ::executorch::runtime::Span<::executorch::runtime::BackendOption>& backend_options) override; + void destroy(::executorch::runtime::DelegateHandle* handle) const override; }; diff --git a/cpp/src/torch_tensorrt/executorch/README.md b/cpp/src/torch_tensorrt/executorch/README.md index e7367f8706..e65f5fa8fc 100644 --- a/cpp/src/torch_tensorrt/executorch/README.md +++ b/cpp/src/torch_tensorrt/executorch/README.md @@ -85,6 +85,14 @@ the removed `CudaStreamGuard`: complete, order any cross-stream producers/consumers with their own events, and synchronize the stream before reading outputs on the host. - With no guard active, the backend falls back to `cudaStreamPerThread`. +- With the `use_shared_activation_scratch` backend option enabled, one buffer per + device backs the activation scratch of every execution context created while it + was on, so an enqueue against that buffer must not overlap another one. The + backend orders consecutive enqueues itself, whether they run on one stream or + on two. What it cannot order is two `execute()` calls submitted concurrently on + one device: the caller must submit them one at a time, whether or not they + share a stream. Contexts created while the option was off keep their own + scratch and are unaffected. - The reference-runner smoke test runs inference inside a caller-stream guard on the discrete-GPU CI configuration, where all inputs and outputs are host-backed and therefore take the synchronized staging path. CI separately asserts that the @@ -104,6 +112,34 @@ the removed `CudaStreamGuard`: asynchronous return described above is still uncovered and the interaction between a green context and the internal completion event remains untested. +## Shared activation scratch + +A TensorRT execution context allocates its own activation scratch and holds it +for as long as the context lives, so a model lowered to N single-layer engines +pays N copies and can run out of device memory on the layer count alone. The +`use_shared_activation_scratch` backend option — a boolean, off by default — +instead backs every context on a device from one buffer, grown to the largest +engine's requirement: + +```cpp +#include + +executorch::runtime::BackendOptions<1> options; +options.set_option("use_shared_activation_scratch", true); +executorch::runtime::set_option("TensorRTBackend", options.view()); +``` + +Check what `executorch::runtime::set_option` returns: `Error::NotFound` means no +backend is registered under that name, which is what a binary that has not linked +the backend archive gets. + +N per-engine copies collapse to one, so the reclaimed memory is `(N-1)` times the +per-engine scratch. Set the option before loading the methods whose engines +should use the pool, and read the `use_shared_activation_scratch` bullet of the +caller-stream contract above first: the pool carries an ordering obligation the +backend cannot discharge for you. The buffer is never released, so the device +keeps the largest scratch it was ever asked for until the process exits. + ## Standalone Backend Archive Use this path only when you need `libexecutorch_trt_backend.a` without building diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index 8408c13e88..91b8d3e7a5 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -6,17 +6,21 @@ */ #include "torch_tensorrt/executorch/TensorRTBackend.h" +#include "torch_tensorrt/executorch/SharedScratchPool.h" #include "torch_tensorrt/executorch/TensorRTBindingNames.h" #include "torch_tensorrt/executorch/TensorRTBlobHeader.h" #include "torch_tensorrt/executorch/WeightStreamingBudget.h" +#include #include #include #include #include #include #include +#include #include +#include #include #include @@ -34,6 +38,8 @@ using ::executorch::aten::SizesType; using ::executorch::runtime::ArrayRef; using ::executorch::runtime::BackendExecutionContext; using ::executorch::runtime::BackendInitContext; +using ::executorch::runtime::BackendOption; +using ::executorch::runtime::BackendOptionContext; using ::executorch::runtime::CompileSpec; using ::executorch::runtime::DelegateHandle; using ::executorch::runtime::Error; @@ -151,6 +157,13 @@ bool infer_binding_names( return true; } +// The setting behind kSharedActivationScratchKey: whether an execution context +// created subsequently draws its activation scratch from the shared per-device +// pool rather than allocating its own. +// +// execute() must read EngineHandle::shared_scratch, never this. +std::atomic scratch_enabled{false}; + Error initialize_engine_io(EngineHandle& handle) { if (handle.input_binding_names.empty() && handle.output_binding_names.empty() && !infer_binding_names(handle.engine.get(), handle.input_binding_names, handle.output_binding_names)) { @@ -161,7 +174,13 @@ Error initialize_engine_io(EngineHandle& handle) { handle.num_inputs = handle.input_binding_names.size(); handle.num_outputs = handle.output_binding_names.size(); - handle.exec_ctx.reset(handle.engine->createExecutionContext()); + // kSTATIC gives the context its own activation scratch; kUSER_MANAGED makes it + // allocate none and take a buffer from execute() instead. The strategy is fixed + // at creation, so it is captured on the handle here rather than read per call. + handle.shared_scratch = scratch_enabled.load(std::memory_order_relaxed); + const auto strategy = handle.shared_scratch ? nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED + : nvinfer1::ExecutionContextAllocationStrategy::kSTATIC; + handle.exec_ctx.reset(handle.engine->createExecutionContext(strategy)); TORCHTRT_ET_CHECK_NOT_NULL( handle.exec_ctx, Error::InvalidProgram, "TensorRTBackend::init: failed to create TensorRT execution context"); @@ -204,6 +223,124 @@ bool is_cuda_accessible_ptr(const void* ptr) { return attrs.type == cudaMemoryTypeDevice || attrs.type == cudaMemoryTypeManaged; } +// Process-wide per-device pool for TensorRT execution-context activation scratch. +// One buffer sized to the largest engine's need serves every context on a device, +// instead of each of N layer-engines pinning its own scratch, which makes device +// memory scale with the layer count and OOMs multi-layer models. +// +// ORDERING: a context reads and writes its scratch for the whole enqueue, which +// can still be in flight when execute() returns, so two enqueues must never hold +// this buffer at the same time. +// +// What the pool's event handoff does NOT cover is concurrent execute() on one +// device: scratch_pool_mu guards the two maps only, and is released before either +// enqueue is submitted, so two threads can interleave their waits and records. +// The requirement is therefore that delegate enqueues on a device are submitted +// one at a time -- they need not share a stream, but they must not be submitted +// concurrently. That is why the pool is opt-in. +// +// The buffers and the events are intentionally never freed. Nothing here runs a +// CUDA call at process exit, which also keeps the pool clear of teardown-order +// hazards against anything else holding device memory. +std::mutex scratch_pool_mu; +std::unordered_map> scratch_pool; +std::unordered_map scratch_pool_markers; + +// Sets out_ptr to a buffer of at least `need` bytes on `device_id` and out_size to +// its capacity, with `stream` ordered after the enqueue that last used the buffer. +// The caller must call mark_shared_scratch_in_flight once it has submitted its own +// enqueue. +// +// Must be called with `device_id` already current: cudaEventCreateWithFlags, +// cudaMalloc, cudaFree and cudaDeviceSynchronize all act on the *current* device +// and nothing in here sets it. +Error get_or_grow_shared_scratch(int device_id, size_t need, cudaStream_t stream, void*& out_ptr, size_t& out_size) { + std::lock_guard lk(scratch_pool_mu); + + const SharedScratchHandoff handoff = shared_scratch_claim_event(scratch_pool_markers, device_id, []() -> cudaEvent_t { + cudaEvent_t event = nullptr; + if (cudaEventCreateWithFlags(&event, cudaEventDisableTiming) != cudaSuccess) { + return nullptr; + } + return event; + }); + if (handoff.event == nullptr) { + ET_LOG( + Error, + "TensorRTBackend::execute: failed to create the shared activation scratch handoff event on device %d", + device_id); + return Error::Internal; + } + if (handoff.needs_wait) { + const cudaError_t err = cudaStreamWaitEvent(stream, handoff.event, 0); + if (err != cudaSuccess) { + ET_LOG( + Error, + "TensorRTBackend::execute: waiting for the enqueue that last used the shared activation scratch failed: %s", + cudaGetErrorString(err)); + return Error::InvalidState; + } + } + + const auto slot = scratch_pool.find(device_id); + const bool first_buffer = slot == scratch_pool.end() || slot->second.first == nullptr; + void* const buffer = shared_scratch_get_or_grow( + scratch_pool, + device_id, + need, + out_size, + [device_id, first_buffer](size_t bytes) -> void* { + void* p = nullptr; + if (cudaMalloc(&p, bytes) != cudaSuccess) { + return nullptr; + } + ET_LOG( + Info, + "TensorRTBackend::execute: shared scratch pool (device %d) %s %zu bytes", + device_id, + first_buffer ? "allocated" : "grew to", + bytes); + return p; + }, + [](void* old) { + // Sync before free so no in-flight enqueue points at the old buffer. + cudaDeviceSynchronize(); + cudaFree(old); + }); + if (buffer == nullptr) { + ET_LOG( + Error, + "TensorRTBackend::execute: failed to allocate %zu bytes of shared activation scratch on device %d", + need, + device_id); + return Error::MemoryAllocationFailed; + } + + out_ptr = buffer; + return Error::Ok; +} + +// Records the enqueue now in flight on `stream` against `device_id`'s shared +// scratch, so the next call to get_or_grow_shared_scratch waits for it. +Error mark_shared_scratch_in_flight(int device_id, cudaStream_t stream) { + std::lock_guard lk(scratch_pool_mu); + + const cudaEvent_t event = shared_scratch_mark_in_flight(scratch_pool_markers, device_id); + if (event == nullptr) { + ET_LOG(Error, "TensorRTBackend::execute: shared activation scratch on device %d has no handoff event", device_id); + return Error::Internal; + } + const cudaError_t err = cudaEventRecord(event, stream); + if (err != cudaSuccess) { + ET_LOG( + Error, + "TensorRTBackend::execute: recording the completion event for the shared activation scratch enqueue failed: %s", + cudaGetErrorString(err)); + return Error::InvalidState; + } + return Error::Ok; +} + } // namespace // --------------------------------------------------------------------------- @@ -906,7 +1043,34 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } // ------------------------------------------------------------------ - // 4. Enqueue inference on the current CUDA stream + // 4. Back activation scratch with the shared per-device pool + // ------------------------------------------------------------------ + // All input shapes are bound by now, so the exact scratch requirement for this + // call is known. The buffer is installed on every call, not once, because a + // larger engine may have grown the pool and moved it since the last one. A + // kSTATIC context owns its private scratch, so setDeviceMemoryV2 must not be + // called on one. + bool scratch_from_pool = false; + if (engine->shared_scratch) { + const size_t need = ctx->updateDeviceMemorySizeForShapes(); + void* pool = nullptr; + size_t pool_size = 0; + // Zero means this call needs no scratch: nothing to claim, and nothing to + // order against the previous user of the buffer. Zero is also what a failed + // query returns; on this context's first call that is caught, because + // enqueueV3 refuses an engine it has never been given scratch for. + if (need > 0) { + const Error scratch_err = get_or_grow_shared_scratch(engine->device_id, need, stream, pool, pool_size); + if (scratch_err != Error::Ok) { + return scratch_err; + } + scratch_from_pool = true; + } + ctx->setDeviceMemoryV2(pool, static_cast(pool_size)); + } + + // ------------------------------------------------------------------ + // 5. Enqueue inference on the current CUDA stream // ------------------------------------------------------------------ if (!ctx->enqueueV3(stream)) { ET_LOG( @@ -918,6 +1082,18 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* return Error::InvalidState; } + // Pairs with get_or_grow_shared_scratch: the next claimant waits on this event. + if (scratch_from_pool) { + const Error mark_err = mark_shared_scratch_in_flight(engine->device_id, stream); + if (mark_err != Error::Ok) { + // Nothing will wait for this enqueue, so wait for it here instead of + // leaving the next user of the buffer to overwrite live scratch. + (void)cudaStreamSynchronize(stream); + engine->inflight_pending = false; + return mark_err; + } + } + // Caller-owned KV: reflect each engine in-place update into its delegate output // EValue (D2D on the same stream, after the engine work). for (const auto& r : aliased_reflects) { @@ -995,6 +1171,26 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* return Error::Ok; } +// --------------------------------------------------------------------------- +// set_option +// --------------------------------------------------------------------------- +Error TensorRTBackend::set_option(ET_UNUSED BackendOptionContext& context, const Span& backend_options) { + for (const auto& option : backend_options) { + // A caller may address one option span to several backends, so a key this + // backend does not read is skipped rather than refused. + if (std::strcmp(option.key, kSharedActivationScratchKey) == 0) { + if (const bool* const val = std::get_if(&option.value)) { + scratch_enabled.store(*val, std::memory_order_relaxed); + } else { + ET_LOG(Error, "TensorRTBackend::set_option: option '%s' must be a boolean", kSharedActivationScratchKey); + return Error::InvalidArgument; + } + } + } + + return Error::Ok; +} + // --------------------------------------------------------------------------- // destroy // diff --git a/tests/cpp/executorch/BUILD b/tests/cpp/executorch/BUILD index 17d2820bf2..d4f7a0c081 100644 --- a/tests/cpp/executorch/BUILD +++ b/tests/cpp/executorch/BUILD @@ -9,6 +9,7 @@ test_suite( ":test_executorch_binding_names", ":test_executorch_blob_header", ":test_executorch_weight_streaming_budget", + ":test_shared_scratch_pool", ], ) @@ -47,3 +48,12 @@ cc_test( "@googletest//:gtest_main", ], ) + +cc_test( + name = "test_shared_scratch_pool", + srcs = ["test_shared_scratch_pool.cpp"], + deps = [ + "//cpp:tensorrt_executorch_shared_scratch_pool", + "@googletest//:gtest_main", + ], +) diff --git a/tests/cpp/executorch/test_shared_scratch_pool.cpp b/tests/cpp/executorch/test_shared_scratch_pool.cpp new file mode 100644 index 0000000000..a65657e5df --- /dev/null +++ b/tests/cpp/executorch/test_shared_scratch_pool.cpp @@ -0,0 +1,286 @@ +#include "torch_tensorrt/executorch/SharedScratchPool.h" + +#include "gtest/gtest.h" + +#include +#include +#include +#include +#include +#include + +namespace torch_tensorrt { +namespace executorch_backend { +namespace { + +// Fake device allocator: hands out distinct non-null pointers and records every +// allocation size and every released pointer, so tests can assert the pool's +// grow/reuse/per-device policy without a CUDA device. +struct FakeAllocator { + std::vector alloc_sizes; + std::vector released; + std::uintptr_t next = 0x1000; + bool fail_next = false; + + void* alloc(std::size_t bytes) { + if (fail_next) { + fail_next = false; + return nullptr; + } + alloc_sizes.push_back(bytes); + void* p = reinterpret_cast(next); + next += 0x1000; + return p; + } + + void release(void* p) { + released.push_back(p); + } + + int alloc_count() const { + return static_cast(alloc_sizes.size()); + } +}; + +using Pool = std::unordered_map>; + +void* call(Pool& pool, FakeAllocator& a, int device_id, std::size_t need, std::size_t& out_size) { + return shared_scratch_get_or_grow( + pool, + device_id, + need, + out_size, + [&a](std::size_t bytes) { return a.alloc(bytes); }, + [&a](void* p) { a.release(p); }); +} + +TEST(SharedScratchPool, FirstRequestAllocatesExactSize) { + Pool pool; + FakeAllocator a; + std::size_t out = 0; + + void* p = call(pool, a, /*device_id=*/0, /*need=*/1024, out); + + EXPECT_NE(p, nullptr); + EXPECT_EQ(out, 1024u); + ASSERT_EQ(a.alloc_count(), 1); + EXPECT_EQ(a.alloc_sizes[0], 1024u); + EXPECT_TRUE(a.released.empty()); +} + +TEST(SharedScratchPool, ReusesWhenExistingBufferIsLargeEnough) { + Pool pool; + FakeAllocator a; + std::size_t out = 0; + + void* first = call(pool, a, 0, 4096, out); + // A smaller and an equal request must both reuse the same buffer (no realloc). + // The smaller one reports into a fresh out2, so what the reuse path writes is + // asserted rather than what the first call left in `out`. + std::size_t out2 = 0; + void* second = call(pool, a, 0, 1000, out2); + void* third = call(pool, a, 0, 4096, out); + + EXPECT_EQ(second, first); + EXPECT_EQ(third, first); + EXPECT_EQ(out, 4096u); + // Reuse reports the buffer's capacity, not the smaller amount asked for. + EXPECT_EQ(out2, 4096u); + EXPECT_EQ(a.alloc_count(), 1); + EXPECT_TRUE(a.released.empty()); +} + +TEST(SharedScratchPool, GrowsMonotonicallyToMaxAndReleasesOldBuffer) { + Pool pool; + FakeAllocator a; + std::size_t out = 0; + + void* small = call(pool, a, 0, 1024, out); + void* big = call(pool, a, 0, 8192, out); + + EXPECT_NE(big, small); + EXPECT_EQ(out, 8192u); + ASSERT_EQ(a.alloc_count(), 2); + EXPECT_EQ(a.alloc_sizes[1], 8192u); + ASSERT_EQ(a.released.size(), 1u); + EXPECT_EQ(a.released[0], small); + + // A subsequent smaller request reuses the grown buffer -- pool never shrinks. + void* reuse = call(pool, a, 0, 512, out); + EXPECT_EQ(reuse, big); + EXPECT_EQ(out, 8192u); + EXPECT_EQ(a.alloc_count(), 2); +} + +TEST(SharedScratchPool, KeepsIndependentBufferPerDevice) { + Pool pool; + FakeAllocator a; + std::size_t out = 0; + + void* dev0 = call(pool, a, /*device_id=*/0, 2048, out); + void* dev1 = call(pool, a, /*device_id=*/1, 2048, out); + + EXPECT_NE(dev0, dev1); + EXPECT_EQ(a.alloc_count(), 2); + EXPECT_TRUE(a.released.empty()); + + // Growing device 1 must not touch device 0's buffer. + void* dev1_big = call(pool, a, 1, 9000, out); + void* dev0_again = call(pool, a, 0, 2048, out); + EXPECT_NE(dev1_big, dev1); + EXPECT_EQ(dev0_again, dev0); + ASSERT_EQ(a.released.size(), 1u); + EXPECT_EQ(a.released[0], dev1); +} + +TEST(SharedScratchPool, AllocationFailureLeavesExistingSlotUntouched) { + Pool pool; + FakeAllocator a; + std::size_t out = 0; + + void* first = call(pool, a, 0, 1024, out); + ASSERT_NE(first, nullptr); + + // A growth whose allocation fails must return nullptr and keep the old buffer, + // so the caller can surface the error without corrupting the pool. + a.fail_next = true; + std::size_t out2 = 0; + void* failed = call(pool, a, 0, 8192, out2); + EXPECT_EQ(failed, nullptr); + EXPECT_TRUE(a.released.empty()); + + // The pool still holds the original buffer and serves it on the next request. + void* again = call(pool, a, 0, 1024, out); + EXPECT_EQ(again, first); + EXPECT_EQ(out, 1024u); +} + +TEST(SharedScratchPool, FirstAllocationFailureReturnsNullAndStoresNothing) { + Pool pool; + FakeAllocator a; + std::size_t out = 0; + + a.fail_next = true; + void* p = call(pool, a, 0, 1024, out); + EXPECT_EQ(p, nullptr); + + // Nothing stored: a later successful request allocates fresh. + void* q = call(pool, a, 0, 1024, out); + EXPECT_NE(q, nullptr); + EXPECT_EQ(a.alloc_count(), 1); +} + +// --------------------------------------------------------------------------- +// Ordering the shared buffer's handoff from one enqueue to the next. +// --------------------------------------------------------------------------- + +// Stands in for the CUDA event factory: hands out distinct non-null handles and +// counts calls, so a test can tell a slot that reuses its event from one that +// creates a new one every call. +struct FakeEventFactory { + int created = 0; + std::uintptr_t next = 0xE000; + bool fail_next = false; + + cudaEvent_t operator()() { + if (fail_next) { + fail_next = false; + return nullptr; + } + ++created; + cudaEvent_t e = reinterpret_cast(next); + next += 0x100; + return e; + } +}; + +using Markers = std::unordered_map; + +TEST(SharedScratchHandoffTest, FirstUseCreatesTheSlotsEventAndWaitsForNothing) { + Markers markers; + FakeEventFactory events; + + const SharedScratchHandoff handoff = shared_scratch_claim_event(markers, /*device_id=*/0, std::ref(events)); + + EXPECT_NE(handoff.event, nullptr); + EXPECT_FALSE(handoff.needs_wait); + EXPECT_EQ(events.created, 1); +} + +TEST(SharedScratchHandoffTest, EveryUseAfterAnEnqueueWaitsOnTheSameEvent) { + Markers markers; + FakeEventFactory events; + const SharedScratchHandoff first = shared_scratch_claim_event(markers, 0, std::ref(events)); + ASSERT_FALSE(first.needs_wait); + + EXPECT_EQ(shared_scratch_mark_in_flight(markers, 0), first.event); + + // Every later enqueue waits, however many there have been and whichever stream + // each of them ran on: the marker records that the buffer was handed out, not + // who it was handed to. Comparing stream handles instead would let a caller + // through whenever its handle matched the recorded one, including when CUDA has + // recycled that value for a different stream. + const SharedScratchHandoff second = shared_scratch_claim_event(markers, 0, std::ref(events)); + EXPECT_TRUE(second.needs_wait); + EXPECT_EQ(second.event, first.event); + + const SharedScratchHandoff third = shared_scratch_claim_event(markers, 0, std::ref(events)); + EXPECT_TRUE(third.needs_wait); + EXPECT_EQ(third.event, first.event); + + // One event serves the slot for its whole life, so the wait never targets an + // event some earlier enqueue was recorded on. + EXPECT_EQ(events.created, 1); +} + +TEST(SharedScratchHandoffTest, KeepsAnIndependentMarkerPerDevice) { + Markers markers; + FakeEventFactory events; + const SharedScratchHandoff dev0 = shared_scratch_claim_event(markers, /*device_id=*/0, std::ref(events)); + ASSERT_EQ(shared_scratch_mark_in_flight(markers, 0), dev0.event); + + // Device 1 has its own buffer, so device 0's enqueue is nothing for it to wait + // on, and it gets its own event. + const SharedScratchHandoff dev1 = shared_scratch_claim_event(markers, /*device_id=*/1, std::ref(events)); + EXPECT_FALSE(dev1.needs_wait); + EXPECT_NE(dev1.event, dev0.event); + EXPECT_EQ(events.created, 2); + + // Marking device 1 does not make device 0 stop waiting, or the other way round. + ASSERT_EQ(shared_scratch_mark_in_flight(markers, 1), dev1.event); + EXPECT_TRUE(shared_scratch_claim_event(markers, 0, std::ref(events)).needs_wait); + EXPECT_TRUE(shared_scratch_claim_event(markers, 1, std::ref(events)).needs_wait); +} + +TEST(SharedScratchHandoffTest, EventCreationFailureIsReportedAndRetried) { + Markers markers; + FakeEventFactory events; + + events.fail_next = true; + const SharedScratchHandoff failed = shared_scratch_claim_event(markers, 0, std::ref(events)); + EXPECT_EQ(failed.event, nullptr); + EXPECT_FALSE(failed.needs_wait); + + // The failure leaves nothing behind, so the next call tries again and succeeds + // rather than serving an unusable slot for the rest of the process. + const SharedScratchHandoff retried = shared_scratch_claim_event(markers, 0, std::ref(events)); + EXPECT_NE(retried.event, nullptr); + EXPECT_FALSE(retried.needs_wait); + EXPECT_EQ(events.created, 1); +} + +TEST(SharedScratchHandoffTest, ASlotWithNoEventIsNotMarked) { + Markers markers; + FakeEventFactory events; + + // Nothing can be recorded without an event, so nothing is claimed to have been. + EXPECT_EQ(shared_scratch_mark_in_flight(markers, 0), nullptr); + + // Otherwise, once an event is finally created for the slot, the next caller + // would wait on it believing an enqueue had been recorded on it that never was. + EXPECT_FALSE(shared_scratch_claim_event(markers, 0, std::ref(events)).needs_wait); +} + +} // namespace +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/third_party/cuda/BUILD b/third_party/cuda/BUILD index 204b9cee23..ed98f4f3c6 100644 --- a/third_party/cuda/BUILD +++ b/third_party/cuda/BUILD @@ -17,6 +17,17 @@ config_setting( ], ) +cc_library( + name = "cuda_headers", + hdrs = glob([ + "include/**/*.h", + "include/**/*.hpp", + "include/**/*.inl", + "include/**/*", + ]), + includes = ["include/"], +) + cc_library( name = "cudart", srcs = select({ From 154a5291c32af77016537d52cc90127e8ecbe46c Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Thu, 27 Aug 2026 14:27:55 -0700 Subject: [PATCH 2/3] fix(executorch): address review of the shared activation-scratch pool Four changes, all against the review of the parent commit. A zero from `updateDeviceMemorySizeForShapes()` is ambiguous: TensorRT answers a failed query and an engine that genuinely needs no activation scratch the same way. The parent treated it as "no scratch needed" and carried on, but `setDeviceMemoryV2(nullptr, 0)` is itself rejected and returns nothing to test, so the context kept whatever buffer it last held -- which a pool growth may already have freed. Forcing that path against a freed buffer reproduces an illegal memory access: the enqueue is submitted, and the fault surfaces at the stream synchronize that follows it. Each engine now records what `ICudaEngine::getDeviceMemorySizeV2()` reports at its own init, and a zero fails the `execute()` only when that recorded requirement is non-zero. An engine that needs no scratch is given no buffer, so it has nothing to claim and nothing for the next claimant to order against. The pool's state and its lock are now per device. The parent held one process-wide mutex across `cudaMalloc` and across a `cudaDeviceSynchronize` in the release path, so a growth on one device blocked a claim on another, and the sync waited on everything queued rather than on the scratch users. A registry lock now finds a device's entry and is held for the lookup alone, never across a CUDA call; the device's own lock covers the claim, the allocation and the wait before a free; and that wait is on the handoff event rather than the device. Entries are never erased and `std::unordered_map` keeps references valid across rehashing, which is what lets the registry lock be dropped before the entry is used. The helper's unit tests drive fakes and reach none of the delegate wiring, so the single-threaded pooled path had no automated coverage at all: the execution context could be reverted to `kSTATIC`, or the `updateDeviceMemorySizeForShapes`/`setDeviceMemoryV2` pair deleted, with every test still green. `tests/cpp/executorch/test_shared_scratch_backend.cpp` links the delegate and covers `set_option`'s three behaviours, the per-engine capture of the setting and of the engine's own scratch requirement, the pooled `execute()` path, the four-contexts-one-allocation claim, the event handoff across two caller streams, and an engine that needs no activation scratch at all. It builds its own TensorRT engine rather than loading a `.pte`. It needs a CUDA device: without one it skips every test and exits zero, and the workflow's `--test_output=errors` keeps the skip reason out of the log, so a green run on a device-less host says nothing about the pool. Finally, four of the parent's statements no longer hold. Its description of what `updateDeviceMemorySizeForShapes()` returns was wrong: it does not report the requirement for the shapes just bound. Whether an engine does that or reports its profile maximum is fixed when the engine is built -- the builder's `PreviewFeature::kRUNTIME_ACTIVATION_RESIZE_10_10` produces the former, and without it either can happen depending on how TensorRT planned the engine -- so the pool can settle well above the live data and a runtime cannot tighten it. The reclaimed memory is likewise the sum of the N per-engine requirements less the largest of them, not `(N-1)` times a uniform per-engine figure. "One buffer serves every context on a device" is true only of contexts created while the option is on; one created while it was off keeps its own scratch. And `set_option` is untested there because no target in `tests/cpp/executorch/` links the backend; `test_shared_scratch_backend.cpp` is such a target and covers all three of its behaviours. --- .../executorch/SharedScratchPool.h | 101 ++- .../executorch/TensorRTBackend.h | 9 + cpp/src/torch_tensorrt/executorch/README.md | 37 +- .../executorch/TensorRTBackend.cpp | 96 +- tests/cpp/BUILD | 2 + tests/cpp/executorch/BUILD | 33 + .../test_shared_scratch_backend.cpp | 851 ++++++++++++++++++ .../executorch/test_shared_scratch_pool.cpp | 375 ++++++-- 8 files changed, 1334 insertions(+), 170 deletions(-) create mode 100644 tests/cpp/executorch/test_shared_scratch_backend.cpp diff --git a/cpp/include/torch_tensorrt/executorch/SharedScratchPool.h b/cpp/include/torch_tensorrt/executorch/SharedScratchPool.h index 62464e3aaf..5f43fdf3ce 100644 --- a/cpp/include/torch_tensorrt/executorch/SharedScratchPool.h +++ b/cpp/include/torch_tensorrt/executorch/SharedScratchPool.h @@ -8,14 +8,15 @@ #pragma once // Bookkeeping for the TensorRT backend's shared per-device activation-scratch -// pool: the grow/reuse/per-device policy and the enqueue-handoff rule. +// pool: the grow/reuse policy, the enqueue-handoff rule, and the lock that scopes +// both to a single device. // Allocation and event creation arrive as callables rather than being made here. #include #include +#include #include -#include namespace torch_tensorrt { namespace executorch_backend { @@ -45,8 +46,42 @@ struct SharedScratchHandoff { bool needs_wait = false; }; +// One device's shared scratch buffer and the marker ordering its handoff, behind +// the lock that covers both. +// +// A claimant holds `mu` from the wait on the previous enqueue through the choice +// of buffer, so it cannot be handed a buffer another claimant is midway through +// replacing, and cannot record its own enqueue against a marker that has since +// moved on. `mu` covers one device, so a growth holds no lock a claim on another +// device has to acquire. +struct SharedScratchDevice { + std::mutex mu; + void* buffer = nullptr; + std::size_t capacity = 0; + SharedScratchMarker marker; +}; + +// Holds one SharedScratchDevice per device id. +// +// `get` locks only long enough to find or create the entry, and the reference it +// returns stays usable once that lock is dropped: std::unordered_map keeps +// references to elements valid across rehashing, and entries are never erased. +// This one lock is shared by every device, which is why nothing but the lookup +// runs under it. +class SharedScratchPool { + public: + SharedScratchDevice& get(int device_id) { + std::lock_guard lk(mu_); + return devices_[device_id]; + } + + private: + std::mutex mu_; + std::unordered_map devices_; +}; + // Claims a device's handoff for a caller about to enqueue against its shared -// scratch, creating the marker's event on first use. +// scratch, creating the marker's event on first use. Call with `dev.mu` held. // // `create_event` returns a CUDA event, or nullptr if one could not be created, // in which case the slot stays empty and the next call retries. @@ -62,59 +97,63 @@ struct SharedScratchHandoff { // it from the stream that recorded it is already satisfied, so the common // single-stream case costs a host call and no device stall. template -SharedScratchHandoff shared_scratch_claim_event( - std::unordered_map& markers, - int device_id, - CreateEvent create_event) { - SharedScratchMarker& marker = markers[device_id]; - if (marker.event == nullptr) { - marker.event = create_event(); +SharedScratchHandoff shared_scratch_claim_event(SharedScratchDevice& dev, CreateEvent create_event) { + if (dev.marker.event == nullptr) { + dev.marker.event = create_event(); } // A slot with no event is never marked, so a failed creation reports nothing to // wait for rather than a wait the caller has no event to perform. - return {marker.event, marker.pending}; + return {dev.marker.event, dev.marker.pending}; } +// Call with `dev.mu` held. +// // The mark precedes the record, so a failed record leaves the slot claiming an // enqueue the event does not cover -- the caller must then synchronize the stream // itself before returning the error. -inline cudaEvent_t shared_scratch_mark_in_flight(std::unordered_map& markers, int device_id) { - SharedScratchMarker& marker = markers[device_id]; - if (marker.event != nullptr) { - marker.pending = true; +inline cudaEvent_t shared_scratch_mark_in_flight(SharedScratchDevice& dev) { + if (dev.marker.event != nullptr) { + dev.marker.pending = true; } - return marker.event; + return dev.marker.event; } -// Bookkeeping for a per-device pool of device-memory buffers that grows -// monotonically to the largest requested size. +// Bookkeeping for a device's scratch buffer, which grows monotonically to the +// largest requested size. Call with `dev.mu` held. // -// `alloc` returns nullptr on failure; the slot is then left untouched. +// `alloc` returns nullptr on failure; the buffer is then left untouched. // Allocating before releasing is what makes that true, and it costs peak -// residency: while a slot grows, the old and the new buffer are both resident. -// `release` must leave no in-flight enqueue pointing at the buffer it frees -- -// the CUDA caller syncs the device first. +// residency: while the buffer grows, the old and the new one are both resident. +// +// `release(old, wait_for)` frees `old`. A non-null `wait_for` is the marker's +// event, on which an enqueue that may still be reading and writing `old` has been +// recorded; the release must wait for that event on the host before freeing. One +// event covers every enqueue the buffer ever served, but only because each of +// them claims the handoff before enqueueing -- which orders its stream after the +// event -- and records on the event afterwards, so the latest recording completes +// only once all the earlier ones have. An enqueue that reaches the buffer without +// doing both is covered by no wait here. A null `wait_for` means nothing was ever +// recorded against this buffer, so there is nothing to wait for. template void* shared_scratch_get_or_grow( - std::unordered_map>& pool, - int device_id, + SharedScratchDevice& dev, std::size_t need, std::size_t& out_size, Alloc alloc, Release release) { - auto& slot = pool[device_id]; - if (slot.first != nullptr && slot.second >= need) { - out_size = slot.second; - return slot.first; + if (dev.buffer != nullptr && dev.capacity >= need) { + out_size = dev.capacity; + return dev.buffer; } void* p = alloc(need); if (p == nullptr) { return nullptr; } - if (slot.first != nullptr) { - release(slot.first); + if (dev.buffer != nullptr) { + release(dev.buffer, dev.marker.pending ? dev.marker.event : nullptr); } - slot = {p, need}; + dev.buffer = p; + dev.capacity = need; out_size = need; return p; } diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index 6713a950a2..e164bd01f9 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -79,6 +79,10 @@ struct EngineHandle { // from the shared per-device pool (kSharedActivationScratchKey, // SharedScratchPool.h). bool shared_scratch = false; + // The activation scratch the engine itself reports needing, read at init when + // shared_scratch is set. execute() needs it to tell a failed per-call query, + // which TensorRT also reports as zero, from an engine that genuinely needs none. + size_t engine_scratch_bytes = 0; std::mutex mu; // Makes the skip-sync fast path safe to reuse: TensorRT forbids reconfiguring or // destroying an execution context while one of its enqueues is in flight, so when @@ -106,6 +110,11 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { // past return, order any other stream against this one, and synchronize the stream // before reading device-resident outputs. The selected stream must be on the engine's // device, and calls on one handle must not overlap each other or its destruction. + // The shared activation scratch pool (kSharedActivationScratchKey) widens that + // across handles: one buffer per device backs every context created while the + // option was on, so calls on two such handles on one device must not overlap + // either. A handle whose context was created while the option was off keeps its + // own scratch and is outside that rule. // Note that other CUDA delegates sharing the same guard may instead synchronize before // returning, so do not assume results are ready on return from this one. ::executorch::runtime::Error execute( diff --git a/cpp/src/torch_tensorrt/executorch/README.md b/cpp/src/torch_tensorrt/executorch/README.md index e65f5fa8fc..b5bd5cc526 100644 --- a/cpp/src/torch_tensorrt/executorch/README.md +++ b/cpp/src/torch_tensorrt/executorch/README.md @@ -85,14 +85,16 @@ the removed `CudaStreamGuard`: complete, order any cross-stream producers/consumers with their own events, and synchronize the stream before reading outputs on the host. - With no guard active, the backend falls back to `cudaStreamPerThread`. -- With the `use_shared_activation_scratch` backend option enabled, one buffer per - device backs the activation scratch of every execution context created while it - was on, so an enqueue against that buffer must not overlap another one. The - backend orders consecutive enqueues itself, whether they run on one stream or - on two. What it cannot order is two `execute()` calls submitted concurrently on - one device: the caller must submit them one at a time, whether or not they - share a stream. Contexts created while the option was off keep their own - scratch and are unaffected. +- With the `use_shared_activation_scratch` backend option enabled, one buffer + per device backs the activation scratch of every execution context created + while it was on, so an enqueue against that buffer must not overlap another + one. The backend orders consecutive enqueues itself, whether they run on one + stream or on two. What it cannot order is two `execute()` calls submitted + concurrently on one device: the caller must submit them one at a time, whether + or not they share a stream. Submitting them concurrently risks one of them + growing the pool and freeing the buffer the other's enqueue is still reading + and writing, not merely reordering them. Contexts created while the option was + off keep their own scratch and are unaffected. - The reference-runner smoke test runs inference inside a caller-stream guard on the discrete-GPU CI configuration, where all inputs and outputs are host-backed and therefore take the synchronized staging path. CI separately asserts that the @@ -133,12 +135,19 @@ Check what `executorch::runtime::set_option` returns: `Error::NotFound` means no backend is registered under that name, which is what a binary that has not linked the backend archive gets. -N per-engine copies collapse to one, so the reclaimed memory is `(N-1)` times the -per-engine scratch. Set the option before loading the methods whose engines -should use the pool, and read the `use_shared_activation_scratch` bullet of the -caller-stream contract above first: the pool carries an ordering obligation the -backend cannot discharge for you. The buffer is never released, so the device -keeps the largest scratch it was ever asked for until the process exits. +N per-engine copies collapse to one, so the reclaimed memory is the sum of the N +requirements less the largest of them. Set the option before loading the methods +whose engines should use the pool, and read the `use_shared_activation_scratch` +bullet of the caller-stream contract above first: the pool carries an ordering +obligation the backend cannot discharge for you. The buffer is never released, so +the device keeps the largest scratch it was ever asked for until the process +exits. + +How much any one engine asks for is fixed when it is built, not when it runs. +The builder's `kRUNTIME_ACTIVATION_RESIZE_10_10` preview feature makes an engine +report what the shapes just bound need; without it, whether an engine does that +or reports its profile maximum depends on how TensorRT planned it. Either way the +pool can settle well above the live data, and nothing the runtime does changes it. ## Standalone Backend Archive diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index 91b8d3e7a5..a362ea7efc 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -18,7 +18,6 @@ #include #include #include -#include #include #include #include @@ -184,6 +183,13 @@ Error initialize_engine_io(EngineHandle& handle) { TORCHTRT_ET_CHECK_NOT_NULL( handle.exec_ctx, Error::InvalidProgram, "TensorRTBackend::init: failed to create TensorRT execution context"); + if (handle.shared_scratch) { + // Read after the weight streaming budget is applied, which the caller does + // before this runs because TensorRT forbids moving the budget once a context + // exists -- and the budget is the one thing that moves this figure. + handle.engine_scratch_bytes = static_cast(handle.engine->getDeviceMemorySizeV2()); + } + return Error::Ok; } @@ -224,27 +230,29 @@ bool is_cuda_accessible_ptr(const void* ptr) { } // Process-wide per-device pool for TensorRT execution-context activation scratch. -// One buffer sized to the largest engine's need serves every context on a device, -// instead of each of N layer-engines pinning its own scratch, which makes device -// memory scale with the layer count and OOMs multi-layer models. +// One buffer sized to the largest engine's need serves every kUSER_MANAGED context +// on a device, instead of each of N layer-engines pinning its own scratch, which +// makes device memory scale with the layer count and OOMs multi-layer models. // // ORDERING: a context reads and writes its scratch for the whole enqueue, which // can still be in flight when execute() returns, so two enqueues must never hold // this buffer at the same time. // // What the pool's event handoff does NOT cover is concurrent execute() on one -// device: scratch_pool_mu guards the two maps only, and is released before either -// enqueue is submitted, so two threads can interleave their waits and records. -// The requirement is therefore that delegate enqueues on a device are submitted -// one at a time -- they need not share a stream, but they must not be submitted -// concurrently. That is why the pool is opt-in. +// device: a device's lock is released before the enqueue is submitted, so an +// enqueue is live for a window before the event carries it, and a second thread +// claiming inside that window is told to wait for the enqueue before it. Such a +// claimant can grow the pool and free the buffer the first thread's enqueue is +// still reading and writing. The requirement is therefore that the enqueues +// drawing on a device's buffer are submitted one at a time, but they need not +// share a stream. That is why the pool is opt-in. The pool's locking does not +// couple two devices: each carries its own lock, and no CUDA call is made under +// the lock that finds it. // // The buffers and the events are intentionally never freed. Nothing here runs a // CUDA call at process exit, which also keeps the pool clear of teardown-order // hazards against anything else holding device memory. -std::mutex scratch_pool_mu; -std::unordered_map> scratch_pool; -std::unordered_map scratch_pool_markers; +SharedScratchPool scratch_pool; // Sets out_ptr to a buffer of at least `need` bytes on `device_id` and out_size to // its capacity, with `stream` ordered after the enqueue that last used the buffer. @@ -252,12 +260,13 @@ std::unordered_map scratch_pool_markers; // enqueue. // // Must be called with `device_id` already current: cudaEventCreateWithFlags, -// cudaMalloc, cudaFree and cudaDeviceSynchronize all act on the *current* device -// and nothing in here sets it. +// cudaMalloc and cudaFree all act on the *current* device and nothing in here +// sets it. Error get_or_grow_shared_scratch(int device_id, size_t need, cudaStream_t stream, void*& out_ptr, size_t& out_size) { - std::lock_guard lk(scratch_pool_mu); + SharedScratchDevice& dev = scratch_pool.get(device_id); + std::lock_guard lk(dev.mu); - const SharedScratchHandoff handoff = shared_scratch_claim_event(scratch_pool_markers, device_id, []() -> cudaEvent_t { + const SharedScratchHandoff handoff = shared_scratch_claim_event(dev, []() -> cudaEvent_t { cudaEvent_t event = nullptr; if (cudaEventCreateWithFlags(&event, cudaEventDisableTiming) != cudaSuccess) { return nullptr; @@ -282,11 +291,9 @@ Error get_or_grow_shared_scratch(int device_id, size_t need, cudaStream_t stream } } - const auto slot = scratch_pool.find(device_id); - const bool first_buffer = slot == scratch_pool.end() || slot->second.first == nullptr; + const bool first_buffer = dev.buffer == nullptr; void* const buffer = shared_scratch_get_or_grow( - scratch_pool, - device_id, + dev, need, out_size, [device_id, first_buffer](size_t bytes) -> void* { @@ -302,9 +309,10 @@ Error get_or_grow_shared_scratch(int device_id, size_t need, cudaStream_t stream bytes); return p; }, - [](void* old) { - // Sync before free so no in-flight enqueue points at the old buffer. - cudaDeviceSynchronize(); + [](void* old, cudaEvent_t wait_for) { + if (wait_for != nullptr) { + cudaEventSynchronize(wait_for); + } cudaFree(old); }); if (buffer == nullptr) { @@ -323,9 +331,10 @@ Error get_or_grow_shared_scratch(int device_id, size_t need, cudaStream_t stream // Records the enqueue now in flight on `stream` against `device_id`'s shared // scratch, so the next call to get_or_grow_shared_scratch waits for it. Error mark_shared_scratch_in_flight(int device_id, cudaStream_t stream) { - std::lock_guard lk(scratch_pool_mu); + SharedScratchDevice& dev = scratch_pool.get(device_id); + std::lock_guard lk(dev.mu); - const cudaEvent_t event = shared_scratch_mark_in_flight(scratch_pool_markers, device_id); + const cudaEvent_t event = shared_scratch_mark_in_flight(dev); if (event == nullptr) { ET_LOG(Error, "TensorRTBackend::execute: shared activation scratch on device %d has no handoff event", device_id); return Error::Internal; @@ -1045,28 +1054,41 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // ------------------------------------------------------------------ // 4. Back activation scratch with the shared per-device pool // ------------------------------------------------------------------ - // All input shapes are bound by now, so the exact scratch requirement for this - // call is known. The buffer is installed on every call, not once, because a - // larger engine may have grown the pool and moved it since the last one. A - // kSTATIC context owns its private scratch, so setDeviceMemoryV2 must not be - // called on one. + // The query requires every input shape to be bound, which they are by here. + // Whatever it answers is binding rather than advisory: setDeviceMemoryV2 refuses + // a smaller buffer, and an engine backed by less than it asked for writes past + // the end. + // + // The buffer is installed on every call, not once, because a larger engine may + // have grown the pool and moved it since the last one. A kSTATIC context owns + // its private scratch, so setDeviceMemoryV2 must not be called on one. + // + // A reported zero is ambiguous: TensorRT answers a failed query and an engine + // that genuinely needs no scratch the same way, and the engine's own + // requirement is what separates them. An engine that needs none is given no + // buffer, so it has nothing to claim and nothing for the next claimant to order + // against. A failed query carried on would instead leave the context enqueueing + // against whatever buffer it last held, because setDeviceMemoryV2(nullptr, 0) + // is rejected and returns nothing to test. bool scratch_from_pool = false; if (engine->shared_scratch) { const size_t need = ctx->updateDeviceMemorySizeForShapes(); - void* pool = nullptr; - size_t pool_size = 0; - // Zero means this call needs no scratch: nothing to claim, and nothing to - // order against the previous user of the buffer. Zero is also what a failed - // query returns; on this context's first call that is caught, because - // enqueueV3 refuses an engine it has never been given scratch for. if (need > 0) { + void* pool = nullptr; + size_t pool_size = 0; const Error scratch_err = get_or_grow_shared_scratch(engine->device_id, need, stream, pool, pool_size); if (scratch_err != Error::Ok) { return scratch_err; } scratch_from_pool = true; + ctx->setDeviceMemoryV2(pool, static_cast(pool_size)); + } else if (engine->engine_scratch_bytes > 0) { + ET_LOG( + Error, + "TensorRTBackend::execute: updateDeviceMemorySizeForShapes returned 0, but the engine needs %zu bytes of activation scratch", + engine->engine_scratch_bytes); + return Error::InvalidState; } - ctx->setDeviceMemoryV2(pool, static_cast(pool_size)); } // ------------------------------------------------------------------ diff --git a/tests/cpp/BUILD b/tests/cpp/BUILD index b5c0c15138..827fa2c409 100644 --- a/tests/cpp/BUILD +++ b/tests/cpp/BUILD @@ -69,6 +69,8 @@ test_suite( "//tests/cpp/executorch:test_executorch_binding_names", "//tests/cpp/executorch:test_executorch_blob_header", "//tests/cpp/executorch:test_executorch_weight_streaming_budget", + "//tests/cpp/executorch:test_shared_scratch_backend", + "//tests/cpp/executorch:test_shared_scratch_pool", ], ) diff --git a/tests/cpp/executorch/BUILD b/tests/cpp/executorch/BUILD index d4f7a0c081..224e270295 100644 --- a/tests/cpp/executorch/BUILD +++ b/tests/cpp/executorch/BUILD @@ -9,6 +9,7 @@ test_suite( ":test_executorch_binding_names", ":test_executorch_blob_header", ":test_executorch_weight_streaming_budget", + ":test_shared_scratch_backend", ":test_shared_scratch_pool", ], ) @@ -57,3 +58,35 @@ cc_test( "@googletest//:gtest_main", ], ) + +# exclusive because the memory comparison reads device-wide free memory, which +# any other GPU target running at the same time would move. +cc_test( + name = "test_shared_scratch_backend", + timeout = "long", + srcs = ["test_shared_scratch_backend.cpp"], + tags = ["exclusive"], + target_compatible_with = select({ + "//cpp:linux_x86_64": [], + "//cpp:sbsa": [], + "//conditions:default": ["@platforms//:incompatible"], + }), + deps = [ + "//cpp:tensorrt_executorch_backend", + "//cpp:tensorrt_executorch_blob_header", + "@executorch//:executorch_core", + "@executorch//:executorch_headers", + "@executorch//:extension_cuda", + "@googletest//:gtest_main", + ] + select({ + "//cpp:linux_x86_64": [ + "@cuda//:cudart", + "@tensorrt//:nvinfer", + ], + "//cpp:sbsa": [ + "@cuda//:cudart", + "@tensorrt_sbsa//:nvinfer", + ], + "//conditions:default": [], + }), +) diff --git a/tests/cpp/executorch/test_shared_scratch_backend.cpp b/tests/cpp/executorch/test_shared_scratch_backend.cpp new file mode 100644 index 0000000000..54842b8db0 --- /dev/null +++ b/tests/cpp/executorch/test_shared_scratch_backend.cpp @@ -0,0 +1,851 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Exercises the shared activation-scratch pool through the delegate that uses +// it: the runtime option that turns it on, the per-engine capture of that +// option, and the single-threaded pooled execute() path -- the kUSER_MANAGED +// context, the updateDeviceMemorySizeForShapes/setDeviceMemoryV2 pair, and the +// enqueue handoff between two caller streams. +// +// The TensorRT engine is built here rather than loaded from a .pte so the target +// carries no exported artifact, at the cost of a few seconds of builder time. +// +// COVERAGE LIMIT: every test below needs a CUDA device and a TensorRT that can +// build an engine. Without one the whole suite skips and covers nothing, so a +// green run on a host with no GPU says nothing about the pool. + +#include "torch_tensorrt/executorch/TensorRTBackend.h" +#include "torch_tensorrt/executorch/TensorRTBlobHeader.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace torch_tensorrt { +namespace executorch_backend { +namespace { + +using ::executorch::aten::ScalarType; +using ::executorch::aten::SizesType; +using ::executorch::runtime::ArrayRef; +using ::executorch::runtime::BackendExecutionContext; +using ::executorch::runtime::BackendInitContext; +using ::executorch::runtime::BackendOption; +using ::executorch::runtime::BackendOptionContext; +using ::executorch::runtime::CompileSpec; +using ::executorch::runtime::DelegateHandle; +using ::executorch::runtime::Error; +using ::executorch::runtime::EValue; +using ::executorch::runtime::FreeableBuffer; +using ::executorch::runtime::MemoryAllocator; +using ::executorch::runtime::Span; + +// Spelled out rather than taken from SharedScratchPool.h: a test that reads the +// key through the production constant cannot pin the key's value. +constexpr char kOptionKey[] = "use_shared_activation_scratch"; + +constexpr int kRows = 2048; +constexpr int kCols = 2048; +constexpr std::size_t kElems = static_cast(kRows) * static_cast(kCols); +constexpr std::size_t kBytes = kElems * sizeof(float); + +// Engines loaded together in the memory test. Four is enough for the private +// case to cost 4x the scratch and the pooled case 1x. +constexpr int kEngineCount = 4; + +// A value neither network below can produce, so an output comparison cannot be +// satisfied by an execute() that never reached the engine. +constexpr float kSentinel = -7.0f; + +// Below this the memory comparison cannot see past allocator granularity, so the +// test reports that its network stopped producing measurable scratch instead of +// passing on a difference it cannot resolve. +constexpr std::size_t kMinMeasurableScratch = 4u << 20; + +// --------------------------------------------------------------------------- +// A TensorRT engine, built here, wrapped in the delegate's blob wire format +// --------------------------------------------------------------------------- + +constexpr char kMagic[4] = {'T', 'R', '0', '1'}; +constexpr std::uint32_t kMetadataOffsetField = 4; +constexpr std::uint32_t kMetadataSizeField = 8; +constexpr std::uint32_t kEngineOffsetField = 12; +constexpr std::uint32_t kEngineSizeField = 16; +constexpr std::uint32_t kHeaderSize = 32; +constexpr std::uint32_t kEngineAlignment = 16; + +class BuilderLogger : public nvinfer1::ILogger { + public: + void log(Severity severity, const char* msg) noexcept override { + if (severity <= Severity::kWARNING) { + std::fprintf(stderr, "[TensorRT] %s\n", msg); + } + } +}; + +template +void write_field(std::vector& blob, std::size_t offset, T value) { + std::memcpy(blob.data() + offset, &value, sizeof(value)); +} + +std::size_t align_up(std::size_t value, std::size_t alignment) { + return ((value + alignment - 1) / alignment) * alignment; +} + +// Two softmaxes over different axes sit between the pointwise layers so the +// chain cannot collapse into a single pass, which is what keeps the engine's +// activation requirement large enough for the memory comparison to resolve. +bool add_scratch_needing_net(nvinfer1::INetworkDefinition& network, nvinfer1::ITensor& input) { + static const float kAddend = 0.125f; + static const float kScale = 1.5f; + + nvinfer1::IConstantLayer* addend = + network.addConstant(nvinfer1::Dims3{1, 1, 1}, nvinfer1::Weights{nvinfer1::DataType::kFLOAT, &kAddend, 1}); + nvinfer1::IConstantLayer* scale = + network.addConstant(nvinfer1::Dims3{1, 1, 1}, nvinfer1::Weights{nvinfer1::DataType::kFLOAT, &kScale, 1}); + if (addend == nullptr || scale == nullptr) { + return false; + } + + nvinfer1::IElementWiseLayer* shifted = + network.addElementWise(input, *addend->getOutput(0), nvinfer1::ElementWiseOperation::kSUM); + nvinfer1::ISoftMaxLayer* over_cols = network.addSoftMax(*shifted->getOutput(0)); + over_cols->setAxes(1u << 2); + nvinfer1::ISoftMaxLayer* over_rows = network.addSoftMax(*over_cols->getOutput(0)); + over_rows->setAxes(1u << 1); + nvinfer1::IElementWiseLayer* scaled = + network.addElementWise(*over_rows->getOutput(0), *scale->getOutput(0), nvinfer1::ElementWiseOperation::kPROD); + scaled->getOutput(0)->setName("output_0"); + network.markOutput(*scaled->getOutput(0)); + return true; +} + +// TensorRT routes a pointwise chain through the I/O tensors alone, so this +// engine's activation requirement is zero -- the same answer it gives for a +// failed query. Every layer is parameterless, because a default alpha or beta +// can collapse a chain to a constant and make an output comparison vacuous. +bool add_scratch_free_net(nvinfer1::INetworkDefinition& network, nvinfer1::ITensor& input) { + static const nvinfer1::ActivationType kChain[] = { + nvinfer1::ActivationType::kSIGMOID, + nvinfer1::ActivationType::kTANH, + nvinfer1::ActivationType::kSOFTSIGN, + nvinfer1::ActivationType::kSIGMOID, + nvinfer1::ActivationType::kTANH, + nvinfer1::ActivationType::kSOFTSIGN, + }; + nvinfer1::ITensor* t = &input; + for (const nvinfer1::ActivationType op : kChain) { + nvinfer1::IActivationLayer* layer = network.addActivation(*t, op); + if (layer == nullptr) { + return false; + } + t = layer->getOutput(0); + } + t->setName("output_0"); + network.markOutput(*t); + return true; +} + +std::vector build_engine_blob(bool needs_scratch) { + static BuilderLogger logger; + + TRTUniquePtr builder(nvinfer1::createInferBuilder(logger)); + if (builder == nullptr) { + return {}; + } + TRTUniquePtr network(builder->createNetworkV2(0)); + if (network == nullptr) { + return {}; + } + + nvinfer1::ITensor* input = network->addInput("input_0", nvinfer1::DataType::kFLOAT, nvinfer1::Dims3{1, kRows, kCols}); + if (input == nullptr) { + return {}; + } + const bool built = needs_scratch ? add_scratch_needing_net(*network, *input) : add_scratch_free_net(*network, *input); + if (!built) { + return {}; + } + + TRTUniquePtr config(builder->createBuilderConfig()); + if (config == nullptr) { + return {}; + } + nvinfer1::IOptimizationProfile* profile = builder->createOptimizationProfile(); + const nvinfer1::Dims3 shape{1, kRows, kCols}; + profile->setDimensions("input_0", nvinfer1::OptProfileSelector::kMIN, shape); + profile->setDimensions("input_0", nvinfer1::OptProfileSelector::kOPT, shape); + profile->setDimensions("input_0", nvinfer1::OptProfileSelector::kMAX, shape); + config->addOptimizationProfile(profile); + + TRTUniquePtr plan(builder->buildSerializedNetwork(*network, *config)); + if (plan == nullptr) { + return {}; + } + + const std::string metadata = + R"({"io_bindings":[{"name":"input_0","is_input":true},{"name":"output_0","is_input":false}],)" + R"("hardware_compatible":false,"device_id":0})"; + const auto metadata_offset = static_cast(kHeaderSize); + const auto metadata_size = static_cast(metadata.size()); + const auto engine_offset = static_cast(align_up(metadata_offset + metadata_size, kEngineAlignment)); + + std::vector blob(static_cast(engine_offset) + plan->size(), 0); + std::memcpy(blob.data(), kMagic, sizeof(kMagic)); + write_field(blob, kMetadataOffsetField, metadata_offset); + write_field(blob, kMetadataSizeField, metadata_size); + write_field(blob, kEngineOffsetField, engine_offset); + write_field(blob, kEngineSizeField, static_cast(plan->size())); + std::memcpy(blob.data() + metadata_offset, metadata.data(), metadata.size()); + std::memcpy(blob.data() + engine_offset, plan->data(), plan->size()); + return blob; +} + +// The activation scratch one context of the shared engine needs, read the way +// execute() reads it. Zero if the engine could not be measured. +std::size_t measure_engine_scratch(const std::vector& blob) { + static BuilderLogger logger; + TensorRTBlobHeader header; + if (!TensorRTBlobHeader::parse(blob.data(), blob.size(), header)) { + return 0; + } + TRTUniquePtr runtime(nvinfer1::createInferRuntime(logger)); + if (runtime == nullptr) { + return 0; + } + TRTUniquePtr engine( + runtime->deserializeCudaEngine(TensorRTBlobHeader::engine_data(blob.data(), header), header.engine_size)); + if (engine == nullptr) { + return 0; + } + // kUSER_MANAGED so the probe context itself allocates no scratch to measure. + TRTUniquePtr ctx( + engine->createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED)); + if (ctx == nullptr) { + return 0; + } + if (!ctx->setInputShape("input_0", nvinfer1::Dims3{1, kRows, kCols})) { + return 0; + } + return ctx->updateDeviceMemorySizeForShapes(); +} + +// What the engine reports it needs, read the way init() reads it. A negative +// result means the blob could not be opened, which no engine reports and which +// no test may mistake for a scratch-free engine. +std::int64_t engine_scratch_requirement(const std::vector& blob) { + static BuilderLogger logger; + TensorRTBlobHeader header; + if (!TensorRTBlobHeader::parse(blob.data(), blob.size(), header)) { + return -1; + } + TRTUniquePtr runtime(nvinfer1::createInferRuntime(logger)); + if (runtime == nullptr) { + return -1; + } + TRTUniquePtr engine( + runtime->deserializeCudaEngine(TensorRTBlobHeader::engine_data(blob.data(), header), header.engine_size)); + if (engine == nullptr) { + return -1; + } + return engine->getDeviceMemorySizeV2(); +} + +// --------------------------------------------------------------------------- +// One loaded delegate handle plus the device-resident I/O its execute() needs +// --------------------------------------------------------------------------- + +// Reproducible on both sides and non-uniform: a constant input would make the +// softmaxes uniform and stop the output depending on the tensor under test. +float pattern(std::size_t index, std::uint32_t seed) { + std::uint32_t h = static_cast(index) * 2654435761u + seed * 40503u; + h ^= h >> 15; + return static_cast(h % 1000u) / 500.0f - 1.0f; +} + +class LoadedEngine { + public: + LoadedEngine() = default; + LoadedEngine(const LoadedEngine&) = delete; + LoadedEngine& operator=(const LoadedEngine&) = delete; + + ~LoadedEngine() { + if (handle_ != nullptr) { + backend_.destroy(handle_); + } + cudaFree(device_in_); + cudaFree(device_out_); + } + + // Loads the blob through the backend, capturing whatever the shared-scratch + // option is set to at this moment. + Error load(const std::vector& blob, std::uint32_t seed) { + std::vector host_in(kElems); + for (std::size_t i = 0; i < kElems; ++i) { + host_in[i] = pattern(i, seed); + } + if (cudaMalloc(&device_in_, kBytes) != cudaSuccess || cudaMalloc(&device_out_, kBytes) != cudaSuccess) { + return Error::MemoryAllocationFailed; + } + if (cudaMemcpy(device_in_, host_in.data(), kBytes, cudaMemcpyHostToDevice) != cudaSuccess) { + return Error::Internal; + } + + arena_storage_.resize(kArenaBytes); + arena_ = std::make_unique(static_cast(kArenaBytes), arena_storage_.data()); + BackendInitContext init_context(arena_.get()); + FreeableBuffer processed(blob.data(), blob.size(), nullptr); + const auto result = backend_.init(init_context, &processed, ArrayRef{}); + if (!result.ok()) { + return result.error(); + } + handle_ = result.get(); + return Error::Ok; + } + + bool fill_output(float value) { + const std::vector host(kElems, value); + return cudaMemcpy(device_out_, host.data(), kBytes, cudaMemcpyHostToDevice) == cudaSuccess; + } + + // Runs one inference on `stream`. Returns without waiting for the enqueue, + // which is the state the pool's handoff exists to order. + Error run(cudaStream_t stream) { + // Separate arrays: execute() resizes the output tensor to the shape TensorRT + // inferred, which writes through whichever array that tensor was given. + SizesType in_sizes[3] = {1, kRows, kCols}; + SizesType out_sizes[3] = {1, kRows, kCols}; + ::executorch::aten::TensorImpl in_impl(ScalarType::Float, 3, in_sizes, device_in_); + ::executorch::aten::TensorImpl out_impl(ScalarType::Float, 3, out_sizes, device_out_); + ::executorch::aten::Tensor in_tensor(&in_impl); + ::executorch::aten::Tensor out_tensor(&out_impl); + EValue in_value(in_tensor); + EValue out_value(out_tensor); + EValue* args[2] = {&in_value, &out_value}; + + BackendExecutionContext exec_context; + ::executorch::extension::cuda::CallerStreamGuard guard(stream); + return backend_.execute(exec_context, handle_, Span(args, 2)); + } + + std::vector read_output() const { + std::vector host_out(kElems); + if (cudaMemcpy(host_out.data(), device_out_, kBytes, cudaMemcpyDeviceToHost) != cudaSuccess) { + host_out.clear(); + } + return host_out; + } + + const EngineHandle* handle() const { + return static_cast(handle_); + } + + private: + // EngineHandle is placement-newed into this arena by init(), and the arena is + // never reset, so it only has to hold one instance. + static constexpr std::size_t kArenaBytes = 4096; + + TensorRTBackend backend_; + std::vector arena_storage_; + std::unique_ptr arena_; + DelegateHandle* handle_ = nullptr; + void* device_in_ = nullptr; + void* device_out_ = nullptr; +}; + +std::size_t device_bytes_in_use() { + std::size_t free_bytes = 0; + std::size_t total_bytes = 0; + if (cudaMemGetInfo(&free_bytes, &total_bytes) != cudaSuccess) { + return 0; + } + return total_bytes - free_bytes; +} + +Error set_shared_scratch(TensorRTBackend& backend, bool enabled) { + BackendOption option; + std::strncpy(option.key, kOptionKey, sizeof(option.key) - 1); + option.value = enabled; + BackendOption options[1] = {option}; + BackendOptionContext context; + return backend.set_option(context, Span(options, 1)); +} + +// --------------------------------------------------------------------------- +// Fixture +// --------------------------------------------------------------------------- + +class SharedScratchBackendTest : public ::testing::Test { + protected: + // Building the engine dominates the runtime of this target, so it is built + // once and every test loads the same blob. + static void SetUpTestSuite() { + ::executorch::runtime::runtime_init(); + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count == 0) { + return; + } + blob_ = build_engine_blob(true); + scratch_free_blob_ = build_engine_blob(false); + if (blob_.empty() || scratch_free_blob_.empty()) { + return; + } + scratch_bytes_ = measure_engine_scratch(blob_); + engine_bytes_ = engine_scratch_requirement(blob_); + scratch_free_engine_bytes_ = engine_scratch_requirement(scratch_free_blob_); + } + + void SetUp() override { + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count == 0) { + GTEST_SKIP() << "no CUDA device: the shared-scratch backend path is not covered by this run"; + } + ASSERT_FALSE(blob_.empty()) << "TensorRT could not build the fixture engine"; + ASSERT_FALSE(scratch_free_blob_.empty()) << "TensorRT could not build the scratch-free fixture engine"; + ASSERT_EQ(set_shared_scratch(backend_, false), Error::Ok); + } + + void TearDown() override { + set_shared_scratch(backend_, false); + } + + const std::vector& blob() const { + return blob_; + } + + const std::vector& scratch_free_blob() const { + return scratch_free_blob_; + } + + TensorRTBackend backend_; + static std::vector blob_; + static std::vector scratch_free_blob_; + static std::size_t scratch_bytes_; + static std::int64_t engine_bytes_; + static std::int64_t scratch_free_engine_bytes_; +}; + +std::vector SharedScratchBackendTest::blob_; +std::vector SharedScratchBackendTest::scratch_free_blob_; +std::size_t SharedScratchBackendTest::scratch_bytes_ = 0; +std::int64_t SharedScratchBackendTest::engine_bytes_ = -1; +std::int64_t SharedScratchBackendTest::scratch_free_engine_bytes_ = -1; + +// --------------------------------------------------------------------------- +// set_option +// --------------------------------------------------------------------------- + +// The foreign key is sent from both settings, because from one of them the test +// cannot tell a key that is ignored from a key that resets the setting to that +// value. +TEST_F(SharedScratchBackendTest, SetOptionAcceptsAKeyThisBackendDoesNotRead) { + BackendOption foreign; + std::strncpy(foreign.key, "some_other_backends_option", sizeof(foreign.key) - 1); + foreign.value = 7; + BackendOption options[1] = {foreign}; + BackendOptionContext context; + + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + EXPECT_EQ(backend_.set_option(context, Span(options, 1)), Error::Ok); + LoadedEngine after_on; + ASSERT_EQ(after_on.load(blob(), 1), Error::Ok); + EXPECT_TRUE(after_on.handle()->shared_scratch) << "a foreign key turned the shared-scratch setting off"; + + ASSERT_EQ(set_shared_scratch(backend_, false), Error::Ok); + EXPECT_EQ(backend_.set_option(context, Span(options, 1)), Error::Ok); + LoadedEngine after_off; + ASSERT_EQ(after_off.load(blob(), 12), Error::Ok); + EXPECT_FALSE(after_off.handle()->shared_scratch) << "a foreign key turned the shared-scratch setting on"; +} + +TEST_F(SharedScratchBackendTest, SetOptionStoresTheBooleanItIsGiven) { + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + LoadedEngine pooled; + ASSERT_EQ(pooled.load(blob(), 2), Error::Ok); + EXPECT_TRUE(pooled.handle()->shared_scratch); + + ASSERT_EQ(set_shared_scratch(backend_, false), Error::Ok); + LoadedEngine priv; + ASSERT_EQ(priv.load(blob(), 3), Error::Ok); + EXPECT_FALSE(priv.handle()->shared_scratch); +} + +TEST_F(SharedScratchBackendTest, SetOptionRejectsANonBooleanAndLeavesTheSettingAlone) { + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + + BackendOption wrong_type; + std::strncpy(wrong_type.key, kOptionKey, sizeof(wrong_type.key) - 1); + // The int has to coerce to the opposite of the setting above: one that coerced + // to the same value would leave the setting exactly where the assertion at the + // end expects to find it, whether it was rejected or not. + wrong_type.value = 0; + BackendOption options[1] = {wrong_type}; + BackendOptionContext context; + EXPECT_EQ(backend_.set_option(context, Span(options, 1)), Error::InvalidArgument); + + LoadedEngine engine; + ASSERT_EQ(engine.load(blob(), 4), Error::Ok); + EXPECT_TRUE(engine.handle()->shared_scratch) << "a rejected option still moved the shared-scratch setting"; +} + +// A context's allocation strategy is fixed when the context is created, so the +// option cannot be re-read per call. +TEST_F(SharedScratchBackendTest, EachEngineCapturesTheSettingInEffectAtItsOwnLoad) { + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + LoadedEngine pooled; + ASSERT_EQ(pooled.load(blob(), 5), Error::Ok); + + ASSERT_EQ(set_shared_scratch(backend_, false), Error::Ok); + LoadedEngine priv; + ASSERT_EQ(priv.load(blob(), 6), Error::Ok); + + EXPECT_TRUE(pooled.handle()->shared_scratch); + EXPECT_FALSE(priv.handle()->shared_scratch); + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + EXPECT_EQ(pooled.run(stream), Error::Ok); + EXPECT_EQ(priv.run(stream), Error::Ok); + EXPECT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + EXPECT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +// --------------------------------------------------------------------------- +// The pooled execute() path +// --------------------------------------------------------------------------- + +TEST_F(SharedScratchBackendTest, APooledEngineProducesWhatAPrivateScratchEngineProduces) { + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + + LoadedEngine priv; + ASSERT_EQ(priv.load(blob(), 7), Error::Ok); + // Two arms on the same setting produce the same bytes whichever setting that + // is, so the comparison at the end is worth nothing unless each arm is pinned + // to the side it stands for. + ASSERT_FALSE(priv.handle()->shared_scratch); + ASSERT_EQ(priv.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::vector expected = priv.read_output(); + + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + LoadedEngine pooled; + ASSERT_EQ(pooled.load(blob(), 7), Error::Ok); + ASSERT_TRUE(pooled.handle()->shared_scratch); + ASSERT_EQ(pooled.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::vector actual = pooled.read_output(); + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); + + ASSERT_EQ(expected.size(), kElems); + ASSERT_EQ(actual.size(), kElems); + // A degenerate output would make the comparison above pass without depending + // on the engine having run. + bool varies = false; + for (std::size_t i = 1; i < kElems && !varies; ++i) { + varies = expected[i] != expected[0]; + } + EXPECT_TRUE(varies) << "the reference output is constant, so the comparison proves nothing"; + EXPECT_EQ(std::memcmp(expected.data(), actual.data(), kBytes), 0); +} + +TEST_F(SharedScratchBackendTest, PooledEnginesShareOneActivationScratchAllocation) { + ASSERT_GE(scratch_bytes_, kMinMeasurableScratch) + << "the fixture engine reports " << scratch_bytes_ + << " bytes of activation scratch, too little for the memory comparison to resolve"; + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + + // One load and run first, so the one-time TensorRT runtime and CUDA module + // allocations land outside both measurements. + { + LoadedEngine warmup; + ASSERT_EQ(warmup.load(blob(), 8), Error::Ok); + ASSERT_EQ(warmup.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + } + + std::size_t private_cost = 0; + { + const std::size_t before = device_bytes_in_use(); + std::vector> engines; + for (int i = 0; i < kEngineCount; ++i) { + engines.push_back(std::make_unique()); + ASSERT_EQ(engines.back()->load(blob(), 9), Error::Ok); + ASSERT_EQ(engines.back()->run(stream), Error::Ok); + } + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::size_t after = device_bytes_in_use(); + // The subtraction is unsigned, so a fall in device-wide usage would wrap it + // to a number that satisfies the comparison at the end for free. + ASSERT_GE(after, before) << "device-wide memory in use fell across the private-scratch measurement, so " + "something outside this test is releasing memory on this device"; + private_cost = after - before; + } + + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + std::size_t pooled_cost = 0; + { + const std::size_t before = device_bytes_in_use(); + std::vector> engines; + for (int i = 0; i < kEngineCount; ++i) { + engines.push_back(std::make_unique()); + ASSERT_EQ(engines.back()->load(blob(), 9), Error::Ok); + ASSERT_EQ(engines.back()->run(stream), Error::Ok); + } + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::size_t after = device_bytes_in_use(); + ASSERT_GE(after, before) << "device-wide memory in use fell across the pooled measurement, so " + "something outside this test is releasing memory on this device"; + pooled_cost = after - before; + } + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); + + // Half the ideal saving, which leaves room for allocator granularity without + // admitting a run in which every context still carries its own scratch. + const std::size_t expected_saving = (kEngineCount - 1) * scratch_bytes_ / 2; + EXPECT_GE(private_cost, pooled_cost + expected_saving) + << kEngineCount << " engines cost " << private_cost << " bytes with private scratch and " << pooled_cost + << " pooled, against " << scratch_bytes_ << " bytes of scratch each"; +} + +// --------------------------------------------------------------------------- +// An engine that needs no activation scratch +// --------------------------------------------------------------------------- + +// updateDeviceMemorySizeForShapes() answers a failed query and an engine that +// needs nothing identically, so execute() separates them on the engine's own +// requirement. Everything below rests on that requirement telling the two +// fixture networks apart, which is why it is asserted on its own first. +TEST_F(SharedScratchBackendTest, TheEngineLevelRequirementSeparatesTheTwoFixtureEngines) { + EXPECT_EQ(scratch_free_engine_bytes_, 0) + << "the pointwise chain reports " << scratch_free_engine_bytes_ + << " bytes of activation scratch, so it no longer covers the scratch-free case"; + EXPECT_GT(engine_bytes_, 0) << "the two-softmax network reports no activation scratch, so it no longer covers the " + "case a failed query has to be told apart from"; +} + +TEST_F(SharedScratchBackendTest, EachEngineRecordsItsOwnActivationScratchRequirement) { + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + LoadedEngine needing; + LoadedEngine scratch_free; + ASSERT_EQ(needing.load(blob(), 12), Error::Ok); + ASSERT_EQ(scratch_free.load(scratch_free_blob(), 13), Error::Ok); + + EXPECT_EQ(static_cast(needing.handle()->engine_scratch_bytes), engine_bytes_); + EXPECT_EQ(scratch_free.handle()->engine_scratch_bytes, 0u); +} + +// Turning the pool on must not turn an engine that legitimately needs no +// activation scratch into a failure. +TEST_F(SharedScratchBackendTest, AnEngineNeedingNoActivationScratchRunsWithThePoolEnabled) { + ASSERT_EQ(scratch_free_engine_bytes_, 0) << "the fixture engine needs scratch, so this test covers nothing"; + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + + LoadedEngine priv; + ASSERT_EQ(priv.load(scratch_free_blob(), 14), Error::Ok); + ASSERT_TRUE(priv.fill_output(kSentinel)); + ASSERT_EQ(priv.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::vector expected = priv.read_output(); + + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + LoadedEngine pooled; + ASSERT_EQ(pooled.load(scratch_free_blob(), 14), Error::Ok); + ASSERT_TRUE(pooled.handle()->shared_scratch); + ASSERT_TRUE(pooled.fill_output(kSentinel)); + EXPECT_EQ(pooled.run(stream), Error::Ok) << "the pool rejected an engine that needs no activation scratch"; + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::vector actual = pooled.read_output(); + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); + + ASSERT_EQ(expected.size(), kElems); + ASSERT_EQ(actual.size(), kElems); + // Without these two the comparison would be satisfied by an execute() that + // wrote nothing, and by a network whose output does not depend on its input. + EXPECT_NE(expected[0], kSentinel) << "the reference output is the sentinel, so a skipped enqueue would pass"; + bool varies = false; + for (std::size_t i = 1; i < kElems && !varies; ++i) { + varies = expected[i] != expected[0]; + } + EXPECT_TRUE(varies) << "the reference output is constant, so the comparison proves nothing"; + EXPECT_EQ(std::memcmp(expected.data(), actual.data(), kBytes), 0); +} + +// --------------------------------------------------------------------------- +// The enqueue handoff, single-threaded, two caller streams +// --------------------------------------------------------------------------- + +struct StreamGate { + std::mutex mu; + std::condition_variable cv; + bool open = false; + // Set when the watchdog, not the test, had to open the gate. + std::atomic forced_open{false}; +}; + +void CUDART_CB hold_stream(void* user_data) { + StreamGate* gate = static_cast(user_data); + std::unique_lock lock(gate->mu); + gate->cv.wait(lock, [gate] { return gate->open; }); +} + +// Long enough that the wait the test performs while the gate is shut, and the +// two enqueues before it, are nowhere near it. +constexpr std::chrono::seconds kGateWatchdog{60}; + +// Opens the gate and waits for the held work to drain, by two routes because two +// different things can go wrong. A held stream outlives any assertion that +// returns early, and every teardown path below -- cudaFree, the delegate +// destructor -- blocks on it, so the destructor opens the gate for a test that +// does not reach its end. That is no help if a delegate call blocks on the held +// stream instead of returning, since the calling thread then never runs the +// destructor either: the watchdog covers that, and records that it had to, so +// the outcome is a failure naming the cause rather than a process that never +// exits. +class GateRelease { + public: + GateRelease(StreamGate& gate, cudaStream_t stream) + : gate_(gate), stream_(stream), deadline_(std::chrono::steady_clock::now() + kGateWatchdog) { + watchdog_ = std::thread([this] { + std::unique_lock lock(gate_.mu); + if (!gate_.cv.wait_until(lock, deadline_, [this] { return gate_.open; })) { + gate_.open = true; + gate_.forced_open.store(true); + lock.unlock(); + gate_.cv.notify_all(); + } + }); + } + + ~GateRelease() { + release(); + watchdog_.join(); + } + + void release() { + if (released_) { + return; + } + released_ = true; + { + std::lock_guard lock(gate_.mu); + gate_.open = true; + } + gate_.cv.notify_all(); + cudaStreamSynchronize(stream_); + } + + private: + StreamGate& gate_; + cudaStream_t stream_; + std::chrono::steady_clock::time_point deadline_; + std::thread watchdog_; + bool released_ = false; +}; + +// Two engines on one device share one scratch buffer, so the second engine's +// enqueue must not start before the first one's has finished with it. The two +// run on different streams, which is what the README permits and what the event +// handoff is for: nothing but the handoff orders them. +TEST_F(SharedScratchBackendTest, ASecondPooledEnqueueWaitsForTheFirstOnAnotherStream) { + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + + cudaStream_t first_stream = nullptr; + cudaStream_t second_stream = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&first_stream, cudaStreamNonBlocking), cudaSuccess); + ASSERT_EQ(cudaStreamCreateWithFlags(&second_stream, cudaStreamNonBlocking), cudaSuccess); + + LoadedEngine first; + LoadedEngine second; + ASSERT_EQ(first.load(blob(), 10), Error::Ok); + ASSERT_EQ(second.load(blob(), 11), Error::Ok); + ASSERT_TRUE(first.handle()->shared_scratch); + ASSERT_TRUE(second.handle()->shared_scratch); + + // Held work at the head of the first stream, so the first enqueue and the + // completion event recorded after it stay pending for as long as the test + // wants them to. + StreamGate gate; + ASSERT_EQ(cudaLaunchHostFunc(first_stream, hold_stream, &gate), cudaSuccess); + GateRelease gate_release(gate, first_stream); + + // Held for the checks below, which take the watchdog flag first: a call that + // blocks on the held stream comes back with an error once the watchdog opens + // the gate, and that error on its own does not say so. + const Error first_error = first.run(first_stream); + const Error second_error = second.run(second_stream); + + bool second_finished_early = false; + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (std::chrono::steady_clock::now() < deadline) { + if (cudaStreamQuery(second_stream) == cudaSuccess) { + second_finished_early = true; + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + + ASSERT_FALSE(gate.forced_open.load()) + << "the watchdog had to open the gate: a call blocked on the held stream rather than returning, " + "so nothing below was measured under the conditions it describes"; + ASSERT_EQ(first_error, Error::Ok); + ASSERT_EQ(second_error, Error::Ok); + + gate_release.release(); + ASSERT_EQ(cudaStreamSynchronize(first_stream), cudaSuccess); + // Rules out the second engine's work having failed rather than been held, + // which would leave the check below false for the wrong reason. + ASSERT_EQ(cudaStreamSynchronize(second_stream), cudaSuccess); + + EXPECT_FALSE(second_finished_early) + << "the second engine ran to completion while the first one's enqueue was still holding the shared buffer"; + + const std::vector first_output = first.read_output(); + const std::vector second_output = second.read_output(); + ASSERT_EQ(first_output.size(), kElems); + ASSERT_EQ(second_output.size(), kElems); + EXPECT_NE(std::memcmp(first_output.data(), second_output.data(), kBytes), 0) + << "the two engines were given different inputs but produced the same output"; + + ASSERT_EQ(cudaStreamDestroy(first_stream), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(second_stream), cudaSuccess); +} + +} // namespace +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/tests/cpp/executorch/test_shared_scratch_pool.cpp b/tests/cpp/executorch/test_shared_scratch_pool.cpp index a65657e5df..0532d0fe40 100644 --- a/tests/cpp/executorch/test_shared_scratch_pool.cpp +++ b/tests/cpp/executorch/test_shared_scratch_pool.cpp @@ -1,11 +1,23 @@ +// Pins the shared scratch pool helper: its grow, reuse and per-device policy and +// its enqueue-handoff rule, driven over fakes so no CUDA device is needed. +// +// This exercises the helper, not the backend: it does not link the delegate, so +// it cannot catch the delegate calling the helper wrongly or ceasing to call it. +// test_shared_scratch_backend covers that, and needs a GPU to do it. + #include "torch_tensorrt/executorch/SharedScratchPool.h" #include "gtest/gtest.h" +#include +#include #include #include #include -#include +#include +#include +#include +#include #include #include @@ -14,11 +26,11 @@ namespace executorch_backend { namespace { // Fake device allocator: hands out distinct non-null pointers and records every -// allocation size and every released pointer, so tests can assert the pool's -// grow/reuse/per-device policy without a CUDA device. +// allocation size and every release, so tests can assert the pool's grow/reuse +// policy and what each release was told to wait for, without a CUDA device. struct FakeAllocator { std::vector alloc_sizes; - std::vector released; + std::vector> released; std::uintptr_t next = 0x1000; bool fail_next = false; @@ -33,8 +45,8 @@ struct FakeAllocator { return p; } - void release(void* p) { - released.push_back(p); + void release(void* p, cudaEvent_t wait_for) { + released.emplace_back(p, wait_for); } int alloc_count() const { @@ -42,24 +54,41 @@ struct FakeAllocator { } }; -using Pool = std::unordered_map>; +// Stands in for the CUDA event factory: hands out distinct non-null handles and +// counts calls, so a test can tell a slot that reuses its event from one that +// creates a new one every call. +struct FakeEventFactory { + int created = 0; + std::uintptr_t next = 0xE000; + bool fail_next = false; + + cudaEvent_t operator()() { + if (fail_next) { + fail_next = false; + return nullptr; + } + ++created; + cudaEvent_t e = reinterpret_cast(next); + next += 0x100; + return e; + } +}; -void* call(Pool& pool, FakeAllocator& a, int device_id, std::size_t need, std::size_t& out_size) { +void* call(SharedScratchDevice& dev, FakeAllocator& a, std::size_t need, std::size_t& out_size) { return shared_scratch_get_or_grow( - pool, - device_id, + dev, need, out_size, [&a](std::size_t bytes) { return a.alloc(bytes); }, - [&a](void* p) { a.release(p); }); + [&a](void* p, cudaEvent_t wait_for) { a.release(p, wait_for); }); } TEST(SharedScratchPool, FirstRequestAllocatesExactSize) { - Pool pool; + SharedScratchDevice dev; FakeAllocator a; std::size_t out = 0; - void* p = call(pool, a, /*device_id=*/0, /*need=*/1024, out); + void* p = call(dev, a, /*need=*/1024, out); EXPECT_NE(p, nullptr); EXPECT_EQ(out, 1024u); @@ -69,17 +98,17 @@ TEST(SharedScratchPool, FirstRequestAllocatesExactSize) { } TEST(SharedScratchPool, ReusesWhenExistingBufferIsLargeEnough) { - Pool pool; + SharedScratchDevice dev; FakeAllocator a; std::size_t out = 0; - void* first = call(pool, a, 0, 4096, out); + void* first = call(dev, a, 4096, out); // A smaller and an equal request must both reuse the same buffer (no realloc). // The smaller one reports into a fresh out2, so what the reuse path writes is // asserted rather than what the first call left in `out`. std::size_t out2 = 0; - void* second = call(pool, a, 0, 1000, out2); - void* third = call(pool, a, 0, 4096, out); + void* second = call(dev, a, 1000, out2); + void* third = call(dev, a, 4096, out); EXPECT_EQ(second, first); EXPECT_EQ(third, first); @@ -91,81 +120,104 @@ TEST(SharedScratchPool, ReusesWhenExistingBufferIsLargeEnough) { } TEST(SharedScratchPool, GrowsMonotonicallyToMaxAndReleasesOldBuffer) { - Pool pool; + SharedScratchDevice dev; FakeAllocator a; std::size_t out = 0; - void* small = call(pool, a, 0, 1024, out); - void* big = call(pool, a, 0, 8192, out); + void* small = call(dev, a, 1024, out); + void* big = call(dev, a, 8192, out); EXPECT_NE(big, small); EXPECT_EQ(out, 8192u); ASSERT_EQ(a.alloc_count(), 2); EXPECT_EQ(a.alloc_sizes[1], 8192u); ASSERT_EQ(a.released.size(), 1u); - EXPECT_EQ(a.released[0], small); + EXPECT_EQ(a.released[0].first, small); // A subsequent smaller request reuses the grown buffer -- pool never shrinks. - void* reuse = call(pool, a, 0, 512, out); + void* reuse = call(dev, a, 512, out); EXPECT_EQ(reuse, big); EXPECT_EQ(out, 8192u); EXPECT_EQ(a.alloc_count(), 2); } -TEST(SharedScratchPool, KeepsIndependentBufferPerDevice) { - Pool pool; +TEST(SharedScratchPool, GrowWaitsOnTheRecordedEnqueueBeforeReleasing) { + SharedScratchDevice dev; FakeAllocator a; + FakeEventFactory events; std::size_t out = 0; - void* dev0 = call(pool, a, /*device_id=*/0, 2048, out); - void* dev1 = call(pool, a, /*device_id=*/1, 2048, out); + void* small = call(dev, a, 1024, out); + ASSERT_NE(small, nullptr); - EXPECT_NE(dev0, dev1); - EXPECT_EQ(a.alloc_count(), 2); - EXPECT_TRUE(a.released.empty()); + // An enqueue against `small` has been submitted and recorded, so the release + // has something specific to outlive. + const SharedScratchHandoff handoff = shared_scratch_claim_event(dev, std::ref(events)); + ASSERT_EQ(shared_scratch_mark_in_flight(dev), handoff.event); + + ASSERT_NE(call(dev, a, 8192, out), nullptr); - // Growing device 1 must not touch device 0's buffer. - void* dev1_big = call(pool, a, 1, 9000, out); - void* dev0_again = call(pool, a, 0, 2048, out); - EXPECT_NE(dev1_big, dev1); - EXPECT_EQ(dev0_again, dev0); ASSERT_EQ(a.released.size(), 1u); - EXPECT_EQ(a.released[0], dev1); + EXPECT_EQ(a.released[0].first, small); + // The release is handed the event that enqueue was recorded on, so it waits for + // that enqueue rather than for everything queued on the device. + EXPECT_EQ(a.released[0].second, handoff.event); } -TEST(SharedScratchPool, AllocationFailureLeavesExistingSlotUntouched) { - Pool pool; +TEST(SharedScratchPool, GrowHasNothingToWaitForWhenNoEnqueueWasRecorded) { + SharedScratchDevice dev; FakeAllocator a; + FakeEventFactory events; std::size_t out = 0; - void* first = call(pool, a, 0, 1024, out); + void* small = call(dev, a, 1024, out); + ASSERT_NE(small, nullptr); + // The slot has an event, but nothing has been recorded on it: claiming the + // handoff is not the same as enqueueing against the buffer. + ASSERT_NE(shared_scratch_claim_event(dev, std::ref(events)).event, nullptr); + + ASSERT_NE(call(dev, a, 8192, out), nullptr); + + ASSERT_EQ(a.released.size(), 1u); + EXPECT_EQ(a.released[0].first, small); + EXPECT_EQ(a.released[0].second, nullptr); +} + +TEST(SharedScratchPool, AllocationFailureLeavesExistingBufferUntouched) { + SharedScratchDevice dev; + FakeAllocator a; + std::size_t out = 0; + + void* first = call(dev, a, 1024, out); ASSERT_NE(first, nullptr); // A growth whose allocation fails must return nullptr and keep the old buffer, // so the caller can surface the error without corrupting the pool. a.fail_next = true; std::size_t out2 = 0; - void* failed = call(pool, a, 0, 8192, out2); + void* failed = call(dev, a, 8192, out2); EXPECT_EQ(failed, nullptr); EXPECT_TRUE(a.released.empty()); - // The pool still holds the original buffer and serves it on the next request. - void* again = call(pool, a, 0, 1024, out); + // The device still holds the original buffer and serves it on the next request. + void* again = call(dev, a, 1024, out); EXPECT_EQ(again, first); EXPECT_EQ(out, 1024u); } TEST(SharedScratchPool, FirstAllocationFailureReturnsNullAndStoresNothing) { - Pool pool; + SharedScratchDevice dev; FakeAllocator a; std::size_t out = 0; a.fail_next = true; - void* p = call(pool, a, 0, 1024, out); + void* p = call(dev, a, 1024, out); EXPECT_EQ(p, nullptr); + EXPECT_EQ(dev.buffer, nullptr); + EXPECT_EQ(dev.capacity, 0u); // Nothing stored: a later successful request allocates fresh. - void* q = call(pool, a, 0, 1024, out); + void* q = call(dev, a, 1024, out); EXPECT_NE(q, nullptr); EXPECT_EQ(a.alloc_count(), 1); } @@ -174,33 +226,11 @@ TEST(SharedScratchPool, FirstAllocationFailureReturnsNullAndStoresNothing) { // Ordering the shared buffer's handoff from one enqueue to the next. // --------------------------------------------------------------------------- -// Stands in for the CUDA event factory: hands out distinct non-null handles and -// counts calls, so a test can tell a slot that reuses its event from one that -// creates a new one every call. -struct FakeEventFactory { - int created = 0; - std::uintptr_t next = 0xE000; - bool fail_next = false; - - cudaEvent_t operator()() { - if (fail_next) { - fail_next = false; - return nullptr; - } - ++created; - cudaEvent_t e = reinterpret_cast(next); - next += 0x100; - return e; - } -}; - -using Markers = std::unordered_map; - TEST(SharedScratchHandoffTest, FirstUseCreatesTheSlotsEventAndWaitsForNothing) { - Markers markers; + SharedScratchDevice dev; FakeEventFactory events; - const SharedScratchHandoff handoff = shared_scratch_claim_event(markers, /*device_id=*/0, std::ref(events)); + const SharedScratchHandoff handoff = shared_scratch_claim_event(dev, std::ref(events)); EXPECT_NE(handoff.event, nullptr); EXPECT_FALSE(handoff.needs_wait); @@ -208,23 +238,23 @@ TEST(SharedScratchHandoffTest, FirstUseCreatesTheSlotsEventAndWaitsForNothing) { } TEST(SharedScratchHandoffTest, EveryUseAfterAnEnqueueWaitsOnTheSameEvent) { - Markers markers; + SharedScratchDevice dev; FakeEventFactory events; - const SharedScratchHandoff first = shared_scratch_claim_event(markers, 0, std::ref(events)); + const SharedScratchHandoff first = shared_scratch_claim_event(dev, std::ref(events)); ASSERT_FALSE(first.needs_wait); - EXPECT_EQ(shared_scratch_mark_in_flight(markers, 0), first.event); + EXPECT_EQ(shared_scratch_mark_in_flight(dev), first.event); // Every later enqueue waits, however many there have been and whichever stream // each of them ran on: the marker records that the buffer was handed out, not // who it was handed to. Comparing stream handles instead would let a caller // through whenever its handle matched the recorded one, including when CUDA has // recycled that value for a different stream. - const SharedScratchHandoff second = shared_scratch_claim_event(markers, 0, std::ref(events)); + const SharedScratchHandoff second = shared_scratch_claim_event(dev, std::ref(events)); EXPECT_TRUE(second.needs_wait); EXPECT_EQ(second.event, first.event); - const SharedScratchHandoff third = shared_scratch_claim_event(markers, 0, std::ref(events)); + const SharedScratchHandoff third = shared_scratch_claim_event(dev, std::ref(events)); EXPECT_TRUE(third.needs_wait); EXPECT_EQ(third.event, first.event); @@ -234,51 +264,220 @@ TEST(SharedScratchHandoffTest, EveryUseAfterAnEnqueueWaitsOnTheSameEvent) { } TEST(SharedScratchHandoffTest, KeepsAnIndependentMarkerPerDevice) { - Markers markers; + SharedScratchPool pool; FakeEventFactory events; - const SharedScratchHandoff dev0 = shared_scratch_claim_event(markers, /*device_id=*/0, std::ref(events)); - ASSERT_EQ(shared_scratch_mark_in_flight(markers, 0), dev0.event); + SharedScratchDevice& dev0 = pool.get(0); + SharedScratchDevice& dev1 = pool.get(1); + const SharedScratchHandoff first = shared_scratch_claim_event(dev0, std::ref(events)); + ASSERT_EQ(shared_scratch_mark_in_flight(dev0), first.event); // Device 1 has its own buffer, so device 0's enqueue is nothing for it to wait // on, and it gets its own event. - const SharedScratchHandoff dev1 = shared_scratch_claim_event(markers, /*device_id=*/1, std::ref(events)); - EXPECT_FALSE(dev1.needs_wait); - EXPECT_NE(dev1.event, dev0.event); + const SharedScratchHandoff second = shared_scratch_claim_event(dev1, std::ref(events)); + EXPECT_FALSE(second.needs_wait); + EXPECT_NE(second.event, first.event); EXPECT_EQ(events.created, 2); // Marking device 1 does not make device 0 stop waiting, or the other way round. - ASSERT_EQ(shared_scratch_mark_in_flight(markers, 1), dev1.event); - EXPECT_TRUE(shared_scratch_claim_event(markers, 0, std::ref(events)).needs_wait); - EXPECT_TRUE(shared_scratch_claim_event(markers, 1, std::ref(events)).needs_wait); + ASSERT_EQ(shared_scratch_mark_in_flight(dev1), second.event); + EXPECT_TRUE(shared_scratch_claim_event(dev0, std::ref(events)).needs_wait); + EXPECT_TRUE(shared_scratch_claim_event(dev1, std::ref(events)).needs_wait); } TEST(SharedScratchHandoffTest, EventCreationFailureIsReportedAndRetried) { - Markers markers; + SharedScratchDevice dev; FakeEventFactory events; events.fail_next = true; - const SharedScratchHandoff failed = shared_scratch_claim_event(markers, 0, std::ref(events)); + const SharedScratchHandoff failed = shared_scratch_claim_event(dev, std::ref(events)); EXPECT_EQ(failed.event, nullptr); EXPECT_FALSE(failed.needs_wait); // The failure leaves nothing behind, so the next call tries again and succeeds // rather than serving an unusable slot for the rest of the process. - const SharedScratchHandoff retried = shared_scratch_claim_event(markers, 0, std::ref(events)); + const SharedScratchHandoff retried = shared_scratch_claim_event(dev, std::ref(events)); EXPECT_NE(retried.event, nullptr); EXPECT_FALSE(retried.needs_wait); EXPECT_EQ(events.created, 1); } TEST(SharedScratchHandoffTest, ASlotWithNoEventIsNotMarked) { - Markers markers; + SharedScratchDevice dev; FakeEventFactory events; // Nothing can be recorded without an event, so nothing is claimed to have been. - EXPECT_EQ(shared_scratch_mark_in_flight(markers, 0), nullptr); + EXPECT_EQ(shared_scratch_mark_in_flight(dev), nullptr); // Otherwise, once an event is finally created for the slot, the next caller // would wait on it believing an enqueue had been recorded on it that never was. - EXPECT_FALSE(shared_scratch_claim_event(markers, 0, std::ref(events)).needs_wait); + EXPECT_FALSE(shared_scratch_claim_event(dev, std::ref(events)).needs_wait); +} + +// --------------------------------------------------------------------------- +// The registry that owns one entry per device. +// --------------------------------------------------------------------------- + +TEST(SharedScratchPoolRegistry, KeepsAnIndependentBufferPerDevice) { + SharedScratchPool pool; + FakeAllocator a; + std::size_t out = 0; + + void* dev0 = call(pool.get(0), a, 2048, out); + void* dev1 = call(pool.get(1), a, 2048, out); + + EXPECT_NE(dev0, dev1); + EXPECT_EQ(a.alloc_count(), 2); + EXPECT_TRUE(a.released.empty()); + + // Growing device 1 must not touch device 0's buffer. + void* dev1_big = call(pool.get(1), a, 9000, out); + void* dev0_again = call(pool.get(0), a, 2048, out); + EXPECT_NE(dev1_big, dev1); + EXPECT_EQ(dev0_again, dev0); + ASSERT_EQ(a.released.size(), 1u); + EXPECT_EQ(a.released[0].first, dev1); +} + +TEST(SharedScratchPoolRegistry, HandsOutOneStableEntryPerDevice) { + SharedScratchPool pool; + + SharedScratchDevice* const seven = &pool.get(7); + EXPECT_EQ(&pool.get(7), seven); + EXPECT_NE(&pool.get(8), seven); + + // Callers keep using an entry after the registry's lock is dropped, and go on + // using it across their CUDA calls, so adding devices must not move it. + std::set distinct; + for (int id = 0; id < 512; ++id) { + distinct.insert(&pool.get(id)); + } + EXPECT_EQ(&pool.get(7), seven); + // Two devices must never land on one entry, or a claimant is handed another + // device's buffer as its own. A bounded or folded key space is a plausible way + // to write this registry and an invisible way to break it. + EXPECT_EQ(distinct.size(), 512u); +} + +TEST(SharedScratchPoolRegistry, AGrowthOnOneDeviceDoesNotBlockAClaimOnAnother) { + SharedScratchPool pool; + // One allocator per thread: the two claims share the registry and nothing else. + FakeAllocator zero; + FakeAllocator one; + + std::promise entered_alloc; + std::promise leave_alloc; + std::future entered = entered_alloc.get_future(); + std::shared_future leave = leave_alloc.get_future().share(); + + SharedScratchDevice& dev0 = pool.get(0); + std::thread grower([&] { + std::lock_guard lk(dev0.mu); + std::size_t out = 0; + shared_scratch_get_or_grow( + dev0, + 4096, + out, + [&](std::size_t bytes) { + entered_alloc.set_value(); + leave.wait(); + return zero.alloc(bytes); + }, + [&](void* p, cudaEvent_t wait_for) { zero.release(p, wait_for); }); + }); + // The cap matters as much as the wait: a growth that takes the reuse path never + // reaches its allocation, so nothing fires this promise and an uncapped wait + // would hang the harness rather than fail the test. + if (entered.wait_for(std::chrono::seconds(10)) != std::future_status::ready) { + leave_alloc.set_value(); + grower.join(); + FAIL() << "the growth on device 0 never reached its allocation"; + } + + // Device 0's growth is stalled inside its allocation with device 0's lock held. + // Without this the rest of the test would pass against any implementation. + if (dev0.mu.try_lock()) { + dev0.mu.unlock(); + ADD_FAILURE() << "device 0's lock was not held across its allocation"; + } + + auto claim = std::async(std::launch::async, [&] { + SharedScratchDevice& dev1 = pool.get(1); + std::lock_guard lk(dev1.mu); + std::size_t out = 0; + return shared_scratch_get_or_grow( + dev1, + 2048, + out, + [&](std::size_t bytes) { return one.alloc(bytes); }, + [&](void* p, cudaEvent_t wait_for) { one.release(p, wait_for); }); + }); + const bool served = claim.wait_for(std::chrono::seconds(10)) == std::future_status::ready; + + leave_alloc.set_value(); + grower.join(); + + ASSERT_TRUE(served) << "a claim on device 1 waited for a growth on device 0"; + EXPECT_NE(claim.get(), nullptr); + EXPECT_EQ(one.alloc_count(), 1); +} + +TEST(SharedScratchPoolRegistry, ConcurrentLookupsKeepTheRegistryIntact) { + // Every other test reaches the registry from one thread at a time, so the + // registry's own lock is the one mechanism here that nothing else exercises: + // without this test it can be deleted outright and the suite stays green. + // + // An unsynchronized std::unordered_map mutated from several threads has no + // defined behaviour, so this cannot assert on a specific corruption. It + // hammers the lookup and then asks the two questions the corruption answers + // wrongly: is every id still where the race left it, and did any two ids land + // on one entry. Each round is an independent chance to observe that; the + // rounds are what make a miss unlikely rather than the assertions. + constexpr int kThreads = 4; + constexpr int kPerThread = 4000; + constexpr int kRounds = 8; + + for (int round = 0; round < kRounds; ++round) { + SharedScratchPool pool; + std::vector> seen(kThreads); + std::atomic ready{0}; + std::atomic go{false}; + + std::vector threads; + threads.reserve(kThreads); + for (int t = 0; t < kThreads; ++t) { + threads.emplace_back([&, t] { + std::vector mine; + mine.reserve(kPerThread); + // Rehashing is what corrupts an unsynchronized map, and it happens on a + // handful of the inserts in a round, so the threads have to be inside + // their loops at the same time. + ready.fetch_add(1); + while (!go.load()) { + } + for (int i = 0; i < kPerThread; ++i) { + mine.push_back(&pool.get(t * kPerThread + i)); + } + seen[t] = std::move(mine); + }); + } + while (ready.load() < kThreads) { + } + go.store(true); + for (std::thread& t : threads) { + t.join(); + } + + std::set distinct; + for (int t = 0; t < kThreads; ++t) { + ASSERT_EQ(seen[t].size(), static_cast(kPerThread)); + for (int i = 0; i < kPerThread; ++i) { + const int id = t * kPerThread + i; + ASSERT_EQ(&pool.get(id), seen[t][i]) << "device " << id << " in round " << round; + distinct.insert(seen[t][i]); + } + } + ASSERT_EQ(distinct.size(), static_cast(kThreads * kPerThread)) << "round " << round; + } } } // namespace From c6795bb15e8c778bfd7542e908179cd27105b98f Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Tue, 1 Sep 2026 11:43:01 -0700 Subject: [PATCH 3/3] fix(executorch): hold the scratch pool's device lock across the enqueue Addresses `:267` and `:314` of the 2026-08-28 review of the shared per-device activation-scratch pool, and the 08-31 follow-up `:316`. The third 08-28 finding, `:164` -- the option cannot be turned on from Python or from a `.pte` -- is deferred. It was flagged as a non-blocker, and closing it means giving the option a load-time runtime spec and a compile-spec fallback the way `weight_streaming_budget` has, which nothing here builds. **Two pooled engines running at once on one device gave wrong output, silently.** `get_or_grow_shared_scratch` dropped the device lock when it returned, so `setDeviceMemoryV2`, `enqueueV3` and the record of that enqueue on the marker's event all ran unlocked. A second thread claiming inside that window was handed the same buffer and told to wait on the enqueue *before* the one now in flight, so nothing ordered the two and both wrote the same scratch. Reproduced on two real engines: one output wrong on every trial, with no CUDA error and no TensorRT error. It needs no growth and no free, so it is the ordinary state once the pool has settled -- unlike the growth hazard the comments and the README did warn about. `execute()` now holds the claim from `get_or_grow_shared_scratch` through `setDeviceMemoryV2`, `enqueueV3` and `mark_shared_scratch_in_flight`, as a `SharedScratchClaim` whose destructor covers the early returns in between. `enqueueV3` is already called under the per-handle `EngineHandle::mu` here, and `core/runtime/execute_engine.cpp` brackets its own enqueue with `compiled_engine->mu` and states that the other `IExecutionContext` calls belong in that scope, so this is a narrower instance of a pattern the runtime already relies on. It nests inside `EngineHandle::mu` and is never taken the other way round. Overlapping `execute()` calls on two pooled handles on one device are therefore now safe -- serialized at submission rather than concurrent. The README and the installed `TensorRTBackend.h` told the caller to stagger them; both now say the backend does it, and that the pool costs the parallelism between them. Holding the lock across `enqueueV3` is cheap here, including for the case the TensorRT header warns about ("If the Engine is streaming weights, enqueueV3 will become synchronous"). Measured on TensorRT 11.2.1.2 and an A100, on a 12-layer engine with 768 MiB of streamable weights, `enqueueV3` stays a submission: | engine | `enqueueV3` | enqueue + drain | |---|---|---| | no weight streaming | 0.079 ms | 0.57 ms | | weight streaming, budget = streamable size (dormant) | 0.079 ms | 0.57 ms | | weight streaming, budget = half | 0.292 ms | 42.4 ms | | weight streaming, budget = 0 | 0.439 ms | 75.7 ms | At budget 0 the enqueue is 0.6% of the inference, and a second context's `enqueueV3` issued while that 75 ms streamed inference was still in flight returned in 0.267 ms. So the lock is held for a submission, not for an inference. The warning still stands in the header, so a different TensorRT version or platform could differ; the pool remains opt-in. `TwoThreadsRunningPooledEnginesOnOneDeviceKeepTheirOwnOutputs` covers the fix: two pooled engines on one device, 60 runs each from two threads on two non-blocking streams, released into their submission together and compared byte-for-byte against what each engine produces with private scratch. Against a build that restores the pre-fix lock scope it reports 118 to 120 of the 120 runs wrong; against this one, 0. It costs 1.8 s. Without the two threads lined up on each submission the host copies around each run serialize them and the same mutant loses only 2 runs of 120, so the rendezvous is what makes the test discriminate. **The growth's `cudaFree` no longer runs under the device lock, and no longer runs at all after a failed wait.** `cudaFree` performs a device-wide synchronization, so swapping the old `cudaDeviceSynchronize` for an event wait did not remove the device-wide wait: measured on an H100 with unrelated work queued on another stream, the event wait returned in under a millisecond and the `cudaFree` took about 1.5 s, matching the queued work -- with the device lock held, so an unrelated claim on that device waited behind it too. Reproduced on an A100 with 1.5 s of unrelated work on a second stream and the handoff event already complete: event wait 0.004 ms, `cudaFree` 1490 ms, over two trials. `shared_scratch_get_or_grow` now reports the displaced buffer through a `RetiredScratch` out-parameter instead of freeing it through a callback. The host wait on the marker's event stays under the lock, because the caller records its own enqueue on that same event before it unlocks and a wait deferred past that point would block on it. The free moves to `SharedScratchClaim::release()`, after the unlock, and reports and clears a CUDA error of its own rather than discarding one: `cudaFree` synchronizes, so its return is often where an earlier asynchronous fault on the device first surfaces. Two consequences of moving the free, both documented in the README: the stall no longer blocks another engine on the device, and it now falls after the growing call's own enqueue, so that one `execute()` waits for its own engine work. A retire list was considered and rejected: deferring frees would make peak device memory the sum of every size the pool ever grew to rather than the maximum, which is the saving the feature exists for. If the `cudaEventSynchronize` before the free fails, the buffer is now leaked with an error logged and the sticky CUDA error cleared rather than freed, since that wait is the only thing keeping the free off a buffer an enqueue may still be reading. The cost is bounded because growth is: it happens only when an engine larger than every engine before it runs for the first time. On Glimmer's 212 engines the pool allocated once, at 7,353.1 MiB, and never grew, because the largest engine ran first. On Gemma's 62 it grew twice before settling, 131.0 -> 142.0 -> 554.0 MiB. The six-engine synthetic fixture, built with deliberately unequal engines, grows four times: 44 -> 76 -> 140 -> 268 MiB. **The growth path had no test.** Every pooled engine in `test_shared_scratch_backend` asked the pool for the same number of bytes, so the reuse branch always won. `ALargerEngineGrowsThePoolAndFreesTheBufferItReplaces` runs a four-times-larger engine after a smaller one (33,554,432 -> 134,217,728 bytes) and brackets the growth's device-memory cost from both sides: the lower bound fails if no growth happened, the upper bound fails if the buffer that was replaced was not freed. It also re-runs the smaller engine afterwards, since the growth freed the address that engine's context was last given. The test was checked against three mutations of the production code, each of which kills it: - never free the displaced buffer: the growth costs 134,217,728 bytes against a 117,440,512-byte upper bound. - never grow (always take the reuse branch): `enqueueV3` refuses the undersized buffer and `execute()` returns `InvalidState`. - install the pool buffer once per context instead of on every call: the smaller engine's next run hits CUDA error 700 on the freed address. Verified on TensorRT 11.2.1.2, CUDA 13, one A100. `test_shared_scratch_pool` 16/16, `test_shared_scratch_backend` 12/12, and 12/12 skipped with no CUDA device visible. The default-off path is unchanged: 9/9 runs (3 programs x 3 repetitions) byte-identical in stdout and canonicalized stderr against a binary built from `main`, with the negative control discriminating. The pool's own A/B is unmoved: 1656 -> 316 MB, 792 -> 316 MB and 1188 -> 372 MB, same outputs. --- .../executorch/SharedScratchPool.h | 55 ++-- .../executorch/TensorRTBackend.h | 11 +- cpp/src/torch_tensorrt/executorch/README.md | 36 ++- .../executorch/TensorRTBackend.cpp | 184 +++++++++--- .../test_shared_scratch_backend.cpp | 275 ++++++++++++++++-- .../executorch/test_shared_scratch_pool.cpp | 77 ++--- 6 files changed, 511 insertions(+), 127 deletions(-) diff --git a/cpp/include/torch_tensorrt/executorch/SharedScratchPool.h b/cpp/include/torch_tensorrt/executorch/SharedScratchPool.h index 5f43fdf3ce..9e72e218bc 100644 --- a/cpp/include/torch_tensorrt/executorch/SharedScratchPool.h +++ b/cpp/include/torch_tensorrt/executorch/SharedScratchPool.h @@ -49,11 +49,13 @@ struct SharedScratchHandoff { // One device's shared scratch buffer and the marker ordering its handoff, behind // the lock that covers both. // -// A claimant holds `mu` from the wait on the previous enqueue through the choice -// of buffer, so it cannot be handed a buffer another claimant is midway through -// replacing, and cannot record its own enqueue against a marker that has since -// moved on. `mu` covers one device, so a growth holds no lock a claim on another -// device has to acquire. +// A claimant holds `mu` from the wait on the previous enqueue through its own +// enqueue and the record of that enqueue on the marker. Holding it that far is +// what makes the marker a complete account of who is using the buffer. Anything +// less leaves an enqueue live in a window the marker does not cover, and a +// claimant entering that window is handed the same buffer with nothing ordering +// the two. `mu` covers one device, so a growth holds no lock a claim on +// another device has to acquire. struct SharedScratchDevice { std::mutex mu; void* buffer = nullptr; @@ -118,6 +120,29 @@ inline cudaEvent_t shared_scratch_mark_in_flight(SharedScratchDevice& dev) { return dev.marker.event; } +// A buffer a growth replaced, handed back for the caller to dispose of. +// +// A non-null `wait_for` is the marker's event, on which an enqueue that may still +// be reading and writing `buffer` has been recorded; the caller must wait for +// that event on the host before it frees, and must do so while it still holds +// `dev.mu`. Once the lock is dropped the next claimant records its own enqueue on +// the same event, and a wait made then would block on work that never touched +// this buffer. A null `wait_for` means nothing was ever recorded against it. +// +// One event covers every enqueue the buffer ever served, but only because each of +// them claims the handoff before enqueueing -- which orders its stream after the +// event -- and records on the event afterwards, so the latest recording completes +// only once all the earlier ones have. An enqueue that reaches the buffer without +// doing both is covered by no wait here. +// +// The free itself belongs outside `dev.mu`: on CUDA it is a device-wide +// synchronization, so performing it under the lock makes an unrelated claim on +// this device wait for every stream on it. +struct RetiredScratch { + void* buffer = nullptr; + cudaEvent_t wait_for = nullptr; +}; + // Bookkeeping for a device's scratch buffer, which grows monotonically to the // largest requested size. Call with `dev.mu` held. // @@ -125,22 +150,17 @@ inline cudaEvent_t shared_scratch_mark_in_flight(SharedScratchDevice& dev) { // Allocating before releasing is what makes that true, and it costs peak // residency: while the buffer grows, the old and the new one are both resident. // -// `release(old, wait_for)` frees `old`. A non-null `wait_for` is the marker's -// event, on which an enqueue that may still be reading and writing `old` has been -// recorded; the release must wait for that event on the host before freeing. One -// event covers every enqueue the buffer ever served, but only because each of -// them claims the handoff before enqueueing -- which orders its stream after the -// event -- and records on the event afterwards, so the latest recording completes -// only once all the earlier ones have. An enqueue that reaches the buffer without -// doing both is covered by no wait here. A null `wait_for` means nothing was ever -// recorded against this buffer, so there is nothing to wait for. -template +// A growth reports the buffer it displaced through `out_retired`; see +// RetiredScratch for what the caller owes it. Nothing is freed here, so a caller +// that ignores `out_retired` leaks rather than frees a buffer an enqueue may +// still be using. +template void* shared_scratch_get_or_grow( SharedScratchDevice& dev, std::size_t need, std::size_t& out_size, Alloc alloc, - Release release) { + RetiredScratch& out_retired) { if (dev.buffer != nullptr && dev.capacity >= need) { out_size = dev.capacity; return dev.buffer; @@ -150,7 +170,8 @@ void* shared_scratch_get_or_grow( return nullptr; } if (dev.buffer != nullptr) { - release(dev.buffer, dev.marker.pending ? dev.marker.event : nullptr); + out_retired.buffer = dev.buffer; + out_retired.wait_for = dev.marker.pending ? dev.marker.event : nullptr; } dev.buffer = p; dev.capacity = need; diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index e164bd01f9..0a9571f22b 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -110,11 +110,12 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { // past return, order any other stream against this one, and synchronize the stream // before reading device-resident outputs. The selected stream must be on the engine's // device, and calls on one handle must not overlap each other or its destruction. - // The shared activation scratch pool (kSharedActivationScratchKey) widens that - // across handles: one buffer per device backs every context created while the - // option was on, so calls on two such handles on one device must not overlap - // either. A handle whose context was created while the option was off keeps its - // own scratch and is outside that rule. + // With the shared activation scratch pool (kSharedActivationScratchKey) one + // buffer per device backs every context created while the option was on. Calls + // on two such handles on one device may overlap: a per-device lock held across + // the enqueue serializes them, so they do not run concurrently on the device. A + // handle whose context was created while the option was off keeps its own + // scratch and is not subject to this. // Note that other CUDA delegates sharing the same guard may instead synchronize before // returning, so do not assume results are ready on return from this one. ::executorch::runtime::Error execute( diff --git a/cpp/src/torch_tensorrt/executorch/README.md b/cpp/src/torch_tensorrt/executorch/README.md index b5bd5cc526..58b8f55255 100644 --- a/cpp/src/torch_tensorrt/executorch/README.md +++ b/cpp/src/torch_tensorrt/executorch/README.md @@ -87,14 +87,15 @@ the removed `CudaStreamGuard`: - With no guard active, the backend falls back to `cudaStreamPerThread`. - With the `use_shared_activation_scratch` backend option enabled, one buffer per device backs the activation scratch of every execution context created - while it was on, so an enqueue against that buffer must not overlap another - one. The backend orders consecutive enqueues itself, whether they run on one - stream or on two. What it cannot order is two `execute()` calls submitted - concurrently on one device: the caller must submit them one at a time, whether - or not they share a stream. Submitting them concurrently risks one of them - growing the pool and freeing the buffer the other's enqueue is still reading - and writing, not merely reordering them. Contexts created while the option was - off keep their own scratch and are unaffected. + while it was on, so no two enqueues against it may overlap. The backend + enforces this itself: it holds a per-device lock from the claim on the buffer + through the enqueue and the completion event recorded on it, so two + `execute()` calls on one device are serialized at submission and the second's + stream waits on the first's enqueue. They may run on one stream or on two, and + they may be submitted concurrently from two threads — but they will not run + concurrently on the device, so the pool costs the parallelism between them. + Contexts created while the option was off keep their own scratch and are + unaffected. - The reference-runner smoke test runs inference inside a caller-stream guard on the discrete-GPU CI configuration, where all inputs and outputs are host-backed and therefore take the synchronized staging path. CI separately asserts that the @@ -138,10 +139,21 @@ the backend archive gets. N per-engine copies collapse to one, so the reclaimed memory is the sum of the N requirements less the largest of them. Set the option before loading the methods whose engines should use the pool, and read the `use_shared_activation_scratch` -bullet of the caller-stream contract above first: the pool carries an ordering -obligation the backend cannot discharge for you. The buffer is never released, so -the device keeps the largest scratch it was ever asked for until the process -exits. +bullet of the caller-stream contract above: engines sharing a buffer do not run +concurrently on the device. The pool never returns memory to the +device, so the largest scratch it was ever asked for stays allocated until the +process exits. + +The buffer grows when an engine asks for more than every engine before it did, +and a growth is not free. It frees the buffer it replaces, and `cudaFree` waits +for everything queued on the device, not only for the enqueues that used that +buffer, so it can stall for far longer than the event wait that precedes it. The +backend keeps that stall out from under the per-device lock, so it does not hold +up another engine on the device, but it does fall after the growing call's own +enqueue, so that one `execute()` waits for its own engine work too. A growth +happens only on an engine's first run and only for an engine larger than every +engine before it, so loading the largest engine first reduces the pool to a +single allocation. How much any one engine asks for is fixed when it is built, not when it runs. The builder's `kRUNTIME_ACTIVATION_RESIZE_10_10` preview feature makes an engine diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index a362ea7efc..3aee341b14 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -236,35 +236,114 @@ bool is_cuda_accessible_ptr(const void* ptr) { // // ORDERING: a context reads and writes its scratch for the whole enqueue, which // can still be in flight when execute() returns, so two enqueues must never hold -// this buffer at the same time. +// this buffer at the same time. A device's lock is what enforces that -- see +// SharedScratchClaim -- and it is held from the claim through the enqueue and the +// record of it, so two execute() calls on one device are serialized at +// submission. The lock does not couple two +// devices: each carries its own, and no CUDA call is made under the one lock the +// registry itself holds. // -// What the pool's event handoff does NOT cover is concurrent execute() on one -// device: a device's lock is released before the enqueue is submitted, so an -// enqueue is live for a window before the event carries it, and a second thread -// claiming inside that window is told to wait for the enqueue before it. Such a -// claimant can grow the pool and free the buffer the first thread's enqueue is -// still reading and writing. The requirement is therefore that the enqueues -// drawing on a device's buffer are submitted one at a time, but they need not -// share a stream. That is why the pool is opt-in. The pool's locking does not -// couple two devices: each carries its own lock, and no CUDA call is made under -// the lock that finds it. -// -// The buffers and the events are intentionally never freed. Nothing here runs a -// CUDA call at process exit, which also keeps the pool clear of teardown-order -// hazards against anything else holding device memory. +// The buffers and the events are intentionally never freed at teardown. Nothing +// here runs a CUDA call at process exit, which keeps the pool clear of +// teardown-order hazards against anything else holding device memory. SharedScratchPool scratch_pool; +// A caller's hold on one device's shared scratch: the device lock, plus the +// buffer a growth displaced, freed once that lock is dropped. +// +// The lock spans the enqueue, not just the choice of buffer. A claimant that +// released it as soon as it had a buffer would leave its enqueue live for a +// window the marker's event does not yet cover, and a second claimant entering +// that window is handed the same buffer and told to wait for the enqueue before +// it -- so nothing orders the two and both write the same scratch. The failure +// is silent: wrong output, no CUDA error, no TensorRT error. +// +// This lock nests inside the per-handle EngineHandle::mu, which already spans +// the enqueue, and is never taken in the other order. +class SharedScratchClaim { + public: + SharedScratchClaim() = default; + SharedScratchClaim(const SharedScratchClaim&) = delete; + SharedScratchClaim& operator=(const SharedScratchClaim&) = delete; + ~SharedScratchClaim() { + release(); + } + + SharedScratchDevice& hold(int device_id) { + dev_ = &scratch_pool.get(device_id); + device_id_ = device_id; + lock_ = std::unique_lock(dev_->mu); + return *dev_; + } + + // Null until hold() runs and null again after release(): non-null exactly while + // this claim holds the device's lock. + SharedScratchDevice* device() const { + return dev_; + } + + int device_id() const { + return device_id_; + } + + // Takes ownership of a buffer a growth displaced, to be freed by release(). + void retire(void* buffer) { + retired_ = buffer; + } + + // Drops the lock, and the device pointer with it so device() cannot hand out a + // pointer this claim no longer holds the lock for. Then frees whatever a growth + // displaced -- after the unlock, because cudaFree waits for every stream on the + // device, which under the lock would stall the next claim on work unrelated to + // the pool. Outside it the stall is this caller's alone and falls after its own + // enqueue, so a growth makes that one execute() wait for its own engine work. + // + // Frees on the current device, which must still be the buffer's. + void release() { + if (lock_.owns_lock()) { + lock_.unlock(); + } + dev_ = nullptr; + if (retired_ != nullptr) { + // cudaFree synchronizes, so an earlier asynchronous fault on this device + // often surfaces here. Report and clear it, or it resurfaces under the + // name of the next CUDA call in execute(). + const cudaError_t err = cudaFree(retired_); + if (err != cudaSuccess) { + ET_LOG( + Error, + "TensorRTBackend::execute: freeing the shared activation scratch buffer that a pool growth replaced on device %d failed: %s", + device_id_, + cudaGetErrorString(err)); + cudaGetLastError(); // clear sticky error; the free is cleanup, so execute() continues + } + retired_ = nullptr; + } + } + + private: + SharedScratchDevice* dev_ = nullptr; + int device_id_ = -1; + std::unique_lock lock_; + void* retired_ = nullptr; +}; + // Sets out_ptr to a buffer of at least `need` bytes on `device_id` and out_size to // its capacity, with `stream` ordered after the enqueue that last used the buffer. -// The caller must call mark_shared_scratch_in_flight once it has submitted its own -// enqueue. +// Returns with `claim` holding the device's lock: the caller must submit its +// enqueue, call mark_shared_scratch_in_flight, and only then release the claim. // // Must be called with `device_id` already current: cudaEventCreateWithFlags, // cudaMalloc and cudaFree all act on the *current* device and nothing in here // sets it. -Error get_or_grow_shared_scratch(int device_id, size_t need, cudaStream_t stream, void*& out_ptr, size_t& out_size) { - SharedScratchDevice& dev = scratch_pool.get(device_id); - std::lock_guard lk(dev.mu); +Error get_or_grow_shared_scratch( + SharedScratchClaim& claim, + int device_id, + size_t need, + cudaStream_t stream, + void*& out_ptr, + size_t& out_size) { + SharedScratchDevice& dev = claim.hold(device_id); const SharedScratchHandoff handoff = shared_scratch_claim_event(dev, []() -> cudaEvent_t { cudaEvent_t event = nullptr; @@ -292,6 +371,7 @@ Error get_or_grow_shared_scratch(int device_id, size_t need, cudaStream_t stream } const bool first_buffer = dev.buffer == nullptr; + RetiredScratch retired; void* const buffer = shared_scratch_get_or_grow( dev, need, @@ -309,12 +389,7 @@ Error get_or_grow_shared_scratch(int device_id, size_t need, cudaStream_t stream bytes); return p; }, - [](void* old, cudaEvent_t wait_for) { - if (wait_for != nullptr) { - cudaEventSynchronize(wait_for); - } - cudaFree(old); - }); + retired); if (buffer == nullptr) { ET_LOG( Error, @@ -324,19 +399,46 @@ Error get_or_grow_shared_scratch(int device_id, size_t need, cudaStream_t stream return Error::MemoryAllocationFailed; } + if (retired.buffer != nullptr) { + // The wait runs here and the free runs at release(), because this caller + // records its own enqueue on the same event before it drops the lock: a wait + // deferred to sit beside the free would block on that enqueue too. + const cudaError_t err = retired.wait_for != nullptr ? cudaEventSynchronize(retired.wait_for) : cudaSuccess; + if (err == cudaSuccess) { + claim.retire(retired.buffer); + } else { + // This wait is the only thing keeping the free off a buffer an enqueue may + // still be reading, so a failed wait leaks it instead. Bounded: at most one + // buffer per growth, and growth is rare -- see SharedScratchClaim::release. + ET_LOG( + Error, + "TensorRTBackend::execute: waiting for the enqueue on the replaced shared activation scratch on device %d failed (%s); leaking that buffer rather than freeing it under a live enqueue", + device_id, + cudaGetErrorString(err)); + cudaGetLastError(); // clear sticky error; execute() continues regardless + } + } + out_ptr = buffer; return Error::Ok; } -// Records the enqueue now in flight on `stream` against `device_id`'s shared -// scratch, so the next call to get_or_grow_shared_scratch waits for it. -Error mark_shared_scratch_in_flight(int device_id, cudaStream_t stream) { - SharedScratchDevice& dev = scratch_pool.get(device_id); - std::lock_guard lk(dev.mu); +// Records the enqueue now in flight on `stream` against the claimed device's +// shared scratch, so the next call to get_or_grow_shared_scratch waits for it. +// Call with `claim` still holding the device's lock. +Error mark_shared_scratch_in_flight(SharedScratchClaim& claim, cudaStream_t stream) { + SharedScratchDevice* const dev = claim.device(); + if (dev == nullptr) { + ET_LOG(Error, "TensorRTBackend::execute: no shared activation scratch claim to record an enqueue against"); + return Error::Internal; + } - const cudaEvent_t event = shared_scratch_mark_in_flight(dev); + const cudaEvent_t event = shared_scratch_mark_in_flight(*dev); if (event == nullptr) { - ET_LOG(Error, "TensorRTBackend::execute: shared activation scratch on device %d has no handoff event", device_id); + ET_LOG( + Error, + "TensorRTBackend::execute: shared activation scratch on device %d has no handoff event", + claim.device_id()); return Error::Internal; } const cudaError_t err = cudaEventRecord(event, stream); @@ -1070,13 +1172,20 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // against. A failed query carried on would instead leave the context enqueueing // against whatever buffer it last held, because setDeviceMemoryV2(nullptr, 0) // is rejected and returns nothing to test. + // + // The claim holds the device's pool lock from here through the record of the + // enqueue below; see SharedScratchClaim for why it spans that far. Every return + // in between drops it through the destructor, which runs ahead of the device + // restore above, so its free lands on the right device. + SharedScratchClaim scratch_claim; bool scratch_from_pool = false; if (engine->shared_scratch) { const size_t need = ctx->updateDeviceMemorySizeForShapes(); if (need > 0) { void* pool = nullptr; size_t pool_size = 0; - const Error scratch_err = get_or_grow_shared_scratch(engine->device_id, need, stream, pool, pool_size); + const Error scratch_err = + get_or_grow_shared_scratch(scratch_claim, engine->device_id, need, stream, pool, pool_size); if (scratch_err != Error::Ok) { return scratch_err; } @@ -1106,7 +1215,7 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // Pairs with get_or_grow_shared_scratch: the next claimant waits on this event. if (scratch_from_pool) { - const Error mark_err = mark_shared_scratch_in_flight(engine->device_id, stream); + const Error mark_err = mark_shared_scratch_in_flight(scratch_claim, stream); if (mark_err != Error::Ok) { // Nothing will wait for this enqueue, so wait for it here instead of // leaving the next user of the buffer to overwrite live scratch. @@ -1115,6 +1224,11 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* return mark_err; } } + // The enqueue is now on the marker's event, so the device's pool is safe to + // hand to the next claimant. Released here rather than at the end of the + // function so the rest of execute() -- the aliased reflects, the D2H copies and + // their synchronizations -- does not hold up another engine on this device. + scratch_claim.release(); // Caller-owned KV: reflect each engine in-place update into its delegate output // EValue (D2D on the same stream, after the engine work). diff --git a/tests/cpp/executorch/test_shared_scratch_backend.cpp b/tests/cpp/executorch/test_shared_scratch_backend.cpp index 54842b8db0..c021e65e46 100644 --- a/tests/cpp/executorch/test_shared_scratch_backend.cpp +++ b/tests/cpp/executorch/test_shared_scratch_backend.cpp @@ -7,9 +7,11 @@ // Exercises the shared activation-scratch pool through the delegate that uses // it: the runtime option that turns it on, the per-engine capture of that -// option, and the single-threaded pooled execute() path -- the kUSER_MANAGED -// context, the updateDeviceMemorySizeForShapes/setDeviceMemoryV2 pair, and the -// enqueue handoff between two caller streams. +// option, and the pooled execute() path -- the kUSER_MANAGED context, the +// updateDeviceMemorySizeForShapes/setDeviceMemoryV2 pair, the enqueue handoff +// between two caller streams, the growth a larger engine forces on a pool a +// smaller one already allocated, and two threads submitting against one pooled +// buffer at once. // // The TensorRT engine is built here rather than loaded from a .pte so the target // carries no exported artifact, at the cost of a few seconds of builder time. @@ -77,6 +79,12 @@ constexpr int kCols = 2048; constexpr std::size_t kElems = static_cast(kRows) * static_cast(kCols); constexpr std::size_t kBytes = kElems * sizeof(float); +// A second scratch-needing engine, four times the elements of the one above, so +// loading it after that one drives the pool's growth path. Every other engine in +// this file asks for the same size, which is why nothing else reaches it. +constexpr int kBigRows = 4096; +constexpr int kBigCols = 4096; + // Engines loaded together in the memory test. Four is enough for the private // case to cost 4x the scratch and the pooled case 1x. constexpr int kEngineCount = 4; @@ -174,7 +182,7 @@ bool add_scratch_free_net(nvinfer1::INetworkDefinition& network, nvinfer1::ITens return true; } -std::vector build_engine_blob(bool needs_scratch) { +std::vector build_engine_blob(bool needs_scratch, int rows = kRows, int cols = kCols) { static BuilderLogger logger; TRTUniquePtr builder(nvinfer1::createInferBuilder(logger)); @@ -186,7 +194,7 @@ std::vector build_engine_blob(bool needs_scratch) { return {}; } - nvinfer1::ITensor* input = network->addInput("input_0", nvinfer1::DataType::kFLOAT, nvinfer1::Dims3{1, kRows, kCols}); + nvinfer1::ITensor* input = network->addInput("input_0", nvinfer1::DataType::kFLOAT, nvinfer1::Dims3{1, rows, cols}); if (input == nullptr) { return {}; } @@ -200,7 +208,7 @@ std::vector build_engine_blob(bool needs_scratch) { return {}; } nvinfer1::IOptimizationProfile* profile = builder->createOptimizationProfile(); - const nvinfer1::Dims3 shape{1, kRows, kCols}; + const nvinfer1::Dims3 shape{1, rows, cols}; profile->setDimensions("input_0", nvinfer1::OptProfileSelector::kMIN, shape); profile->setDimensions("input_0", nvinfer1::OptProfileSelector::kOPT, shape); profile->setDimensions("input_0", nvinfer1::OptProfileSelector::kMAX, shape); @@ -231,7 +239,7 @@ std::vector build_engine_blob(bool needs_scratch) { // The activation scratch one context of the shared engine needs, read the way // execute() reads it. Zero if the engine could not be measured. -std::size_t measure_engine_scratch(const std::vector& blob) { +std::size_t measure_engine_scratch(const std::vector& blob, int rows = kRows, int cols = kCols) { static BuilderLogger logger; TensorRTBlobHeader header; if (!TensorRTBlobHeader::parse(blob.data(), blob.size(), header)) { @@ -252,7 +260,7 @@ std::size_t measure_engine_scratch(const std::vector& blob) { if (ctx == nullptr) { return 0; } - if (!ctx->setInputShape("input_0", nvinfer1::Dims3{1, kRows, kCols})) { + if (!ctx->setInputShape("input_0", nvinfer1::Dims3{1, rows, cols})) { return 0; } return ctx->updateDeviceMemorySizeForShapes(); @@ -306,16 +314,19 @@ class LoadedEngine { } // Loads the blob through the backend, capturing whatever the shared-scratch - // option is set to at this moment. - Error load(const std::vector& blob, std::uint32_t seed) { - std::vector host_in(kElems); - for (std::size_t i = 0; i < kElems; ++i) { + // option is set to at this moment. `rows`/`cols` must be the shape the blob was + // built for. + Error load(const std::vector& blob, std::uint32_t seed, int rows = kRows, int cols = kCols) { + rows_ = static_cast(rows); + cols_ = static_cast(cols); + std::vector host_in(elems()); + for (std::size_t i = 0; i < elems(); ++i) { host_in[i] = pattern(i, seed); } - if (cudaMalloc(&device_in_, kBytes) != cudaSuccess || cudaMalloc(&device_out_, kBytes) != cudaSuccess) { + if (cudaMalloc(&device_in_, bytes()) != cudaSuccess || cudaMalloc(&device_out_, bytes()) != cudaSuccess) { return Error::MemoryAllocationFailed; } - if (cudaMemcpy(device_in_, host_in.data(), kBytes, cudaMemcpyHostToDevice) != cudaSuccess) { + if (cudaMemcpy(device_in_, host_in.data(), bytes(), cudaMemcpyHostToDevice) != cudaSuccess) { return Error::Internal; } @@ -332,8 +343,8 @@ class LoadedEngine { } bool fill_output(float value) { - const std::vector host(kElems, value); - return cudaMemcpy(device_out_, host.data(), kBytes, cudaMemcpyHostToDevice) == cudaSuccess; + const std::vector host(elems(), value); + return cudaMemcpy(device_out_, host.data(), bytes(), cudaMemcpyHostToDevice) == cudaSuccess; } // Runs one inference on `stream`. Returns without waiting for the enqueue, @@ -341,8 +352,8 @@ class LoadedEngine { Error run(cudaStream_t stream) { // Separate arrays: execute() resizes the output tensor to the shape TensorRT // inferred, which writes through whichever array that tensor was given. - SizesType in_sizes[3] = {1, kRows, kCols}; - SizesType out_sizes[3] = {1, kRows, kCols}; + SizesType in_sizes[3] = {1, rows_, cols_}; + SizesType out_sizes[3] = {1, rows_, cols_}; ::executorch::aten::TensorImpl in_impl(ScalarType::Float, 3, in_sizes, device_in_); ::executorch::aten::TensorImpl out_impl(ScalarType::Float, 3, out_sizes, device_out_); ::executorch::aten::Tensor in_tensor(&in_impl); @@ -357,8 +368,8 @@ class LoadedEngine { } std::vector read_output() const { - std::vector host_out(kElems); - if (cudaMemcpy(host_out.data(), device_out_, kBytes, cudaMemcpyDeviceToHost) != cudaSuccess) { + std::vector host_out(elems()); + if (cudaMemcpy(host_out.data(), device_out_, bytes(), cudaMemcpyDeviceToHost) != cudaSuccess) { host_out.clear(); } return host_out; @@ -368,6 +379,14 @@ class LoadedEngine { return static_cast(handle_); } + std::size_t elems() const { + return static_cast(rows_) * static_cast(cols_); + } + + std::size_t bytes() const { + return elems() * sizeof(float); + } + private: // EngineHandle is placement-newed into this arena by init(), and the arena is // never reset, so it only has to hold one instance. @@ -379,6 +398,8 @@ class LoadedEngine { DelegateHandle* handle_ = nullptr; void* device_in_ = nullptr; void* device_out_ = nullptr; + SizesType rows_ = kRows; + SizesType cols_ = kCols; }; std::size_t device_bytes_in_use() { @@ -415,10 +436,12 @@ class SharedScratchBackendTest : public ::testing::Test { } blob_ = build_engine_blob(true); scratch_free_blob_ = build_engine_blob(false); - if (blob_.empty() || scratch_free_blob_.empty()) { + big_blob_ = build_engine_blob(true, kBigRows, kBigCols); + if (blob_.empty() || scratch_free_blob_.empty() || big_blob_.empty()) { return; } scratch_bytes_ = measure_engine_scratch(blob_); + big_scratch_bytes_ = measure_engine_scratch(big_blob_, kBigRows, kBigCols); engine_bytes_ = engine_scratch_requirement(blob_); scratch_free_engine_bytes_ = engine_scratch_requirement(scratch_free_blob_); } @@ -430,6 +453,7 @@ class SharedScratchBackendTest : public ::testing::Test { } ASSERT_FALSE(blob_.empty()) << "TensorRT could not build the fixture engine"; ASSERT_FALSE(scratch_free_blob_.empty()) << "TensorRT could not build the scratch-free fixture engine"; + ASSERT_FALSE(big_blob_.empty()) << "TensorRT could not build the larger fixture engine"; ASSERT_EQ(set_shared_scratch(backend_, false), Error::Ok); } @@ -445,17 +469,25 @@ class SharedScratchBackendTest : public ::testing::Test { return scratch_free_blob_; } + const std::vector& big_blob() const { + return big_blob_; + } + TensorRTBackend backend_; static std::vector blob_; static std::vector scratch_free_blob_; + static std::vector big_blob_; static std::size_t scratch_bytes_; + static std::size_t big_scratch_bytes_; static std::int64_t engine_bytes_; static std::int64_t scratch_free_engine_bytes_; }; std::vector SharedScratchBackendTest::blob_; std::vector SharedScratchBackendTest::scratch_free_blob_; +std::vector SharedScratchBackendTest::big_blob_; std::size_t SharedScratchBackendTest::scratch_bytes_ = 0; +std::size_t SharedScratchBackendTest::big_scratch_bytes_ = 0; std::int64_t SharedScratchBackendTest::engine_bytes_ = -1; std::int64_t SharedScratchBackendTest::scratch_free_engine_bytes_ = -1; @@ -640,6 +672,99 @@ TEST_F(SharedScratchBackendTest, PooledEnginesShareOneActivationScratchAllocatio << " pooled, against " << scratch_bytes_ << " bytes of scratch each"; } +// --------------------------------------------------------------------------- +// Growing the pool +// --------------------------------------------------------------------------- + +// Runs a four-times-larger engine after a smaller one to reach the growth path, +// which nothing else in this file does. +// +// The bounds cover the second allocation and the free of the buffer it replaces, +// not the host wait before that free: cudaFree synchronizes device-wide anyway, +// so deleting the wait leaves this test green. The wait stays as the explicit +// guarantee rather than a reliance on cudaFree's implicit one. +// +// The lower bound also fails if an earlier test left the pool already large +// enough, which is how this test could otherwise pass vacuously. +TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItReplaces) { + ASSERT_GE(big_scratch_bytes_, scratch_bytes_ + kMinMeasurableScratch) + << "the two fixture engines ask for " << scratch_bytes_ << " and " << big_scratch_bytes_ + << " bytes of activation scratch, too close for the growth to be measurable"; + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + + // Private-scratch references for both engines, and the run that pays the larger + // engine's one-time TensorRT and CUDA module costs so they land outside the + // measurement below. + std::vector small_expected; + std::vector big_expected; + { + LoadedEngine small_priv; + LoadedEngine big_priv; + ASSERT_EQ(small_priv.load(blob(), 15), Error::Ok); + ASSERT_EQ(big_priv.load(big_blob(), 16, kBigRows, kBigCols), Error::Ok); + ASSERT_FALSE(small_priv.handle()->shared_scratch); + ASSERT_FALSE(big_priv.handle()->shared_scratch); + ASSERT_EQ(small_priv.run(stream), Error::Ok); + ASSERT_EQ(big_priv.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + small_expected = small_priv.read_output(); + big_expected = big_priv.read_output(); + } + + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + LoadedEngine small; + ASSERT_EQ(small.load(blob(), 15), Error::Ok); + ASSERT_TRUE(small.handle()->shared_scratch); + ASSERT_EQ(small.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + // Loaded before the measurement starts: its weights and its I/O are not part of + // what the growth costs. + LoadedEngine big; + ASSERT_EQ(big.load(big_blob(), 16, kBigRows, kBigCols), Error::Ok); + ASSERT_TRUE(big.handle()->shared_scratch); + + const std::size_t before = device_bytes_in_use(); + ASSERT_EQ(big.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::size_t after = device_bytes_in_use(); + ASSERT_GE(after, before) << "device-wide memory in use fell across the growth, so something outside this test is " + "releasing memory on this device"; + const std::size_t growth_cost = after - before; + const std::size_t difference = big_scratch_bytes_ - scratch_bytes_; + + // The pool must still serve the smaller engine after the growth moved the + // buffer: its context holds the address it was given on its previous call, and + // that address has been freed. + ASSERT_TRUE(small.fill_output(kSentinel)); + ASSERT_EQ(small.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::vector small_actual = small.read_output(); + const std::vector big_actual = big.read_output(); + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); + + EXPECT_GE(growth_cost, difference / 2) << "the larger engine cost " << growth_cost << " bytes against a " + << difference + << "-byte difference in requirement, so the pool did not grow for it"; + EXPECT_LE(growth_cost, difference + scratch_bytes_ / 2) + << "the larger engine cost " << growth_cost << " bytes, about the whole " << big_scratch_bytes_ + << "-byte buffer rather than the " << difference << "-byte difference, so the buffer it replaced was not freed"; + + ASSERT_EQ(big_expected.size(), big_actual.size()); + ASSERT_FALSE(big_expected.empty()); + EXPECT_EQ(std::memcmp(big_expected.data(), big_actual.data(), big_expected.size() * sizeof(float)), 0) + << "the engine that grew the pool did not produce what it produces with its own scratch"; + + ASSERT_EQ(small_expected.size(), kElems); + ASSERT_EQ(small_actual.size(), kElems); + EXPECT_NE(small_expected[0], kSentinel) << "the reference output is the sentinel, so a skipped enqueue would pass"; + EXPECT_EQ(std::memcmp(small_expected.data(), small_actual.data(), kBytes), 0) + << "the smaller engine stopped producing its own output once the growth moved the shared buffer"; +} + // --------------------------------------------------------------------------- // An engine that needs no activation scratch // --------------------------------------------------------------------------- @@ -846,6 +971,114 @@ TEST_F(SharedScratchBackendTest, ASecondPooledEnqueueWaitsForTheFirstOnAnotherSt ASSERT_EQ(cudaStreamDestroy(second_stream), cudaSuccess); } +// --------------------------------------------------------------------------- +// The pooled path under two concurrent callers +// --------------------------------------------------------------------------- + +// The window a dropped lock would leave open is a few microseconds wide, so one +// pair of runs would find it only by luck. At this count a build that leaves the +// window open loses most of the runs, and the test costs about two seconds. +constexpr int kConcurrentRunsPerThread = 60; + +// Two pooled engines on one device, submitted from two threads on two streams, +// with nothing but the backend ordering them. Both are backed by the same buffer, +// so a claim that ends before the enqueue is recorded hands a second caller the +// same scratch with nothing ordering the two -- silently wrong output, no CUDA +// error, no TensorRT error. Each thread compares byte-for-byte against what its own +// engine produces with private scratch. +TEST_F(SharedScratchBackendTest, TwoThreadsRunningPooledEnginesOnOneDeviceKeepTheirOwnOutputs) { + cudaStream_t reference_stream = nullptr; + ASSERT_EQ(cudaStreamCreate(&reference_stream), cudaSuccess); + + std::vector first_expected; + std::vector second_expected; + { + LoadedEngine first_priv; + LoadedEngine second_priv; + ASSERT_EQ(first_priv.load(blob(), 17), Error::Ok); + ASSERT_EQ(second_priv.load(blob(), 18), Error::Ok); + ASSERT_FALSE(first_priv.handle()->shared_scratch); + ASSERT_FALSE(second_priv.handle()->shared_scratch); + ASSERT_EQ(first_priv.run(reference_stream), Error::Ok); + ASSERT_EQ(second_priv.run(reference_stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(reference_stream), cudaSuccess); + first_expected = first_priv.read_output(); + second_expected = second_priv.read_output(); + } + ASSERT_EQ(cudaStreamDestroy(reference_stream), cudaSuccess); + + ASSERT_EQ(first_expected.size(), kElems); + ASSERT_EQ(second_expected.size(), kElems); + // Two engines producing the same bytes would let each thread pass on the other + // one's output, which is the outcome this test exists to catch. + ASSERT_NE(std::memcmp(first_expected.data(), second_expected.data(), kBytes), 0) + << "the two engines were given different inputs but produced the same output"; + ASSERT_NE(first_expected[0], kSentinel) << "the reference output is the sentinel, so a skipped enqueue would pass"; + ASSERT_NE(second_expected[0], kSentinel) << "the reference output is the sentinel, so a skipped enqueue would pass"; + + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + LoadedEngine first; + LoadedEngine second; + ASSERT_EQ(first.load(blob(), 17), Error::Ok); + ASSERT_EQ(second.load(blob(), 18), Error::Ok); + ASSERT_TRUE(first.handle()->shared_scratch); + ASSERT_TRUE(second.handle()->shared_scratch); + + cudaStream_t first_stream = nullptr; + cudaStream_t second_stream = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&first_stream, cudaStreamNonBlocking), cudaSuccess); + ASSERT_EQ(cudaStreamCreateWithFlags(&second_stream, cudaStreamNonBlocking), cudaSuccess); + + // The host copies that bracket each run synchronize the whole device, so two + // threads left to themselves take turns rather than overlap. Against a build + // that leaves the window open, taking turns caught it in 2 of the 120 runs + // below; releasing both threads together caught nearly all of them. Neither + // thread can strand the other here -- both run the same fixed + // number of iterations and neither leaves the loop early. + std::atomic arrived{0}; + auto submit_together = [&arrived](int iteration) { + arrived.fetch_add(1); + while (arrived.load() < 2 * (iteration + 1)) { + std::this_thread::yield(); + } + }; + + std::atomic wrong_outputs{0}; + std::atomic failures{0}; + auto run_repeatedly = [&](LoadedEngine& engine, const std::vector& expected, cudaStream_t stream) { + for (int i = 0; i < kConcurrentRunsPerThread; ++i) { + // Rewritten every iteration, so a run whose enqueue never reached the engine + // leaves the sentinel behind rather than the previous iteration's output. + if (!engine.fill_output(kSentinel)) { + failures.fetch_add(1); + } + submit_together(i); + if (engine.run(stream) != Error::Ok || cudaStreamSynchronize(stream) != cudaSuccess) { + failures.fetch_add(1); + continue; + } + const std::vector actual = engine.read_output(); + if (actual.size() != expected.size() || std::memcmp(actual.data(), expected.data(), kBytes) != 0) { + wrong_outputs.fetch_add(1); + } + } + }; + + std::thread first_thread([&] { run_repeatedly(first, first_expected, first_stream); }); + std::thread second_thread([&] { run_repeatedly(second, second_expected, second_stream); }); + first_thread.join(); + second_thread.join(); + + ASSERT_EQ(cudaStreamDestroy(first_stream), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(second_stream), cudaSuccess); + + EXPECT_EQ(failures.load(), 0) << "a run failed outright, so fewer than " << (2 * kConcurrentRunsPerThread) + << " runs reached the comparison below"; + EXPECT_EQ(wrong_outputs.load(), 0) << wrong_outputs.load() << " of " << (2 * kConcurrentRunsPerThread) + << " concurrent pooled runs did not produce what the same engine produces with " + "its own scratch"; +} + } // namespace } // namespace executorch_backend } // namespace torch_tensorrt diff --git a/tests/cpp/executorch/test_shared_scratch_pool.cpp b/tests/cpp/executorch/test_shared_scratch_pool.cpp index 0532d0fe40..329c7b6758 100644 --- a/tests/cpp/executorch/test_shared_scratch_pool.cpp +++ b/tests/cpp/executorch/test_shared_scratch_pool.cpp @@ -26,11 +26,12 @@ namespace executorch_backend { namespace { // Fake device allocator: hands out distinct non-null pointers and records every -// allocation size and every release, so tests can assert the pool's grow/reuse -// policy and what each release was told to wait for, without a CUDA device. +// allocation size and every buffer a growth retired, so tests can assert the +// pool's grow/reuse policy and what each retirement has to wait for, without a +// CUDA device. struct FakeAllocator { std::vector alloc_sizes; - std::vector> released; + std::vector> retirements; std::uintptr_t next = 0x1000; bool fail_next = false; @@ -45,8 +46,8 @@ struct FakeAllocator { return p; } - void release(void* p, cudaEvent_t wait_for) { - released.emplace_back(p, wait_for); + void retire(void* p, cudaEvent_t wait_for) { + retirements.emplace_back(p, wait_for); } int alloc_count() const { @@ -74,13 +75,16 @@ struct FakeEventFactory { } }; +// Stands in for the backend: passes the allocator through and records whatever +// the call retired, the way execute() hands a retired buffer to its claim. void* call(SharedScratchDevice& dev, FakeAllocator& a, std::size_t need, std::size_t& out_size) { - return shared_scratch_get_or_grow( - dev, - need, - out_size, - [&a](std::size_t bytes) { return a.alloc(bytes); }, - [&a](void* p, cudaEvent_t wait_for) { a.release(p, wait_for); }); + RetiredScratch retired; + void* const p = shared_scratch_get_or_grow( + dev, need, out_size, [&a](std::size_t bytes) { return a.alloc(bytes); }, retired); + if (retired.buffer != nullptr) { + a.retire(retired.buffer, retired.wait_for); + } + return p; } TEST(SharedScratchPool, FirstRequestAllocatesExactSize) { @@ -94,7 +98,7 @@ TEST(SharedScratchPool, FirstRequestAllocatesExactSize) { EXPECT_EQ(out, 1024u); ASSERT_EQ(a.alloc_count(), 1); EXPECT_EQ(a.alloc_sizes[0], 1024u); - EXPECT_TRUE(a.released.empty()); + EXPECT_TRUE(a.retirements.empty()); } TEST(SharedScratchPool, ReusesWhenExistingBufferIsLargeEnough) { @@ -116,10 +120,10 @@ TEST(SharedScratchPool, ReusesWhenExistingBufferIsLargeEnough) { // Reuse reports the buffer's capacity, not the smaller amount asked for. EXPECT_EQ(out2, 4096u); EXPECT_EQ(a.alloc_count(), 1); - EXPECT_TRUE(a.released.empty()); + EXPECT_TRUE(a.retirements.empty()); } -TEST(SharedScratchPool, GrowsMonotonicallyToMaxAndReleasesOldBuffer) { +TEST(SharedScratchPool, GrowsMonotonicallyToMaxAndRetiresOldBuffer) { SharedScratchDevice dev; FakeAllocator a; std::size_t out = 0; @@ -131,8 +135,8 @@ TEST(SharedScratchPool, GrowsMonotonicallyToMaxAndReleasesOldBuffer) { EXPECT_EQ(out, 8192u); ASSERT_EQ(a.alloc_count(), 2); EXPECT_EQ(a.alloc_sizes[1], 8192u); - ASSERT_EQ(a.released.size(), 1u); - EXPECT_EQ(a.released[0].first, small); + ASSERT_EQ(a.retirements.size(), 1u); + EXPECT_EQ(a.retirements[0].first, small); // A subsequent smaller request reuses the grown buffer -- pool never shrinks. void* reuse = call(dev, a, 512, out); @@ -141,7 +145,7 @@ TEST(SharedScratchPool, GrowsMonotonicallyToMaxAndReleasesOldBuffer) { EXPECT_EQ(a.alloc_count(), 2); } -TEST(SharedScratchPool, GrowWaitsOnTheRecordedEnqueueBeforeReleasing) { +TEST(SharedScratchPool, GrowRetiresTheOldBufferWithTheEventToWaitOn) { SharedScratchDevice dev; FakeAllocator a; FakeEventFactory events; @@ -150,18 +154,19 @@ TEST(SharedScratchPool, GrowWaitsOnTheRecordedEnqueueBeforeReleasing) { void* small = call(dev, a, 1024, out); ASSERT_NE(small, nullptr); - // An enqueue against `small` has been submitted and recorded, so the release - // has something specific to outlive. + // An enqueue against `small` has been submitted and recorded, so its + // retirement has something specific to outlive. const SharedScratchHandoff handoff = shared_scratch_claim_event(dev, std::ref(events)); ASSERT_EQ(shared_scratch_mark_in_flight(dev), handoff.event); ASSERT_NE(call(dev, a, 8192, out), nullptr); - ASSERT_EQ(a.released.size(), 1u); - EXPECT_EQ(a.released[0].first, small); - // The release is handed the event that enqueue was recorded on, so it waits for - // that enqueue rather than for everything queued on the device. - EXPECT_EQ(a.released[0].second, handoff.event); + ASSERT_EQ(a.retirements.size(), 1u); + EXPECT_EQ(a.retirements[0].first, small); + // The retirement carries the event that enqueue was recorded on, so the caller + // has one specific enqueue to wait for, rather than needing a device-wide + // synchronize to be correct. + EXPECT_EQ(a.retirements[0].second, handoff.event); } TEST(SharedScratchPool, GrowHasNothingToWaitForWhenNoEnqueueWasRecorded) { @@ -178,9 +183,9 @@ TEST(SharedScratchPool, GrowHasNothingToWaitForWhenNoEnqueueWasRecorded) { ASSERT_NE(call(dev, a, 8192, out), nullptr); - ASSERT_EQ(a.released.size(), 1u); - EXPECT_EQ(a.released[0].first, small); - EXPECT_EQ(a.released[0].second, nullptr); + ASSERT_EQ(a.retirements.size(), 1u); + EXPECT_EQ(a.retirements[0].first, small); + EXPECT_EQ(a.retirements[0].second, nullptr); } TEST(SharedScratchPool, AllocationFailureLeavesExistingBufferUntouched) { @@ -197,7 +202,7 @@ TEST(SharedScratchPool, AllocationFailureLeavesExistingBufferUntouched) { std::size_t out2 = 0; void* failed = call(dev, a, 8192, out2); EXPECT_EQ(failed, nullptr); - EXPECT_TRUE(a.released.empty()); + EXPECT_TRUE(a.retirements.empty()); // The device still holds the original buffer and serves it on the next request. void* again = call(dev, a, 1024, out); @@ -327,15 +332,15 @@ TEST(SharedScratchPoolRegistry, KeepsAnIndependentBufferPerDevice) { EXPECT_NE(dev0, dev1); EXPECT_EQ(a.alloc_count(), 2); - EXPECT_TRUE(a.released.empty()); + EXPECT_TRUE(a.retirements.empty()); // Growing device 1 must not touch device 0's buffer. void* dev1_big = call(pool.get(1), a, 9000, out); void* dev0_again = call(pool.get(0), a, 2048, out); EXPECT_NE(dev1_big, dev1); EXPECT_EQ(dev0_again, dev0); - ASSERT_EQ(a.released.size(), 1u); - EXPECT_EQ(a.released[0].first, dev1); + ASSERT_EQ(a.retirements.size(), 1u); + EXPECT_EQ(a.retirements[0].first, dev1); } TEST(SharedScratchPoolRegistry, HandsOutOneStableEntryPerDevice) { @@ -373,6 +378,7 @@ TEST(SharedScratchPoolRegistry, AGrowthOnOneDeviceDoesNotBlockAClaimOnAnother) { std::thread grower([&] { std::lock_guard lk(dev0.mu); std::size_t out = 0; + RetiredScratch retired; shared_scratch_get_or_grow( dev0, 4096, @@ -382,7 +388,7 @@ TEST(SharedScratchPoolRegistry, AGrowthOnOneDeviceDoesNotBlockAClaimOnAnother) { leave.wait(); return zero.alloc(bytes); }, - [&](void* p, cudaEvent_t wait_for) { zero.release(p, wait_for); }); + retired); }); // The cap matters as much as the wait: a growth that takes the reuse path never // reaches its allocation, so nothing fires this promise and an uncapped wait @@ -404,12 +410,9 @@ TEST(SharedScratchPoolRegistry, AGrowthOnOneDeviceDoesNotBlockAClaimOnAnother) { SharedScratchDevice& dev1 = pool.get(1); std::lock_guard lk(dev1.mu); std::size_t out = 0; + RetiredScratch retired; return shared_scratch_get_or_grow( - dev1, - 2048, - out, - [&](std::size_t bytes) { return one.alloc(bytes); }, - [&](void* p, cudaEvent_t wait_for) { one.release(p, wait_for); }); + dev1, 2048, out, [&](std::size_t bytes) { return one.alloc(bytes); }, retired); }); const bool served = claim.wait_for(std::chrono::seconds(10)) == std::future_status::ready;