From 6c8c12eef42a9e12db5f4c414fc8aabb7d6f6f53 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Tue, 18 Aug 2026 17:33:31 -0700 Subject: [PATCH 01/13] 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 A100, 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 30619cda923..5c237010608 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 00000000000..62464e3aaf7 --- /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 b33d712d403..6713a950a24 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 e7367f87065..e65f5fa8fc2 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 8408c13e881..91b8d3e7a5f 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 17d2820bf25..d4f7a0c0811 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 00000000000..a65657e5df2 --- /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 204b9cee23f..ed98f4f3c69 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 ba430a3c44d7ea33a4d73d9eb70411cdd6ec9f71 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Thu, 27 Aug 2026 14:27:55 -0700 Subject: [PATCH 02/13] fix(executorch): make the scratch pool per device, fail an ambiguous zero-size query, cover the pooled path Four changes to the shared activation-scratch pool the parent commit added. 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 62464e3aaf7..5f43fdf3ce0 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 6713a950a24..e164bd01f96 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 e65f5fa8fc2..b5bd5cc526d 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 91b8d3e7a5f..a362ea7efcf 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 b5c0c151387..827fa2c4099 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 d4f7a0c0811..224e270295c 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 00000000000..54842b8db08 --- /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 a65657e5df2..0532d0fe404 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 a67cdf31bdced13e3f0db87c7009f7f8c8df71c6 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Tue, 1 Sep 2026 11:43:01 -0700 Subject: [PATCH 03/13] fix(executorch): hold the scratch pool's device lock across the enqueue Three fixes to the shared per-device activation-scratch pool: the device lock now spans the enqueue, the growth's `cudaFree` moves out from under it, and the growth path gets its first test. Left open: the option cannot be turned on from Python or from a `.pte`. 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 a 212-engine model the pool allocated once, at 7,353.1 MiB, and never grew, because the largest engine ran first. On a 62-engine one 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 5f43fdf3ce0..9e72e218bce 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 e164bd01f96..0a9571f22bb 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 b5bd5cc526d..58b8f552557 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 a362ea7efcf..3aee341b140 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 54842b8db08..c021e65e46a 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 0532d0fe404..329c7b67588 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; From 8e774ebbadf36c6f1042abecf13c38efa80629c5 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Wed, 2 Sep 2026 23:17:53 -0700 Subject: [PATCH 04/13] fix(executorch): unpublish the pool header and scope the growth claims Two corrections to the shared per-device activation-scratch pool: its header leaves the installed set, and the claim that a growth happens only on an engine's first run is scoped to the engines it holds for. Not addressed here: the ExecuTorch runtime CI job has been skipped on every run at this head, which needs a full CI lane. **`SharedScratchPool.h` was in the installed header set for no reason.** `TensorRTBackend.h` does not include it and does not use anything it declares: it named `kSharedActivationScratchKey` in three comments and nowhere else, and the README example and the backend test both spell the key out as a string literal. So the installed set carried a mutex, a per-device registry and three functions whose only locking rule is a sentence saying to call them with the device lock held -- no assert, no annotation, and the lock and the state are separate public members, so the type cannot enforce it -- to deliver one string constant. The constant moves next to `TensorRTBackend::set_option`, the one thing that reads it, and the header moves from `cpp/include/torch_tensorrt/executorch/` to `cpp/src/torch_tensorrt/executorch/`, beside the source that includes it. **The header has to stay in the package; the move itself is convention.** Dropping it from `executorch_api_headers` alone would have dropped it from the release too. `libtorchtrt.tar.gz` ships `TensorRTBackend.cpp` as buildable source -- the "Standalone Backend Archive" section of the README tells a caller to build `libexecutorch_trt_backend.a` from it -- and that source includes the pool header; reproducing the released layout without it fails with `fatal error: torch_tensorrt/executorch/SharedScratchPool.h: No such file or directory`. `executorch_api_headers` is also part of the source set that feeds the `rules_foreign_cc` CMake build in `py/torch-tensorrt-executorch-runtime/native`. Listing the header in `executorch_backend_source_files` covers both. That listing would have been enough on its own, wherever the file sat: `pkg_files` writes each file to its `prefix` plus the file's basename (`BUILD.bazel:134`, `:211`), so where a file lands in the tarball follows filegroup membership, not its path in the repository -- the header could have stayed under `cpp/include/` and still shipped under `src/`. It is moved for what `cpp/include/` means here: every header under it belongs to an installed set, the four ExecuTorch ones by name in `executorch_api_headers` and the rest by the `api_headers` glob over `include/**/*.h`. A private header left there would be the first exception, and would read as public to the next person who found it. `cpp/src/` held no headers before this, so the placement sets a convention rather than following one. The two CMake builds that compile `TensorRTBackend.cpp` gain the directory the header now sits under as an include root; `#include "torch_tensorrt/executorch/SharedScratchPool.h"` is unchanged in all three build systems. **"A growth happens only on an engine's first run" is false for a dynamic-shape engine.** `execute()` calls `updateDeviceMemorySizeForShapes()` once the input shapes are bound, on every call, and an engine that answers with what those shapes need answers differently as they change. Measured on TensorRT 11.2.1.2, one engine over a `[1..4, 512, 512]` profile with a `kUSER_MANAGED` context: | built with | batch 1 | batch 2 | batch 3 | batch 4 | |---|---|---|---|---| | `kRUNTIME_ACTIVATION_RESIZE_10_10` | 2 MiB | 4 MiB | 6 MiB | 8 MiB | | no preview feature | 8 MiB | 8 MiB | 8 MiB | 8 MiB | The first row is four growths out of one engine, so no ordering of the engines bounds the number of allocations; the second reports the profile maximum every time and grows the pool once. Which one an engine is, is fixed when it is built, which the README paragraph three lines below the claim already said. The statements that rested on the claim are corrected: - the README's growth-frequency guidance, and the advice that putting the largest engine first reduces the pool to a single allocation -- now scoped to the engines it holds for, and to run order rather than load order, since the pool allocates nothing until an engine runs. The measurement above stays in this message rather than going into the README: it pins a TensorRT version and the rule it supports does not; - the comment above the leak on a failed pre-free wait, which called the leak bounded because growth is rare. It is one buffer per growth whose wait fails, and the number of growths is not bounded by the engine count; - three descriptions of the pool as sized to "the largest engine's need", each of which assumes a per-engine constant: the README sentence introducing the option, the comment over the process-wide `scratch_pool`, and the comment over the per-call `setDeviceMemoryV2`, which put a move of the buffer down to "a larger engine". Verified: `//tests/cpp/executorch:test_shared_scratch_pool` 16/16, and the 12 cases of `test_shared_scratch_backend.cpp` 12/12 on an A100 with TensorRT 11.2.1.2, none skipped. The standalone CMake archive builds both from this tree and from a reproduction of the released package layout. --- cpp/BUILD | 9 ++-- .../executorch/TensorRTBackend.h | 14 ++++-- .../torch_tensorrt/executorch/CMakeLists.txt | 6 +++ cpp/src/torch_tensorrt/executorch/README.md | 48 +++++++++++-------- .../executorch/SharedScratchPool.h | 9 ---- .../executorch/TensorRTBackend.cpp | 21 ++++---- .../native/CMakeLists.txt | 3 ++ .../test_shared_scratch_backend.cpp | 2 +- 8 files changed, 67 insertions(+), 45 deletions(-) rename cpp/{include => src}/torch_tensorrt/executorch/SharedScratchPool.h (92%) diff --git a/cpp/BUILD b/cpp/BUILD index 5c237010608..71af9b8f4b7 100644 --- a/cpp/BUILD +++ b/cpp/BUILD @@ -191,12 +191,15 @@ cc_library( ], ) +# Implementation detail of :tensorrt_executorch_backend, so the header sits under +# src/ and ships with the sources rather than in the installed API set. A caller +# turning the pool on needs only the option key, which is in TensorRTBackend.h. cc_library( name = "tensorrt_executorch_shared_scratch_pool", hdrs = [ - "include/torch_tensorrt/executorch/SharedScratchPool.h", + "src/torch_tensorrt/executorch/SharedScratchPool.h", ], - strip_include_prefix = "include", + strip_include_prefix = "src", target_compatible_with = select({ ":linux_x86_64": [], ":sbsa": [], @@ -258,6 +261,7 @@ filegroup( srcs = [ "src/torch_tensorrt/executorch/CMakeLists.txt", "src/torch_tensorrt/executorch/README.md", + "src/torch_tensorrt/executorch/SharedScratchPool.h", "src/torch_tensorrt/executorch/TensorRTBackend.cpp", "src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp", "src/torch_tensorrt/executorch/WeightStreamingBudget.cpp", @@ -277,7 +281,6 @@ 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/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index 0a9571f22bb..2ebbad25aea 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -76,8 +76,7 @@ struct EngineHandle { 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). + // from the shared per-device pool (kSharedActivationScratchKey). 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, @@ -95,6 +94,15 @@ struct EngineHandle { ~EngineHandle(); }; +// 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. Read by TensorRTBackend::set_option below, and 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"; + class TensorRTBackend final : public ::executorch::runtime::BackendInterface { public: bool is_available() const override; @@ -125,7 +133,7 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { // Applies the runtime backend options a caller passes to // executorch::runtime::set_option("TensorRTBackend", ...). The only key read is - // kSharedActivationScratchKey (SharedScratchPool.h), a boolean. + // kSharedActivationScratchKey, a boolean. ::executorch::runtime::Error set_option( ET_UNUSED ::executorch::runtime::BackendOptionContext& context, const ::executorch::runtime::Span<::executorch::runtime::BackendOption>& backend_options) override; diff --git a/cpp/src/torch_tensorrt/executorch/CMakeLists.txt b/cpp/src/torch_tensorrt/executorch/CMakeLists.txt index 1b503c567da..6a291685259 100644 --- a/cpp/src/torch_tensorrt/executorch/CMakeLists.txt +++ b/cpp/src/torch_tensorrt/executorch/CMakeLists.txt @@ -36,6 +36,12 @@ set_target_properties(executorch_trt_backend target_include_directories(executorch_trt_backend PUBLIC "${CMAKE_CURRENT_LIST_DIR}/../../../include" + PRIVATE + # SharedScratchPool.h is an implementation detail, so it sits beside the + # sources instead of in the installed include tree. This directory is + # .../torch_tensorrt/executorch in both the repository and the released + # package, so two levels up is the root its #include path is relative to. + "${CMAKE_CURRENT_LIST_DIR}/../.." ) get_filename_component(_torchtrt_repo_root "${CMAKE_CURRENT_LIST_DIR}/../../../.." ABSOLUTE) diff --git a/cpp/src/torch_tensorrt/executorch/README.md b/cpp/src/torch_tensorrt/executorch/README.md index 58b8f552557..0088e04a5b9 100644 --- a/cpp/src/torch_tensorrt/executorch/README.md +++ b/cpp/src/torch_tensorrt/executorch/README.md @@ -122,7 +122,7 @@ 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: +requirement any call on that device has asked for: ```cpp #include @@ -140,26 +140,32 @@ 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: 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 -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. +concurrently on the device. The pool never shrinks, so the largest scratch it was +ever asked for stays allocated until the process exits. + +The buffer grows when a call asks for more than every call 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. + +What an engine answers when asked how much it needs is decided 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. + +How often the pool grows follows from that. The backend asks afresh on every +`execute()`, after the input shapes are bound, so an engine whose answer does not +vary with the bound shapes grows the pool at most once, on its first run — for a +program built only from such engines, running the largest one first leaves the +pool with a single allocation. An engine whose answer does vary can grow it on any +call whose shapes need more than every call before them, so with one of those in +the program no run order bounds the number of allocations. One engine over a +`[1..4, 512, 512]` profile is either, depending on how it was built. ## Standalone Backend Archive diff --git a/cpp/include/torch_tensorrt/executorch/SharedScratchPool.h b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h similarity index 92% rename from cpp/include/torch_tensorrt/executorch/SharedScratchPool.h rename to cpp/src/torch_tensorrt/executorch/SharedScratchPool.h index 9e72e218bce..6e19b4e6325 100644 --- a/cpp/include/torch_tensorrt/executorch/SharedScratchPool.h +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h @@ -21,15 +21,6 @@ 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 { diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index 3aee341b140..f480ae68bd0 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -230,9 +230,10 @@ 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 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. +// One buffer, grown to the largest requirement any call on that device has asked +// for, 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 @@ -408,8 +409,11 @@ Error get_or_grow_shared_scratch( 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. + // still be reading, so a failed wait leaks it instead -- one buffer for each + // growth whose wait fails. How many growths a run sees is not bounded by the + // engine count: the requirement is re-queried every execute() below, so an + // engine that answers with what the bound shapes need can grow the pool on + // any call. 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", @@ -1161,9 +1165,10 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // 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. + // The buffer is installed on every call, not once, because any call needing more + // than the pool holds -- this engine on other shapes, or another one -- grows it + // and moves it. 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 diff --git a/py/torch-tensorrt-executorch-runtime/native/CMakeLists.txt b/py/torch-tensorrt-executorch-runtime/native/CMakeLists.txt index 4dadbf56f5f..287ec853d53 100644 --- a/py/torch-tensorrt-executorch-runtime/native/CMakeLists.txt +++ b/py/torch-tensorrt-executorch-runtime/native/CMakeLists.txt @@ -173,6 +173,9 @@ target_compile_definitions(torch_tensorrt_executorch_backend target_include_directories(torch_tensorrt_executorch_backend PRIVATE "${TORCH_TENSORRT_SOURCE_DIR}" "${TORCH_TENSORRT_SOURCE_DIR}/cpp/include" + # Headers the backend sources include but that are not part of the installed + # API, such as SharedScratchPool.h. + "${TORCH_TENSORRT_SOURCE_DIR}/cpp/src" "${EXECUTORCH_SOURCE_DIR}/.." "${EXECUTORCH_SOURCE_DIR}/runtime/core/portable_type/c10") if(NOT TARGET extension_cuda) diff --git a/tests/cpp/executorch/test_shared_scratch_backend.cpp b/tests/cpp/executorch/test_shared_scratch_backend.cpp index c021e65e46a..741ce94478a 100644 --- a/tests/cpp/executorch/test_shared_scratch_backend.cpp +++ b/tests/cpp/executorch/test_shared_scratch_backend.cpp @@ -70,7 +70,7 @@ 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 +// Spelled out rather than taken from TensorRTBackend.h: a test that reads the // key through the production constant cannot pin the key's value. constexpr char kOptionKey[] = "use_shared_activation_scratch"; From 45f9990f14eba7ff89a0f8887d664032c06029fc Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Fri, 4 Sep 2026 15:25:05 -0700 Subject: [PATCH 05/13] fix(executorch): install the engine-wide activation scratch when the per-shape query answers zero `updateDeviceMemorySizeForShapes()` returns 0 for a valid empty input. A dynamic engine over a `[0..128]` batch profile, bound to batch 0, answers 0 with `allInputDimensionsSpecified()` true while `getDeviceMemorySizeV2()` is not zero, so the guard in `execute()` read that pair as a failed query and returned `Error::InvalidState`. With `use_shared_activation_scratch` on, such a call never reached the engine. Measured on TensorRT 11.2.1.2, the two-softmax network over `[0..128, 64, 64]` used by the test reports 4194304 bytes engine-wide and 0, 32768, 131072 and 4194304 bytes for batches 0, 1, 4 and 128. A zero has three causes and nothing at that point tells them apart: a failed query, an engine that needs no scratch under any shape, and a call whose bound shapes need none. The engine's own requirement covers all three, so the guard is replaced by a substitution rather than loosened. Where `engine_scratch_bytes` is zero no shape needs scratch and no buffer is installed, which is what the scratch-free path already did. Where it is not, it is an upper bound over every shape the engine accepts: a conservatively over-sized but valid buffer for a call whose own requirement could not be read, which is better than the stale, possibly freed pointer the guard existed to prevent. A `Debug` log records that a zero was seen and what was substituted for it. The comments that described the old discrimination -- on `engine_scratch_bytes` in the public header and on `TheEngineLevelRequirementSeparatesTheTwoFixtureEngines` -- are rewritten to match. `AnEmptyInputRunsWithThePoolEnabled` covers the empty input, on a fourth fixture engine built over a dynamic batch whose profile minimum is empty. Three assertions pin the test to that cause rather than the other two: the per-shape query answers zero at batch 0, the engine's own requirement is not zero, and the same engine does report a requirement at batch 4. Run against the previous guard the test fails with the guard's own message, and the other twelve in the target pass. Three smaller fixes in the same area: - The pool's handoff event is created with `cudaEventBlockingSync` as well as `cudaEventDisableTiming`, matching the engine-completion event in `init()`. The one host wait on it runs with the per-device lock held, so a spin burned a core for the whole of a wait that was already serializing the device. - The `get_or_grow_shared_scratch` precondition no longer lists `cudaFree` among the calls it makes on the current device; that free moved to `SharedScratchClaim::release()` earlier in this branch. - The README's `set_option` example stores and handles the `Error` it returns instead of dropping it, and the four em dashes this branch added, on three lines, become a comma, a pair of parentheses and a full stop. --- .../executorch/TensorRTBackend.h | 6 +- cpp/src/torch_tensorrt/executorch/README.md | 19 +- .../executorch/TensorRTBackend.cpp | 41 +-- .../test_shared_scratch_backend.cpp | 238 +++++++++++++++--- 4 files changed, 237 insertions(+), 67 deletions(-) diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index 2ebbad25aea..6bdd64267bb 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -79,8 +79,10 @@ struct EngineHandle { // from the shared per-device pool (kSharedActivationScratchKey). 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. + // shared_scratch is set. It bounds every shape the engine accepts, so execute() + // substitutes it whenever the per-call query answers zero -- which TensorRT does + // for a failed query, for an engine that genuinely needs none, and for shapes + // that need none. size_t engine_scratch_bytes = 0; std::mutex mu; // Makes the skip-sync fast path safe to reuse: TensorRT forbids reconfiguring or diff --git a/cpp/src/torch_tensorrt/executorch/README.md b/cpp/src/torch_tensorrt/executorch/README.md index 0088e04a5b9..2d77186cb29 100644 --- a/cpp/src/torch_tensorrt/executorch/README.md +++ b/cpp/src/torch_tensorrt/executorch/README.md @@ -92,7 +92,7 @@ the removed `CudaStreamGuard`: 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 + 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. @@ -120,7 +120,7 @@ the removed `CudaStreamGuard`: 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 — +`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 requirement any call on that device has asked for: @@ -129,12 +129,17 @@ requirement any call on that device has asked for: executorch::runtime::BackendOptions<1> options; options.set_option("use_shared_activation_scratch", true); -executorch::runtime::set_option("TensorRTBackend", options.view()); +const executorch::runtime::Error err = + executorch::runtime::set_option("TensorRTBackend", options.view()); +if (err != executorch::runtime::Error::Ok) { + return err; // nothing was set: every context still allocates its own scratch +} ``` -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. +`Error::NotFound` means no backend is registered under that name, which is what a +binary that has not linked the backend archive gets. Nothing forces the check: the +free `executorch::runtime::set_option` is not `ET_NODISCARD`, so dropping its +return compiles. 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 @@ -160,7 +165,7 @@ runtime does changes it. How often the pool grows follows from that. The backend asks afresh on every `execute()`, after the input shapes are bound, so an engine whose answer does not -vary with the bound shapes grows the pool at most once, on its first run — for a +vary with the bound shapes grows the pool at most once, on its first run. For a program built only from such engines, running the largest one first leaves the pool with a single allocation. An engine whose answer does vary can grow it on any call whose shapes need more than every call before them, so with one of those in diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index f480ae68bd0..b8a92fc6764 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -334,9 +334,8 @@ class SharedScratchClaim { // 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. +// Must be called with `device_id` already current: cudaEventCreateWithFlags and +// cudaMalloc both act on the *current* device and nothing in here sets it. Error get_or_grow_shared_scratch( SharedScratchClaim& claim, int device_id, @@ -348,7 +347,11 @@ Error get_or_grow_shared_scratch( const SharedScratchHandoff handoff = shared_scratch_claim_event(dev, []() -> cudaEvent_t { cudaEvent_t event = nullptr; - if (cudaEventCreateWithFlags(&event, cudaEventDisableTiming) != cudaSuccess) { + // Blocking-sync so the host yields instead of busy-spinning. The one host wait + // on this event -- the one below, before a displaced buffer is freed -- runs + // with the device's lock held, so it already holds off every other pooled + // engine on the device; spinning would burn a core for that whole time as well. + if (cudaEventCreateWithFlags(&event, cudaEventDisableTiming | cudaEventBlockingSync) != cudaSuccess) { return nullptr; } return event; @@ -1170,13 +1173,14 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // and moves it. 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. + // A reported zero has more than one cause, and nothing here tells them apart: a + // failed query, an engine that needs no scratch under any shape, and a call + // whose bound shapes need none -- an empty input inside a profile that admits + // one. The engine's own requirement stands in for all three. Zero there means + // no shape needs scratch, so no buffer is installed and the context has nothing + // to claim; anything else is an upper bound over every shape the engine + // accepts, which is a valid buffer for a call whose own requirement could not + // be read and an over-sized one for a call that needs less. // // 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 @@ -1185,7 +1189,14 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* SharedScratchClaim scratch_claim; bool scratch_from_pool = false; if (engine->shared_scratch) { - const size_t need = ctx->updateDeviceMemorySizeForShapes(); + size_t need = ctx->updateDeviceMemorySizeForShapes(); + if (need == 0 && engine->engine_scratch_bytes > 0) { + ET_LOG( + Debug, + "TensorRTBackend::execute: updateDeviceMemorySizeForShapes returned 0; using the engine-wide activation scratch requirement of %zu bytes instead", + engine->engine_scratch_bytes); + need = engine->engine_scratch_bytes; + } if (need > 0) { void* pool = nullptr; size_t pool_size = 0; @@ -1196,12 +1207,6 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } 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; } } diff --git a/tests/cpp/executorch/test_shared_scratch_backend.cpp b/tests/cpp/executorch/test_shared_scratch_backend.cpp index 741ce94478a..3876ceae98c 100644 --- a/tests/cpp/executorch/test_shared_scratch_backend.cpp +++ b/tests/cpp/executorch/test_shared_scratch_backend.cpp @@ -10,8 +10,8 @@ // 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. +// smaller one already allocated, the two cases that make the per-shape query +// answer zero, 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. @@ -85,6 +85,17 @@ constexpr std::size_t kBytes = kElems * sizeof(float); constexpr int kBigRows = 4096; constexpr int kBigCols = 4096; +// A dynamic-batch engine whose profile admits an empty batch. Binding one is a +// valid call that needs no activation scratch, from an engine that needs some, +// which is one of the three things a zero from updateDeviceMemorySizeForShapes +// can mean. The dimensions are small because no test compares this engine's +// memory against a figure; only whether each answer it gives is zero matters. +constexpr int kDynRows = 64; +constexpr int kDynCols = 64; +constexpr int kDynMaxBatch = 128; +// Inside the same profile and not empty, so one engine covers both answers. +constexpr int kDynBatch = 4; + // 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; @@ -182,6 +193,26 @@ bool add_scratch_free_net(nvinfer1::INetworkDefinition& network, nvinfer1::ITens return true; } +// Wraps a serialized engine in the delegate's blob wire format. +std::vector wrap_engine_plan(const nvinfer1::IHostMemory& plan) { + 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; +} + std::vector build_engine_blob(bool needs_scratch, int rows = kRows, int cols = kCols) { static BuilderLogger logger; @@ -218,28 +249,56 @@ std::vector build_engine_blob(bool needs_scratch, int rows = kRows if (plan == nullptr) { return {}; } + return wrap_engine_plan(*plan); +} - 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)); +// The scratch-needing network again, over a leading dimension the caller chooses +// per call, with an empty batch as the profile minimum. +std::vector build_dynamic_batch_engine_blob() { + static BuilderLogger logger; - 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; + 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, kDynRows, kDynCols}); + if (input == nullptr || !add_scratch_needing_net(*network, *input)) { + return {}; + } + + TRTUniquePtr config(builder->createBuilderConfig()); + if (config == nullptr) { + return {}; + } + nvinfer1::IOptimizationProfile* profile = builder->createOptimizationProfile(); + profile->setDimensions("input_0", nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims3{0, kDynRows, kDynCols}); + profile->setDimensions("input_0", nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims3{kDynBatch, kDynRows, kDynCols}); + profile->setDimensions( + "input_0", nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims3{kDynMaxBatch, kDynRows, kDynCols}); + config->addOptimizationProfile(profile); + + TRTUniquePtr plan(builder->buildSerializedNetwork(*network, *config)); + if (plan == nullptr) { + return {}; + } + return wrap_engine_plan(*plan); } -// 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, int rows = kRows, int cols = kCols) { +// The activation scratch one context of the shared engine needs for the given +// shape, read the way execute() reads it. Zero if the engine could not be +// measured, which is also what an empty batch answers, so a caller reading a zero +// as meaningful has to rule the failure out by some other measurement. +std::size_t measure_engine_scratch( + const std::vector& blob, + int rows = kRows, + int cols = kCols, + int batch = 1) { static BuilderLogger logger; TensorRTBlobHeader header; if (!TensorRTBlobHeader::parse(blob.data(), blob.size(), header)) { @@ -260,7 +319,7 @@ std::size_t measure_engine_scratch(const std::vector& blob, int ro if (ctx == nullptr) { return 0; } - if (!ctx->setInputShape("input_0", nvinfer1::Dims3{1, rows, cols})) { + if (!ctx->setInputShape("input_0", nvinfer1::Dims3{batch, rows, cols})) { return 0; } return ctx->updateDeviceMemorySizeForShapes(); @@ -314,20 +373,29 @@ class LoadedEngine { } // Loads the blob through the backend, capturing whatever the shared-scratch - // 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) { + // option is set to at this moment. `batch`/`rows`/`cols` must be a shape the + // blob's profile admits; `batch` may be 0, which allocates nothing and leaves + // both device pointers null, a state an empty ExecuTorch tensor can arrive in. + Error load( + const std::vector& blob, + std::uint32_t seed, + int rows = kRows, + int cols = kCols, + int batch = 1) { + batch_ = static_cast(batch); 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_, bytes()) != cudaSuccess || cudaMalloc(&device_out_, bytes()) != cudaSuccess) { - return Error::MemoryAllocationFailed; - } - if (cudaMemcpy(device_in_, host_in.data(), bytes(), cudaMemcpyHostToDevice) != cudaSuccess) { - return Error::Internal; + if (bytes() > 0) { + if (cudaMalloc(&device_in_, bytes()) != cudaSuccess || cudaMalloc(&device_out_, bytes()) != cudaSuccess) { + return Error::MemoryAllocationFailed; + } + if (cudaMemcpy(device_in_, host_in.data(), bytes(), cudaMemcpyHostToDevice) != cudaSuccess) { + return Error::Internal; + } } arena_storage_.resize(kArenaBytes); @@ -343,6 +411,9 @@ class LoadedEngine { } bool fill_output(float value) { + if (bytes() == 0) { + return true; + } const std::vector host(elems(), value); return cudaMemcpy(device_out_, host.data(), bytes(), cudaMemcpyHostToDevice) == cudaSuccess; } @@ -352,8 +423,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, rows_, cols_}; - SizesType out_sizes[3] = {1, rows_, cols_}; + SizesType in_sizes[3] = {batch_, rows_, cols_}; + SizesType out_sizes[3] = {batch_, 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); @@ -369,7 +440,7 @@ class LoadedEngine { std::vector read_output() const { std::vector host_out(elems()); - if (cudaMemcpy(host_out.data(), device_out_, bytes(), cudaMemcpyDeviceToHost) != cudaSuccess) { + if (bytes() > 0 && cudaMemcpy(host_out.data(), device_out_, bytes(), cudaMemcpyDeviceToHost) != cudaSuccess) { host_out.clear(); } return host_out; @@ -380,7 +451,7 @@ class LoadedEngine { } std::size_t elems() const { - return static_cast(rows_) * static_cast(cols_); + return static_cast(batch_) * static_cast(rows_) * static_cast(cols_); } std::size_t bytes() const { @@ -398,6 +469,7 @@ class LoadedEngine { DelegateHandle* handle_ = nullptr; void* device_in_ = nullptr; void* device_out_ = nullptr; + SizesType batch_ = 1; SizesType rows_ = kRows; SizesType cols_ = kCols; }; @@ -437,13 +509,17 @@ class SharedScratchBackendTest : public ::testing::Test { blob_ = build_engine_blob(true); scratch_free_blob_ = build_engine_blob(false); big_blob_ = build_engine_blob(true, kBigRows, kBigCols); - if (blob_.empty() || scratch_free_blob_.empty() || big_blob_.empty()) { + dynamic_blob_ = build_dynamic_batch_engine_blob(); + if (blob_.empty() || scratch_free_blob_.empty() || big_blob_.empty() || dynamic_blob_.empty()) { return; } scratch_bytes_ = measure_engine_scratch(blob_); big_scratch_bytes_ = measure_engine_scratch(big_blob_, kBigRows, kBigCols); + empty_batch_scratch_bytes_ = measure_engine_scratch(dynamic_blob_, kDynRows, kDynCols, 0); + dynamic_batch_scratch_bytes_ = measure_engine_scratch(dynamic_blob_, kDynRows, kDynCols, kDynBatch); engine_bytes_ = engine_scratch_requirement(blob_); scratch_free_engine_bytes_ = engine_scratch_requirement(scratch_free_blob_); + dynamic_engine_bytes_ = engine_scratch_requirement(dynamic_blob_); } void SetUp() override { @@ -454,6 +530,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_FALSE(dynamic_blob_.empty()) << "TensorRT could not build the dynamic-batch fixture engine"; ASSERT_EQ(set_shared_scratch(backend_, false), Error::Ok); } @@ -473,23 +550,35 @@ class SharedScratchBackendTest : public ::testing::Test { return big_blob_; } + const std::vector& dynamic_blob() const { + return dynamic_blob_; + } + TensorRTBackend backend_; static std::vector blob_; static std::vector scratch_free_blob_; static std::vector big_blob_; + static std::vector dynamic_blob_; static std::size_t scratch_bytes_; static std::size_t big_scratch_bytes_; + static std::size_t empty_batch_scratch_bytes_; + static std::size_t dynamic_batch_scratch_bytes_; static std::int64_t engine_bytes_; static std::int64_t scratch_free_engine_bytes_; + static std::int64_t dynamic_engine_bytes_; }; std::vector SharedScratchBackendTest::blob_; std::vector SharedScratchBackendTest::scratch_free_blob_; std::vector SharedScratchBackendTest::big_blob_; +std::vector SharedScratchBackendTest::dynamic_blob_; std::size_t SharedScratchBackendTest::scratch_bytes_ = 0; std::size_t SharedScratchBackendTest::big_scratch_bytes_ = 0; +std::size_t SharedScratchBackendTest::empty_batch_scratch_bytes_ = 0; +std::size_t SharedScratchBackendTest::dynamic_batch_scratch_bytes_ = 0; std::int64_t SharedScratchBackendTest::engine_bytes_ = -1; std::int64_t SharedScratchBackendTest::scratch_free_engine_bytes_ = -1; +std::int64_t SharedScratchBackendTest::dynamic_engine_bytes_ = -1; // --------------------------------------------------------------------------- // set_option @@ -769,16 +858,18 @@ TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItRep // 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. +// A zero from updateDeviceMemorySizeForShapes() does not say which of its causes +// it is, so execute() substitutes the engine's own requirement and installs a +// buffer unless that is zero too. 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"; + EXPECT_GT(engine_bytes_, 0) << "the two-softmax network reports no activation scratch of its own, so it no longer " + "covers the case where the engine's own requirement stands in for a zero from the " + "per-shape query"; } TEST_F(SharedScratchBackendTest, EachEngineRecordsItsOwnActivationScratchRequirement) { @@ -831,6 +922,73 @@ TEST_F(SharedScratchBackendTest, AnEngineNeedingNoActivationScratchRunsWithThePo EXPECT_EQ(std::memcmp(expected.data(), actual.data(), kBytes), 0); } +// --------------------------------------------------------------------------- +// An empty input to an engine that does need activation scratch +// --------------------------------------------------------------------------- + +// An empty batch inside the profile is a valid call, and the per-shape query +// answers it with the same zero it gives for a failed query, from an engine whose +// own requirement is not zero. The three assertions at the top are what make this +// the empty-input case rather than either of the others: an engine that needs +// nothing would fail the second, and a measurement that could not read the engine +// at all would fail the third. +TEST_F(SharedScratchBackendTest, AnEmptyInputRunsWithThePoolEnabled) { + ASSERT_EQ(empty_batch_scratch_bytes_, 0u) + << "an empty batch reports " << empty_batch_scratch_bytes_ + << " bytes of activation scratch, so this test no longer covers a call whose per-shape query answers zero"; + ASSERT_GT(dynamic_engine_bytes_, 0) + << "the dynamic fixture engine reports no activation scratch of its own, so the zero above is the " + "scratch-free case rather than the empty-input one"; + ASSERT_GT(dynamic_batch_scratch_bytes_, 0u) << "the same engine reports no activation scratch at batch " << kDynBatch + << " either, so the zero above says nothing about the batch being empty"; + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + + // The non-empty reference is taken with private scratch, so the pooled run + // below has something to be compared against that the pool did not produce. + LoadedEngine priv; + ASSERT_EQ(priv.load(dynamic_blob(), 19, kDynRows, kDynCols, kDynBatch), Error::Ok); + ASSERT_FALSE(priv.handle()->shared_scratch); + 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 empty; + ASSERT_EQ(empty.load(dynamic_blob(), 19, kDynRows, kDynCols, 0), Error::Ok); + ASSERT_TRUE(empty.handle()->shared_scratch); + EXPECT_EQ(empty.run(stream), Error::Ok) << "the pool rejected an empty input to an engine that needs scratch"; + EXPECT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + // The same engine on a shape that does need scratch, after the empty call, so a + // fallback that installed an unusable buffer is not left untested. + LoadedEngine pooled; + ASSERT_EQ(pooled.load(dynamic_blob(), 19, kDynRows, kDynCols, kDynBatch), Error::Ok); + ASSERT_TRUE(pooled.handle()->shared_scratch); + ASSERT_TRUE(pooled.fill_output(kSentinel)); + EXPECT_EQ(pooled.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::vector actual = pooled.read_output(); + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); + + const std::size_t dyn_elems = + static_cast(kDynBatch) * static_cast(kDynRows) * static_cast(kDynCols); + ASSERT_EQ(expected.size(), dyn_elems); + ASSERT_EQ(actual.size(), dyn_elems); + // 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 < dyn_elems && !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(), dyn_elems * sizeof(float)), 0); +} + // --------------------------------------------------------------------------- // The enqueue handoff, single-threaded, two caller streams // --------------------------------------------------------------------------- From f594f82fc5da2a3dd20295d4c498bf18cfd7849d Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Sat, 5 Sep 2026 17:53:44 -0700 Subject: [PATCH 06/13] fix(executorch): stop a zero-scratch call sizing the pool for the whole profile `updateDeviceMemorySizeForShapes()` answers zero for a call whose bound shapes need no activation scratch, and the previous commit stood the engine's own `getDeviceMemorySizeV2()` in for that zero. That figure covers every shape in the profile. Measured on the test's dynamic `[0..128, 64, 64]` engine: batch 0 answers 0 and installed 4194304 bytes, while the next call, at batch 4, needs 131072. The pool never shrinks, so one empty batch pinned it at 32x the next call's requirement for the process lifetime, on the path the pool exists to keep small. A zero now asks for whatever the pool already holds, at its current size, which grows nothing: any live buffer is large enough for a call that needs none. Handing such a call no buffer at all does not work, which is worth recording because it is the obvious reading of "it needs nothing". `enqueueV3` refuses a `kUSER_MANAGED` context with no device memory installed as soon as the engine reports needing some under any shape, whatever the shapes actually bound need: `Parameter check failed, condition: noDeviceMemory. The engine requires 4194304 device memory.` So against an empty pool the call allocates `kMinPooledScratchBytes`, and the first call with a real requirement grows it; the same engine's batch-4 call then takes the pool to 131072 rather than 4 MiB. `EngineHandle::engine_scratch_bytes` becomes `engine_needs_scratch`, a bool. The size had no remaining reader, but the predicate has one: an engine that needs no scratch under any shape is left out of the pool entirely, so it takes no per-device lock and does not serialize against the engines that do. Two hazards the pool did not address: - Stream capture is now refused with `Error::NotSupported`. The handoff waits on an event recorded outside the capture, which `cudaStreamWaitEvent` fails with `cudaErrorStreamCaptureIsolation` and which invalidates the capture under every capture mode; that wait, and not the allocation, is what makes `cudaStreamEndCapture` hand back 901 and a null graph far from the cause. A growth's `cudaMalloc`, its `cudaEventSynchronize` and its `cudaFree` invalidate a capture too, but only outside `cudaStreamCaptureModeRelaxed`, which permits all three and then runs them uncaptured. The check is one `cudaStreamIsCapturing` ahead of the device lock. - `cudaDeviceReset()` is documented rather than guarded, in the installed header and the README. The pool holds its buffer and its handoff event for the process lifetime and a reset destroys both; catching that would mean revalidating both on every call, which costs as much as the work it protects. The installed `execute()` contract also now says that a growth turns an otherwise asynchronous call into a device-wide `cudaFree` wait. The README, the installed header and the test comments are brought into agreement on the shape of the ambiguity: a zero from the per-shape query has three causes, the engine-level one is settled at init and never reaches `execute()`, and the other two are indistinguishable where the zero is read. Test changes: - The fixture resets the process-wide pool in setup and teardown, through test-only hooks declared in `SharedScratchPool.h`, which ships under `src/` and is not installed. Nothing reset the pool before, and most of these tests depend on what it holds when they start, so the suite passed in declaration order and in no other. Running `AnEmptyInputRunsWithThePoolEnabled` before `ASecondPooledEnqueueWaitsForTheFirstOnAnotherStream` turned the latter's first run into a growth, whose `cudaFree` waits device-wide and so waited on the blocking host function that test parks on its own stream, from the one thread that could release it: a deadlock ending at the 60 s watchdog. `--gtest_repeat=2` on the growth case failed on its own second iteration. Both now pass. - `AnEmptyInputRunsWithThePoolEnabled` asserts what the pool holds rather than only that the call succeeds, at three points: after an empty call against an empty pool, after a call that needs scratch, and after a second empty call. Against the previous behaviour it fails on the first and third. The first is pinned to `kMinPooledScratchBytes` exactly rather than bounded below the next call's requirement, so a minimum raised to anything under 128 KiB fails it too. - `ResetHandsBackEverySlotAndLeavesItEmpty` covers `reset_for_testing` over the pool's fakes, which nothing did: that each slot's buffer and event both reach the disposer, that a slot left without an event is handled, that the reference `get` handed out survives, and that the next claim allocates rather than reusing the capacity of a buffer the reset has already released. - `AnEngineNeedingNoActivationScratchRunsWithThePoolEnabled` also checks the pool still holds nothing afterwards, which is what pins such an engine being left out of the pool rather than merely surviving it. - `device_bytes_in_use()` reports a failed `cudaMemGetInfo` instead of returning zero. Both memory comparisons subtract two of its readings, and a zero for the first of a pair made the second look like the whole cost of what was measured between them. - `APooledEngineRefusesToRunWhileItsStreamIsCapturing` covers the capture guard, and checks the refusal did not itself spoil the capture. - `EachEngineRecordsItsOwnActivationScratchRequirement` and `TheEngineLevelRequirementSeparatesTheTwoFixtureEngines` are replaced by `EachEngineRecordsWhetherItNeedsActivationScratch`, which pins the predicate the pooled path now branches on. --- .../executorch/TensorRTBackend.h | 36 +++- cpp/src/torch_tensorrt/executorch/README.md | 54 +++++- .../executorch/SharedScratchPool.h | 40 ++++ .../executorch/TensorRTBackend.cpp | 135 +++++++++---- tests/cpp/executorch/BUILD | 1 + .../test_shared_scratch_backend.cpp | 179 ++++++++++++++---- .../executorch/test_shared_scratch_pool.cpp | 71 +++++++ 7 files changed, 438 insertions(+), 78 deletions(-) diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index 6bdd64267bb..2c65d0e0220 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -78,12 +78,14 @@ struct EngineHandle { // Whether exec_ctx was created kUSER_MANAGED and draws its activation scratch // from the shared per-device pool (kSharedActivationScratchKey). bool shared_scratch = false; - // The activation scratch the engine itself reports needing, read at init when - // shared_scratch is set. It bounds every shape the engine accepts, so execute() - // substitutes it whenever the per-call query answers zero -- which TensorRT does - // for a failed query, for an engine that genuinely needs none, and for shapes - // that need none. - size_t engine_scratch_bytes = 0; + // Whether the engine reports needing activation scratch under any shape, read + // at init when shared_scratch is set. It is a predicate and not a size because + // the size is never the right amount to install: enqueueV3 refuses a + // kUSER_MANAGED context with no device memory once this is true, however little + // the shapes bound to a given call need, so execute() must hand such a call + // some buffer -- but the engine's own figure covers every shape in the profile, + // and installing it would size the pool for the largest of them. + bool engine_needs_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 @@ -125,7 +127,27 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { // 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. + // scratch and is not subject to this, nor is one whose engine reports needing no + // activation scratch under any shape, which is left out of the pool. For the + // handles that do draw on it, three further consequences: + // - A call needing more scratch than the pool holds grows it, and the growth + // frees the buffer it replaces. cudaFree waits for every stream on the + // device, so that one call blocks until the device is idle however + // asynchronous the rest of this contract makes it -- an unbounded wait on + // work this call did not submit. Which calls grow the pool is not knowable + // from here; see the README. + // - Capturing a CUDA graph from the selected stream is refused with + // Error::NotSupported. The pool's event handoff waits on an event recorded + // outside the capture, which invalidates it under every capture mode. A + // growth's allocation, its wait on the replaced buffer and the free of that + // buffer invalidate it under every mode but cudaStreamCaptureModeRelaxed, + // which permits all three but does not record them, leaving a replay pointed + // at a buffer the pool may since have freed. The alternative to refusing is a + // capture that silently comes back invalidated. + // - cudaDeviceReset() invalidates the pool without emptying it. The buffer + // and the handoff event it still holds are destroyed with the primary + // context, and the next call on that device uses both. There is no guard: + // do not reset a device this backend has run a pooled engine on. // 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 2d77186cb29..4708adf2bee 100644 --- a/cpp/src/torch_tensorrt/executorch/README.md +++ b/cpp/src/torch_tensorrt/executorch/README.md @@ -68,8 +68,7 @@ type. The upstream `CallerStreamGuard` documents the generic contract (per-thread, nested scoping; the caller owns the stream for the guard's lifetime; the caller manages host-data lifetime for async work). The TensorRT backend adds these -requirements, which previously lived on -the removed `CudaStreamGuard`: +requirements, which previously lived on the removed `CudaStreamGuard`: - The selected stream must be on the TensorRT engine's device. - Calls using one delegate handle must not overlap, and must not overlap with @@ -95,7 +94,8 @@ the removed `CudaStreamGuard`: 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. + unaffected, and so is a context whose engine needs no activation scratch under + any shape: that one is left out of the pool and serializes against nothing. - 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 @@ -148,6 +148,27 @@ bullet of the caller-stream contract above: engines sharing a buffer do not run concurrently on the device. The pool never shrinks, so the largest scratch it was ever asked for stays allocated until the process exits. +A call that asks for nothing is handed whatever the pool already holds, at its +current size, so it never grows it. Two of the things a zero from the per-shape +query can mean reach this point, and are indistinguishable where it is read: the +shapes bound to this call need none -- an empty batch inside a profile that admits +one -- or the query failed. Neither wants a buffer of its own, and any buffer the +pool already holds is large enough for either. + +Such a call cannot be left with no buffer at all. `enqueueV3` refuses a +`kUSER_MANAGED` context with no device memory installed as soon as the engine +reports needing some under *any* shape, however little the shapes actually bound +need. So where the pool holds nothing yet, a zero allocates the smallest buffer +that satisfies that check, and the first call with a real requirement grows it. +Standing the engine's own figure in for the zero would satisfy the check too, but +that figure covers the whole profile: on the dynamic engine the tests use it is +4 MiB against the 128 KiB the next call needs, and the pool never shrinks. + +The third thing a zero can mean is an engine that needs no scratch under *any* +shape, and that one is settled before a call is ever made: such an engine is left +out of the pool entirely, so it takes no per-device lock and does not serialize +against the engines that do. + The buffer grows when a call asks for more than every call 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, @@ -172,6 +193,33 @@ call whose shapes need more than every call before them, so with one of those in the program no run order bounds the number of allocations. One engine over a `[1..4, 512, 512]` profile is either, depending on how it was built. +Two things the pool does not support. + +Capturing a CUDA graph from the stream its `execute()` runs on is refused with +`Error::NotSupported`. The handoff between one enqueue and the next waits on an +event recorded outside the capture, and `cudaStreamWaitEvent` on such an event +fails with `cudaErrorStreamCaptureIsolation` and invalidates the capture under +every capture mode, `cudaStreamCaptureModeRelaxed` included; the caller learns of +that only when `cudaStreamEndCapture` hands back `cudaErrorStreamCaptureInvalidated` +and a null graph, long after the call that caused it. A growth is mode-dependent +rather than always fatal. Its `cudaMalloc`, the `cudaEventSynchronize` it makes +before releasing the buffer it replaces, and the `cudaFree` that releases it each +return `cudaErrorStreamCaptureUnsupported` and invalidate the capture under the +`Global` and `ThreadLocal` modes; under `Relaxed` all three are permitted and the +capture survives them, but they then run uncaptured, so a replayed graph would use +whatever buffer was installed when it was captured -- which by then the pool may +have freed. The backend checks the stream before it makes any of these calls and +returns instead. Load the engine with the option off to capture it: a context that +owns its scratch makes none of them. + +`cudaDeviceReset()` is not survivable and is not guarded against. The pool holds +its buffer and its handoff event for the process lifetime, and a reset destroys +the primary context under both. The next pooled `execute()` on that device then +waits on a destroyed event and hands the engine a pointer that is no longer a +device allocation. Guarding this would mean revalidating both on every call, and +the check that would catch it is as expensive as the work it protects. Treat a +device the backend has run a pooled engine on as one that must not be reset. + ## Standalone Backend Archive Use this path only when you need `libexecutorch_trt_backend.a` without building diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h index 6e19b4e6325..bdb7df592f6 100644 --- a/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h @@ -68,6 +68,28 @@ class SharedScratchPool { return devices_[device_id]; } + // Test-only. Returns every device's slot to the state it had before anything + // claimed it, handing what the slot held to `dispose(device_id, buffer, event)` + // so the caller can release it. Entries stay in the map, so a reference `get` + // handed out remains valid. + // + // Nothing here waits for an enqueue: the buffer it frees may still be in use by + // one, and the caller is responsible for there being none. `dispose` runs under + // both this registry's lock and the device's, so it must not claim either. + template + void reset_for_testing(Dispose dispose) { + std::lock_guard lk(mu_); + for (auto& entry : devices_) { + SharedScratchDevice& dev = entry.second; + std::lock_guard dev_lk(dev.mu); + dispose(entry.first, dev.buffer, dev.marker.event); + dev.buffer = nullptr; + dev.capacity = 0; + dev.marker.event = nullptr; + dev.marker.pending = false; + } + } + private: std::mutex mu_; std::unordered_map devices_; @@ -170,5 +192,23 @@ void* shared_scratch_get_or_grow( return p; } +// Test-only views of the process-wide pool the TensorRT backend runs on. Declared +// here, in a header that ships under src/, so they add nothing to the installed +// API; defined in TensorRTBackend.cpp, which owns the pool instance they act on. +// A binary that links the pool header alone -- the pool's own unit test -- never +// names them, so the missing definition costs it nothing. +// +// Both are safe only with no claim outstanding and no enqueue in flight against a +// pooled buffer; a test earns that by synchronizing every stream it submitted on. +// The reset also frees what the pool holds, so a test should destroy its delegate +// handles before that one as well. + +// The bytes the pool holds for `device_id` right now; zero if it holds nothing. +std::size_t shared_scratch_capacity_for_testing(int device_id); + +// Frees every device's buffer, destroys its handoff event, and clears the marker, +// so one test does not inherit a pool an earlier one grew. +void reset_shared_scratch_pool_for_testing(); + } // namespace executorch_backend } // namespace torch_tensorrt diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index b8a92fc6764..38355ee1963 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -187,7 +187,7 @@ Error initialize_engine_io(EngineHandle& handle) { // 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()); + handle.engine_needs_scratch = handle.engine->getDeviceMemorySizeV2() > 0; } return Error::Ok; @@ -239,10 +239,9 @@ bool is_cuda_accessible_ptr(const void* ptr) { // can still be in flight when execute() returns, so two enqueues must never hold // 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. +// 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. // // 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 @@ -329,11 +328,26 @@ class SharedScratchClaim { void* retired_ = nullptr; }; +// What a call needing no activation scratch is given when the pool holds nothing +// yet. It cannot be given nothing: enqueueV3 refuses a kUSER_MANAGED context with +// no device memory installed as soon as the engine reports needing any under some +// shape, whatever the shapes actually bound need. So the pool starts at the +// smallest allocation that satisfies that check and the first call with a real +// requirement grows it -- as against standing the engine's profile-wide figure in +// for the zero, which would pin the pool at the largest shape the engine admits +// on the strength of a call that uses none of it. +constexpr size_t kMinPooledScratchBytes = 1; + // 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. // 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. // +// A `need` of zero asks for whatever the pool already holds rather than for a +// buffer of its own: any live buffer is large enough, so nothing is allocated and +// nothing grows. Only against an empty pool does it allocate, and then at +// kMinPooledScratchBytes. +// // Must be called with `device_id` already current: cudaEventCreateWithFlags and // cudaMalloc both act on the *current* device and nothing in here sets it. Error get_or_grow_shared_scratch( @@ -343,8 +357,37 @@ Error get_or_grow_shared_scratch( cudaStream_t stream, void*& out_ptr, size_t& out_size) { + // Checked before the lock, and so ahead of every CUDA call the pooled path makes + // that a capture cannot take. The handoff's event wait is on an event recorded + // outside the capture: cudaStreamWaitEvent fails it with + // cudaErrorStreamCaptureIsolation and invalidates the capture under every capture + // mode. The allocation, the cudaEventSynchronize on the replaced buffer and the + // cudaFree of it invalidate the capture too, but only outside + // cudaStreamCaptureModeRelaxed, which permits all three -- and then runs them + // uncaptured, leaving a replay pointed at a buffer the pool may since have freed. + // None of it fails cleanly: the caller learns of an invalidation only when + // cudaStreamEndCapture hands back an error and a null graph. Refusing names the + // cause instead. A failed query is refused with the rest: what it reports on the + // legacy stream while another stream captures is cudaErrorStreamCaptureImplicit, + // which is the hazard rather than an unrelated fault. + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + const cudaError_t capture_err = cudaStreamIsCapturing(stream, &capture); + if (capture_err != cudaSuccess || capture != cudaStreamCaptureStatusNone) { + ET_LOG( + Error, + "TensorRTBackend::execute: the selected stream is capturing a CUDA graph (%s), which the shared activation scratch pool on device %d does not support. Load this engine with '%s' off so its context keeps its own scratch.", + capture_err == cudaSuccess ? "capture in progress" : cudaGetErrorString(capture_err), + device_id, + kSharedActivationScratchKey); + return Error::NotSupported; + } + SharedScratchDevice& dev = claim.hold(device_id); + if (need == 0) { + need = dev.capacity > 0 ? dev.capacity : kMinPooledScratchBytes; + } + const SharedScratchHandoff handoff = shared_scratch_claim_event(dev, []() -> cudaEvent_t { cudaEvent_t event = nullptr; // Blocking-sync so the host yields instead of busy-spinning. The one host wait @@ -461,6 +504,40 @@ Error mark_shared_scratch_in_flight(SharedScratchClaim& claim, cudaStream_t stre } // namespace +// Declared in SharedScratchPool.h, which ships under src/ and is not installed; +// defined here because this is where the pool instance lives. See that header for +// what a caller owes them. +std::size_t shared_scratch_capacity_for_testing(int device_id) { + SharedScratchDevice& dev = scratch_pool.get(device_id); + std::lock_guard lk(dev.mu); + return dev.capacity; +} + +void reset_shared_scratch_pool_for_testing() { + int restore_to = 0; + const bool have_current = cudaGetDevice(&restore_to) == cudaSuccess; + scratch_pool.reset_for_testing([](int device_id, void* buffer, cudaEvent_t event) { + if (buffer == nullptr && event == nullptr) { + return; + } + // cudaFree and cudaEventDestroy both act on the current device, and a slot is + // keyed by the device its buffer came from. + if (cudaSetDevice(device_id) != cudaSuccess) { + return; + } + if (buffer != nullptr) { + (void)cudaFree(buffer); + } + if (event != nullptr) { + (void)cudaEventDestroy(event); + } + (void)cudaGetLastError(); // a reset is cleanup; do not leave a sticky error for the next call + }); + if (have_current) { + (void)cudaSetDevice(restore_to); + } +} + // --------------------------------------------------------------------------- // is_available // --------------------------------------------------------------------------- @@ -1173,14 +1250,17 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // and moves it. A kSTATIC context owns its private scratch, so setDeviceMemoryV2 // must not be called on one. // - // A reported zero has more than one cause, and nothing here tells them apart: a - // failed query, an engine that needs no scratch under any shape, and a call - // whose bound shapes need none -- an empty input inside a profile that admits - // one. The engine's own requirement stands in for all three. Zero there means - // no shape needs scratch, so no buffer is installed and the context has nothing - // to claim; anything else is an upper bound over every shape the engine - // accepts, which is a valid buffer for a call whose own requirement could not - // be read and an over-sized one for a call that needs less. + // An engine that needs no scratch under any shape is left out of the pool + // entirely: enqueueV3 accepts it with no device memory installed, so it need not + // claim the device and does not serialize against the engines that do. + // + // For the rest, a reported zero has more than one cause and nothing here tells + // them apart: a failed query, and a call whose bound shapes need none -- an + // empty input inside a profile that admits one. Neither asks for a buffer of its + // own, so neither grows the pool: get_or_grow_shared_scratch hands a zero + // whatever the pool already holds, and only where it holds nothing does it + // allocate, at the minimum TensorRT will accept rather than at the engine's + // profile-wide figure. // // 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 @@ -1188,26 +1268,17 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // restore above, so its free lands on the right device. SharedScratchClaim scratch_claim; bool scratch_from_pool = false; - if (engine->shared_scratch) { - size_t need = ctx->updateDeviceMemorySizeForShapes(); - if (need == 0 && engine->engine_scratch_bytes > 0) { - ET_LOG( - Debug, - "TensorRTBackend::execute: updateDeviceMemorySizeForShapes returned 0; using the engine-wide activation scratch requirement of %zu bytes instead", - engine->engine_scratch_bytes); - need = engine->engine_scratch_bytes; - } - if (need > 0) { - void* pool = nullptr; - size_t pool_size = 0; - 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; - } - scratch_from_pool = true; - ctx->setDeviceMemoryV2(pool, static_cast(pool_size)); + if (engine->shared_scratch && engine->engine_needs_scratch) { + const size_t need = ctx->updateDeviceMemorySizeForShapes(); + void* pool = nullptr; + size_t pool_size = 0; + 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; } + scratch_from_pool = true; + ctx->setDeviceMemoryV2(pool, static_cast(pool_size)); } // ------------------------------------------------------------------ diff --git a/tests/cpp/executorch/BUILD b/tests/cpp/executorch/BUILD index 224e270295c..a64447a9446 100644 --- a/tests/cpp/executorch/BUILD +++ b/tests/cpp/executorch/BUILD @@ -74,6 +74,7 @@ cc_test( deps = [ "//cpp:tensorrt_executorch_backend", "//cpp:tensorrt_executorch_blob_header", + "//cpp:tensorrt_executorch_shared_scratch_pool", "@executorch//:executorch_core", "@executorch//:executorch_headers", "@executorch//:extension_cuda", diff --git a/tests/cpp/executorch/test_shared_scratch_backend.cpp b/tests/cpp/executorch/test_shared_scratch_backend.cpp index 3876ceae98c..dd731b9d348 100644 --- a/tests/cpp/executorch/test_shared_scratch_backend.cpp +++ b/tests/cpp/executorch/test_shared_scratch_backend.cpp @@ -10,8 +10,9 @@ // 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, the two cases that make the per-shape query -// answer zero, and two threads submitting against one pooled buffer at once. +// smaller one already allocated, two of the three things that make the per-shape +// query answer zero -- the third, a failed query, cannot be induced from inside +// this process -- 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. @@ -20,6 +21,7 @@ // 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/SharedScratchPool.h" #include "torch_tensorrt/executorch/TensorRTBackend.h" #include "torch_tensorrt/executorch/TensorRTBlobHeader.h" @@ -290,7 +292,7 @@ std::vector build_dynamic_batch_engine_blob() { return wrap_engine_plan(*plan); } -// The activation scratch one context of the shared engine needs for the given +// The activation scratch one context of `blob`'s engine needs for the given // shape, read the way execute() reads it. Zero if the engine could not be // measured, which is also what an empty batch answers, so a caller reading a zero // as meaningful has to rule the failure out by some other measurement. @@ -474,13 +476,18 @@ class LoadedEngine { SizesType cols_ = kCols; }; -std::size_t device_bytes_in_use() { +// Sets `out` to the device-wide bytes in use, or returns false leaving it alone. +// It reports the failure rather than substituting a figure because both callers +// subtract two of these: a zero for the first reading of a pair makes the second +// look like the whole cost of what was measured between them, which is a pass. +bool device_bytes_in_use(std::size_t& out) { std::size_t free_bytes = 0; std::size_t total_bytes = 0; if (cudaMemGetInfo(&free_bytes, &total_bytes) != cudaSuccess) { - return 0; + return false; } - return total_bytes - free_bytes; + out = total_bytes - free_bytes; + return true; } Error set_shared_scratch(TensorRTBackend& backend, bool enabled) { @@ -517,7 +524,6 @@ class SharedScratchBackendTest : public ::testing::Test { big_scratch_bytes_ = measure_engine_scratch(big_blob_, kBigRows, kBigCols); empty_batch_scratch_bytes_ = measure_engine_scratch(dynamic_blob_, kDynRows, kDynCols, 0); dynamic_batch_scratch_bytes_ = measure_engine_scratch(dynamic_blob_, kDynRows, kDynCols, kDynBatch); - engine_bytes_ = engine_scratch_requirement(blob_); scratch_free_engine_bytes_ = engine_scratch_requirement(scratch_free_blob_); dynamic_engine_bytes_ = engine_scratch_requirement(dynamic_blob_); } @@ -532,10 +538,21 @@ class SharedScratchBackendTest : public ::testing::Test { ASSERT_FALSE(big_blob_.empty()) << "TensorRT could not build the larger fixture engine"; ASSERT_FALSE(dynamic_blob_.empty()) << "TensorRT could not build the dynamic-batch fixture engine"; ASSERT_EQ(set_shared_scratch(backend_, false), Error::Ok); + // The pool outlives every test in this file, and most of them depend on what + // it holds when they start: one expects a growth, one expects none, one + // expects a first allocation. Leaving it alone made the suite pass in + // declaration order and in no other -- running the empty-input case first + // turned the handoff test's first run into a growth, whose cudaFree waits + // device-wide and so waited on the blocked host function that test parks on + // its own stream, from the thread that alone could release it. + reset_shared_scratch_pool_for_testing(); } void TearDown() override { set_shared_scratch(backend_, false); + // Also here, so a buffer this test grew is not still resident while the next + // one measures device-wide memory. + reset_shared_scratch_pool_for_testing(); } const std::vector& blob() const { @@ -563,7 +580,6 @@ class SharedScratchBackendTest : public ::testing::Test { static std::size_t big_scratch_bytes_; static std::size_t empty_batch_scratch_bytes_; static std::size_t dynamic_batch_scratch_bytes_; - static std::int64_t engine_bytes_; static std::int64_t scratch_free_engine_bytes_; static std::int64_t dynamic_engine_bytes_; }; @@ -576,7 +592,6 @@ std::size_t SharedScratchBackendTest::scratch_bytes_ = 0; std::size_t SharedScratchBackendTest::big_scratch_bytes_ = 0; std::size_t SharedScratchBackendTest::empty_batch_scratch_bytes_ = 0; std::size_t SharedScratchBackendTest::dynamic_batch_scratch_bytes_ = 0; -std::int64_t SharedScratchBackendTest::engine_bytes_ = -1; std::int64_t SharedScratchBackendTest::scratch_free_engine_bytes_ = -1; std::int64_t SharedScratchBackendTest::dynamic_engine_bytes_ = -1; @@ -718,7 +733,8 @@ TEST_F(SharedScratchBackendTest, PooledEnginesShareOneActivationScratchAllocatio std::size_t private_cost = 0; { - const std::size_t before = device_bytes_in_use(); + std::size_t before = 0; + ASSERT_TRUE(device_bytes_in_use(before)) << "cudaMemGetInfo failed, so this test measured nothing"; std::vector> engines; for (int i = 0; i < kEngineCount; ++i) { engines.push_back(std::make_unique()); @@ -726,7 +742,8 @@ TEST_F(SharedScratchBackendTest, PooledEnginesShareOneActivationScratchAllocatio ASSERT_EQ(engines.back()->run(stream), Error::Ok); } ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); - const std::size_t after = device_bytes_in_use(); + std::size_t after = 0; + ASSERT_TRUE(device_bytes_in_use(after)) << "cudaMemGetInfo failed, so this test measured nothing"; // 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 " @@ -737,7 +754,8 @@ TEST_F(SharedScratchBackendTest, PooledEnginesShareOneActivationScratchAllocatio 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::size_t before = 0; + ASSERT_TRUE(device_bytes_in_use(before)) << "cudaMemGetInfo failed, so this test measured nothing"; std::vector> engines; for (int i = 0; i < kEngineCount; ++i) { engines.push_back(std::make_unique()); @@ -745,7 +763,8 @@ TEST_F(SharedScratchBackendTest, PooledEnginesShareOneActivationScratchAllocatio ASSERT_EQ(engines.back()->run(stream), Error::Ok); } ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); - const std::size_t after = device_bytes_in_use(); + std::size_t after = 0; + ASSERT_TRUE(device_bytes_in_use(after)) << "cudaMemGetInfo failed, so this test measured nothing"; 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; @@ -773,8 +792,10 @@ TEST_F(SharedScratchBackendTest, PooledEnginesShareOneActivationScratchAllocatio // 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. +// The lower bound also fails if the pool were already large enough for the second +// engine, which is how this test could otherwise pass vacuously. The fixture +// empties the pool before each test, so it is the smaller engine's run below that +// establishes the size the growth has to exceed. TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItReplaces) { ASSERT_GE(big_scratch_bytes_, scratch_bytes_ + kMinMeasurableScratch) << "the two fixture engines ask for " << scratch_bytes_ << " and " << big_scratch_bytes_ @@ -815,10 +836,12 @@ TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItRep 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(); + std::size_t before = 0; + ASSERT_TRUE(device_bytes_in_use(before)) << "cudaMemGetInfo failed, so this test measured nothing"; ASSERT_EQ(big.run(stream), Error::Ok); ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); - const std::size_t after = device_bytes_in_use(); + std::size_t after = 0; + ASSERT_TRUE(device_bytes_in_use(after)) << "cudaMemGetInfo failed, so this test measured nothing"; 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; @@ -858,33 +881,24 @@ TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItRep // An engine that needs no activation scratch // --------------------------------------------------------------------------- -// A zero from updateDeviceMemorySizeForShapes() does not say which of its causes -// it is, so execute() substitutes the engine's own requirement and installs a -// buffer unless that is zero too. 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 of its own, so it no longer " - "covers the case where the engine's own requirement stands in for a zero from the " - "per-shape query"; -} - -TEST_F(SharedScratchBackendTest, EachEngineRecordsItsOwnActivationScratchRequirement) { +// The pooled path branches on this flag, and it must separate the two fixture +// networks or the branch is only ever taken one way below. +TEST_F(SharedScratchBackendTest, EachEngineRecordsWhetherItNeedsActivationScratch) { 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); + EXPECT_TRUE(needing.handle()->engine_needs_scratch) + << "the two-softmax network reports no activation scratch of its own, so no test here reaches the pooled path"; + EXPECT_FALSE(scratch_free.handle()->engine_needs_scratch) + << "the pointwise chain reports activation scratch, so it no longer covers the scratch-free case"; } // Turning the pool on must not turn an engine that legitimately needs no -// activation scratch into a failure. +// activation scratch into a failure. Such an engine skips the pool altogether, +// which is what the capacity check below pins. TEST_F(SharedScratchBackendTest, AnEngineNeedingNoActivationScratchRunsWithThePoolEnabled) { ASSERT_EQ(scratch_free_engine_bytes_, 0) << "the fixture engine needs scratch, so this test covers nothing"; @@ -905,6 +919,8 @@ TEST_F(SharedScratchBackendTest, AnEngineNeedingNoActivationScratchRunsWithThePo 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); + EXPECT_EQ(shared_scratch_capacity_for_testing(pooled.handle()->device_id), 0u) + << "an engine that needs no activation scratch under any shape claimed the pool anyway"; const std::vector actual = pooled.read_output(); ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); @@ -932,6 +948,16 @@ TEST_F(SharedScratchBackendTest, AnEngineNeedingNoActivationScratchRunsWithThePo // the empty-input case rather than either of the others: an engine that needs // nothing would fail the second, and a measurement that could not read the engine // at all would fail the third. +// +// What the empty call must not do is size the pool for itself. The engine's own +// profile-wide requirement, `dynamic_engine_bytes_`, is the figure a zero invites +// installing, and it stands far above the `dynamic_batch_scratch_bytes_` the next +// call actually needs; the pool never shrinks, so installing it would pin the pool +// at many times that for the rest of the process. The empty call is therefore +// checked twice, once against an empty pool and once against a pool holding a real +// requirement, and neither may leave the pool above what a non-empty call asked +// for. The capacity is read directly rather than inferred from device-wide +// memory, which cannot resolve the difference reliably. TEST_F(SharedScratchBackendTest, AnEmptyInputRunsWithThePoolEnabled) { ASSERT_EQ(empty_batch_scratch_bytes_, 0u) << "an empty batch reports " << empty_batch_scratch_bytes_ @@ -959,11 +985,25 @@ TEST_F(SharedScratchBackendTest, AnEmptyInputRunsWithThePoolEnabled) { LoadedEngine empty; ASSERT_EQ(empty.load(dynamic_blob(), 19, kDynRows, kDynCols, 0), Error::Ok); ASSERT_TRUE(empty.handle()->shared_scratch); + const int device = empty.handle()->device_id; + ASSERT_EQ(shared_scratch_capacity_for_testing(device), 0u) + << "the fixture left the pool holding something, so an empty call adding nothing to it proves nothing"; EXPECT_EQ(empty.run(stream), Error::Ok) << "the pool rejected an empty input to an engine that needs scratch"; EXPECT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + // Not zero: enqueueV3 refuses this context with no device memory installed + // whatever the bound shapes need, so the pool has to allocate something. What it + // must not do is take the engine's word for how much. + const std::size_t after_empty = shared_scratch_capacity_for_testing(device); + // Pinned to the exact figure rather than bounded below the 131072 bytes the next + // call needs, because a bound that loose is met by any minimum up to 128 KiB and + // the point of the minimum is that it is the smallest thing enqueueV3 accepts. + EXPECT_EQ(after_empty, 1u) << "an empty call against an empty pool took " << after_empty + << " bytes rather than the one-byte kMinPooledScratchBytes; a zero means no buffer was " + "installed, which TensorRT refuses for this engine, and anything larger means the " + "call sized the pool from a figure of its own"; // The same engine on a shape that does need scratch, after the empty call, so a - // fallback that installed an unusable buffer is not left untested. + // pool the empty call left in an unusable state is not left untested. LoadedEngine pooled; ASSERT_EQ(pooled.load(dynamic_blob(), 19, kDynRows, kDynCols, kDynBatch), Error::Ok); ASSERT_TRUE(pooled.handle()->shared_scratch); @@ -971,6 +1011,18 @@ TEST_F(SharedScratchBackendTest, AnEmptyInputRunsWithThePoolEnabled) { EXPECT_EQ(pooled.run(stream), Error::Ok); ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); const std::vector actual = pooled.read_output(); + const std::size_t after_non_empty = shared_scratch_capacity_for_testing(device); + EXPECT_EQ(after_non_empty, dynamic_batch_scratch_bytes_) + << "the pool holds " << after_non_empty << " bytes after a call needing " << dynamic_batch_scratch_bytes_ + << ", so it was not sized to what that call asked for"; + + // A second empty call, now that the pool holds a real requirement: it must be + // handed that buffer rather than grow the pool to the engine's profile-wide one. + ASSERT_EQ(empty.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + EXPECT_EQ(shared_scratch_capacity_for_testing(device), after_non_empty) + << "the empty call grew the pool from " << after_non_empty << " bytes, against the " << dynamic_engine_bytes_ + << " bytes this engine reports over its whole profile"; ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); @@ -989,6 +1041,61 @@ TEST_F(SharedScratchBackendTest, AnEmptyInputRunsWithThePoolEnabled) { EXPECT_EQ(std::memcmp(expected.data(), actual.data(), dyn_elems * sizeof(float)), 0); } +// --------------------------------------------------------------------------- +// Stream capture +// --------------------------------------------------------------------------- + +// The pool's handoff waits on an event recorded outside the capture, which +// cudaStreamWaitEvent refuses with cudaErrorStreamCaptureIsolation and which +// invalidates the capture under every capture mode. Left to run, the caller learns +// about it only when cudaStreamEndCapture returns an error and a null graph, far +// from the cause. execute() refuses instead. A growth's allocation and free +// invalidate a capture as well, but only outside cudaStreamCaptureModeRelaxed, so +// this test -- relaxed mode, and starting from a pool already large enough -- +// pins the wait rather than them. +// +// The capture is ended either way: a capture left open belongs to the stream, and +// destroying that stream would abandon it. +TEST_F(SharedScratchBackendTest, APooledEngineRefusesToRunWhileItsStreamIsCapturing) { + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + + LoadedEngine pooled; + ASSERT_EQ(pooled.load(blob(), 20), Error::Ok); + ASSERT_TRUE(pooled.handle()->shared_scratch); + + // One run outside the capture, so the pool holds a buffer and an event with a + // recording against it. Against an empty pool the handoff has nothing to wait + // for, and that wait is the call here that invalidates a capture, so the refusal + // below would be refusing a run that was never a hazard. + ASSERT_EQ(pooled.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + ASSERT_GT(shared_scratch_capacity_for_testing(pooled.handle()->device_id), 0u) + << "the pool holds nothing, so the run below has no handoff to wait on and the capture it is refused for was " + "never at risk"; + + // Relaxed, so the host-side waits execute() makes before it reaches the pool -- + // the one on the previous enqueue's completion event -- are permitted. This test + // is about what the pool does under capture, not about which of the surrounding + // calls a stricter mode would reject first. + ASSERT_EQ(cudaStreamBeginCapture(stream, cudaStreamCaptureModeRelaxed), cudaSuccess); + const Error captured = pooled.run(stream); + cudaGraph_t graph = nullptr; + const cudaError_t end_err = cudaStreamEndCapture(stream, &graph); + if (graph != nullptr) { + cudaGraphDestroy(graph); + } + + EXPECT_EQ(captured, Error::NotSupported) << "execute() did not refuse a pooled run on a capturing stream"; + // The refusal is worth nothing if it still spoiled the capture on the way out: + // a caller told no is expected to end the capture and get its graph. + EXPECT_EQ(end_err, cudaSuccess) << "the refused run left the capture invalidated: " << cudaGetErrorString(end_err); + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + // --------------------------------------------------------------------------- // The enqueue handoff, single-threaded, two caller streams // --------------------------------------------------------------------------- diff --git a/tests/cpp/executorch/test_shared_scratch_pool.cpp b/tests/cpp/executorch/test_shared_scratch_pool.cpp index 329c7b67588..ff6d84ce257 100644 --- a/tests/cpp/executorch/test_shared_scratch_pool.cpp +++ b/tests/cpp/executorch/test_shared_scratch_pool.cpp @@ -363,6 +363,77 @@ TEST(SharedScratchPoolRegistry, HandsOutOneStableEntryPerDevice) { EXPECT_EQ(distinct.size(), 512u); } +TEST(SharedScratchPoolRegistry, ResetHandsBackEverySlotAndLeavesItEmpty) { + // The backend test fixture runs this between cases, so every case that reads + // what the pool holds depends on it. The two things it has to get right are + // handing the caller both the buffer and the event to release -- neither the + // pool nor the disposer can find them afterwards -- and clearing the slot, so + // the next claim allocates instead of reusing a pointer that has been freed. + SharedScratchPool pool; + FakeAllocator zero; + FakeAllocator one; + FakeEventFactory events; + std::size_t out = 0; + + SharedScratchDevice& dev0 = pool.get(0); + SharedScratchDevice& dev1 = pool.get(1); + void* const dev0_buffer = call(dev0, zero, 4096, out); + void* const dev1_buffer = call(dev1, one, 2048, out); + const cudaEvent_t dev0_event = shared_scratch_claim_event(dev0, std::ref(events)).event; + shared_scratch_mark_in_flight(dev0); + // Device 1 is left with a buffer and no event, which is the state of a slot + // whose event creation failed: the reset has to cope with a null there. + ASSERT_NE(dev0_buffer, nullptr); + ASSERT_NE(dev1_buffer, nullptr); + ASSERT_NE(dev0_event, nullptr); + ASSERT_EQ(dev1.marker.event, nullptr); + + struct Disposal { + int device_id; + void* buffer; + cudaEvent_t event; + }; + std::vector disposed; + pool.reset_for_testing([&](int device_id, void* buffer, cudaEvent_t event) { + disposed.push_back({device_id, buffer, event}); + }); + + // The registry iterates in unspecified order, so each device is looked up. + const auto disposal_for = [&disposed](int device_id) -> const Disposal* { + for (const Disposal& d : disposed) { + if (d.device_id == device_id) { + return &d; + } + } + return nullptr; + }; + ASSERT_EQ(disposed.size(), 2u) << "the reset skipped a device's slot, whose buffer is then never freed"; + ASSERT_NE(disposal_for(0), nullptr); + ASSERT_NE(disposal_for(1), nullptr); + EXPECT_EQ(disposal_for(0)->buffer, dev0_buffer); + EXPECT_EQ(disposal_for(0)->event, dev0_event); + EXPECT_EQ(disposal_for(1)->buffer, dev1_buffer); + EXPECT_EQ(disposal_for(1)->event, nullptr); + + // The reference handed out before the reset stays valid, which is what lets the + // backend hold one across a reset without re-looking it up. + EXPECT_EQ(&pool.get(0), &dev0); + EXPECT_EQ(dev0.buffer, nullptr); + EXPECT_EQ(dev0.capacity, 0u); + EXPECT_EQ(dev0.marker.event, nullptr); + EXPECT_FALSE(dev0.marker.pending); + + // A smaller request than the freed buffer served: reuse would satisfy it from + // the stale capacity and allocate nothing, so this is what distinguishes a + // cleared slot from one the reset only emptied of its event. + std::size_t after = 0; + void* const fresh = call(dev0, zero, 1024, after); + EXPECT_NE(fresh, nullptr); + EXPECT_EQ(after, 1024u); + EXPECT_EQ(zero.alloc_count(), 2) << "the slot was not cleared, so the request reused the freed buffer"; + EXPECT_TRUE(zero.retirements.empty()) << "the cleared slot retired a buffer the reset had already handed back"; +} + TEST(SharedScratchPoolRegistry, AGrowthOnOneDeviceDoesNotBlockAClaimOnAnother) { SharedScratchPool pool; // One allocator per thread: the two claims share the registry and nothing else. From dfc97a2d11e3d6b8a60c5786ad10878e4e31e089 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Sun, 6 Sep 2026 00:37:27 -0700 Subject: [PATCH 07/13] fix(executorch): refuse a pooled scratch install TensorRT rejected A zero from `updateDeviceMemorySizeForShapes()` reaches `execute()` for two reasons and the pooled path treated both as "any buffer will do". Only one of them is safe with a buffer that small. Where the bound shapes genuinely need nothing, the engine expects nothing and the pool's minimum is accepted. Where the query *failed*, the engine still expects what it always did: `setDeviceMemoryV2` refuses the undersized buffer, returns `void`, and leaves the context pointed at whatever it was last given -- which a pool growth may already have freed. Measured on TensorRT 11.2.1.2 with this branch's own two-softmax net (expected 8388608): installing 1 byte, `enqueueV3` returns true, `cudaStreamSynchronize` reports no error, the output matches byte for byte, and the engine writes the stale buffer. Free it, let the next `cudaMalloc` take the address, and 4090 of 4096 bytes of an unrelated allocation are overwritten with no error anywhere; when the block goes back to the driver instead, `cudaErrorIllegalAddress`. The refusal is observable, through an `IErrorRecorder` attached to the context for the duration of that one call. Measured on the same TensorRT: a correct install records nothing, an undersized one records `kINVALID_ARGUMENT` naming both sizes, and the genuinely-empty case -- expected size zero, one byte installed -- records nothing, so it keeps working. `install_pooled_scratch` makes the install and reads that back; a refusal ends the call with `Error::InvalidState` and nothing reaches `enqueueV3`. The engine requirement does not have to come back as a size, so the pool is still never pinned at the profile maximum. Also here: - The pool is a leaked singleton rather than a namespace-scope object. Static destruction destroyed the registry mutex, every device mutex and the map nodes a live reference points into while a thread could still be between the lookup and its unlock: 7 of 18 runs of a driver over the real header segfaulted at exit, against 0 of 18 for the same code leaked. - The capture guard runs before `execute()` makes any call a capture cannot take, rather than after the wait on a previous enqueue and the `cudaMalloc` that grows a host-input staging buffer. Under `cudaStreamCaptureModeGlobal` those two fail first, so the refusal used to arrive with the caller's capture already invalidated -- reproduced through the delegate: `cudaEventSynchronize` returns `cudaErrorStreamCaptureUnsupported`, `execute()` answers `InvalidProgram`, and `cudaStreamEndCapture` hands back 901 and a null graph. The staging copy is not one of them: a pageable `cudaMemcpyAsync` is captured cleanly under all three modes, measured on CUDA 13.0 at 4 KiB and 64 MiB, and the device query and the device switch are captured cleanly too. - The guard still sees only the selected stream, and the header and README now say so instead of promising more. Widening it is not available: `cudaStreamIsCapturing(cudaStreamLegacy)` reports a capture on another stream only when that stream is a blocking one, and then reports it under `Relaxed` and under another thread's `ThreadLocal` too, where the pool's calls are permitted and refusing would be wrong. CUDA has no process-wide query. - `reset_for_testing` empties every slot under the locks and disposes of what it took with neither held. The production disposer runs `cudaSetDevice`, `cudaFree` and `cudaEventDestroy`, and a device-wide free blocks on everything queued on that device -- a parked host function included, which the suite parks deliberately -- so under the registry lock one device's teardown held up every other device's claims. - A failed `cudaEventSynchronize` on the buffer a growth replaces now fails the call. It is a device synchronization, so what it reports is usually an asynchronous fault already raised on that device, and continuing enqueued against a context on a device in that state. - The four comments claiming a `cudaGetLastError` stops an error resurfacing under the name of the next call said the opposite of what happens: a sticky error survives the clear and the next `cudaMalloc` still returns it. Only non-sticky errors recover that way, and the comments say that now. The free path's message no longer blames the pool for a fault the free merely surfaced. - `need = dev.capacity > 0 ? dev.capacity : kMinPooledScratchBytes` collapses to `kMinPooledScratchBytes`: the reuse branch already answers with the existing capacity for any request it covers, so the two are identical over every reachable pool state (6 of 6 measured). It read like the profile-wide sizing policy the previous commit removed. - `SharedScratchMarker`'s comment said the event is never destroyed; the reset hook destroys it. It says what is true. Tests, 70 -> 74 cases over the same 6 targets: - `AnUndersizedScratchInstallIsSeenAsRefused` drives the same TensorRT call `execute()` makes over a live `kUSER_MANAGED` context: a correct install is accepted, a one-byte install for an engine expecting more is refused, and the empty-batch context accepts the one byte. A failed per-shape query cannot be induced from inside the process, so what is pinned is the check that catches its consequence. - `APooledEngineRefusesACaptureBeforeAnythingCanInvalidateIt` captures in the default `Global` mode and requires the refused run to leave the capture intact, which the existing relaxed-mode case cannot see. - `ResetLeavesEveryEntryWhereItWas` looks up 512 entries either side of a reset. The old single-address check passes against an erasing reset, because the allocator hands the freed node straight back; over 512 entries it does not. - `ResetDisposesWithNoLockHeld` observes, from another thread and while the disposer runs, that the slot's device lock is free and that a lookup for a device the reset never touched completes. - Without a CUDA device the backend suite prints how many of its cases it skipped, and `TORCHTRT_EXECUTORCH_REQUIRE_CUDA=1` turns the skip into a failure. Getting the CI job onto a GPU runner is a separate matter. --- .../executorch/TensorRTBackend.h | 20 +- cpp/src/torch_tensorrt/executorch/README.md | 36 ++- .../executorch/SharedScratchPool.h | 69 ++++- .../executorch/TensorRTBackend.cpp | 280 +++++++++++++----- .../test_shared_scratch_backend.cpp | 166 ++++++++++- .../executorch/test_shared_scratch_pool.cpp | 95 +++++- 6 files changed, 571 insertions(+), 95 deletions(-) diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index 2c65d0e0220..0158291f4f5 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -137,13 +137,19 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { // work this call did not submit. Which calls grow the pool is not knowable // from here; see the README. // - Capturing a CUDA graph from the selected stream is refused with - // Error::NotSupported. The pool's event handoff waits on an event recorded - // outside the capture, which invalidates it under every capture mode. A - // growth's allocation, its wait on the replaced buffer and the free of that - // buffer invalidate it under every mode but cudaStreamCaptureModeRelaxed, - // which permits all three but does not record them, leaving a replay pointed - // at a buffer the pool may since have freed. The alternative to refusing is a - // capture that silently comes back invalidated. + // Error::NotSupported, ahead of every CUDA call it makes that a capture + // cannot take -- only the device query and the device switch run first. The + // pool's event handoff waits on an event recorded outside the capture, which + // invalidates it under every capture mode. A growth's allocation, its wait on + // the replaced buffer and the free of that buffer invalidate it under every + // mode but cudaStreamCaptureModeRelaxed, which permits all three but does not + // record them, leaving a replay pointed at a buffer the pool may since have + // freed. The alternative to refusing is a capture that silently comes back + // invalidated. Only a capture on the selected stream is caught. A capture + // running on any other stream under cudaStreamCaptureModeGlobal, or under + // cudaStreamCaptureModeThreadLocal from this thread, is invalidated by the + // same calls and is not refused, because CUDA offers no query for it: do not + // run a pooled engine while capturing anywhere in the process. // - cudaDeviceReset() invalidates the pool without emptying it. The buffer // and the handoff event it still holds are destroyed with the primary // context, and the next call on that device uses both. There is no guard: diff --git a/cpp/src/torch_tensorrt/executorch/README.md b/cpp/src/torch_tensorrt/executorch/README.md index 4708adf2bee..6e97250ceeb 100644 --- a/cpp/src/torch_tensorrt/executorch/README.md +++ b/cpp/src/torch_tensorrt/executorch/README.md @@ -152,8 +152,7 @@ A call that asks for nothing is handed whatever the pool already holds, at its current size, so it never grows it. Two of the things a zero from the per-shape query can mean reach this point, and are indistinguishable where it is read: the shapes bound to this call need none -- an empty batch inside a profile that admits -one -- or the query failed. Neither wants a buffer of its own, and any buffer the -pool already holds is large enough for either. +one -- or the query failed. Neither wants a buffer of its own. Such a call cannot be left with no buffer at all. `enqueueV3` refuses a `kUSER_MANAGED` context with no device memory installed as soon as the engine @@ -164,6 +163,18 @@ Standing the engine's own figure in for the zero would satisfy the check too, bu that figure covers the whole profile: on the dynamic engine the tests use it is 4 MiB against the 128 KiB the next call needs, and the pool never shrinks. +Only one of the two causes is safe to hand a buffer that small, and the install is +what separates them. `setDeviceMemoryV2` refuses a buffer smaller than the bound +shapes need, and it returns `void`: where the shapes genuinely need nothing the +expected size is zero and the one-byte buffer is accepted, but where the query +failed the engine expects what it always did, the install is refused, and the +context keeps the buffer it was last given -- which a growth may already have +freed, so the enqueue would read and write memory the pool no longer owns while +`enqueueV3` reports success. The backend reads the refusal back through a TensorRT +`IErrorRecorder` scoped to that one call and fails with `Error::InvalidState` +rather than enqueueing. Nothing reaches `enqueueV3` on an install TensorRT +rejected. + The third thing a zero can mean is an engine that needs no scratch under *any* shape, and that one is settled before a call is ever made: such an engine is left out of the pool entirely, so it takes no per-device lock and does not serialize @@ -208,9 +219,24 @@ return `cudaErrorStreamCaptureUnsupported` and invalidate the capture under the `Global` and `ThreadLocal` modes; under `Relaxed` all three are permitted and the capture survives them, but they then run uncaptured, so a replayed graph would use whatever buffer was installed when it was captured -- which by then the pool may -have freed. The backend checks the stream before it makes any of these calls and -returns instead. Load the engine with the option off to capture it: a context that -owns its scratch makes none of them. +have freed. The backend checks the selected stream and returns instead, and it +checks ahead of everything a capture cannot take -- not merely ahead of the pool's +own calls, since the wait on a previous enqueue and the `cudaMalloc` that grows a +host-input staging buffer come before those and would invalidate the capture +first. Only the device query and the device switch run before the check, and a +capture takes both. Load the engine with the option off to capture it: a context +that owns its scratch makes none of them. + +What that check does not cover is a capture running on some *other* stream. Under +`Global`, and under `ThreadLocal` from the thread that calls `execute()`, the same +calls are prohibited and invalidate that capture, and the call is not refused: +`cudaStreamIsCapturing` answers about one stream, and asking it about +`cudaStreamLegacy` instead does not close the gap. Measured on CUDA 13.0, that +reports the capture only when the capturing stream is a blocking one, and then +reports it under `Relaxed` and under another thread's `ThreadLocal` as well, where +the pool's calls are permitted and refusing would be wrong. CUDA has no query for +"is a capture live in this process", so the rule is the caller's to keep: do not +run a pooled engine while any capture is open anywhere in the process. `cudaDeviceReset()` is not survivable and is not guarded against. The pool holds its buffer and its handoff event for the process lifetime, and a reset destroys diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h index bdb7df592f6..2ad867351ba 100644 --- a/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h @@ -16,7 +16,15 @@ #include #include +#include #include +#include + +namespace nvinfer1 { +// Named by install_pooled_scratch below. Forward-declared, as TensorRT's own +// headers do, so this header still builds against the CUDA headers alone. +class IExecutionContext; +} // namespace nvinfer1 namespace torch_tensorrt { namespace executorch_backend { @@ -24,7 +32,9 @@ namespace executorch_backend { // 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 + // Nothing in the normal path destroys this event; only reset_for_testing below + // hands it to a disposer that does. + cudaEvent_t event = nullptr; bool pending = false; // an enqueue against the buffer has been recorded on `event` }; @@ -59,8 +69,9 @@ struct SharedScratchDevice { // `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. +// This one lock is shared by every device, so holding it couples every device to +// whoever holds it: it covers the lookup, and in reset_for_testing the snapshot +// that empties the slots, and nothing else. No CUDA call is made under it. class SharedScratchPool { public: SharedScratchDevice& get(int device_id) { @@ -73,20 +84,32 @@ class SharedScratchPool { // so the caller can release it. Entries stay in the map, so a reference `get` // handed out remains valid. // + // Every slot is emptied under the locks, and what came out of it is disposed of + // afterwards with neither held, because the disposer frees device memory and a + // device-wide free blocks on everything queued on that device -- a parked host + // function included. Under the device's lock that would hold off that device's + // next claimant; under this registry's, every device's. + // // Nothing here waits for an enqueue: the buffer it frees may still be in use by - // one, and the caller is responsible for there being none. `dispose` runs under - // both this registry's lock and the device's, so it must not claim either. + // one, and the caller is responsible for there being none. template void reset_for_testing(Dispose dispose) { - std::lock_guard lk(mu_); - for (auto& entry : devices_) { - SharedScratchDevice& dev = entry.second; - std::lock_guard dev_lk(dev.mu); - dispose(entry.first, dev.buffer, dev.marker.event); - dev.buffer = nullptr; - dev.capacity = 0; - dev.marker.event = nullptr; - dev.marker.pending = false; + std::vector> taken; + { + std::lock_guard lk(mu_); + taken.reserve(devices_.size()); + for (auto& entry : devices_) { + SharedScratchDevice& dev = entry.second; + std::lock_guard dev_lk(dev.mu); + taken.emplace_back(entry.first, dev.buffer, dev.marker.event); + dev.buffer = nullptr; + dev.capacity = 0; + dev.marker.event = nullptr; + dev.marker.pending = false; + } + } + for (const auto& slot : taken) { + dispose(std::get<0>(slot), std::get<1>(slot), std::get<2>(slot)); } } @@ -192,6 +215,24 @@ void* shared_scratch_get_or_grow( return p; } +// Installs `bytes` of `buffer` as the activation scratch of `ctx`, a +// kUSER_MANAGED context, and reports whether TensorRT accepted it. Logs the +// refusal, naming `device_id`, when it did not. +// +// The check is the point. setDeviceMemoryV2 returns void and refuses a buffer +// smaller than the bound shapes need, so a caller that does not ask cannot tell +// an accepted install from a refused one -- and a refused one leaves the context +// pointed at the buffer it was last given, which a pool growth may since have +// freed. The engine then reads and writes freed memory with enqueueV3 reporting +// success. The refusal is read back through an IErrorRecorder scoped to this one +// call, the only channel that hands it to the caller: with no recorder attached +// TensorRT writes it to the runtime's ILogger and this function has nothing to +// return. +// +// Defined in TensorRTBackend.cpp, which owns the TensorRT dependency; a binary +// that links this header alone never names it. +bool install_pooled_scratch(nvinfer1::IExecutionContext& ctx, void* buffer, std::size_t bytes, int device_id); + // Test-only views of the process-wide pool the TensorRT backend runs on. Declared // here, in a header that ships under src/, so they add nothing to the installed // API; defined in TensorRTBackend.cpp, which owns the pool instance they act on. diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index 38355ee1963..108d2be5d1f 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -246,7 +246,16 @@ bool is_cuda_accessible_ptr(const void* ptr) { // 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; +// +// The C++ object is never destroyed either, and that is not the same claim. +// Static destruction would destroy the registry's mutex, every device's mutex +// and the map nodes a live reference points into, while a thread anywhere +// between the lookup and the release of its device lock still holds them. +// Leaking it costs one allocation and removes that race. +SharedScratchPool& scratch_pool() { + static SharedScratchPool* const pool = new SharedScratchPool(); + return *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. @@ -270,7 +279,7 @@ class SharedScratchClaim { } SharedScratchDevice& hold(int device_id) { - dev_ = &scratch_pool.get(device_id); + dev_ = &scratch_pool().get(device_id); device_id_ = device_id; lock_ = std::unique_lock(dev_->mu); return *dev_; @@ -305,17 +314,20 @@ class SharedScratchClaim { } 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) { + // cudaFree synchronizes, so what it reports is more often an earlier + // asynchronous fault on this device than a fault in the free -- which is + // why the message does not call it one. ET_LOG( Error, - "TensorRTBackend::execute: freeing the shared activation scratch buffer that a pool growth replaced on device %d failed: %s", + "TensorRTBackend::execute: freeing the shared activation scratch buffer that a pool growth replaced on device %d reported %s; a device-wide free reports whatever fault this device is already in, so this need not be the pool's", device_id_, cudaGetErrorString(err)); - cudaGetLastError(); // clear sticky error; the free is cleanup, so execute() continues + // Clears a non-sticky error so it does not resurface under the name of the + // next CUDA call in execute(). A sticky one survives the clear and will + // resurface anyway; the caller learns of it from that call. + cudaGetLastError(); } retired_ = nullptr; } @@ -338,15 +350,70 @@ class SharedScratchClaim { // on the strength of a call that uses none of it. constexpr size_t kMinPooledScratchBytes = 1; +// Refuses a pooled call whose stream is capturing a CUDA graph. +// +// The handoff's event wait is on an event recorded outside the capture: +// cudaStreamWaitEvent fails it with cudaErrorStreamCaptureIsolation and +// invalidates the capture under every capture mode. A growth's allocation, the +// cudaEventSynchronize on the buffer it replaces and the cudaFree of that buffer +// invalidate it too, but only outside cudaStreamCaptureModeRelaxed, which permits +// all three -- and then runs them uncaptured, leaving a replay pointed at a +// buffer the pool may since have freed. None of it fails cleanly: the caller +// learns of an invalidation only when cudaStreamEndCapture hands back an error +// and a null graph. Refusing names the cause instead. +// +// execute() calls this ahead of every CUDA call it makes that a capture cannot +// take, not just the pool's own: the wait on a previous enqueue invalidates a +// capture under every mode, and the cudaMalloc that grows a host-input staging +// buffer under every mode but Relaxed. Either leaves a refusal made after it +// nothing to save. Only the device query and the device switch run earlier, and +// a capture takes both. +// +// It sees only `stream`. A capture live on some *other* stream is invalidated by +// the pool's calls just the same under the Global mode, and under ThreadLocal +// from the same thread, and this does not catch it: cudaStreamIsCapturing +// answers per stream, and asking it about cudaStreamLegacy instead does not +// close the gap -- measured on CUDA 13.0, that reports the capture only when the +// capturing stream is a blocking one, and then reports it under Relaxed and +// under another thread's ThreadLocal too, where the pool's calls are permitted +// and refusing would be wrong. There is no query for "is any capture live in +// this process", so the gap stays; the header and the README say so. +// +// A failed query is refused with the rest. Where `stream` is the legacy stream +// and some other stream is capturing, what it reports is +// cudaErrorStreamCaptureImplicit -- a capture somewhere, not an unrelated fault. +// Refusing on that is over-conservative where the capture is one the pool's calls +// would have survived, and never the other way round. +Error refuse_pooled_call_on_a_capturing_stream(cudaStream_t stream, int device_id) { + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + const cudaError_t capture_err = cudaStreamIsCapturing(stream, &capture); + if (capture_err == cudaSuccess && capture == cudaStreamCaptureStatusNone) { + return Error::Ok; + } + ET_LOG( + Error, + "TensorRTBackend::execute: the selected stream is capturing a CUDA graph (%s), which the shared activation scratch pool on device %d does not support. Load this engine with '%s' off so its context keeps its own scratch.", + capture_err == cudaSuccess ? "capture in progress" : cudaGetErrorString(capture_err), + device_id, + kSharedActivationScratchKey); + cudaGetLastError(); // the query's own failure is this call's, not the next one's + return Error::NotSupported; +} + // 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. // 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. // // A `need` of zero asks for whatever the pool already holds rather than for a -// buffer of its own: any live buffer is large enough, so nothing is allocated and -// nothing grows. Only against an empty pool does it allocate, and then at -// kMinPooledScratchBytes. +// buffer of its own: any live buffer is large enough for a call that needs +// nothing, so nothing is allocated and nothing grows. Only against an empty pool +// does it allocate, and then at kMinPooledScratchBytes. +// +// The caller must already have refused a capturing `stream`: the handoff's wait +// invalidates a capture under every mode, and a growth's allocation, its wait on +// the buffer it replaces and its free of that buffer under every mode but +// Relaxed. // // Must be called with `device_id` already current: cudaEventCreateWithFlags and // cudaMalloc both act on the *current* device and nothing in here sets it. @@ -357,35 +424,14 @@ Error get_or_grow_shared_scratch( cudaStream_t stream, void*& out_ptr, size_t& out_size) { - // Checked before the lock, and so ahead of every CUDA call the pooled path makes - // that a capture cannot take. The handoff's event wait is on an event recorded - // outside the capture: cudaStreamWaitEvent fails it with - // cudaErrorStreamCaptureIsolation and invalidates the capture under every capture - // mode. The allocation, the cudaEventSynchronize on the replaced buffer and the - // cudaFree of it invalidate the capture too, but only outside - // cudaStreamCaptureModeRelaxed, which permits all three -- and then runs them - // uncaptured, leaving a replay pointed at a buffer the pool may since have freed. - // None of it fails cleanly: the caller learns of an invalidation only when - // cudaStreamEndCapture hands back an error and a null graph. Refusing names the - // cause instead. A failed query is refused with the rest: what it reports on the - // legacy stream while another stream captures is cudaErrorStreamCaptureImplicit, - // which is the hazard rather than an unrelated fault. - cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; - const cudaError_t capture_err = cudaStreamIsCapturing(stream, &capture); - if (capture_err != cudaSuccess || capture != cudaStreamCaptureStatusNone) { - ET_LOG( - Error, - "TensorRTBackend::execute: the selected stream is capturing a CUDA graph (%s), which the shared activation scratch pool on device %d does not support. Load this engine with '%s' off so its context keeps its own scratch.", - capture_err == cudaSuccess ? "capture in progress" : cudaGetErrorString(capture_err), - device_id, - kSharedActivationScratchKey); - return Error::NotSupported; - } - SharedScratchDevice& dev = claim.hold(device_id); + // kMinPooledScratchBytes rather than the capacity the pool already holds: the + // reuse branch below answers with that capacity for any request it covers, so + // the two ask for the same buffer, and the smallest request says plainly that + // this call is not sizing the pool. if (need == 0) { - need = dev.capacity > 0 ? dev.capacity : kMinPooledScratchBytes; + need = kMinPooledScratchBytes; } const SharedScratchHandoff handoff = shared_scratch_claim_event(dev, []() -> cudaEvent_t { @@ -455,17 +501,19 @@ Error get_or_grow_shared_scratch( 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 -- one buffer for each - // growth whose wait fails. How many growths a run sees is not bounded by the - // engine count: the requirement is re-queried every execute() below, so an - // engine that answers with what the bound shapes need can grow the pool on - // any call. + // still be reading, so a failed wait leaks it instead. + // + // The call fails with it. cudaEventSynchronize waits on device work, so what + // it reports is usually an asynchronous fault raised by earlier work on this + // device; running on regardless would enqueue against a context on a device + // already in that state and report the same fault later, under the name of + // whichever call surfaced it next. 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", + "TensorRTBackend::execute: waiting for the enqueue on the replaced shared activation scratch on device %d failed (%s), which for a wait on device work is usually an earlier asynchronous fault on this device surfacing here; leaking that buffer rather than freeing it under a live enqueue", device_id, cudaGetErrorString(err)); - cudaGetLastError(); // clear sticky error; execute() continues regardless + return Error::InvalidProgram; } } @@ -502,13 +550,89 @@ Error mark_shared_scratch_in_flight(SharedScratchClaim& claim, cudaStream_t stre return Error::Ok; } +// Collects what TensorRT reports while it is attached. +// +// install_pooled_scratch attaches one across a single API call and detaches it +// again, so anything in it came from that call. The detach is not tidiness: while +// a recorder is attached TensorRT reports to it *instead of* to the ILogger, so +// one left in place would divert the diagnostics for setInputShape, +// setTensorAddress and enqueueV3 away from the backend's logger. +class ScopedErrorRecorder final : public nvinfer1::IErrorRecorder { + public: + int32_t getNbErrors() const noexcept override { + return static_cast(errors_.size()); + } + nvinfer1::ErrorCode getErrorCode(int32_t index) const noexcept override { + return in_range(index) ? errors_[static_cast(index)].first : nvinfer1::ErrorCode::kSUCCESS; + } + ErrorDesc getErrorDesc(int32_t index) const noexcept override { + return in_range(index) ? errors_[static_cast(index)].second.c_str() : ""; + } + bool hasOverflowed() const noexcept override { + return overflowed_; + } + void clear() noexcept override { + errors_.clear(); + } + bool reportError(nvinfer1::ErrorCode code, ErrorDesc desc) noexcept override { + // noexcept, and a throwing push_back here would terminate: an allocation + // failure while reporting an error costs the description, not the process. + try { + errors_.emplace_back(code, std::string(desc == nullptr ? "" : desc)); + } catch (...) { + overflowed_ = true; + } + return false; // false asks TensorRT to keep going; the caller decides + } + RefCount incRefCount() noexcept override { + return ++refs_; + } + RefCount decRefCount() noexcept override { + return --refs_; + } + + bool anything_reported() const { + return !errors_.empty() || overflowed_; + } + + private: + bool in_range(int32_t index) const { + return index >= 0 && static_cast(index) < errors_.size(); + } + + std::vector> errors_; + bool overflowed_ = false; + RefCount refs_ = 1; +}; + } // namespace +// Declared in SharedScratchPool.h. See that header for why the install is +// checked rather than made and trusted. +bool install_pooled_scratch(nvinfer1::IExecutionContext& ctx, void* buffer, size_t bytes, int device_id) { + ScopedErrorRecorder recorder; + nvinfer1::IErrorRecorder* const previous = ctx.getErrorRecorder(); + ctx.setErrorRecorder(&recorder); + ctx.setDeviceMemoryV2(buffer, static_cast(bytes)); + ctx.setErrorRecorder(previous); + + if (!recorder.anything_reported()) { + return true; + } + ET_LOG( + Error, + "TensorRTBackend::execute: TensorRT refused the %zu-byte shared activation scratch buffer for device %d: %s", + bytes, + device_id, + recorder.getNbErrors() > 0 ? recorder.getErrorDesc(0) : "the refusal could not be recorded"); + return false; +} + // Declared in SharedScratchPool.h, which ships under src/ and is not installed; // defined here because this is where the pool instance lives. See that header for // what a caller owes them. std::size_t shared_scratch_capacity_for_testing(int device_id) { - SharedScratchDevice& dev = scratch_pool.get(device_id); + SharedScratchDevice& dev = scratch_pool().get(device_id); std::lock_guard lk(dev.mu); return dev.capacity; } @@ -516,7 +640,7 @@ std::size_t shared_scratch_capacity_for_testing(int device_id) { void reset_shared_scratch_pool_for_testing() { int restore_to = 0; const bool have_current = cudaGetDevice(&restore_to) == cudaSuccess; - scratch_pool.reset_for_testing([](int device_id, void* buffer, cudaEvent_t event) { + scratch_pool().reset_for_testing([](int device_id, void* buffer, cudaEvent_t event) { if (buffer == nullptr && event == nullptr) { return; } @@ -531,7 +655,9 @@ void reset_shared_scratch_pool_for_testing() { if (event != nullptr) { (void)cudaEventDestroy(event); } - (void)cudaGetLastError(); // a reset is cleanup; do not leave a sticky error for the next call + // Clears a non-sticky error so a reset does not leave one for the next call to + // report. A sticky one survives the clear, and no cleanup here recovers it. + (void)cudaGetLastError(); }); if (have_current) { (void)cudaSetDevice(restore_to); @@ -943,6 +1069,27 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* nvinfer1::IExecutionContext* ctx = engine->exec_ctx.get(); TORCHTRT_ET_CHECK_NOT_NULL(ctx, Error::InvalidState, "TensorRTBackend::execute: backend is not initialized"); + const auto caller_stream = ::executorch::extension::cuda::getCallerStream(); + const bool caller_stream_set = caller_stream.has_value(); + cudaStream_t stream = caller_stream.value_or(cudaStreamPerThread); + + // An engine that needs no scratch under any shape is left out of the pool + // entirely: enqueueV3 accepts it with no device memory installed, so it need not + // claim the device and does not serialize against the engines that do. + const bool pooled_scratch = engine->shared_scratch && engine->engine_needs_scratch; + + // Refused here, ahead of every call this function makes that a capture cannot + // take. The pool's own are not the only ones: the wait on a previous enqueue just + // below, and the cudaMalloc that grows a host-input staging buffer further down, + // would invalidate the capture before the pooled path was ever reached, so a + // refusal any later would arrive after the thing it exists to protect was gone. + if (pooled_scratch) { + const Error capture_err = refuse_pooled_call_on_a_capturing_stream(stream, engine->device_id); + if (capture_err != Error::Ok) { + return capture_err; + } + } + // A prior fast-path execute() may have returned with its enqueue still in flight // on the shared exec_ctx. Wait for it before reconfiguring the context below: // TensorRT forbids mutating a context while one of its enqueues is in flight, and @@ -955,9 +1102,6 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* return Error::InvalidProgram; } } - const auto caller_stream = ::executorch::extension::cuda::getCallerStream(); - const bool caller_stream_set = caller_stream.has_value(); - cudaStream_t stream = caller_stream.value_or(cudaStreamPerThread); bool output_staged_to_host = false; bool input_staged_from_host = false; @@ -1241,26 +1385,28 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // 4. Back activation scratch with the shared per-device pool // ------------------------------------------------------------------ // 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. + // Whatever it answers is binding rather than advisory: an engine backed by less + // than it asked for writes past the end. // // The buffer is installed on every call, not once, because any call needing more // than the pool holds -- this engine on other shapes, or another one -- grows it - // and moves it. A kSTATIC context owns its private scratch, so setDeviceMemoryV2 - // must not be called on one. + // and moves it. A kSTATIC context owns its private scratch, so the install must + // not be made on one. // - // An engine that needs no scratch under any shape is left out of the pool - // entirely: enqueueV3 accepts it with no device memory installed, so it need not - // claim the device and does not serialize against the engines that do. + // A reported zero has two causes that reach here and nothing distinguishes them: + // a call whose bound shapes need none -- an empty input inside a profile that + // admits one -- and a query that failed. Neither is asked to size the pool: + // get_or_grow_shared_scratch hands a zero whatever the pool already holds, and + // only where it holds nothing does it allocate, at the minimum TensorRT will + // accept rather than at the engine's profile-wide figure. // - // For the rest, a reported zero has more than one cause and nothing here tells - // them apart: a failed query, and a call whose bound shapes need none -- an - // empty input inside a profile that admits one. Neither asks for a buffer of its - // own, so neither grows the pool: get_or_grow_shared_scratch hands a zero - // whatever the pool already holds, and only where it holds nothing does it - // allocate, at the minimum TensorRT will accept rather than at the engine's - // profile-wide figure. + // The two causes are told apart by the install rather than by the query, which + // is the only place the difference is observable. Where the shapes genuinely + // need nothing the engine expects nothing and the minimum is accepted; where the + // query failed the engine still expects what it always did, and a pool holding + // less than that is refused -- and this call ends here rather than enqueueing + // against whatever pointer the context was last given, which a growth may since + // have freed. // // 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 @@ -1268,7 +1414,7 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // restore above, so its free lands on the right device. SharedScratchClaim scratch_claim; bool scratch_from_pool = false; - if (engine->shared_scratch && engine->engine_needs_scratch) { + if (pooled_scratch) { const size_t need = ctx->updateDeviceMemorySizeForShapes(); void* pool = nullptr; size_t pool_size = 0; @@ -1277,8 +1423,10 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* if (scratch_err != Error::Ok) { return scratch_err; } + if (!install_pooled_scratch(*ctx, pool, pool_size, engine->device_id)) { + return Error::InvalidState; + } scratch_from_pool = true; - ctx->setDeviceMemoryV2(pool, static_cast(pool_size)); } // ------------------------------------------------------------------ diff --git a/tests/cpp/executorch/test_shared_scratch_backend.cpp b/tests/cpp/executorch/test_shared_scratch_backend.cpp index dd731b9d348..a8d82034b4a 100644 --- a/tests/cpp/executorch/test_shared_scratch_backend.cpp +++ b/tests/cpp/executorch/test_shared_scratch_backend.cpp @@ -19,7 +19,12 @@ // // 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. +// green run on a host with no GPU says nothing about the pool -- and a skipped +// gtest case exits zero, which Bazel reports as a passing target. A run that is +// meant to have a device says so through TORCHTRT_EXECUTORCH_REQUIRE_CUDA, and +// then a skip is a failure instead; the count of skipped cases is printed either +// way. Making the CI job actually run on a GPU runner is a separate matter, and +// until it does the strongest evidence for this file is hand measurement. #include "torch_tensorrt/executorch/SharedScratchPool.h" #include "torch_tensorrt/executorch/TensorRTBackend.h" @@ -46,6 +51,7 @@ #include #include #include +#include #include #include #include @@ -490,6 +496,46 @@ bool device_bytes_in_use(std::size_t& out) { return true; } +// Opens a live kUSER_MANAGED context over `blob`'s engine with one shape bound, +// the state execute() installs a scratch buffer into. Held by the caller, unlike +// measure_engine_scratch's, which is gone by the time it returns its figure. +struct LiveContext { + TRTUniquePtr runtime; + TRTUniquePtr engine; + TRTUniquePtr ctx; +}; + +bool open_user_managed_context(const std::vector& blob, int rows, int cols, int batch, LiveContext& out) { + static BuilderLogger logger; + TensorRTBlobHeader header; + if (!TensorRTBlobHeader::parse(blob.data(), blob.size(), header)) { + return false; + } + out.runtime.reset(nvinfer1::createInferRuntime(logger)); + if (out.runtime == nullptr) { + return false; + } + out.engine.reset( + out.runtime->deserializeCudaEngine(TensorRTBlobHeader::engine_data(blob.data(), header), header.engine_size)); + if (out.engine == nullptr) { + return false; + } + out.ctx.reset(out.engine->createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED)); + if (out.ctx == nullptr) { + return false; + } + return out.ctx->setInputShape("input_0", nvinfer1::Dims3{batch, rows, cols}); +} + +// Set by a run that is supposed to have a GPU. A skip is then a failure, rather +// than a green target that covered nothing. +constexpr char kRequireCudaEnvVar[] = "TORCHTRT_EXECUTORCH_REQUIRE_CUDA"; + +bool cuda_device_is_required() { + const char* const value = std::getenv(kRequireCudaEnvVar); + return value != nullptr && value[0] != '\0' && std::strcmp(value, "0") != 0; +} + Error set_shared_scratch(TensorRTBackend& backend, bool enabled) { BackendOption option; std::strncpy(option.key, kOptionKey, sizeof(option.key) - 1); @@ -528,9 +574,30 @@ class SharedScratchBackendTest : public ::testing::Test { dynamic_engine_bytes_ = engine_scratch_requirement(dynamic_blob_); } + static void TearDownTestSuite() { + if (skipped_for_no_device_ == 0) { + return; + } + const ::testing::TestSuite* const suite = ::testing::UnitTest::GetInstance()->current_test_suite(); + std::fprintf( + stderr, + "[ SKIPPED ] %d of %d cases in this file: no CUDA device. Nothing here ran, so this target passing says " + "nothing about the shared activation scratch pool. Set %s=1 on a run that is meant to have a device and a " + "skip becomes a failure.\n", + skipped_for_no_device_, + suite == nullptr ? skipped_for_no_device_ : suite->total_test_count(), + kRequireCudaEnvVar); + } + void SetUp() override { int device_count = 0; if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count == 0) { + ++skipped_for_no_device_; + if (cuda_device_is_required()) { + FAIL() << "no CUDA device, and " << kRequireCudaEnvVar + << " says this run must have one. Every case in this file needs a device, so without one the binary " + "would exit zero having covered nothing."; + } 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"; @@ -582,6 +649,7 @@ class SharedScratchBackendTest : public ::testing::Test { static std::size_t dynamic_batch_scratch_bytes_; static std::int64_t scratch_free_engine_bytes_; static std::int64_t dynamic_engine_bytes_; + static int skipped_for_no_device_; }; std::vector SharedScratchBackendTest::blob_; @@ -594,6 +662,7 @@ std::size_t SharedScratchBackendTest::empty_batch_scratch_bytes_ = 0; std::size_t SharedScratchBackendTest::dynamic_batch_scratch_bytes_ = 0; std::int64_t SharedScratchBackendTest::scratch_free_engine_bytes_ = -1; std::int64_t SharedScratchBackendTest::dynamic_engine_bytes_ = -1; +int SharedScratchBackendTest::skipped_for_no_device_ = 0; // --------------------------------------------------------------------------- // set_option @@ -1041,6 +1110,52 @@ TEST_F(SharedScratchBackendTest, AnEmptyInputRunsWithThePoolEnabled) { EXPECT_EQ(std::memcmp(expected.data(), actual.data(), dyn_elems * sizeof(float)), 0); } +// The install is what separates the two causes of a zero from the per-shape +// query. Both reach execute() as a zero and neither is asked to size the pool, so +// both can be handed a buffer as small as one byte; only one of them is safe with +// it. Where the bound shapes genuinely need nothing the engine expects nothing +// and the one-byte buffer is accepted. Where the query failed the engine expects what it +// always did, TensorRT refuses the install through a call that returns void, and +// the context keeps the buffer it was last given -- which a growth may already +// have freed, and enqueueV3 then reports success while the engine reads and +// writes memory the pool no longer owns. +// +// The failed query itself cannot be induced from inside this process, so what is +// pinned here is the check that catches its consequence, over the same TensorRT +// call execute() makes. +TEST_F(SharedScratchBackendTest, AnUndersizedScratchInstallIsSeenAsRefused) { + LiveContext needs_scratch; + ASSERT_TRUE(open_user_managed_context(blob(), kRows, kCols, 1, needs_scratch)) + << "could not open a user-managed context over the fixture engine"; + const std::size_t need = needs_scratch.ctx->updateDeviceMemorySizeForShapes(); + ASSERT_GT(need, 1u) << "the fixture engine needs no more than one byte of activation scratch for this shape, so an " + "undersized install cannot be built out of it"; + + void* enough = nullptr; + ASSERT_EQ(cudaMalloc(&enough, need), cudaSuccess); + EXPECT_TRUE(install_pooled_scratch(*needs_scratch.ctx, enough, need, 0)) + << "a buffer of exactly what the bound shapes need was reported as refused, which would fail every pooled call"; + + void* one_byte = nullptr; + ASSERT_EQ(cudaMalloc(&one_byte, 1), cudaSuccess); + EXPECT_FALSE(install_pooled_scratch(*needs_scratch.ctx, one_byte, 1, 0)) + << "TensorRT refused a one-byte buffer for a context expecting " << need + << " bytes and the backend read the install as accepted, so the enqueue would run on the buffer installed " + "before it"; + + // The safe cause, which has to keep working: an empty batch expects nothing, so + // the one-byte buffer is accepted and no refusal may be invented for it. + LiveContext empty_batch; + ASSERT_TRUE(open_user_managed_context(dynamic_blob(), kDynRows, kDynCols, 0, empty_batch)); + ASSERT_EQ(empty_batch.ctx->updateDeviceMemorySizeForShapes(), 0u) + << "an empty batch does not answer zero for this engine, so this half covers nothing"; + EXPECT_TRUE(install_pooled_scratch(*empty_batch.ctx, one_byte, 1, 0)) + << "the one-byte buffer an empty call is handed was reported as refused, which would fail a call that is safe"; + + EXPECT_EQ(cudaFree(enough), cudaSuccess); + EXPECT_EQ(cudaFree(one_byte), cudaSuccess); +} + // --------------------------------------------------------------------------- // Stream capture // --------------------------------------------------------------------------- @@ -1096,6 +1211,55 @@ TEST_F(SharedScratchBackendTest, APooledEngineRefusesToRunWhileItsStreamIsCaptur ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); } +// Where in execute() the refusal comes, which the relaxed-mode case above cannot +// see. Under cudaStreamCaptureModeRelaxed everything execute() does before it +// reaches the pool is permitted, so a guard sitting anywhere ahead of the pool +// looks the same. Under the default Global mode it does not: the host wait on the +// previous enqueue, and further down the cudaMalloc that grows a host-input +// staging buffer, are prohibited and invalidate the capture where they stand. A +// refusal after either is a clean error handed to a caller whose capture is +// already dead, which is the outcome the guard exists to prevent. +// +// The first run is what arms that wait -- it returns with its enqueue still in +// flight, and the next call waits for it before touching the context. +TEST_F(SharedScratchBackendTest, APooledEngineRefusesACaptureBeforeAnythingCanInvalidateIt) { + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + // A node of the caller's own, so a null graph below means the capture was + // invalidated rather than that nothing was ever captured. + void* captured_target = nullptr; + ASSERT_EQ(cudaMalloc(&captured_target, 16), cudaSuccess); + + LoadedEngine pooled; + ASSERT_EQ(pooled.load(blob(), 21), Error::Ok); + ASSERT_TRUE(pooled.handle()->shared_scratch); + ASSERT_EQ(pooled.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + ASSERT_TRUE(pooled.handle()->inflight_pending) + << "the first run left no enqueue in flight, so the next call makes no host wait and this case pins nothing"; + + ASSERT_EQ(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal), cudaSuccess); + ASSERT_EQ(cudaMemsetAsync(captured_target, 0, 16, stream), cudaSuccess); + const Error captured = pooled.run(stream); + cudaGraph_t graph = nullptr; + const cudaError_t end_err = cudaStreamEndCapture(stream, &graph); + const bool have_graph = graph != nullptr; + if (have_graph) { + cudaGraphDestroy(graph); + } + + EXPECT_EQ(captured, Error::NotSupported) + << "execute() did not refuse a pooled run on a stream capturing in the default mode"; + EXPECT_EQ(end_err, cudaSuccess) << "the refused run had already made a call the capture could not take: " + << cudaGetErrorString(end_err); + EXPECT_TRUE(have_graph) << "the capture ended with no graph, so the run invalidated it before refusing"; + + EXPECT_EQ(cudaFree(captured_target), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + // --------------------------------------------------------------------------- // The enqueue handoff, single-threaded, two caller streams // --------------------------------------------------------------------------- diff --git a/tests/cpp/executorch/test_shared_scratch_pool.cpp b/tests/cpp/executorch/test_shared_scratch_pool.cpp index ff6d84ce257..51b300ba26e 100644 --- a/tests/cpp/executorch/test_shared_scratch_pool.cpp +++ b/tests/cpp/executorch/test_shared_scratch_pool.cpp @@ -415,8 +415,10 @@ TEST(SharedScratchPoolRegistry, ResetHandsBackEverySlotAndLeavesItEmpty) { EXPECT_EQ(disposal_for(1)->buffer, dev1_buffer); EXPECT_EQ(disposal_for(1)->event, nullptr); - // The reference handed out before the reset stays valid, which is what lets the - // backend hold one across a reset without re-looking it up. + // The slot the reset cleared is the one a later lookup finds, so the reads below + // are of the entry the backend would go on using. One address cannot tell a + // surviving entry from a recycled allocation, so the invariant itself is pinned + // by ResetLeavesEveryEntryWhereItWas rather than here. EXPECT_EQ(&pool.get(0), &dev0); EXPECT_EQ(dev0.buffer, nullptr); EXPECT_EQ(dev0.capacity, 0u); @@ -434,6 +436,95 @@ TEST(SharedScratchPoolRegistry, ResetHandsBackEverySlotAndLeavesItEmpty) { EXPECT_TRUE(zero.retirements.empty()) << "the cleared slot retired a buffer the reset had already handed back"; } +// The registry hands out a reference and drops its lock, so callers go on using +// that reference; the reset has to empty a slot without moving it. Checking one +// address cannot pin that. An erasing reset returns the node to the allocator, +// which hands the very same address back to the next lookup, so the comparison +// still holds -- and every read through the old reference between the two is of +// freed memory. Enough entries and the allocator cannot reproduce them all. +TEST(SharedScratchPoolRegistry, ResetLeavesEveryEntryWhereItWas) { + constexpr int kDevices = 512; + SharedScratchPool pool; + FakeAllocator alloc; + FakeEventFactory events; + std::size_t out = 0; + + std::vector before; + before.reserve(kDevices); + for (int id = 0; id < kDevices; ++id) { + SharedScratchDevice& dev = pool.get(id); + // Give each slot something, so the reset has work to do on all of them rather + // than skipping past empty ones. + ASSERT_NE(call(dev, alloc, 1024, out), nullptr); + ASSERT_NE(shared_scratch_claim_event(dev, std::ref(events)).event, nullptr); + before.push_back(&dev); + } + + int disposed = 0; + pool.reset_for_testing([&disposed](int, void*, cudaEvent_t) { ++disposed; }); + ASSERT_EQ(disposed, kDevices); + + int moved = 0; + for (int id = 0; id < kDevices; ++id) { + if (&pool.get(id) != before[static_cast(id)]) { + ++moved; + } + } + EXPECT_EQ(moved, 0) << moved << " of " << kDevices + << " entries moved, so a reference the registry handed out before the reset names freed memory " + "or another device's slot"; +} + +// The reset's disposer frees device memory, and a device-wide free waits on +// everything queued on that device -- a parked host function included. Under +// either lock that would hold up whatever the lock covers: the device's own next +// claimant under the device lock, and every device's claimants under the +// registry's. So each slot is emptied under the locks and what came out of it is +// disposed of with neither held. +TEST(SharedScratchPoolRegistry, ResetDisposesWithNoLockHeld) { + constexpr int kUntouchedDevice = 99; + SharedScratchPool pool; + FakeAllocator alloc; + std::size_t out = 0; + + SharedScratchDevice& dev0 = pool.get(0); + ASSERT_NE(call(dev0, alloc, 1024, out), nullptr); + + std::atomic disposing{false}; + std::atomic claimed{false}; + std::atomic device_lock_free{false}; + std::thread claimer([&] { + while (!disposing.load()) { + std::this_thread::yield(); + } + if (dev0.mu.try_lock()) { + device_lock_free.store(true); + dev0.mu.unlock(); + } + pool.get(kUntouchedDevice); + claimed.store(true); + }); + + // Read inside the disposer: the claim completes once the reset returns either + // way, so only what was true while the disposer ran tells the two apart. + bool claimed_during_dispose = false; + pool.reset_for_testing([&](int, void*, cudaEvent_t) { + disposing.store(true); + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (!claimed.load() && std::chrono::steady_clock::now() < deadline) { + std::this_thread::yield(); + } + claimed_during_dispose = claimed.load(); + }); + claimer.join(); + + EXPECT_TRUE(device_lock_free.load()) << "the disposer ran holding the slot's own device lock, so a claim on that " + "device waits for a free that waits on the whole device"; + EXPECT_TRUE(claimed_during_dispose) << "a lookup for a device the reset never touched could not complete while the " + "disposer ran, so the registry's lock was held across it and one device's " + "teardown blocks every other device"; +} + TEST(SharedScratchPoolRegistry, AGrowthOnOneDeviceDoesNotBlockAClaimOnAnother) { SharedScratchPool pool; // One allocator per thread: the two claims share the registry and nothing else. From 070456f1f44a0c52b0e223ac10cf5267779b7261 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Sun, 6 Sep 2026 23:51:37 -0700 Subject: [PATCH 08/13] fix(executorch): correct the capture advice, arm the CI CUDA check, unship the test hooks Five changes to the shared activation-scratch pool. **The capture advice was wrong.** The README's closing sentence and the refusal's log text both told a caller to load with the option off in order to capture. The guard is inside the `if (pooled_scratch)` branch, so with the option off nothing refuses -- and `execute()` still makes calls a capture cannot take. Measured on CUDA 13.0, one process per cell, on a non-blocking stream: `cudaEventSynchronize`, `cudaMalloc` and `cudaFree` are prohibited under `Global` and `ThreadLocal` and permitted under `Relaxed`, while `cudaStreamSynchronize` on the capturing stream is prohibited under **every** mode, `Relaxed` included. So with the option off a call captures cleanly when it is submitted on the capturing stream through a `CallerStreamGuard`, binds only non-empty device-resident tensors, aliases no output, and follows no call that left an enqueue in flight -- and since such a call leaves its own enqueue in flight, under `Global` and `ThreadLocal` a handle is good for one captured call rather than for capture in general. The README, the installed header and the refusal message now say that, and two new tests pin both outcomes. Moving the guard out of the pooled branch was the alternative. It is declined: it would refuse exactly the calls above, for contexts that never touch the pool. Mutating the guard to be unconditional fails both new tests, which is that behaviour change made visible. **`cudaFree`'s device-wide wait is documented, not fixed.** Re-measured: with a host function parked on a stream neither the pool nor the growing engine had used, the growth's `cudaMalloc` and its wait on the retired buffer each returned at once and the `cudaFree` returned after 3000 ms. Nothing later on the path already pays that wait, so deferring the free only moves it, and `cudaFreeAsync` pairs only with `cudaMallocAsync`, so a stream-ordered free would move every pool allocation onto a different allocator. The installed contract already called this "an unbounded wait on work this call did not submit"; it now says that unbounded is literal -- if some queued stream is waiting on work only this thread submits after `execute()` returns, the call does not come back. The backend's own tests deadlocked on exactly that once. **The CI check was armed by nothing.** `TORCHTRT_EXECUTORCH_REQUIRE_CUDA` turns a device-less skip into a failure, and no workflow set it. With `CUDA_VISIBLE_DEVICES=''` the backend target reported 19 of 19 skipped and exited 0; with the variable set it exits 1 with 19 failures. The CI invocation now passes it. **The pool's test hooks no longer ship.** `shared_scratch_capacity_for_testing` and `reset_shared_scratch_pool_for_testing` were defined in `TensorRTBackend.cpp` and so exported from the archive the README tells a consumer to build, one of them freeing the live pool with no wait for work in flight. They move to `SharedScratchPoolTestHooks.{h,cpp}`, compiled by a `testonly` Bazel target and by neither released target, so a release build has no definition and no symbol -- checked with `nm`/`readelf` on the release object. Reaching the pool from a separate translation unit is why `scratch_pool()` is now `inline` in `SharedScratchPool.h`; that was cheaper than a macro guard, which would have needed a build-system flag or a second copy of the backend target. **Three of this branch's earlier fixes had no discriminating test.** Each now has one, and each was proven by mutating the fix away: - `ExecuteInstallsPooledScratchThroughTheCheckedHelper`. Forcing the refusal through `execute()` is not constructible -- the size installed is the pool's capacity, which is never below what the query just returned. What is observable is that `install_pooled_scratch` scopes an `IErrorRecorder` over the install: a recorder the test leaves on the delegate's context sees itself replaced and restored exactly once per pooled run. Replacing the checked call with a bare `setDeviceMemoryV2` takes both counts to zero. - `TheProcessPoolOutlivesStaticDestructionUnderALiveClaimant`, in the CPU-only pool target. Forks children that leave a thread inside the pooled path and then `exit()`. Putting the registry back on a function-local static object killed 24 of 24 children on a signal; the leaked singleton kills 0 of 24. - `AGrowthFreesTheBufferItReplacesWithTheDeviceLockDropped`. Parks work on an unrelated stream so the growth's free stalls, then takes the device's pool lock while the growth is still inside it, reading the grown capacity under that same `try_lock` so a lock that is free only because the growth has not started cannot satisfy it. Moving the unlock below the `cudaFree` fails it. 6 targets, 74 -> 79 cases, all passing on an A100 with `--nocache_test_results` and `TORCHTRT_EXECUTORCH_REQUIRE_CUDA=1` set, so nothing skipped. --- .github/workflows/executorch-build-linux.yml | 9 +- cpp/BUILD | 34 ++ .../executorch/TensorRTBackend.h | 22 +- cpp/src/torch_tensorrt/executorch/README.md | 75 ++++- .../executorch/SharedScratchPool.h | 55 ++-- .../executorch/SharedScratchPoolTestHooks.cpp | 53 +++ .../executorch/SharedScratchPoolTestHooks.h | 39 +++ .../executorch/TensorRTBackend.cpp | 86 +---- tests/cpp/executorch/BUILD | 1 + .../test_shared_scratch_backend.cpp | 302 +++++++++++++++++- .../executorch/test_shared_scratch_pool.cpp | 84 +++++ 11 files changed, 658 insertions(+), 102 deletions(-) create mode 100644 cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp create mode 100644 cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.h diff --git a/.github/workflows/executorch-build-linux.yml b/.github/workflows/executorch-build-linux.yml index f3eab9c2379..1e9923a4d7d 100644 --- a/.github/workflows/executorch-build-linux.yml +++ b/.github/workflows/executorch-build-linux.yml @@ -119,9 +119,16 @@ jobs: fi export LD_LIBRARY_PATH="${_dir}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" done + # TORCHTRT_EXECUTORCH_REQUIRE_CUDA turns a case that skips for want of a + # CUDA device into a failure. This job runs in a container started with + # all GPUs attached, so a skip here means the device went missing rather + # than that the runner never had one -- and without the variable the + # backend suite would skip every case, exit zero and be reported as a + # passing target that covered nothing. bazel test //tests/cpp/executorch:executorch_backend_tests \ --compilation_mode opt --config=linux --test_output=errors \ - --test_env=LD_LIBRARY_PATH + --test_env=LD_LIBRARY_PATH \ + --test_env=TORCHTRT_EXECUTORCH_REQUIRE_CUDA=1 executorch_cmake_location="$(bazel query @executorch//:executorch/CMakeLists.txt --output=location)" export EXECUTORCH_SOURCE_DIR="$(dirname "${executorch_cmake_location%%:*}")" export EXECUTORCH_ROOT="${EXECUTORCH_SOURCE_DIR}" diff --git a/cpp/BUILD b/cpp/BUILD index 71af9b8f4b7..a8f0eca220e 100644 --- a/cpp/BUILD +++ b/cpp/BUILD @@ -216,6 +216,40 @@ cc_library( }), ) +# The pool's test-only entry points. A separate target, and testonly, so the +# released archive -- Bazel's :tensorrt_executorch_backend and CMake's +# executorch_trt_backend, neither of which compiles this source -- carries no +# definition of them and exports no symbol for them. One of them frees the live +# pool with no wait for work in flight, which is not something a release build +# should offer a caller. +cc_library( + name = "tensorrt_executorch_shared_scratch_pool_test_hooks", + testonly = True, + srcs = [ + "src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp", + ], + hdrs = [ + "src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.h", + ], + strip_include_prefix = "src", + target_compatible_with = select({ + ":linux_x86_64": [], + ":sbsa": [], + "//conditions:default": ["@platforms//:incompatible"], + }), + deps = [ + ":tensorrt_executorch_shared_scratch_pool", + ] + select({ + ":linux_x86_64": [ + "@cuda//:cudart", + ], + ":sbsa": [ + "@cuda//:cudart", + ], + "//conditions:default": [], + }), +) + cc_library( name = "tensorrt_executorch_backend", srcs = [ diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index 0158291f4f5..39b7b688fc9 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -134,8 +134,12 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { // frees the buffer it replaces. cudaFree waits for every stream on the // device, so that one call blocks until the device is idle however // asynchronous the rest of this contract makes it -- an unbounded wait on - // work this call did not submit. Which calls grow the pool is not knowable - // from here; see the README. + // work this call did not submit. Unbounded is meant literally: if any of + // that work is itself waiting on something only this thread supplies once + // execute() returns -- a host function it will release, a copy it will + // enqueue next -- the call does not return, and the thread that would + // unblock it is the one inside cudaFree. Which calls grow the pool is not + // knowable from here; see the README. // - Capturing a CUDA graph from the selected stream is refused with // Error::NotSupported, ahead of every CUDA call it makes that a capture // cannot take -- only the device query and the device switch run first. The @@ -149,7 +153,19 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { // running on any other stream under cudaStreamCaptureModeGlobal, or under // cudaStreamCaptureModeThreadLocal from this thread, is invalidated by the // same calls and is not refused, because CUDA offers no query for it: do not - // run a pooled engine while capturing anywhere in the process. + // run a pooled engine while capturing anywhere in the process. Turning the + // option off removes the pool's calls, and the refusal with them -- the + // check is inside the pooled path. It does not make execute() capturable in + // general: the wait on a previous enqueue, a staging buffer's cudaMalloc and + // the cudaStreamSynchronize that ends any call staging through host memory, + // aliasing an output or running with no caller stream all remain, the last + // of them prohibited under every capture mode. A call submitted on the + // capturing stream through a CallerStreamGuard that binds only non-empty + // device-resident tensors, aliases no output and follows no call that left + // an enqueue in flight makes none of them -- and then leaves its own enqueue + // in flight, so under the Global and ThreadLocal modes the next call on that + // handle fails the last condition: one captured call per handle, not + // capture. The README states the rule in full. // - cudaDeviceReset() invalidates the pool without emptying it. The buffer // and the handoff event it still holds are destroyed with the primary // context, and the next call on that device uses both. There is no guard: diff --git a/cpp/src/torch_tensorrt/executorch/README.md b/cpp/src/torch_tensorrt/executorch/README.md index 6e97250ceeb..f673d0f7a5a 100644 --- a/cpp/src/torch_tensorrt/executorch/README.md +++ b/cpp/src/torch_tensorrt/executorch/README.md @@ -183,10 +183,31 @@ against the engines that do. The buffer grows when a call asks for more than every call 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. +so it can stall for far longer than the event wait that precedes it. It is the +only call on the growth path that does: measured with a host function parked on a +stream neither the pool nor the growing engine had ever used, the `cudaMalloc` for +the new buffer and the wait on the retired one each returned at once and the +`cudaFree` returned after 3000 ms, exactly when that host function was released. +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. + +The wait has no upper bound, and a caller can turn it into one that never ends. +Anything queued anywhere on the device is enough to hold it, so if some stream is +waiting on work only this thread will submit -- a host function it will release +after `execute()` returns, a copy it will enqueue next -- the growing call does +not return and the thread that would unblock it is the one inside `cudaFree`. +This is not hypothetical: the backend's own tests deadlocked on it once, when a +change of test order turned a case that parks a host function on its own stream +into the one that grew the pool. A program that parks work like that has to keep +the pool's growths away from it, and the paragraph below on how often the pool +grows is what says whether a run order can do that. + +A stream-ordered free would bound the wait to the buffer's own users, but +`cudaFreeAsync` pairs only with `cudaMallocAsync`, so it would move the pool onto +the stream-ordered allocator for every allocation it makes, not just the free. +Deferring the free instead only moves the wait, since nothing later on the path +pays it. The pool keeps the plain allocator and this section is the warning. What an engine answers when asked how much it needs is decided when it is built, not when it runs. The builder's `kRUNTIME_ACTIVATION_RESIZE_10_10` preview feature @@ -222,10 +243,9 @@ whatever buffer was installed when it was captured -- which by then the pool may have freed. The backend checks the selected stream and returns instead, and it checks ahead of everything a capture cannot take -- not merely ahead of the pool's own calls, since the wait on a previous enqueue and the `cudaMalloc` that grows a -host-input staging buffer come before those and would invalidate the capture -first. Only the device query and the device switch run before the check, and a -capture takes both. Load the engine with the option off to capture it: a context -that owns its scratch makes none of them. +host-input staging buffer come before those and, outside `Relaxed`, would +invalidate the capture first. Only the device query and the device switch run +before the check, and a capture takes both. What that check does not cover is a capture running on some *other* stream. Under `Global`, and under `ThreadLocal` from the thread that calls `execute()`, the same @@ -238,6 +258,45 @@ the pool's calls are permitted and refusing would be wrong. CUDA has no query fo "is a capture live in this process", so the rule is the caller's to keep: do not run a pooled engine while any capture is open anywhere in the process. +Turning the option off takes the pool's calls out of the way, and the refusal with +them: the check lives inside the pooled path, so a context that keeps its own +scratch is never asked. That does not make `execute()` capturable in general. +Three kinds of call remain that a capture may not survive, and both the modes and +the conditions are worth stating exactly, because with nothing refusing them the +failure is silent: + +- The host wait on a previous enqueue, made when the last call on this handle + returned with its enqueue still in flight. Prohibited under `Global` and + `ThreadLocal`, permitted under `Relaxed`. +- A staging buffer's `cudaMalloc`, made when a tensor is empty or host-resident + and this handle holds no cached buffer for it that is already large enough, and + the `cudaFree` of a cached buffer a larger tensor has outgrown. Prohibited under + `Global` and `ThreadLocal`, permitted under `Relaxed`. +- The `cudaStreamSynchronize` on the selected stream that ends the call, made + when any tensor is staged through host memory, when an output is aliased, or + when no caller stream is set. Prohibited under **every** mode, `Relaxed` + included: a capturing stream cannot be synchronized at all. + +So a call captures cleanly with the option off when it runs on the capturing +stream under a `CallerStreamGuard`, binds only non-empty device-resident tensors, +has no aliased outputs, and follows no call that left an enqueue in flight. A +first call on a freshly loaded handle over device-resident inputs and outputs +meets all four, and that is measured rather than argued. But such a call then +leaves its own enqueue in flight, so under `Global` and `ThreadLocal` the next one +on the same handle fails the last condition: with the option off a handle is good +for one captured call, not for capture in general. `Relaxed` permits that wait, so +the limit is not one there -- it permits the staging calls too, and runs each of +them for real and outside the graph, which is a hazard of its own rather than a +way to capture more. `AnUnpooledEngineCapturesOnAFirstCallThatStagesNothing` +and `AnUnpooledEngineInvalidatesACaptureOnceAnEnqueueIsInFlight` in +`tests/cpp/executorch/test_shared_scratch_backend.cpp` are those two outcomes. + +Moving the pool's check out of the pooled path so that it refused both ways would +make the error honest but would also refuse the calls above that do capture, for +contexts that never touch the pool. That is a behaviour change for callers who +have not turned the option on, so the check stays where it is and this section is +what the refusal message points at. + `cudaDeviceReset()` is not survivable and is not guarded against. The pool holds its buffer and its handoff event for the process lifetime, and a reset destroys the primary context under both. The next pooled `execute()` on that device then diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h index 2ad867351ba..eb719a115d5 100644 --- a/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h @@ -84,6 +84,10 @@ class SharedScratchPool { // so the caller can release it. Entries stay in the map, so a reference `get` // handed out remains valid. // + // A template, so nothing is emitted for it until something instantiates it, and + // in the released archive nothing does: the only entry point that calls it is + // compiled into the test target alone. See SharedScratchPoolTestHooks.h. + // // Every slot is emptied under the locks, and what came out of it is disposed of // afterwards with neither held, because the disposer frees device memory and a // device-wide free blocks on everything queued on that device -- a parked host @@ -118,6 +122,39 @@ class SharedScratchPool { std::unordered_map devices_; }; +// The process-wide per-device pool the TensorRT backend runs on. One buffer per +// device, grown to the largest requirement any call on that device has asked +// for, serves every kUSER_MANAGED context on it, 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 +// one buffer at the same time. A device's lock is what enforces that -- see the +// backend's 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. +// +// 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. +// +// The C++ object is never destroyed either, and that is not the same claim. +// Static destruction would destroy the registry's mutex, every device's mutex +// and the map nodes a live reference points into, while a thread anywhere +// between the lookup and the release of its device lock still holds them. +// Leaking it costs one allocation and removes that race. +// +// Inline, so the pool's test hooks can reach the same instance from the +// translation unit that defines them without the backend having to export an +// accessor for them; see SharedScratchPoolTestHooks.h. One instance still, for +// the usual reason a function-local static in an inline function is one. +inline SharedScratchPool& scratch_pool() { + static SharedScratchPool* const pool = new SharedScratchPool(); + return *pool; +} + // Claims a device's handoff for a caller about to enqueue against its shared // scratch, creating the marker's event on first use. Call with `dev.mu` held. // @@ -233,23 +270,5 @@ void* shared_scratch_get_or_grow( // that links this header alone never names it. bool install_pooled_scratch(nvinfer1::IExecutionContext& ctx, void* buffer, std::size_t bytes, int device_id); -// Test-only views of the process-wide pool the TensorRT backend runs on. Declared -// here, in a header that ships under src/, so they add nothing to the installed -// API; defined in TensorRTBackend.cpp, which owns the pool instance they act on. -// A binary that links the pool header alone -- the pool's own unit test -- never -// names them, so the missing definition costs it nothing. -// -// Both are safe only with no claim outstanding and no enqueue in flight against a -// pooled buffer; a test earns that by synchronizing every stream it submitted on. -// The reset also frees what the pool holds, so a test should destroy its delegate -// handles before that one as well. - -// The bytes the pool holds for `device_id` right now; zero if it holds nothing. -std::size_t shared_scratch_capacity_for_testing(int device_id); - -// Frees every device's buffer, destroys its handoff event, and clears the marker, -// so one test does not inherit a pool an earlier one grew. -void reset_shared_scratch_pool_for_testing(); - } // namespace executorch_backend } // namespace torch_tensorrt diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp new file mode 100644 index 00000000000..9cf80e8e1ef --- /dev/null +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp @@ -0,0 +1,53 @@ +/* + * 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. + */ + +#include "torch_tensorrt/executorch/SharedScratchPoolTestHooks.h" + +#include "torch_tensorrt/executorch/SharedScratchPool.h" + +#include + +#include + +namespace torch_tensorrt { +namespace executorch_backend { + +std::size_t shared_scratch_capacity_for_testing(int device_id) { + SharedScratchDevice& dev = scratch_pool().get(device_id); + std::lock_guard lk(dev.mu); + return dev.capacity; +} + +void reset_shared_scratch_pool_for_testing() { + int restore_to = 0; + const bool have_current = cudaGetDevice(&restore_to) == cudaSuccess; + scratch_pool().reset_for_testing([](int device_id, void* buffer, cudaEvent_t event) { + if (buffer == nullptr && event == nullptr) { + return; + } + // cudaFree and cudaEventDestroy both act on the current device, and a slot is + // keyed by the device its buffer came from. + if (cudaSetDevice(device_id) != cudaSuccess) { + return; + } + if (buffer != nullptr) { + (void)cudaFree(buffer); + } + if (event != nullptr) { + (void)cudaEventDestroy(event); + } + // Clears a non-sticky error so a reset does not leave one for the next call to + // report. A sticky one survives the clear, and no cleanup here recovers it. + (void)cudaGetLastError(); + }); + if (have_current) { + (void)cudaSetDevice(restore_to); + } +} + +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.h b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.h new file mode 100644 index 00000000000..c9716d1523f --- /dev/null +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.h @@ -0,0 +1,39 @@ +/* + * 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 + +// Views of the shared activation-scratch pool that only a test has any business +// calling: one reads the pool's capacity, the other frees everything it holds +// with no wait for work in flight against it. +// +// They live in this header and SharedScratchPoolTestHooks.cpp rather than beside +// the pool, and neither file is compiled into libexecutorch_trt_backend or +// shipped in the source package, so a released build contains no definition of +// either and exports no symbol for them. The Bazel target that carries them is +// testonly. Reaching the pool from a separate translation unit is what +// scratch_pool() is inline in SharedScratchPool.h for. +// +// Both are safe only with no claim outstanding and no enqueue in flight against a +// pooled buffer; a test earns that by synchronizing every stream it submitted on. +// The reset also frees what the pool holds, so a test should destroy its delegate +// handles before calling it. + +#include + +namespace torch_tensorrt { +namespace executorch_backend { + +// The bytes the pool holds for `device_id` right now; zero if it holds nothing. +std::size_t shared_scratch_capacity_for_testing(int device_id); + +// Frees every device's buffer, destroys its handoff event, and clears the marker, +// so one test does not inherit a pool an earlier one grew. +void reset_shared_scratch_pool_for_testing(); + +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index 108d2be5d1f..e9c8aebe7c6 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -229,34 +229,6 @@ 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, grown to the largest requirement any call on that device has asked -// for, 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. 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. -// -// 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. -// -// The C++ object is never destroyed either, and that is not the same claim. -// Static destruction would destroy the registry's mutex, every device's mutex -// and the map nodes a live reference points into, while a thread anywhere -// between the lookup and the release of its device lock still holds them. -// Leaking it costs one allocation and removes that race. -SharedScratchPool& scratch_pool() { - static SharedScratchPool* const pool = new SharedScratchPool(); - return *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. // @@ -307,6 +279,15 @@ class SharedScratchClaim { // 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. // + // Unlocking bounds who waits, not how long. This free is still the one call on + // the growth path that waits device-wide -- measured, the allocation and the + // marker wait do not -- and a caller whose own later work some queued stream is + // waiting on will not get this call back at all. TensorRTBackend.h says so in + // the execute() contract, and the README says what to do about it. Making the + // free stream-ordered is not a local change: cudaFreeAsync pairs only with + // cudaMallocAsync, so it would move every pool allocation onto the + // stream-ordered allocator to bound this one wait. + // // Frees on the current device, which must still be the buffer's. void release() { if (lock_.owns_lock()) { @@ -363,11 +344,12 @@ constexpr size_t kMinPooledScratchBytes = 1; // and a null graph. Refusing names the cause instead. // // execute() calls this ahead of every CUDA call it makes that a capture cannot -// take, not just the pool's own: the wait on a previous enqueue invalidates a -// capture under every mode, and the cudaMalloc that grows a host-input staging -// buffer under every mode but Relaxed. Either leaves a refusal made after it -// nothing to save. Only the device query and the device switch run earlier, and -// a capture takes both. +// take, not just the pool's own: the cudaEventSynchronize on a previous enqueue +// and the cudaMalloc that grows a host-input staging buffer each invalidate a +// capture under every mode but Relaxed, and the cudaStreamSynchronize that ends +// a call this backend does not let return early does so under every mode. Any of +// them leaves a refusal made after it nothing to save. Only the device query and +// the device switch run earlier, and a capture takes both. // // It sees only `stream`. A capture live on some *other* stream is invalidated by // the pool's calls just the same under the Global mode, and under ThreadLocal @@ -392,7 +374,7 @@ Error refuse_pooled_call_on_a_capturing_stream(cudaStream_t stream, int device_i } ET_LOG( Error, - "TensorRTBackend::execute: the selected stream is capturing a CUDA graph (%s), which the shared activation scratch pool on device %d does not support. Load this engine with '%s' off so its context keeps its own scratch.", + "TensorRTBackend::execute: the selected stream is capturing a CUDA graph (%s), which the shared activation scratch pool on device %d does not support. Loading this engine with '%s' off takes the pool's calls out of the way, but not every call a capture cannot take, and nothing refuses on that path: the capture section of the backend README says which calls remain and when a call survives them.", capture_err == cudaSuccess ? "capture in progress" : cudaGetErrorString(capture_err), device_id, kSharedActivationScratchKey); @@ -628,42 +610,6 @@ bool install_pooled_scratch(nvinfer1::IExecutionContext& ctx, void* buffer, size return false; } -// Declared in SharedScratchPool.h, which ships under src/ and is not installed; -// defined here because this is where the pool instance lives. See that header for -// what a caller owes them. -std::size_t shared_scratch_capacity_for_testing(int device_id) { - SharedScratchDevice& dev = scratch_pool().get(device_id); - std::lock_guard lk(dev.mu); - return dev.capacity; -} - -void reset_shared_scratch_pool_for_testing() { - int restore_to = 0; - const bool have_current = cudaGetDevice(&restore_to) == cudaSuccess; - scratch_pool().reset_for_testing([](int device_id, void* buffer, cudaEvent_t event) { - if (buffer == nullptr && event == nullptr) { - return; - } - // cudaFree and cudaEventDestroy both act on the current device, and a slot is - // keyed by the device its buffer came from. - if (cudaSetDevice(device_id) != cudaSuccess) { - return; - } - if (buffer != nullptr) { - (void)cudaFree(buffer); - } - if (event != nullptr) { - (void)cudaEventDestroy(event); - } - // Clears a non-sticky error so a reset does not leave one for the next call to - // report. A sticky one survives the clear, and no cleanup here recovers it. - (void)cudaGetLastError(); - }); - if (have_current) { - (void)cudaSetDevice(restore_to); - } -} - // --------------------------------------------------------------------------- // is_available // --------------------------------------------------------------------------- diff --git a/tests/cpp/executorch/BUILD b/tests/cpp/executorch/BUILD index a64447a9446..b12c8091543 100644 --- a/tests/cpp/executorch/BUILD +++ b/tests/cpp/executorch/BUILD @@ -75,6 +75,7 @@ cc_test( "//cpp:tensorrt_executorch_backend", "//cpp:tensorrt_executorch_blob_header", "//cpp:tensorrt_executorch_shared_scratch_pool", + "//cpp:tensorrt_executorch_shared_scratch_pool_test_hooks", "@executorch//:executorch_core", "@executorch//:executorch_headers", "@executorch//:extension_cuda", diff --git a/tests/cpp/executorch/test_shared_scratch_backend.cpp b/tests/cpp/executorch/test_shared_scratch_backend.cpp index a8d82034b4a..9a944c1344d 100644 --- a/tests/cpp/executorch/test_shared_scratch_backend.cpp +++ b/tests/cpp/executorch/test_shared_scratch_backend.cpp @@ -23,10 +23,14 @@ // gtest case exits zero, which Bazel reports as a passing target. A run that is // meant to have a device says so through TORCHTRT_EXECUTORCH_REQUIRE_CUDA, and // then a skip is a failure instead; the count of skipped cases is printed either -// way. Making the CI job actually run on a GPU runner is a separate matter, and -// until it does the strongest evidence for this file is hand measurement. +// way. The CI invocation in .github/workflows/executorch-build-linux.yml passes +// that variable, and the job it sits in asks for a GPU runner and starts its +// container with every GPU attached. What still keeps these cases off most runs +// is the lane gate on the whole ExecuTorch job in ci-linux-x86_64.yml, not the +// runner it would land on. #include "torch_tensorrt/executorch/SharedScratchPool.h" +#include "torch_tensorrt/executorch/SharedScratchPoolTestHooks.h" #include "torch_tensorrt/executorch/TensorRTBackend.h" #include "torch_tensorrt/executorch/TensorRTBlobHeader.h" @@ -458,6 +462,12 @@ class LoadedEngine { return static_cast(handle_); } + // The execution context execute() configures and enqueues on, so a test can + // watch what the delegate does to it. + nvinfer1::IExecutionContext* context() const { + return static_cast(handle_)->exec_ctx.get(); + } + std::size_t elems() const { return static_cast(batch_) * static_cast(rows_) * static_cast(cols_); } @@ -527,6 +537,56 @@ bool open_user_managed_context(const std::vector& blob, int rows, return out.ctx->setInputShape("input_0", nvinfer1::Dims3{batch, rows, cols}); } +// Counts the reference-count calls TensorRT makes on a recorder as it attaches +// and detaches one, which is how a test attached to the delegate's own context +// sees install_pooled_scratch scope its recorder over the install. +// +// TensorRT documents setErrorRecorder as calling incRefCount on the recorder it +// takes and decRefCount on the one it replaces, and it does: measured, attaching +// a second recorder over this one and then restoring it moves detached to 1 and +// reattached to 1, and nothing else in a pooled execute() touches either. +class CountingErrorRecorder final : public nvinfer1::IErrorRecorder { + public: + int32_t getNbErrors() const noexcept override { + return 0; + } + nvinfer1::ErrorCode getErrorCode(int32_t) const noexcept override { + return nvinfer1::ErrorCode::kSUCCESS; + } + ErrorDesc getErrorDesc(int32_t) const noexcept override { + return ""; + } + bool hasOverflowed() const noexcept override { + return false; + } + void clear() noexcept override {} + bool reportError(nvinfer1::ErrorCode, ErrorDesc) noexcept override { + reported.fetch_add(1); + return false; + } + RefCount incRefCount() noexcept override { + reattached.fetch_add(1); + return ++refs_; + } + RefCount decRefCount() noexcept override { + detached.fetch_add(1); + return --refs_; + } + + void forget() { + reattached.store(0); + detached.store(0); + reported.store(0); + } + + std::atomic reattached{0}; + std::atomic detached{0}; + std::atomic reported{0}; + + private: + RefCount refs_ = 1; +}; + // Set by a run that is supposed to have a GPU. A skip is then a failure, rather // than a green target that covered nothing. constexpr char kRequireCudaEnvVar[] = "TORCHTRT_EXECUTORCH_REQUIRE_CUDA"; @@ -1156,6 +1216,57 @@ TEST_F(SharedScratchBackendTest, AnUndersizedScratchInstallIsSeenAsRefused) { EXPECT_EQ(cudaFree(one_byte), cudaSuccess); } +// The case above drives install_pooled_scratch directly, which leaves execute() +// free to stop calling it: replacing the checked call with a bare +// setDeviceMemoryV2 installs the same buffer, produces the same output and drops +// only the refusal, so every other case here stays green. +// +// Forcing the refusal itself through execute() is not available. The size +// execute() installs is the pool's capacity, and the pool is grown to at least +// what the query just returned, so the installed buffer is never short of what +// the context expects. Only a failed query makes the two disagree, and that +// cannot be induced from inside this process. +// +// What is observable is the recorder. install_pooled_scratch attaches one for +// the duration of the install and restores the previous one after, because that +// is the only channel TensorRT reports a refusal on. A recorder this test leaves +// on the delegate's context therefore sees itself replaced and put back exactly +// once per pooled run, and sees nothing at all if the install stops going +// through the helper. +TEST_F(SharedScratchBackendTest, ExecuteInstallsPooledScratchThroughTheCheckedHelper) { + // Declared before the engine so it outlives the context TensorRT attaches it + // to: the context decrements this recorder's count as it is destroyed. + CountingErrorRecorder recorder; + + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + + LoadedEngine pooled; + ASSERT_EQ(pooled.load(blob(), 25), Error::Ok); + ASSERT_TRUE(pooled.handle()->shared_scratch); + ASSERT_TRUE(pooled.handle()->engine_needs_scratch) + << "this engine skips the pool, so its execute() makes no install for this case to watch"; + pooled.context()->setErrorRecorder(&recorder); + + recorder.forget(); + ASSERT_EQ(pooled.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + EXPECT_EQ(recorder.detached.load(), 1) + << "a pooled execute() did not replace the context's error recorder, so its activation-scratch install was not " + "the checked one and a refusal by TensorRT would go unread"; + EXPECT_EQ(recorder.reattached.load(), 1) + << "the recorder the caller had attached was not put back after the install, so TensorRT's diagnostics for the " + "rest of the call go somewhere the caller did not ask for"; + EXPECT_EQ(pooled.context()->getErrorRecorder(), &recorder) + << "execute() left an error recorder of its own attached to the context"; + EXPECT_EQ(recorder.reported.load(), 0) + << "TensorRT reported an error to the caller's recorder during a run that succeeded"; + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + // --------------------------------------------------------------------------- // Stream capture // --------------------------------------------------------------------------- @@ -1260,6 +1371,98 @@ TEST_F(SharedScratchBackendTest, APooledEngineRefusesACaptureBeforeAnythingCanIn ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); } +// The two cases above are the pool refusing. These two are what a caller gets +// when it takes the refusal's advice and loads with the option off, which the +// README and the refusal message both describe and neither may overstate. +// +// Nothing refuses then -- the guard sits inside the pooled branch -- so the +// question is only whether execute()'s own calls are ones a capture can take. +// On a first call over non-empty device-resident tensors, submitted on the +// capturing stream through a CallerStreamGuard, they are: the device query, the +// device switch, the shape and address binding, enqueueV3 and the completion +// record. Nothing here allocates, waits or synchronizes. +TEST_F(SharedScratchBackendTest, AnUnpooledEngineCapturesOnAFirstCallThatStagesNothing) { + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + // A node of the caller's own, so a null graph below means the capture was + // invalidated rather than that nothing was ever captured. + void* captured_target = nullptr; + ASSERT_EQ(cudaMalloc(&captured_target, 16), cudaSuccess); + + LoadedEngine unpooled; + ASSERT_EQ(unpooled.load(blob(), 22), Error::Ok); + ASSERT_FALSE(unpooled.handle()->shared_scratch) + << "this handle took the pool, so it says nothing about what capturing with the option off gets"; + ASSERT_FALSE(unpooled.handle()->inflight_pending) + << "a freshly loaded handle already reports an enqueue in flight, so the host wait this case is built to avoid " + "would run anyway"; + + ASSERT_EQ(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal), cudaSuccess); + ASSERT_EQ(cudaMemsetAsync(captured_target, 0, 16, stream), cudaSuccess); + const Error captured = unpooled.run(stream); + cudaGraph_t graph = nullptr; + const cudaError_t end_err = cudaStreamEndCapture(stream, &graph); + const bool have_graph = graph != nullptr; + if (have_graph) { + cudaGraphDestroy(graph); + } + + EXPECT_EQ(captured, Error::Ok) << "a call with the option off was refused or failed under capture"; + EXPECT_EQ(end_err, cudaSuccess) << "a call with the option off invalidated the capture: " + << cudaGetErrorString(end_err); + EXPECT_TRUE(have_graph) << "the capture ended with no graph, so the call invalidated it"; + + EXPECT_EQ(cudaFree(captured_target), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +// The other half, and the reason the advice cannot be left at "load with the +// option off to capture it". The call that captured cleanly returns with its +// enqueue still in flight, and the next call on that handle waits for it on the +// host before it may touch the context. That wait is prohibited under the +// default capture mode, and nothing refuses it: the caller gets an error from +// execute() and a capture that ends with no graph. +// +// So under this mode the option being off buys one capture per handle, not +// capture in general. +TEST_F(SharedScratchBackendTest, AnUnpooledEngineInvalidatesACaptureOnceAnEnqueueIsInFlight) { + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + void* captured_target = nullptr; + ASSERT_EQ(cudaMalloc(&captured_target, 16), cudaSuccess); + + LoadedEngine unpooled; + ASSERT_EQ(unpooled.load(blob(), 23), Error::Ok); + ASSERT_FALSE(unpooled.handle()->shared_scratch); + ASSERT_EQ(unpooled.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + ASSERT_TRUE(unpooled.handle()->inflight_pending) + << "the first run left no enqueue in flight, so the next call makes no host wait and this case pins nothing"; + + ASSERT_EQ(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal), cudaSuccess); + ASSERT_EQ(cudaMemsetAsync(captured_target, 0, 16, stream), cudaSuccess); + const Error captured = unpooled.run(stream); + cudaGraph_t graph = nullptr; + const cudaError_t end_err = cudaStreamEndCapture(stream, &graph); + const bool have_graph = graph != nullptr; + if (have_graph) { + cudaGraphDestroy(graph); + } + // The wait's own cudaErrorStreamCaptureUnsupported is not sticky and execute() + // does not clear it, so it would otherwise surface under the name of whatever + // the next case calls first. + (void)cudaGetLastError(); + + EXPECT_EQ(captured, Error::InvalidProgram) + << "a second call on this handle did not fail on the host wait for the previous enqueue, so the documented " + "one-captured-call-per-handle limit with the option off is wrong"; + EXPECT_NE(end_err, cudaSuccess) << "the capture survived, so the documented caveat about the second call is stale"; + EXPECT_FALSE(have_graph) << "the capture ended with a graph, so the documented caveat about the second call is stale"; + + EXPECT_EQ(cudaFree(captured_target), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + // --------------------------------------------------------------------------- // The enqueue handoff, single-threaded, two caller streams // --------------------------------------------------------------------------- @@ -1400,6 +1603,101 @@ TEST_F(SharedScratchBackendTest, ASecondPooledEnqueueWaitsForTheFirstOnAnotherSt ASSERT_EQ(cudaStreamDestroy(second_stream), cudaSuccess); } +// --------------------------------------------------------------------------- +// Where the growth's free falls relative to the device lock +// --------------------------------------------------------------------------- + +// The growth test earlier runs alone, so it cannot see whether the buffer it +// retires is freed with the device's lock held. Moving the unlock in +// SharedScratchClaim::release() below the cudaFree leaves it green. +// +// It matters because cudaFree waits for every stream on the device, not only for +// the ones that touched the buffer. Measured on this box: with a host function +// parked on a stream that never saw the buffer, cudaFree returned after 3000 ms +// while the cudaMalloc and the marker wait on the same growth path each returned +// at once. Under the lock, that wait is one every other pooled engine on the +// device has to sit through. +// +// So: park work on an unrelated stream to stall the free, then check the pool +// lock while the growth is still inside it. The capacity read under that same +// try_lock is what stops the case passing on a lock that is free because the +// growth has not started -- past a growth it is above what the smaller engine +// left, and the growth has not returned. +TEST_F(SharedScratchBackendTest, AGrowthFreesTheBufferItReplacesWithTheDeviceLockDropped) { + 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 second to be sure of growing the pool"; + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + + cudaStream_t stream = nullptr; + cudaStream_t unrelated = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + ASSERT_EQ(cudaStreamCreateWithFlags(&unrelated, cudaStreamNonBlocking), cudaSuccess); + + LoadedEngine small; + LoadedEngine big; + ASSERT_EQ(small.load(blob(), 26), Error::Ok); + ASSERT_EQ(big.load(big_blob(), 27, kBigRows, kBigCols), Error::Ok); + ASSERT_TRUE(small.handle()->shared_scratch); + ASSERT_TRUE(big.handle()->shared_scratch); + const int device_id = big.handle()->device_id; + + // Without this the larger engine allocates rather than grows, and a growth that + // retires nothing frees nothing. + ASSERT_EQ(small.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::size_t before = shared_scratch_capacity_for_testing(device_id); + ASSERT_GT(before, 0u) << "the smaller engine left the pool empty, so the larger one has nothing to retire"; + + // Neither engine submits to this stream, so nothing on the growth path but the + // device-wide free has any reason to wait for it. + StreamGate gate; + ASSERT_EQ(cudaLaunchHostFunc(unrelated, hold_stream, &gate), cudaSuccess); + GateRelease gate_release(gate, unrelated); + + std::atomic growth_returned{false}; + Error growth_error = Error::Internal; + std::thread grower([&] { + growth_error = big.run(stream); + growth_returned.store(true); + }); + + SharedScratchDevice& dev = scratch_pool().get(device_id); + bool lock_free_after_the_growth = false; + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (!growth_returned.load() && std::chrono::steady_clock::now() < deadline) { + if (dev.mu.try_lock()) { + const std::size_t capacity_now = dev.capacity; + dev.mu.unlock(); + if (capacity_now > before && !growth_returned.load()) { + lock_free_after_the_growth = true; + break; + } + } + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + const bool still_inside_execute = !growth_returned.load(); + + gate_release.release(); + grower.join(); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(unrelated), cudaSuccess); + + 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(growth_error, Error::Ok); + ASSERT_TRUE(still_inside_execute) + << "the growing call returned before the gate opened, so its free never stalled and there was no window in " + "which to observe the lock"; + ASSERT_GT(shared_scratch_capacity_for_testing(device_id), before) + << "the pool did not grow, so no buffer was retired and no free was made"; + EXPECT_TRUE(lock_free_after_the_growth) + << "the device's pool lock stayed held for the whole of a stalled growth free, so every other pooled engine on " + "this device waits out a device-wide free it has nothing to do with"; +} + // --------------------------------------------------------------------------- // The pooled path under two concurrent callers // --------------------------------------------------------------------------- diff --git a/tests/cpp/executorch/test_shared_scratch_pool.cpp b/tests/cpp/executorch/test_shared_scratch_pool.cpp index 51b300ba26e..d5c0ab2f0e9 100644 --- a/tests/cpp/executorch/test_shared_scratch_pool.cpp +++ b/tests/cpp/executorch/test_shared_scratch_pool.cpp @@ -4,6 +4,11 @@ // 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. +// +// One case at the end is the exception to "over fakes": the lifetime of the +// process-wide registry scratch_pool() hands out is a property of that object and +// not of any fake, and it is observable only after main returns. It forks to see +// it, and still makes no CUDA call. #include "torch_tensorrt/executorch/SharedScratchPool.h" @@ -13,6 +18,9 @@ #include #include #include +#include +#include +#include #include #include #include @@ -21,6 +29,13 @@ #include #include +#if defined(__unix__) +#include +#include + +#include +#endif + namespace torch_tensorrt { namespace executorch_backend { namespace { @@ -645,6 +660,75 @@ TEST(SharedScratchPoolRegistry, ConcurrentLookupsKeepTheRegistryIntact) { } } +#if defined(__unix__) + +// scratch_pool() hands back a registry that is allocated once and never +// destroyed. The alternative -- an object with static storage duration -- is +// destroyed after main returns, which destroys the registry's mutex, every +// device's mutex and the map nodes a live reference points into, while a thread +// anywhere between the lookup and the release of its device lock still holds +// them. +// +// Nothing gtest asserts inside this process can see that: it happens after the +// last test has finished. So these run in forked children that put a thread +// inside the pooled path and then exit, and read the status the child died with. +// +// Probabilistic in one direction only. A child that is torn down under a live +// claimant does not have to crash, so the count below is a floor on how often it +// would; a child that is not torn down under one cannot crash at all, so a +// single non-zero status is a real failure and not a flake. +constexpr int kTeardownChildren = 24; + +[[noreturn]] void hold_the_pool_open_then_exit() { + static std::atomic inside{false}; + std::thread claimant([] { + for (unsigned i = 0;; ++i) { + SharedScratchDevice& dev = scratch_pool().get(static_cast(i % 8)); + std::lock_guard lk(dev.mu); + dev.capacity += 1; + inside.store(true); + } + }); + // Never joined: the point is for it to still be in there at teardown. + claimant.detach(); + while (!inside.load()) { + std::this_thread::yield(); + } + std::this_thread::sleep_for(std::chrono::milliseconds(30)); + // exit() and not _exit(): static destruction is the event under test. + std::exit(0); +} + +TEST(SharedScratchPoolRegistry, TheProcessPoolOutlivesStaticDestructionUnderALiveClaimant) { + int killed_by_signal = 0; + int exited_nonzero = 0; + for (int i = 0; i < kTeardownChildren; ++i) { + // Anything gtest has buffered would otherwise be written twice, once by each + // side of the fork. + std::fflush(nullptr); + const pid_t child = fork(); + ASSERT_NE(child, -1) << "fork failed: " << std::strerror(errno); + if (child == 0) { + hold_the_pool_open_then_exit(); + } + int status = 0; + ASSERT_EQ(waitpid(child, &status, 0), child) << "waitpid failed: " << std::strerror(errno); + if (WIFSIGNALED(status)) { + ++killed_by_signal; + } else if (WEXITSTATUS(status) != 0) { + ++exited_nonzero; + } + } + + EXPECT_EQ(killed_by_signal, 0) << killed_by_signal << " of " << kTeardownChildren + << " children died on a signal at exit with a thread still inside the pool, so the " + "registry is being destroyed out from under a live claimant"; + EXPECT_EQ(exited_nonzero, 0) << exited_nonzero << " of " << kTeardownChildren + << " children exited non-zero at teardown with a thread still inside the pool"; +} + +#endif // defined(__unix__) + } // namespace } // namespace executorch_backend } // namespace torch_tensorrt From 55a859afb742061a3823f58be2213e38a3551220 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Mon, 7 Sep 2026 12:58:10 -0700 Subject: [PATCH 09/13] fix(executorch): retract the capture advice, install each call's own scratch, drop the lock across a growth's wait Every figure below is from local runs on one A100 host; the CI job that would run these tests has not run them at this head. **The capture advice was wrong in a way three rounds of correcting it missed, so it is withdrawn rather than corrected again.** Measured through the delegate: with the option off, a first call over device-resident tensors on the capturing stream captures, and the graph instantiates, launches and reproduces the bytes the same engine produces outside a capture. The next ordinary `execute()` on that handle then fails `Error::InvalidProgram`, logging `cudaEventSynchronize failed: invalid argument`, because the completion event `execute()` records for its asynchronous return was recorded inside the graph. No call shape avoids it: the alternative to recording that event is the `cudaStreamSynchronize` that ends a staging call, and a capturing stream cannot be synchronized under any mode. So the README, the installed header and the refusal message now say that this delegate does not support CUDA graph capture, with the option on or off, instead of pointing a caller at the option-off path. The shipped test hid this by destroying the captured graph without launching it, which also left the poisoned handle's failure to surface as an unasserted log line from `~EngineHandle`. It now launches the graph, checks the engine's output against a reference run, and asserts the poisoning as the measurement the documentation rests on. Making the option-off path work means reworking the completion event, which the asynchronous-return contract rests on, and refusing capture on every path would refuse calls that do capture today for contexts that never touch the pool; both are recorded in the README as declined. **A pooled call now installs its own requirement, not the pool's capacity.** They differ after any growth -- measured, 3 of 139 pooled installs were oversized, one of them telling an engine needing 33,554,432 bytes that it owned 134,217,728, of which the sampled tail still held 4078 of 4096 bytes of the previous engine's activations. The larger figure hands a context write access to another engine's scratch, keeps an overrun past its own requirement inside the region TensorRT thinks it owns, and blunts the install check this branch added: after a growth the capacity can cover what a failed per-shape query concealed. The pool never clears the buffer between users, which the README now says. **A growth's host wait moves out from under the device lock, beside the free.** The lock was held across a wait for the previous inference to complete, so another pooled engine on the device blocked 2,119 ms -- serialized at completion rather than at submission, which is what the installed contract promises. The wait now runs in `SharedScratchClaim::release()`, after the unlock and before the device-wide `cudaFree` it guards. Deferring it can pick up enqueues submitted after the unlock, and that costs nothing, because the free that follows waits for every stream on the device anyway; it stays sufficient for the retired buffer because every claimant orders its stream after the marker event before enqueueing, so the latest recording completes only once the earlier ones have. A failed wait still leaks the buffer rather than freeing it under a live enqueue, and still fails the call, now after draining the enqueue this call submitted. **A refused option span no longer applies the entries before the bad one.** `set_option` stored each value as it read it, so a valid `use_shared_activation_scratch=true` followed by a wrong-typed duplicate returned `Error::InvalidArgument` and left the pool on for every engine loaded afterwards. The span is now validated in full before anything is stored; where it names the key more than once the last one wins, as before. **The `cudaGetLastError` in the capture refusal moves inside the failed-query branch.** On the ordinary refusal it was clearing an error the caller was already carrying -- measured with a real `cudaErrorMemoryAllocation` left pending, which the refusal swallowed -- and that error is not sticky, so nothing surfaced it later either. Also: two comments that described every `kUSER_MANAGED` context as drawing on the pool, which the exclusion of scratch-free engines makes false; the declaration of `install_pooled_scratch` moved out of the pool header, which is meant to build against the CUDA headers alone, into the header of the target that defines it; the destructor's remaining `clear sticky error` comment, which the rest of the branch disproves; the README's `set_option` example, which dropped one of the two errors it demonstrates and did not compile as printed; `shared_scratch_get_or_grow` clearing its retirement output on every path, so a caller reusing one is not handed a buffer that was already disposed of; and the NVIDIA copyright header the new CPU-side pool test was missing. `mark_shared_scratch_in_flight` loses its two error branches, which were dead from its only caller: the claim is holding the device's lock, so its device pointer cannot be null, and `get_or_grow_shared_scratch` has already failed the call with `Error::Internal` if the marker had no event, which only the test-only reset clears again and that needs the lock the claim holds. Instrumenting both branches over a whole suite run reached neither. The contract that makes them unnecessary is now written above the function instead, and the accessor whose only caller was one of them goes with them. What can still fail there is the `cudaEventRecord`. The suite's device-less skip banner no longer prints when `TORCHTRT_EXECUTORCH_REQUIRE_CUDA` is set. The counter it reads is incremented before the arming turns the missing device into a failure, so the banner appeared underneath those failures, called the run a passing target and advised setting the variable that had just produced them. With the variable unset it is unchanged, and that is the only run on which a case skips, which the file's coverage note now says rather than promising the count either way. Three descriptions the changes above left behind: the handoff event's blocking-sync comment, which still had its one host wait running under the device lock; the comment at the capture refusal's call site, which had the host wait and the staging `cudaMalloc` invalidating a capture unconditionally, without the `cudaStreamCaptureModeRelaxed` carve-out the same file states thirty lines earlier; and the second capture case, whose comment and assertion messages still described the option off as buying one capture per handle. Five test gaps in the pool's coverage are closed. The reset cases now cover a slot holding an event and no buffer, the state a claim whose allocation failed leaves behind. The reset rendezvous has a deadline, so a reset that disposes of nothing fails in under a second instead of hanging until the target times out. The growth case takes a control reading over a window as long as its measurement and skips rather than blaming the pool when another process moved device-wide memory in it; that is a sample of a different window and not a guarantee -- measured under a process cycling 256 MiB allocations, 8 runs gave 1 skip and 1 failure of the bound -- so the bound's own message now names unrelated device activity too. A new case drives a pooled call whose allocation fails, by taking the device's memory first, and pins that it leaves the device's pool lock free and the pool usable -- the only one of the pooled early returns that can be reached from inside the process. And the reuse case asserts on an output variable the call under test is the only writer of. 6 targets, 79 -> 85 cases, all passing on a real GPU with `--nocache_test_results` and `TORCHTRT_EXECUTORCH_REQUIRE_CUDA=1`. Every behavioural fix was proven by mutating it away: storing before validating the span, the growth's wait back under the lock, the error clear back outside its branch, the retirement output left uncleared, a reset that skips buffer-less slots, a reset that disposes of nothing, a reuse that does not write its output size, and a manually locked device mutex whose only unlock is the explicit release -- each fails the case named for it and, where the target could be run whole, only that case. --- .../executorch/TensorRTBackend.h | 79 ++-- cpp/src/torch_tensorrt/executorch/README.md | 151 +++--- .../executorch/SharedScratchPool.h | 53 +-- .../executorch/TensorRTBackend.cpp | 255 +++++++---- .../test_shared_scratch_backend.cpp | 430 ++++++++++++++++-- .../executorch/test_shared_scratch_pool.cpp | 93 +++- 6 files changed, 797 insertions(+), 264 deletions(-) diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index 39b7b688fc9..cfc110f7978 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -17,6 +17,7 @@ #include +#include #include #include #include @@ -75,8 +76,11 @@ 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). + // Whether exec_ctx was created kUSER_MANAGED, because + // kSharedActivationScratchKey was on at this engine's load. Such a context takes + // its activation scratch from the shared per-device pool on every call, unless + // engine_needs_scratch below is false, in which case it needs none and never + // claims from the pool at all. bool shared_scratch = false; // Whether the engine reports needing activation scratch under any shape, read // at init when shared_scratch is set. It is a predicate and not a size because @@ -98,6 +102,27 @@ struct EngineHandle { ~EngineHandle(); }; +// Installs `bytes` of `buffer` as the activation scratch of `ctx`, a +// kUSER_MANAGED context, and reports whether TensorRT accepted it. Logs the +// refusal, naming `device_id`, when it did not. +// +// The check is the point. setDeviceMemoryV2 returns void and refuses a buffer +// smaller than the bound shapes need, so a caller that does not ask cannot tell an +// accepted install from a refused one -- and a refused one leaves the context +// pointed at the buffer it was last given, which a shared-scratch pool growth may +// since have freed. The engine then reads and writes freed memory with enqueueV3 +// reporting success. The refusal is read back through an IErrorRecorder scoped to +// this one call, the only channel that hands it to the caller: with no recorder +// attached TensorRT writes it to the runtime's ILogger and this function has +// nothing to return. +// +// An implementation detail of execute(), declared here rather than beside the +// pool's bookkeeping because this is the header of the target that defines it and +// the only one that may name TensorRT. The backend's own test drives it directly, +// over the same TensorRT call, to cover a refusal that cannot be induced through +// execute(). +bool install_pooled_scratch(nvinfer1::IExecutionContext& ctx, void* buffer, std::size_t bytes, int device_id); + // 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. Read by TensorRTBackend::set_option below, and delivered as @@ -140,32 +165,34 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { // enqueue next -- the call does not return, and the thread that would // unblock it is the one inside cudaFree. Which calls grow the pool is not // knowable from here; see the README. - // - Capturing a CUDA graph from the selected stream is refused with - // Error::NotSupported, ahead of every CUDA call it makes that a capture - // cannot take -- only the device query and the device switch run first. The - // pool's event handoff waits on an event recorded outside the capture, which - // invalidates it under every capture mode. A growth's allocation, its wait on - // the replaced buffer and the free of that buffer invalidate it under every - // mode but cudaStreamCaptureModeRelaxed, which permits all three but does not - // record them, leaving a replay pointed at a buffer the pool may since have - // freed. The alternative to refusing is a capture that silently comes back - // invalidated. Only a capture on the selected stream is caught. A capture - // running on any other stream under cudaStreamCaptureModeGlobal, or under - // cudaStreamCaptureModeThreadLocal from this thread, is invalidated by the - // same calls and is not refused, because CUDA offers no query for it: do not - // run a pooled engine while capturing anywhere in the process. Turning the - // option off removes the pool's calls, and the refusal with them -- the - // check is inside the pooled path. It does not make execute() capturable in - // general: the wait on a previous enqueue, a staging buffer's cudaMalloc and - // the cudaStreamSynchronize that ends any call staging through host memory, - // aliasing an output or running with no caller stream all remain, the last - // of them prohibited under every capture mode. A call submitted on the + // - Capturing a CUDA graph around this delegate is not supported, with the + // option on or off. With it on, a call whose selected stream is capturing is + // refused with Error::NotSupported, ahead of every CUDA call it makes that a + // capture cannot take -- only the device query and the device switch run + // first. The pool's event handoff waits on an event recorded outside the + // capture, which invalidates it under every capture mode. A growth's + // allocation, its wait on the replaced buffer and the free of that buffer + // invalidate it under every mode but cudaStreamCaptureModeRelaxed, which + // permits all three but does not record them, leaving a replay pointed at a + // buffer the pool may since have freed. The alternative to refusing is a + // capture that silently comes back invalidated. Only a capture on the + // selected stream is caught. A capture running on any other stream under + // cudaStreamCaptureModeGlobal, or under cudaStreamCaptureModeThreadLocal + // from this thread, is invalidated by the same calls and is not refused, + // because CUDA offers no query for it: do not run a pooled engine while + // capturing anywhere in the process. + // With the option off nothing refuses, because the check is inside the + // pooled path, and a call can then be captured: one submitted on the // capturing stream through a CallerStreamGuard that binds only non-empty // device-resident tensors, aliases no output and follows no call that left - // an enqueue in flight makes none of them -- and then leaves its own enqueue - // in flight, so under the Global and ThreadLocal modes the next call on that - // handle fails the last condition: one captured call per handle, not - // capture. The README states the rule in full. + // an enqueue in flight makes no call a capture cannot take, and the graph + // replays the engine. It costs the handle. Such a call ends by recording the + // completion event this contract describes, and under capture that record is + // a node of the graph rather than an event the host can wait on, so every + // later execute() on the handle fails Error::InvalidProgram on the wait for + // it. A call that ends by synchronizing the stream instead is one no capture + // mode permits, so there is no call shape that captures and leaves the + // handle usable. The README states this in full. // - cudaDeviceReset() invalidates the pool without emptying it. The buffer // and the handoff event it still holds are destroyed with the primary // context, and the next call on that device uses both. There is no guard: diff --git a/cpp/src/torch_tensorrt/executorch/README.md b/cpp/src/torch_tensorrt/executorch/README.md index f673d0f7a5a..e79ff1f6b78 100644 --- a/cpp/src/torch_tensorrt/executorch/README.md +++ b/cpp/src/torch_tensorrt/executorch/README.md @@ -127,12 +127,18 @@ requirement any call on that device has asked for: ```cpp #include -executorch::runtime::BackendOptions<1> options; -options.set_option("use_shared_activation_scratch", true); -const executorch::runtime::Error err = - executorch::runtime::set_option("TensorRTBackend", options.view()); -if (err != executorch::runtime::Error::Ok) { - return err; // nothing was set: every context still allocates its own scratch +executorch::runtime::Error enable_shared_activation_scratch() { + executorch::runtime::BackendOptions<1> options; + // Returns Error::InvalidArgument when the object has no room left. One key in a + // BackendOptions<1> always fits, so this cannot fail as written; the check is + // what keeps it honest when a second key is added beside it. + const executorch::runtime::Error stored = + options.set_option("use_shared_activation_scratch", true); + if (stored != executorch::runtime::Error::Ok) { + return stored; + } + // Nothing was set if this fails: every context still allocates its own scratch. + return executorch::runtime::set_option("TensorRTBackend", options.view()); } ``` @@ -175,6 +181,16 @@ freed, so the enqueue would read and write memory the pool no longer owns while rather than enqueueing. Nothing reaches `enqueueV3` on an install TensorRT rejected. +What each call installs is its own requirement, not the capacity the pool holds, +which is larger whenever an earlier call asked for more. Handing a context the +larger figure would say it owns bytes that belong to whatever ran before it, and +the pool never clears the buffer between users -- activation scratch is written +before it is read, so a context that stays inside its own requirement never sees +those bytes, and one that runs past it sees another engine's activations rather +than zeros. The smaller figure also keeps the refusal above sharp: after a growth +the capacity may well cover what a failed query concealed, and an install sized +from it would be accepted. + The third thing a zero can mean is an engine that needs no scratch under *any* shape, and that one is settled before a call is ever made: such an engine is left out of the pool entirely, so it takes no per-device lock and does not serialize @@ -188,9 +204,11 @@ only call on the growth path that does: measured with a host function parked on stream neither the pool nor the growing engine had ever used, the `cudaMalloc` for the new buffer and the wait on the retired one each returned at once and the `cudaFree` returned after 3000 ms, exactly when that host function was released. -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. +The backend keeps both that stall and the event wait before it out from under the +per-device lock, so neither holds up another pooled engine on the device: two +pooled calls still serialize at submission, as the caller-stream contract above +says, and a growth is not an exception to it. Both do fall after the growing +call's own enqueue, so that one `execute()` waits for its own engine work too. The wait has no upper bound, and a caller can turn it into one that never ends. Anything queued anywhere on the device is enough to hold it, so if some stream is @@ -225,12 +243,16 @@ call whose shapes need more than every call before them, so with one of those in the program no run order bounds the number of allocations. One engine over a `[1..4, 512, 512]` profile is either, depending on how it was built. -Two things the pool does not support. +### CUDA graph capture is not supported + +A pooled call whose selected stream is capturing is refused with +`Error::NotSupported`. Capture with the option off is not a way round that, for a +different reason given at the end of this section, so the refusal is a diagnosis +and not a redirection. -Capturing a CUDA graph from the stream its `execute()` runs on is refused with -`Error::NotSupported`. The handoff between one enqueue and the next waits on an -event recorded outside the capture, and `cudaStreamWaitEvent` on such an event -fails with `cudaErrorStreamCaptureIsolation` and invalidates the capture under +The handoff between one enqueue and the next waits on an event recorded outside +the capture, and `cudaStreamWaitEvent` on such an event fails with +`cudaErrorStreamCaptureIsolation` and invalidates the capture under every capture mode, `cudaStreamCaptureModeRelaxed` included; the caller learns of that only when `cudaStreamEndCapture` hands back `cudaErrorStreamCaptureInvalidated` and a null graph, long after the call that caused it. A growth is mode-dependent @@ -258,52 +280,61 @@ the pool's calls are permitted and refusing would be wrong. CUDA has no query fo "is a capture live in this process", so the rule is the caller's to keep: do not run a pooled engine while any capture is open anywhere in the process. -Turning the option off takes the pool's calls out of the way, and the refusal with -them: the check lives inside the pooled path, so a context that keeps its own -scratch is never asked. That does not make `execute()` capturable in general. -Three kinds of call remain that a capture may not survive, and both the modes and -the conditions are worth stating exactly, because with nothing refusing them the -failure is silent: - -- The host wait on a previous enqueue, made when the last call on this handle - returned with its enqueue still in flight. Prohibited under `Global` and - `ThreadLocal`, permitted under `Relaxed`. -- A staging buffer's `cudaMalloc`, made when a tensor is empty or host-resident - and this handle holds no cached buffer for it that is already large enough, and - the `cudaFree` of a cached buffer a larger tensor has outgrown. Prohibited under - `Global` and `ThreadLocal`, permitted under `Relaxed`. -- The `cudaStreamSynchronize` on the selected stream that ends the call, made - when any tensor is staged through host memory, when an output is aliased, or - when no caller stream is set. Prohibited under **every** mode, `Relaxed` - included: a capturing stream cannot be synchronized at all. - -So a call captures cleanly with the option off when it runs on the capturing -stream under a `CallerStreamGuard`, binds only non-empty device-resident tensors, -has no aliased outputs, and follows no call that left an enqueue in flight. A -first call on a freshly loaded handle over device-resident inputs and outputs -meets all four, and that is measured rather than argued. But such a call then -leaves its own enqueue in flight, so under `Global` and `ThreadLocal` the next one -on the same handle fails the last condition: with the option off a handle is good -for one captured call, not for capture in general. `Relaxed` permits that wait, so -the limit is not one there -- it permits the staging calls too, and runs each of -them for real and outside the graph, which is a hazard of its own rather than a -way to capture more. `AnUnpooledEngineCapturesOnAFirstCallThatStagesNothing` -and `AnUnpooledEngineInvalidatesACaptureOnceAnEnqueueIsInFlight` in -`tests/cpp/executorch/test_shared_scratch_backend.cpp` are those two outcomes. - -Moving the pool's check out of the pooled path so that it refused both ways would -make the error honest but would also refuse the calls above that do capture, for -contexts that never touch the pool. That is a behaviour change for callers who -have not turned the option on, so the check stays where it is and this section is -what the refusal message points at. - -`cudaDeviceReset()` is not survivable and is not guarded against. The pool holds -its buffer and its handoff event for the process lifetime, and a reset destroys -the primary context under both. The next pooled `execute()` on that device then -waits on a destroyed event and hands the engine a pointer that is no longer a -device allocation. Guarding this would mean revalidating both on every call, and -the check that would catch it is as expensive as the work it protects. Treat a -device the backend has run a pooled engine on as one that must not be reset. +**With the option off nothing refuses, and capture still does not work.** +Turning the option off does take the pool's calls out of the way, and the refusal +with them, because the check lives inside the pooled path. What it leaves is a +capture that succeeds and a handle that cannot be used again. + +A call captures cleanly with the option off when it runs on the capturing stream +under a `CallerStreamGuard`, binds only non-empty device-resident tensors, has no +aliased outputs, and follows no call that left an enqueue in flight. A first call +on a freshly loaded handle over device-resident inputs and outputs meets all four: +`execute()` then makes only the device query, the device switch, the shape and +address binding, `enqueueV3` and the completion record, and none of those is a +call a capture cannot take. Measured, not argued -- the graph instantiates, +replays, and reproduces the bytes the same engine produces outside a capture. + +The completion record is what costs the handle. `execute()` records that event so +the next call and `~EngineHandle` can wait for an enqueue that outlived the +return; under capture the record becomes a node of the graph rather than an event +the host can wait on. The next `execute()` on that handle fails +`Error::InvalidProgram` on the wait at its head -- measured, logging +`cudaEventSynchronize failed: invalid argument` -- and so does the destructor's +wait. No call shape avoids this: the alternative to recording the event is the +`cudaStreamSynchronize` that ends any call staging through host memory, aliasing +an output or running with no caller stream, and a capturing stream cannot be +synchronized under any mode, `Relaxed` included. So the only call a capture can +take is the one that arms the event. +`AnUnpooledEngineCaptureReplaysAndLeavesTheHandleUnusable` in +`tests/cpp/executorch/test_shared_scratch_backend.cpp` is that whole sequence. + +Two further calls invalidate a capture with the option off, before the handle +question arises, and nothing refuses them either: the host wait on a previous +enqueue, and a staging buffer's `cudaMalloc` or the `cudaFree` of a cached buffer +a larger tensor has outgrown. Both are prohibited under `Global` and +`ThreadLocal` and permitted under `Relaxed`, where they run for real and outside +the graph. +`AnUnpooledEngineInvalidatesACaptureOnceAnEnqueueIsInFlight` covers the first. + +Two changes would be needed to make capture work rather than documenting it, and +neither is local. The completion event would have to be reworked, and the +asynchronous return described in the caller-stream contract above rests on it. +And moving the pool's check out of the pooled path so that capture were refused +both ways would refuse the calls above, which do capture, for contexts that never +touch the pool -- a behaviour change for callers who have not turned the option +on. So the option-off path stays as it is, and this section, the option key's +documentation in `TensorRTBackend.h` and the refusal message all say that capture +is unsupported rather than pointing a caller at it. + +### cudaDeviceReset() is not survivable + +The pool is not guarded against it: it holds its buffer and its handoff event for +the process lifetime, and a reset destroys the primary context under both. The +next pooled `execute()` on that device then waits on a destroyed event and hands +the engine a pointer that is no longer a device allocation. Guarding this would +mean revalidating both on every call, and the check that would catch it is as +expensive as the work it protects. Treat a device the backend has run a pooled +engine on as one that must not be reset. ## Standalone Backend Archive diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h index eb719a115d5..545f4556fe0 100644 --- a/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h @@ -20,12 +20,6 @@ #include #include -namespace nvinfer1 { -// Named by install_pooled_scratch below. Forward-declared, as TensorRT's own -// headers do, so this header still builds against the CUDA headers alone. -class IExecutionContext; -} // namespace nvinfer1 - namespace torch_tensorrt { namespace executorch_backend { @@ -124,9 +118,10 @@ class SharedScratchPool { // The process-wide per-device pool the TensorRT backend runs on. One buffer per // device, grown to the largest requirement any call on that device has asked -// for, serves every kUSER_MANAGED context on it, 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. +// for, serves every kUSER_MANAGED context on it that needs activation scratch -- +// a context whose engine needs none under any shape claims nothing here -- 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 @@ -196,11 +191,9 @@ inline cudaEvent_t shared_scratch_mark_in_flight(SharedScratchDevice& dev) { // 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. +// be reading and writing `buffer` has been recorded; the caller must wait for that +// event on the host before it frees. 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 @@ -208,9 +201,13 @@ inline cudaEvent_t shared_scratch_mark_in_flight(SharedScratchDevice& dev) { // 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. +// Both the wait and the free belong outside `dev.mu`, and for the same reason: +// each blocks on device work, and under the lock that makes an unrelated claim on +// this device wait for it. The free is a device-wide synchronization; the wait is +// for an inference to finish. That the wait may then also pick up enqueues made +// after the lock was dropped costs nothing, because the free that follows it waits +// for every stream on the device anyway, and it stays sufficient for this buffer +// by the ordering argument above. struct RetiredScratch { void* buffer = nullptr; cudaEvent_t wait_for = nullptr; @@ -226,7 +223,8 @@ struct RetiredScratch { // 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. +// still be using. Every path clears `out_retired` first, so a caller reusing one +// across calls is not handed a buffer an earlier call already disposed of. template void* shared_scratch_get_or_grow( SharedScratchDevice& dev, @@ -234,6 +232,7 @@ void* shared_scratch_get_or_grow( std::size_t& out_size, Alloc alloc, RetiredScratch& out_retired) { + out_retired = RetiredScratch{}; if (dev.buffer != nullptr && dev.capacity >= need) { out_size = dev.capacity; return dev.buffer; @@ -252,23 +251,5 @@ void* shared_scratch_get_or_grow( return p; } -// Installs `bytes` of `buffer` as the activation scratch of `ctx`, a -// kUSER_MANAGED context, and reports whether TensorRT accepted it. Logs the -// refusal, naming `device_id`, when it did not. -// -// The check is the point. setDeviceMemoryV2 returns void and refuses a buffer -// smaller than the bound shapes need, so a caller that does not ask cannot tell -// an accepted install from a refused one -- and a refused one leaves the context -// pointed at the buffer it was last given, which a pool growth may since have -// freed. The engine then reads and writes freed memory with enqueueV3 reporting -// success. The refusal is read back through an IErrorRecorder scoped to this one -// call, the only channel that hands it to the caller: with no recorder attached -// TensorRT writes it to the runtime's ILogger and this function has nothing to -// return. -// -// Defined in TensorRTBackend.cpp, which owns the TensorRT dependency; a binary -// that links this header alone never names it. -bool install_pooled_scratch(nvinfer1::IExecutionContext& ctx, void* buffer, std::size_t bytes, int device_id); - } // namespace executorch_backend } // namespace torch_tensorrt diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index e9c8aebe7c6..201cc49d7ac 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -90,7 +90,9 @@ EngineHandle::~EngineHandle() { cudaError_t err = cudaEventSynchronize(inflight_event); if (err != cudaSuccess) { ET_LOG(Error, "EngineHandle::~EngineHandle: cudaEventSynchronize failed: %s", cudaGetErrorString(err)); - cudaGetLastError(); // clear sticky error; tear down regardless + // Clears a non-sticky error so it does not resurface under the name of the + // next call here; a sticky one survives the clear. Tear down regardless. + cudaGetLastError(); } inflight_pending = false; } @@ -247,7 +249,9 @@ class SharedScratchClaim { SharedScratchClaim(const SharedScratchClaim&) = delete; SharedScratchClaim& operator=(const SharedScratchClaim&) = delete; ~SharedScratchClaim() { - release(); + // The return is the caller's to report, and on this path there is no caller + // left to report it to: an execute() that returned early has already failed. + (void)release(); } SharedScratchDevice& hold(int device_id) { @@ -263,55 +267,87 @@ class SharedScratchClaim { 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) { + // Takes ownership of a buffer a growth displaced, to be waited for and freed by + // release(). `wait_for` is the marker event the enqueues that used it were + // recorded on, or null if none were. + void retire(void* buffer, cudaEvent_t wait_for) { retired_ = buffer; + retired_wait_ = wait_for; } // 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. + // pointer this claim no longer holds the lock for. Then waits for the enqueues + // that used whatever a growth displaced and frees it -- both after the unlock, + // because both block on device work that another pooled engine on this device + // has nothing to do with. The free waits for every stream on the device; + // the wait waits for an inference to finish, so under the lock it would move the + // point at which two pooled calls serialize from submission, which is what the + // execute() contract promises, to the completion of the call before. + // + // Waiting here rather than under the lock can wait for more than the retired + // buffer's own users: this claim records its own enqueue on the same event + // before it unlocks, and a claimant entering afterwards records on it again. It + // costs nothing. The free below waits for the whole device, so it already covers + // everything the wait could have picked up, and the wait is still sufficient -- + // every claimant orders its stream after this event before enqueueing, so the + // latest recording completes only once all the earlier ones have. // - // Unlocking bounds who waits, not how long. This free is still the one call on - // the growth path that waits device-wide -- measured, the allocation and the - // marker wait do not -- and a caller whose own later work some queued stream is - // waiting on will not get this call back at all. TensorRTBackend.h says so in - // the execute() contract, and the README says what to do about it. Making the - // free stream-ordered is not a local change: cudaFreeAsync pairs only with - // cudaMallocAsync, so it would move every pool allocation onto the - // stream-ordered allocator to bound this one wait. + // Unlocking bounds who waits, not how long. The free is still the one call on + // the growth path that waits device-wide -- measured, the allocation does not -- + // and a caller whose own later work some queued stream is waiting on will not + // get this call back at all. TensorRTBackend.h says so in the execute() contract, + // and the README says what to do about it. Making the free stream-ordered is not + // a local change: cudaFreeAsync pairs only with cudaMallocAsync, so it would move + // every pool allocation onto the stream-ordered allocator to bound this one wait. // - // Frees on the current device, which must still be the buffer's. - void release() { + // Returns false when the wait failed, which leaves the buffer leaked rather than + // freed under a live enqueue; the caller reports it. Frees on the current device, + // which must still be the buffer's. + bool release() { if (lock_.owns_lock()) { lock_.unlock(); } dev_ = nullptr; - if (retired_ != nullptr) { - const cudaError_t err = cudaFree(retired_); - if (err != cudaSuccess) { - // cudaFree synchronizes, so what it reports is more often an earlier - // asynchronous fault on this device than a fault in the free -- which is - // why the message does not call it one. + if (retired_ == nullptr) { + return true; + } + void* const retired = retired_; + const cudaEvent_t wait_for = retired_wait_; + retired_ = nullptr; + retired_wait_ = nullptr; + + if (wait_for != nullptr) { + const cudaError_t wait_err = cudaEventSynchronize(wait_for); + if (wait_err != cudaSuccess) { + // 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. What it reports is + // usually an asynchronous fault raised by earlier work on this device. ET_LOG( Error, - "TensorRTBackend::execute: freeing the shared activation scratch buffer that a pool growth replaced on device %d reported %s; a device-wide free reports whatever fault this device is already in, so this need not be the pool's", + "TensorRTBackend::execute: waiting for the enqueue on the replaced shared activation scratch on device %d failed (%s), which for a wait on device work is usually an earlier asynchronous fault on this device surfacing here; leaking that buffer rather than freeing it under a live enqueue", device_id_, - cudaGetErrorString(err)); - // Clears a non-sticky error so it does not resurface under the name of the - // next CUDA call in execute(). A sticky one survives the clear and will - // resurface anyway; the caller learns of it from that call. + cudaGetErrorString(wait_err)); cudaGetLastError(); + return false; } - retired_ = nullptr; } + + const cudaError_t err = cudaFree(retired); + if (err != cudaSuccess) { + // cudaFree synchronizes, so what it reports is more often an earlier + // asynchronous fault on this device than a fault in the free -- which is + // why the message does not call it one. + ET_LOG( + Error, + "TensorRTBackend::execute: freeing the shared activation scratch buffer that a pool growth replaced on device %d reported %s; a device-wide free reports whatever fault this device is already in, so this need not be the pool's", + device_id_, + cudaGetErrorString(err)); + // Clears a non-sticky error so it does not resurface under the name of the + // next CUDA call in execute(). A sticky one survives the clear and will + // resurface anyway; the caller learns of it from that call. + cudaGetLastError(); + } + return true; } private: @@ -319,6 +355,7 @@ class SharedScratchClaim { int device_id_ = -1; std::unique_lock lock_; void* retired_ = nullptr; + cudaEvent_t retired_wait_ = nullptr; }; // What a call needing no activation scratch is given when the pool holds nothing @@ -374,18 +411,27 @@ Error refuse_pooled_call_on_a_capturing_stream(cudaStream_t stream, int device_i } ET_LOG( Error, - "TensorRTBackend::execute: the selected stream is capturing a CUDA graph (%s), which the shared activation scratch pool on device %d does not support. Loading this engine with '%s' off takes the pool's calls out of the way, but not every call a capture cannot take, and nothing refuses on that path: the capture section of the backend README says which calls remain and when a call survives them.", + "TensorRTBackend::execute: the selected stream is capturing a CUDA graph (%s), which the shared activation scratch pool on device %d does not support. Neither does the rest of this delegate: loading the engine with '%s' off removes this refusal, but a call captured that way records the handle's completion event into the graph, and every later call on that handle then fails. The capture section of the backend README has the detail.", capture_err == cudaSuccess ? "capture in progress" : cudaGetErrorString(capture_err), device_id, kSharedActivationScratchKey); - cudaGetLastError(); // the query's own failure is this call's, not the next one's + if (capture_err != cudaSuccess) { + // The query's own failure is this call's, not the next one's. Only here: on + // the ordinary refusal the query succeeded, so anything pending on this thread + // was left by earlier work and belongs to whoever calls CUDA next. + cudaGetLastError(); + } return Error::NotSupported; } -// 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. -// 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. +// Sets out_ptr to a buffer of at least `need` bytes on `device_id`, with `stream` +// ordered after the enqueue that last used the buffer. 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. +// +// The buffer's capacity is not reported, because no caller has any use for it: +// what a call installs on its context is its own requirement, not whatever the +// pool grew to for someone else. // // A `need` of zero asks for whatever the pool already holds rather than for a // buffer of its own: any live buffer is large enough for a call that needs @@ -404,8 +450,7 @@ Error get_or_grow_shared_scratch( int device_id, size_t need, cudaStream_t stream, - void*& out_ptr, - size_t& out_size) { + void*& out_ptr) { SharedScratchDevice& dev = claim.hold(device_id); // kMinPooledScratchBytes rather than the capacity the pool already holds: the @@ -419,9 +464,10 @@ Error get_or_grow_shared_scratch( const SharedScratchHandoff handoff = shared_scratch_claim_event(dev, []() -> cudaEvent_t { cudaEvent_t event = nullptr; // Blocking-sync so the host yields instead of busy-spinning. The one host wait - // on this event -- the one below, before a displaced buffer is freed -- runs - // with the device's lock held, so it already holds off every other pooled - // engine on the device; spinning would burn a core for that whole time as well. + // on this event is the one SharedScratchClaim::release() makes before freeing a + // buffer a growth displaced, and it waits for a whole inference; spinning would + // burn a core for that time and be no faster, since it is followed by a + // device-wide cudaFree. if (cudaEventCreateWithFlags(&event, cudaEventDisableTiming | cudaEventBlockingSync) != cudaSuccess) { return nullptr; } @@ -447,10 +493,12 @@ Error get_or_grow_shared_scratch( const bool first_buffer = dev.buffer == nullptr; RetiredScratch retired; + // Written by the pool and read by nobody, for the reason given above. + size_t capacity = 0; void* const buffer = shared_scratch_get_or_grow( dev, need, - out_size, + capacity, [device_id, first_buffer](size_t bytes) -> void* { void* p = nullptr; if (cudaMalloc(&p, bytes) != cudaSuccess) { @@ -474,30 +522,10 @@ Error get_or_grow_shared_scratch( 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. - // - // The call fails with it. cudaEventSynchronize waits on device work, so what - // it reports is usually an asynchronous fault raised by earlier work on this - // device; running on regardless would enqueue against a context on a device - // already in that state and report the same fault later, under the name of - // whichever call surfaced it next. - ET_LOG( - Error, - "TensorRTBackend::execute: waiting for the enqueue on the replaced shared activation scratch on device %d failed (%s), which for a wait on device work is usually an earlier asynchronous fault on this device surfacing here; leaking that buffer rather than freeing it under a live enqueue", - device_id, - cudaGetErrorString(err)); - return Error::InvalidProgram; - } - } + // Both the wait for the enqueues that used the retired buffer and the free of it + // happen at release(), with the device's lock dropped; see SharedScratchClaim. + // Nothing here makes a CUDA call that blocks on device work under that lock. + claim.retire(retired.buffer, retired.wait_for); out_ptr = buffer; return Error::Ok; @@ -505,22 +533,15 @@ Error get_or_grow_shared_scratch( // 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. +// +// Call on a `claim` that get_or_grow_shared_scratch returned Error::Ok on and that +// still holds the device's lock. Both halves matter, and together they are why +// neither the device nor the event below is checked: a claim's device is non-null +// exactly while it holds the lock, and get_or_grow_shared_scratch has already +// failed the call with Error::Internal if the marker had no event, which only the +// test-only reset clears again and that needs the lock this claim is holding. 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); - if (event == nullptr) { - ET_LOG( - Error, - "TensorRTBackend::execute: shared activation scratch on device %d has no handoff event", - claim.device_id()); - return Error::Internal; - } + const cudaEvent_t event = shared_scratch_mark_in_flight(*claim.device()); const cudaError_t err = cudaEventRecord(event, stream); if (err != cudaSuccess) { ET_LOG( @@ -589,8 +610,8 @@ class ScopedErrorRecorder final : public nvinfer1::IErrorRecorder { } // namespace -// Declared in SharedScratchPool.h. See that header for why the install is -// checked rather than made and trusted. +// Declared in TensorRTBackend.h. See that header for why the install is checked +// rather than made and trusted. bool install_pooled_scratch(nvinfer1::IExecutionContext& ctx, void* buffer, size_t bytes, int device_id) { ScopedErrorRecorder recorder; nvinfer1::IErrorRecorder* const previous = ctx.getErrorRecorder(); @@ -1027,8 +1048,10 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // Refused here, ahead of every call this function makes that a capture cannot // take. The pool's own are not the only ones: the wait on a previous enqueue just // below, and the cudaMalloc that grows a host-input staging buffer further down, - // would invalidate the capture before the pooled path was ever reached, so a - // refusal any later would arrive after the thing it exists to protect was gone. + // are prohibited under every mode but cudaStreamCaptureModeRelaxed, so outside + // that mode they would invalidate the capture before the pooled path was ever + // reached and a refusal any later would arrive after the thing it exists to + // protect was gone. if (pooled_scratch) { const Error capture_err = refuse_pooled_call_on_a_capturing_stream(stream, engine->device_id); if (capture_err != Error::Ok) { @@ -1349,11 +1372,20 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // The two causes are told apart by the install rather than by the query, which // is the only place the difference is observable. Where the shapes genuinely // need nothing the engine expects nothing and the minimum is accepted; where the - // query failed the engine still expects what it always did, and a pool holding - // less than that is refused -- and this call ends here rather than enqueueing + // query failed the engine still expects what it always did, is offered that same + // minimum, and refuses it -- and this call ends here rather than enqueueing // against whatever pointer the context was last given, which a growth may since // have freed. // + // What is installed is this call's own requirement, not the capacity the pool + // happens to hold. They differ whenever an earlier call asked for more, and + // handing this context the larger figure would tell TensorRT it owns bytes + // holding another engine's activations -- the pool never clears the buffer -- + // and would leave an engine overrunning its own requirement inside the region it + // was told it owns, so nothing catches it. It would also blunt the refusal + // above: after any growth, the capacity may cover what a failed query concealed, + // and the install that is meant to catch it succeeds. + // // 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 @@ -1363,13 +1395,15 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* if (pooled_scratch) { const size_t need = ctx->updateDeviceMemorySizeForShapes(); void* pool = nullptr; - size_t pool_size = 0; - const Error scratch_err = - get_or_grow_shared_scratch(scratch_claim, 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); if (scratch_err != Error::Ok) { return scratch_err; } - if (!install_pooled_scratch(*ctx, pool, pool_size, engine->device_id)) { + // The same substitution get_or_grow_shared_scratch makes for a zero: a + // context whose engine needs scratch under some shape has to be given a + // buffer whatever this call's shapes need, or enqueueV3 refuses it. + const size_t install_bytes = need == 0 ? kMinPooledScratchBytes : need; + if (!install_pooled_scratch(*ctx, pool, install_bytes, engine->device_id)) { return Error::InvalidState; } scratch_from_pool = true; @@ -1403,7 +1437,17 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // 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(); + // The release is where a growth pays for the buffer it replaced: it waits for + // the enqueues that used it and frees it, both with the lock dropped. + if (!scratch_claim.release()) { + // The wait failed, which for a wait on device work means this device is + // already in a faulted state. The enqueue above is submitted and this handle's + // completion marker is not armed yet, so drain before reporting, or a later + // call reconfigures exec_ctx while that enqueue is still running. + (void)cudaStreamSynchronize(stream); + engine->inflight_pending = false; + return Error::InvalidProgram; + } // Caller-owned KV: reflect each engine in-place update into its delegate output // EValue (D2D on the same stream, after the engine work). @@ -1486,19 +1530,30 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // set_option // --------------------------------------------------------------------------- Error TensorRTBackend::set_option(ET_UNUSED BackendOptionContext& context, const Span& backend_options) { + // The whole span is read before anything is stored. A span is one request, so a + // caller told it was refused must not find part of it applied -- and the part + // that would be applied here is process-wide and governs every engine loaded + // after it. Where a span names this key more than once the last one wins, which + // is what applying each in turn did. + bool requested = false; + bool have_request = false; 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 { + const bool* const val = std::get_if(&option.value); + if (val == nullptr) { ET_LOG(Error, "TensorRTBackend::set_option: option '%s' must be a boolean", kSharedActivationScratchKey); return Error::InvalidArgument; } + requested = *val; + have_request = true; } } + if (have_request) { + scratch_enabled.store(requested, std::memory_order_relaxed); + } return Error::Ok; } diff --git a/tests/cpp/executorch/test_shared_scratch_backend.cpp b/tests/cpp/executorch/test_shared_scratch_backend.cpp index 9a944c1344d..03c959ed745 100644 --- a/tests/cpp/executorch/test_shared_scratch_backend.cpp +++ b/tests/cpp/executorch/test_shared_scratch_backend.cpp @@ -22,8 +22,9 @@ // green run on a host with no GPU says nothing about the pool -- and a skipped // gtest case exits zero, which Bazel reports as a passing target. A run that is // meant to have a device says so through TORCHTRT_EXECUTORCH_REQUIRE_CUDA, and -// then a skip is a failure instead; the count of skipped cases is printed either -// way. The CI invocation in .github/workflows/executorch-build-linux.yml passes +// then a skip is a failure instead, which each failing case says for itself. Where +// that variable is unset, a count of the cases that skipped is printed at the end +// of the suite. The CI invocation in .github/workflows/executorch-build-linux.yml passes // that variable, and the job it sits in asks for a GPU runner and starts its // container with every GPU attached. What still keeps these cases off most runs // is the lane gate on the whole ExecuTorch job in ci-linux-x86_64.yml, not the @@ -506,6 +507,54 @@ bool device_bytes_in_use(std::size_t& out) { return true; } +// Takes device memory until an allocation of the size a caller names can no longer +// succeed, and gives it all back when it goes out of scope. The only way to reach +// the pool's allocation-failure path from inside the process, and the reason it is +// held for as short a window as possible: while it is up, every other process on +// this device is out of memory too. +class DeviceMemoryHog { + public: + DeviceMemoryHog() = default; + DeviceMemoryHog(const DeviceMemoryHog&) = delete; + DeviceMemoryHog& operator=(const DeviceMemoryHog&) = delete; + + ~DeviceMemoryHog() { + release(); + } + + // Chunked from large to small, so a device with tens of gigabytes free is filled + // in a few dozen allocations rather than thousands. + bool leave_less_free_than(std::size_t bytes) { + for (std::size_t chunk = std::size_t{1} << 30; chunk >= (std::size_t{1} << 20); chunk /= 4) { + while (free_bytes_above(bytes)) { + void* block = nullptr; + if (cudaMalloc(&block, chunk) != cudaSuccess) { + cudaGetLastError(); + break; + } + blocks_.push_back(block); + } + } + return !free_bytes_above(bytes); + } + + void release() { + for (void* block : blocks_) { + cudaFree(block); + } + blocks_.clear(); + } + + private: + static bool free_bytes_above(std::size_t bytes) { + std::size_t free_bytes = 0; + std::size_t total_bytes = 0; + return cudaMemGetInfo(&free_bytes, &total_bytes) == cudaSuccess && free_bytes > bytes; + } + + std::vector blocks_; +}; + // Opens a live kUSER_MANAGED context over `blob`'s engine with one shape bound, // the state execute() installs a scratch buffer into. Held by the caller, unlike // measure_engine_scratch's, which is gone by the time it returns its figure. @@ -638,6 +687,12 @@ class SharedScratchBackendTest : public ::testing::Test { if (skipped_for_no_device_ == 0) { return; } + // Nothing skipped when the requirement is on -- SetUp counts the missing device + // and then fails the case. Printing here would put a skip banner under those + // failures and tell the reader to set the very variable that produced them. + if (cuda_device_is_required()) { + return; + } const ::testing::TestSuite* const suite = ::testing::UnitTest::GetInstance()->current_test_suite(); std::fprintf( stderr, @@ -781,6 +836,32 @@ TEST_F(SharedScratchBackendTest, SetOptionRejectsANonBooleanAndLeavesTheSettingA EXPECT_TRUE(engine.handle()->shared_scratch) << "a rejected option still moved the shared-scratch setting"; } +// The case above sends the bad entry on its own, where returning before the store +// and storing before the return look the same. A span carrying a good entry ahead +// of a bad one tells them apart, and it is the span a caller writes when it sets +// several options at once. +TEST_F(SharedScratchBackendTest, SetOptionRejectsASpanWithoutApplyingTheEntriesBeforeTheBadOne) { + ASSERT_EQ(set_shared_scratch(backend_, false), Error::Ok); + + BackendOption turn_on; + std::strncpy(turn_on.key, kOptionKey, sizeof(turn_on.key) - 1); + turn_on.value = true; + // The same key again, so what the span asks for is unambiguous: this entry + // cannot be read as addressed to some other backend. + BackendOption wrong_type; + std::strncpy(wrong_type.key, kOptionKey, sizeof(wrong_type.key) - 1); + wrong_type.value = 1; + BackendOption options[2] = {turn_on, wrong_type}; + BackendOptionContext context; + + EXPECT_EQ(backend_.set_option(context, Span(options, 2)), Error::InvalidArgument); + + LoadedEngine engine; + ASSERT_EQ(engine.load(blob(), 28), Error::Ok); + EXPECT_FALSE(engine.handle()->shared_scratch) + << "a span that was refused still turned the pool on for every engine loaded after it"; +} + // 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) { @@ -967,10 +1048,28 @@ TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItRep std::size_t before = 0; ASSERT_TRUE(device_bytes_in_use(before)) << "cudaMemGetInfo failed, so this test measured nothing"; + const auto measurement_began = std::chrono::steady_clock::now(); ASSERT_EQ(big.run(stream), Error::Ok); ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); std::size_t after = 0; ASSERT_TRUE(device_bytes_in_use(after)) << "cudaMemGetInfo failed, so this test measured nothing"; + // A control window of the same length as the measurement, with nothing of this + // test's running in it. On a device this test has to itself, device-wide usage + // is the same at both ends of it; anything else means another process is moving + // memory on the same timescale as the growth, which is what makes the bounds + // below report a figure the pool did not produce. The same length matters: two + // readings taken back to back would sample microseconds against the growth's + // milliseconds and would miss almost everything. It is still a sample of a + // different window, so it narrows the misdiagnosis rather than removing it, and + // the upper bound names the cause as well. Measured with another process cycling + // 256 MiB allocations on this device: of 8 runs, 1 skipped here, 1 failed that + // bound and 1 failed the lower-direction assertion below, whose message is + // accurate already. The exclusive tag keeps other Bazel actions off this device, + // not other processes. + std::this_thread::sleep_for(std::chrono::steady_clock::now() - measurement_began); + std::size_t settled = 0; + ASSERT_TRUE(device_bytes_in_use(settled)) << "cudaMemGetInfo failed, so this test measured nothing"; + const bool device_was_quiet = settled == after; 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; @@ -978,7 +1077,9 @@ TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItRep // 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. + // that address has been freed. This run is also the one place the suite installs + // a size below the pool's capacity, since that is what this engine's shapes + // need, so a size TensorRT refuses fails here. ASSERT_TRUE(small.fill_output(kSentinel)); ASSERT_EQ(small.run(stream), Error::Ok); ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); @@ -987,13 +1088,6 @@ TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItRep 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) @@ -1004,6 +1098,21 @@ TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItRep 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"; + + // Last, so the output comparisons above are made either way. + if (!device_was_quiet) { + GTEST_SKIP() << "device-wide memory in use moved from " << after << " to " << settled + << " bytes with nothing of this test's running in between, so another process is allocating on this " + "device and the two bounds below would measure it rather than the growth"; + } + 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 either the buffer it replaced was not freed or another process allocated on this " + "device across the measurement, which the control window above samples for and cannot rule out"; } // --------------------------------------------------------------------------- @@ -1222,11 +1331,19 @@ TEST_F(SharedScratchBackendTest, AnUndersizedScratchInstallIsSeenAsRefused) { // only the refusal, so every other case here stays green. // // Forcing the refusal itself through execute() is not available. The size -// execute() installs is the pool's capacity, and the pool is grown to at least -// what the query just returned, so the installed buffer is never short of what +// execute() installs is what the per-shape query just returned, over a buffer the +// pool has grown to at least that, so the installed size is never short of what // the context expects. Only a failed query makes the two disagree, and that // cannot be induced from inside this process. // +// Nor is the choice between that size and the pool's capacity observable through +// TensorRT: an install larger than the context expects is accepted in silence, so +// no black-box case here can tell the two apart, and the argument for the smaller +// one is the hazard it removes rather than anything a test can read back. What the +// suite does see is the other direction -- ALargerEngineGrowsThePoolAndFreesTheBufferItReplaces +// runs the smaller engine again after the growth, so its install is smaller than +// the capacity, and a size TensorRT would not accept fails that case. +// // What is observable is the recorder. install_pooled_scratch attaches one for // the duration of the install and restores the previous one after, because that // is the only channel TensorRT reports a refusal on. A recorder this test leaves @@ -1267,6 +1384,71 @@ TEST_F(SharedScratchBackendTest, ExecuteInstallsPooledScratchThroughTheCheckedHe ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); } +// Every pooled path out of execute() that is not the happy one returns while the +// claim still holds the device's lock, and leaves the claim's destructor to drop +// it. Nothing else here takes one of those paths, and the failure a regression +// there produces is not a wrong answer: it is every later pooled call on the +// device blocking forever with nothing logged. +// +// Only one of them can be reached from inside the process. A refused install needs +// the pool to hold less than the context expects, which it never does; a failed +// enqueue and a failed completion record need TensorRT or CUDA to fail a call that +// is correct as made. The pool's own allocation can be made to fail by taking the +// device's memory first, and it returns through the same destructor as the others. +TEST_F(SharedScratchBackendTest, AFailedPooledAllocationLeavesTheDeviceLockFree) { + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + + LoadedEngine pooled; + ASSERT_EQ(pooled.load(blob(), 31), Error::Ok); + ASSERT_TRUE(pooled.handle()->shared_scratch); + ASSERT_TRUE(pooled.handle()->engine_needs_scratch) << "this engine skips the pool, so it never allocates from it"; + const int device_id = pooled.handle()->device_id; + ASSERT_EQ(shared_scratch_capacity_for_testing(device_id), 0u) + << "the pool already holds a buffer, so this call would reuse it rather than allocate"; + + Error failed_run = Error::Ok; + { + DeviceMemoryHog hog; + if (!hog.leave_less_free_than(scratch_bytes_)) { + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); + GTEST_SKIP() << "could not take the device below the " << scratch_bytes_ + << " bytes this engine's scratch needs, so its allocation would have succeeded"; + } + failed_run = pooled.run(stream); + } + + if (failed_run == Error::Ok) { + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); + GTEST_SKIP() << "the pooled allocation succeeded with the device full, so this run did not take the path under " + "test"; + } + ASSERT_EQ(failed_run, Error::MemoryAllocationFailed) + << "the run failed somewhere other than the pool's allocation, so it says nothing about that path"; + + SharedScratchDevice& dev = scratch_pool().get(device_id); + const bool lock_free = dev.mu.try_lock(); + if (lock_free) { + dev.mu.unlock(); + } + ASSERT_TRUE(lock_free) << "a pooled call that returned early left the device's pool lock held, so every later pooled " + "call on this device blocks forever"; + + // And the slot is still usable, not merely unlocked: with the memory back, the + // next call allocates and runs. + ASSERT_TRUE(pooled.fill_output(kSentinel)); + EXPECT_EQ(pooled.run(stream), Error::Ok) << "the pool was left unusable by a call whose allocation failed"; + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::vector actual = pooled.read_output(); + ASSERT_EQ(actual.size(), kElems); + EXPECT_NE(actual[0], kSentinel) << "the engine did not write its output, so the run above did not reach it"; + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + // --------------------------------------------------------------------------- // Stream capture // --------------------------------------------------------------------------- @@ -1322,6 +1504,46 @@ TEST_F(SharedScratchBackendTest, APooledEngineRefusesToRunWhileItsStreamIsCaptur ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); } +// The refusal clears the CUDA error state only when the capture query itself +// failed, because then the error is the query's own. On the ordinary refusal -- +// the query succeeded and reported a capture -- whatever is pending was left by +// earlier work on this thread, and clearing it takes it away from the caller, +// which learns of it from no later call either: the error is not sticky. +TEST_F(SharedScratchBackendTest, ARefusedPooledCallLeavesTheCallersPendingCudaErrorAlone) { + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + + LoadedEngine pooled; + ASSERT_EQ(pooled.load(blob(), 32), Error::Ok); + ASSERT_TRUE(pooled.handle()->shared_scratch); + + // An error of the caller's own, unread. Nothing but cudaGetLastError clears one, + // so it is still there when the refused call returns unless that call took it. + void* impossible = nullptr; + ASSERT_EQ(cudaMalloc(&impossible, static_cast(1) << 62), cudaErrorMemoryAllocation); + + ASSERT_EQ(cudaStreamBeginCapture(stream, cudaStreamCaptureModeRelaxed), cudaSuccess); + const Error refused = pooled.run(stream); + cudaGraph_t graph = nullptr; + const cudaError_t end_err = cudaStreamEndCapture(stream, &graph); + if (graph != nullptr) { + cudaGraphDestroy(graph); + } + // Read here, before any assertion can return early with it still pending and + // hand it to the next case. + const cudaError_t pending = cudaGetLastError(); + + ASSERT_EQ(refused, Error::NotSupported) << "the run was not refused, so it did not take the path under test"; + EXPECT_EQ(end_err, cudaSuccess) << "the refused run invalidated the capture: " << cudaGetErrorString(end_err); + EXPECT_EQ(pending, cudaErrorMemoryAllocation) + << "the refusal cleared an error this backend did not cause, so the caller never sees it: got " + << cudaGetErrorString(pending); + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + // Where in execute() the refusal comes, which the relaxed-mode case above cannot // see. Under cudaStreamCaptureModeRelaxed everything execute() does before it // reaches the pool is permitted, so a guard sitting anywhere ahead of the pool @@ -1372,16 +1594,26 @@ TEST_F(SharedScratchBackendTest, APooledEngineRefusesACaptureBeforeAnythingCanIn } // The two cases above are the pool refusing. These two are what a caller gets -// when it takes the refusal's advice and loads with the option off, which the -// README and the refusal message both describe and neither may overstate. +// with the option off, where nothing refuses because the guard sits inside the +// pooled branch. Between them they are why the README and the refusal message +// say this delegate does not support capture at all, rather than offering the +// option-off path as the way to capture. // -// Nothing refuses then -- the guard sits inside the pooled branch -- so the -// question is only whether execute()'s own calls are ones a capture can take. -// On a first call over non-empty device-resident tensors, submitted on the -// capturing stream through a CallerStreamGuard, they are: the device query, the +// This one takes the capture as far as it goes. A first call over non-empty +// device-resident tensors, submitted on the capturing stream through a +// CallerStreamGuard, makes no call a capture cannot take: the device query, the // device switch, the shape and address binding, enqueueV3 and the completion -// record. Nothing here allocates, waits or synchronizes. -TEST_F(SharedScratchBackendTest, AnUnpooledEngineCapturesOnAFirstCallThatStagesNothing) { +// record, none of which allocates, waits or synchronizes. The graph that comes +// out of it instantiates and replays the engine. +// +// The completion record is the problem. execute() records it so that the next +// call and the destructor can wait for an enqueue that outlived the return, and +// under capture that record becomes a node of the graph rather than an event the +// host can wait on. The handle is then unusable: the wait at the head of the next +// execute() fails, and so does the one in ~EngineHandle. The assertions on that +// below describe what the delegate does, not what it should do; they are the +// measurement the capture documentation rests on. +TEST_F(SharedScratchBackendTest, AnUnpooledEngineCaptureReplaysAndLeavesTheHandleUnusable) { cudaStream_t stream = nullptr; ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); // A node of the caller's own, so a null graph below means the capture was @@ -1389,6 +1621,18 @@ TEST_F(SharedScratchBackendTest, AnUnpooledEngineCapturesOnAFirstCallThatStagesN void* captured_target = nullptr; ASSERT_EQ(cudaMalloc(&captured_target, 16), cudaSuccess); + // The same engine on the same input, run outside any capture, so the replay + // below has something to be compared against. A separate handle, because a run + // on the capturing handle would leave an enqueue in flight and the capture + // would then fail on the wait for it, which is the other case. + LoadedEngine reference; + ASSERT_EQ(reference.load(blob(), 22), Error::Ok); + ASSERT_EQ(reference.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::vector expected = reference.read_output(); + ASSERT_EQ(expected.size(), kElems); + EXPECT_NE(expected[0], kSentinel) << "the reference output is the sentinel, so a graph that ran nothing would pass"; + LoadedEngine unpooled; ASSERT_EQ(unpooled.load(blob(), 22), Error::Ok); ASSERT_FALSE(unpooled.handle()->shared_scratch) @@ -1402,29 +1646,52 @@ TEST_F(SharedScratchBackendTest, AnUnpooledEngineCapturesOnAFirstCallThatStagesN const Error captured = unpooled.run(stream); cudaGraph_t graph = nullptr; const cudaError_t end_err = cudaStreamEndCapture(stream, &graph); - const bool have_graph = graph != nullptr; - if (have_graph) { - cudaGraphDestroy(graph); - } EXPECT_EQ(captured, Error::Ok) << "a call with the option off was refused or failed under capture"; EXPECT_EQ(end_err, cudaSuccess) << "a call with the option off invalidated the capture: " << cudaGetErrorString(end_err); - EXPECT_TRUE(have_graph) << "the capture ended with no graph, so the call invalidated it"; + ASSERT_NE(graph, nullptr) << "the capture ended with no graph, so the call invalidated it"; + + // The launch is not optional here. A graph destroyed unexecuted says only that + // the capture was not invalidated; it shows neither what the graph computes nor + // what capturing it cost the handle, which is the whole of what this case is for. + cudaGraphExec_t graph_exec = nullptr; + ASSERT_EQ(cudaGraphInstantiate(&graph_exec, graph, 0), cudaSuccess); + ASSERT_TRUE(unpooled.fill_output(kSentinel)); + ASSERT_EQ(cudaGraphLaunch(graph_exec, stream), cudaSuccess); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::vector replayed = unpooled.read_output(); + ASSERT_EQ(replayed.size(), kElems); + EXPECT_EQ(std::memcmp(expected.data(), replayed.data(), kBytes), 0) + << "replaying the captured graph did not reproduce what the same engine produces outside a capture"; + + // And what the capture cost. The completion event was recorded into the graph, + // so the host wait the next call makes on it fails, and the call with it. + const Error after_replay = unpooled.run(stream); + EXPECT_EQ(after_replay, Error::InvalidProgram) + << "a handle whose completion event was recorded inside a captured graph ran again; if that is now supported, " + "the capture section of the README and the refusal message are stale"; + // Not sticky, but left pending it would surface under the name of whatever the + // next case calls first. + (void)cudaGetLastError(); + EXPECT_EQ(cudaGraphExecDestroy(graph_exec), cudaSuccess); + EXPECT_EQ(cudaGraphDestroy(graph), cudaSuccess); EXPECT_EQ(cudaFree(captured_target), cudaSuccess); ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); } -// The other half, and the reason the advice cannot be left at "load with the -// option off to capture it". The call that captured cleanly returns with its -// enqueue still in flight, and the next call on that handle waits for it on the -// host before it may touch the context. That wait is prohibited under the -// default capture mode, and nothing refuses it: the caller gets an error from -// execute() and a capture that ends with no graph. +// The other half. The case above ends with a handle the completion event has made +// unusable; this one is the failure a caller meets first if the handle already ran +// once. A call that returned with its enqueue in flight leaves the next call on +// that handle waiting for it on the host before it may touch the context, and that +// wait is prohibited under the default capture mode. Nothing refuses it: the +// caller gets an error from execute() and a capture that ends with no graph. // -// So under this mode the option being off buys one capture per handle, not -// capture in general. +// Between the two, a handle with the option off has no state in which capturing it +// is useful: a capture that reaches the completion record succeeds and poisons the +// handle, and one made behind a call that left an enqueue in flight is destroyed +// before it gets that far. TEST_F(SharedScratchBackendTest, AnUnpooledEngineInvalidatesACaptureOnceAnEnqueueIsInFlight) { cudaStream_t stream = nullptr; ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); @@ -1454,10 +1721,13 @@ TEST_F(SharedScratchBackendTest, AnUnpooledEngineInvalidatesACaptureOnceAnEnqueu (void)cudaGetLastError(); EXPECT_EQ(captured, Error::InvalidProgram) - << "a second call on this handle did not fail on the host wait for the previous enqueue, so the documented " - "one-captured-call-per-handle limit with the option off is wrong"; - EXPECT_NE(end_err, cudaSuccess) << "the capture survived, so the documented caveat about the second call is stale"; - EXPECT_FALSE(have_graph) << "the capture ended with a graph, so the documented caveat about the second call is stale"; + << "a call following one that left an enqueue in flight did not fail on the host wait for it, so the capture " + "documentation is stale in saying that wait is what a caller meets on a handle that has already run"; + EXPECT_NE(end_err, cudaSuccess) + << "that wait did not invalidate the capture, so the capture documentation is stale in listing it among the " + "calls a capture cannot take"; + EXPECT_FALSE(have_graph) << "the capture ended with a graph despite the host wait, so the capture documentation is " + "stale in listing that wait among the calls a capture cannot take"; EXPECT_EQ(cudaFree(captured_target), cudaSuccess); ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); @@ -1698,6 +1968,92 @@ TEST_F(SharedScratchBackendTest, AGrowthFreesTheBufferItReplacesWithTheDeviceLoc "this device waits out a device-wide free it has nothing to do with"; } +// The case above stalls the free. This one stalls the wait that comes before it: +// a growth waits for the enqueue against the buffer it retires, and that wait is +// for one inference to finish, not for the device to drain. Held under the device +// lock it would move where two pooled calls serialize -- from submission, which +// the execute() contract promises, to the completion of the call before -- so +// another pooled engine on the device would sit through an inference it has +// nothing to do with -- measured at 2119 ms with the wait moved back under the +// lock. +// +// The smaller engine's own enqueue is parked behind a host function, so the event +// it records stays unsignalled and the growth's wait for it cannot return. As in +// the case above, the capacity read under the same try_lock is what stops this +// passing on a lock that is free because the growth has not started yet. +TEST_F(SharedScratchBackendTest, AGrowthWaitsForTheRetiredBuffersEnqueueWithTheDeviceLockDropped) { + 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 second to be sure of growing the pool"; + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + + cudaStream_t held = nullptr; + cudaStream_t growing = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&held, cudaStreamNonBlocking), cudaSuccess); + ASSERT_EQ(cudaStreamCreateWithFlags(&growing, cudaStreamNonBlocking), cudaSuccess); + + LoadedEngine small; + LoadedEngine big; + ASSERT_EQ(small.load(blob(), 29), Error::Ok); + ASSERT_EQ(big.load(big_blob(), 30, kBigRows, kBigCols), Error::Ok); + ASSERT_TRUE(small.handle()->shared_scratch); + ASSERT_TRUE(big.handle()->shared_scratch); + const int device_id = big.handle()->device_id; + + // Parked ahead of the smaller engine's enqueue, so that enqueue and the event + // recorded after it both stay pending for as long as this test wants them to. + StreamGate gate; + ASSERT_EQ(cudaLaunchHostFunc(held, hold_stream, &gate), cudaSuccess); + GateRelease gate_release(gate, held); + + ASSERT_EQ(small.run(held), Error::Ok); + const std::size_t before = shared_scratch_capacity_for_testing(device_id); + ASSERT_GT(before, 0u) << "the smaller engine left the pool empty, so the larger one has nothing to retire"; + + std::atomic growth_returned{false}; + Error growth_error = Error::Internal; + std::thread grower([&] { + growth_error = big.run(growing); + growth_returned.store(true); + }); + + SharedScratchDevice& dev = scratch_pool().get(device_id); + bool lock_free_during_the_wait = false; + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (!growth_returned.load() && std::chrono::steady_clock::now() < deadline) { + if (dev.mu.try_lock()) { + const std::size_t capacity_now = dev.capacity; + dev.mu.unlock(); + if (capacity_now > before && !growth_returned.load()) { + lock_free_during_the_wait = true; + break; + } + } + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + const bool still_inside_execute = !growth_returned.load(); + + gate_release.release(); + grower.join(); + ASSERT_EQ(cudaStreamSynchronize(held), cudaSuccess); + ASSERT_EQ(cudaStreamSynchronize(growing), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(held), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(growing), cudaSuccess); + + 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(growth_error, Error::Ok); + ASSERT_TRUE(still_inside_execute) + << "the growing call returned before the gate opened, so nothing about it stalled and there was no window in " + "which to observe the lock"; + ASSERT_GT(shared_scratch_capacity_for_testing(device_id), before) + << "the pool did not grow, so no buffer was retired and there was nothing to wait for"; + EXPECT_TRUE(lock_free_during_the_wait) + << "the device's pool lock stayed held while a growth waited for the enqueue on the buffer it retired, so every " + "other pooled engine on this device waits out that inference"; +} + // --------------------------------------------------------------------------- // The pooled path under two concurrent callers // --------------------------------------------------------------------------- diff --git a/tests/cpp/executorch/test_shared_scratch_pool.cpp b/tests/cpp/executorch/test_shared_scratch_pool.cpp index d5c0ab2f0e9..ec8ea58c817 100644 --- a/tests/cpp/executorch/test_shared_scratch_pool.cpp +++ b/tests/cpp/executorch/test_shared_scratch_pool.cpp @@ -1,3 +1,10 @@ +/* + * 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. + */ + // 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. // @@ -123,15 +130,17 @@ TEST(SharedScratchPool, ReusesWhenExistingBufferIsLargeEnough) { 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`. + // Each reports into an output variable of its own, so what the reuse path writes + // is asserted rather than what the first call left in `out`, which already holds + // the value both of them are expected to write. std::size_t out2 = 0; + std::size_t out3 = 0; void* second = call(dev, a, 1000, out2); - void* third = call(dev, a, 4096, out); + void* third = call(dev, a, 4096, out3); EXPECT_EQ(second, first); EXPECT_EQ(third, first); - EXPECT_EQ(out, 4096u); + EXPECT_EQ(out3, 4096u); // Reuse reports the buffer's capacity, not the smaller amount asked for. EXPECT_EQ(out2, 4096u); EXPECT_EQ(a.alloc_count(), 1); @@ -242,6 +251,39 @@ TEST(SharedScratchPool, FirstAllocationFailureReturnsNullAndStoresNothing) { EXPECT_EQ(a.alloc_count(), 1); } +// The backend passes one RetiredScratch through a call and acts on what it holds, +// so a path that leaves the previous growth's buffer in it hands that buffer to a +// second caller to free. +TEST(SharedScratchPool, EveryPathClearsTheRetirementItReports) { + SharedScratchDevice dev; + FakeAllocator a; + std::size_t out = 0; + // Deliberately reused across the calls below, which is the state the clearing is + // for; `call` above gives each of its calls a fresh one. + RetiredScratch retired; + const auto request = [&](std::size_t need) { + return shared_scratch_get_or_grow( + dev, need, out, [&a](std::size_t bytes) { return a.alloc(bytes); }, retired); + }; + + void* const first = request(1024); + ASSERT_NE(first, nullptr); + ASSERT_EQ(retired.buffer, nullptr); + ASSERT_NE(request(8192), nullptr); + ASSERT_EQ(retired.buffer, first) << "the growth did not report the buffer it replaced"; + + EXPECT_NE(request(512), nullptr); + EXPECT_EQ(retired.buffer, nullptr) << "a reuse left the previous growth's buffer in the result, so a caller acting " + "on it frees a buffer that was already handed back once"; + + void* const second = dev.buffer; + ASSERT_NE(request(16384), nullptr); + ASSERT_EQ(retired.buffer, second); + a.fail_next = true; + EXPECT_EQ(request(65536), nullptr); + EXPECT_EQ(retired.buffer, nullptr) << "a failed allocation left the previous growth's buffer in the result"; +} + // --------------------------------------------------------------------------- // Ordering the shared buffer's handoff from one enqueue to the next. // --------------------------------------------------------------------------- @@ -451,6 +493,33 @@ TEST(SharedScratchPoolRegistry, ResetHandsBackEverySlotAndLeavesItEmpty) { EXPECT_TRUE(zero.retirements.empty()) << "the cleared slot retired a buffer the reset had already handed back"; } +// The mirror of device 1 in the case above: an event and no buffer. That state is +// reachable, because the handoff event is created before the allocation, so a +// claim whose allocation fails leaves one behind. A reset that skipped slots with +// no buffer would leak that event and pass every other case in this file. +TEST(SharedScratchPoolRegistry, ResetHandsBackTheEventOfASlotThatNeverGotABuffer) { + SharedScratchPool pool; + FakeAllocator a; + FakeEventFactory events; + std::size_t out = 0; + + SharedScratchDevice& dev = pool.get(3); + const cudaEvent_t event = shared_scratch_claim_event(dev, std::ref(events)).event; + ASSERT_NE(event, nullptr); + a.fail_next = true; + ASSERT_EQ(call(dev, a, 1024, out), nullptr); + ASSERT_EQ(dev.buffer, nullptr) << "the allocation did not fail, so this slot is not the state under test"; + + std::vector> disposed; + pool.reset_for_testing([&](int, void* buffer, cudaEvent_t slot_event) { disposed.emplace_back(buffer, slot_event); }); + + ASSERT_EQ(disposed.size(), 1u) << "the reset skipped a slot holding an event and no buffer, so nothing ever destroys " + "that event"; + EXPECT_EQ(disposed[0].first, nullptr); + EXPECT_EQ(disposed[0].second, event); + EXPECT_EQ(dev.marker.event, nullptr) << "the slot kept the event the reset handed to the disposer"; +} + // The registry hands out a reference and drops its lock, so callers go on using // that reference; the reset has to empty a slot without moving it. Checking one // address cannot pin that. An erasing reset returns the node to the allocator, @@ -506,12 +575,23 @@ TEST(SharedScratchPoolRegistry, ResetDisposesWithNoLockHeld) { ASSERT_NE(call(dev0, alloc, 1024, out), nullptr); std::atomic disposing{false}; + std::atomic reset_returned{false}; + std::atomic nothing_disposed{false}; std::atomic claimed{false}; std::atomic device_lock_free{false}; std::thread claimer([&] { - while (!disposing.load()) { + // The regression this case exists to catch is a reset that disposes of + // nothing, and this wait is where that regression arrives. Waiting for + // `disposing` alone would turn it into a target that hangs until Bazel kills + // it, with no assertion to say what broke. + const auto rendezvous_deadline = std::chrono::steady_clock::now() + std::chrono::seconds(10); + while (!disposing.load() && !reset_returned.load() && std::chrono::steady_clock::now() < rendezvous_deadline) { std::this_thread::yield(); } + if (!disposing.load()) { + nothing_disposed.store(true); + return; + } if (dev0.mu.try_lock()) { device_lock_free.store(true); dev0.mu.unlock(); @@ -531,8 +611,11 @@ TEST(SharedScratchPoolRegistry, ResetDisposesWithNoLockHeld) { } claimed_during_dispose = claimed.load(); }); + reset_returned.store(true); claimer.join(); + ASSERT_FALSE(nothing_disposed.load()) << "the reset returned without disposing of the slot that holds a buffer, so " + "there was no window in which to observe either lock"; EXPECT_TRUE(device_lock_free.load()) << "the disposer ran holding the slot's own device lock, so a claim on that " "device waits for a free that waits on the whole device"; EXPECT_TRUE(claimed_during_dispose) << "a lookup for a device the reset never touched could not complete while the " From 45ea07cc91b385d08f1d9444992956af2ba49650 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Mon, 7 Sep 2026 22:02:06 -0700 Subject: [PATCH 10/13] fix(executorch): drain a failing call's host-input copy, queue a growth's free on its stream, unship the install helper Every figure below is from local runs on one A100 host with CUDA 13.0 and driver 13.0; the CI job that would run these tests has not run them at this head. **A pooled call that fails after staging a host-resident input no longer returns with that copy still running.** The staging copy is a `cudaMemcpyAsync` reading the caller's own memory, and the caller owns that memory again the moment `execute()` returns. Measured, from pinned host memory: the copy returns in 0.0 ms without having read anything, and 100% of the bytes the device receives are what the caller wrote after the return -- while the same copy from pageable memory blocks for 1005.8 ms and takes the bytes it was called with, so the hazard is specific to pinned sources. The success path already synchronizes whenever anything was staged; a scope guard declared where the staging flag is now covers every error return between the first copy and that point -- the two this branch added in the pooled scratch block, and, all predating it, the copy's own failure, `setTensorAddress`, `inferShapes`, the whole output-binding loop, `enqueueV3` itself and the two returns after the enqueue -- so no instance of this class is left anywhere in `execute()`. It adds no synchronize to a path that did not already make one, because its condition is the staging flag and the success path's own test includes that flag. `AFailedPooledCallDrainsTheHostInputCopyItQueued` drives the one reachable failure -- the pool's allocation, with the device full -- behind a parked stream and checks what the device ends up holding; without the drain it holds the sentinel the caller wrote after the return. **A growth's disposal of the buffer it replaces is now stream-ordered, so the growing call does not stall.** The branch documented `cudaFree`'s device-wide wait as unavoidable because `cudaFreeAsync` "pairs only with cudaMallocAsync". That is false as measured: on a 256 MiB `cudaMalloc` buffer with a host function parked on an unrelated stream, `cudaFreeAsync` returned `cudaSuccess` in 0.2 ms without waiting for it, deferred the free until the stream reached it, returned the bytes to the device, and `compute-sanitizer --tool memcheck` reported no errors; the same `cudaFree` did not return until the parked work did, at 902.6 ms. `SharedScratchClaim::release()` now queues the free on the stream the claim enqueued on. Ordering rather than a wait is what makes that safe: every claimant makes its stream wait on the marker event before it enqueues and records on it afterwards, so a free queued on this stream is behind every enqueue that ever used the retired buffer. Where a device has no stream-ordered allocator, `cudaFreeAsync` reports `cudaErrorNotSupported` and the old host wait and `cudaFree` stay as the fallback, which the installed contract and the README now describe as the platform-dependent case rather than as what always happens. The two growth cases are rewritten around the new behaviour; forcing the fallback fails exactly those two, each after 60 s of parked work, and leaves the other 22 backend cases passing. **The capture documentation overstated what a capture costs the handle.** It said every later call on the handle fails. Measured through the delegate: the next call fails `Error::InvalidProgram` on the wait for the completion event, that call clears the in-flight flag before returning, and the call after it records the event again outside the capture and produces the bytes the same engine produces uncaptured. The installed header, the README and the refusal message now say the next call fails, and say what is not recovered -- a replay of the captured graph enqueues on the same context and no `execute()` waits for it. The case is renamed `AnUnpooledEngineCaptureReplaysAndCostsTheHandleItsNextCall` and asserts the recovery. **Two detected failures no longer become hangs.** The pool's test-only reset took every device lock unconditionally, so the leaked-lock defect that `AFailedPooledAllocationLeavesTheDeviceLockFree` exists to catch would hang the fixture that resets before and after every case -- a target timeout, with the message naming the leak never printed. The reset now waits a bounded 5 s per run and reports the slots it could not take; the fixture fails on that report. And in the fork-based static-destruction case, a child that throws on its way to the state under test unwound back into gtest and ran on through the rest of the binary -- one run printed 25 teardown banners, one per child -- while the parent's `waitpid` had no deadline, so a wedged child parked it until the target timed out. The child body is now wrapped in a catch-all that `_exit`s with a distinct status, and the parent reaps with a 30 s deadline, killing and reporting an overrun. Mutating each failure back in fails the case that hunts it, with one teardown banner rather than 25. **`install_pooled_scratch` is out of the installed header set.** It is described in its own comment as an implementation detail of `execute()`, and the header it was declared in ships in the release tarball, so the next release would freeze it as public API. The declaration moves to `PooledScratchInstall.h` beside the sources, which ships with them and not in the include tree, following what this branch already did for the pool header. Built `//:include_executorch`, the header half of the release tarball, and read the tar: none of its five headers mentions it. The symbol itself is still a global `T` in the backend object, because the backend's own test links against it to cover a refusal that cannot be induced through `execute()`; nothing short of dropping that test removes it, and the declaration is what decides whether the next release has to keep the entry point. **The installed `execute()` contract now names the pooled path's error returns.** Beyond the `Error::NotSupported` it already documented for a capturing stream, a pooled call can fail four ways an unpooled one cannot, and the contract described none of them: `Error::Internal` where the handoff event cannot be created, `Error::MemoryAllocationFailed` where the buffer cannot be allocated or grown, `Error::InvalidState` from the wait that orders this call behind the previous enqueue, from an install TensorRT refuses, or from the record of this enqueue, and `Error::InvalidProgram` from the fallback disposal's host wait, which leaks the retired buffer rather than freeing it under a live enqueue. The contract now lists them, says each logs at `Error` first, and names the only two reached after the enqueue is submitted -- both of which synchronize the stream before returning, so neither leaves engine work in flight. Also: the zero-to-minimum substitution is made once at the call site instead of twice, so the size the pool guarantees and the size the context is told it owns cannot drift apart -- `setDeviceMemoryV2` refuses an undersized install and says nothing about an oversized one, so nothing downstream would catch it; `shared_scratch_get_or_grow` drops the size out-parameter no production caller read, and the tests read `dev.capacity`, which they already did elsewhere; `get_or_grow_shared_scratch` and `mark_shared_scratch_in_flight` become `claim_shared_scratch` and `record_shared_scratch_enqueue`, so the wrappers are no longer their wrappees' words in a different order; `scratch_from_pool` is gone, being `pooled_scratch` on every path that reaches its read; the growth case's direction check joins the interference skip below it rather than failing on the unrelated device activity that skip exists to absorb; the two-thread rendezvous gets a 30 s deadline, so a pooled call that blocks reports here instead of spinning until the target times out; and the file's coverage note no longer claims that `TORCHTRT_EXECUTORCH_REQUIRE_CUDA` turns every skip into a failure -- it covers the missing device, and three cases skip for reasons it does not reach. Three descriptions the stream-ordered free left behind, each of which had the old synchronous disposal as the normal path: the capture refusal's own doc comment, which still had a growth's `cudaEventSynchronize` and `cudaFree` invalidating a capture outside `Relaxed` and permitted under it, where the header and the README at this head both say the `cudaFreeAsync` is refused under all three modes and drops the growth onto the fallback; the handoff event's blocking-sync rationale, whose "the one host wait on this event" is now a wait only the fallback path makes; and the growth test's bounds comment, which explained which of a host wait and a device-wide free the bounds cover when neither is on the path any more. 6 targets, 85 -> 87 cases, all passing on a real GPU with `--nocache_test_results` and `TORCHTRT_EXECUTORCH_REQUIRE_CUDA=1`, no skips. Every behavioural change was proven by mutating it away. --- cpp/BUILD | 28 + .../executorch/TensorRTBackend.h | 85 ++-- .../torch_tensorrt/executorch/CMakeLists.txt | 5 +- .../executorch/PooledScratchInstall.h | 42 ++ cpp/src/torch_tensorrt/executorch/README.md | 108 ++-- .../executorch/SharedScratchPool.h | 67 ++- .../executorch/SharedScratchPoolTestHooks.cpp | 5 +- .../executorch/SharedScratchPoolTestHooks.h | 7 +- .../executorch/TensorRTBackend.cpp | 235 +++++---- tests/cpp/executorch/BUILD | 1 + .../test_shared_scratch_backend.cpp | 479 ++++++++++++------ .../executorch/test_shared_scratch_pool.cpp | 213 +++++--- 12 files changed, 851 insertions(+), 424 deletions(-) create mode 100644 cpp/src/torch_tensorrt/executorch/PooledScratchInstall.h diff --git a/cpp/BUILD b/cpp/BUILD index a8f0eca220e..d2db81912ee 100644 --- a/cpp/BUILD +++ b/cpp/BUILD @@ -216,6 +216,32 @@ cc_library( }), ) +# The declaration of the one internal entry point of :tensorrt_executorch_backend +# that its test calls directly. Under src/ for the same reason as the pool header +# above: it is an implementation detail of execute(), and an installed header +# declaring it would make it API from the next release. +cc_library( + name = "tensorrt_executorch_pooled_scratch_install", + hdrs = [ + "src/torch_tensorrt/executorch/PooledScratchInstall.h", + ], + strip_include_prefix = "src", + target_compatible_with = select({ + ":linux_x86_64": [], + ":sbsa": [], + "//conditions:default": ["@platforms//:incompatible"], + }), + deps = select({ + ":linux_x86_64": [ + "@tensorrt//:nvinfer", + ], + ":sbsa": [ + "@tensorrt_sbsa//:nvinfer", + ], + "//conditions:default": [], + }), +) + # The pool's test-only entry points. A separate target, and testonly, so the # released archive -- Bazel's :tensorrt_executorch_backend and CMake's # executorch_trt_backend, neither of which compiles this source -- carries no @@ -270,6 +296,7 @@ cc_library( deps = [ ":tensorrt_executorch_binding_names", ":tensorrt_executorch_blob_header", + ":tensorrt_executorch_pooled_scratch_install", ":tensorrt_executorch_shared_scratch_pool", ":tensorrt_executorch_weight_streaming_budget", ] + select({ @@ -294,6 +321,7 @@ filegroup( name = "executorch_backend_source_files", srcs = [ "src/torch_tensorrt/executorch/CMakeLists.txt", + "src/torch_tensorrt/executorch/PooledScratchInstall.h", "src/torch_tensorrt/executorch/README.md", "src/torch_tensorrt/executorch/SharedScratchPool.h", "src/torch_tensorrt/executorch/TensorRTBackend.cpp", diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index cfc110f7978..7489f08a93d 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -102,27 +102,6 @@ struct EngineHandle { ~EngineHandle(); }; -// Installs `bytes` of `buffer` as the activation scratch of `ctx`, a -// kUSER_MANAGED context, and reports whether TensorRT accepted it. Logs the -// refusal, naming `device_id`, when it did not. -// -// The check is the point. setDeviceMemoryV2 returns void and refuses a buffer -// smaller than the bound shapes need, so a caller that does not ask cannot tell an -// accepted install from a refused one -- and a refused one leaves the context -// pointed at the buffer it was last given, which a shared-scratch pool growth may -// since have freed. The engine then reads and writes freed memory with enqueueV3 -// reporting success. The refusal is read back through an IErrorRecorder scoped to -// this one call, the only channel that hands it to the caller: with no recorder -// attached TensorRT writes it to the runtime's ILogger and this function has -// nothing to return. -// -// An implementation detail of execute(), declared here rather than beside the -// pool's bookkeeping because this is the header of the target that defines it and -// the only one that may name TensorRT. The backend's own test drives it directly, -// over the same TensorRT call, to cover a refusal that cannot be induced through -// execute(). -bool install_pooled_scratch(nvinfer1::IExecutionContext& ctx, void* buffer, std::size_t bytes, int device_id); - // 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. Read by TensorRTBackend::set_option below, and delivered as @@ -154,28 +133,34 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { // handle whose context was created while the option was off keeps its own // scratch and is not subject to this, nor is one whose engine reports needing no // activation scratch under any shape, which is left out of the pool. For the - // handles that do draw on it, three further consequences: + // handles that do draw on it, four further consequences: // - A call needing more scratch than the pool holds grows it, and the growth - // frees the buffer it replaces. cudaFree waits for every stream on the - // device, so that one call blocks until the device is idle however - // asynchronous the rest of this contract makes it -- an unbounded wait on - // work this call did not submit. Unbounded is meant literally: if any of - // that work is itself waiting on something only this thread supplies once - // execute() returns -- a host function it will release, a copy it will - // enqueue next -- the call does not return, and the thread that would - // unblock it is the one inside cudaFree. Which calls grow the pool is not - // knowable from here; see the README. + // gets rid of the buffer it replaces with a stream-ordered cudaFreeAsync on + // the stream it enqueued on, so the call does not wait for that free and the + // bytes come back when the stream reaches it. Where the device has no + // stream-ordered allocator that call reports cudaErrorNotSupported and the + // growth falls back to a host wait and cudaFree, which waits for every + // stream on the device: on such a device that one call blocks until the + // device is idle however asynchronous the rest of this contract makes it -- + // an unbounded wait on work this call did not submit. Unbounded is meant + // literally: if any of that work is itself waiting on something only this + // thread supplies once execute() returns -- a host function it will release, + // a copy it will enqueue next -- the call does not return, and the thread + // that would unblock it is the one inside cudaFree. Which calls grow the + // pool is not knowable from here; see the README. // - Capturing a CUDA graph around this delegate is not supported, with the // option on or off. With it on, a call whose selected stream is capturing is // refused with Error::NotSupported, ahead of every CUDA call it makes that a // capture cannot take -- only the device query and the device switch run // first. The pool's event handoff waits on an event recorded outside the // capture, which invalidates it under every capture mode. A growth's - // allocation, its wait on the replaced buffer and the free of that buffer - // invalidate it under every mode but cudaStreamCaptureModeRelaxed, which - // permits all three but does not record them, leaving a replay pointed at a - // buffer the pool may since have freed. The alternative to refusing is a - // capture that silently comes back invalidated. Only a capture on the + // allocation invalidates it under every mode but cudaStreamCaptureModeRelaxed, + // which permits it but does not record it, leaving a replay pointed at a + // buffer the pool may since have freed; and a growth's stream-ordered free is + // refused under every mode -- a free made while capturing has to be given a + // graph allocation and this is not one -- which drops the growth onto the + // fallback wait and cudaFree that a capture cannot take either. The + // alternative to refusing is a capture that silently comes back invalidated. Only a capture on the // selected stream is caught. A capture running on any other stream under // cudaStreamCaptureModeGlobal, or under cudaStreamCaptureModeThreadLocal // from this thread, is invalidated by the same calls and is not refused, @@ -188,15 +173,33 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { // an enqueue in flight makes no call a capture cannot take, and the graph // replays the engine. It costs the handle. Such a call ends by recording the // completion event this contract describes, and under capture that record is - // a node of the graph rather than an event the host can wait on, so every - // later execute() on the handle fails Error::InvalidProgram on the wait for - // it. A call that ends by synchronizing the stream instead is one no capture - // mode permits, so there is no call shape that captures and leaves the - // handle usable. The README states this in full. + // a node of the graph rather than an event the host can wait on, so the next + // execute() on the handle fails Error::InvalidProgram on the wait for it, + // and so does the destructor's wait. That failing call clears the flag + // before it returns, so the call after it records the event again outside + // the capture and runs -- but nothing waits for a replay of the graph, which + // enqueues on the same context, and reconfiguring a context under a live + // enqueue is what the event was there to prevent. A call that ends by + // synchronizing the stream instead is one no capture mode permits, so there + // is no call shape that captures and leaves the handle safe to go on using. + // The README states this in full. // - cudaDeviceReset() invalidates the pool without emptying it. The buffer // and the handoff event it still holds are destroyed with the primary // context, and the next call on that device uses both. There is no guard: // do not reset a device this backend has run a pooled engine on. + // - The pool gives a pooled call four error returns an unpooled one never + // makes, on top of Error::NotSupported for the capture above. Claiming the + // device's scratch answers Error::Internal if its handoff event cannot be + // created, Error::MemoryAllocationFailed if the buffer cannot be allocated + // or grown, and Error::InvalidState if the wait ordering this call behind + // the previous enqueue fails; installing the buffer answers + // Error::InvalidState if TensorRT refuses it, and recording this call's + // enqueue for the next claimant answers the same. Error::InvalidProgram + // means the fallback disposal's host wait failed, which leaks the retired + // buffer rather than freeing it under an enqueue that may still be reading + // it. Every one of them logs at Error first. The last two are the only ones + // reached after this call's enqueue is submitted, and both synchronize the + // stream before returning, so neither leaves engine work in flight. // 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/CMakeLists.txt b/cpp/src/torch_tensorrt/executorch/CMakeLists.txt index 6a291685259..ec75e2e18cf 100644 --- a/cpp/src/torch_tensorrt/executorch/CMakeLists.txt +++ b/cpp/src/torch_tensorrt/executorch/CMakeLists.txt @@ -37,8 +37,9 @@ target_include_directories(executorch_trt_backend PUBLIC "${CMAKE_CURRENT_LIST_DIR}/../../../include" PRIVATE - # SharedScratchPool.h is an implementation detail, so it sits beside the - # sources instead of in the installed include tree. This directory is + # SharedScratchPool.h and PooledScratchInstall.h are implementation + # details, so they sit beside the sources instead of in the installed + # include tree. This directory is # .../torch_tensorrt/executorch in both the repository and the released # package, so two levels up is the root its #include path is relative to. "${CMAKE_CURRENT_LIST_DIR}/../.." diff --git a/cpp/src/torch_tensorrt/executorch/PooledScratchInstall.h b/cpp/src/torch_tensorrt/executorch/PooledScratchInstall.h new file mode 100644 index 00000000000..83708012187 --- /dev/null +++ b/cpp/src/torch_tensorrt/executorch/PooledScratchInstall.h @@ -0,0 +1,42 @@ +/* + * 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 + +// The one step of execute()'s pooled path that the backend's test drives +// directly, declared beside the source that defines it rather than in the +// installed header set: it is an implementation detail of execute(), and a +// declaration in an installed header is API the next release has to keep. +// SharedScratchPool.h sits here for the same reason. + +#include + +#include + +namespace torch_tensorrt { +namespace executorch_backend { + +// Installs `bytes` of `buffer` as the activation scratch of `ctx`, a +// kUSER_MANAGED context, and reports whether TensorRT accepted it. Logs the +// refusal, naming `device_id`, when it did not. +// +// The check is the point. setDeviceMemoryV2 returns void and refuses a buffer +// smaller than the bound shapes need, so a caller that does not ask cannot tell an +// accepted install from a refused one -- and a refused one leaves the context +// pointed at the buffer it was last given, which a shared-scratch pool growth may +// since have freed. The engine then reads and writes freed memory with enqueueV3 +// reporting success. The refusal is read back through an IErrorRecorder scoped to +// this one call, the only channel that hands it to the caller: with no recorder +// attached TensorRT writes it to the runtime's ILogger and this function has +// nothing to return. +// +// The backend's own test calls it over the same TensorRT call to cover a refusal +// that cannot be induced through execute(). +bool install_pooled_scratch(nvinfer1::IExecutionContext& ctx, void* buffer, std::size_t bytes, int device_id); + +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/cpp/src/torch_tensorrt/executorch/README.md b/cpp/src/torch_tensorrt/executorch/README.md index e79ff1f6b78..86577b6a045 100644 --- a/cpp/src/torch_tensorrt/executorch/README.md +++ b/cpp/src/torch_tensorrt/executorch/README.md @@ -196,36 +196,48 @@ shape, and that one is settled before a call is ever made: such an engine is lef out of the pool entirely, so it takes no per-device lock and does not serialize against the engines that do. -The buffer grows when a call asks for more than every call 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. It is the -only call on the growth path that does: measured with a host function parked on a -stream neither the pool nor the growing engine had ever used, the `cudaMalloc` for -the new buffer and the wait on the retired one each returned at once and the -`cudaFree` returned after 3000 ms, exactly when that host function was released. -The backend keeps both that stall and the event wait before it out from under the -per-device lock, so neither holds up another pooled engine on the device: two -pooled calls still serialize at submission, as the caller-stream contract above -says, and a growth is not an exception to it. Both do fall after the growing -call's own enqueue, so that one `execute()` waits for its own engine work too. - -The wait has no upper bound, and a caller can turn it into one that never ends. -Anything queued anywhere on the device is enough to hold it, so if some stream is -waiting on work only this thread will submit -- a host function it will release -after `execute()` returns, a copy it will enqueue next -- the growing call does -not return and the thread that would unblock it is the one inside `cudaFree`. -This is not hypothetical: the backend's own tests deadlocked on it once, when a -change of test order turned a case that parks a host function on its own stream -into the one that grew the pool. A program that parks work like that has to keep -the pool's growths away from it, and the paragraph below on how often the pool -grows is what says whether a run order can do that. - -A stream-ordered free would bound the wait to the buffer's own users, but -`cudaFreeAsync` pairs only with `cudaMallocAsync`, so it would move the pool onto -the stream-ordered allocator for every allocation it makes, not just the free. -Deferring the free instead only moves the wait, since nothing later on the path -pays it. The pool keeps the plain allocator and this section is the warning. +The buffer grows when a call asks for more than every call before it did, and the +growth has to get rid of the buffer it replaces. It queues that free rather than +making it: `cudaFreeAsync` on the stream the growing call enqueued on, with the +per-device lock already dropped. The call returns without waiting for it, and the +bytes come back when the stream reaches the free. Ordering is what makes this +safe, not a wait: every pooled call makes its stream wait on the device's handoff +event before it enqueues and records on that event afterwards, so a free queued on +this call's stream sits behind every enqueue that ever used the retired buffer, +and behind this call's own, which uses the new one. + +`cudaFreeAsync` takes a pointer `cudaMalloc` returned; the pool does not move to +the stream-ordered allocator to use it. Measured on CUDA 13.0 with driver 13.0, on +a 256 MiB `cudaMalloc` buffer with a host function parked on an unrelated stream: +`cudaFreeAsync` returned `cudaSuccess` in 0.2 ms without waiting for the parked +work, the bytes returned to the device once the streams drained, and +`compute-sanitizer --tool memcheck` reported no errors; `cudaFree` on the same +buffer with the same work parked did not return until that work did. + +Where a device has no stream-ordered allocator (`cudaDevAttrMemoryPoolsSupported`), +`cudaFreeAsync` reports `cudaErrorNotSupported` instead of freeing, and the +backend falls back to what it did before: a host wait on the handoff event, then +`cudaFree`. That path is the one to know about. `cudaFree` waits for everything +queued on the device, not only for the enqueues that used the buffer -- measured +with a host function parked on a stream neither the pool nor the growing engine +had ever used, the `cudaMalloc` for the new buffer and the wait on the retired one +each returned at once and the `cudaFree` returned after 3000 ms, exactly when that +host function was released. Both are still made with the per-device lock dropped, +so neither holds up another pooled engine on the device: two pooled calls +serialize at submission, as the caller-stream contract above says, and a growth is +not an exception to it. + +On that fallback the wait has no upper bound, and a caller can turn it into one +that never ends. Anything queued anywhere on the device is enough to hold it, so +if some stream is waiting on work only this thread will submit -- a host function +it will release after `execute()` returns, a copy it will enqueue next -- the +growing call does not return and the thread that would unblock it is the one +inside `cudaFree`. This is not hypothetical: the backend's own tests deadlocked on +it once, when a change of test order turned a case that parks a host function on +its own stream into the one that grew the pool. A program on such a device that +parks work like that has to keep the pool's growths away from it, and the +paragraph below on how often the pool grows is what says whether a run order can +do that. What an engine answers when asked how much it needs is decided when it is built, not when it runs. The builder's `kRUNTIME_ACTIVATION_RESIZE_10_10` preview feature @@ -255,14 +267,18 @@ the capture, and `cudaStreamWaitEvent` on such an event fails with `cudaErrorStreamCaptureIsolation` and invalidates the capture under every capture mode, `cudaStreamCaptureModeRelaxed` included; the caller learns of that only when `cudaStreamEndCapture` hands back `cudaErrorStreamCaptureInvalidated` -and a null graph, long after the call that caused it. A growth is mode-dependent -rather than always fatal. Its `cudaMalloc`, the `cudaEventSynchronize` it makes -before releasing the buffer it replaces, and the `cudaFree` that releases it each -return `cudaErrorStreamCaptureUnsupported` and invalidate the capture under the -`Global` and `ThreadLocal` modes; under `Relaxed` all three are permitted and the -capture survives them, but they then run uncaptured, so a replayed graph would use -whatever buffer was installed when it was captured -- which by then the pool may -have freed. The backend checks the selected stream and returns instead, and it +and a null graph, long after the call that caused it. A growth adds more of the +same. Its `cudaMalloc` returns `cudaErrorStreamCaptureUnsupported` and invalidates +the capture under `Global` and `ThreadLocal`, and under `Relaxed` is permitted but +runs uncaptured, so a replayed graph would use whatever buffer was installed when +it was captured -- which by then the pool may have freed. Its `cudaFreeAsync` is +refused under all three modes, with `cudaErrorInvalidValue` and without +invalidating the capture -- measured; a free made while capturing has to be given +a graph allocation, and the pool's buffers come from `cudaMalloc` -- and the +growth then falls back to a host wait on the handoff event and a `cudaFree`, and +those two return `cudaErrorStreamCaptureUnsupported` and invalidate the capture +under `Global` and `ThreadLocal`, and under `Relaxed` are permitted but run +uncaptured. The backend checks the selected stream and returns instead, and it checks ahead of everything a capture cannot take -- not merely ahead of the pool's own calls, since the wait on a previous enqueue and the `cudaMalloc` that grows a host-input staging buffer come before those and, outside `Relaxed`, would @@ -283,7 +299,8 @@ run a pooled engine while any capture is open anywhere in the process. **With the option off nothing refuses, and capture still does not work.** Turning the option off does take the pool's calls out of the way, and the refusal with them, because the check lives inside the pooled path. What it leaves is a -capture that succeeds and a handle that cannot be used again. +capture that succeeds, a handle whose next call fails, and a graph whose replays +nothing orders against the delegate's own calls. A call captures cleanly with the option off when it runs on the capturing stream under a `CallerStreamGuard`, binds only non-empty device-resident tensors, has no @@ -305,7 +322,18 @@ wait. No call shape avoids this: the alternative to recording the event is the an output or running with no caller stream, and a capturing stream cannot be synchronized under any mode, `Relaxed` included. So the only call a capture can take is the one that arms the event. -`AnUnpooledEngineCaptureReplaysAndLeavesTheHandleUnusable` in + +It is one call that fails, not every call after the capture. The failing call +clears the in-flight flag before returning the error, so the call after it skips +the wait, records the event again outside any capture, and runs -- measured, +producing what the same engine produces uncaptured. That recovery is not a reason +to capture. What the handle does not get back is the ordering the event carried +over the captured work: a `cudaGraphLaunch` of that graph enqueues on the same +execution context, no `execute()` waits for it, and TensorRT forbids +reconfiguring a context while one of its enqueues is in flight -- so a caller who +replays the graph and calls `execute()` again is reconfiguring a live context with +nothing to say when it stops being live. +`AnUnpooledEngineCaptureReplaysAndCostsTheHandleItsNextCall` in `tests/cpp/executorch/test_shared_scratch_backend.cpp` is that whole sequence. Two further calls invalidate a capture with the option off, before the handle diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h index 545f4556fe0..e895d9c4a0d 100644 --- a/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h @@ -14,8 +14,10 @@ #include +#include #include #include +#include #include #include #include @@ -73,11 +75,23 @@ class SharedScratchPool { return devices_[device_id]; } + // How long a reset waits, in total, for the device locks. A test calls this + // between cases with nothing claimed, so anything held is a leak the run has + // already failed over; the wait only has to outlast a claim still winding down. + static constexpr std::chrono::seconds kResetLockWait{5}; + // Test-only. Returns every device's slot to the state it had before anything // claimed it, handing what the slot held to `dispose(device_id, buffer, event)` // so the caller can release it. Entries stay in the map, so a reference `get` // handed out remains valid. // + // Answers with the number of devices whose lock it could not take within + // kResetLockWait, whose slots it left alone. That case is a leaked claim, which + // is a defect a case here exists to catch -- and taking the locks unconditionally + // would hang on exactly that defect, in a fixture that resets both before and + // after every case, so the run would end in a target timeout with nothing said + // about the cause. The caller reports it instead. + // // A template, so nothing is emitted for it until something instantiates it, and // in the released archive nothing does: the only entry point that calls it is // compiled into the test target alone. See SharedScratchPoolTestHooks.h. @@ -91,14 +105,27 @@ class SharedScratchPool { // Nothing here waits for an enqueue: the buffer it frees may still be in use by // one, and the caller is responsible for there being none. template - void reset_for_testing(Dispose dispose) { + std::size_t reset_for_testing(Dispose dispose) { std::vector> taken; + std::size_t still_locked = 0; { std::lock_guard lk(mu_); taken.reserve(devices_.size()); + // One deadline for the whole reset rather than one per device, so a run + // with a leaked lock costs the same however many devices the pool has seen. + const auto deadline = std::chrono::steady_clock::now() + kResetLockWait; for (auto& entry : devices_) { SharedScratchDevice& dev = entry.second; - std::lock_guard dev_lk(dev.mu); + bool locked = dev.mu.try_lock(); + while (!locked && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + locked = dev.mu.try_lock(); + } + if (!locked) { + ++still_locked; + continue; + } + std::lock_guard dev_lk(dev.mu, std::adopt_lock); taken.emplace_back(entry.first, dev.buffer, dev.marker.event); dev.buffer = nullptr; dev.capacity = 0; @@ -109,6 +136,7 @@ class SharedScratchPool { for (const auto& slot : taken) { dispose(std::get<0>(slot), std::get<1>(slot), std::get<2>(slot)); } + return still_locked; } private: @@ -191,23 +219,21 @@ inline cudaEvent_t shared_scratch_mark_in_flight(SharedScratchDevice& dev) { // 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. A null `wait_for` means nothing was ever -// recorded against it. +// be reading and writing `buffer` has been recorded; the caller must not free the +// buffer ahead of that enqueue. A null `wait_for` means nothing was ever recorded +// against it, so nothing is using 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. +// doing both is ordered against nothing here. // -// Both the wait and the free belong outside `dev.mu`, and for the same reason: -// each blocks on device work, and under the lock that makes an unrelated claim on -// this device wait for it. The free is a device-wide synchronization; the wait is -// for an inference to finish. That the wait may then also pick up enqueues made -// after the lock was dropped costs nothing, because the free that follows it waits -// for every stream on the device anyway, and it stays sufficient for this buffer -// by the ordering argument above. +// The caller has two ways to honour that, and both belong outside `dev.mu`: a +// free queued on a stream already ordered after the event, which is what the +// backend does, or a host wait on the event followed by a device-wide free, which +// is what it falls back to. Under the lock either one makes an unrelated claim on +// this device wait for work it has nothing to do with. struct RetiredScratch { void* buffer = nullptr; cudaEvent_t wait_for = nullptr; @@ -225,16 +251,16 @@ struct RetiredScratch { // that ignores `out_retired` leaks rather than frees a buffer an enqueue may // still be using. Every path clears `out_retired` first, so a caller reusing one // across calls is not handed a buffer an earlier call already disposed of. +// +// What the buffer ended up sized at is `dev.capacity`, which the caller holds the +// lock for anyway. It is not reported separately, because a caller reading it +// would be reading what an earlier call grew the pool to and not its own request: +// the only figure that belongs on a caller's execution context is the one it +// asked for here. template -void* shared_scratch_get_or_grow( - SharedScratchDevice& dev, - std::size_t need, - std::size_t& out_size, - Alloc alloc, - RetiredScratch& out_retired) { +void* shared_scratch_get_or_grow(SharedScratchDevice& dev, std::size_t need, Alloc alloc, RetiredScratch& out_retired) { out_retired = RetiredScratch{}; if (dev.buffer != nullptr && dev.capacity >= need) { - out_size = dev.capacity; return dev.buffer; } void* p = alloc(need); @@ -247,7 +273,6 @@ void* shared_scratch_get_or_grow( } dev.buffer = p; dev.capacity = need; - out_size = need; return p; } diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp index 9cf80e8e1ef..de6bf3f32c9 100644 --- a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp @@ -22,10 +22,10 @@ std::size_t shared_scratch_capacity_for_testing(int device_id) { return dev.capacity; } -void reset_shared_scratch_pool_for_testing() { +bool reset_shared_scratch_pool_for_testing() { int restore_to = 0; const bool have_current = cudaGetDevice(&restore_to) == cudaSuccess; - scratch_pool().reset_for_testing([](int device_id, void* buffer, cudaEvent_t event) { + const std::size_t still_locked = scratch_pool().reset_for_testing([](int device_id, void* buffer, cudaEvent_t event) { if (buffer == nullptr && event == nullptr) { return; } @@ -47,6 +47,7 @@ void reset_shared_scratch_pool_for_testing() { if (have_current) { (void)cudaSetDevice(restore_to); } + return still_locked == 0; } } // namespace executorch_backend diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.h b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.h index c9716d1523f..01ad3172277 100644 --- a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.h +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.h @@ -33,7 +33,12 @@ std::size_t shared_scratch_capacity_for_testing(int device_id); // Frees every device's buffer, destroys its handoff event, and clears the marker, // so one test does not inherit a pool an earlier one grew. -void reset_shared_scratch_pool_for_testing(); +// +// False when a device's lock was still held after the pool's reset deadline: that +// slot is left as it was. A caller should report it rather than carry on, and +// must not wait for the lock itself -- the state it describes is a claim that was +// never released, so nothing is going to release it. +bool reset_shared_scratch_pool_for_testing(); } // namespace executorch_backend } // namespace torch_tensorrt diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index 201cc49d7ac..a7e5f77de67 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -6,6 +6,7 @@ */ #include "torch_tensorrt/executorch/TensorRTBackend.h" +#include "torch_tensorrt/executorch/PooledScratchInstall.h" #include "torch_tensorrt/executorch/SharedScratchPool.h" #include "torch_tensorrt/executorch/TensorRTBindingNames.h" #include "torch_tensorrt/executorch/TensorRTBlobHeader.h" @@ -267,42 +268,46 @@ class SharedScratchClaim { return dev_; } - // Takes ownership of a buffer a growth displaced, to be waited for and freed by - // release(). `wait_for` is the marker event the enqueues that used it were - // recorded on, or null if none were. - void retire(void* buffer, cudaEvent_t wait_for) { + // Takes ownership of a buffer a growth displaced, to be freed by release() on + // `stream`. `wait_for` is the marker event the enqueues that used it were + // recorded on, or null if none were; `stream` is the one this claim is about to + // enqueue on, which release() needs because the free it makes is ordered on it. + void retire(void* buffer, cudaEvent_t wait_for, cudaStream_t stream) { retired_ = buffer; retired_wait_ = wait_for; + stream_ = stream; } // 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 waits for the enqueues - // that used whatever a growth displaced and frees it -- both after the unlock, - // because both block on device work that another pooled engine on this device - // has nothing to do with. The free waits for every stream on the device; - // the wait waits for an inference to finish, so under the lock it would move the - // point at which two pooled calls serialize from submission, which is what the - // execute() contract promises, to the completion of the call before. + // pointer this claim no longer holds the lock for. Then disposes of whatever a + // growth displaced, after the unlock, because either way of doing that blocks or + // queues on device work another pooled engine on this device has nothing to do + // with. // - // Waiting here rather than under the lock can wait for more than the retired - // buffer's own users: this claim records its own enqueue on the same event - // before it unlocks, and a claimant entering afterwards records on it again. It - // costs nothing. The free below waits for the whole device, so it already covers - // everything the wait could have picked up, and the wait is still sufficient -- - // every claimant orders its stream after this event before enqueueing, so the - // latest recording completes only once all the earlier ones have. + // The free is stream-ordered: cudaFreeAsync queued on the stream this claim + // enqueued on. That is what keeps a growth from stalling. Measured on CUDA 13.0 + // with driver 13.0: cudaFreeAsync takes a pointer cudaMalloc returned, does not + // block on unrelated work parked on another stream, defers the free until the + // stream reaches it, and gives the bytes back to the device rather than holding + // them in the allocator's pool -- while cudaFree on the same buffer, with the + // same work parked, does not return until that work does. // - // Unlocking bounds who waits, not how long. The free is still the one call on - // the growth path that waits device-wide -- measured, the allocation does not -- - // and a caller whose own later work some queued stream is waiting on will not - // get this call back at all. TensorRTBackend.h says so in the execute() contract, - // and the README says what to do about it. Making the free stream-ordered is not - // a local change: cudaFreeAsync pairs only with cudaMallocAsync, so it would move - // every pool allocation onto the stream-ordered allocator to bound this one wait. + // Ordering, not the wait, is what makes it safe. Every claimant makes its stream + // wait on the marker event before it enqueues, and this claim did so while + // holding the lock, so this stream is already behind every enqueue that ever used + // the retired buffer -- the marker covers all of them, because each records on it + // afterwards. Queuing the free on this stream therefore puts it after the last of + // them, and after this claim's own enqueue, which uses the new buffer. // - // Returns false when the wait failed, which leaves the buffer leaked rather than - // freed under a live enqueue; the caller reports it. Frees on the current device, - // which must still be the buffer's. + // cudaFreeAsync needs the device's stream-ordered allocator, which not every + // platform has; where it is missing the call fails rather than freeing, so the + // host wait and the device-wide cudaFree stay as the fallback. On that path a + // growth blocks until the device is idle, which is what the execute() contract + // and the README describe as the platform-dependent case. + // + // Returns false when the fallback's wait failed, which leaves the buffer leaked + // rather than freed under a live enqueue; the caller reports it. Frees on the + // current device, which must still be the buffer's. bool release() { if (lock_.owns_lock()) { lock_.unlock(); @@ -313,8 +318,24 @@ class SharedScratchClaim { } void* const retired = retired_; const cudaEvent_t wait_for = retired_wait_; + const cudaStream_t stream = stream_; retired_ = nullptr; retired_wait_ = nullptr; + stream_ = nullptr; + + const cudaError_t async_err = cudaFreeAsync(retired, stream); + if (async_err == cudaSuccess) { + return true; + } + // Not this device's allocator, then. The error is this call's own, so it is + // cleared here rather than left for the next CUDA call in execute() to report + // under its own name. + cudaGetLastError(); + ET_LOG( + Info, + "TensorRTBackend::execute: the stream-ordered free of the shared activation scratch buffer a pool growth replaced on device %d is not available here (%s); falling back to a host wait and a device-wide free, which blocks this call until the device is idle", + device_id_, + cudaGetErrorString(async_err)); if (wait_for != nullptr) { const cudaError_t wait_err = cudaEventSynchronize(wait_for); @@ -356,6 +377,7 @@ class SharedScratchClaim { std::unique_lock lock_; void* retired_ = nullptr; cudaEvent_t retired_wait_ = nullptr; + cudaStream_t stream_ = nullptr; }; // What a call needing no activation scratch is given when the pool holds nothing @@ -372,13 +394,16 @@ constexpr size_t kMinPooledScratchBytes = 1; // // The handoff's event wait is on an event recorded outside the capture: // cudaStreamWaitEvent fails it with cudaErrorStreamCaptureIsolation and -// invalidates the capture under every capture mode. A growth's allocation, the -// cudaEventSynchronize on the buffer it replaces and the cudaFree of that buffer -// invalidate it too, but only outside cudaStreamCaptureModeRelaxed, which permits -// all three -- and then runs them uncaptured, leaving a replay pointed at a -// buffer the pool may since have freed. None of it fails cleanly: the caller -// learns of an invalidation only when cudaStreamEndCapture hands back an error -// and a null graph. Refusing names the cause instead. +// invalidates the capture under every capture mode. A growth adds more. Its +// cudaMalloc invalidates the capture outside cudaStreamCaptureModeRelaxed, and +// under Relaxed is permitted but runs uncaptured, leaving a replay pointed at a +// buffer the pool may since have freed. Its cudaFreeAsync of the buffer it +// replaces is refused under all three modes, without invalidating the capture -- +// a free made while capturing has to be given a graph allocation and the pool's +// buffers come from cudaMalloc -- which drops the growth onto the fallback host +// wait and cudaFree, and those two behave like the cudaMalloc. The invalidations +// do not fail cleanly: the caller learns of one only when cudaStreamEndCapture +// hands back an error and a null graph. Refusing names the cause instead. // // execute() calls this ahead of every CUDA call it makes that a capture cannot // take, not just the pool's own: the cudaEventSynchronize on a previous enqueue @@ -411,7 +436,7 @@ Error refuse_pooled_call_on_a_capturing_stream(cudaStream_t stream, int device_i } ET_LOG( Error, - "TensorRTBackend::execute: the selected stream is capturing a CUDA graph (%s), which the shared activation scratch pool on device %d does not support. Neither does the rest of this delegate: loading the engine with '%s' off removes this refusal, but a call captured that way records the handle's completion event into the graph, and every later call on that handle then fails. The capture section of the backend README has the detail.", + "TensorRTBackend::execute: the selected stream is capturing a CUDA graph (%s), which the shared activation scratch pool on device %d does not support. Neither does the rest of this delegate: loading the engine with '%s' off removes this refusal, but a call captured that way records the handle's completion event into the graph, the next call on that handle then fails, and nothing afterwards waits for a replay of the graph. The capture section of the backend README has the detail.", capture_err == cudaSuccess ? "capture in progress" : cudaGetErrorString(capture_err), device_id, kSharedActivationScratchKey); @@ -427,47 +452,35 @@ Error refuse_pooled_call_on_a_capturing_stream(cudaStream_t stream, int device_i // Sets out_ptr to a buffer of at least `need` bytes on `device_id`, with `stream` // ordered after the enqueue that last used the buffer. 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. +// record_shared_scratch_enqueue, and only then release the claim. // // The buffer's capacity is not reported, because no caller has any use for it: // what a call installs on its context is its own requirement, not whatever the // pool grew to for someone else. // -// A `need` of zero asks for whatever the pool already holds rather than for a -// buffer of its own: any live buffer is large enough for a call that needs -// nothing, so nothing is allocated and nothing grows. Only against an empty pool -// does it allocate, and then at kMinPooledScratchBytes. +// `need` is a real request and never zero -- execute() substitutes +// kMinPooledScratchBytes for a zero before calling, so that the size the pool +// guarantees and the size the context is told it owns are one figure and not two. // // The caller must already have refused a capturing `stream`: the handoff's wait -// invalidates a capture under every mode, and a growth's allocation, its wait on -// the buffer it replaces and its free of that buffer under every mode but -// Relaxed. +// invalidates a capture under every mode, a growth's allocation under every mode +// but Relaxed, and a growth's stream-ordered free of the buffer it replaces is +// refused outright under all three. // // Must be called with `device_id` already current: cudaEventCreateWithFlags and // cudaMalloc both act on the *current* device and nothing in here sets it. -Error get_or_grow_shared_scratch( - SharedScratchClaim& claim, - int device_id, - size_t need, - cudaStream_t stream, - void*& out_ptr) { +Error claim_shared_scratch(SharedScratchClaim& claim, int device_id, size_t need, cudaStream_t stream, void*& out_ptr) { SharedScratchDevice& dev = claim.hold(device_id); - // kMinPooledScratchBytes rather than the capacity the pool already holds: the - // reuse branch below answers with that capacity for any request it covers, so - // the two ask for the same buffer, and the smallest request says plainly that - // this call is not sizing the pool. - if (need == 0) { - need = kMinPooledScratchBytes; - } - const SharedScratchHandoff handoff = shared_scratch_claim_event(dev, []() -> cudaEvent_t { cudaEvent_t event = nullptr; - // Blocking-sync so the host yields instead of busy-spinning. The one host wait - // on this event is the one SharedScratchClaim::release() makes before freeing a - // buffer a growth displaced, and it waits for a whole inference; spinning would - // burn a core for that time and be no faster, since it is followed by a - // device-wide cudaFree. + // Blocking-sync so the host yields instead of busy-spinning. The only host + // wait ever made on this event is SharedScratchClaim::release()'s, on the + // fallback disposal path where cudaFreeAsync is unavailable, and it waits for a + // whole inference; spinning would burn a core for that time and be no faster, + // since it is followed by a device-wide cudaFree. Where cudaFreeAsync is + // available nothing waits on this event from the host at all, and the flag + // costs nothing. if (cudaEventCreateWithFlags(&event, cudaEventDisableTiming | cudaEventBlockingSync) != cudaSuccess) { return nullptr; } @@ -493,12 +506,9 @@ Error get_or_grow_shared_scratch( const bool first_buffer = dev.buffer == nullptr; RetiredScratch retired; - // Written by the pool and read by nobody, for the reason given above. - size_t capacity = 0; void* const buffer = shared_scratch_get_or_grow( dev, need, - capacity, [device_id, first_buffer](size_t bytes) -> void* { void* p = nullptr; if (cudaMalloc(&p, bytes) != cudaSuccess) { @@ -522,25 +532,26 @@ Error get_or_grow_shared_scratch( return Error::MemoryAllocationFailed; } - // Both the wait for the enqueues that used the retired buffer and the free of it - // happen at release(), with the device's lock dropped; see SharedScratchClaim. - // Nothing here makes a CUDA call that blocks on device work under that lock. - claim.retire(retired.buffer, retired.wait_for); + // The retired buffer is disposed of at release(), with the device's lock + // dropped, by a free queued on `stream`; see SharedScratchClaim for why that is + // ordered after every enqueue that used it. Nothing here makes a CUDA call that + // blocks on device work under that lock. + claim.retire(retired.buffer, retired.wait_for, stream); out_ptr = buffer; return Error::Ok; } // 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. +// shared scratch, so the next call to claim_shared_scratch waits for it. // -// Call on a `claim` that get_or_grow_shared_scratch returned Error::Ok on and that +// Call on a `claim` that claim_shared_scratch returned Error::Ok on and that // still holds the device's lock. Both halves matter, and together they are why // neither the device nor the event below is checked: a claim's device is non-null -// exactly while it holds the lock, and get_or_grow_shared_scratch has already -// failed the call with Error::Internal if the marker had no event, which only the -// test-only reset clears again and that needs the lock this claim is holding. -Error mark_shared_scratch_in_flight(SharedScratchClaim& claim, cudaStream_t stream) { +// exactly while it holds the lock, and claim_shared_scratch has already failed the +// call with Error::Internal if the marker had no event, which only the test-only +// reset clears again and that needs the lock this claim is holding. +Error record_shared_scratch_enqueue(SharedScratchClaim& claim, cudaStream_t stream) { const cudaEvent_t event = shared_scratch_mark_in_flight(*claim.device()); const cudaError_t err = cudaEventRecord(event, stream); if (err != cudaSuccess) { @@ -610,8 +621,8 @@ class ScopedErrorRecorder final : public nvinfer1::IErrorRecorder { } // namespace -// Declared in TensorRTBackend.h. See that header for why the install is checked -// rather than made and trusted. +// Declared in PooledScratchInstall.h. See that header for why the install is +// checked rather than made and trusted. bool install_pooled_scratch(nvinfer1::IExecutionContext& ctx, void* buffer, size_t bytes, int device_id) { ScopedErrorRecorder recorder; nvinfer1::IErrorRecorder* const previous = ctx.getErrorRecorder(); @@ -1074,6 +1085,26 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* bool output_staged_to_host = false; bool input_staged_from_host = false; + // A host-resident input is staged with an asynchronous copy that reads the + // caller's own memory, and that copy must not still be running when execute() + // returns: the caller is free to write that buffer again, and a copy from + // pinned memory has not read it yet -- measured, every byte the device + // received was the value written after the return. The success path already + // synchronizes whenever anything was staged (must_sync below); this covers the + // returns between the copy and that point, which would otherwise leave it + // live. It waits only on what this call queued, because must_sync means the + // stream is synchronized before the call ends either way. + struct StagedInputDrain { + cudaStream_t stream; + const bool& staged; + bool done = false; + ~StagedInputDrain() { + if (staged && !done) { + (void)cudaStreamSynchronize(stream); + } + } + } staged_input_drain{stream, input_staged_from_host}; + if (engine->cached_input_ptrs.empty()) { engine->cached_input_ptrs.resize(num_inputs, nullptr); engine->cached_input_sizes.resize(num_inputs, 0); @@ -1364,10 +1395,10 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // // A reported zero has two causes that reach here and nothing distinguishes them: // a call whose bound shapes need none -- an empty input inside a profile that - // admits one -- and a query that failed. Neither is asked to size the pool: - // get_or_grow_shared_scratch hands a zero whatever the pool already holds, and - // only where it holds nothing does it allocate, at the minimum TensorRT will - // accept rather than at the engine's profile-wide figure. + // admits one -- and a query that failed. Neither is asked to size the pool: a + // zero asks for kMinPooledScratchBytes, which any live buffer already covers, so + // only against an empty pool does it allocate anything, and then the minimum + // TensorRT will accept rather than the engine's profile-wide figure. // // The two causes are told apart by the install rather than by the query, which // is the only place the difference is observable. Where the shapes genuinely @@ -1391,22 +1422,24 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // 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 (pooled_scratch) { const size_t need = ctx->updateDeviceMemorySizeForShapes(); + // The substitution for a zero is made here and once: the same figure has to + // be the size the pool guarantees and the size the context is told it owns, + // and nothing downstream checks that two copies of it still agree -- + // setDeviceMemoryV2 refuses an install smaller than the bound shapes need and + // says nothing about one that is larger. A context whose engine needs scratch + // under some shape has to be given a buffer whatever this call's shapes need, + // or enqueueV3 refuses it. + const size_t scratch_bytes = need == 0 ? kMinPooledScratchBytes : need; void* pool = nullptr; - const Error scratch_err = get_or_grow_shared_scratch(scratch_claim, engine->device_id, need, stream, pool); + const Error scratch_err = claim_shared_scratch(scratch_claim, engine->device_id, scratch_bytes, stream, pool); if (scratch_err != Error::Ok) { return scratch_err; } - // The same substitution get_or_grow_shared_scratch makes for a zero: a - // context whose engine needs scratch under some shape has to be given a - // buffer whatever this call's shapes need, or enqueueV3 refuses it. - const size_t install_bytes = need == 0 ? kMinPooledScratchBytes : need; - if (!install_pooled_scratch(*ctx, pool, install_bytes, engine->device_id)) { + if (!install_pooled_scratch(*ctx, pool, scratch_bytes, engine->device_id)) { return Error::InvalidState; } - scratch_from_pool = true; } // ------------------------------------------------------------------ @@ -1422,9 +1455,9 @@ 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(scratch_claim, stream); + // Pairs with claim_shared_scratch: the next claimant waits on this event. + if (pooled_scratch) { + const Error mark_err = record_shared_scratch_enqueue(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. @@ -1437,13 +1470,14 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // 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. - // The release is where a growth pays for the buffer it replaced: it waits for - // the enqueues that used it and frees it, both with the lock dropped. + // The release is also where a growth disposes of the buffer it replaced, by + // queueing a free on this stream with the lock dropped. if (!scratch_claim.release()) { - // The wait failed, which for a wait on device work means this device is - // already in a faulted state. The enqueue above is submitted and this handle's - // completion marker is not armed yet, so drain before reporting, or a later - // call reconfigures exec_ctx while that enqueue is still running. + // Only the fallback disposal reports a failure, and only from its host wait, + // which for a wait on device work means this device is already in a faulted + // state. The enqueue above is submitted and this handle's completion marker is + // not armed yet, so drain before reporting, or a later call reconfigures + // exec_ctx while that enqueue is still running. (void)cudaStreamSynchronize(stream); engine->inflight_pending = false; return Error::InvalidProgram; @@ -1484,6 +1518,9 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* const bool must_sync = output_staged_to_host || input_staged_from_host || aliased_reflect_pending || !caller_stream_set; if (must_sync) { + // Every return from here on is behind the cudaStreamSynchronize below, so + // the staging drain has nothing left to do. + staged_input_drain.done = true; Error copy_err = Error::Ok; for (auto& output : outputs_needing_copy) { exec_aten::Tensor et_out = args[output.first]->toTensor(); diff --git a/tests/cpp/executorch/BUILD b/tests/cpp/executorch/BUILD index b12c8091543..86443f894a5 100644 --- a/tests/cpp/executorch/BUILD +++ b/tests/cpp/executorch/BUILD @@ -74,6 +74,7 @@ cc_test( deps = [ "//cpp:tensorrt_executorch_backend", "//cpp:tensorrt_executorch_blob_header", + "//cpp:tensorrt_executorch_pooled_scratch_install", "//cpp:tensorrt_executorch_shared_scratch_pool", "//cpp:tensorrt_executorch_shared_scratch_pool_test_hooks", "@executorch//:executorch_core", diff --git a/tests/cpp/executorch/test_shared_scratch_backend.cpp b/tests/cpp/executorch/test_shared_scratch_backend.cpp index 03c959ed745..1be6a5faec8 100644 --- a/tests/cpp/executorch/test_shared_scratch_backend.cpp +++ b/tests/cpp/executorch/test_shared_scratch_backend.cpp @@ -22,14 +22,24 @@ // green run on a host with no GPU says nothing about the pool -- and a skipped // gtest case exits zero, which Bazel reports as a passing target. A run that is // meant to have a device says so through TORCHTRT_EXECUTORCH_REQUIRE_CUDA, and -// then a skip is a failure instead, which each failing case says for itself. Where -// that variable is unset, a count of the cases that skipped is printed at the end -// of the suite. The CI invocation in .github/workflows/executorch-build-linux.yml passes -// that variable, and the job it sits in asks for a GPU runner and starts its -// container with every GPU attached. What still keeps these cases off most runs +// then the missing device is a failure instead, which each failing case says for +// itself. Where that variable is unset, a count of the cases that skipped for that +// reason is printed at the end of the suite. +// +// That variable covers the missing device and nothing else. Three cases skip for a +// second reason: two need the device full and one needs its memory quiet, and +// neither is a state a test can insist on while sharing the device with other +// processes. Those skips stand whether the variable is set or not, so a green +// required-CUDA run says every case ran, not that every case covered what it is +// named for; the skip messages say which did not. +// +// The CI invocation in .github/workflows/executorch-build-linux.yml passes that +// variable, and the job it sits in asks for a GPU runner and starts its container +// with every GPU attached. What still keeps these cases off most runs // is the lane gate on the whole ExecuTorch job in ci-linux-x86_64.yml, not the // runner it would land on. +#include "torch_tensorrt/executorch/PooledScratchInstall.h" #include "torch_tensorrt/executorch/SharedScratchPool.h" #include "torch_tensorrt/executorch/SharedScratchPoolTestHooks.h" #include "torch_tensorrt/executorch/TensorRTBackend.h" @@ -434,21 +444,21 @@ class LoadedEngine { // 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] = {batch_, rows_, cols_}; - SizesType out_sizes[3] = {batch_, 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); - ::executorch::aten::Tensor out_tensor(&out_impl); - EValue in_value(in_tensor); - EValue out_value(out_tensor); - EValue* args[2] = {&in_value, &out_value}; + return run_with_input(stream, device_in_); + } - BackendExecutionContext exec_context; - ::executorch::extension::cuda::CallerStreamGuard guard(stream); - return backend_.execute(exec_context, handle_, Span(args, 2)); + // The same run with the input bound to caller-owned host memory, which + // execute() stages through a device buffer of its own instead of binding + // directly. The output stays device-resident, so the only staging is the + // input's. `host_in` must hold bytes() bytes. + Error run_from_host_input(cudaStream_t stream, void* host_in) { + return run_with_input(stream, host_in); + } + + // Where execute() staged this handle's input, or null if it never had to. + void* staging_buffer_for_input_0() const { + const EngineHandle* const h = handle(); + return h->cached_input_ptrs.empty() ? nullptr : h->cached_input_ptrs[0]; } std::vector read_output() const { @@ -478,6 +488,24 @@ class LoadedEngine { } private: + Error run_with_input(cudaStream_t stream, void* in_ptr) { + // 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] = {batch_, rows_, cols_}; + SizesType out_sizes[3] = {batch_, rows_, cols_}; + ::executorch::aten::TensorImpl in_impl(ScalarType::Float, 3, in_sizes, in_ptr); + ::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)); + } + // 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; @@ -722,19 +750,25 @@ class SharedScratchBackendTest : public ::testing::Test { ASSERT_EQ(set_shared_scratch(backend_, false), Error::Ok); // The pool outlives every test in this file, and most of them depend on what // it holds when they start: one expects a growth, one expects none, one - // expects a first allocation. Leaving it alone made the suite pass in - // declaration order and in no other -- running the empty-input case first - // turned the handoff test's first run into a growth, whose cudaFree waits - // device-wide and so waited on the blocked host function that test parks on - // its own stream, from the thread that alone could release it. - reset_shared_scratch_pool_for_testing(); + // expects a first allocation. Leaving it alone makes each of those depend on + // which cases ran before it, so the suite passes in declaration order and in + // no other -- running the empty-input case first turns the handoff test's + // first run into a growth, which is not the state that case is written for. + // + // Reported rather than waited out: a lock still held here was leaked by the + // case before, and this case cannot run without the reset, so it fails now + // and says which state it found. + ASSERT_TRUE(reset_shared_scratch_pool_for_testing()) + << "a device's pool lock was still held when this case started, so an earlier one returned without releasing " + "its claim and the pool could not be reset"; } void TearDown() override { set_shared_scratch(backend_, false); // Also here, so a buffer this test grew is not still resident while the next // one measures device-wide memory. - reset_shared_scratch_pool_for_testing(); + ASSERT_TRUE(reset_shared_scratch_pool_for_testing()) + << "this case left a device's pool lock held, so the pool could not be reset for the next one"; } const std::vector& blob() const { @@ -997,10 +1031,12 @@ TEST_F(SharedScratchBackendTest, PooledEnginesShareOneActivationScratchAllocatio // 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 bounds cover the second allocation and the disposal of the buffer it +// replaces. That disposal is a cudaFreeAsync queued on the run's stream, so the +// bytes come back only once the stream drains; the synchronize below the run is +// what makes the second reading see them. The bounds say nothing about when the +// disposal happens relative to the return, which is what the two cases further +// down measure. // // The lower bound also fails if the pool were already large enough for the second // engine, which is how this test could otherwise pass vacuously. The fixture @@ -1063,16 +1099,19 @@ TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItRep // different window, so it narrows the misdiagnosis rather than removing it, and // the upper bound names the cause as well. Measured with another process cycling // 256 MiB allocations on this device: of 8 runs, 1 skipped here, 1 failed that - // bound and 1 failed the lower-direction assertion below, whose message is - // accurate already. The exclusive tag keeps other Bazel actions off this device, - // not other processes. + // bound and 1 tripped the direction check below, which skips with the rest now + // rather than failing. The exclusive tag keeps other Bazel actions off this + // device, not other processes. std::this_thread::sleep_for(std::chrono::steady_clock::now() - measurement_began); std::size_t settled = 0; ASSERT_TRUE(device_bytes_in_use(settled)) << "cudaMemGetInfo failed, so this test measured nothing"; const bool device_was_quiet = settled == after; - 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; + // Only something outside this test can make device-wide usage fall across a + // growth: the growth allocates the larger buffer before it disposes of the + // smaller one. It belongs with the skip below rather than being a failure of its + // own, for the reason that skip exists -- and it has to be answered before the + // subtraction, which is unsigned. + const bool device_gave_memory_back = 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 @@ -1100,11 +1139,17 @@ TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItRep << "the smaller engine stopped producing its own output once the growth moved the shared buffer"; // Last, so the output comparisons above are made either way. + if (device_gave_memory_back) { + GTEST_SKIP() << "device-wide memory in use fell from " << before << " to " << after + << " bytes across the growth, which a growth cannot do, so another process released memory on this " + "device and the two bounds below would measure it rather than the growth"; + } if (!device_was_quiet) { GTEST_SKIP() << "device-wide memory in use moved from " << after << " to " << settled << " bytes with nothing of this test's running in between, so another process is allocating on this " "device and the two bounds below would measure it rather than the growth"; } + const std::size_t growth_cost = after - before; 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"; @@ -1609,11 +1654,11 @@ TEST_F(SharedScratchBackendTest, APooledEngineRefusesACaptureBeforeAnythingCanIn // The completion record is the problem. execute() records it so that the next // call and the destructor can wait for an enqueue that outlived the return, and // under capture that record becomes a node of the graph rather than an event the -// host can wait on. The handle is then unusable: the wait at the head of the next -// execute() fails, and so does the one in ~EngineHandle. The assertions on that -// below describe what the delegate does, not what it should do; they are the -// measurement the capture documentation rests on. -TEST_F(SharedScratchBackendTest, AnUnpooledEngineCaptureReplaysAndLeavesTheHandleUnusable) { +// host can wait on. It costs the handle its next call, and the destructor's wait: +// the call after the failing one records the event again outside the capture and +// works. The assertions below describe what the delegate does, not what it should +// do; they are the measurement the capture documentation rests on. +TEST_F(SharedScratchBackendTest, AnUnpooledEngineCaptureReplaysAndCostsTheHandleItsNextCall) { cudaStream_t stream = nullptr; ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); // A node of the caller's own, so a null graph below means the capture was @@ -1675,23 +1720,40 @@ TEST_F(SharedScratchBackendTest, AnUnpooledEngineCaptureReplaysAndLeavesTheHandl // next case calls first. (void)cudaGetLastError(); + // And what it does not cost. The failing call clears the in-flight flag before + // it returns the error, so the call after it skips the wait, records the + // completion event again outside any capture and runs normally -- measured + // here, byte-for-byte against the uncaptured reference. The handle recovers; + // what does not is the ordering the event carried, because a replay of the + // graph enqueues on this same context and no execute() waits for it. + ASSERT_TRUE(unpooled.fill_output(kSentinel)); + const Error after_the_failure = unpooled.run(stream); + EXPECT_EQ(after_the_failure, Error::Ok) + << "the call after the failing one did not recover; the capture section of the README and the option key's " + "documentation say it does, so one of the two is now wrong"; + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::vector recovered = unpooled.read_output(); + ASSERT_EQ(recovered.size(), kElems); + EXPECT_EQ(std::memcmp(expected.data(), recovered.data(), kBytes), 0) + << "the recovered call returned Ok without producing what the same engine produces outside a capture"; + EXPECT_EQ(cudaGraphExecDestroy(graph_exec), cudaSuccess); EXPECT_EQ(cudaGraphDestroy(graph), cudaSuccess); EXPECT_EQ(cudaFree(captured_target), cudaSuccess); ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); } -// The other half. The case above ends with a handle the completion event has made -// unusable; this one is the failure a caller meets first if the handle already ran -// once. A call that returned with its enqueue in flight leaves the next call on +// The other half. The case above ends with a handle whose next call the completion +// event has cost it; this one is the failure a caller meets first if the handle +// already ran once. A call that returned with its enqueue in flight leaves the next call on // that handle waiting for it on the host before it may touch the context, and that // wait is prohibited under the default capture mode. Nothing refuses it: the // caller gets an error from execute() and a capture that ends with no graph. // // Between the two, a handle with the option off has no state in which capturing it -// is useful: a capture that reaches the completion record succeeds and poisons the -// handle, and one made behind a call that left an enqueue in flight is destroyed -// before it gets that far. +// is useful: a capture that reaches the completion record succeeds and costs the +// handle its next call, and one made behind a call that left an enqueue in flight +// is destroyed before it gets that far. TEST_F(SharedScratchBackendTest, AnUnpooledEngineInvalidatesACaptureOnceAnEnqueueIsInFlight) { cudaStream_t stream = nullptr; ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); @@ -1874,26 +1936,25 @@ TEST_F(SharedScratchBackendTest, ASecondPooledEnqueueWaitsForTheFirstOnAnotherSt } // --------------------------------------------------------------------------- -// Where the growth's free falls relative to the device lock +// What a growth waits for // --------------------------------------------------------------------------- -// The growth test earlier runs alone, so it cannot see whether the buffer it -// retires is freed with the device's lock held. Moving the unlock in -// SharedScratchClaim::release() below the cudaFree leaves it green. +// A growth gets rid of the buffer it replaces with a free queued on the stream it +// enqueued on, so it waits for nothing on the host. These two cases park the two +// things a synchronous disposal waits for, which is what the fallback path does +// where cudaFreeAsync is unavailable: everything queued on the device, because +// cudaFree synchronizes it, and the enqueue against the retired buffer, because +// the host waits on the marker event before freeing. Both are parked here, and the +// growing call still has to come back. // -// It matters because cudaFree waits for every stream on the device, not only for -// the ones that touched the buffer. Measured on this box: with a host function -// parked on a stream that never saw the buffer, cudaFree returned after 3000 ms -// while the cudaMalloc and the marker wait on the same growth path each returned -// at once. Under the lock, that wait is one every other pooled engine on the -// device has to sit through. -// -// So: park work on an unrelated stream to stall the free, then check the pool -// lock while the growth is still inside it. The capacity read under that same -// try_lock is what stops the case passing on a lock that is free because the -// growth has not started -- past a growth it is above what the smaller engine -// left, and the growth has not returned. -TEST_F(SharedScratchBackendTest, AGrowthFreesTheBufferItReplacesWithTheDeviceLockDropped) { +// Each measures "came back while the work was still parked" through the gate's +// watchdog, which is the only other thing that can open a gate: if the call had +// waited, the watchdog would have had to open the gate to end the test, and the +// case says so. + +// Nothing the growing engine or the pool ever submitted to runs on the parked +// stream, so only a device-wide synchronization has any reason to wait for it. +TEST_F(SharedScratchBackendTest, AGrowthDoesNotWaitForUnrelatedWorkQueuedOnTheDevice) { 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 second to be sure of growing the pool"; @@ -1913,85 +1974,69 @@ TEST_F(SharedScratchBackendTest, AGrowthFreesTheBufferItReplacesWithTheDeviceLoc const int device_id = big.handle()->device_id; // Without this the larger engine allocates rather than grows, and a growth that - // retires nothing frees nothing. + // retires nothing disposes of nothing. ASSERT_EQ(small.run(stream), Error::Ok); ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); const std::size_t before = shared_scratch_capacity_for_testing(device_id); ASSERT_GT(before, 0u) << "the smaller engine left the pool empty, so the larger one has nothing to retire"; - // Neither engine submits to this stream, so nothing on the growth path but the - // device-wide free has any reason to wait for it. StreamGate gate; ASSERT_EQ(cudaLaunchHostFunc(unrelated, hold_stream, &gate), cudaSuccess); GateRelease gate_release(gate, unrelated); - std::atomic growth_returned{false}; - Error growth_error = Error::Internal; - std::thread grower([&] { - growth_error = big.run(stream); - growth_returned.store(true); - }); - - SharedScratchDevice& dev = scratch_pool().get(device_id); - bool lock_free_after_the_growth = false; - const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); - while (!growth_returned.load() && std::chrono::steady_clock::now() < deadline) { - if (dev.mu.try_lock()) { - const std::size_t capacity_now = dev.capacity; - dev.mu.unlock(); - if (capacity_now > before && !growth_returned.load()) { - lock_free_after_the_growth = true; - break; - } - } - std::this_thread::sleep_for(std::chrono::milliseconds(5)); - } - const bool still_inside_execute = !growth_returned.load(); + ASSERT_TRUE(big.fill_output(kSentinel)); + const Error growth_error = big.run(stream); + const bool came_back_with_the_device_held = !gate.forced_open.load(); gate_release.release(); - grower.join(); ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + EXPECT_EQ(growth_error, Error::Ok); + EXPECT_TRUE(came_back_with_the_device_held) + << "the growing call did not return until the watchdog released a host function parked on a stream it never " + "submitted to, so its disposal of the retired buffer waits for the whole device"; + EXPECT_GT(shared_scratch_capacity_for_testing(device_id), before) + << "the pool did not grow, so no buffer was retired and nothing was disposed of"; + const std::vector grown_output = big.read_output(); + ASSERT_FALSE(grown_output.empty()); + EXPECT_NE(grown_output[0], kSentinel) + << "the growing engine did not write its output, so it never reached the engine"; + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); ASSERT_EQ(cudaStreamDestroy(unrelated), cudaSuccess); - - 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(growth_error, Error::Ok); - ASSERT_TRUE(still_inside_execute) - << "the growing call returned before the gate opened, so its free never stalled and there was no window in " - "which to observe the lock"; - ASSERT_GT(shared_scratch_capacity_for_testing(device_id), before) - << "the pool did not grow, so no buffer was retired and no free was made"; - EXPECT_TRUE(lock_free_after_the_growth) - << "the device's pool lock stayed held for the whole of a stalled growth free, so every other pooled engine on " - "this device waits out a device-wide free it has nothing to do with"; } -// The case above stalls the free. This one stalls the wait that comes before it: -// a growth waits for the enqueue against the buffer it retires, and that wait is -// for one inference to finish, not for the device to drain. Held under the device -// lock it would move where two pooled calls serialize -- from submission, which -// the execute() contract promises, to the completion of the call before -- so -// another pooled engine on the device would sit through an inference it has -// nothing to do with -- measured at 2119 ms with the wait moved back under the -// lock. +// The other thing the old disposal waited for. The smaller engine's own enqueue is +// parked behind a host function, so the marker event it recorded cannot signal -- +// and the retired buffer is the one that enqueue is using. The growth may not free +// it ahead of that enqueue, and it does not have to wait for it either: the free is +// queued on a stream that already waits on the same marker. // -// The smaller engine's own enqueue is parked behind a host function, so the event -// it records stays unsignalled and the growth's wait for it cannot return. As in -// the case above, the capacity read under the same try_lock is what stops this -// passing on a lock that is free because the growth has not started yet. -TEST_F(SharedScratchBackendTest, AGrowthWaitsForTheRetiredBuffersEnqueueWithTheDeviceLockDropped) { +// The smaller engine's output is checked against the same engine run with private +// scratch, because it is the one whose buffer was retired underneath it: a free +// that landed early would take its activations with it. +TEST_F(SharedScratchBackendTest, AGrowthDoesNotWaitForTheEnqueueOnTheBufferItRetires) { 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 second to be sure of growing the pool"; - ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); cudaStream_t held = nullptr; cudaStream_t growing = nullptr; ASSERT_EQ(cudaStreamCreateWithFlags(&held, cudaStreamNonBlocking), cudaSuccess); ASSERT_EQ(cudaStreamCreateWithFlags(&growing, cudaStreamNonBlocking), cudaSuccess); + // Loaded with the option off, so this one keeps its own scratch and its output + // is what the pooled run below has to reproduce. + LoadedEngine reference; + ASSERT_EQ(set_shared_scratch(backend_, false), Error::Ok); + ASSERT_EQ(reference.load(blob(), 29), Error::Ok); + ASSERT_FALSE(reference.handle()->shared_scratch); + ASSERT_EQ(reference.run(held), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(held), cudaSuccess); + const std::vector expected = reference.read_output(); + ASSERT_EQ(expected.size(), kElems); + + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); LoadedEngine small; LoadedEngine big; ASSERT_EQ(small.load(blob(), 29), Error::Ok); @@ -1999,6 +2044,7 @@ TEST_F(SharedScratchBackendTest, AGrowthWaitsForTheRetiredBuffersEnqueueWithTheD ASSERT_TRUE(small.handle()->shared_scratch); ASSERT_TRUE(big.handle()->shared_scratch); const int device_id = big.handle()->device_id; + ASSERT_TRUE(small.fill_output(kSentinel)); // Parked ahead of the smaller engine's enqueue, so that enqueue and the event // recorded after it both stay pending for as long as this test wants them to. @@ -2010,48 +2056,156 @@ TEST_F(SharedScratchBackendTest, AGrowthWaitsForTheRetiredBuffersEnqueueWithTheD const std::size_t before = shared_scratch_capacity_for_testing(device_id); ASSERT_GT(before, 0u) << "the smaller engine left the pool empty, so the larger one has nothing to retire"; - std::atomic growth_returned{false}; - Error growth_error = Error::Internal; - std::thread grower([&] { - growth_error = big.run(growing); - growth_returned.store(true); - }); - - SharedScratchDevice& dev = scratch_pool().get(device_id); - bool lock_free_during_the_wait = false; - const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); - while (!growth_returned.load() && std::chrono::steady_clock::now() < deadline) { - if (dev.mu.try_lock()) { - const std::size_t capacity_now = dev.capacity; - dev.mu.unlock(); - if (capacity_now > before && !growth_returned.load()) { - lock_free_during_the_wait = true; - break; - } - } - std::this_thread::sleep_for(std::chrono::milliseconds(5)); - } - const bool still_inside_execute = !growth_returned.load(); + const Error growth_error = big.run(growing); + const bool came_back_with_the_enqueue_parked = !gate.forced_open.load(); gate_release.release(); - grower.join(); ASSERT_EQ(cudaStreamSynchronize(held), cudaSuccess); ASSERT_EQ(cudaStreamSynchronize(growing), cudaSuccess); + + EXPECT_EQ(growth_error, Error::Ok); + EXPECT_TRUE(came_back_with_the_enqueue_parked) + << "the growing call did not return until the watchdog released the enqueue against the buffer it retired, so a " + "growth waits out an inference it has nothing to do with"; + EXPECT_GT(shared_scratch_capacity_for_testing(device_id), before) + << "the pool did not grow, so no buffer was retired and there was nothing to order the free against"; + const std::vector small_output = small.read_output(); + ASSERT_EQ(small_output.size(), kElems); + EXPECT_EQ(std::memcmp(expected.data(), small_output.data(), kBytes), 0) + << "the engine whose scratch buffer the growth retired did not produce what the same engine produces with " + "private scratch, so the buffer went away while its enqueue was still using it"; + ASSERT_EQ(cudaStreamDestroy(held), cudaSuccess); ASSERT_EQ(cudaStreamDestroy(growing), cudaSuccess); +} + +// --------------------------------------------------------------------------- +// A failing call and the host-input copy it queued +// --------------------------------------------------------------------------- + +// An input that is not device-resident is staged with cudaMemcpyAsync from the +// caller's own memory, and from pinned memory that copy reads the caller's bytes +// when the stream reaches it, not when it is queued. So a call that fails after +// queueing one must not return while it is still pending: the caller owns that +// buffer again the moment execute() returns, and what it writes there is what +// the device then receives. +// +// The failure driven here is the pool's allocation, the one early return in the +// pooled scratch block that can be reached from inside this process. The stream +// is parked ahead of the copy so it cannot complete on its own, and the caller +// overwrites its buffer as soon as the call comes back -- which is exactly what a +// caller may do. +// +// Three outcomes are separated at the end, by giving the failing call a +// different input from the successful one before it: the sentinel means the copy +// took the caller's post-return write, the first call's pattern means it never +// ran at all, and the second call's pattern is the only pass. +TEST_F(SharedScratchBackendTest, AFailedPooledCallDrainsTheHostInputCopyItQueued) { + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + + LoadedEngine pooled; + ASSERT_EQ(pooled.load(blob(), 41), Error::Ok); + ASSERT_TRUE(pooled.handle()->shared_scratch); + const int device_id = pooled.handle()->device_id; + + // Pinned: cudaMemcpyAsync from pageable memory does not return until the + // driver has taken the bytes, so a later write cannot reach the device and + // there is no hazard to pin. + float* host_in = nullptr; + ASSERT_EQ(cudaHostAlloc(reinterpret_cast(&host_in), pooled.bytes(), cudaHostAllocDefault), cudaSuccess); + for (std::size_t i = 0; i < pooled.elems(); ++i) { + host_in[i] = pattern(i, 41); + } + + // One good run first, so the failing one reaches the copy: the staging buffer + // is allocated on the call that first needs it, and with the device full that + // allocation would fail ahead of everything this case is about. + ASSERT_EQ(pooled.run_from_host_input(stream, host_in), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + void* const staging = pooled.staging_buffer_for_input_0(); + ASSERT_NE(staging, nullptr) << "the run bound the caller's host memory directly, so nothing was staged and there is " + "no asynchronous copy to leave running"; + + // The second call's input, so what the device holds at the end says which of + // the three outcomes happened. + for (std::size_t i = 0; i < pooled.elems(); ++i) { + host_in[i] = pattern(i, 42); + } + + // Empties the pool, so the failing run allocates instead of reusing what the + // run above left in it. + ASSERT_TRUE(reset_shared_scratch_pool_for_testing()); + ASSERT_EQ(shared_scratch_capacity_for_testing(device_id), 0u); + + DeviceMemoryHog hog; + if (!hog.leave_less_free_than(scratch_bytes_)) { + ASSERT_EQ(cudaFreeHost(host_in), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); + GTEST_SKIP() << "could not take the device below the " << scratch_bytes_ + << " bytes this engine's scratch needs, so its allocation would have succeeded"; + } + + StreamGate gate; + ASSERT_EQ(cudaLaunchHostFunc(stream, hold_stream, &gate), cudaSuccess); + GateRelease gate_release(gate, stream); + + std::atomic call_returned{false}; + Error failed_run = Error::Ok; + std::thread caller([&] { + failed_run = pooled.run_from_host_input(stream, host_in); + call_returned.store(true); + }); + + // Long enough that a call which does not wait for its own copy has returned; + // the stream cannot move until the gate opens, so nothing else ends the wait. + std::this_thread::sleep_for(std::chrono::milliseconds(500)); + const bool returned_with_the_copy_still_queued = call_returned.load(); + if (returned_with_the_copy_still_queued) { + for (std::size_t i = 0; i < pooled.elems(); ++i) { + host_in[i] = kSentinel; + } + } + gate_release.release(); + caller.join(); + if (!returned_with_the_copy_still_queued) { + for (std::size_t i = 0; i < pooled.elems(); ++i) { + host_in[i] = kSentinel; + } + } + // After the gate, because a free is device-wide and would otherwise block on it. + hog.release(); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + std::vector staged(pooled.elems(), 0.0f); + ASSERT_EQ(cudaMemcpy(staged.data(), staging, pooled.bytes(), cudaMemcpyDeviceToHost), cudaSuccess); + ASSERT_EQ(cudaFreeHost(host_in), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); 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(growth_error, Error::Ok); - ASSERT_TRUE(still_inside_execute) - << "the growing call returned before the gate opened, so nothing about it stalled and there was no window in " - "which to observe the lock"; - ASSERT_GT(shared_scratch_capacity_for_testing(device_id), before) - << "the pool did not grow, so no buffer was retired and there was nothing to wait for"; - EXPECT_TRUE(lock_free_during_the_wait) - << "the device's pool lock stayed held while a growth waited for the enqueue on the buffer it retired, so every " - "other pooled engine on this device waits out that inference"; + << "the watchdog had to open the gate, so the timings below were not measured under the conditions they describe"; + if (failed_run == Error::Ok) { + GTEST_SKIP() << "the pooled allocation succeeded with the device full, so this run did not fail where the case " + "needs it to"; + } + ASSERT_EQ(failed_run, Error::MemoryAllocationFailed) + << "the run failed somewhere other than the pool's allocation, so it says nothing about that early return"; + + EXPECT_FALSE(returned_with_the_copy_still_queued) + << "the failing call returned while the copy reading the caller's host input was still queued behind a parked " + "stream, so whatever the caller writes next is what the device receives"; + std::size_t bytes_the_device_did_not_get = 0; + for (std::size_t i = 0; i < pooled.elems(); ++i) { + if (staged[i] != pattern(i, 42)) { + ++bytes_the_device_did_not_get; + } + } + EXPECT_EQ(bytes_the_device_did_not_get, 0u) + << "the device does not hold the input this call was made with. It holds " << staged[0] + << " at element 0: " << kSentinel << " is what the caller wrote after execute() returned, " << pattern(0, 41) + << " is the previous call's input and means the copy never ran, and " << pattern(0, 42) << " is a pass"; } // --------------------------------------------------------------------------- @@ -2115,13 +2269,24 @@ TEST_F(SharedScratchBackendTest, TwoThreadsRunningPooledEnginesOnOneDeviceKeepTh // 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. + // below; releasing both threads together caught nearly all of them. + // + // The wait has a deadline because the loop it sits in contains a delegate call, + // and a pooled call that blocks is one of the failures this case exists to + // catch. Without one the thread whose partner is stuck spins here until the + // target's own timeout, which reports as a killed binary rather than as this + // case; with one, the run ends and the assertion below names it. + constexpr std::chrono::seconds kRendezvousDeadline{30}; std::atomic arrived{0}; - auto submit_together = [&arrived](int iteration) { + std::atomic partner_never_arrived{false}; + auto submit_together = [&](int iteration) { arrived.fetch_add(1); + const auto deadline = std::chrono::steady_clock::now() + kRendezvousDeadline; while (arrived.load() < 2 * (iteration + 1)) { + if (std::chrono::steady_clock::now() >= deadline) { + partner_never_arrived.store(true); + return; + } std::this_thread::yield(); } }; @@ -2155,6 +2320,10 @@ TEST_F(SharedScratchBackendTest, TwoThreadsRunningPooledEnginesOnOneDeviceKeepTh ASSERT_EQ(cudaStreamDestroy(first_stream), cudaSuccess); ASSERT_EQ(cudaStreamDestroy(second_stream), cudaSuccess); + EXPECT_FALSE(partner_never_arrived.load()) + << "one thread waited " << kRendezvousDeadline.count() + << " seconds at the rendezvous without its partner arriving, so a pooled call blocked instead of returning and " + "the runs after that point were not submitted together"; 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) diff --git a/tests/cpp/executorch/test_shared_scratch_pool.cpp b/tests/cpp/executorch/test_shared_scratch_pool.cpp index ec8ea58c817..223b103b701 100644 --- a/tests/cpp/executorch/test_shared_scratch_pool.cpp +++ b/tests/cpp/executorch/test_shared_scratch_pool.cpp @@ -41,6 +41,7 @@ #include #include +#include #endif namespace torch_tensorrt { @@ -99,10 +100,10 @@ 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) { +void* call(SharedScratchDevice& dev, FakeAllocator& a, std::size_t need) { RetiredScratch retired; void* const p = shared_scratch_get_or_grow( - dev, need, out_size, [&a](std::size_t bytes) { return a.alloc(bytes); }, retired); + dev, need, [&a](std::size_t bytes) { return a.alloc(bytes); }, retired); if (retired.buffer != nullptr) { a.retire(retired.buffer, retired.wait_for); } @@ -112,12 +113,11 @@ void* call(SharedScratchDevice& dev, FakeAllocator& a, std::size_t need, std::si TEST(SharedScratchPool, FirstRequestAllocatesExactSize) { SharedScratchDevice dev; FakeAllocator a; - std::size_t out = 0; - void* p = call(dev, a, /*need=*/1024, out); + void* p = call(dev, a, /*need=*/1024); EXPECT_NE(p, nullptr); - EXPECT_EQ(out, 1024u); + EXPECT_EQ(dev.capacity, 1024u); ASSERT_EQ(a.alloc_count(), 1); EXPECT_EQ(a.alloc_sizes[0], 1024u); EXPECT_TRUE(a.retirements.empty()); @@ -126,23 +126,16 @@ TEST(SharedScratchPool, FirstRequestAllocatesExactSize) { TEST(SharedScratchPool, ReusesWhenExistingBufferIsLargeEnough) { SharedScratchDevice dev; FakeAllocator a; - std::size_t out = 0; - void* first = call(dev, a, 4096, out); + void* first = call(dev, a, 4096); // A smaller and an equal request must both reuse the same buffer (no realloc). - // Each reports into an output variable of its own, so what the reuse path writes - // is asserted rather than what the first call left in `out`, which already holds - // the value both of them are expected to write. - std::size_t out2 = 0; - std::size_t out3 = 0; - void* second = call(dev, a, 1000, out2); - void* third = call(dev, a, 4096, out3); + void* second = call(dev, a, 1000); + void* third = call(dev, a, 4096); EXPECT_EQ(second, first); EXPECT_EQ(third, first); - EXPECT_EQ(out3, 4096u); - // Reuse reports the buffer's capacity, not the smaller amount asked for. - EXPECT_EQ(out2, 4096u); + // The smaller request did not shrink the pool to what it asked for. + EXPECT_EQ(dev.capacity, 4096u); EXPECT_EQ(a.alloc_count(), 1); EXPECT_TRUE(a.retirements.empty()); } @@ -150,22 +143,21 @@ TEST(SharedScratchPool, ReusesWhenExistingBufferIsLargeEnough) { TEST(SharedScratchPool, GrowsMonotonicallyToMaxAndRetiresOldBuffer) { SharedScratchDevice dev; FakeAllocator a; - std::size_t out = 0; - void* small = call(dev, a, 1024, out); - void* big = call(dev, a, 8192, out); + void* small = call(dev, a, 1024); + void* big = call(dev, a, 8192); EXPECT_NE(big, small); - EXPECT_EQ(out, 8192u); + EXPECT_EQ(dev.capacity, 8192u); ASSERT_EQ(a.alloc_count(), 2); EXPECT_EQ(a.alloc_sizes[1], 8192u); 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); + void* reuse = call(dev, a, 512); EXPECT_EQ(reuse, big); - EXPECT_EQ(out, 8192u); + EXPECT_EQ(dev.capacity, 8192u); EXPECT_EQ(a.alloc_count(), 2); } @@ -173,9 +165,8 @@ TEST(SharedScratchPool, GrowRetiresTheOldBufferWithTheEventToWaitOn) { SharedScratchDevice dev; FakeAllocator a; FakeEventFactory events; - std::size_t out = 0; - void* small = call(dev, a, 1024, out); + void* small = call(dev, a, 1024); ASSERT_NE(small, nullptr); // An enqueue against `small` has been submitted and recorded, so its @@ -183,7 +174,7 @@ TEST(SharedScratchPool, GrowRetiresTheOldBufferWithTheEventToWaitOn) { 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_NE(call(dev, a, 8192), nullptr); ASSERT_EQ(a.retirements.size(), 1u); EXPECT_EQ(a.retirements[0].first, small); @@ -197,15 +188,14 @@ TEST(SharedScratchPool, GrowHasNothingToWaitForWhenNoEnqueueWasRecorded) { SharedScratchDevice dev; FakeAllocator a; FakeEventFactory events; - std::size_t out = 0; - void* small = call(dev, a, 1024, out); + void* small = call(dev, a, 1024); 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_NE(call(dev, a, 8192), nullptr); ASSERT_EQ(a.retirements.size(), 1u); EXPECT_EQ(a.retirements[0].first, small); @@ -215,38 +205,35 @@ TEST(SharedScratchPool, GrowHasNothingToWaitForWhenNoEnqueueWasRecorded) { TEST(SharedScratchPool, AllocationFailureLeavesExistingBufferUntouched) { SharedScratchDevice dev; FakeAllocator a; - std::size_t out = 0; - void* first = call(dev, a, 1024, out); + void* first = call(dev, a, 1024); 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(dev, a, 8192, out2); + void* failed = call(dev, a, 8192); EXPECT_EQ(failed, nullptr); 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); + void* again = call(dev, a, 1024); EXPECT_EQ(again, first); - EXPECT_EQ(out, 1024u); + EXPECT_EQ(dev.capacity, 1024u); } TEST(SharedScratchPool, FirstAllocationFailureReturnsNullAndStoresNothing) { SharedScratchDevice dev; FakeAllocator a; - std::size_t out = 0; a.fail_next = true; - void* p = call(dev, a, 1024, out); + void* p = call(dev, a, 1024); 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(dev, a, 1024, out); + void* q = call(dev, a, 1024); EXPECT_NE(q, nullptr); EXPECT_EQ(a.alloc_count(), 1); } @@ -257,13 +244,12 @@ TEST(SharedScratchPool, FirstAllocationFailureReturnsNullAndStoresNothing) { TEST(SharedScratchPool, EveryPathClearsTheRetirementItReports) { SharedScratchDevice dev; FakeAllocator a; - std::size_t out = 0; // Deliberately reused across the calls below, which is the state the clearing is // for; `call` above gives each of its calls a fresh one. RetiredScratch retired; const auto request = [&](std::size_t need) { return shared_scratch_get_or_grow( - dev, need, out, [&a](std::size_t bytes) { return a.alloc(bytes); }, retired); + dev, need, [&a](std::size_t bytes) { return a.alloc(bytes); }, retired); }; void* const first = request(1024); @@ -382,18 +368,17 @@ TEST(SharedScratchHandoffTest, ASlotWithNoEventIsNotMarked) { 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); + void* dev0 = call(pool.get(0), a, 2048); + void* dev1 = call(pool.get(1), a, 2048); EXPECT_NE(dev0, dev1); EXPECT_EQ(a.alloc_count(), 2); 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); + void* dev1_big = call(pool.get(1), a, 9000); + void* dev0_again = call(pool.get(0), a, 2048); EXPECT_NE(dev1_big, dev1); EXPECT_EQ(dev0_again, dev0); ASSERT_EQ(a.retirements.size(), 1u); @@ -430,12 +415,11 @@ TEST(SharedScratchPoolRegistry, ResetHandsBackEverySlotAndLeavesItEmpty) { FakeAllocator zero; FakeAllocator one; FakeEventFactory events; - std::size_t out = 0; SharedScratchDevice& dev0 = pool.get(0); SharedScratchDevice& dev1 = pool.get(1); - void* const dev0_buffer = call(dev0, zero, 4096, out); - void* const dev1_buffer = call(dev1, one, 2048, out); + void* const dev0_buffer = call(dev0, zero, 4096); + void* const dev1_buffer = call(dev1, one, 2048); const cudaEvent_t dev0_event = shared_scratch_claim_event(dev0, std::ref(events)).event; shared_scratch_mark_in_flight(dev0); // Device 1 is left with a buffer and no event, which is the state of a slot @@ -485,10 +469,9 @@ TEST(SharedScratchPoolRegistry, ResetHandsBackEverySlotAndLeavesItEmpty) { // A smaller request than the freed buffer served: reuse would satisfy it from // the stale capacity and allocate nothing, so this is what distinguishes a // cleared slot from one the reset only emptied of its event. - std::size_t after = 0; - void* const fresh = call(dev0, zero, 1024, after); + void* const fresh = call(dev0, zero, 1024); EXPECT_NE(fresh, nullptr); - EXPECT_EQ(after, 1024u); + EXPECT_EQ(dev0.capacity, 1024u); EXPECT_EQ(zero.alloc_count(), 2) << "the slot was not cleared, so the request reused the freed buffer"; EXPECT_TRUE(zero.retirements.empty()) << "the cleared slot retired a buffer the reset had already handed back"; } @@ -501,13 +484,12 @@ TEST(SharedScratchPoolRegistry, ResetHandsBackTheEventOfASlotThatNeverGotABuffer SharedScratchPool pool; FakeAllocator a; FakeEventFactory events; - std::size_t out = 0; SharedScratchDevice& dev = pool.get(3); const cudaEvent_t event = shared_scratch_claim_event(dev, std::ref(events)).event; ASSERT_NE(event, nullptr); a.fail_next = true; - ASSERT_EQ(call(dev, a, 1024, out), nullptr); + ASSERT_EQ(call(dev, a, 1024), nullptr); ASSERT_EQ(dev.buffer, nullptr) << "the allocation did not fail, so this slot is not the state under test"; std::vector> disposed; @@ -531,7 +513,6 @@ TEST(SharedScratchPoolRegistry, ResetLeavesEveryEntryWhereItWas) { SharedScratchPool pool; FakeAllocator alloc; FakeEventFactory events; - std::size_t out = 0; std::vector before; before.reserve(kDevices); @@ -539,7 +520,7 @@ TEST(SharedScratchPoolRegistry, ResetLeavesEveryEntryWhereItWas) { SharedScratchDevice& dev = pool.get(id); // Give each slot something, so the reset has work to do on all of them rather // than skipping past empty ones. - ASSERT_NE(call(dev, alloc, 1024, out), nullptr); + ASSERT_NE(call(dev, alloc, 1024), nullptr); ASSERT_NE(shared_scratch_claim_event(dev, std::ref(events)).event, nullptr); before.push_back(&dev); } @@ -569,10 +550,9 @@ TEST(SharedScratchPoolRegistry, ResetDisposesWithNoLockHeld) { constexpr int kUntouchedDevice = 99; SharedScratchPool pool; FakeAllocator alloc; - std::size_t out = 0; SharedScratchDevice& dev0 = pool.get(0); - ASSERT_NE(call(dev0, alloc, 1024, out), nullptr); + ASSERT_NE(call(dev0, alloc, 1024), nullptr); std::atomic disposing{false}; std::atomic reset_returned{false}; @@ -623,6 +603,57 @@ TEST(SharedScratchPoolRegistry, ResetDisposesWithNoLockHeld) { "teardown blocks every other device"; } +// A device lock still held between cases is a claim that was never released -- +// the defect the backend suite's own case hunts. Waiting for it would hang the +// fixture that resets before and after every case, so the run would end in a +// target timeout and the message naming the leak would never be printed. The +// reset reports the slot and leaves it instead. +// +// The holder below releases on its own deadline, so a reset that waits for the +// lock rather than reporting it ends this case in a failure rather than a hang. +TEST(SharedScratchPoolRegistry, ResetReportsALockedSlotRatherThanWaitingForIt) { + constexpr std::chrono::seconds kHolderDeadline{30}; + + SharedScratchPool pool; + FakeAllocator alloc; + + SharedScratchDevice& leaked = pool.get(0); + SharedScratchDevice& ordinary = pool.get(1); + ASSERT_NE(call(leaked, alloc, 1024), nullptr); + ASSERT_NE(call(ordinary, alloc, 2048), nullptr); + + std::atomic holding{false}; + std::atomic release_it{false}; + std::thread holder([&] { + leaked.mu.lock(); + holding.store(true); + const auto deadline = std::chrono::steady_clock::now() + kHolderDeadline; + while (!release_it.load() && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + leaked.mu.unlock(); + }); + while (!holding.load()) { + std::this_thread::yield(); + } + + int disposed = 0; + const auto started = std::chrono::steady_clock::now(); + const std::size_t still_locked = pool.reset_for_testing([&](int, void*, cudaEvent_t) { ++disposed; }); + const auto took = std::chrono::steady_clock::now() - started; + release_it.store(true); + holder.join(); + + EXPECT_EQ(still_locked, 1u) << "the reset did not report the one device whose lock it could not take"; + EXPECT_LT(took, kHolderDeadline) << "the reset returned only once the lock was released, so it waited for a claim " + "that a leak would never release"; + EXPECT_EQ(disposed, 1) << "the reset disposed of " << disposed + << " slots: it should hand back the unlocked one and leave the locked one alone"; + EXPECT_NE(leaked.buffer, nullptr) << "the reset emptied a slot whose lock it never held, so it handed the disposer a " + "buffer a live claimant is still using"; + EXPECT_EQ(ordinary.buffer, nullptr) << "one locked slot stopped the reset clearing the others"; +} + TEST(SharedScratchPoolRegistry, AGrowthOnOneDeviceDoesNotBlockAClaimOnAnother) { SharedScratchPool pool; // One allocator per thread: the two claims share the registry and nothing else. @@ -637,12 +668,10 @@ TEST(SharedScratchPoolRegistry, AGrowthOnOneDeviceDoesNotBlockAClaimOnAnother) { SharedScratchDevice& dev0 = pool.get(0); std::thread grower([&] { std::lock_guard lk(dev0.mu); - std::size_t out = 0; RetiredScratch retired; shared_scratch_get_or_grow( dev0, 4096, - out, [&](std::size_t bytes) { entered_alloc.set_value(); leave.wait(); @@ -669,10 +698,9 @@ TEST(SharedScratchPoolRegistry, AGrowthOnOneDeviceDoesNotBlockAClaimOnAnother) { auto claim = std::async(std::launch::async, [&] { 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); }, retired); + dev1, 2048, [&](std::size_t bytes) { return one.alloc(bytes); }, retired); }); const bool served = claim.wait_for(std::chrono::seconds(10)) == std::future_status::ready; @@ -762,6 +790,44 @@ TEST(SharedScratchPoolRegistry, ConcurrentLookupsKeepTheRegistryIntact) { // single non-zero status is a real failure and not a flake. constexpr int kTeardownChildren = 24; +// What a child exits with when it could not get as far as the state under test -- +// starting a thread, two dozen forks in, is the way that happens. Distinct from +// zero and from the teardown crash this case counts, so a run that could not set +// itself up is not read as evidence either way. +constexpr int kChildCouldNotStart = 121; + +// Long enough that the children below -- a thread, a 30 ms sleep and exit() -- +// are nowhere near it. +constexpr std::chrono::seconds kChildDeadline{30}; + +// Reaps `child`, killing it if it overruns `kChildDeadline`; returns false when +// it had to. Without a deadline a wedged child parks the parent here until the +// whole target times out, which reports as a timeout on the binary rather than as +// this case failing, and this target sets no timeout of its own. An interrupted +// wait is retried, since a signal can arrive at any point and says nothing about +// the child. +bool reap_child(pid_t child, int& status) { + const auto deadline = std::chrono::steady_clock::now() + kChildDeadline; + for (;;) { + const pid_t reaped = waitpid(child, &status, WNOHANG); + if (reaped == child) { + return true; + } + if (reaped == -1 && errno != EINTR) { + return false; + } + if (std::chrono::steady_clock::now() >= deadline) { + kill(child, SIGKILL); + // Reaped even so: a zombie left behind would be inherited by the next fork + // in the loop as a child that never reports. + while (waitpid(child, &status, 0) == -1 && errno == EINTR) { + } + return false; + } + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } +} + [[noreturn]] void hold_the_pool_open_then_exit() { static std::atomic inside{false}; std::thread claimant([] { @@ -785,6 +851,8 @@ constexpr int kTeardownChildren = 24; TEST(SharedScratchPoolRegistry, TheProcessPoolOutlivesStaticDestructionUnderALiveClaimant) { int killed_by_signal = 0; int exited_nonzero = 0; + int could_not_start = 0; + int overran = 0; for (int i = 0; i < kTeardownChildren; ++i) { // Anything gtest has buffered would otherwise be written twice, once by each // side of the fork. @@ -792,17 +860,36 @@ TEST(SharedScratchPoolRegistry, TheProcessPoolOutlivesStaticDestructionUnderALiv const pid_t child = fork(); ASSERT_NE(child, -1) << "fork failed: " << std::strerror(errno); if (child == 0) { - hold_the_pool_open_then_exit(); + // Nothing may leave the child by any route but exit. Left to unwind, a + // failed thread construction returns into the test body and the child + // carries on through the rest of the binary as a second gtest process, + // forking children of its own and writing over the parent's output. + try { + hold_the_pool_open_then_exit(); + } catch (...) { + _exit(kChildCouldNotStart); + } } int status = 0; - ASSERT_EQ(waitpid(child, &status, 0), child) << "waitpid failed: " << std::strerror(errno); + if (!reap_child(child, status)) { + ++overran; + continue; + } if (WIFSIGNALED(status)) { ++killed_by_signal; + } else if (WEXITSTATUS(status) == kChildCouldNotStart) { + ++could_not_start; } else if (WEXITSTATUS(status) != 0) { ++exited_nonzero; } } + EXPECT_EQ(overran, 0) << overran << " of " << kTeardownChildren << " children were still running after " + << kChildDeadline.count() + << " seconds and were killed, so they neither reached teardown nor reported anything about it"; + EXPECT_EQ(could_not_start, 0) << could_not_start << " of " << kTeardownChildren + << " children could not start their claimant thread, so those runs never put anything " + "inside the pool and say nothing about what teardown does to a live claimant"; EXPECT_EQ(killed_by_signal, 0) << killed_by_signal << " of " << kTeardownChildren << " children died on a signal at exit with a thread still inside the pool, so the " "registry is being destroyed out from under a live claimant"; From 949efb5c467b9e7b3937756a0aa47b9619a3ac3a Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Tue, 8 Sep 2026 02:44:19 -0700 Subject: [PATCH 11/13] fix(executorch): stop a growth's queued free holding the bytes it retired `cudaFreeAsync` does not give the bytes back when the stream reaches the free. Measured on an A100 with CUDA 13.0 and driver 13.0: on a `cudaMalloc`'d pointer it returns `cudaSuccess` and defers, and the bytes come back at the next `cudaStreamSynchronize` of that stream and at no point before it -- a stream `cudaStreamQuery` reports as drained still holds them, `cudaMemGetInfo` does not move, and a `cudaMalloc` of the same size fails with out of memory until the synchronize. The previous commit queued that free on every growth, including the caller-stream, device-resident path that returns with its enqueue still running and never synchronizes. Nothing there would ever have released the bytes, so peak memory became the sum of every size the pool grew to rather than the largest of them. Replaying the pool's own growth sequence on one stream -- 8, 16, 24 and 32 GiB, no synchronize between them -- four growths that all succeed with `cudaFree` fail on the fourth with out of memory. That is a memory regression in the multi-layer-engine case the pool exists to fix. The disposal now turns on `must_sync`, which `execute()` already computes and which is now decided before the claim rather than after the enqueue. A call that synchronizes the stream queues the free and its own synchronization returns the bytes, waiting for nothing it did not submit -- around 10 ms against the 1505 ms the host wait took, with a host function parked for 1500 ms on an unrelated stream. A call that does not makes a host wait on the handoff event and a device-wide `cudaFree`, the previously documented behaviour, whose unbounded wait is back in the installed contract and the README with it. The same host wait is also the fallback where a device has no stream-ordered allocator, so it is no longer a path only a mutation reaches. `ALargerEngineGrowsThePoolAndFreesTheBufferItReplaces` reads device memory the moment `execute()` returns and samples its control window before any synchronize of its own, because a synchronize in between is exactly what releases a queued free -- that blind spot is why the growth test could not see the regression. It is now a fixture helper run on both paths. `AGrowthDoesNotWaitForUnrelatedWorkQueuedOnTheDevice` moves to the synchronized path, where the queued free is, and skips on a device reporting no `cudaDevAttrMemoryPoolsSupported` rather than reporting a defect the fallback does not have. `AGrowthDoesNotWaitForTheEnqueueOnTheBufferItRetires` becomes `AGrowthOnACallerStreamWaitsForTheEnqueueOnTheBufferItRetires`: a growth there does wait for that enqueue, and what has to hold is that it waits rather than freeing under it. Also in this commit: - The two-thread case stops at the first rendezvous that expires. Armed afresh for 60 iterations a wedged partner cost the survivor 60 x 30 s against a 900 s target timeout, so the assertion naming the failure never printed. Measured at a 1 s deadline: 1.8 s now, 61.7 s before. - `test_shared_scratch_pool` takes `timeout = "long"`. Its fork case budgets 24 children x 30 s = 720 s against Bazel's 300 s default for a target that names neither size nor timeout, so the per-child deadline could not deliver what it was added for. - The reset's body moves out of `SharedScratchPool.h`, which ships in the source package, into a testonly `SharedScratchPoolReset.h`. The shipped header keeps only the friend declaration, so the three-line consumer that freed the live pool no longer compiles. - A second translation unit in the pool test pins that `scratch_pool()` is one registry. `inline` -> `static` used to leave the target fully green. - `ResetReportsALockedSlotRatherThanWaitingForIt` gets a second unlocked slot, so the count the reset answers with is no longer equal to the number it cleared. - The four post-enqueue drains become one helper, which also clears the staged-input guard's flag: three error returns synchronized the stream and then let the guard synchronize it again. - The growth's fallback log names a missing stream-ordered allocator only for `cudaErrorNotSupported`, not for a fault on a device that has one. - The installed header no longer claims the pool adds four error returns an unpooled call never makes. It adds six failure points and one new code. - Prose: the README's opening sentence carries the two exceptions the header already states; the zero-size call is described as being told one byte rather than the pool's size; the capture section points at the `execute()` contract rather than at the option key's comment; the test file's coverage note stops attributing its absence from CI to a lane gate that does not apply to an ordinary pull-request push, and names the device-memory blast radius of the two cases that fill the device; the stale "only outside Relaxed" half of a capture comment is corrected; the four places that reasoned from "a capturing call is never one that synchronizes the stream" say instead that no capture mode permits such a call, which is the fact the code rests on; `RetiredScratch` names which disposal each path takes rather than calling one of them the fallback; and a zombie is described as a process-table entry the parent holds rather than something the next fork inherits. - `third_party/cuda` gets a headers target that is the only place the toolkit glob is written, one catch-all pattern rather than a list of extensions the catch-all already matched, and the three library targets beside it depend on it instead of repeating it; `tests/cpp`'s duplicate `test_suite` is the label for the one the workflow runs rather than a second list. 6 targets, 87 -> 89 cases, all passing on an A100 with `--nocache_test_results` and `TORCHTRT_EXECUTORCH_REQUIRE_CUDA=1`, no skips. --- cpp/BUILD | 22 ++ .../executorch/TensorRTBackend.h | 88 +++-- cpp/src/torch_tensorrt/executorch/README.md | 143 +++++--- .../executorch/SharedScratchPool.h | 95 ++---- .../executorch/SharedScratchPoolReset.h | 91 +++++ .../executorch/SharedScratchPoolTestHooks.cpp | 40 +-- .../executorch/TensorRTBackend.cpp | 288 ++++++++++------ tests/cpp/BUILD | 10 +- tests/cpp/executorch/BUILD | 12 +- .../test_shared_scratch_backend.cpp | 319 +++++++++++++----- .../executorch/test_shared_scratch_pool.cpp | 57 +++- .../test_shared_scratch_pool_other_tu.cpp | 29 ++ third_party/cuda/BUILD | 36 +- 13 files changed, 828 insertions(+), 402 deletions(-) create mode 100644 cpp/src/torch_tensorrt/executorch/SharedScratchPoolReset.h create mode 100644 tests/cpp/executorch/test_shared_scratch_pool_other_tu.cpp diff --git a/cpp/BUILD b/cpp/BUILD index d2db81912ee..1335303cfc7 100644 --- a/cpp/BUILD +++ b/cpp/BUILD @@ -242,6 +242,27 @@ cc_library( }), ) +# The body of the pool's test-only reset. testonly and under src/, so it is in +# neither the packaged include tree nor executorch_backend_source_files below: a +# release build has no definition of it anywhere. SharedScratchPool.h keeps only +# the friend declaration that lets this reach the registry's internals. +cc_library( + name = "tensorrt_executorch_shared_scratch_pool_reset", + testonly = True, + hdrs = [ + "src/torch_tensorrt/executorch/SharedScratchPoolReset.h", + ], + strip_include_prefix = "src", + target_compatible_with = select({ + ":linux_x86_64": [], + ":sbsa": [], + "//conditions:default": ["@platforms//:incompatible"], + }), + deps = [ + ":tensorrt_executorch_shared_scratch_pool", + ], +) + # The pool's test-only entry points. A separate target, and testonly, so the # released archive -- Bazel's :tensorrt_executorch_backend and CMake's # executorch_trt_backend, neither of which compiles this source -- carries no @@ -265,6 +286,7 @@ cc_library( }), deps = [ ":tensorrt_executorch_shared_scratch_pool", + ":tensorrt_executorch_shared_scratch_pool_reset", ] + select({ ":linux_x86_64": [ "@cuda//:cudart", diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index 7489f08a93d..abf9ee8fd5f 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -133,21 +133,33 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { // handle whose context was created while the option was off keeps its own // scratch and is not subject to this, nor is one whose engine reports needing no // activation scratch under any shape, which is left out of the pool. For the - // handles that do draw on it, four further consequences: + // handles that do draw on it, five further consequences: // - A call needing more scratch than the pool holds grows it, and the growth - // gets rid of the buffer it replaces with a stream-ordered cudaFreeAsync on - // the stream it enqueued on, so the call does not wait for that free and the - // bytes come back when the stream reaches it. Where the device has no - // stream-ordered allocator that call reports cudaErrorNotSupported and the - // growth falls back to a host wait and cudaFree, which waits for every - // stream on the device: on such a device that one call blocks until the - // device is idle however asynchronous the rest of this contract makes it -- - // an unbounded wait on work this call did not submit. Unbounded is meant - // literally: if any of that work is itself waiting on something only this - // thread supplies once execute() returns -- a host function it will release, - // a copy it will enqueue next -- the call does not return, and the thread - // that would unblock it is the one inside cudaFree. Which calls grow the - // pool is not knowable from here; see the README. + // has to get rid of the buffer it replaces before it returns, or the bytes + // of every size the pool ever grew to stay resident at once. What that costs + // depends on whether this call is one that synchronizes the stream. + // A call that does -- anything staging through host memory, aliasing an + // output, or running with no caller stream -- queues a stream-ordered + // cudaFreeAsync and its own synchronization returns the bytes, waiting for + // nothing it did not submit. A call that does not -- the caller-stream, + // device-resident case the rest of this contract is about -- makes a host + // wait on the previous enqueue and a device-wide cudaFree instead, because a + // queued free on a stream nothing ever synchronizes never returns the bytes + // at all. So that one call blocks until the device is idle however + // asynchronous the rest of this contract makes it -- an unbounded wait on + // work this call did not submit. Unbounded is meant literally: if any of + // that work is itself waiting on something only this thread supplies once + // execute() returns -- a host function it will release, a copy it will + // enqueue next -- the call does not return, and the thread that would + // unblock it is the one inside cudaFree. The same wait is also the fallback + // for the first case, on a device with no stream-ordered allocator, where + // cudaFreeAsync reports cudaErrorNotSupported rather than freeing. Which + // calls grow the pool is not knowable from here; see the README. + // - A call that synchronizes the stream waits, through the handoff, for every + // pooled enqueue submitted before it on that device. That is what one buffer + // costs, and it is another unbounded wait: park a host function ahead of a + // pooled enqueue and release it only after a later execute() returns, and + // the later call is the one that does not return. // - Capturing a CUDA graph around this delegate is not supported, with the // option on or off. With it on, a call whose selected stream is capturing is // refused with Error::NotSupported, ahead of every CUDA call it makes that a @@ -156,16 +168,16 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { // capture, which invalidates it under every capture mode. A growth's // allocation invalidates it under every mode but cudaStreamCaptureModeRelaxed, // which permits it but does not record it, leaving a replay pointed at a - // buffer the pool may since have freed; and a growth's stream-ordered free is - // refused under every mode -- a free made while capturing has to be given a - // graph allocation and this is not one -- which drops the growth onto the - // fallback wait and cudaFree that a capture cannot take either. The - // alternative to refusing is a capture that silently comes back invalidated. Only a capture on the - // selected stream is caught. A capture running on any other stream under - // cudaStreamCaptureModeGlobal, or under cudaStreamCaptureModeThreadLocal - // from this thread, is invalidated by the same calls and is not refused, - // because CUDA offers no query for it: do not run a pooled engine while - // capturing anywhere in the process. + // buffer the pool may since have freed; and a growth's disposal of the buffer + // it replaces is prohibited too -- a call that synchronizes the stream is one + // no capture mode permits, so the only disposal a capture could reach is the + // host wait and the device-wide cudaFree, which a capture cannot take outside + // Relaxed either. The alternative to refusing is a capture that silently + // comes back invalidated. Only a capture on the selected stream is caught. A + // capture running on any other stream under cudaStreamCaptureModeGlobal, or + // under cudaStreamCaptureModeThreadLocal from this thread, is invalidated by + // the same calls and is not refused, because CUDA offers no query for it: do + // not run a pooled engine while capturing anywhere in the process. // With the option off nothing refuses, because the check is inside the // pooled path, and a call can then be captured: one submitted on the // capturing stream through a CallerStreamGuard that binds only non-empty @@ -187,19 +199,25 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { // and the handoff event it still holds are destroyed with the primary // context, and the next call on that device uses both. There is no guard: // do not reset a device this backend has run a pooled engine on. - // - The pool gives a pooled call four error returns an unpooled one never - // makes, on top of Error::NotSupported for the capture above. Claiming the - // device's scratch answers Error::Internal if its handoff event cannot be - // created, Error::MemoryAllocationFailed if the buffer cannot be allocated - // or grown, and Error::InvalidState if the wait ordering this call behind - // the previous enqueue fails; installing the buffer answers + // - The pool adds six failure points to a call, on top of Error::NotSupported + // for the capture above. Only one of them has a code of its own: + // Error::Internal, which a pooled call returns when the device's handoff + // event cannot be created and an unpooled call never returns at all. The + // other five reuse codes an unpooled call already returns from elsewhere in + // execute(), so the code alone does not say the pool was involved -- the log + // line does, and every one of them logs at Error first. Claiming the + // device's scratch answers Error::MemoryAllocationFailed if the buffer + // cannot be allocated or grown and Error::InvalidState if the wait ordering + // this call behind the previous enqueue fails; installing the buffer answers // Error::InvalidState if TensorRT refuses it, and recording this call's // enqueue for the next claimant answers the same. Error::InvalidProgram - // means the fallback disposal's host wait failed, which leaks the retired - // buffer rather than freeing it under an enqueue that may still be reading - // it. Every one of them logs at Error first. The last two are the only ones - // reached after this call's enqueue is submitted, and both synchronize the - // stream before returning, so neither leaves engine work in flight. + // means the disposal's host wait failed, which leaks the retired buffer + // rather than freeing it under an enqueue that may still be reading it. The + // last two are the only ones reached after this call's enqueue is submitted, + // and both synchronize the stream before returning, so neither leaves engine + // work in flight -- and the record failure makes that wait with the device's + // pool lock still held, which is the one place a pooled call holds it across + // a host wait. // 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 86577b6a045..76335de1bb8 100644 --- a/cpp/src/torch_tensorrt/executorch/README.md +++ b/cpp/src/torch_tensorrt/executorch/README.md @@ -121,8 +121,11 @@ 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 -requirement any call on that device has asked for: +instead backs a device's contexts from one buffer, grown to the largest +requirement any call on that device has asked for. Which contexts: those created +while the option was on, less any whose engine needs no scratch under any shape, +which are left out of the pool entirely. The two exceptions are spelled out +below. ```cpp #include @@ -154,8 +157,11 @@ bullet of the caller-stream contract above: engines sharing a buffer do not run concurrently on the device. The pool never shrinks, so the largest scratch it was ever asked for stays allocated until the process exits. -A call that asks for nothing is handed whatever the pool already holds, at its -current size, so it never grows it. Two of the things a zero from the per-shape +A call that asks for nothing is handed whatever buffer the pool already holds and +never grows it. It is not handed the pool's size with it: what the context is +told it owns is the figure this call asked for, so a zero becomes the one-byte +minimum below and TensorRT is told the buffer is one byte, whatever the pool's +capacity is. Two of the things a zero from the per-shape query can mean reach this point, and are indistinguishable where it is read: the shapes bound to this call need none -- an empty batch inside a profile that admits one -- or the query failed. Neither wants a buffer of its own. @@ -197,47 +203,68 @@ out of the pool entirely, so it takes no per-device lock and does not serialize against the engines that do. The buffer grows when a call asks for more than every call before it did, and the -growth has to get rid of the buffer it replaces. It queues that free rather than -making it: `cudaFreeAsync` on the stream the growing call enqueued on, with the -per-device lock already dropped. The call returns without waiting for it, and the -bytes come back when the stream reaches the free. Ordering is what makes this -safe, not a wait: every pooled call makes its stream wait on the device's handoff -event before it enqueues and records on that event afterwards, so a free queued on -this call's stream sits behind every enqueue that ever used the retired buffer, -and behind this call's own, which uses the new one. - -`cudaFreeAsync` takes a pointer `cudaMalloc` returned; the pool does not move to -the stream-ordered allocator to use it. Measured on CUDA 13.0 with driver 13.0, on -a 256 MiB `cudaMalloc` buffer with a host function parked on an unrelated stream: -`cudaFreeAsync` returned `cudaSuccess` in 0.2 ms without waiting for the parked -work, the bytes returned to the device once the streams drained, and -`compute-sanitizer --tool memcheck` reported no errors; `cudaFree` on the same -buffer with the same work parked did not return until that work did. - -Where a device has no stream-ordered allocator (`cudaDevAttrMemoryPoolsSupported`), -`cudaFreeAsync` reports `cudaErrorNotSupported` instead of freeing, and the -backend falls back to what it did before: a host wait on the handoff event, then -`cudaFree`. That path is the one to know about. `cudaFree` waits for everything -queued on the device, not only for the enqueues that used the buffer -- measured -with a host function parked on a stream neither the pool nor the growing engine -had ever used, the `cudaMalloc` for the new buffer and the wait on the retired one -each returned at once and the `cudaFree` returned after 3000 ms, exactly when that -host function was released. Both are still made with the per-device lock dropped, -so neither holds up another pooled engine on the device: two pooled calls -serialize at submission, as the caller-stream contract above says, and a growth is -not an exception to it. - -On that fallback the wait has no upper bound, and a caller can turn it into one -that never ends. Anything queued anywhere on the device is enough to hold it, so -if some stream is waiting on work only this thread will submit -- a host function -it will release after `execute()` returns, a copy it will enqueue next -- the -growing call does not return and the thread that would unblock it is the one -inside `cudaFree`. This is not hypothetical: the backend's own tests deadlocked on -it once, when a change of test order turned a case that parks a host function on -its own stream into the one that grew the pool. A program on such a device that -parks work like that has to keep the pool's growths away from it, and the -paragraph below on how often the pool grows is what says whether a run order can -do that. +growth has to get rid of the buffer it replaces *before `execute()` returns*. +Not disposing of it inside the call is not an option: the pool grows monotonically, +so a run that let each retired buffer outlive its growth would end up holding the +sum of every size the pool was ever grown to rather than the largest of them. + +There are two ways to get rid of it, and which one a call uses is decided by +whether that call synchronizes the stream before it returns -- that is, whether it +stages anything through host memory, aliases an output, or runs with no caller +stream. Both are made with the per-device lock dropped, so neither holds up +another pooled engine on the device: two pooled calls serialize at submission, as +the caller-stream contract above says, and a growth is not an exception to it. + +**A call that synchronizes** queues the free: `cudaFreeAsync` on the stream it +enqueued on. Ordering is what makes that safe, not a wait -- every pooled call +makes its stream wait on the device's handoff event before it enqueues and records +on that event afterwards, so a free queued on this call's stream sits behind every +enqueue that ever used the retired buffer, and behind this call's own, which uses +the new one. The call's own `cudaStreamSynchronize` is then what returns the +bytes, and nothing waits for work the call did not submit. Measured on an A100 +with CUDA 13.0 and driver 13.0, with a host function parked for 1500 ms on a stream +neither the pool nor the growing engine had ever used: around 10 ms, against the +1505 ms the disposal below took, which returned exactly when that host function was +released. + +**A call that does not synchronize** -- the caller-stream, device-resident case, +which is the one this option exists for -- makes a host wait on the handoff event +and then a device-wide `cudaFree`. It cannot queue the free, because a queued free +does not return the bytes when the stream reaches it. Measured on the same +machine: `cudaFreeAsync` on a `cudaMalloc`'d pointer returns `cudaSuccess` and +defers the free, but the bytes come back at the next `cudaStreamSynchronize` of +that stream and at no point before it. A stream `cudaStreamQuery` reports as +drained still holds them, `cudaMemGetInfo` does not move, and a `cudaMalloc` of the +same size fails with out of memory until the synchronize. On this path nothing +this backend does will ever synchronize that stream, so replaying the pool's own +growth sequence on one stream -- 8, 16, 24 and 32 GiB, no synchronize between +them -- four growths that all succeed with `cudaFree` fail on the fourth with out +of memory with `cudaFreeAsync`. + +The same host wait and `cudaFree` are also the fallback for the first case: where a +device has no stream-ordered allocator (`cudaDevAttrMemoryPoolsSupported`), +`cudaFreeAsync` reports `cudaErrorNotSupported` instead of freeing. + +**That `cudaFree` is the thing to know about**, because it is what every growth on +the asynchronous path pays. It waits for everything queued on the device, not only +for the enqueues that used the buffer, so on such a growth the asynchronous return +this backend otherwise promises does not happen: the call blocks until the device +is idle. The wait has no upper bound, and a caller can +turn it into one that never ends. Anything queued anywhere on the device is enough +to hold it, so if some stream is waiting on work only this thread will submit -- a +host function it will release after `execute()` returns, a copy it will enqueue +next -- the growing call does not return and the thread that would unblock it is +the one inside `cudaFree`. This is not hypothetical: the backend's own tests +deadlocked on it once, when a change of test order turned a case that parks a host +function on its own stream into the one that grew the pool. A program that parks +work like that has to keep the pool's growths away from it, and the paragraph +below on how often the pool grows is what says whether a run order can do that. + +Sharing one buffer has an unbounded wait of its own, growth or no growth. A call +that synchronizes waits, through the handoff, for every pooled enqueue submitted +before it on that device. Park a host function ahead of one pooled call's enqueue +and release it only once a later pooled call has returned, and the later call is +the one that never returns. What an engine answers when asked how much it needs is decided when it is built, not when it runs. The builder's `kRUNTIME_ACTIVATION_RESIZE_10_10` preview feature @@ -271,14 +298,18 @@ and a null graph, long after the call that caused it. A growth adds more of the same. Its `cudaMalloc` returns `cudaErrorStreamCaptureUnsupported` and invalidates the capture under `Global` and `ThreadLocal`, and under `Relaxed` is permitted but runs uncaptured, so a replayed graph would use whatever buffer was installed when -it was captured -- which by then the pool may have freed. Its `cudaFreeAsync` is -refused under all three modes, with `cudaErrorInvalidValue` and without -invalidating the capture -- measured; a free made while capturing has to be given -a graph allocation, and the pool's buffers come from `cudaMalloc` -- and the -growth then falls back to a host wait on the handoff event and a `cudaFree`, and -those two return `cudaErrorStreamCaptureUnsupported` and invalidate the capture -under `Global` and `ThreadLocal`, and under `Relaxed` are permitted but run -uncaptured. The backend checks the selected stream and returns instead, and it +it was captured -- which by then the pool may have freed. So does its disposal of +the buffer it replaces. A call that synchronizes the stream is one no capture mode +permits, so the only disposal a capture could reach is the host wait on the handoff +event and the `cudaFree`, and those two return `cudaErrorStreamCaptureUnsupported` +and invalidate the capture under `Global` and `ThreadLocal`, and under `Relaxed` +are permitted but run uncaptured. The queued free would be no way round it either: +a `cudaFreeAsync` made while capturing is refused under all three modes, with +`cudaErrorInvalidValue` and without invalidating the capture -- measured; a free +made while capturing has to be given a graph allocation, and the pool's buffers +come from `cudaMalloc`. + +The backend checks the selected stream and returns instead, and it checks ahead of everything a capture cannot take -- not merely ahead of the pool's own calls, since the wait on a previous enqueue and the `cudaMalloc` that grows a host-input staging buffer come before those and, outside `Relaxed`, would @@ -350,9 +381,9 @@ asynchronous return described in the caller-stream contract above rests on it. And moving the pool's check out of the pooled path so that capture were refused both ways would refuse the calls above, which do capture, for contexts that never touch the pool -- a behaviour change for callers who have not turned the option -on. So the option-off path stays as it is, and this section, the option key's -documentation in `TensorRTBackend.h` and the refusal message all say that capture -is unsupported rather than pointing a caller at it. +on. So the option-off path stays as it is, and this section, the `execute()` +contract in `TensorRTBackend.h` and the refusal message all say that capture is +unsupported rather than pointing a caller at it. ### cudaDeviceReset() is not survivable diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h index e895d9c4a0d..42ebb707767 100644 --- a/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h @@ -17,10 +17,7 @@ #include #include #include -#include -#include #include -#include namespace torch_tensorrt { namespace executorch_backend { @@ -28,7 +25,7 @@ namespace executorch_backend { // 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 { - // Nothing in the normal path destroys this event; only reset_for_testing below + // Nothing in the normal path destroys this event; only the test-only reset // hands it to a disposer that does. cudaEvent_t event = nullptr; bool pending = false; // an enqueue against the buffer has been recorded on `event` @@ -60,13 +57,20 @@ struct SharedScratchDevice { SharedScratchMarker marker; }; +class SharedScratchPool; + +// Defined in SharedScratchPoolReset.h; see the friend declaration inside +// SharedScratchPool for why the body is not here. +template +std::size_t reset_shared_scratch_pool_slots(SharedScratchPool& pool, Dispose dispose); + // 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, so holding it couples every device to -// whoever holds it: it covers the lookup, and in reset_for_testing the snapshot +// whoever holds it: it covers the lookup, and in the test-only reset the snapshot // that empties the slots, and nothing else. No CUDA call is made under it. class SharedScratchPool { public: @@ -75,71 +79,23 @@ class SharedScratchPool { return devices_[device_id]; } - // How long a reset waits, in total, for the device locks. A test calls this + // How long a reset waits, in total, for the device locks. A test calls the reset // between cases with nothing claimed, so anything held is a leak the run has // already failed over; the wait only has to outlast a claim still winding down. static constexpr std::chrono::seconds kResetLockWait{5}; - // Test-only. Returns every device's slot to the state it had before anything - // claimed it, handing what the slot held to `dispose(device_id, buffer, event)` - // so the caller can release it. Entries stay in the map, so a reference `get` - // handed out remains valid. - // - // Answers with the number of devices whose lock it could not take within - // kResetLockWait, whose slots it left alone. That case is a leaked claim, which - // is a defect a case here exists to catch -- and taking the locks unconditionally - // would hang on exactly that defect, in a fixture that resets both before and - // after every case, so the run would end in a target timeout with nothing said - // about the cause. The caller reports it instead. - // - // A template, so nothing is emitted for it until something instantiates it, and - // in the released archive nothing does: the only entry point that calls it is - // compiled into the test target alone. See SharedScratchPoolTestHooks.h. - // - // Every slot is emptied under the locks, and what came out of it is disposed of - // afterwards with neither held, because the disposer frees device memory and a - // device-wide free blocks on everything queued on that device -- a parked host - // function included. Under the device's lock that would hold off that device's - // next claimant; under this registry's, every device's. - // - // Nothing here waits for an enqueue: the buffer it frees may still be in use by - // one, and the caller is responsible for there being none. + private: + // The reset is test-only and its body is not here. It frees everything the live + // pool holds without waiting for work in flight against it, which is not a thing + // this header should hand a consumer of the released source package -- and this + // header does ship with those sources. Only the grant of access is here; the + // body is in SharedScratchPoolReset.h, which is testonly and ships nowhere. What + // is left reachable from a release build is the grant, so reaching the pool that + // way means writing the traversal, the locks and the deadline again rather than + // calling something. template - std::size_t reset_for_testing(Dispose dispose) { - std::vector> taken; - std::size_t still_locked = 0; - { - std::lock_guard lk(mu_); - taken.reserve(devices_.size()); - // One deadline for the whole reset rather than one per device, so a run - // with a leaked lock costs the same however many devices the pool has seen. - const auto deadline = std::chrono::steady_clock::now() + kResetLockWait; - for (auto& entry : devices_) { - SharedScratchDevice& dev = entry.second; - bool locked = dev.mu.try_lock(); - while (!locked && std::chrono::steady_clock::now() < deadline) { - std::this_thread::sleep_for(std::chrono::milliseconds(5)); - locked = dev.mu.try_lock(); - } - if (!locked) { - ++still_locked; - continue; - } - std::lock_guard dev_lk(dev.mu, std::adopt_lock); - taken.emplace_back(entry.first, dev.buffer, dev.marker.event); - dev.buffer = nullptr; - dev.capacity = 0; - dev.marker.event = nullptr; - dev.marker.pending = false; - } - } - for (const auto& slot : taken) { - dispose(std::get<0>(slot), std::get<1>(slot), std::get<2>(slot)); - } - return still_locked; - } + friend std::size_t reset_shared_scratch_pool_slots(SharedScratchPool& pool, Dispose dispose); - private: std::mutex mu_; std::unordered_map devices_; }; @@ -230,10 +186,13 @@ inline cudaEvent_t shared_scratch_mark_in_flight(SharedScratchDevice& dev) { // doing both is ordered against nothing here. // // The caller has two ways to honour that, and both belong outside `dev.mu`: a -// free queued on a stream already ordered after the event, which is what the -// backend does, or a host wait on the event followed by a device-wide free, which -// is what it falls back to. Under the lock either one makes an unrelated claim on -// this device wait for work it has nothing to do with. +// free queued on a stream already ordered after the event, or a host wait on the +// event followed by a device-wide free. Neither is the one the backend always +// takes. A queued free returns the bytes only at the next synchronize of that +// stream, so the backend queues it on a call that synchronizes before returning +// and makes the host wait on a call that does not; the host wait is also where a +// device has no stream-ordered allocator to queue onto. Under the lock either one +// makes an unrelated claim on this device wait for work it has nothing to do with. struct RetiredScratch { void* buffer = nullptr; cudaEvent_t wait_for = nullptr; diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolReset.h b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolReset.h new file mode 100644 index 00000000000..785b0321b1a --- /dev/null +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolReset.h @@ -0,0 +1,91 @@ +/* + * 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 + +// The body of the shared scratch pool's test-only reset. +// +// It lives here rather than in SharedScratchPool.h because that header ships in +// the released source package and this operation frees everything the live pool +// holds without waiting for work in flight against it. This file is carried by a +// testonly Bazel target, is in neither the packaged include tree nor the packaged +// source set, and is compiled by nothing a release build produces. What +// SharedScratchPool leaves reachable is the friend declaration alone, so a +// consumer of the released sources who wants this has to write it. +// +// Safe only with no claim outstanding and no enqueue in flight against a pooled +// buffer; the caller earns that by synchronizing every stream it submitted on. + +#include "torch_tensorrt/executorch/SharedScratchPool.h" + +#include +#include +#include +#include +#include +#include + +namespace torch_tensorrt { +namespace executorch_backend { + +// Returns every device's slot to the state it had before anything claimed it, +// handing what the slot held to `dispose(device_id, buffer, event)` so the caller +// can release it. Entries stay in the map, so a reference `get` handed out +// remains valid. +// +// Answers with the number of devices whose lock it could not take within +// SharedScratchPool::kResetLockWait, whose slots it left alone. That case is a +// leaked claim, which is a defect a case in the pool's own test exists to catch -- +// and taking the locks unconditionally would hang on exactly that defect, in a +// fixture that resets both before and after every case, so the run would end in a +// target timeout with nothing said about the cause. The caller reports it instead. +// +// Every slot is emptied under the locks, and what came out of it is disposed of +// afterwards with neither held, because the disposer frees device memory and a +// device-wide free blocks on everything queued on that device -- a parked host +// function included. Under the device's lock that would hold off that device's +// next claimant; under the registry's, every device's. +// +// Nothing here waits for an enqueue: the buffer it frees may still be in use by +// one, and the caller is responsible for there being none. +template +std::size_t reset_shared_scratch_pool_slots(SharedScratchPool& pool, Dispose dispose) { + std::vector> taken; + std::size_t still_locked = 0; + { + std::lock_guard lk(pool.mu_); + taken.reserve(pool.devices_.size()); + // One deadline for the whole reset rather than one per device, so a run with a + // leaked lock costs the same however many devices the pool has seen. + const auto deadline = std::chrono::steady_clock::now() + SharedScratchPool::kResetLockWait; + for (auto& entry : pool.devices_) { + SharedScratchDevice& dev = entry.second; + bool locked = dev.mu.try_lock(); + while (!locked && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + locked = dev.mu.try_lock(); + } + if (!locked) { + ++still_locked; + continue; + } + std::lock_guard dev_lk(dev.mu, std::adopt_lock); + taken.emplace_back(entry.first, dev.buffer, dev.marker.event); + dev.buffer = nullptr; + dev.capacity = 0; + dev.marker.event = nullptr; + dev.marker.pending = false; + } + } + for (const auto& slot : taken) { + dispose(std::get<0>(slot), std::get<1>(slot), std::get<2>(slot)); + } + return still_locked; +} + +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp index de6bf3f32c9..bbeaea9837d 100644 --- a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp @@ -8,6 +8,7 @@ #include "torch_tensorrt/executorch/SharedScratchPoolTestHooks.h" #include "torch_tensorrt/executorch/SharedScratchPool.h" +#include "torch_tensorrt/executorch/SharedScratchPoolReset.h" #include @@ -25,25 +26,26 @@ std::size_t shared_scratch_capacity_for_testing(int device_id) { bool reset_shared_scratch_pool_for_testing() { int restore_to = 0; const bool have_current = cudaGetDevice(&restore_to) == cudaSuccess; - const std::size_t still_locked = scratch_pool().reset_for_testing([](int device_id, void* buffer, cudaEvent_t event) { - if (buffer == nullptr && event == nullptr) { - return; - } - // cudaFree and cudaEventDestroy both act on the current device, and a slot is - // keyed by the device its buffer came from. - if (cudaSetDevice(device_id) != cudaSuccess) { - return; - } - if (buffer != nullptr) { - (void)cudaFree(buffer); - } - if (event != nullptr) { - (void)cudaEventDestroy(event); - } - // Clears a non-sticky error so a reset does not leave one for the next call to - // report. A sticky one survives the clear, and no cleanup here recovers it. - (void)cudaGetLastError(); - }); + const std::size_t still_locked = + reset_shared_scratch_pool_slots(scratch_pool(), [](int device_id, void* buffer, cudaEvent_t event) { + if (buffer == nullptr && event == nullptr) { + return; + } + // cudaFree and cudaEventDestroy both act on the current device, and a slot is + // keyed by the device its buffer came from. + if (cudaSetDevice(device_id) != cudaSuccess) { + return; + } + if (buffer != nullptr) { + (void)cudaFree(buffer); + } + if (event != nullptr) { + (void)cudaEventDestroy(event); + } + // Clears a non-sticky error so a reset does not leave one for the next call to + // report. A sticky one survives the clear, and no cleanup here recovers it. + (void)cudaGetLastError(); + }); if (have_current) { (void)cudaSetDevice(restore_to); } diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index a7e5f77de67..bf7d39badbc 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -268,14 +268,17 @@ class SharedScratchClaim { return dev_; } - // Takes ownership of a buffer a growth displaced, to be freed by release() on - // `stream`. `wait_for` is the marker event the enqueues that used it were - // recorded on, or null if none were; `stream` is the one this claim is about to - // enqueue on, which release() needs because the free it makes is ordered on it. - void retire(void* buffer, cudaEvent_t wait_for, cudaStream_t stream) { + // Takes ownership of a buffer a growth displaced, to be freed by release(). + // `wait_for` is the marker event the enqueues that used it were recorded on, or + // null if none were; `stream` is the one this claim is about to enqueue on. + // `stream_is_synchronized_before_return` says whether execute() will + // synchronize that stream before it returns, which is what decides between the + // two disposals release() has -- see there. + void retire(void* buffer, cudaEvent_t wait_for, cudaStream_t stream, bool stream_is_synchronized_before_return) { retired_ = buffer; retired_wait_ = wait_for; stream_ = stream; + stream_is_synchronized_before_return_ = stream_is_synchronized_before_return; } // Drops the lock, and the device pointer with it so device() cannot hand out a @@ -284,30 +287,52 @@ class SharedScratchClaim { // queues on device work another pooled engine on this device has nothing to do // with. // - // The free is stream-ordered: cudaFreeAsync queued on the stream this claim - // enqueued on. That is what keeps a growth from stalling. Measured on CUDA 13.0 - // with driver 13.0: cudaFreeAsync takes a pointer cudaMalloc returned, does not - // block on unrelated work parked on another stream, defers the free until the - // stream reaches it, and gives the bytes back to the device rather than holding - // them in the allocator's pool -- while cudaFree on the same buffer, with the - // same work parked, does not return until that work does. + // Which disposal, because a stream-ordered free does not return the bytes when + // the stream reaches it. Measured on an A100 with CUDA 13.0 and driver 13.0: + // cudaFreeAsync accepts a pointer cudaMalloc returned and defers the free, but + // the bytes come back at the next cudaStreamSynchronize of that stream and at + // no point before it. A stream cudaStreamQuery reports as drained still holds + // them -- a cudaMalloc of the same size fails with out of memory until the + // synchronize, and cudaMemGetInfo does not move. // - // Ordering, not the wait, is what makes it safe. Every claimant makes its stream - // wait on the marker event before it enqueues, and this claim did so while - // holding the lock, so this stream is already behind every enqueue that ever used - // the retired buffer -- the marker covers all of them, because each records on it - // afterwards. Queuing the free on this stream therefore puts it after the last of - // them, and after this claim's own enqueue, which uses the new buffer. + // So the free is queued only where this call synchronizes the stream itself: // - // cudaFreeAsync needs the device's stream-ordered allocator, which not every - // platform has; where it is missing the call fails rather than freeing, so the - // host wait and the device-wide cudaFree stay as the fallback. On that path a - // growth blocks until the device is idle, which is what the execute() contract - // and the README describe as the platform-dependent case. + // - stream_is_synchronized_before_return_. cudaFreeAsync on the stream this + // claim enqueued on, and execute()'s own cudaStreamSynchronize returns the + // bytes before the call ends. Nothing waits for work this call did not + // submit -- measured with a host function parked for 1500 ms on a stream + // neither the pool nor this engine had ever used, around 10 ms against the + // 1505 ms the synchronous disposal took, which came back when that host + // function was released. // - // Returns false when the fallback's wait failed, which leaves the buffer leaked - // rather than freed under a live enqueue; the caller reports it. Frees on the - // current device, which must still be the buffer's. + // - Otherwise -- the caller-stream, device-resident path that returns with its + // enqueue still running -- nothing this backend does will ever synchronize + // that stream, so a queued free would hold the bytes until the caller + // happened to. Every buffer a run retired would stay resident and peak + // memory would be the sum of every size the pool grew to rather than the + // largest of them: measured, replaying this pool's growth sequence on one + // stream with no synchronize, four growths that all succeed under the + // synchronous disposal fail on the fourth with out of memory. That path + // waits on the marker event and makes a device-wide cudaFree, which has the + // bytes back before execute() returns and costs a host wait for everything + // queued on the device. The execute() contract and the README say so. + // + // Ordering, not a wait, is what makes the queued free safe. Every claimant makes + // its stream wait on the marker event before it enqueues, and this claim did so + // while holding the lock, so this stream is already behind every enqueue that + // ever used the retired buffer -- the marker covers all of them, because each + // records on it afterwards. Queuing the free on this stream therefore puts it + // after the last of them, and after this claim's own enqueue, which uses the new + // buffer. + // + // cudaFreeAsync also needs the device's stream-ordered allocator, which not + // every platform has; where it is missing the call fails rather than freeing and + // the synchronous disposal runs instead, so that one is both the other path's + // disposal and this path's fallback. + // + // Returns false when the synchronous disposal's wait failed, which leaves the + // buffer leaked rather than freed under a live enqueue; the caller reports it. + // Frees on the current device, which must still be the buffer's. bool release() { if (lock_.owns_lock()) { lock_.unlock(); @@ -319,24 +344,44 @@ class SharedScratchClaim { void* const retired = retired_; const cudaEvent_t wait_for = retired_wait_; const cudaStream_t stream = stream_; + const bool stream_is_synchronized = stream_is_synchronized_before_return_; retired_ = nullptr; retired_wait_ = nullptr; stream_ = nullptr; + stream_is_synchronized_before_return_ = false; - const cudaError_t async_err = cudaFreeAsync(retired, stream); - if (async_err == cudaSuccess) { - return true; + if (stream_is_synchronized) { + const cudaError_t async_err = cudaFreeAsync(retired, stream); + if (async_err == cudaSuccess) { + return true; + } + // The error is this call's own, so it is cleared here rather than left for + // the next CUDA call in execute() to report under its own name. + cudaGetLastError(); + if (async_err == cudaErrorNotSupported) { + ET_LOG( + Info, + "TensorRTBackend::execute: device %d has no stream-ordered allocator, so the free of the shared activation scratch buffer a pool growth replaced falls back to a host wait and a device-wide free, which blocks this call until the device is idle", + device_id_); + } else { + // Any other code is a fault on a device that does have the allocator, so + // this does not say the platform lacks one. + ET_LOG( + Info, + "TensorRTBackend::execute: the stream-ordered free of the shared activation scratch buffer a pool growth replaced on device %d returned %s, so it falls back to a host wait and a device-wide free, which blocks this call until the device is idle", + device_id_, + cudaGetErrorString(async_err)); + } } - // Not this device's allocator, then. The error is this call's own, so it is - // cleared here rather than left for the next CUDA call in execute() to report - // under its own name. - cudaGetLastError(); - ET_LOG( - Info, - "TensorRTBackend::execute: the stream-ordered free of the shared activation scratch buffer a pool growth replaced on device %d is not available here (%s); falling back to a host wait and a device-wide free, which blocks this call until the device is idle", - device_id_, - cudaGetErrorString(async_err)); + return dispose_with_a_host_wait(retired, wait_for); + } + private: + // The disposal that has the bytes back before it returns: wait on the host for + // the enqueue that last used the buffer, then free it device-wide. Both waits + // are unbounded -- see the execute() contract -- which is why the queued free is + // preferred wherever it can return the bytes. + bool dispose_with_a_host_wait(void* retired, cudaEvent_t wait_for) { if (wait_for != nullptr) { const cudaError_t wait_err = cudaEventSynchronize(wait_for); if (wait_err != cudaSuccess) { @@ -371,13 +416,13 @@ class SharedScratchClaim { return true; } - private: SharedScratchDevice* dev_ = nullptr; int device_id_ = -1; std::unique_lock lock_; void* retired_ = nullptr; cudaEvent_t retired_wait_ = nullptr; cudaStream_t stream_ = nullptr; + bool stream_is_synchronized_before_return_ = false; }; // What a call needing no activation scratch is given when the pool holds nothing @@ -397,13 +442,16 @@ constexpr size_t kMinPooledScratchBytes = 1; // invalidates the capture under every capture mode. A growth adds more. Its // cudaMalloc invalidates the capture outside cudaStreamCaptureModeRelaxed, and // under Relaxed is permitted but runs uncaptured, leaving a replay pointed at a -// buffer the pool may since have freed. Its cudaFreeAsync of the buffer it -// replaces is refused under all three modes, without invalidating the capture -- -// a free made while capturing has to be given a graph allocation and the pool's -// buffers come from cudaMalloc -- which drops the growth onto the fallback host -// wait and cudaFree, and those two behave like the cudaMalloc. The invalidations -// do not fail cleanly: the caller learns of one only when cudaStreamEndCapture -// hands back an error and a null graph. Refusing names the cause instead. +// buffer the pool may since have freed. Its disposal of the buffer it replaces +// behaves the same way: a call that synchronizes the stream is one no capture mode +// permits, so the only disposal a capture could reach is the host wait and the +// device-wide cudaFree, and both are prohibited outside Relaxed. Even the queued +// free would be no way round it -- a cudaFreeAsync made while capturing has to be +// given a graph allocation and the pool's buffers come from cudaMalloc, so it is +// refused under all three modes, though without invalidating the capture. The +// invalidations do not fail cleanly: +// the caller learns of one only when cudaStreamEndCapture hands back an error and +// a null graph. Refusing names the cause instead. // // execute() calls this ahead of every CUDA call it makes that a capture cannot // take, not just the pool's own: the cudaEventSynchronize on a previous enqueue @@ -464,23 +512,32 @@ Error refuse_pooled_call_on_a_capturing_stream(cudaStream_t stream, int device_i // // The caller must already have refused a capturing `stream`: the handoff's wait // invalidates a capture under every mode, a growth's allocation under every mode -// but Relaxed, and a growth's stream-ordered free of the buffer it replaces is -// refused outright under all three. +// but Relaxed, and a growth's disposal of the buffer it replaces under every mode +// but Relaxed as well. +// +// `stream_is_synchronized_before_return` is execute()'s must_sync, passed down +// because a growth's disposal of the buffer it replaces turns on it; see +// SharedScratchClaim::release(). // // Must be called with `device_id` already current: cudaEventCreateWithFlags and // cudaMalloc both act on the *current* device and nothing in here sets it. -Error claim_shared_scratch(SharedScratchClaim& claim, int device_id, size_t need, cudaStream_t stream, void*& out_ptr) { +Error claim_shared_scratch( + SharedScratchClaim& claim, + int device_id, + size_t need, + cudaStream_t stream, + bool stream_is_synchronized_before_return, + void*& out_ptr) { SharedScratchDevice& dev = claim.hold(device_id); const SharedScratchHandoff handoff = shared_scratch_claim_event(dev, []() -> cudaEvent_t { cudaEvent_t event = nullptr; // Blocking-sync so the host yields instead of busy-spinning. The only host - // wait ever made on this event is SharedScratchClaim::release()'s, on the - // fallback disposal path where cudaFreeAsync is unavailable, and it waits for a - // whole inference; spinning would burn a core for that time and be no faster, - // since it is followed by a device-wide cudaFree. Where cudaFreeAsync is - // available nothing waits on this event from the host at all, and the flag - // costs nothing. + // wait ever made on this event is SharedScratchClaim::release()'s synchronous + // disposal, and it waits for a whole inference; spinning would burn a core for + // that time and be no faster, since it is followed by a device-wide cudaFree. + // Where the free is queued on the stream instead nothing waits on this event + // from the host at all, and the flag costs nothing. if (cudaEventCreateWithFlags(&event, cudaEventDisableTiming | cudaEventBlockingSync) != cudaSuccess) { return nullptr; } @@ -533,10 +590,10 @@ Error claim_shared_scratch(SharedScratchClaim& claim, int device_id, size_t need } // The retired buffer is disposed of at release(), with the device's lock - // dropped, by a free queued on `stream`; see SharedScratchClaim for why that is - // ordered after every enqueue that used it. Nothing here makes a CUDA call that - // blocks on device work under that lock. - claim.retire(retired.buffer, retired.wait_for, stream); + // dropped, either by a free queued on `stream` or by a host wait and a + // device-wide free; see SharedScratchClaim::release() for which and why. + // Nothing here makes a CUDA call that blocks on device work under that lock. + claim.retire(retired.buffer, retired.wait_for, stream, stream_is_synchronized_before_return); out_ptr = buffer; return Error::Ok; @@ -1091,9 +1148,17 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // pinned memory has not read it yet -- measured, every byte the device // received was the value written after the return. The success path already // synchronizes whenever anything was staged (must_sync below); this covers the - // returns between the copy and that point, which would otherwise leave it - // live. It waits only on what this call queued, because must_sync means the - // stream is synchronized before the call ends either way. + // returns between the copy and that point, which would otherwise leave it live. + // + // The wait it makes is a whole-stream one and not a wait on the copy alone: + // cudaStreamSynchronize returns when everything queued on the caller's stream + // has run, including work the caller queued before calling. So an error return + // that staged a host input can block on work this delegate never submitted. + // Waiting for the copy alone would mean recording an event after the last + // staging copy of every call that makes one and waiting on that instead, which + // costs a per-call event record on the success path to narrow a wait only error + // returns make. Staging already implies must_sync, so no successful call reaches + // the destructor with this pending. struct StagedInputDrain { cudaStream_t stream; const bool& staged; @@ -1381,6 +1446,28 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } } + // Whether this call ends by waiting for its own enqueue. Decided here rather + // than at the wait itself, because a growth's disposal of the buffer it + // replaces turns on it -- see SharedScratchClaim::release() -- and the claim is + // made below. Everything it reads is settled by this point: the two staging + // flags and the aliased reflects are set while the bindings are built above, and + // whether a caller stream is active was read at the top. + // + // must_sync = an output is staged to host (the caller reads the D2H result on + // return), an input was staged from host (its async H2D read the caller's host + // buffer, which the caller may reuse once we return), an aliased reflect is + // queued (ExecuTorch's buffer-mutation copy_ reads that EValue after execute() + // returns, so the reflect must complete first), or no caller stream is active + // (preserve the historical "results ready on return" behavior). + // + // Otherwise -- caller stream, all I/O device-resident, no alias -- the engine + // work is left enqueued so it composes with the caller's later GPU work, and a + // completion event is recorded instead so the next execute() and the destructor + // wait before reusing or freeing exec_ctx. + const bool aliased_reflect_pending = !aliased_reflects.empty(); + const bool must_sync = + output_staged_to_host || input_staged_from_host || aliased_reflect_pending || !caller_stream_set; + // ------------------------------------------------------------------ // 4. Back activation scratch with the shared per-device pool // ------------------------------------------------------------------ @@ -1433,7 +1520,8 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // or enqueueV3 refuses it. const size_t scratch_bytes = need == 0 ? kMinPooledScratchBytes : need; void* pool = nullptr; - const Error scratch_err = claim_shared_scratch(scratch_claim, engine->device_id, scratch_bytes, stream, pool); + const Error scratch_err = + claim_shared_scratch(scratch_claim, engine->device_id, scratch_bytes, stream, must_sync, pool); if (scratch_err != Error::Ok) { return scratch_err; } @@ -1455,14 +1543,33 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* return Error::InvalidState; } + // What every early return past this point owes. The enqueue is submitted and + // this handle's completion marker is not armed, so a return that left it running + // would have a later execute() or ~EngineHandle reconfigure or free exec_ctx + // underneath it -- which TensorRT forbids. Where the pool is in play the same + // wait is also what keeps the next claimant from overwriting scratch this + // enqueue is still using. The drain guard's flag is cleared with it, since this + // wait covers a staged input's copy as well. + const auto drain_the_enqueue = [&]() { + (void)cudaStreamSynchronize(stream); + engine->inflight_pending = false; + staged_input_drain.done = true; + }; + // Pairs with claim_shared_scratch: the next claimant waits on this event. if (pooled_scratch) { const Error mark_err = record_shared_scratch_enqueue(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. - (void)cudaStreamSynchronize(stream); - engine->inflight_pending = false; + // Nothing will wait for this enqueue otherwise: the record that would have + // put it on the marker is the call that just failed. + // + // The wait is made with the device's pool lock still held -- the claim is + // released below -- which is the one place a pooled call holds it across a + // host wait, so another pooled engine on this device waits out this + // inference. Releasing first would be worse: it would hand the buffer to a + // claimant with nothing ordering it against the enqueue this call just + // submitted, which is the silent-corruption case the lock exists for. + drain_the_enqueue(); return mark_err; } } @@ -1470,16 +1577,13 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // 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. - // The release is also where a growth disposes of the buffer it replaced, by - // queueing a free on this stream with the lock dropped. + // The release is also where a growth disposes of the buffer it replaced, with + // the lock dropped. if (!scratch_claim.release()) { - // Only the fallback disposal reports a failure, and only from its host wait, - // which for a wait on device work means this device is already in a faulted - // state. The enqueue above is submitted and this handle's completion marker is - // not armed yet, so drain before reporting, or a later call reconfigures - // exec_ctx while that enqueue is still running. - (void)cudaStreamSynchronize(stream); - engine->inflight_pending = false; + // Only the synchronous disposal reports a failure, and only from its host + // wait, which for a wait on device work means this device is already in a + // faulted state. + drain_the_enqueue(); return Error::InvalidProgram; } @@ -1490,33 +1594,15 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* if (cuda_err != cudaSuccess) { ET_LOG( Error, "TensorRTBackend::execute: aliased-output reflect D2D copy failed: %s", cudaGetErrorString(cuda_err)); - // enqueueV3 already submitted engine work to `stream`, and inflight_pending - // is not armed until the end of the happy path -- drain now so a later - // execute() or the destructor never reconfigures/frees exec_ctx while this - // enqueue is still running. - (void)cudaStreamSynchronize(stream); - engine->inflight_pending = false; + drain_the_enqueue(); return Error::InvalidProgram; } } - // The engine work is now in flight on `stream`. Decide whether to wait for it: - // must_sync = an output is staged to host (the caller reads the D2H result on - // return), an input was staged from host (its async H2D read the caller's host - // buffer, which the caller may reuse once we return), or no caller stream is - // active (preserve the historical "results ready on return" behavior). - // Otherwise (caller stream + all I/O device-resident) leave the work enqueued so - // it composes with the caller's later GPU work, and record inflight_event so the - // next execute() and the destructor wait before reusing/freeing exec_ctx. The D2H - // copies live in the must_sync branch: an output staged to host always sets + // The engine work is now in flight on `stream`, and must_sync -- decided above, + // where the scratch claim needed it -- says whether to wait for it. The D2H + // copies live in this branch: an output staged to host always sets // output_staged_to_host, so outputs_needing_copy is empty on the skip path. - // An aliased reflect enqueues the engine's in-place update into the delegate - // output EValue on `stream`; ExecuTorch's buffer-mutation copy_ reads that EValue - // after execute() returns, so the reflect must complete first. A model with - // aliased outputs therefore always syncs here. - const bool aliased_reflect_pending = !aliased_reflects.empty(); - const bool must_sync = - output_staged_to_host || input_staged_from_host || aliased_reflect_pending || !caller_stream_set; if (must_sync) { // Every return from here on is behind the cudaStreamSynchronize below, so // the staging drain has nothing left to do. @@ -1551,11 +1637,9 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } else { cuda_err = cudaEventRecord(engine->inflight_event, stream); if (cuda_err != cudaSuccess) { - // Could not arm the completion marker; drain now so a later execute() or the - // destructor never reconfigures or frees exec_ctx while this enqueue runs. + // Could not arm the completion marker, so nothing downstream would wait. ET_LOG(Error, "TensorRTBackend::execute: cudaEventRecord failed: %s", cudaGetErrorString(cuda_err)); - (void)cudaStreamSynchronize(stream); - engine->inflight_pending = false; + drain_the_enqueue(); return Error::InvalidProgram; } engine->inflight_pending = true; diff --git a/tests/cpp/BUILD b/tests/cpp/BUILD index 827fa2c4099..b8d18994af3 100644 --- a/tests/cpp/BUILD +++ b/tests/cpp/BUILD @@ -62,15 +62,13 @@ test_suite( ], ) +# The suite in the executorch package is the one the workflow runs; this is the +# label for it from here. Listing the tests again would be a second list to keep +# in step, for a name nothing refers to. test_suite( name = "executorch_backend_tests", tests = [ - "//tests/cpp/executorch:test_caller_stream", - "//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", + "//tests/cpp/executorch:executorch_backend_tests", ], ) diff --git a/tests/cpp/executorch/BUILD b/tests/cpp/executorch/BUILD index 86443f894a5..1b8f123b593 100644 --- a/tests/cpp/executorch/BUILD +++ b/tests/cpp/executorch/BUILD @@ -50,11 +50,21 @@ cc_test( ], ) +# The long timeout is what makes the fork case's per-child deadline mean anything: +# it reaps 24 children at 30 s each, so the path where they wedge budgets 720 s +# against Bazel's 300 s default for a target that names neither size nor timeout. +# Two source files because one of them is how the target sees that scratch_pool() +# is one registry across translation units; see the second file's comment. cc_test( name = "test_shared_scratch_pool", - srcs = ["test_shared_scratch_pool.cpp"], + timeout = "long", + srcs = [ + "test_shared_scratch_pool.cpp", + "test_shared_scratch_pool_other_tu.cpp", + ], deps = [ "//cpp:tensorrt_executorch_shared_scratch_pool", + "//cpp:tensorrt_executorch_shared_scratch_pool_reset", "@googletest//:gtest_main", ], ) diff --git a/tests/cpp/executorch/test_shared_scratch_backend.cpp b/tests/cpp/executorch/test_shared_scratch_backend.cpp index 1be6a5faec8..db3e97985b9 100644 --- a/tests/cpp/executorch/test_shared_scratch_backend.cpp +++ b/tests/cpp/executorch/test_shared_scratch_backend.cpp @@ -26,18 +26,32 @@ // itself. Where that variable is unset, a count of the cases that skipped for that // reason is printed at the end of the suite. // -// That variable covers the missing device and nothing else. Three cases skip for a -// second reason: two need the device full and one needs its memory quiet, and -// neither is a state a test can insist on while sharing the device with other -// processes. Those skips stand whether the variable is set or not, so a green -// required-CUDA run says every case ran, not that every case covered what it is -// named for; the skip messages say which did not. +// That variable covers the missing device and nothing else. Five cases skip for a +// second reason: two need the device full, two need its memory quiet, and one +// needs the device's stream-ordered allocator. None of those is a state a test can +// insist on while sharing the device with other processes. Those skips stand +// whether the variable is set or not, so a green required-CUDA run says every case +// ran, not that every case covered what it is named for; the skip messages say +// which did not. +// +// SHARED-DEVICE WARNING: the two cases that need the device full take it to +// essentially zero free bytes for as long as they hold their DeviceMemoryHog -- +// measured, 16 MiB free of 81151 MiB, with a neighbour process getting +// out-of-memory on 256 MiB allocations throughout the window. That is the only way +// to reach the pool's allocation-failure path from inside this process: nothing in +// the delegate takes an allocator a test could substitute, and the other early +// returns on that path need TensorRT or CUDA to fail a call that is correct as +// made. Bazel's exclusive tag keeps other actions in the same build off this +// device; it cannot keep anything else off it. // // The CI invocation in .github/workflows/executorch-build-linux.yml passes that // variable, and the job it sits in asks for a GPU runner and starts its container -// with every GPU attached. What still keeps these cases off most runs -// is the lane gate on the whole ExecuTorch job in ci-linux-x86_64.yml, not the -// runner it would land on. +// with every GPU attached. The ExecuTorch job's own gate in ci-linux-x86_64.yml +// does not keep these cases off an ordinary pull-request push: it sits out only +// when the lane is skip or the backend is RTX, and _decide.yml resolves an +// unlabelled pull_request to lane=fast and backend=standard. What does drop the +// job is the standard channel having been cancelled, which the same workflow +// already carries a comment about. #include "torch_tensorrt/executorch/PooledScratchInstall.h" #include "torch_tensorrt/executorch/SharedScratchPool.h" @@ -447,6 +461,18 @@ class LoadedEngine { return run_with_input(stream, device_in_); } + // The same run with no CallerStreamGuard, which is one of the four things that + // make execute() synchronize the stream before returning. It is the one a memory + // measurement can use: the other three each need a buffer this fixture does not + // otherwise allocate -- a host-backed input or output, which execute() stages + // through a device buffer of its own, or an extra output EValue for an aliased + // output -- and the growth measurement would read those alongside the pool. The + // enqueue goes on cudaStreamPerThread, which is not ordered against any stream + // created cudaStreamNonBlocking. + Error run_on_the_synchronized_path() { + return run_with_input(nullptr, device_in_, /*scope_a_caller_stream=*/false); + } + // The same run with the input bound to caller-owned host memory, which // execute() stages through a device buffer of its own instead of binding // directly. The output stays device-resident, so the only staging is the @@ -488,7 +514,10 @@ class LoadedEngine { } private: - Error run_with_input(cudaStream_t stream, void* in_ptr) { + // With `scope_a_caller_stream` false no CallerStreamGuard is scoped and `stream` + // is unread: execute() then finds no caller stream, enqueues on + // cudaStreamPerThread and synchronizes before returning. + Error run_with_input(cudaStream_t stream, void* in_ptr, bool scope_a_caller_stream = true) { // 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] = {batch_, rows_, cols_}; @@ -502,6 +531,9 @@ class LoadedEngine { EValue* args[2] = {&in_value, &out_value}; BackendExecutionContext exec_context; + if (!scope_a_caller_stream) { + return backend_.execute(exec_context, handle_, Span(args, 2)); + } ::executorch::extension::cuda::CallerStreamGuard guard(stream); return backend_.execute(exec_context, handle_, Span(args, 2)); } @@ -787,6 +819,9 @@ class SharedScratchBackendTest : public ::testing::Test { return dynamic_blob_; } + // Defined below, next to the two cases that call it. + void run_a_growth_and_bound_what_it_cost(bool synchronized_path); + TensorRTBackend backend_; static std::vector blob_; static std::vector scratch_free_blob_; @@ -1029,20 +1064,32 @@ TEST_F(SharedScratchBackendTest, PooledEnginesShareOneActivationScratchAllocatio // --------------------------------------------------------------------------- // Runs a four-times-larger engine after a smaller one to reach the growth path, -// which nothing else in this file does. +// which nothing else in this file does, and bounds what the growth cost the +// device. // -// The bounds cover the second allocation and the disposal of the buffer it -// replaces. That disposal is a cudaFreeAsync queued on the run's stream, so the -// bytes come back only once the stream drains; the synchronize below the run is -// what makes the second reading see them. The bounds say nothing about when the -// disposal happens relative to the return, which is what the two cases further -// down measure. +// The reading is taken the moment execute() returns and *before* any synchronize +// of this test's, which is what makes the upper bound mean something. A growth has +// two disposals to choose from and only one of them has the bytes back by then; a +// synchronize in between hides the difference, because it is exactly what releases +// a queued free. Measured on an A100 with CUDA 13.0: cudaFreeAsync on a +// cudaMalloc'd pointer returns cudaSuccess and defers the free, and the bytes come +// back at the next cudaStreamSynchronize of that stream and at no point before it +// -- a stream cudaStreamQuery reports as drained still holds them. So a growth +// that queued the free on a stream nothing synchronizes costs the whole new buffer +// here rather than the difference, and the upper bound below fails. // // The lower bound also fails if the pool were already large enough for the second // engine, which is how this test could otherwise pass vacuously. The fixture // empties the pool before each test, so it is the smaller engine's run below that // establishes the size the growth has to exceed. -TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItReplaces) { +// +// `synchronized_path` picks which of the two disposals runs, since the choice is +// execute()'s must_sync and not the pool's: a call with no caller stream ends by +// synchronizing and queues the free, a call on a caller stream with device-resident +// I/O does not and makes a host wait and a device-wide free instead. Both have to +// have the bytes back before the call returns, and this is the same measurement +// pointed at each. +void SharedScratchBackendTest::run_a_growth_and_bound_what_it_cost(bool synchronized_path) { 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"; @@ -1050,6 +1097,10 @@ TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItRep cudaStream_t stream = nullptr; ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + const auto run_one = [synchronized_path, stream](LoadedEngine& engine) { + return synchronized_path ? engine.run_on_the_synchronized_path() : engine.run(stream); + }; + // 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. @@ -1062,8 +1113,8 @@ TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItRep 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(run_one(small_priv), Error::Ok); + ASSERT_EQ(run_one(big_priv), Error::Ok); ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); small_expected = small_priv.read_output(); big_expected = big_priv.read_output(); @@ -1073,7 +1124,7 @@ TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItRep 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(run_one(small), Error::Ok); ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); // Loaded before the measurement starts: its weights and its I/O are not part of @@ -1085,27 +1136,33 @@ TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItRep std::size_t before = 0; ASSERT_TRUE(device_bytes_in_use(before)) << "cudaMemGetInfo failed, so this test measured nothing"; const auto measurement_began = std::chrono::steady_clock::now(); - ASSERT_EQ(big.run(stream), Error::Ok); - ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + ASSERT_EQ(run_one(big), Error::Ok); + // No synchronize between the growth and this reading: see the note above. Both + // the allocation and both disposals are host calls execute() makes and returns + // from, so what this reads is settled whether or not the enqueue has finished. std::size_t after = 0; ASSERT_TRUE(device_bytes_in_use(after)) << "cudaMemGetInfo failed, so this test measured nothing"; // A control window of the same length as the measurement, with nothing of this - // test's running in it. On a device this test has to itself, device-wide usage - // is the same at both ends of it; anything else means another process is moving - // memory on the same timescale as the growth, which is what makes the bounds - // below report a figure the pool did not produce. The same length matters: two - // readings taken back to back would sample microseconds against the growth's - // milliseconds and would miss almost everything. It is still a sample of a - // different window, so it narrows the misdiagnosis rather than removing it, and - // the upper bound names the cause as well. Measured with another process cycling - // 256 MiB allocations on this device: of 8 runs, 1 skipped here, 1 failed that - // bound and 1 tripped the direction check below, which skips with the rest now - // rather than failing. The exclusive tag keeps other Bazel actions off this - // device, not other processes. + // test's running in it. It is sampled before the synchronize below for the same + // reason the reading above is: a synchronize releases a free the growth only + // queued, so one taken inside this window would move `settled` away from `after` + // and turn the failure the bounds are meant to report into a skip. On a device + // this test has to itself, device-wide usage is the same at both ends of it; + // anything else means another process is moving memory on the same timescale as + // the growth, which is what makes the bounds below report a figure the pool did + // not produce. The same length matters: two readings taken back to back would + // sample microseconds against the growth's milliseconds and would miss almost + // everything. It is still a sample of a different window, so it narrows the + // misdiagnosis rather than removing it, and the upper bound names the cause as + // well. Measured with another process cycling 256 MiB allocations on this device: + // of 8 runs, 1 skipped here, 1 failed that bound and 1 tripped the direction + // check below, which skips with the rest rather than failing. The exclusive tag + // keeps other Bazel actions off this device, not other processes. std::this_thread::sleep_for(std::chrono::steady_clock::now() - measurement_began); std::size_t settled = 0; ASSERT_TRUE(device_bytes_in_use(settled)) << "cudaMemGetInfo failed, so this test measured nothing"; const bool device_was_quiet = settled == after; + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); // Only something outside this test can make device-wide usage fall across a // growth: the growth allocates the larger buffer before it disposes of the // smaller one. It belongs with the skip below rather than being a failure of its @@ -1120,7 +1177,7 @@ TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItRep // a size below the pool's capacity, since that is what this engine's shapes // need, so a size TensorRT refuses fails here. ASSERT_TRUE(small.fill_output(kSentinel)); - ASSERT_EQ(small.run(stream), Error::Ok); + ASSERT_EQ(run_one(small), Error::Ok); ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); const std::vector small_actual = small.read_output(); const std::vector big_actual = big.read_output(); @@ -1156,8 +1213,25 @@ TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItRep 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 either the buffer it replaced was not freed or another process allocated on this " - "device across the measurement, which the control window above samples for and cannot rule out"; + << "-byte difference, so the buffer it replaced was still resident when execute() returned -- a free the growth " + "only queued, or none at all -- unless another process allocated on this device across the measurement, " + "which the control window above samples for and cannot rule out"; +} + +// The path this option exists for: a caller stream with device-resident I/O, so +// execute() returns with the enqueue still running and never synchronizes the +// stream. A growth here cannot queue its free, because nothing would ever make the +// bytes come back. +TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItReplaces) { + run_a_growth_and_bound_what_it_cost(/*synchronized_path=*/false); +} + +// The other path, where execute() ends by synchronizing the stream and the growth +// queues its free on it. The bytes have to be back by the time the call returns +// here too -- the same bound, and it is the call's own synchronize that satisfies +// it rather than anything this test does. +TEST_F(SharedScratchBackendTest, AGrowthOnASynchronizedCallFreesTheBufferItReplaces) { + run_a_growth_and_bound_what_it_cost(/*synchronized_path=*/true); } // --------------------------------------------------------------------------- @@ -1502,10 +1576,13 @@ TEST_F(SharedScratchBackendTest, AFailedPooledAllocationLeavesTheDeviceLockFree) // cudaStreamWaitEvent refuses with cudaErrorStreamCaptureIsolation and which // invalidates the capture under every capture mode. Left to run, the caller learns // about it only when cudaStreamEndCapture returns an error and a null graph, far -// from the cause. execute() refuses instead. A growth's allocation and free -// invalidate a capture as well, but only outside cudaStreamCaptureModeRelaxed, so -// this test -- relaxed mode, and starting from a pool already large enough -- -// pins the wait rather than them. +// from the cause. execute() refuses instead. A growth adds more of the same: its +// allocation invalidates a capture outside cudaStreamCaptureModeRelaxed, and its +// disposal of the buffer it replaces does too -- a call that synchronizes the +// stream is one no capture mode permits, so the only disposal a capture could +// reach is the host wait and the device-wide free. This test starts from a pool +// already large enough and runs in relaxed mode, so it pins the wait rather than +// either of them. // // The capture is ended either way: a capture left open belongs to the stream, and // destroying that stream would abandon it. @@ -1867,6 +1944,28 @@ class GateRelease { bool released_ = false; }; +// A thread that is joined however the case leaves its scope, so an assertion that +// returns early does not leave one joinable and take the process down with it. +class JoinAtScopeExit { + public: + explicit JoinAtScopeExit(std::thread t) : t_(std::move(t)) {} + JoinAtScopeExit(const JoinAtScopeExit&) = delete; + JoinAtScopeExit& operator=(const JoinAtScopeExit&) = delete; + + ~JoinAtScopeExit() { + join(); + } + + void join() { + if (t_.joinable()) { + t_.join(); + } + } + + private: + std::thread t_; +}; + // 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 @@ -1939,30 +2038,49 @@ TEST_F(SharedScratchBackendTest, ASecondPooledEnqueueWaitsForTheFirstOnAnotherSt // What a growth waits for // --------------------------------------------------------------------------- -// A growth gets rid of the buffer it replaces with a free queued on the stream it -// enqueued on, so it waits for nothing on the host. These two cases park the two -// things a synchronous disposal waits for, which is what the fallback path does -// where cudaFreeAsync is unavailable: everything queued on the device, because -// cudaFree synchronizes it, and the enqueue against the retired buffer, because -// the host waits on the marker event before freeing. Both are parked here, and the -// growing call still has to come back. +// A growth's disposal of the buffer it replaces is one of two things, and each of +// these two cases parks the work one of them waits for. // -// Each measures "came back while the work was still parked" through the gate's -// watchdog, which is the only other thing that can open a gate: if the call had -// waited, the watchdog would have had to open the gate to end the test, and the +// A call that ends by synchronizing the stream queues the free on that stream, so +// it waits for nothing it did not submit -- in particular not for work parked on +// some unrelated stream, which is what the first case pins. A call that does not +// synchronize cannot queue the free, because the bytes would never come back, so +// it makes a host wait on the marker event and a device-wide cudaFree: it does +// wait for the enqueue against the buffer it retires, and the second case pins +// that it waits rather than freeing under a live enqueue. +// +// Each measures what the call came back before through the gate's watchdog, which +// is the only other thing that can open a gate: if a call waited when it should +// not have, the watchdog would have had to open the gate to end the test, and the // case says so. // Nothing the growing engine or the pool ever submitted to runs on the parked // stream, so only a device-wide synchronization has any reason to wait for it. +// This is the queued free's whole point, so the case runs on the path that queues +// it: with no caller stream, execute() ends by synchronizing its own stream, which +// is what returns the bytes. +// +// Gated on the device's stream-ordered allocator, and gated before the gate is +// parked so a device without one costs milliseconds rather than the watchdog +// interval. Where there is none, cudaFreeAsync reports cudaErrorNotSupported, the +// backend correctly falls back to the host wait and the device-wide free, and this +// assertion would report a defect that is not there -- the target is built for +// sbsa as well as x86_64, and nothing in the delegate asks the device about memory +// pools, so the fallback is what a whole platform would take. +// +// That fallback is not left untested by the skip. It is the same host wait and +// device-wide free that every growth on a caller stream makes on this device -- +// AGrowthOnACallerStreamWaitsForTheEnqueueOnTheBufferItRetires and +// ALargerEngineGrowsThePoolAndFreesTheBufferItReplaces both run it and check what +// it costs and what it frees. What a device with a memory pool cannot exercise is +// the one branch that enters it, cudaFreeAsync answering cudaErrorNotSupported. TEST_F(SharedScratchBackendTest, AGrowthDoesNotWaitForUnrelatedWorkQueuedOnTheDevice) { 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 second to be sure of growing the pool"; ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); - cudaStream_t stream = nullptr; cudaStream_t unrelated = nullptr; - ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); ASSERT_EQ(cudaStreamCreateWithFlags(&unrelated, cudaStreamNonBlocking), cudaSuccess); LoadedEngine small; @@ -1973,10 +2091,18 @@ TEST_F(SharedScratchBackendTest, AGrowthDoesNotWaitForUnrelatedWorkQueuedOnTheDe ASSERT_TRUE(big.handle()->shared_scratch); const int device_id = big.handle()->device_id; + int memory_pools = 0; + ASSERT_EQ(cudaDeviceGetAttribute(&memory_pools, cudaDevAttrMemoryPoolsSupported, device_id), cudaSuccess); + if (memory_pools == 0) { + GTEST_SKIP() << "device " << device_id + << " reports no stream-ordered allocator (cudaDevAttrMemoryPoolsSupported = 0), so a growth here " + "correctly takes the host wait and the device-wide free, which does wait for work parked " + "anywhere on the device"; + } + // Without this the larger engine allocates rather than grows, and a growth that // retires nothing disposes of nothing. - ASSERT_EQ(small.run(stream), Error::Ok); - ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + ASSERT_EQ(small.run_on_the_synchronized_path(), Error::Ok); const std::size_t before = shared_scratch_capacity_for_testing(device_id); ASSERT_GT(before, 0u) << "the smaller engine left the pool empty, so the larger one has nothing to retire"; @@ -1985,16 +2111,12 @@ TEST_F(SharedScratchBackendTest, AGrowthDoesNotWaitForUnrelatedWorkQueuedOnTheDe GateRelease gate_release(gate, unrelated); ASSERT_TRUE(big.fill_output(kSentinel)); - const Error growth_error = big.run(stream); + const Error growth_error = big.run_on_the_synchronized_path(); const bool came_back_with_the_device_held = !gate.forced_open.load(); gate_release.release(); - ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); EXPECT_EQ(growth_error, Error::Ok); - EXPECT_TRUE(came_back_with_the_device_held) - << "the growing call did not return until the watchdog released a host function parked on a stream it never " - "submitted to, so its disposal of the retired buffer waits for the whole device"; EXPECT_GT(shared_scratch_capacity_for_testing(device_id), before) << "the pool did not grow, so no buffer was retired and nothing was disposed of"; const std::vector grown_output = big.read_output(); @@ -2002,24 +2124,37 @@ TEST_F(SharedScratchBackendTest, AGrowthDoesNotWaitForUnrelatedWorkQueuedOnTheDe EXPECT_NE(grown_output[0], kSentinel) << "the growing engine did not write its output, so it never reached the engine"; - ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); + EXPECT_TRUE(came_back_with_the_device_held) + << "the growing call did not return until the watchdog released a host function parked on a stream it never " + "submitted to, so its disposal of the retired buffer waits for the whole device"; + ASSERT_EQ(cudaStreamDestroy(unrelated), cudaSuccess); } -// The other thing the old disposal waited for. The smaller engine's own enqueue is -// parked behind a host function, so the marker event it recorded cannot signal -- -// and the retired buffer is the one that enqueue is using. The growth may not free -// it ahead of that enqueue, and it does not have to wait for it either: the free is -// queued on a stream that already waits on the same marker. +// The other disposal. On a caller stream with device-resident I/O nothing will +// ever synchronize the stream, so the growth cannot queue the free -- the bytes +// would stay resident until the caller happened to synchronize, and every buffer a +// run retired would be resident at once. It waits on the marker event and frees +// device-wide instead, so it *does* wait for the enqueue against the buffer it +// retires, and what has to hold is that it waits rather than freeing under it. +// +// The smaller engine's own enqueue is parked behind a host function, and a thread +// releases it a short while after the growing call starts. A growth that waited +// comes back after that release; one that freed the retired buffer under the live +// enqueue comes back before it. // // The smaller engine's output is checked against the same engine run with private // scratch, because it is the one whose buffer was retired underneath it: a free // that landed early would take its activations with it. -TEST_F(SharedScratchBackendTest, AGrowthDoesNotWaitForTheEnqueueOnTheBufferItRetires) { +TEST_F(SharedScratchBackendTest, AGrowthOnACallerStreamWaitsForTheEnqueueOnTheBufferItRetires) { 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 second to be sure of growing the pool"; + // Long enough that the growing call cannot come back after it by accident, short + // enough not to lengthen the suite. + constexpr std::chrono::milliseconds kHoldTheEnqueueFor{500}; + cudaStream_t held = nullptr; cudaStream_t growing = nullptr; ASSERT_EQ(cudaStreamCreateWithFlags(&held, cudaStreamNonBlocking), cudaSuccess); @@ -2056,17 +2191,36 @@ TEST_F(SharedScratchBackendTest, AGrowthDoesNotWaitForTheEnqueueOnTheBufferItRet const std::size_t before = shared_scratch_capacity_for_testing(device_id); ASSERT_GT(before, 0u) << "the smaller engine left the pool empty, so the larger one has nothing to retire"; + // Opens the gate, without the stream synchronization GateRelease::release() also + // makes: that one runs on this thread at scope exit, after the join below, so the + // two never touch GateRelease at once. + std::atomic the_enqueue_was_released{false}; + JoinAtScopeExit opener{std::thread([&] { + std::this_thread::sleep_for(kHoldTheEnqueueFor); + the_enqueue_was_released.store(true); + { + std::lock_guard lock(gate.mu); + gate.open = true; + } + gate.cv.notify_all(); + })}; + const Error growth_error = big.run(growing); - const bool came_back_with_the_enqueue_parked = !gate.forced_open.load(); + const bool waited_for_the_parked_enqueue = the_enqueue_was_released.load(); + opener.join(); gate_release.release(); ASSERT_EQ(cudaStreamSynchronize(held), cudaSuccess); ASSERT_EQ(cudaStreamSynchronize(growing), cudaSuccess); EXPECT_EQ(growth_error, Error::Ok); - EXPECT_TRUE(came_back_with_the_enqueue_parked) - << "the growing call did not return until the watchdog released the enqueue against the buffer it retired, so a " - "growth waits out an inference it has nothing to do with"; + ASSERT_FALSE(gate.forced_open.load()) + << "the watchdog had to open the gate, so the growing call blocked for the whole watchdog interval rather than " + "for the enqueue it had to wait for"; + EXPECT_TRUE(waited_for_the_parked_enqueue) + << "the growing call returned while the enqueue against the buffer it retired was still parked, so it did not " + "wait for that enqueue -- either it freed the buffer under it or it only queued the free, which on this path " + "never gives the bytes back"; EXPECT_GT(shared_scratch_capacity_for_testing(device_id), before) << "the pool did not grow, so no buffer was retired and there was nothing to order the free against"; const std::vector small_output = small.read_output(); @@ -2276,6 +2430,15 @@ TEST_F(SharedScratchBackendTest, TwoThreadsRunningPooledEnginesOnOneDeviceKeepTh // catch. Without one the thread whose partner is stuck spins here until the // target's own timeout, which reports as a killed binary rather than as this // case; with one, the run ends and the assertion below names it. + // + // The loop below stops at the first rendezvous that expires, which is what makes + // the deadline affordable. It is armed per rendezvous -- a whole-case wall clock + // would have to cover 60 real runs and would fail a slow machine instead -- and + // paid at most once per thread, so the worst case is a bounded 30 s rather than + // the 60 x 30 s a loop that carried on would cost against this target's 900 s + // timeout. Stopping also keeps the surviving thread out of the delegate call + // that follows the rendezvous, which its wedged partner may be holding the + // device's pool lock across. constexpr std::chrono::seconds kRendezvousDeadline{30}; std::atomic arrived{0}; std::atomic partner_never_arrived{false}; @@ -2283,7 +2446,7 @@ TEST_F(SharedScratchBackendTest, TwoThreadsRunningPooledEnginesOnOneDeviceKeepTh arrived.fetch_add(1); const auto deadline = std::chrono::steady_clock::now() + kRendezvousDeadline; while (arrived.load() < 2 * (iteration + 1)) { - if (std::chrono::steady_clock::now() >= deadline) { + if (partner_never_arrived.load() || std::chrono::steady_clock::now() >= deadline) { partner_never_arrived.store(true); return; } @@ -2301,6 +2464,9 @@ TEST_F(SharedScratchBackendTest, TwoThreadsRunningPooledEnginesOnOneDeviceKeepTh failures.fetch_add(1); } submit_together(i); + if (partner_never_arrived.load()) { + break; + } if (engine.run(stream) != Error::Ok || cudaStreamSynchronize(stream) != cudaSuccess) { failures.fetch_add(1); continue; @@ -2321,9 +2487,10 @@ TEST_F(SharedScratchBackendTest, TwoThreadsRunningPooledEnginesOnOneDeviceKeepTh ASSERT_EQ(cudaStreamDestroy(second_stream), cudaSuccess); EXPECT_FALSE(partner_never_arrived.load()) - << "one thread waited " << kRendezvousDeadline.count() - << " seconds at the rendezvous without its partner arriving, so a pooled call blocked instead of returning and " - "the runs after that point were not submitted together"; + << "a thread waited out the " << kRendezvousDeadline.count() + << "-second rendezvous deadline without its partner arriving, so a pooled call blocked instead of returning and " + "both threads stopped short of " + << kConcurrentRunsPerThread << " runs"; 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) diff --git a/tests/cpp/executorch/test_shared_scratch_pool.cpp b/tests/cpp/executorch/test_shared_scratch_pool.cpp index 223b103b701..41aa4fe6f83 100644 --- a/tests/cpp/executorch/test_shared_scratch_pool.cpp +++ b/tests/cpp/executorch/test_shared_scratch_pool.cpp @@ -18,6 +18,7 @@ // it, and still makes no CUDA call. #include "torch_tensorrt/executorch/SharedScratchPool.h" +#include "torch_tensorrt/executorch/SharedScratchPoolReset.h" #include "gtest/gtest.h" @@ -46,6 +47,11 @@ namespace torch_tensorrt { namespace executorch_backend { + +// Defined in test_shared_scratch_pool_other_tu.cpp; see +// EveryTranslationUnitSeesOneProcessPool below. +const SharedScratchPool* scratch_pool_seen_from_another_translation_unit(); + namespace { // Fake device allocator: hands out distinct non-null pointers and records every @@ -405,6 +411,21 @@ TEST(SharedScratchPoolRegistry, HandsOutOneStableEntryPerDevice) { EXPECT_EQ(distinct.size(), 512u); } +// scratch_pool() is inline so that the backend and the pool's test hooks, which +// are separate translation units, reach one registry. Nothing else here can see +// that: every case above builds its own SharedScratchPool. Give the accessor +// internal linkage instead and each translation unit gets a registry of its own, +// at which point the reset hook clears one nobody uses, the capacity hook always +// answers zero, and several of the backend suite's pool assertions pass without +// testing anything -- on a GPU host, which is the only place they run. +// +// The second translation unit is test_shared_scratch_pool_other_tu.cpp. +TEST(SharedScratchPoolRegistry, EveryTranslationUnitSeesOneProcessPool) { + EXPECT_EQ(scratch_pool_seen_from_another_translation_unit(), &scratch_pool()) + << "two translation units got different process pools, so scratch_pool() no longer has external linkage and " + "the pool's test hooks operate on a registry the backend never uses"; +} + TEST(SharedScratchPoolRegistry, ResetHandsBackEverySlotAndLeavesItEmpty) { // The backend test fixture runs this between cases, so every case that reads // what the pool holds depends on it. The two things it has to get right are @@ -435,7 +456,7 @@ TEST(SharedScratchPoolRegistry, ResetHandsBackEverySlotAndLeavesItEmpty) { cudaEvent_t event; }; std::vector disposed; - pool.reset_for_testing([&](int device_id, void* buffer, cudaEvent_t event) { + reset_shared_scratch_pool_slots(pool, [&](int device_id, void* buffer, cudaEvent_t event) { disposed.push_back({device_id, buffer, event}); }); @@ -493,7 +514,8 @@ TEST(SharedScratchPoolRegistry, ResetHandsBackTheEventOfASlotThatNeverGotABuffer ASSERT_EQ(dev.buffer, nullptr) << "the allocation did not fail, so this slot is not the state under test"; std::vector> disposed; - pool.reset_for_testing([&](int, void* buffer, cudaEvent_t slot_event) { disposed.emplace_back(buffer, slot_event); }); + reset_shared_scratch_pool_slots( + pool, [&](int, void* buffer, cudaEvent_t slot_event) { disposed.emplace_back(buffer, slot_event); }); ASSERT_EQ(disposed.size(), 1u) << "the reset skipped a slot holding an event and no buffer, so nothing ever destroys " "that event"; @@ -526,7 +548,7 @@ TEST(SharedScratchPoolRegistry, ResetLeavesEveryEntryWhereItWas) { } int disposed = 0; - pool.reset_for_testing([&disposed](int, void*, cudaEvent_t) { ++disposed; }); + reset_shared_scratch_pool_slots(pool, [&disposed](int, void*, cudaEvent_t) { ++disposed; }); ASSERT_EQ(disposed, kDevices); int moved = 0; @@ -583,7 +605,7 @@ TEST(SharedScratchPoolRegistry, ResetDisposesWithNoLockHeld) { // Read inside the disposer: the claim completes once the reset returns either // way, so only what was true while the disposer ran tells the two apart. bool claimed_during_dispose = false; - pool.reset_for_testing([&](int, void*, cudaEvent_t) { + reset_shared_scratch_pool_slots(pool, [&](int, void*, cudaEvent_t) { disposing.store(true); const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2); while (!claimed.load() && std::chrono::steady_clock::now() < deadline) { @@ -617,10 +639,16 @@ TEST(SharedScratchPoolRegistry, ResetReportsALockedSlotRatherThanWaitingForIt) { SharedScratchPool pool; FakeAllocator alloc; + // Two unlocked slots against one locked one, so the number the reset answers + // with -- devices it could not take -- differs from the number it cleared. With + // one of each the two are both 1 and a reset that reported the wrong one of them + // would pass. SharedScratchDevice& leaked = pool.get(0); SharedScratchDevice& ordinary = pool.get(1); + SharedScratchDevice& also_ordinary = pool.get(2); ASSERT_NE(call(leaked, alloc, 1024), nullptr); ASSERT_NE(call(ordinary, alloc, 2048), nullptr); + ASSERT_NE(call(also_ordinary, alloc, 4096), nullptr); std::atomic holding{false}; std::atomic release_it{false}; @@ -639,7 +667,7 @@ TEST(SharedScratchPoolRegistry, ResetReportsALockedSlotRatherThanWaitingForIt) { int disposed = 0; const auto started = std::chrono::steady_clock::now(); - const std::size_t still_locked = pool.reset_for_testing([&](int, void*, cudaEvent_t) { ++disposed; }); + const std::size_t still_locked = reset_shared_scratch_pool_slots(pool, [&](int, void*, cudaEvent_t) { ++disposed; }); const auto took = std::chrono::steady_clock::now() - started; release_it.store(true); holder.join(); @@ -647,11 +675,12 @@ TEST(SharedScratchPoolRegistry, ResetReportsALockedSlotRatherThanWaitingForIt) { EXPECT_EQ(still_locked, 1u) << "the reset did not report the one device whose lock it could not take"; EXPECT_LT(took, kHolderDeadline) << "the reset returned only once the lock was released, so it waited for a claim " "that a leak would never release"; - EXPECT_EQ(disposed, 1) << "the reset disposed of " << disposed - << " slots: it should hand back the unlocked one and leave the locked one alone"; + EXPECT_EQ(disposed, 2) << "the reset disposed of " << disposed + << " slots: it should hand back the two unlocked ones and leave the locked one alone"; EXPECT_NE(leaked.buffer, nullptr) << "the reset emptied a slot whose lock it never held, so it handed the disposer a " "buffer a live claimant is still using"; EXPECT_EQ(ordinary.buffer, nullptr) << "one locked slot stopped the reset clearing the others"; + EXPECT_EQ(also_ordinary.buffer, nullptr) << "one locked slot stopped the reset clearing the others"; } TEST(SharedScratchPoolRegistry, AGrowthOnOneDeviceDoesNotBlockAClaimOnAnother) { @@ -803,9 +832,12 @@ constexpr std::chrono::seconds kChildDeadline{30}; // Reaps `child`, killing it if it overruns `kChildDeadline`; returns false when // it had to. Without a deadline a wedged child parks the parent here until the // whole target times out, which reports as a timeout on the binary rather than as -// this case failing, and this target sets no timeout of its own. An interrupted -// wait is retried, since a signal can arrive at any point and says nothing about -// the child. +// this case failing. The deadline only delivers that if the whole reaping loop +// fits inside the target's timeout -- kTeardownChildren children at kChildDeadline +// each is 720 s -- which is why the BUILD file gives this target the long (900 s) +// timeout rather than leaving it Bazel's 300 s default. An interrupted wait is +// retried, since a signal can arrive at any point and says nothing about the +// child. bool reap_child(pid_t child, int& status) { const auto deadline = std::chrono::steady_clock::now() + kChildDeadline; for (;;) { @@ -818,8 +850,9 @@ bool reap_child(pid_t child, int& status) { } if (std::chrono::steady_clock::now() >= deadline) { kill(child, SIGKILL); - // Reaped even so: a zombie left behind would be inherited by the next fork - // in the loop as a child that never reports. + // Reaped even so. A zombie is not inherited by anything -- it stays this + // process's child until this process reaps it or exits -- so what leaving it + // costs is a process-table entry per overrun, held for the rest of the run. while (waitpid(child, &status, 0) == -1 && errno == EINTR) { } return false; diff --git a/tests/cpp/executorch/test_shared_scratch_pool_other_tu.cpp b/tests/cpp/executorch/test_shared_scratch_pool_other_tu.cpp new file mode 100644 index 00000000000..37c7809d9c7 --- /dev/null +++ b/tests/cpp/executorch/test_shared_scratch_pool_other_tu.cpp @@ -0,0 +1,29 @@ +/* + * 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. + */ + +// A second translation unit, so the pool's test target can see whether +// scratch_pool() hands the same registry to every one of them. +// +// It has to be a separate file: the property is a linkage one, and a single-file +// target cannot observe it. scratch_pool() is inline for exactly this reason -- +// the backend and the pool's test hooks are different translation units and both +// have to reach the one registry -- and with internal linkage instead each would +// get a registry of its own, the reset hook would clear one nobody uses and the +// capacity hook would always answer zero, which the backend suite's pool +// assertions would then pass without testing anything. + +#include "torch_tensorrt/executorch/SharedScratchPool.h" + +namespace torch_tensorrt { +namespace executorch_backend { + +const SharedScratchPool* scratch_pool_seen_from_another_translation_unit() { + return &scratch_pool(); +} + +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/third_party/cuda/BUILD b/third_party/cuda/BUILD index ed98f4f3c69..49869f4f0f0 100644 --- a/third_party/cuda/BUILD +++ b/third_party/cuda/BUILD @@ -17,14 +17,14 @@ config_setting( ], ) +# Every header in the toolkit, and the only place the glob is written: the library +# targets below take their headers from here. One catch-all pattern rather than a +# list of extensions, because `include/**/*` already matches every .h, .hpp and +# .inl under it. Keeping the headers in a target of their own is the shape +# third_party/cublas and third_party/cudnn/archive already use. cc_library( name = "cuda_headers", - hdrs = glob([ - "include/**/*.h", - "include/**/*.hpp", - "include/**/*.inl", - "include/**/*", - ]), + hdrs = glob(["include/**/*"]), includes = ["include/"], ) @@ -41,13 +41,7 @@ cc_library( "lib64/libcudart.so", ], }), - hdrs = glob([ - "include/**/*.h", - "include/**/*.hpp", - "include/**/*.inl", - "include/**/*", - ]), - includes = ["include/"], + deps = [":cuda_headers"], ) cc_library( @@ -82,14 +76,8 @@ cc_library( allow_empty = True, ), }), - hdrs = glob([ - "include/**/*.h", - "include/**/*.hpp", - "include/**/*.inl", - "include/**/*", - ]), - includes = ["include/"], linkopts = ["-Wl,-rpath,lib/"], + deps = [":cuda_headers"], ) cc_library( @@ -108,12 +96,6 @@ cc_library( allow_empty = True, ), }), - hdrs = glob([ - "include/**/*cublas*.h", - "include/**/*.hpp", - "include/**/*.inl", - "include/**/*", - ]), - includes = ["include/"], linkopts = ["-Wl,-rpath,lib/"], + deps = [":cuda_headers"], ) From c1004de600a322d1b92b8fd856bb5add296de601 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Tue, 8 Sep 2026 13:52:09 -0700 Subject: [PATCH 12/13] fix(executorch): decide a growth's disposal where it happens, not where it is claimed `SharedScratchClaim` picked between its two disposals on a flag recorded when the buffer was retired: whether `execute()` would synchronize the stream before returning. That is a promise about what the caller has still to do, and the returns between the claim and the release break it. A failed `enqueueV3` and a failed record of an enqueue already submitted both leave the destructor to queue the free with no `cudaStreamSynchronize` of that stream still to come -- the record failure most clearly, since it drains the stream and only then returns, so the destructor queues the free after the last synchronization the call makes. Measured in the delegate, with the enqueue's return forced to fail after a growth: a bail-out cost 134217728 bytes of device memory, the whole freshly grown buffer, against the 100663296-byte difference in requirement it costs when the disposal returns the retired buffer. Measured at the CUDA level on an A100 with CUDA 13.0 and driver 13.0, the bytes do not come back later either: a `cudaFreeAsync`'d 4 GiB block is still resident after `cudaStreamDestroy`, after another stream's `cudaStreamSynchronize`, after `cudaDeviceSynchronize` and after `cudaMemPoolTrimTo(pool, 0)`. A caller stream is the caller's to destroy, so a failed pooled growth could hold the largest buffer the pool ever had for the life of the process. `release()` now takes the answer as an argument, so it is supplied by the one caller that still knows it -- and the destructor, which is by definition the bail-out, answers it false and takes the host wait. The flag, and the parameter that carried it down through `claim_shared_scratch`, are gone; a return added between the claim and the release cannot leave a promise behind it. Both earlier defects in this path on this branch were the same mistake -- a disposal choosing for a call it could not see the end of -- which is why the choice moves to the point of use rather than the condition being narrowed once more. No test reaches this. The two returns that carry a retired buffer need TensorRT or CUDA to fail a call that is correct as made, and the in-process injections tried for it are each worse than what they would cover: a fabricated or destroyed `cudaEvent_t` segfaults rather than answering `cudaErrorInvalidResourceHandle`, a misaligned device binding gets past `setTensorAddress` and `enqueueV3` and then raises a sticky asynchronous fault that poisons every teardown after it, and a foreign-device stream or event needs a second GPU and so a conditional skip. What carries the fix is that the state cannot be built: with the parameter gone there is nothing to record at the claim and so nothing for a later return to leave stale. A test target that `#include`s `TensorRTBackend.cpp` would give the class a named case, at the cost of a Bazel split so the test does not also link the backend library; not done here. Also in the pooled path: - The pool's own failed `cudaEventCreateWithFlags`, `cudaStreamWaitEvent`, `cudaMalloc` and `cudaEventRecord` now clear the CUDA error where they make it, as every other site in this file already did. Left pending, the next CUDA call any caller makes on that thread collects the pool's error as its own. `AFailedPooledAllocationLeavesTheDeviceLockFreeAndNothingPending` pins it. - `reset_shared_scratch_pool_slots` retries the slots it has not taken instead of spending the whole budget on the first leaked lock and giving every slot after it a single `try_lock`. Measured before the change, a leaked lock and a busy neighbour cost the reset two slots and left the neighbour's buffer in place; the budget is still one deadline for the whole reset. CI: - The GPU suite runs last in the build script rather than ahead of `verify-executorch-reference-runner.sh`. Running the workflow's own script text with a `bazel` whose `test` fails, the reference-runner check -- the only thing in CI that unpacks `libtorchtrt.tar.gz` and builds the packaged CMake backend out of it -- never ran; it now does, and the script still exits 3. The wheel upload is not recovered and the comment says so: `linux-test.yml` guards that step on `upload-artifact` alone with no `always()`, so any failure in this script still skips it. - `--test_summary=detailed`, so Bazel names every case that skipped. With `--test_output=errors` a passing target prints nothing, and five cases in the backend suite skip for reasons the CUDA requirement does not cover. Tests: - `AGrowthOnACallerStreamDisposesWithTheDeviceLockDropped` restores the coverage a rewrite dropped: nothing else takes the device's pool lock while a growth is inside `execute()`, so moving the unlock below the disposal was an unkilled mutation. It fails with the unlock moved. - The backend fixture collects the skips the missing-device counter does not see and lists them by name and reason at teardown. - The growth measurement's control window is quiet within what the bounds themselves absorb rather than to the byte, which threw the measurement away for a single byte of movement. - `AFailedPooledCallDrainsTheHostInputCopyItQueued` has the caller overwrite its buffer in the thread that made the call, the moment the call returns, so what the device ends up holding decides the case rather than a fixed sleep. A build with no drain now fails on the device's contents. Documentation: - The installed contract's sentence about other CUDA delegates sharing the guard is back beside the caller-stream paragraph it qualifies. - The contract and the README said `cudaErrorNotSupported` was the only route to the blocking fallback; the code takes it on any code `cudaFreeAsync` reports, and on an error return after a growth. Two places still carrying the narrower rule go with them: the pool header's account of the two disposals, and the comment on `AGrowthDoesNotWaitForUnrelatedWorkQueuedOnTheDevice`, which called that one code the only branch into the fallback. - The bail-out is a return after the growth, not one before the enqueue: one of the two is a failed record of an enqueue already submitted. The claim's destructor, the installed contract and the README each said "between the growth and the enqueue"; all three now name both returns, and say that the record failure drains the stream and then returns, so the disposal falls after the last synchronization it makes. - The allocation-failure case now says what it does not reach. A failed allocation retires nothing, so it drives the destructor's lock drop and not its disposal, and nothing in the suite drives that disposal at all. - The README said neither disposal holds up another pooled engine on the device, 34 lines above saying the device-wide free waits for everything queued on it. What dropping the lock buys is that the disposal is not serialized behind the pool's own lock. - `SharedScratchPool`'s friend grant is a template, so it admits any definition of that name a consumer writes, including one that erases entries live references point into. The header says so rather than reading as a barrier. - `test_shared_scratch_pool_other_tu.cpp` said internal linkage would leave the backend suite's pool assertions passing vacuously. Six of them across five cases would fail; what the separate translation unit buys is a case that names the cause. 6 targets, 89 -> 91 cases, 0 skipped, on an A100 with `--nocache_test_results` and `TORCHTRT_EXECUTORCH_REQUIRE_CUDA=1`. --- .github/workflows/executorch-build-linux.yml | 60 ++-- .../executorch/TensorRTBackend.h | 14 +- cpp/src/torch_tensorrt/executorch/README.md | 23 +- .../executorch/SharedScratchPool.h | 18 +- .../executorch/SharedScratchPoolReset.h | 47 +-- .../executorch/TensorRTBackend.cpp | 88 +++--- .../test_shared_scratch_backend.cpp | 271 ++++++++++++++++-- .../executorch/test_shared_scratch_pool.cpp | 86 ++++++ .../test_shared_scratch_pool_other_tu.cpp | 10 +- 9 files changed, 502 insertions(+), 115 deletions(-) diff --git a/.github/workflows/executorch-build-linux.yml b/.github/workflows/executorch-build-linux.yml index 1e9923a4d7d..33fb1071e8e 100644 --- a/.github/workflows/executorch-build-linux.yml +++ b/.github/workflows/executorch-build-linux.yml @@ -101,12 +101,33 @@ jobs: python -m pip wheel --no-build-isolation --no-deps --wheel-dir dist py/torch-tensorrt-executorch-runtime # this is to build the libtorchtrt.tar.gz bazel build //:libtorchtrt --compilation_mode opt --config=linux + executorch_cmake_location="$(bazel query @executorch//:executorch/CMakeLists.txt --output=location)" + export EXECUTORCH_SOURCE_DIR="$(dirname "${executorch_cmake_location%%:*}")" + export EXECUTORCH_ROOT="${EXECUTORCH_SOURCE_DIR}" + # this is to verify the end user's workflow + python -m pip install pyyaml "executorch>=1.4.1,<1.5" + python examples/torchtrt_executorch_example/export_static_shape.py \ + --model_path="${RUNNER_TEMP}/torchtrt-reference-runner.pte" + .github/scripts/verify-executorch-reference-runner.sh \ + "${RUNNER_TEMP}/torchtrt-reference-runner.pte" # Run the ExecuTorch backend C++ unit tests, which are otherwise only - # built, never executed. These tests link TensorRT and CUDA, but Bazel's - # cc_import ships the unversioned libnvinfer.so and libcudart.so while the - # loader asks for the versioned sonames, so add the directories holding - # those to LD_LIBRARY_PATH. Append rather than replace: the toolchain - # entries already present are still needed. + # built, never executed. + # + # Last in this script, and not next to the Bazel build above, because + # verify-executorch-reference-runner.sh between the two is the only check + # anywhere in CI that the release archive still carries every header the + # packaged CMake backend needs, and that the backend still builds out of it. + # `set -e` ends the script at the first failure, so a GPU suite placed above + # that check trades it away whenever a device is missing or flaky. Run last, + # a failure here no longer costs that check. It still costs the wheel: + # linux-test.yml guards its upload step on `inputs.upload-artifact` alone, + # with no always(), so any failure in this script skips the upload. + # + # These tests link TensorRT and CUDA, but Bazel's cc_import ships the + # unversioned libnvinfer.so and libcudart.so while the loader asks for the + # versioned sonames, so add the directories holding those to + # LD_LIBRARY_PATH. Append rather than replace: the toolchain entries + # already present are still needed. bazel_external="$(bazel info output_base)/external" for _soname in libnvinfer.so libcudart.so; do _dir="$( @@ -119,22 +140,23 @@ jobs: fi export LD_LIBRARY_PATH="${_dir}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" done + # --test_summary=detailed names every case that skipped, in Bazel's own + # summary rather than in the test log, which --test_output=errors prints + # nothing of for a target that passes. Cases in that suite skip for reasons + # the variable below does not cover -- its own coverage note lists them -- + # and without those names a run in which the pool's growth and + # allocation-failure paths never executed looks like one in which they did. + # # TORCHTRT_EXECUTORCH_REQUIRE_CUDA turns a case that skips for want of a - # CUDA device into a failure. This job runs in a container started with - # all GPUs attached, so a skip here means the device went missing rather - # than that the runner never had one -- and without the variable the - # backend suite would skip every case, exit zero and be reported as a - # passing target that covered nothing. + # CUDA device into a failure. Every row of this job's matrix is a CUDA row + # -- filter-matrix.py drops any whose desired_cuda is not a TensorRT CUDA + # version -- and linux-test.yml starts the container with `--gpus all` for + # those, so a skip here means the device went missing rather than that the + # runner never had one. Without the variable the backend suite would skip + # every case, exit zero and be reported as a passing target that covered + # nothing. bazel test //tests/cpp/executorch:executorch_backend_tests \ --compilation_mode opt --config=linux --test_output=errors \ + --test_summary=detailed \ --test_env=LD_LIBRARY_PATH \ --test_env=TORCHTRT_EXECUTORCH_REQUIRE_CUDA=1 - executorch_cmake_location="$(bazel query @executorch//:executorch/CMakeLists.txt --output=location)" - export EXECUTORCH_SOURCE_DIR="$(dirname "${executorch_cmake_location%%:*}")" - export EXECUTORCH_ROOT="${EXECUTORCH_SOURCE_DIR}" - # this is to verify the end user's workflow - python -m pip install pyyaml "executorch>=1.4.1,<1.5" - python examples/torchtrt_executorch_example/export_static_shape.py \ - --model_path="${RUNNER_TEMP}/torchtrt-reference-runner.pte" - .github/scripts/verify-executorch-reference-runner.sh \ - "${RUNNER_TEMP}/torchtrt-reference-runner.pte" diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index abf9ee8fd5f..c9bbe6ec577 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -126,6 +126,8 @@ 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. + // Other CUDA delegates sharing the same guard may instead synchronize before + // returning, so do not assume results are ready on return from this one. // 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 @@ -153,8 +155,14 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { // enqueue next -- the call does not return, and the thread that would // unblock it is the one inside cudaFree. The same wait is also the fallback // for the first case, on a device with no stream-ordered allocator, where - // cudaFreeAsync reports cudaErrorNotSupported rather than freeing. Which - // calls grow the pool is not knowable from here; see the README. + // cudaFreeAsync reports cudaErrorNotSupported rather than freeing, and on + // any other code it reports, since a queued free that did not happen is not + // a free. It is also what a call that fails after the growth takes -- a + // refused enqueue, or a failed record of one already submitted -- whichever + // kind of call it is, because by the time such a call disposes of the + // buffer it has no synchronization of that stream left to make, so a queued + // free's bytes would never come back. Which calls grow the pool is not + // knowable from here; see the README. // - A call that synchronizes the stream waits, through the handoff, for every // pooled enqueue submitted before it on that device. That is what one buffer // costs, and it is another unbounded wait: park a host function ahead of a @@ -218,8 +226,6 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { // work in flight -- and the record failure makes that wait with the device's // pool lock still held, which is the one place a pooled call holds it across // a host wait. - // 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( ::executorch::runtime::BackendExecutionContext& context, ::executorch::runtime::DelegateHandle* handle, diff --git a/cpp/src/torch_tensorrt/executorch/README.md b/cpp/src/torch_tensorrt/executorch/README.md index 76335de1bb8..7e40c5efefd 100644 --- a/cpp/src/torch_tensorrt/executorch/README.md +++ b/cpp/src/torch_tensorrt/executorch/README.md @@ -211,9 +211,23 @@ sum of every size the pool was ever grown to rather than the largest of them. There are two ways to get rid of it, and which one a call uses is decided by whether that call synchronizes the stream before it returns -- that is, whether it stages anything through host memory, aliases an output, or runs with no caller -stream. Both are made with the per-device lock dropped, so neither holds up -another pooled engine on the device: two pooled calls serialize at submission, as -the caller-stream contract above says, and a growth is not an exception to it. +stream. Both are made with the per-device lock dropped, so neither is serialized +behind the pool's own lock and neither holds the next claimant off while it runs: +two pooled calls serialize at submission, as the caller-stream contract above says, +and a growth is not an exception to it. That is not the same as holding nobody up, +and the paragraph below on the device-wide free says what it does hold up. + +Which of the two a call uses is settled where the buffer is disposed of and not +where it is claimed, because it is a claim about what the call has still to do. A +call that fails after the growth -- a refused `enqueueV3`, or a failed record of an +enqueue already submitted -- takes the host wait whichever kind of call it is. The +failed record does drain the stream, but it drains and then returns, so the +disposal falls after the last synchronization that call makes; the refused enqueue +makes none at all. A free queued on a stream that is never synchronized again does +not come back: not at `cudaStreamDestroy`, not at `cudaDeviceSynchronize`, not from +another stream, and not at `cudaMemPoolTrimTo`. Measured on an A100 with CUDA 13.0, +a 4 GiB block queued that way and then abandoned with the stream is gone for the +life of the process, on a device with tens of gigabytes still free. **A call that synchronizes** queues the free: `cudaFreeAsync` on the stream it enqueued on. Ordering is what makes that safe, not a wait -- every pooled call @@ -243,7 +257,8 @@ of memory with `cudaFreeAsync`. The same host wait and `cudaFree` are also the fallback for the first case: where a device has no stream-ordered allocator (`cudaDevAttrMemoryPoolsSupported`), -`cudaFreeAsync` reports `cudaErrorNotSupported` instead of freeing. +`cudaFreeAsync` reports `cudaErrorNotSupported` instead of freeing -- and on any +other code it reports, since a queued free that did not happen is not a free. **That `cudaFree` is the thing to know about**, because it is what every growth on the asynchronous path pays. It waits for everything queued on the device, not only diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h index 42ebb707767..9fb538036bb 100644 --- a/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h @@ -93,6 +93,15 @@ class SharedScratchPool { // is left reachable from a release build is the grant, so reaching the pool that // way means writing the traversal, the locks and the deadline again rather than // calling something. + // + // It is an open door and not a lock: the grant is on a template, so it admits + // any definition of that name a consumer writes in this namespace, not only the + // one shipped alongside it, and what such a definition reaches is both members + // below. That includes erasing entries a live `get` reference points into, which + // the stability rule above forbids. C++ offers no narrower grant here -- a + // non-template friend admits a definition of its own just the same -- so the + // alternatives are putting the body back in this header, which is what moving it + // out was for, or making the members public, which is worse. template friend std::size_t reset_shared_scratch_pool_slots(SharedScratchPool& pool, Dispose dispose); @@ -190,9 +199,12 @@ inline cudaEvent_t shared_scratch_mark_in_flight(SharedScratchDevice& dev) { // event followed by a device-wide free. Neither is the one the backend always // takes. A queued free returns the bytes only at the next synchronize of that // stream, so the backend queues it on a call that synchronizes before returning -// and makes the host wait on a call that does not; the host wait is also where a -// device has no stream-ordered allocator to queue onto. Under the lock either one -// makes an unrelated claim on this device wait for work it has nothing to do with. +// and makes the host wait everywhere else: on a call that does not synchronize, +// on a call that fails after the growth and so has no synchronization left to +// make, and wherever the queued free is refused -- a device with no +// stream-ordered allocator to queue onto, or any other code cudaFreeAsync +// reports. Under the lock either one makes an unrelated claim on this device wait +// for work it has nothing to do with. struct RetiredScratch { void* buffer = nullptr; cudaEvent_t wait_for = nullptr; diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolReset.h b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolReset.h index 785b0321b1a..8b96d0ce2d2 100644 --- a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolReset.h +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolReset.h @@ -27,6 +27,7 @@ #include #include #include +#include #include namespace torch_tensorrt { @@ -44,6 +45,13 @@ namespace executorch_backend { // fixture that resets both before and after every case, so the run would end in a // target timeout with nothing said about the cause. The caller reports it instead. // +// The deadline covers the whole reset rather than each device, so a leaked lock +// costs the same however many devices the pool has seen. What a leak must not do +// is spend the budget on behalf of the slots after it: the retry below sweeps +// every slot it has not taken yet on each pass, so a slot whose claim is merely +// winding down is tried again inside the window it is free, rather than once at +// whatever instant a leaked neighbour leaves in the budget. +// // Every slot is emptied under the locks, and what came out of it is disposed of // afterwards with neither held, because the disposer frees device memory and a // device-wide free blocks on everything queued on that device -- a parked host @@ -59,27 +67,34 @@ std::size_t reset_shared_scratch_pool_slots(SharedScratchPool& pool, Dispose dis { std::lock_guard lk(pool.mu_); taken.reserve(pool.devices_.size()); - // One deadline for the whole reset rather than one per device, so a run with a - // leaked lock costs the same however many devices the pool has seen. const auto deadline = std::chrono::steady_clock::now() + SharedScratchPool::kResetLockWait; + std::vector> not_taken_yet; + not_taken_yet.reserve(pool.devices_.size()); for (auto& entry : pool.devices_) { - SharedScratchDevice& dev = entry.second; - bool locked = dev.mu.try_lock(); - while (!locked && std::chrono::steady_clock::now() < deadline) { - std::this_thread::sleep_for(std::chrono::milliseconds(5)); - locked = dev.mu.try_lock(); + not_taken_yet.emplace_back(entry.first, &entry.second); + } + while (!not_taken_yet.empty()) { + std::vector> still_to_try; + for (const auto& slot : not_taken_yet) { + SharedScratchDevice& dev = *slot.second; + if (!dev.mu.try_lock()) { + still_to_try.push_back(slot); + continue; + } + std::lock_guard dev_lk(dev.mu, std::adopt_lock); + taken.emplace_back(slot.first, dev.buffer, dev.marker.event); + dev.buffer = nullptr; + dev.capacity = 0; + dev.marker.event = nullptr; + dev.marker.pending = false; } - if (!locked) { - ++still_locked; - continue; + not_taken_yet.swap(still_to_try); + if (not_taken_yet.empty() || std::chrono::steady_clock::now() >= deadline) { + break; } - std::lock_guard dev_lk(dev.mu, std::adopt_lock); - taken.emplace_back(entry.first, dev.buffer, dev.marker.event); - dev.buffer = nullptr; - dev.capacity = 0; - dev.marker.event = nullptr; - dev.marker.pending = false; + std::this_thread::sleep_for(std::chrono::milliseconds(5)); } + still_locked = not_taken_yet.size(); } for (const auto& slot : taken) { dispose(std::get<0>(slot), std::get<1>(slot), std::get<2>(slot)); diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index bf7d39badbc..fa1c4f69462 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -250,9 +250,26 @@ class SharedScratchClaim { SharedScratchClaim(const SharedScratchClaim&) = delete; SharedScratchClaim& operator=(const SharedScratchClaim&) = delete; ~SharedScratchClaim() { + // A claim that still holds a retired buffer here was never released, so + // execute() returned somewhere between the claim and the release below -- a + // refused enqueue, or a failed record of one already submitted -- and no + // synchronization of this stream is still to come. The record failure does + // drain the stream, but it drains and then returns, so a free queued here + // would be queued after the last synchronization the call makes; the refused + // enqueue makes none at all. Such a free would come back at nothing: measured + // on an A100 with CUDA 13.0 and driver 13.0, a cudaFreeAsync'd 4 GiB block is + // still resident after the stream is destroyed, after cudaDeviceSynchronize + // and after another stream's synchronize -- and a caller stream is the + // caller's to destroy, so those bytes are gone for the life of the process. A + // bail-out therefore takes the disposal that has the bytes back before it + // returns, whatever the call would have done had it run to the end. Passing + // that in here rather than storing it at retire() is what keeps the two in + // step: a return added between the claim and the release cannot leave a + // promise to synchronize behind it. + // // The return is the caller's to report, and on this path there is no caller // left to report it to: an execute() that returned early has already failed. - (void)release(); + (void)release(/*the_call_synchronizes_the_stream=*/false); } SharedScratchDevice& hold(int device_id) { @@ -271,14 +288,10 @@ class SharedScratchClaim { // Takes ownership of a buffer a growth displaced, to be freed by release(). // `wait_for` is the marker event the enqueues that used it were recorded on, or // null if none were; `stream` is the one this claim is about to enqueue on. - // `stream_is_synchronized_before_return` says whether execute() will - // synchronize that stream before it returns, which is what decides between the - // two disposals release() has -- see there. - void retire(void* buffer, cudaEvent_t wait_for, cudaStream_t stream, bool stream_is_synchronized_before_return) { + void retire(void* buffer, cudaEvent_t wait_for, cudaStream_t stream) { retired_ = buffer; retired_wait_ = wait_for; stream_ = stream; - stream_is_synchronized_before_return_ = stream_is_synchronized_before_return; } // Drops the lock, and the device pointer with it so device() cannot hand out a @@ -295,9 +308,13 @@ class SharedScratchClaim { // them -- a cudaMalloc of the same size fails with out of memory until the // synchronize, and cudaMemGetInfo does not move. // - // So the free is queued only where this call synchronizes the stream itself: + // So the free is queued only where this call synchronizes the stream itself, + // which `the_call_synchronizes_the_stream` says. It is passed in here rather + // than recorded at retire() because it is a claim about what the caller has + // still to do, and only the caller, at the point it releases, knows that it is + // going to: the destructor answers it false for exactly that reason. // - // - stream_is_synchronized_before_return_. cudaFreeAsync on the stream this + // - the_call_synchronizes_the_stream. cudaFreeAsync on the stream this // claim enqueued on, and execute()'s own cudaStreamSynchronize returns the // bytes before the call ends. Nothing waits for work this call did not // submit -- measured with a host function parked for 1500 ms on a stream @@ -333,7 +350,7 @@ class SharedScratchClaim { // Returns false when the synchronous disposal's wait failed, which leaves the // buffer leaked rather than freed under a live enqueue; the caller reports it. // Frees on the current device, which must still be the buffer's. - bool release() { + bool release(bool the_call_synchronizes_the_stream) { if (lock_.owns_lock()) { lock_.unlock(); } @@ -344,13 +361,11 @@ class SharedScratchClaim { void* const retired = retired_; const cudaEvent_t wait_for = retired_wait_; const cudaStream_t stream = stream_; - const bool stream_is_synchronized = stream_is_synchronized_before_return_; + // Clearing the pointer is what stops the destructor's release from disposing + // of the same buffer twice. retired_ = nullptr; - retired_wait_ = nullptr; - stream_ = nullptr; - stream_is_synchronized_before_return_ = false; - if (stream_is_synchronized) { + if (the_call_synchronizes_the_stream) { const cudaError_t async_err = cudaFreeAsync(retired, stream); if (async_err == cudaSuccess) { return true; @@ -422,7 +437,6 @@ class SharedScratchClaim { void* retired_ = nullptr; cudaEvent_t retired_wait_ = nullptr; cudaStream_t stream_ = nullptr; - bool stream_is_synchronized_before_return_ = false; }; // What a call needing no activation scratch is given when the pool holds nothing @@ -515,19 +529,9 @@ Error refuse_pooled_call_on_a_capturing_stream(cudaStream_t stream, int device_i // but Relaxed, and a growth's disposal of the buffer it replaces under every mode // but Relaxed as well. // -// `stream_is_synchronized_before_return` is execute()'s must_sync, passed down -// because a growth's disposal of the buffer it replaces turns on it; see -// SharedScratchClaim::release(). -// // Must be called with `device_id` already current: cudaEventCreateWithFlags and // cudaMalloc both act on the *current* device and nothing in here sets it. -Error claim_shared_scratch( - SharedScratchClaim& claim, - int device_id, - size_t need, - cudaStream_t stream, - bool stream_is_synchronized_before_return, - void*& out_ptr) { +Error claim_shared_scratch(SharedScratchClaim& claim, int device_id, size_t need, cudaStream_t stream, void*& out_ptr) { SharedScratchDevice& dev = claim.hold(device_id); const SharedScratchHandoff handoff = shared_scratch_claim_event(dev, []() -> cudaEvent_t { @@ -539,6 +543,10 @@ Error claim_shared_scratch( // Where the free is queued on the stream instead nothing waits on this event // from the host at all, and the flag costs nothing. if (cudaEventCreateWithFlags(&event, cudaEventDisableTiming | cudaEventBlockingSync) != cudaSuccess) { + // The pool's own failure, cleared where it is made: the caller is told by + // the return, and leaving it pending would surface it under the name of + // whatever this thread calls next. + cudaGetLastError(); return nullptr; } return event; @@ -557,6 +565,7 @@ Error claim_shared_scratch( Error, "TensorRTBackend::execute: waiting for the enqueue that last used the shared activation scratch failed: %s", cudaGetErrorString(err)); + cudaGetLastError(); return Error::InvalidState; } } @@ -569,6 +578,7 @@ Error claim_shared_scratch( [device_id, first_buffer](size_t bytes) -> void* { void* p = nullptr; if (cudaMalloc(&p, bytes) != cudaSuccess) { + cudaGetLastError(); return nullptr; } ET_LOG( @@ -593,7 +603,7 @@ Error claim_shared_scratch( // dropped, either by a free queued on `stream` or by a host wait and a // device-wide free; see SharedScratchClaim::release() for which and why. // Nothing here makes a CUDA call that blocks on device work under that lock. - claim.retire(retired.buffer, retired.wait_for, stream, stream_is_synchronized_before_return); + claim.retire(retired.buffer, retired.wait_for, stream); out_ptr = buffer; return Error::Ok; @@ -616,6 +626,7 @@ Error record_shared_scratch_enqueue(SharedScratchClaim& claim, cudaStream_t stre Error, "TensorRTBackend::execute: recording the completion event for the shared activation scratch enqueue failed: %s", cudaGetErrorString(err)); + cudaGetLastError(); return Error::InvalidState; } return Error::Ok; @@ -1446,12 +1457,12 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } } - // Whether this call ends by waiting for its own enqueue. Decided here rather - // than at the wait itself, because a growth's disposal of the buffer it - // replaces turns on it -- see SharedScratchClaim::release() -- and the claim is - // made below. Everything it reads is settled by this point: the two staging - // flags and the aliased reflects are set while the bindings are built above, and - // whether a caller stream is active was read at the top. + // Whether this call ends by waiting for its own enqueue. Read twice below -- at + // the release of the scratch claim, whose disposal of a retired buffer turns on + // it, and at the wait itself -- so it is settled once here, which is the first + // point everything it reads is final: the two staging flags and the aliased + // reflects are set while the bindings are built above, and whether a caller + // stream is active was read at the top. // // must_sync = an output is staged to host (the caller reads the D2H result on // return), an input was staged from host (its async H2D read the caller's host @@ -1520,8 +1531,7 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // or enqueueV3 refuses it. const size_t scratch_bytes = need == 0 ? kMinPooledScratchBytes : need; void* pool = nullptr; - const Error scratch_err = - claim_shared_scratch(scratch_claim, engine->device_id, scratch_bytes, stream, must_sync, pool); + const Error scratch_err = claim_shared_scratch(scratch_claim, engine->device_id, scratch_bytes, stream, pool); if (scratch_err != Error::Ok) { return scratch_err; } @@ -1578,8 +1588,12 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // function so the rest of execute() -- the aliased reflects, the D2H copies and // their synchronizations -- does not hold up another engine on this device. // The release is also where a growth disposes of the buffer it replaced, with - // the lock dropped. - if (!scratch_claim.release()) { + // the lock dropped. This is the only place that release is made rather than + // left to the claim's destructor, and it is the only place must_sync is a + // promise this function can still keep: every return past it synchronizes the + // stream when must_sync is set -- the branch below does, and the aliased + // reflects drain the enqueue before returning. + if (!scratch_claim.release(must_sync)) { // Only the synchronous disposal reports a failure, and only from its host // wait, which for a wait on device work means this device is already in a // faulted state. diff --git a/tests/cpp/executorch/test_shared_scratch_backend.cpp b/tests/cpp/executorch/test_shared_scratch_backend.cpp index db3e97985b9..da3ea2ef7f1 100644 --- a/tests/cpp/executorch/test_shared_scratch_backend.cpp +++ b/tests/cpp/executorch/test_shared_scratch_backend.cpp @@ -31,8 +31,12 @@ // needs the device's stream-ordered allocator. None of those is a state a test can // insist on while sharing the device with other processes. Those skips stand // whether the variable is set or not, so a green required-CUDA run says every case -// ran, not that every case covered what it is named for; the skip messages say -// which did not. +// ran, not that every case covered what it is named for. So the suite counts these +// separately from the missing device and lists them by name and reason at the end +// of the file, whether or not the variable is set. That list is in the test log, +// which --test_output=errors prints nothing of for a target that passes; the CI +// invocation therefore also passes --test_summary=detailed, which names every +// skipped case in Bazel's own summary. // // SHARED-DEVICE WARNING: the two cases that need the device full take it to // essentially zero free bytes for as long as they hold their DeviceMemoryHog -- @@ -86,6 +90,7 @@ #include #include #include +#include #include namespace torch_tensorrt { @@ -744,6 +749,7 @@ class SharedScratchBackendTest : public ::testing::Test { } static void TearDownTestSuite() { + report_the_second_reason_skips(); if (skipped_for_no_device_ == 0) { return; } @@ -768,6 +774,7 @@ class SharedScratchBackendTest : public ::testing::Test { int device_count = 0; if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count == 0) { ++skipped_for_no_device_; + this_case_skipped_for_no_device_ = true; if (cuda_device_is_required()) { FAIL() << "no CUDA device, and " << kRequireCudaEnvVar << " says this run must have one. Every case in this file needs a device, so without one the binary " @@ -796,6 +803,7 @@ class SharedScratchBackendTest : public ::testing::Test { } void TearDown() override { + record_a_second_reason_skip(); set_shared_scratch(backend_, false); // Also here, so a buffer this test grew is not still resident while the next // one measures device-wide memory. @@ -822,6 +830,62 @@ class SharedScratchBackendTest : public ::testing::Test { // Defined below, next to the two cases that call it. void run_a_growth_and_bound_what_it_cost(bool synchronized_path); + // A case can also skip for a reason the device being present says nothing about + // -- the device would not fill, its memory would not stay still, it has no + // stream-ordered allocator. Those skips are the difference between a green + // target and a run that covered what the cases are named for, and with + // --test_output=errors a passing target prints not one of them. So they are + // collected here and printed at the end whether or not the CUDA requirement is + // armed, which is the case the missing-device banner cannot cover. + void record_a_second_reason_skip() { + if (this_case_skipped_for_no_device_ || !::testing::Test::IsSkipped()) { + return; + } + const ::testing::TestInfo* const info = ::testing::UnitTest::GetInstance()->current_test_info(); + if (info == nullptr) { + return; + } + std::string reason; + const ::testing::TestResult* const result = info->result(); + for (int i = 0; result != nullptr && i < result->total_part_count(); ++i) { + const ::testing::TestPartResult& part = result->GetTestPartResult(i); + if (part.type() == ::testing::TestPartResult::kSkip) { + reason = part.summary(); + break; + } + } + // The reason a case gives usually wraps, and one line per skip is what makes + // the report below readable. + for (char& c : reason) { + if (c == '\n') { + c = ' '; + } + } + while (!reason.empty() && reason.back() == ' ') { + reason.pop_back(); + } + if (reason.empty()) { + reason = "no reason recorded"; + } + second_reason_skips_.emplace_back(info->name(), reason); + } + + static void report_the_second_reason_skips() { + if (second_reason_skips_.empty()) { + return; + } + const ::testing::TestSuite* const suite = ::testing::UnitTest::GetInstance()->current_test_suite(); + std::fprintf( + stderr, + "[ SKIPPED ] %zu of %d cases in this file ran on a device but did not cover what they are named for, so this " + "target passing does not say the pool's growth, allocation-failure and queued-free paths were exercised:\n", + second_reason_skips_.size(), + suite == nullptr ? static_cast(second_reason_skips_.size()) : suite->total_test_count()); + for (const auto& skip : second_reason_skips_) { + std::fprintf(stderr, "[ SKIPPED ] %s: %s\n", skip.first.c_str(), skip.second.c_str()); + } + } + TensorRTBackend backend_; static std::vector blob_; static std::vector scratch_free_blob_; @@ -834,6 +898,8 @@ class SharedScratchBackendTest : public ::testing::Test { static std::int64_t scratch_free_engine_bytes_; static std::int64_t dynamic_engine_bytes_; static int skipped_for_no_device_; + static std::vector> second_reason_skips_; + bool this_case_skipped_for_no_device_ = false; }; std::vector SharedScratchBackendTest::blob_; @@ -847,6 +913,7 @@ std::size_t SharedScratchBackendTest::dynamic_batch_scratch_bytes_ = 0; std::int64_t SharedScratchBackendTest::scratch_free_engine_bytes_ = -1; std::int64_t SharedScratchBackendTest::dynamic_engine_bytes_ = -1; int SharedScratchBackendTest::skipped_for_no_device_ = 0; +std::vector> SharedScratchBackendTest::second_reason_skips_; // --------------------------------------------------------------------------- // set_option @@ -1147,21 +1214,31 @@ void SharedScratchBackendTest::run_a_growth_and_bound_what_it_cost(bool synchron // reason the reading above is: a synchronize releases a free the growth only // queued, so one taken inside this window would move `settled` away from `after` // and turn the failure the bounds are meant to report into a skip. On a device - // this test has to itself, device-wide usage is the same at both ends of it; - // anything else means another process is moving memory on the same timescale as - // the growth, which is what makes the bounds below report a figure the pool did - // not produce. The same length matters: two readings taken back to back would - // sample microseconds against the growth's milliseconds and would miss almost + // this test has to itself, device-wide usage does not move across it; anything + // else means another process is moving memory on the same timescale as the + // growth, which is what makes the bounds below report a figure the pool did not + // produce. The same length matters: two readings taken back to back would sample + // microseconds against the growth's milliseconds and would miss almost // everything. It is still a sample of a different window, so it narrows the // misdiagnosis rather than removing it, and the upper bound names the cause as - // well. Measured with another process cycling 256 MiB allocations on this device: - // of 8 runs, 1 skipped here, 1 failed that bound and 1 tripped the direction - // check below, which skips with the rest rather than failing. The exclusive tag - // keeps other Bazel actions off this device, not other processes. + // well. The exclusive tag keeps other Bazel actions off this device, not other + // processes. + // + // What counts as quiet is what the bounds below can absorb, rather than exact + // equality, which threw the whole measurement away for a byte. The tighter of + // the two bounds allows scratch_bytes_/2 of slack, so a control window that + // moved by at most half of that cannot turn what they report into something + // else, while anything large enough to matter still skips. The tolerance is for + // allocator granularity and not for a busy device: on a device this suite had to + // itself the window was measured moving 0 bytes against a tolerance of 8 MiB for + // these fixture engines, and a neighbour allocating at the scale they do moves it + // by far more than that. std::this_thread::sleep_for(std::chrono::steady_clock::now() - measurement_began); std::size_t settled = 0; ASSERT_TRUE(device_bytes_in_use(settled)) << "cudaMemGetInfo failed, so this test measured nothing"; - const bool device_was_quiet = settled == after; + const std::size_t quiet_tolerance = scratch_bytes_ / 4; + const std::size_t control_window_moved_by = settled > after ? settled - after : after - settled; + const bool device_was_quiet = control_window_moved_by <= quiet_tolerance; ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); // Only something outside this test can make device-wide usage fall across a // growth: the growth allocates the larger buffer before it disposes of the @@ -1202,9 +1279,10 @@ void SharedScratchBackendTest::run_a_growth_and_bound_what_it_cost(bool synchron "device and the two bounds below would measure it rather than the growth"; } if (!device_was_quiet) { - GTEST_SKIP() << "device-wide memory in use moved from " << after << " to " << settled - << " bytes with nothing of this test's running in between, so another process is allocating on this " - "device and the two bounds below would measure it rather than the growth"; + GTEST_SKIP() << "device-wide memory in use moved from " << after << " to " << settled << " bytes -- " + << control_window_moved_by << " against a " << quiet_tolerance + << "-byte tolerance -- with nothing of this test's running in between, so another process is " + "allocating on this device and the two bounds below would measure it rather than the growth"; } const std::size_t growth_cost = after - before; EXPECT_GE(growth_cost, difference / 2) << "the larger engine cost " << growth_cost << " bytes against a " @@ -1514,7 +1592,16 @@ TEST_F(SharedScratchBackendTest, ExecuteInstallsPooledScratchThroughTheCheckedHe // enqueue and a failed completion record need TensorRT or CUDA to fail a call that // is correct as made. The pool's own allocation can be made to fail by taking the // device's memory first, and it returns through the same destructor as the others. -TEST_F(SharedScratchBackendTest, AFailedPooledAllocationLeavesTheDeviceLockFree) { +// +// What it reaches is the lock the destructor drops, not the buffer it disposes +// of: an allocation that failed retired nothing, so this case leaves the +// destructor with none to free. Nothing in this suite reaches a bail-out that +// still holds a retired buffer, since the two returns that can are the two named +// above. What keeps that path right is the shape of release() rather than a case: +// it takes "does this call still synchronize the stream?" as an argument, so a +// bail-out answers it where the bail-out happens and there is no recorded promise +// for a later return to leave stale. +TEST_F(SharedScratchBackendTest, AFailedPooledAllocationLeavesTheDeviceLockFreeAndNothingPending) { ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); cudaStream_t stream = nullptr; @@ -1529,6 +1616,7 @@ TEST_F(SharedScratchBackendTest, AFailedPooledAllocationLeavesTheDeviceLockFree) << "the pool already holds a buffer, so this call would reuse it rather than allocate"; Error failed_run = Error::Ok; + cudaError_t pending_after_the_failure = cudaSuccess; { DeviceMemoryHog hog; if (!hog.leave_less_free_than(scratch_bytes_)) { @@ -1536,7 +1624,12 @@ TEST_F(SharedScratchBackendTest, AFailedPooledAllocationLeavesTheDeviceLockFree) GTEST_SKIP() << "could not take the device below the " << scratch_bytes_ << " bytes this engine's scratch needs, so its allocation would have succeeded"; } + // The hog clears the last failed allocation of its own, and nothing else has + // made a CUDA call since, so whatever is pending after the run below was left + // by the run. + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << "something before the call under test left a CUDA error pending"; failed_run = pooled.run(stream); + pending_after_the_failure = cudaPeekAtLastError(); } if (failed_run == Error::Ok) { @@ -1556,6 +1649,17 @@ TEST_F(SharedScratchBackendTest, AFailedPooledAllocationLeavesTheDeviceLockFree) ASSERT_TRUE(lock_free) << "a pooled call that returned early left the device's pool lock held, so every later pooled " "call on this device blocks forever"; + // The failure is reported by the return value, so the CUDA error the pool's own + // cudaMalloc left has to be cleared where it was made. Left pending it is + // collected by whichever caller makes the next CUDA call on this thread and + // reported under that call's name -- which is what the capture refusal has its + // own case for, from the other direction. + EXPECT_EQ(pending_after_the_failure, cudaSuccess) + << "the pooled call left '" << cudaGetErrorString(pending_after_the_failure) + << "' pending on this thread after reporting the failure through its return value, so the next CUDA call any " + "caller makes collects the pool's error as its own"; + cudaGetLastError(); + // And the slot is still usable, not merely unlocked: with the memory back, the // next call allocates and runs. ASSERT_TRUE(pooled.fill_output(kSentinel)); @@ -2073,7 +2177,8 @@ TEST_F(SharedScratchBackendTest, ASecondPooledEnqueueWaitsForTheFirstOnAnotherSt // AGrowthOnACallerStreamWaitsForTheEnqueueOnTheBufferItRetires and // ALargerEngineGrowsThePoolAndFreesTheBufferItReplaces both run it and check what // it costs and what it frees. What a device with a memory pool cannot exercise is -// the one branch that enters it, cudaFreeAsync answering cudaErrorNotSupported. +// either branch that enters it from a call that does synchronize: cudaFreeAsync +// answering cudaErrorNotSupported, or answering any other code. TEST_F(SharedScratchBackendTest, AGrowthDoesNotWaitForUnrelatedWorkQueuedOnTheDevice) { ASSERT_GE(big_scratch_bytes_, scratch_bytes_ + kMinMeasurableScratch) << "the two fixture engines ask for " << scratch_bytes_ << " and " << big_scratch_bytes_ @@ -2233,6 +2338,102 @@ TEST_F(SharedScratchBackendTest, AGrowthOnACallerStreamWaitsForTheEnqueueOnTheBu ASSERT_EQ(cudaStreamDestroy(growing), cudaSuccess); } +// Neither growth case above can see whether the disposal is made with the device's +// lock held: both run their growth alone, so a lock held across it holds nothing +// up and both stay green with the unlock moved below the free. This one is what +// covers that. +// +// It matters because the disposal on this path is a host wait and a cudaFree, and +// cudaFree waits for every stream on the device rather than for the ones that +// touched the buffer -- the 1505 ms against a host function parked for 1500 ms on +// an unrelated stream that the execute() contract cites. Under the lock, that wait +// is one every other pooled engine on the device has to sit through, which is the +// serialization the contract says a pooled call does not impose past its own +// submission. +// +// So: park work on an unrelated stream to stall the free, then check the pool lock +// while the growth is still inside it. The capacity read under that same try_lock +// is what stops the case passing on a lock that is free because the growth has not +// started -- past a growth it is above what the smaller engine left, and the growth +// has not returned. +TEST_F(SharedScratchBackendTest, AGrowthOnACallerStreamDisposesWithTheDeviceLockDropped) { + 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 second to be sure of growing the pool"; + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + + cudaStream_t stream = nullptr; + cudaStream_t unrelated = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + ASSERT_EQ(cudaStreamCreateWithFlags(&unrelated, cudaStreamNonBlocking), cudaSuccess); + + LoadedEngine small; + LoadedEngine big; + ASSERT_EQ(small.load(blob(), 32), Error::Ok); + ASSERT_EQ(big.load(big_blob(), 33, kBigRows, kBigCols), Error::Ok); + ASSERT_TRUE(small.handle()->shared_scratch); + ASSERT_TRUE(big.handle()->shared_scratch); + const int device_id = big.handle()->device_id; + + // Without this the larger engine allocates rather than grows, and a growth that + // retires nothing frees nothing. Synchronized, so the growth's wait on the + // marker returns at once and the free is the only thing left to stall it. + ASSERT_EQ(small.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::size_t before = shared_scratch_capacity_for_testing(device_id); + ASSERT_GT(before, 0u) << "the smaller engine left the pool empty, so the larger one has nothing to retire"; + + // Neither engine submits to this stream, so nothing on the growth path but the + // device-wide free has any reason to wait for it. + StreamGate gate; + ASSERT_EQ(cudaLaunchHostFunc(unrelated, hold_stream, &gate), cudaSuccess); + GateRelease gate_release(gate, unrelated); + + std::atomic growth_returned{false}; + Error growth_error = Error::Internal; + JoinAtScopeExit grower{std::thread([&] { + // A caller stream with device-resident I/O, which is the path whose disposal + // is the host wait and the device-wide free. + growth_error = big.run(stream); + growth_returned.store(true); + })}; + + SharedScratchDevice& dev = scratch_pool().get(device_id); + bool lock_free_during_the_disposal = false; + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (!growth_returned.load() && std::chrono::steady_clock::now() < deadline) { + if (dev.mu.try_lock()) { + const std::size_t capacity_now = dev.capacity; + dev.mu.unlock(); + if (capacity_now > before && !growth_returned.load()) { + lock_free_during_the_disposal = true; + break; + } + } + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + const bool still_inside_execute = !growth_returned.load(); + + gate_release.release(); + grower.join(); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(unrelated), cudaSuccess); + + 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(growth_error, Error::Ok); + ASSERT_TRUE(still_inside_execute) + << "the growing call returned before the gate opened, so its free never stalled and there was no window in " + "which to observe the lock"; + ASSERT_GT(shared_scratch_capacity_for_testing(device_id), before) + << "the pool did not grow, so no buffer was retired and no free was made"; + EXPECT_TRUE(lock_free_during_the_disposal) + << "the device's pool lock stayed held for the whole of a stalled growth free, so every other pooled engine on " + "this device waits out a device-wide free it has nothing to do with"; +} + // --------------------------------------------------------------------------- // A failing call and the host-input copy it queued // --------------------------------------------------------------------------- @@ -2253,7 +2454,10 @@ TEST_F(SharedScratchBackendTest, AGrowthOnACallerStreamWaitsForTheEnqueueOnTheBu // Three outcomes are separated at the end, by giving the failing call a // different input from the successful one before it: the sentinel means the copy // took the caller's post-return write, the first call's pattern means it never -// ran at all, and the second call's pattern is the only pass. +// ran at all, and the second call's pattern is the only pass. Which of the three +// happens is decided by where the caller's write falls relative to the copy and +// not by any interval this case picks, so a slow return cannot turn the first +// outcome into the third. TEST_F(SharedScratchBackendTest, AFailedPooledCallDrainsTheHostInputCopyItQueued) { ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); @@ -2306,29 +2510,36 @@ TEST_F(SharedScratchBackendTest, AFailedPooledCallDrainsTheHostInputCopyItQueued ASSERT_EQ(cudaLaunchHostFunc(stream, hold_stream, &gate), cudaSuccess); GateRelease gate_release(gate, stream); + // The caller overwrites its own buffer the moment the call comes back, in the + // thread that made the call, which is what a caller is entitled to do and is + // what decides this case. A build that returns with the copy still queued does + // that write before the gate opens, so the copy reads the sentinel and the + // comparison at the end sees it. A build that drains cannot reach the write + // until the copy has run. Neither outcome depends on how long anything takes: + // the wait below is a bound on this case, not the thing that discriminates. std::atomic call_returned{false}; + std::atomic caller_overwrote_its_buffer{false}; Error failed_run = Error::Ok; std::thread caller([&] { failed_run = pooled.run_from_host_input(stream, host_in); call_returned.store(true); - }); - - // Long enough that a call which does not wait for its own copy has returned; - // the stream cannot move until the gate opens, so nothing else ends the wait. - std::this_thread::sleep_for(std::chrono::milliseconds(500)); - const bool returned_with_the_copy_still_queued = call_returned.load(); - if (returned_with_the_copy_still_queued) { for (std::size_t i = 0; i < pooled.elems(); ++i) { host_in[i] = kSentinel; } + caller_overwrote_its_buffer.store(true); + }); + + // Waited for rather than slept through, so a build that returns early is not + // handed the rest of the interval to finish its write in. Reaching the deadline + // is the passing outcome: the stream cannot move until the gate opens, so a call + // that waits for its own copy is still inside execute() here. + const auto give_up_waiting_at = std::chrono::steady_clock::now() + std::chrono::milliseconds(500); + while (!caller_overwrote_its_buffer.load() && std::chrono::steady_clock::now() < give_up_waiting_at) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); } + const bool returned_with_the_copy_still_queued = call_returned.load(); gate_release.release(); caller.join(); - if (!returned_with_the_copy_still_queued) { - for (std::size_t i = 0; i < pooled.elems(); ++i) { - host_in[i] = kSentinel; - } - } // After the gate, because a free is device-wide and would otherwise block on it. hog.release(); ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); diff --git a/tests/cpp/executorch/test_shared_scratch_pool.cpp b/tests/cpp/executorch/test_shared_scratch_pool.cpp index 41aa4fe6f83..acd446b3025 100644 --- a/tests/cpp/executorch/test_shared_scratch_pool.cpp +++ b/tests/cpp/executorch/test_shared_scratch_pool.cpp @@ -683,6 +683,92 @@ TEST(SharedScratchPoolRegistry, ResetReportsALockedSlotRatherThanWaitingForIt) { EXPECT_EQ(also_ordinary.buffer, nullptr) << "one locked slot stopped the reset clearing the others"; } +// The budget covers the whole reset, so a leaked lock has to be prevented from +// spending it on behalf of the devices it shares the pool with. A device that is +// merely busy -- claims arriving and finishing, which is every device the pool is +// for -- must still be taken, and a single try_lock made once the leak has run the +// clock down only takes it if it happens to be idle at that one instant. +// +// So the neighbour below is locked when the reset starts, free for half a second +// in the middle, and locked again long before the deadline. A reset that retries +// takes it inside that window; one that spends its budget on the leak first and +// then tries once reports it as locked and never hands its buffer to the disposer +// -- which, in the fixture that resets between cases, means the next case starts +// on a device still holding what an earlier one grew. +// +// Which slot is which is decided by the map rather than by this case, and only a +// leak the reset reaches *before* the busy slot can cost it anything. So the order +// is read off a reset of the same three entries first -- entries are never erased, +// so the second walk takes them in the same order -- and the roles are handed out +// from that. +TEST(SharedScratchPoolRegistry, ALeakedLockDoesNotCostTheResetTheSlotsItSharesThePoolWith) { + constexpr std::chrono::milliseconds kNeighbourFreeFrom{200}; + constexpr std::chrono::milliseconds kNeighbourBusyAgainFrom{700}; + const auto holder_deadline = SharedScratchPool::kResetLockWait + std::chrono::seconds{25}; + + SharedScratchPool pool; + FakeAllocator alloc; + // Three slots, so the number the reset answers with and the number it cleared + // stay different figures: one leaked, one busy, one free throughout. + for (const int device_id : {0, 1, 2}) { + ASSERT_NE(call(pool.get(device_id), alloc, 1024), nullptr); + } + std::vector walk_order; + ASSERT_EQ(reset_shared_scratch_pool_slots(pool, [&](int id, void*, cudaEvent_t) { walk_order.push_back(id); }), 0u) + << "a reset of three unlocked slots reported one it could not take"; + ASSERT_EQ(walk_order.size(), 3u); + + SharedScratchDevice& leaked = pool.get(walk_order[0]); + SharedScratchDevice& busy = pool.get(walk_order[1]); + SharedScratchDevice& ordinary = pool.get(walk_order[2]); + ASSERT_NE(call(leaked, alloc, 1024), nullptr); + ASSERT_NE(call(busy, alloc, 2048), nullptr); + ASSERT_NE(call(ordinary, alloc, 4096), nullptr); + + std::atomic both_held{false}; + std::atomic release_it{false}; + const auto started = std::chrono::steady_clock::now(); + std::thread holder([&] { + leaked.mu.lock(); + busy.mu.lock(); + both_held.store(true); + std::this_thread::sleep_until(started + kNeighbourFreeFrom); + busy.mu.unlock(); + std::this_thread::sleep_until(started + kNeighbourBusyAgainFrom); + // Locked again well inside the deadline, so only a reset that looked during + // the window above ever had it. + busy.mu.lock(); + const auto deadline = std::chrono::steady_clock::now() + holder_deadline; + while (!release_it.load() && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + busy.mu.unlock(); + leaked.mu.unlock(); + }); + while (!both_held.load()) { + std::this_thread::yield(); + } + + int disposed = 0; + const std::size_t still_locked = reset_shared_scratch_pool_slots(pool, [&](int, void*, cudaEvent_t) { ++disposed; }); + release_it.store(true); + holder.join(); + + EXPECT_EQ(still_locked, 1u) + << "the reset reported " << still_locked + << " locked slots: only the leaked one is unreachable, and the busy one was free for " + << (kNeighbourBusyAgainFrom - kNeighbourFreeFrom).count() << " ms of the " + << std::chrono::duration_cast(SharedScratchPool::kResetLockWait).count() + << " ms it had to look"; + EXPECT_EQ(disposed, 2) << "the reset disposed of " << disposed + << " slots: the leaked one it must leave alone, and the busy one it must come back to rather " + "than give up on because another device's lock spent the budget"; + EXPECT_EQ(busy.buffer, nullptr) << "one leaked lock left a busy device's buffer resident, so the pool a case is " + "reset for still holds what an earlier one grew"; + EXPECT_EQ(ordinary.buffer, nullptr) << "the reset left a slot nothing ever locked"; + EXPECT_NE(leaked.buffer, nullptr) << "the reset emptied a slot whose lock it never held"; +} + TEST(SharedScratchPoolRegistry, AGrowthOnOneDeviceDoesNotBlockAClaimOnAnother) { SharedScratchPool pool; // One allocator per thread: the two claims share the registry and nothing else. diff --git a/tests/cpp/executorch/test_shared_scratch_pool_other_tu.cpp b/tests/cpp/executorch/test_shared_scratch_pool_other_tu.cpp index 37c7809d9c7..18cebba6194 100644 --- a/tests/cpp/executorch/test_shared_scratch_pool_other_tu.cpp +++ b/tests/cpp/executorch/test_shared_scratch_pool_other_tu.cpp @@ -13,8 +13,14 @@ // the backend and the pool's test hooks are different translation units and both // have to reach the one registry -- and with internal linkage instead each would // get a registry of its own, the reset hook would clear one nobody uses and the -// capacity hook would always answer zero, which the backend suite's pool -// assertions would then pass without testing anything. +// capacity hook would always answer zero. +// +// The backend suite does not let that through: six of its assertions across five +// cases read the capacity back and expect a figure above zero, and every one of +// them would fail. What it cannot do is say why, since those cases are named for +// the empty input, the capture refusal and the three growths rather than for +// linkage. That is what this file buys -- one named case in the pool's own target, +// which needs no GPU. #include "torch_tensorrt/executorch/SharedScratchPool.h" From 225299748be513a3aa6967369a227bf9b726a946 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Tue, 8 Sep 2026 17:14:45 -0700 Subject: [PATCH 13/13] fix(executorch): free a growth's retired buffer on a stream the pool owns `SharedScratchClaim::release()` took a boolean saying whether the calling `execute()` would synchronize its stream before returning, and picked between two disposals on it. That boolean is the shared root of the three defects this path has had on this branch. Each was the same mistake in a different place: a disposal whose correctness depended on what the caller was going to do next. - `45ea07cc91` queued the free on the caller's stream unconditionally. - `949efb5c46` found that on the caller-stream, device-resident path nothing ever synchronizes that stream, so the bytes never came back, and gated the choice on `must_sync` recorded when the buffer was retired. - `c1004de600` found that a call failing after the growth had recorded a promise to synchronize that it no longer had anything left to keep, and moved the answer to a `release()` parameter. The parameter is now gone, and with it the last piece of caller state the claim held: `retire()` is no longer given the calling stream either. The pool keeps a non-blocking CUDA stream per device, created by the first growth that has a buffer to dispose of. Nothing destroys it, for the same reason nothing destroys the buffer or the handoff event: the pool runs no CUDA call at process exit. It is a new leaked handle, one per device that ever grew. Every disposal is now the same three steps -- wait on the host for the enqueue the handoff event names, `cudaFreeAsync` on the pool's stream, synchronize that stream. The pool can promise that synchronize because the stream is its own; a caller cannot, which is the whole reason the choice existed. A return added between the claim and the release now cannot leave anything stale behind it, because there is nothing recorded at the claim, and a new kind of call cannot be given the wrong disposal, because the disposal never asks what kind of call it is. Measured on an A100 with CUDA 13.0, freeing a 512 MiB buffer behind a 10 ms enqueue with a host function parked for 1500 ms on a stream nothing else in the measurement submitted to, six runs each: | disposal | ms | bytes still held on return | | --- | --- | --- | | host wait + device-wide `cudaFree` (what a caller-stream growth did) | 1500.8-1501.4 | 0 of 512 MiB | | host wait + `cudaFreeAsync` on the pool's stream + synchronize it | 10.6-11.7 | 0 of 512 MiB | | `cudaFreeAsync` on the caller's stream, return (what a synchronizing growth did) | 0.0 | 512 MiB of 512 MiB | | the same, plus the call's own synchronize | 10.4-11.8 | -- | A caller-stream growth therefore goes from 1501 ms to 10.6 ms, because the device-wide free waits out the parked host function and the stream-ordered one does not. A synchronizing growth pays 10.6 ms in the disposal where it used to pay nothing and arrives at the same point at the same time, because the wait for its own enqueue moves out of its own `cudaStreamSynchronize` and into the disposal; what that costs it is that its aliased reflects and D2H copies are queued after the enqueue has finished rather than while it runs. A call that does not grow the pool pays nothing either way, because `release()` with nothing retired makes no CUDA call at all. The host wait for the enqueue is unchanged and is still what a growth mostly costs. The pool's stream does not avoid it and is not meant to: `record_shared_scratch_enqueue` re-records the handoff event on the calling stream before `release()` runs, so the event the disposal waits on names this call's own enqueue as well as the one that last used the buffer. A growth that reached its own enqueue therefore returns with its engine work finished rather than in flight. What the pool's stream avoids is the rest of the device. A device-wide `cudaFree` is still the fallback -- where the device has no stream-ordered allocator and `cudaFreeAsync` reports `cudaErrorNotSupported`, where it reports anything else, and where the pool's stream could not be created at all. The test-only reset hands the disposal stream back with the buffer and the event, so a slot it clears is the slot a fresh process would have. Tests: - `AGrowthOnACallerStreamDoesNotWaitForUnrelatedWorkQueuedOnTheDevice` is new and is the case the old disposal could not pass: a growth on a caller stream returns while a host function is still parked on a stream it never submitted to. Forcing the disposal device-wide fails it, and its synchronized-path twin, at the gate's 60 s watchdog. - Both twins also assert that the pool created its disposal stream. Without that, a disposal that queued the free on the calling engine's stream would satisfy every other assertion on the synchronized path. - `AGrowthOnACallerStreamDisposesWithTheDeviceLockDropped` stalled the growth with a device-wide free waiting on an unrelated stream, which no longer happens. It parks the host function ahead of the growing engine's own enqueue instead, which is what the disposal's wait is for. Skipping that wait fails it, and fails `AGrowthOnACallerStreamWaitsForTheEnqueueOnTheBufferItRetires` with it -- there through `cudaErrorIllegalAddress`, the engine reading a buffer freed under it. - `ALargerEngineGrowsThePoolAndFreesTheBufferItReplaces` and its synchronized twin are unchanged in substance: both paths must still have the bytes back before `execute()` returns. - The pool's own suite gains `SharedScratchDisposalStream`: one stream per device, reused by every later growth, retried after a creation failure. Nothing covers the device-wide fallback any more. Reaching it needs `cudaFreeAsync` to fail a call that is correct as made, or the stream creation to fail, and neither can be induced from inside the test process; the skip in those cases says so rather than pointing at coverage that no longer exists. `README.md` and the `execute()` contract in `TensorRTBackend.h` described the two disposals and which call got which, in the growth bullet and the capture bullet; both are rewritten around the one disposal. The `cudaDeviceReset()` bullet in each gains the disposal stream as a third thing the primary context takes with it. Local run on one A100 with CUDA 13.0: `//tests/cpp/executorch:executorch_backend_tests`, 95 of 95 cases passing. --- .../executorch/TensorRTBackend.h | 57 ++-- cpp/src/torch_tensorrt/executorch/README.md | 169 +++++----- .../executorch/SharedScratchPool.h | 53 +++- .../executorch/SharedScratchPoolReset.h | 13 +- .../executorch/SharedScratchPoolTestHooks.cpp | 13 +- .../executorch/SharedScratchPoolTestHooks.h | 5 +- .../executorch/TensorRTBackend.cpp | 296 ++++++++++-------- .../test_shared_scratch_backend.cpp | 227 ++++++++------ .../executorch/test_shared_scratch_pool.cpp | 134 ++++++-- 9 files changed, 561 insertions(+), 406 deletions(-) diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index c9bbe6ec577..93ff207bf14 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -138,31 +138,24 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { // handles that do draw on it, five further consequences: // - A call needing more scratch than the pool holds grows it, and the growth // has to get rid of the buffer it replaces before it returns, or the bytes - // of every size the pool ever grew to stay resident at once. What that costs - // depends on whether this call is one that synchronizes the stream. - // A call that does -- anything staging through host memory, aliasing an - // output, or running with no caller stream -- queues a stream-ordered - // cudaFreeAsync and its own synchronization returns the bytes, waiting for - // nothing it did not submit. A call that does not -- the caller-stream, - // device-resident case the rest of this contract is about -- makes a host - // wait on the previous enqueue and a device-wide cudaFree instead, because a - // queued free on a stream nothing ever synchronizes never returns the bytes - // at all. So that one call blocks until the device is idle however - // asynchronous the rest of this contract makes it -- an unbounded wait on - // work this call did not submit. Unbounded is meant literally: if any of - // that work is itself waiting on something only this thread supplies once + // of every size the pool ever grew to stay resident at once. Every growth + // pays the same thing for that, whatever kind of call it is and whether or + // not it goes on to fail: a host wait on the handoff event, which is a whole + // inference and has no upper bound. By the time a call that reached its own + // enqueue makes that wait, the event names that enqueue as well as the one + // that last used the buffer, so such a growth returns with its engine work + // finished rather than in flight. Unbounded is meant literally: if that + // enqueue is itself waiting on something only this thread supplies once // execute() returns -- a host function it will release, a copy it will // enqueue next -- the call does not return, and the thread that would - // unblock it is the one inside cudaFree. The same wait is also the fallback - // for the first case, on a device with no stream-ordered allocator, where - // cudaFreeAsync reports cudaErrorNotSupported rather than freeing, and on - // any other code it reports, since a queued free that did not happen is not - // a free. It is also what a call that fails after the growth takes -- a - // refused enqueue, or a failed record of one already submitted -- whichever - // kind of call it is, because by the time such a call disposes of the - // buffer it has no synchronization of that stream left to make, so a queued - // free's bytes would never come back. Which calls grow the pool is not - // knowable from here; see the README. + // unblock it is the one inside the wait. What the growth does not wait for + // is the rest of the device. The free is queued on a stream the pool owns + // and synchronizes itself, so work on streams this call never submitted to + // is not in the wait, and the bytes are back before it returns. On a device + // with no stream-ordered allocator, where cudaFreeAsync reports + // cudaErrorNotSupported rather than freeing, the free is a device-wide + // cudaFree instead and the call does block until the device is idle. Which + // calls grow the pool is not knowable from here; see the README. // - A call that synchronizes the stream waits, through the handoff, for every // pooled enqueue submitted before it on that device. That is what one buffer // costs, and it is another unbounded wait: park a host function ahead of a @@ -177,11 +170,10 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { // allocation invalidates it under every mode but cudaStreamCaptureModeRelaxed, // which permits it but does not record it, leaving a replay pointed at a // buffer the pool may since have freed; and a growth's disposal of the buffer - // it replaces is prohibited too -- a call that synchronizes the stream is one - // no capture mode permits, so the only disposal a capture could reach is the - // host wait and the device-wide cudaFree, which a capture cannot take outside - // Relaxed either. The alternative to refusing is a capture that silently - // comes back invalidated. Only a capture on the selected stream is caught. A + // it replaces is prohibited too, at its first step: the host wait on the + // handoff event is, outside Relaxed, and so is the cudaFree it falls back to. + // The alternative to refusing is a capture that silently comes back + // invalidated. Only a capture on the selected stream is caught. A // capture running on any other stream under cudaStreamCaptureModeGlobal, or // under cudaStreamCaptureModeThreadLocal from this thread, is invalidated by // the same calls and is not refused, because CUDA offers no query for it: do @@ -203,10 +195,11 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { // synchronizing the stream instead is one no capture mode permits, so there // is no call shape that captures and leaves the handle safe to go on using. // The README states this in full. - // - cudaDeviceReset() invalidates the pool without emptying it. The buffer - // and the handoff event it still holds are destroyed with the primary - // context, and the next call on that device uses both. There is no guard: - // do not reset a device this backend has run a pooled engine on. + // - cudaDeviceReset() invalidates the pool without emptying it. The buffer, + // the handoff event and the disposal stream it still holds are destroyed + // with the primary context, and the next call on that device uses all three. + // There is no guard: do not reset a device this backend has run a pooled + // engine on. // - The pool adds six failure points to a call, on top of Error::NotSupported // for the capture above. Only one of them has a code of its own: // Error::Internal, which a pooled call returns when the device's handoff diff --git a/cpp/src/torch_tensorrt/executorch/README.md b/cpp/src/torch_tensorrt/executorch/README.md index 7e40c5efefd..008960a613d 100644 --- a/cpp/src/torch_tensorrt/executorch/README.md +++ b/cpp/src/torch_tensorrt/executorch/README.md @@ -208,72 +208,29 @@ Not disposing of it inside the call is not an option: the pool grows monotonical so a run that let each retired buffer outlive its growth would end up holding the sum of every size the pool was ever grown to rather than the largest of them. -There are two ways to get rid of it, and which one a call uses is decided by -whether that call synchronizes the stream before it returns -- that is, whether it -stages anything through host memory, aliases an output, or runs with no caller -stream. Both are made with the per-device lock dropped, so neither is serialized -behind the pool's own lock and neither holds the next claimant off while it runs: -two pooled calls serialize at submission, as the caller-stream contract above says, -and a growth is not an exception to it. That is not the same as holding nobody up, -and the paragraph below on the device-wide free says what it does hold up. - -Which of the two a call uses is settled where the buffer is disposed of and not -where it is claimed, because it is a claim about what the call has still to do. A -call that fails after the growth -- a refused `enqueueV3`, or a failed record of an -enqueue already submitted -- takes the host wait whichever kind of call it is. The -failed record does drain the stream, but it drains and then returns, so the -disposal falls after the last synchronization that call makes; the refused enqueue -makes none at all. A free queued on a stream that is never synchronized again does -not come back: not at `cudaStreamDestroy`, not at `cudaDeviceSynchronize`, not from -another stream, and not at `cudaMemPoolTrimTo`. Measured on an A100 with CUDA 13.0, -a 4 GiB block queued that way and then abandoned with the stream is gone for the -life of the process, on a device with tens of gigabytes still free. - -**A call that synchronizes** queues the free: `cudaFreeAsync` on the stream it -enqueued on. Ordering is what makes that safe, not a wait -- every pooled call -makes its stream wait on the device's handoff event before it enqueues and records -on that event afterwards, so a free queued on this call's stream sits behind every -enqueue that ever used the retired buffer, and behind this call's own, which uses -the new one. The call's own `cudaStreamSynchronize` is then what returns the -bytes, and nothing waits for work the call did not submit. Measured on an A100 -with CUDA 13.0 and driver 13.0, with a host function parked for 1500 ms on a stream -neither the pool nor the growing engine had ever used: around 10 ms, against the -1505 ms the disposal below took, which returned exactly when that host function was -released. - -**A call that does not synchronize** -- the caller-stream, device-resident case, -which is the one this option exists for -- makes a host wait on the handoff event -and then a device-wide `cudaFree`. It cannot queue the free, because a queued free -does not return the bytes when the stream reaches it. Measured on the same -machine: `cudaFreeAsync` on a `cudaMalloc`'d pointer returns `cudaSuccess` and -defers the free, but the bytes come back at the next `cudaStreamSynchronize` of -that stream and at no point before it. A stream `cudaStreamQuery` reports as -drained still holds them, `cudaMemGetInfo` does not move, and a `cudaMalloc` of the -same size fails with out of memory until the synchronize. On this path nothing -this backend does will ever synchronize that stream, so replaying the pool's own -growth sequence on one stream -- 8, 16, 24 and 32 GiB, no synchronize between -them -- four growths that all succeed with `cudaFree` fail on the fourth with out -of memory with `cudaFreeAsync`. - -The same host wait and `cudaFree` are also the fallback for the first case: where a -device has no stream-ordered allocator (`cudaDevAttrMemoryPoolsSupported`), -`cudaFreeAsync` reports `cudaErrorNotSupported` instead of freeing -- and on any -other code it reports, since a queued free that did not happen is not a free. - -**That `cudaFree` is the thing to know about**, because it is what every growth on -the asynchronous path pays. It waits for everything queued on the device, not only -for the enqueues that used the buffer, so on such a growth the asynchronous return -this backend otherwise promises does not happen: the call blocks until the device -is idle. The wait has no upper bound, and a caller can -turn it into one that never ends. Anything queued anywhere on the device is enough -to hold it, so if some stream is waiting on work only this thread will submit -- a -host function it will release after `execute()` returns, a copy it will enqueue -next -- the growing call does not return and the thread that would unblock it is -the one inside `cudaFree`. This is not hypothetical: the backend's own tests -deadlocked on it once, when a change of test order turned a case that parks a host -function on its own stream into the one that grew the pool. A program that parks -work like that has to keep the pool's growths away from it, and the paragraph -below on how often the pool grows is what says whether a run order can do that. +Every growth gets rid of it the same way, and none of the three steps reads +anything about the call the growth was made from: wait on the host for the enqueue +the handoff event names, queue a `cudaFreeAsync` on a stream the pool owns for that +device, and synchronize that stream, which is what returns the bytes. All three are +made with the per-device lock dropped, so none of them is serialized behind the +pool's own lock and none holds the next claimant off while it runs: two pooled calls +serialize at submission, as the caller-stream contract above says, and a growth is +not an exception to it. + +**The host wait is the thing to know about**, because it is what every growth pays. +It is a whole inference, and it has no upper bound: whatever the enqueue is waiting +for holds it too, so if a host function ahead of that enqueue will be released only +by this thread once `execute()` returns, the growing call does not return and the +thread that would release it is the one inside the wait. + +The event it waits on is the handoff, and a call that reached its own enqueue has +recorded that enqueue there before it releases the claim, so the wait covers this +call's inference as well as the one before it. A growth therefore returns with its +engine work finished rather than in flight, however asynchronous the caller-stream +contract above lets a call that does not grow the pool be, and anything `execute()` +queues after the release -- an aliased reflect, a D2H copy -- goes behind a +completed enqueue rather than a running one. Only a growth pays that; a call with +nothing retired makes no CUDA call in the release at all. Sharing one buffer has an unbounded wait of its own, growth or no growth. A call that synchronizes waits, through the handoff, for every pooled enqueue submitted @@ -281,6 +238,56 @@ before it on that device. Park a host function ahead of one pooled call's enqueu and release it only once a later pooled call has returned, and the later call is the one that never returns. +What a growth does *not* wait for is the rest of the device, and that is what the +pool's own stream buys. A device-wide `cudaFree` waits for everything queued on the +device rather than for the work that touched the buffer. Measured on an A100 with +CUDA 13.0 and driver 13.0, freeing a 512 MiB buffer behind a 10 ms enqueue with a +host function parked for 1500 ms on a stream nothing else in the measurement +submitted to: 1500.8-1501.4 ms device-wide against 10.6-11.7 ms on the pool's +stream, both with every byte back by the time the disposal returned. That wait is +not hypothetical: the backend's own tests deadlocked on the device-wide free once, +when a change of test order turned a case that parks a host function on its own +stream into the one that grew the pool. + +The stream has to be the pool's rather than the calling engine's. A stream-ordered +free hands the bytes back at the next `cudaStreamSynchronize` of the stream it was +queued on and at no point before it -- measured on the same machine, `cudaFreeAsync` +on a `cudaMalloc`'d pointer returns `cudaSuccess` and defers the free, a stream +`cudaStreamQuery` reports as drained still holds the bytes, `cudaMemGetInfo` does +not move, and a `cudaMalloc` of the same size fails with out of memory until the +synchronize -- so whoever queues the free has to be able to promise that +synchronize, and only the pool can. Queuing it on the caller's stream costs the +disposal nothing, 0.0 ms in the same measurement, because all it does is queue; the +whole 512 MiB was still resident when it returned, and the caller's own synchronize +is what gives it back. Where the caller makes one, that is no faster overall: +10.4-11.8 ms to the same point. Where it does not -- the caller-stream, +device-resident case, which is the one this option exists for -- the bytes never +come back at all: nothing this backend does will ever synchronize that stream, and +the stream is the caller's to destroy. Measured, a 4 GiB block queued +that way and then abandoned with the stream is gone for the life of the process, on +a device with tens of gigabytes still free: not at `cudaStreamDestroy`, not at +`cudaDeviceSynchronize`, not from another stream, and not at `cudaMemPoolTrimTo`. +Replaying the pool's own growth sequence on one such stream -- 8, 16, 24 and 32 GiB, +no synchronize between them -- four growths that all succeed when the disposal +returns the bytes fail on the fourth with out of memory. + +Which calls do synchronize their stream is knowable, and the disposal does not ask. +It is a claim about what the caller has still to do rather than about anything the +pool holds, and three defects on this path in a row were a disposal getting that +claim wrong -- the last of them a call that failed after the growth and so had no +synchronization left to make, having recorded at the claim that it would make one. +So the disposal takes no argument: there is nothing to record when the buffer is +retired and nothing a later return can leave stale. + +A device-wide `cudaFree`, and the wait for the whole device that comes with it, is +still the fallback. It is what a growth makes where the device has no +stream-ordered allocator (`cudaDevAttrMemoryPoolsSupported`) and `cudaFreeAsync` +reports `cudaErrorNotSupported` instead of freeing -- and on any other code it +reports, since a queued free that did not happen is not a free -- and where the +pool's stream could not be created at all. A program that parks work across an +`execute()` has to keep such growths away from it, and the paragraph below on how +often the pool grows is what says whether a run order can do that. + What an engine answers when asked how much it needs is decided 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 @@ -314,15 +321,12 @@ same. Its `cudaMalloc` returns `cudaErrorStreamCaptureUnsupported` and invalidat the capture under `Global` and `ThreadLocal`, and under `Relaxed` is permitted but runs uncaptured, so a replayed graph would use whatever buffer was installed when it was captured -- which by then the pool may have freed. So does its disposal of -the buffer it replaces. A call that synchronizes the stream is one no capture mode -permits, so the only disposal a capture could reach is the host wait on the handoff -event and the `cudaFree`, and those two return `cudaErrorStreamCaptureUnsupported` -and invalidate the capture under `Global` and `ThreadLocal`, and under `Relaxed` -are permitted but run uncaptured. The queued free would be no way round it either: -a `cudaFreeAsync` made while capturing is refused under all three modes, with -`cudaErrorInvalidValue` and without invalidating the capture -- measured; a free -made while capturing has to be given a graph allocation, and the pool's buffers -come from `cudaMalloc`. +the buffer it replaces, at its first step: the host wait on the handoff event +returns `cudaErrorStreamCaptureUnsupported` and invalidates the capture under +`Global` and `ThreadLocal`, and under `Relaxed` is permitted but runs uncaptured. +So does the `cudaFree` it falls back to where the stream-ordered free is refused. +Queuing the free on the pool's own stream is no way round any of that, because the +wait ahead of it has already happened. The backend checks the selected stream and returns instead, and it checks ahead of everything a capture cannot take -- not merely ahead of the pool's @@ -402,13 +406,14 @@ unsupported rather than pointing a caller at it. ### cudaDeviceReset() is not survivable -The pool is not guarded against it: it holds its buffer and its handoff event for -the process lifetime, and a reset destroys the primary context under both. The -next pooled `execute()` on that device then waits on a destroyed event and hands -the engine a pointer that is no longer a device allocation. Guarding this would -mean revalidating both on every call, and the check that would catch it is as -expensive as the work it protects. Treat a device the backend has run a pooled -engine on as one that must not be reset. +The pool is not guarded against it: it holds its buffer, its handoff event and the +stream its growths dispose on for the process lifetime, and a reset destroys the +primary context under all three. The next pooled `execute()` on that device then +waits on a destroyed event and hands the engine a pointer that is no longer a +device allocation, and the next growth after that queues its free on a stream that +no longer exists. Guarding this would mean revalidating all three on every call, +and the check that would catch it is as expensive as the work it protects. Treat a +device the backend has run a pooled engine on as one that must not be reset. ## Standalone Backend Archive diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h index 9fb538036bb..df2e5a0f242 100644 --- a/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h @@ -8,9 +8,10 @@ #pragma once // Bookkeeping for the TensorRT backend's shared per-device activation-scratch -// 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. +// pool: the grow/reuse policy, the enqueue-handoff rule, the stream a growth +// disposes on, and the lock that scopes them to a single device. +// Allocation, event creation and stream creation arrive as callables rather than +// being made here. #include @@ -55,6 +56,10 @@ struct SharedScratchDevice { void* buffer = nullptr; std::size_t capacity = 0; SharedScratchMarker marker; + // The stream a growth queues its free of the buffer it replaced on. Created on + // the first growth that has one to dispose of; nothing in the normal path + // destroys it, only the test-only reset hands it to a disposer that does. + cudaStream_t disposal_stream = nullptr; }; class SharedScratchPool; @@ -124,9 +129,10 @@ class SharedScratchPool { // 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. // -// 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. +// The buffers, the events and the disposal streams are intentionally never +// released 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. // // The C++ object is never destroyed either, and that is not the same claim. // Static destruction would destroy the registry's mutex, every device's mutex @@ -181,6 +187,27 @@ inline cudaEvent_t shared_scratch_mark_in_flight(SharedScratchDevice& dev) { return dev.marker.event; } +// The stream this device's growths queue their frees on, creating it on first +// use. Call with `dev.mu` held. +// +// `create_stream` returns a CUDA stream, or nullptr if one could not be created, +// in which case the slot stays empty and the next growth retries. +// +// It is the pool's own stream and not the claimant's, because a stream-ordered +// free hands the bytes back at the next synchronize of the stream it was queued +// on and at no point before it, so whoever queues it has to be able to promise +// that synchronize. A claimant cannot: on the path this pool exists for it never +// synchronizes its stream at all, and that stream is the caller's to destroy the +// moment the call returns. The pool synchronizes this one itself, in the same +// call that queued the free. +template +cudaStream_t shared_scratch_disposal_stream(SharedScratchDevice& dev, CreateStream create_stream) { + if (dev.disposal_stream == nullptr) { + dev.disposal_stream = create_stream(); + } + return dev.disposal_stream; +} + // 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 @@ -194,17 +221,9 @@ inline cudaEvent_t shared_scratch_mark_in_flight(SharedScratchDevice& dev) { // only once all the earlier ones have. An enqueue that reaches the buffer without // doing both is ordered against nothing here. // -// The caller has two ways to honour that, and both belong outside `dev.mu`: a -// free queued on a stream already ordered after the event, or a host wait on the -// event followed by a device-wide free. Neither is the one the backend always -// takes. A queued free returns the bytes only at the next synchronize of that -// stream, so the backend queues it on a call that synchronizes before returning -// and makes the host wait everywhere else: on a call that does not synchronize, -// on a call that fails after the growth and so has no synchronization left to -// make, and wherever the queued free is refused -- a device with no -// stream-ordered allocator to queue onto, or any other code cudaFreeAsync -// reports. Under the lock either one makes an unrelated claim on this device wait -// for work it has nothing to do with. +// The disposal belongs outside `dev.mu`: it waits on the host for that enqueue, +// and under the lock that wait is one an unrelated claim on this device would sit +// through for nothing. struct RetiredScratch { void* buffer = nullptr; cudaEvent_t wait_for = nullptr; diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolReset.h b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolReset.h index 8b96d0ce2d2..35ed8b77eb0 100644 --- a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolReset.h +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolReset.h @@ -34,9 +34,9 @@ namespace torch_tensorrt { namespace executorch_backend { // Returns every device's slot to the state it had before anything claimed it, -// handing what the slot held to `dispose(device_id, buffer, event)` so the caller -// can release it. Entries stay in the map, so a reference `get` handed out -// remains valid. +// handing what the slot held to `dispose(device_id, buffer, event, disposal_stream)` +// so the caller can release it. Entries stay in the map, so a reference `get` +// handed out remains valid. // // Answers with the number of devices whose lock it could not take within // SharedScratchPool::kResetLockWait, whose slots it left alone. That case is a @@ -62,7 +62,7 @@ namespace executorch_backend { // one, and the caller is responsible for there being none. template std::size_t reset_shared_scratch_pool_slots(SharedScratchPool& pool, Dispose dispose) { - std::vector> taken; + std::vector> taken; std::size_t still_locked = 0; { std::lock_guard lk(pool.mu_); @@ -82,11 +82,12 @@ std::size_t reset_shared_scratch_pool_slots(SharedScratchPool& pool, Dispose dis continue; } std::lock_guard dev_lk(dev.mu, std::adopt_lock); - taken.emplace_back(slot.first, dev.buffer, dev.marker.event); + taken.emplace_back(slot.first, dev.buffer, dev.marker.event, dev.disposal_stream); dev.buffer = nullptr; dev.capacity = 0; dev.marker.event = nullptr; dev.marker.pending = false; + dev.disposal_stream = nullptr; } not_taken_yet.swap(still_to_try); if (not_taken_yet.empty() || std::chrono::steady_clock::now() >= deadline) { @@ -97,7 +98,7 @@ std::size_t reset_shared_scratch_pool_slots(SharedScratchPool& pool, Dispose dis still_locked = not_taken_yet.size(); } for (const auto& slot : taken) { - dispose(std::get<0>(slot), std::get<1>(slot), std::get<2>(slot)); + dispose(std::get<0>(slot), std::get<1>(slot), std::get<2>(slot), std::get<3>(slot)); } return still_locked; } diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp index bbeaea9837d..d8bc1b09b0d 100644 --- a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp @@ -26,13 +26,13 @@ std::size_t shared_scratch_capacity_for_testing(int device_id) { bool reset_shared_scratch_pool_for_testing() { int restore_to = 0; const bool have_current = cudaGetDevice(&restore_to) == cudaSuccess; - const std::size_t still_locked = - reset_shared_scratch_pool_slots(scratch_pool(), [](int device_id, void* buffer, cudaEvent_t event) { - if (buffer == nullptr && event == nullptr) { + const std::size_t still_locked = reset_shared_scratch_pool_slots( + scratch_pool(), [](int device_id, void* buffer, cudaEvent_t event, cudaStream_t disposal_stream) { + if (buffer == nullptr && event == nullptr && disposal_stream == nullptr) { return; } - // cudaFree and cudaEventDestroy both act on the current device, and a slot is - // keyed by the device its buffer came from. + // cudaFree, cudaEventDestroy and cudaStreamDestroy all act on the current + // device, and a slot is keyed by the device its buffer came from. if (cudaSetDevice(device_id) != cudaSuccess) { return; } @@ -42,6 +42,9 @@ bool reset_shared_scratch_pool_for_testing() { if (event != nullptr) { (void)cudaEventDestroy(event); } + if (disposal_stream != nullptr) { + (void)cudaStreamDestroy(disposal_stream); + } // Clears a non-sticky error so a reset does not leave one for the next call to // report. A sticky one survives the clear, and no cleanup here recovers it. (void)cudaGetLastError(); diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.h b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.h index 01ad3172277..c09895e2d07 100644 --- a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.h +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.h @@ -31,8 +31,9 @@ namespace executorch_backend { // The bytes the pool holds for `device_id` right now; zero if it holds nothing. std::size_t shared_scratch_capacity_for_testing(int device_id); -// Frees every device's buffer, destroys its handoff event, and clears the marker, -// so one test does not inherit a pool an earlier one grew. +// Frees every device's buffer, destroys its handoff event and its disposal +// stream, and clears the marker, so one test does not inherit a pool an earlier +// one grew. // // False when a device's lock was still held after the pool's reset deadline: that // slot is left as it was. A caller should report it rather than carry on, and diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index fa1c4f69462..78d6a06217b 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -251,25 +251,15 @@ class SharedScratchClaim { SharedScratchClaim& operator=(const SharedScratchClaim&) = delete; ~SharedScratchClaim() { // A claim that still holds a retired buffer here was never released, so - // execute() returned somewhere between the claim and the release below -- a - // refused enqueue, or a failed record of one already submitted -- and no - // synchronization of this stream is still to come. The record failure does - // drain the stream, but it drains and then returns, so a free queued here - // would be queued after the last synchronization the call makes; the refused - // enqueue makes none at all. Such a free would come back at nothing: measured - // on an A100 with CUDA 13.0 and driver 13.0, a cudaFreeAsync'd 4 GiB block is - // still resident after the stream is destroyed, after cudaDeviceSynchronize - // and after another stream's synchronize -- and a caller stream is the - // caller's to destroy, so those bytes are gone for the life of the process. A - // bail-out therefore takes the disposal that has the bytes back before it - // returns, whatever the call would have done had it run to the end. Passing - // that in here rather than storing it at retire() is what keeps the two in - // step: a return added between the claim and the release cannot leave a - // promise to synchronize behind it. + // execute() returned between the claim and the release below -- a refused + // enqueue, or a failed record of one already submitted. The disposal it makes + // is the one every other path makes, and it has the bytes back before it + // returns, so a bail-out after a growth costs the retired buffer and not the + // whole pool. // // The return is the caller's to report, and on this path there is no caller // left to report it to: an execute() that returned early has already failed. - (void)release(/*the_call_synchronizes_the_stream=*/false); + (void)release(); } SharedScratchDevice& hold(int device_id) { @@ -287,70 +277,63 @@ class SharedScratchClaim { // Takes ownership of a buffer a growth displaced, to be freed by release(). // `wait_for` is the marker event the enqueues that used it were recorded on, or - // null if none were; `stream` is the one this claim is about to enqueue on. - void retire(void* buffer, cudaEvent_t wait_for, cudaStream_t stream) { + // null if none were; `disposal_stream` is the pool's own stream for this device, + // or null if one could not be created. + void retire(void* buffer, cudaEvent_t wait_for, cudaStream_t disposal_stream) { retired_ = buffer; retired_wait_ = wait_for; - stream_ = stream; + disposal_stream_ = disposal_stream; } // 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 disposes of whatever a - // growth displaced, after the unlock, because either way of doing that blocks or - // queues on device work another pooled engine on this device has nothing to do - // with. + // growth displaced, after the unlock, because that disposal waits for an enqueue + // another pooled engine on this device has nothing to do with. // - // Which disposal, because a stream-ordered free does not return the bytes when - // the stream reaches it. Measured on an A100 with CUDA 13.0 and driver 13.0: - // cudaFreeAsync accepts a pointer cudaMalloc returned and defers the free, but - // the bytes come back at the next cudaStreamSynchronize of that stream and at - // no point before it. A stream cudaStreamQuery reports as drained still holds - // them -- a cudaMalloc of the same size fails with out of memory until the - // synchronize, and cudaMemGetInfo does not move. + // The disposal reads nothing the caller supplies and asks nothing about it: wait + // on the host for the enqueue the marker names, queue the free on the pool's own + // stream for this device, and synchronize that stream, which is what has the + // bytes back before this returns. It is the same three steps on every path, + // including the destructor's, so a return added between the claim and the + // release cannot change what the disposal does. // - // So the free is queued only where this call synchronizes the stream itself, - // which `the_call_synchronizes_the_stream` says. It is passed in here rather - // than recorded at retire() because it is a claim about what the caller has - // still to do, and only the caller, at the point it releases, knows that it is - // going to: the destructor answers it false for exactly that reason. + // The free is stream-ordered because a device-wide cudaFree waits for everything + // queued on the device rather than for the work that touched the buffer, and + // that wait has no upper bound: measured on an A100 with CUDA 13.0 and driver + // 13.0, freeing a 512 MiB buffer behind a 10 ms enqueue with a host function + // parked for 1500 ms on a stream nothing else here submitted to, 1500.8-1501.4 ms + // device-wide against 10.6-11.7 ms on the pool's stream, both with every byte + // back by the time the disposal returned. A caller can make the device-wide wait + // one that never ends, by parking work only this thread will release once + // execute() returns. // - // - the_call_synchronizes_the_stream. cudaFreeAsync on the stream this - // claim enqueued on, and execute()'s own cudaStreamSynchronize returns the - // bytes before the call ends. Nothing waits for work this call did not - // submit -- measured with a host function parked for 1500 ms on a stream - // neither the pool nor this engine had ever used, around 10 ms against the - // 1505 ms the synchronous disposal took, which came back when that host - // function was released. + // The stream has to be the pool's. A stream-ordered free hands the bytes back at + // the next synchronize of the stream it was queued on and at no point before it + // -- measured on the same machine, a stream cudaStreamQuery reports as drained + // still holds them, cudaMemGetInfo does not move, and a cudaMalloc of the same + // size fails with out of memory until the synchronize -- so whoever queues the + // free has to be able to promise that synchronize. Queuing it on the caller's + // stream makes the disposal itself cost nothing, 0.0 ms in the same measurement, + // because all it does is queue; the whole 512 MiB was still resident when it + // returned, and it is the caller's own synchronize that gives it back. Where the + // caller makes one that is not slower overall -- 10.4-11.8 ms to the same point, + // against the 10.6-11.7 above -- and where it does not, the bytes never come + // back: a 4 GiB block queued on a stream and abandoned is still resident after + // cudaStreamDestroy, after cudaDeviceSynchronize, after another stream's + // cudaStreamSynchronize and after cudaMemPoolTrimTo, and replaying this pool's + // growth sequence on one unsynchronized stream turns four growths that all + // succeed into an out-of-memory on the fourth. Telling those calls apart is a + // claim about what the caller has still to do; three defects on this path have + // been a disposal getting that claim wrong, so it is not made. // - // - Otherwise -- the caller-stream, device-resident path that returns with its - // enqueue still running -- nothing this backend does will ever synchronize - // that stream, so a queued free would hold the bytes until the caller - // happened to. Every buffer a run retired would stay resident and peak - // memory would be the sum of every size the pool grew to rather than the - // largest of them: measured, replaying this pool's growth sequence on one - // stream with no synchronize, four growths that all succeed under the - // synchronous disposal fail on the fourth with out of memory. That path - // waits on the marker event and makes a device-wide cudaFree, which has the - // bytes back before execute() returns and costs a host wait for everything - // queued on the device. The execute() contract and the README say so. + // cudaFreeAsync needs the device's stream-ordered allocator, which not every + // platform has. Where it is missing, or the pool's stream could not be created, + // the free is device-wide instead and costs the wait above. // - // Ordering, not a wait, is what makes the queued free safe. Every claimant makes - // its stream wait on the marker event before it enqueues, and this claim did so - // while holding the lock, so this stream is already behind every enqueue that - // ever used the retired buffer -- the marker covers all of them, because each - // records on it afterwards. Queuing the free on this stream therefore puts it - // after the last of them, and after this claim's own enqueue, which uses the new - // buffer. - // - // cudaFreeAsync also needs the device's stream-ordered allocator, which not - // every platform has; where it is missing the call fails rather than freeing and - // the synchronous disposal runs instead, so that one is both the other path's - // disposal and this path's fallback. - // - // Returns false when the synchronous disposal's wait failed, which leaves the - // buffer leaked rather than freed under a live enqueue; the caller reports it. - // Frees on the current device, which must still be the buffer's. - bool release(bool the_call_synchronizes_the_stream) { + // Returns false when the wait for the enqueue failed, which leaves the buffer + // leaked rather than freed under a live enqueue; the caller reports it. Frees on + // the current device, which must still be the buffer's. + bool release() { if (lock_.owns_lock()) { lock_.unlock(); } @@ -360,59 +343,89 @@ class SharedScratchClaim { } void* const retired = retired_; const cudaEvent_t wait_for = retired_wait_; - const cudaStream_t stream = stream_; + const cudaStream_t disposal_stream = disposal_stream_; // Clearing the pointer is what stops the destructor's release from disposing // of the same buffer twice. retired_ = nullptr; - if (the_call_synchronizes_the_stream) { - const cudaError_t async_err = cudaFreeAsync(retired, stream); - if (async_err == cudaSuccess) { - return true; - } + if (!wait_for_the_enqueue_that_used_it(wait_for)) { + return false; + } + if (disposal_stream == nullptr || !free_on_the_pools_stream(retired, disposal_stream)) { + free_device_wide(retired); + } + return true; + } + + private: + // Waits for the enqueue that last used the retired buffer, so the free below is + // not made under one still reading it. `wait_for` is the device's handoff marker, + // so where the calling execute() has already recorded its own enqueue there, the + // wait covers that one as well. Unbounded either way -- it is a whole inference, + // and the execute() contract says as much. + bool wait_for_the_enqueue_that_used_it(cudaEvent_t wait_for) { + if (wait_for == nullptr) { + return true; + } + const cudaError_t wait_err = cudaEventSynchronize(wait_for); + if (wait_err != cudaSuccess) { + // 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. What it reports is + // usually an asynchronous fault raised by earlier work on this device. + ET_LOG( + Error, + "TensorRTBackend::execute: waiting for the enqueue on the replaced shared activation scratch on device %d failed (%s), which for a wait on device work is usually an earlier asynchronous fault on this device surfacing here; leaking that buffer rather than freeing it under a live enqueue", + device_id_, + cudaGetErrorString(wait_err)); + cudaGetLastError(); + return false; + } + return true; + } + + // Queues the free on the pool's stream for this device and synchronizes it, + // which is what returns the bytes. Reports whether the buffer was freed, so a + // refusal falls back to the device-wide free; a synchronize that fails after the + // free was accepted is not one, and freeing again would be freeing twice. + bool free_on_the_pools_stream(void* retired, cudaStream_t disposal_stream) { + const cudaError_t free_err = cudaFreeAsync(retired, disposal_stream); + if (free_err != cudaSuccess) { // The error is this call's own, so it is cleared here rather than left for // the next CUDA call in execute() to report under its own name. cudaGetLastError(); - if (async_err == cudaErrorNotSupported) { + if (free_err == cudaErrorNotSupported) { ET_LOG( Info, - "TensorRTBackend::execute: device %d has no stream-ordered allocator, so the free of the shared activation scratch buffer a pool growth replaced falls back to a host wait and a device-wide free, which blocks this call until the device is idle", + "TensorRTBackend::execute: device %d has no stream-ordered allocator, so the free of the shared activation scratch buffer a pool growth replaced falls back to a device-wide free, which blocks this call until the device is idle", device_id_); } else { // Any other code is a fault on a device that does have the allocator, so // this does not say the platform lacks one. ET_LOG( Info, - "TensorRTBackend::execute: the stream-ordered free of the shared activation scratch buffer a pool growth replaced on device %d returned %s, so it falls back to a host wait and a device-wide free, which blocks this call until the device is idle", + "TensorRTBackend::execute: the stream-ordered free of the shared activation scratch buffer a pool growth replaced on device %d returned %s, so it falls back to a device-wide free, which blocks this call until the device is idle", device_id_, - cudaGetErrorString(async_err)); + cudaGetErrorString(free_err)); } + return false; } - return dispose_with_a_host_wait(retired, wait_for); - } - private: - // The disposal that has the bytes back before it returns: wait on the host for - // the enqueue that last used the buffer, then free it device-wide. Both waits - // are unbounded -- see the execute() contract -- which is why the queued free is - // preferred wherever it can return the bytes. - bool dispose_with_a_host_wait(void* retired, cudaEvent_t wait_for) { - if (wait_for != nullptr) { - const cudaError_t wait_err = cudaEventSynchronize(wait_for); - if (wait_err != cudaSuccess) { - // 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. What it reports is - // usually an asynchronous fault raised by earlier work on this device. - ET_LOG( - Error, - "TensorRTBackend::execute: waiting for the enqueue on the replaced shared activation scratch on device %d failed (%s), which for a wait on device work is usually an earlier asynchronous fault on this device surfacing here; leaking that buffer rather than freeing it under a live enqueue", - device_id_, - cudaGetErrorString(wait_err)); - cudaGetLastError(); - return false; - } + const cudaError_t sync_err = cudaStreamSynchronize(disposal_stream); + if (sync_err != cudaSuccess) { + // Nothing was on this stream but the free, and the enqueue it had to follow + // was already waited for, so a failure here is a fault this device was + // already in. Whether the bytes came back is not knowable from it. + ET_LOG( + Error, + "TensorRTBackend::execute: synchronizing the shared activation scratch pool's disposal stream on device %d reported %s, so the free of the buffer a growth replaced may not have returned its bytes", + device_id_, + cudaGetErrorString(sync_err)); + cudaGetLastError(); } + return true; + } + void free_device_wide(void* retired) { const cudaError_t err = cudaFree(retired); if (err != cudaSuccess) { // cudaFree synchronizes, so what it reports is more often an earlier @@ -428,7 +441,6 @@ class SharedScratchClaim { // resurface anyway; the caller learns of it from that call. cudaGetLastError(); } - return true; } SharedScratchDevice* dev_ = nullptr; @@ -436,7 +448,7 @@ class SharedScratchClaim { std::unique_lock lock_; void* retired_ = nullptr; cudaEvent_t retired_wait_ = nullptr; - cudaStream_t stream_ = nullptr; + cudaStream_t disposal_stream_ = nullptr; }; // What a call needing no activation scratch is given when the pool holds nothing @@ -457,13 +469,10 @@ constexpr size_t kMinPooledScratchBytes = 1; // cudaMalloc invalidates the capture outside cudaStreamCaptureModeRelaxed, and // under Relaxed is permitted but runs uncaptured, leaving a replay pointed at a // buffer the pool may since have freed. Its disposal of the buffer it replaces -// behaves the same way: a call that synchronizes the stream is one no capture mode -// permits, so the only disposal a capture could reach is the host wait and the -// device-wide cudaFree, and both are prohibited outside Relaxed. Even the queued -// free would be no way round it -- a cudaFreeAsync made while capturing has to be -// given a graph allocation and the pool's buffers come from cudaMalloc, so it is -// refused under all three modes, though without invalidating the capture. The -// invalidations do not fail cleanly: +// behaves the same way at its first step: the host wait on the handoff event is +// prohibited outside Relaxed, and so is the cudaFree it falls back to. Queuing the +// free on the pool's own stream is no way round that, because the wait ahead of it +// has already happened. The invalidations do not fail cleanly: // the caller learns of one only when cudaStreamEndCapture hands back an error and // a null graph. Refusing names the cause instead. // @@ -537,11 +546,11 @@ Error claim_shared_scratch(SharedScratchClaim& claim, int device_id, size_t need const SharedScratchHandoff handoff = shared_scratch_claim_event(dev, []() -> cudaEvent_t { cudaEvent_t event = nullptr; // Blocking-sync so the host yields instead of busy-spinning. The only host - // wait ever made on this event is SharedScratchClaim::release()'s synchronous - // disposal, and it waits for a whole inference; spinning would burn a core for - // that time and be no faster, since it is followed by a device-wide cudaFree. - // Where the free is queued on the stream instead nothing waits on this event - // from the host at all, and the flag costs nothing. + // wait ever made on this event is the one a growth's disposal makes in + // SharedScratchClaim::release(), and it waits for a whole inference; spinning + // would burn a core for that time and be no faster. Nothing on a call that + // does not grow the pool waits on it from the host, and the flag costs + // nothing there. if (cudaEventCreateWithFlags(&event, cudaEventDisableTiming | cudaEventBlockingSync) != cudaSuccess) { // The pool's own failure, cleared where it is made: the caller is told by // the return, and leaving it pending would surface it under the name of @@ -600,10 +609,32 @@ Error claim_shared_scratch(SharedScratchClaim& claim, int device_id, size_t need } // The retired buffer is disposed of at release(), with the device's lock - // dropped, either by a free queued on `stream` or by a host wait and a - // device-wide free; see SharedScratchClaim::release() for which and why. - // Nothing here makes a CUDA call that blocks on device work under that lock. - claim.retire(retired.buffer, retired.wait_for, stream); + // dropped; see SharedScratchClaim::release() for what that costs. Nothing here + // makes a CUDA call that blocks on device work under that lock -- the stream is + // created and not waited on, and only a growth that displaced a buffer creates + // one at all. + if (retired.buffer != nullptr) { + const cudaStream_t disposal_stream = shared_scratch_disposal_stream(dev, [device_id]() -> cudaStream_t { + cudaStream_t stream_for_disposals = nullptr; + // Non-blocking, so a free queued here is not ordered against the legacy + // default stream: the disposal waits for the enqueue that used the buffer + // and for nothing else, and the legacy stream would add whatever any other + // library on this device happens to have queued on it. + if (cudaStreamCreateWithFlags(&stream_for_disposals, cudaStreamNonBlocking) != cudaSuccess) { + // The pool's own failure, cleared where it is made: the disposal falls + // back to a device-wide free and says so, and leaving this pending would + // surface it under the name of whatever this thread calls next. + cudaGetLastError(); + ET_LOG( + Info, + "TensorRTBackend::execute: could not create the shared activation scratch pool's disposal stream on device %d, so this growth's free of the buffer it replaced is device-wide", + device_id); + return nullptr; + } + return stream_for_disposals; + }); + claim.retire(retired.buffer, retired.wait_for, disposal_stream); + } out_ptr = buffer; return Error::Ok; @@ -1457,9 +1488,7 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } } - // Whether this call ends by waiting for its own enqueue. Read twice below -- at - // the release of the scratch claim, whose disposal of a retired buffer turns on - // it, and at the wait itself -- so it is settled once here, which is the first + // Whether this call ends by waiting for its own enqueue. Settled here, the first // point everything it reads is final: the two staging flags and the aliased // reflects are set while the bindings are built above, and whether a caller // stream is active was read at the top. @@ -1588,15 +1617,12 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // function so the rest of execute() -- the aliased reflects, the D2H copies and // their synchronizations -- does not hold up another engine on this device. // The release is also where a growth disposes of the buffer it replaced, with - // the lock dropped. This is the only place that release is made rather than - // left to the claim's destructor, and it is the only place must_sync is a - // promise this function can still keep: every return past it synchronizes the - // stream when must_sync is set -- the branch below does, and the aliased - // reflects drain the enqueue before returning. - if (!scratch_claim.release(must_sync)) { - // Only the synchronous disposal reports a failure, and only from its host - // wait, which for a wait on device work means this device is already in a - // faulted state. + // the lock dropped; this is the only place that release is made rather than + // left to the claim's destructor. + if (!scratch_claim.release()) { + // The only failure it reports is its wait for the enqueue on the buffer a + // growth retired, which for a wait on device work means this device is + // already in a faulted state. drain_the_enqueue(); return Error::InvalidProgram; } @@ -1613,9 +1639,9 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } } - // The engine work is now in flight on `stream`, and must_sync -- decided above, - // where the scratch claim needed it -- says whether to wait for it. The D2H - // copies live in this branch: an output staged to host always sets + // The engine work is on `stream` -- still in flight, unless a growth's disposal + // above waited out the enqueue -- and must_sync says whether to wait for it. The + // D2H copies live in this branch: an output staged to host always sets // output_staged_to_host, so outputs_needing_copy is empty on the skip path. if (must_sync) { // Every return from here on is behind the cudaStreamSynchronize below, so @@ -1632,7 +1658,7 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* "TensorRTBackend::execute: D2H copy failed for output %zu: %s", output.first, cudaGetErrorString(cuda_err)); - // The enqueue already succeeded, so the engine is still running on the + // The enqueue already succeeded, so the engine may still be running on the // stream. Drain below before returning, or the next call mutates a live // execution context, which TensorRT forbids. copy_err = Error::InvalidProgram; diff --git a/tests/cpp/executorch/test_shared_scratch_backend.cpp b/tests/cpp/executorch/test_shared_scratch_backend.cpp index da3ea2ef7f1..e8d5e1d905c 100644 --- a/tests/cpp/executorch/test_shared_scratch_backend.cpp +++ b/tests/cpp/executorch/test_shared_scratch_backend.cpp @@ -827,8 +827,9 @@ class SharedScratchBackendTest : public ::testing::Test { return dynamic_blob_; } - // Defined below, next to the two cases that call it. + // Defined below, next to the cases that call them. void run_a_growth_and_bound_what_it_cost(bool synchronized_path); + void run_a_growth_beside_parked_work(bool synchronized_path); // A case can also skip for a reason the device being present says nothing about // -- the device would not fill, its memory would not stay still, it has no @@ -1135,27 +1136,25 @@ TEST_F(SharedScratchBackendTest, PooledEnginesShareOneActivationScratchAllocatio // device. // // The reading is taken the moment execute() returns and *before* any synchronize -// of this test's, which is what makes the upper bound mean something. A growth has -// two disposals to choose from and only one of them has the bytes back by then; a -// synchronize in between hides the difference, because it is exactly what releases -// a queued free. Measured on an A100 with CUDA 13.0: cudaFreeAsync on a -// cudaMalloc'd pointer returns cudaSuccess and defers the free, and the bytes come -// back at the next cudaStreamSynchronize of that stream and at no point before it -// -- a stream cudaStreamQuery reports as drained still holds them. So a growth -// that queued the free on a stream nothing synchronizes costs the whole new buffer -// here rather than the difference, and the upper bound below fails. +// of this test's, which is what makes the upper bound mean something. A +// stream-ordered free hands the bytes back only at the next synchronize of the +// stream it was queued on, so a synchronize in between is exactly what would hide +// a disposal that queued the free and left it. Measured on an A100 with CUDA 13.0: +// cudaFreeAsync on a cudaMalloc'd pointer returns cudaSuccess and defers the free, +// and the bytes come back at that synchronize and at no point before it -- a +// stream cudaStreamQuery reports as drained still holds them. So a growth that +// queued the free and returned costs the whole new buffer here rather than the +// difference, and the upper bound below fails. // // The lower bound also fails if the pool were already large enough for the second // engine, which is how this test could otherwise pass vacuously. The fixture // empties the pool before each test, so it is the smaller engine's run below that // establishes the size the growth has to exceed. // -// `synchronized_path` picks which of the two disposals runs, since the choice is -// execute()'s must_sync and not the pool's: a call with no caller stream ends by -// synchronizing and queues the free, a call on a caller stream with device-resident -// I/O does not and makes a host wait and a device-wide free instead. Both have to -// have the bytes back before the call returns, and this is the same measurement -// pointed at each. +// `synchronized_path` picks which kind of call grows the pool. The disposal is the +// same either way and this is the same measurement pointed at each, which is the +// point: the bound holds on a call whose own synchronize would have covered for a +// free the pool did not finish, and on one that makes no synchronize at all. void SharedScratchBackendTest::run_a_growth_and_bound_what_it_cost(bool synchronized_path) { ASSERT_GE(big_scratch_bytes_, scratch_bytes_ + kMinMeasurableScratch) << "the two fixture engines ask for " << scratch_bytes_ << " and " << big_scratch_bytes_ @@ -1204,8 +1203,8 @@ void SharedScratchBackendTest::run_a_growth_and_bound_what_it_cost(bool synchron ASSERT_TRUE(device_bytes_in_use(before)) << "cudaMemGetInfo failed, so this test measured nothing"; const auto measurement_began = std::chrono::steady_clock::now(); ASSERT_EQ(run_one(big), Error::Ok); - // No synchronize between the growth and this reading: see the note above. Both - // the allocation and both disposals are host calls execute() makes and returns + // No synchronize between the growth and this reading: see the note above. The + // allocation and the whole disposal are host calls execute() makes and returns // from, so what this reads is settled whether or not the enqueue has finished. std::size_t after = 0; ASSERT_TRUE(device_bytes_in_use(after)) << "cudaMemGetInfo failed, so this test measured nothing"; @@ -1297,17 +1296,17 @@ void SharedScratchBackendTest::run_a_growth_and_bound_what_it_cost(bool synchron } // The path this option exists for: a caller stream with device-resident I/O, so -// execute() returns with the enqueue still running and never synchronizes the -// stream. A growth here cannot queue its free, because nothing would ever make the -// bytes come back. +// execute() never synchronizes the stream. Nothing this test or the caller does +// would ever release a free the pool left queued, so the bound here is the one +// that says the pool finished it. TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItReplaces) { run_a_growth_and_bound_what_it_cost(/*synchronized_path=*/false); } -// The other path, where execute() ends by synchronizing the stream and the growth -// queues its free on it. The bytes have to be back by the time the call returns -// here too -- the same bound, and it is the call's own synchronize that satisfies -// it rather than anything this test does. +// The other path, where execute() ends by synchronizing the stream. The bytes have +// to be back by the time the call returns here too, and the same bound says so -- +// what this case adds is a growth beside a synchronize the disposal does not lean +// on, since the disposal returns the bytes itself and runs before it. TEST_F(SharedScratchBackendTest, AGrowthOnASynchronizedCallFreesTheBufferItReplaces) { run_a_growth_and_bound_what_it_cost(/*synchronized_path=*/true); } @@ -1681,12 +1680,10 @@ TEST_F(SharedScratchBackendTest, AFailedPooledAllocationLeavesTheDeviceLockFreeA // invalidates the capture under every capture mode. Left to run, the caller learns // about it only when cudaStreamEndCapture returns an error and a null graph, far // from the cause. execute() refuses instead. A growth adds more of the same: its -// allocation invalidates a capture outside cudaStreamCaptureModeRelaxed, and its -// disposal of the buffer it replaces does too -- a call that synchronizes the -// stream is one no capture mode permits, so the only disposal a capture could -// reach is the host wait and the device-wide free. This test starts from a pool -// already large enough and runs in relaxed mode, so it pins the wait rather than -// either of them. +// allocation invalidates a capture outside cudaStreamCaptureModeRelaxed, and so +// does the host wait its disposal of the buffer it replaces begins with. This test +// starts from a pool already large enough and runs in relaxed mode, so it pins the +// handoff's wait rather than either of them. // // The capture is ended either way: a capture left open belongs to the stream, and // destroying that stream would abandon it. @@ -2142,16 +2139,15 @@ TEST_F(SharedScratchBackendTest, ASecondPooledEnqueueWaitsForTheFirstOnAnotherSt // What a growth waits for // --------------------------------------------------------------------------- -// A growth's disposal of the buffer it replaces is one of two things, and each of -// these two cases parks the work one of them waits for. +// A growth's disposal of the buffer it replaces waits for one thing and not the +// other, and these three cases park each of them in turn. // -// A call that ends by synchronizing the stream queues the free on that stream, so -// it waits for nothing it did not submit -- in particular not for work parked on -// some unrelated stream, which is what the first case pins. A call that does not -// synchronize cannot queue the free, because the bytes would never come back, so -// it makes a host wait on the marker event and a device-wide cudaFree: it does -// wait for the enqueue against the buffer it retires, and the second case pins -// that it waits rather than freeing under a live enqueue. +// It waits for the enqueue against the buffer it retires, because freeing under a +// live enqueue is what the marker event is there to stop; the third case pins that +// it waits rather than freeing under it. It does not wait for anything else the +// device is running, because the free is queued on a stream the pool owns rather +// than made device-wide; the first two pin that, one on each kind of call. Both +// kinds, because what kind of call it is must not decide which disposal it gets. // // Each measures what the call came back before through the gate's watchdog, which // is the only other thing that can open a gate: if a call waited when it should @@ -2159,35 +2155,37 @@ TEST_F(SharedScratchBackendTest, ASecondPooledEnqueueWaitsForTheFirstOnAnotherSt // case says so. // Nothing the growing engine or the pool ever submitted to runs on the parked -// stream, so only a device-wide synchronization has any reason to wait for it. -// This is the queued free's whole point, so the case runs on the path that queues -// it: with no caller stream, execute() ends by synchronizing its own stream, which -// is what returns the bytes. +// stream, so only a device-wide free has any reason to wait for it. // // Gated on the device's stream-ordered allocator, and gated before the gate is // parked so a device without one costs milliseconds rather than the watchdog // interval. Where there is none, cudaFreeAsync reports cudaErrorNotSupported, the -// backend correctly falls back to the host wait and the device-wide free, and this -// assertion would report a defect that is not there -- the target is built for -// sbsa as well as x86_64, and nothing in the delegate asks the device about memory -// pools, so the fallback is what a whole platform would take. +// backend correctly falls back to a device-wide free, and this assertion would +// report a defect that is not there -- the target is built for sbsa as well as +// x86_64, and nothing in the delegate asks the device about memory pools, so the +// fallback is what a whole platform would take. // -// That fallback is not left untested by the skip. It is the same host wait and -// device-wide free that every growth on a caller stream makes on this device -- -// AGrowthOnACallerStreamWaitsForTheEnqueueOnTheBufferItRetires and -// ALargerEngineGrowsThePoolAndFreesTheBufferItReplaces both run it and check what -// it costs and what it frees. What a device with a memory pool cannot exercise is -// either branch that enters it from a call that does synchronize: cudaFreeAsync -// answering cudaErrorNotSupported, or answering any other code. -TEST_F(SharedScratchBackendTest, AGrowthDoesNotWaitForUnrelatedWorkQueuedOnTheDevice) { +// On a device that does have one, nothing in this suite covers that fallback. +// Reaching it needs cudaFreeAsync to fail a call that is correct as made, or the +// pool's stream creation to fail, and neither can be induced from inside this +// process. +void SharedScratchBackendTest::run_a_growth_beside_parked_work(bool synchronized_path) { 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 second to be sure of growing the pool"; ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + cudaStream_t stream = nullptr; cudaStream_t unrelated = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); ASSERT_EQ(cudaStreamCreateWithFlags(&unrelated, cudaStreamNonBlocking), cudaSuccess); + // `stream` is unread on the synchronized path: with no CallerStreamGuard + // execute() enqueues on cudaStreamPerThread and synchronizes before returning. + const auto run_one = [synchronized_path, stream](LoadedEngine& engine) { + return synchronized_path ? engine.run_on_the_synchronized_path() : engine.run(stream); + }; + LoadedEngine small; LoadedEngine big; ASSERT_EQ(small.load(blob(), 26), Error::Ok); @@ -2201,13 +2199,16 @@ TEST_F(SharedScratchBackendTest, AGrowthDoesNotWaitForUnrelatedWorkQueuedOnTheDe if (memory_pools == 0) { GTEST_SKIP() << "device " << device_id << " reports no stream-ordered allocator (cudaDevAttrMemoryPoolsSupported = 0), so a growth here " - "correctly takes the host wait and the device-wide free, which does wait for work parked " - "anywhere on the device"; + "correctly falls back to a device-wide free, which does wait for work parked anywhere on the " + "device"; } // Without this the larger engine allocates rather than grows, and a growth that - // retires nothing disposes of nothing. - ASSERT_EQ(small.run_on_the_synchronized_path(), Error::Ok); + // retires nothing disposes of nothing. Drained, so the only enqueue the growth's + // wait can be held by is the one it submits itself, which nothing here parks; + // the parked stream is the only other thing that could hold the growth up. + ASSERT_EQ(run_one(small), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); const std::size_t before = shared_scratch_capacity_for_testing(device_id); ASSERT_GT(before, 0u) << "the smaller engine left the pool empty, so the larger one has nothing to retire"; @@ -2216,14 +2217,25 @@ TEST_F(SharedScratchBackendTest, AGrowthDoesNotWaitForUnrelatedWorkQueuedOnTheDe GateRelease gate_release(gate, unrelated); ASSERT_TRUE(big.fill_output(kSentinel)); - const Error growth_error = big.run_on_the_synchronized_path(); + const Error growth_error = run_one(big); const bool came_back_with_the_device_held = !gate.forced_open.load(); gate_release.release(); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); EXPECT_EQ(growth_error, Error::Ok); EXPECT_GT(shared_scratch_capacity_for_testing(device_id), before) << "the pool did not grow, so no buffer was retired and nothing was disposed of"; + { + // The stream the free went on is the pool's, and this growth is what created + // it: a disposal that queued the free on the calling engine's stream instead + // would leave this slot empty and still return the bytes on the synchronized + // path, so nothing else here would notice. + SharedScratchDevice& dev = scratch_pool().get(device_id); + std::lock_guard lock(dev.mu); + EXPECT_NE(dev.disposal_stream, nullptr) + << "the growth disposed of the buffer it retired without the pool's own stream for this device"; + } const std::vector grown_output = big.read_output(); ASSERT_FALSE(grown_output.empty()); EXPECT_NE(grown_output[0], kSentinel) @@ -2233,15 +2245,29 @@ TEST_F(SharedScratchBackendTest, AGrowthDoesNotWaitForUnrelatedWorkQueuedOnTheDe << "the growing call did not return until the watchdog released a host function parked on a stream it never " "submitted to, so its disposal of the retired buffer waits for the whole device"; + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); ASSERT_EQ(cudaStreamDestroy(unrelated), cudaSuccess); } -// The other disposal. On a caller stream with device-resident I/O nothing will -// ever synchronize the stream, so the growth cannot queue the free -- the bytes -// would stay resident until the caller happened to synchronize, and every buffer a -// run retired would be resident at once. It waits on the marker event and frees -// device-wide instead, so it *does* wait for the enqueue against the buffer it -// retires, and what has to hold is that it waits rather than freeing under it. +// With no caller stream, execute() ends by synchronizing its own stream. +TEST_F(SharedScratchBackendTest, AGrowthDoesNotWaitForUnrelatedWorkQueuedOnTheDevice) { + run_a_growth_beside_parked_work(/*synchronized_path=*/true); +} + +// The path this option exists for, and the one a disposal that read the caller +// could not serve: nothing on it ever synchronizes the caller's stream, so a free +// queued there would never come back and the only alternative to the pool's own +// stream is a device-wide free. This case is that difference -- the growth here +// returns while a host function is still parked on an unrelated stream, which a +// device-wide free waits out. +TEST_F(SharedScratchBackendTest, AGrowthOnACallerStreamDoesNotWaitForUnrelatedWorkQueuedOnTheDevice) { + run_a_growth_beside_parked_work(/*synchronized_path=*/false); +} + +// The wait the disposal does make. Freeing the retired buffer while an enqueue is +// still reading it corrupts that engine's output and reports nothing, so the +// disposal waits on the marker event first, and what has to hold is that it waits +// rather than freeing under the enqueue. // // The smaller engine's own enqueue is parked behind a host function, and a thread // releases it a short while after the growing call starts. A growth that waited @@ -2323,11 +2349,10 @@ TEST_F(SharedScratchBackendTest, AGrowthOnACallerStreamWaitsForTheEnqueueOnTheBu << "the watchdog had to open the gate, so the growing call blocked for the whole watchdog interval rather than " "for the enqueue it had to wait for"; EXPECT_TRUE(waited_for_the_parked_enqueue) - << "the growing call returned while the enqueue against the buffer it retired was still parked, so it did not " - "wait for that enqueue -- either it freed the buffer under it or it only queued the free, which on this path " - "never gives the bytes back"; + << "the growing call returned while the enqueue against the buffer it retired was still parked, so it freed that " + "buffer under a live enqueue"; EXPECT_GT(shared_scratch_capacity_for_testing(device_id), before) - << "the pool did not grow, so no buffer was retired and there was nothing to order the free against"; + << "the pool did not grow, so no buffer was retired and there was nothing to wait for"; const std::vector small_output = small.read_output(); ASSERT_EQ(small_output.size(), kElems); EXPECT_EQ(std::memcmp(expected.data(), small_output.data(), kBytes), 0) @@ -2338,24 +2363,22 @@ TEST_F(SharedScratchBackendTest, AGrowthOnACallerStreamWaitsForTheEnqueueOnTheBu ASSERT_EQ(cudaStreamDestroy(growing), cudaSuccess); } -// Neither growth case above can see whether the disposal is made with the device's -// lock held: both run their growth alone, so a lock held across it holds nothing -// up and both stay green with the unlock moved below the free. This one is what -// covers that. +// None of the growth cases above can see whether the disposal is made with the +// device's lock held: each runs its growth alone, so a lock held across it holds +// nothing up and they all stay green with the unlock moved below the disposal. +// This one is what covers that. // -// It matters because the disposal on this path is a host wait and a cudaFree, and -// cudaFree waits for every stream on the device rather than for the ones that -// touched the buffer -- the 1505 ms against a host function parked for 1500 ms on -// an unrelated stream that the execute() contract cites. Under the lock, that wait -// is one every other pooled engine on the device has to sit through, which is the -// serialization the contract says a pooled call does not impose past its own -// submission. +// It matters because the disposal waits on the host for the enqueue against the +// buffer it retires, which is a whole inference. Under the lock, that wait is one +// every other pooled engine on the device has to sit through before it can even +// submit, which is the serialization the execute() contract says a pooled call +// does not impose past its own submission. // -// So: park work on an unrelated stream to stall the free, then check the pool lock -// while the growth is still inside it. The capacity read under that same try_lock -// is what stops the case passing on a lock that is free because the growth has not -// started -- past a growth it is above what the smaller engine left, and the growth -// has not returned. +// So: park a host function ahead of the growing engine's own enqueue, which is +// what that wait is for, then check the pool lock while the growth is still inside +// it. The capacity read under that same try_lock is what stops the case passing on +// a lock that is free because the growth has not started -- past a growth it is +// above what the smaller engine left, and the growth has not returned. TEST_F(SharedScratchBackendTest, AGrowthOnACallerStreamDisposesWithTheDeviceLockDropped) { ASSERT_GE(big_scratch_bytes_, scratch_bytes_ + kMinMeasurableScratch) << "the two fixture engines ask for " << scratch_bytes_ << " and " << big_scratch_bytes_ @@ -2363,9 +2386,7 @@ TEST_F(SharedScratchBackendTest, AGrowthOnACallerStreamDisposesWithTheDeviceLock ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); cudaStream_t stream = nullptr; - cudaStream_t unrelated = nullptr; ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); - ASSERT_EQ(cudaStreamCreateWithFlags(&unrelated, cudaStreamNonBlocking), cudaSuccess); LoadedEngine small; LoadedEngine big; @@ -2376,24 +2397,22 @@ TEST_F(SharedScratchBackendTest, AGrowthOnACallerStreamDisposesWithTheDeviceLock const int device_id = big.handle()->device_id; // Without this the larger engine allocates rather than grows, and a growth that - // retires nothing frees nothing. Synchronized, so the growth's wait on the - // marker returns at once and the free is the only thing left to stall it. + // retires nothing frees nothing. Drained, so the only enqueue the growth's wait + // can be held by is the one it submits itself, below the gate. ASSERT_EQ(small.run(stream), Error::Ok); ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); const std::size_t before = shared_scratch_capacity_for_testing(device_id); ASSERT_GT(before, 0u) << "the smaller engine left the pool empty, so the larger one has nothing to retire"; - // Neither engine submits to this stream, so nothing on the growth path but the - // device-wide free has any reason to wait for it. + // Ahead of the growing engine's enqueue, so that enqueue and the marker recorded + // after it stay pending and the disposal's wait for them does too. StreamGate gate; - ASSERT_EQ(cudaLaunchHostFunc(unrelated, hold_stream, &gate), cudaSuccess); - GateRelease gate_release(gate, unrelated); + ASSERT_EQ(cudaLaunchHostFunc(stream, hold_stream, &gate), cudaSuccess); + GateRelease gate_release(gate, stream); std::atomic growth_returned{false}; Error growth_error = Error::Internal; JoinAtScopeExit grower{std::thread([&] { - // A caller stream with device-resident I/O, which is the path whose disposal - // is the host wait and the device-wide free. growth_error = big.run(stream); growth_returned.store(true); })}; @@ -2418,20 +2437,20 @@ TEST_F(SharedScratchBackendTest, AGrowthOnACallerStreamDisposesWithTheDeviceLock grower.join(); ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); - ASSERT_EQ(cudaStreamDestroy(unrelated), cudaSuccess); 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"; + << "the watchdog had to open the gate rather than this test, so more than the watchdog interval passed between " + "parking the host function ahead of the growing engine's enqueue and releasing it, and nothing here was " + "measured under the conditions it describes"; ASSERT_EQ(growth_error, Error::Ok); ASSERT_TRUE(still_inside_execute) - << "the growing call returned before the gate opened, so its free never stalled and there was no window in " + << "the growing call returned before the gate opened, so its disposal never stalled and there was no window in " "which to observe the lock"; ASSERT_GT(shared_scratch_capacity_for_testing(device_id), before) - << "the pool did not grow, so no buffer was retired and no free was made"; + << "the pool did not grow, so no buffer was retired and nothing was disposed of"; EXPECT_TRUE(lock_free_during_the_disposal) - << "the device's pool lock stayed held for the whole of a stalled growth free, so every other pooled engine on " - "this device waits out a device-wide free it has nothing to do with"; + << "the device's pool lock stayed held for the whole of a stalled growth disposal, so every other pooled engine " + "on this device waits out an inference it has nothing to do with before it can submit"; } // --------------------------------------------------------------------------- diff --git a/tests/cpp/executorch/test_shared_scratch_pool.cpp b/tests/cpp/executorch/test_shared_scratch_pool.cpp index acd446b3025..a6e8693a78d 100644 --- a/tests/cpp/executorch/test_shared_scratch_pool.cpp +++ b/tests/cpp/executorch/test_shared_scratch_pool.cpp @@ -5,8 +5,9 @@ * LICENSE file in the root directory of this source tree. */ -// 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. +// Pins the shared scratch pool helper: its grow, reuse and per-device policy, its +// enqueue-handoff rule and the stream a growth disposes on, 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. @@ -104,12 +105,31 @@ struct FakeEventFactory { } }; +// Stands in for the CUDA stream factory, the way FakeEventFactory stands in for +// the event one: distinct non-null handles and a call count, so a test can tell a +// slot that keeps one stream from one that creates a stream per growth. +struct FakeStreamFactory { + int created = 0; + std::uintptr_t next = 0x5000; + bool fail_next = false; + + cudaStream_t operator()() { + if (fail_next) { + fail_next = false; + return nullptr; + } + ++created; + cudaStream_t s = reinterpret_cast(next); + next += 0x100; + return s; + } +}; + // 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) { RetiredScratch retired; - void* const p = shared_scratch_get_or_grow( - dev, need, [&a](std::size_t bytes) { return a.alloc(bytes); }, retired); + void* const p = shared_scratch_get_or_grow(dev, need, [&a](std::size_t bytes) { return a.alloc(bytes); }, retired); if (retired.buffer != nullptr) { a.retire(retired.buffer, retired.wait_for); } @@ -254,8 +274,7 @@ TEST(SharedScratchPool, EveryPathClearsTheRetirementItReports) { // for; `call` above gives each of its calls a fresh one. RetiredScratch retired; const auto request = [&](std::size_t need) { - return shared_scratch_get_or_grow( - dev, need, [&a](std::size_t bytes) { return a.alloc(bytes); }, retired); + return shared_scratch_get_or_grow(dev, need, [&a](std::size_t bytes) { return a.alloc(bytes); }, retired); }; void* const first = request(1024); @@ -367,6 +386,59 @@ TEST(SharedScratchHandoffTest, ASlotWithNoEventIsNotMarked) { EXPECT_FALSE(shared_scratch_claim_event(dev, std::ref(events)).needs_wait); } +// --------------------------------------------------------------------------- +// The stream a growth queues its free on. +// --------------------------------------------------------------------------- + +TEST(SharedScratchDisposalStream, FirstUseCreatesTheSlotsStreamAndEveryLaterOneReusesIt) { + SharedScratchDevice dev; + FakeStreamFactory streams; + + const cudaStream_t first = shared_scratch_disposal_stream(dev, std::ref(streams)); + EXPECT_NE(first, nullptr); + EXPECT_EQ(streams.created, 1); + + // A stream per growth would be a stream leaked per growth: nothing destroys one + // in the normal path, by the same argument that leaves the buffers and the + // events alone at teardown. + EXPECT_EQ(shared_scratch_disposal_stream(dev, std::ref(streams)), first); + EXPECT_EQ(shared_scratch_disposal_stream(dev, std::ref(streams)), first); + EXPECT_EQ(streams.created, 1); +} + +TEST(SharedScratchDisposalStream, StreamCreationFailureIsReportedAndRetried) { + SharedScratchDevice dev; + FakeStreamFactory streams; + + streams.fail_next = true; + // Null is what the disposal reads as "no stream to queue the free on", so it + // falls back to a device-wide free rather than passing this to cudaFreeAsync, + // where it would name the legacy default stream. + EXPECT_EQ(shared_scratch_disposal_stream(dev, std::ref(streams)), nullptr); + + // The failure leaves nothing behind, so the next growth tries again rather than + // taking the fallback for the rest of the process. + EXPECT_NE(shared_scratch_disposal_stream(dev, std::ref(streams)), nullptr); + EXPECT_EQ(streams.created, 1); +} + +// A stream belongs to the device that was current when it was created, and the +// free queued on it is of that device's memory. One stream shared across devices +// would put every growth's free on whichever device grew first. +TEST(SharedScratchDisposalStream, KeepsAnIndependentStreamPerDevice) { + SharedScratchPool pool; + FakeStreamFactory streams; + + const cudaStream_t zero = shared_scratch_disposal_stream(pool.get(0), std::ref(streams)); + const cudaStream_t one = shared_scratch_disposal_stream(pool.get(1), std::ref(streams)); + + EXPECT_NE(zero, nullptr); + EXPECT_NE(one, nullptr); + EXPECT_NE(zero, one); + EXPECT_EQ(streams.created, 2); + EXPECT_EQ(shared_scratch_disposal_stream(pool.get(0), std::ref(streams)), zero); +} + // --------------------------------------------------------------------------- // The registry that owns one entry per device. // --------------------------------------------------------------------------- @@ -429,13 +501,15 @@ TEST(SharedScratchPoolRegistry, EveryTranslationUnitSeesOneProcessPool) { TEST(SharedScratchPoolRegistry, ResetHandsBackEverySlotAndLeavesItEmpty) { // The backend test fixture runs this between cases, so every case that reads // what the pool holds depends on it. The two things it has to get right are - // handing the caller both the buffer and the event to release -- neither the - // pool nor the disposer can find them afterwards -- and clearing the slot, so - // the next claim allocates instead of reusing a pointer that has been freed. + // handing the caller everything the slot held -- the buffer, the event and the + // disposal stream, none of which the pool or the disposer can find afterwards -- + // and clearing the slot, so the next claim allocates instead of reusing a + // pointer that has been freed. SharedScratchPool pool; FakeAllocator zero; FakeAllocator one; FakeEventFactory events; + FakeStreamFactory streams; SharedScratchDevice& dev0 = pool.get(0); SharedScratchDevice& dev1 = pool.get(1); @@ -443,22 +517,28 @@ TEST(SharedScratchPoolRegistry, ResetHandsBackEverySlotAndLeavesItEmpty) { void* const dev1_buffer = call(dev1, one, 2048); const cudaEvent_t dev0_event = shared_scratch_claim_event(dev0, std::ref(events)).event; shared_scratch_mark_in_flight(dev0); - // Device 1 is left with a buffer and no event, which is the state of a slot - // whose event creation failed: the reset has to cope with a null there. + const cudaStream_t dev0_stream = shared_scratch_disposal_stream(dev0, std::ref(streams)); + // Device 1 is left with a buffer and neither an event nor a disposal stream, + // which is the state of a slot whose event creation failed and which never grew: + // the reset has to cope with nulls there. ASSERT_NE(dev0_buffer, nullptr); ASSERT_NE(dev1_buffer, nullptr); ASSERT_NE(dev0_event, nullptr); + ASSERT_NE(dev0_stream, nullptr); ASSERT_EQ(dev1.marker.event, nullptr); + ASSERT_EQ(dev1.disposal_stream, nullptr); struct Disposal { int device_id; void* buffer; cudaEvent_t event; + cudaStream_t disposal_stream; }; std::vector disposed; - reset_shared_scratch_pool_slots(pool, [&](int device_id, void* buffer, cudaEvent_t event) { - disposed.push_back({device_id, buffer, event}); - }); + reset_shared_scratch_pool_slots( + pool, [&](int device_id, void* buffer, cudaEvent_t event, cudaStream_t disposal_stream) { + disposed.push_back({device_id, buffer, event, disposal_stream}); + }); // The registry iterates in unspecified order, so each device is looked up. const auto disposal_for = [&disposed](int device_id) -> const Disposal* { @@ -474,8 +554,10 @@ TEST(SharedScratchPoolRegistry, ResetHandsBackEverySlotAndLeavesItEmpty) { ASSERT_NE(disposal_for(1), nullptr); EXPECT_EQ(disposal_for(0)->buffer, dev0_buffer); EXPECT_EQ(disposal_for(0)->event, dev0_event); + EXPECT_EQ(disposal_for(0)->disposal_stream, dev0_stream); EXPECT_EQ(disposal_for(1)->buffer, dev1_buffer); EXPECT_EQ(disposal_for(1)->event, nullptr); + EXPECT_EQ(disposal_for(1)->disposal_stream, nullptr); // The slot the reset cleared is the one a later lookup finds, so the reads below // are of the entry the backend would go on using. One address cannot tell a @@ -486,6 +568,7 @@ TEST(SharedScratchPoolRegistry, ResetHandsBackEverySlotAndLeavesItEmpty) { EXPECT_EQ(dev0.capacity, 0u); EXPECT_EQ(dev0.marker.event, nullptr); EXPECT_FALSE(dev0.marker.pending); + EXPECT_EQ(dev0.disposal_stream, nullptr); // A smaller request than the freed buffer served: reuse would satisfy it from // the stale capacity and allocate nothing, so this is what distinguishes a @@ -514,8 +597,9 @@ TEST(SharedScratchPoolRegistry, ResetHandsBackTheEventOfASlotThatNeverGotABuffer ASSERT_EQ(dev.buffer, nullptr) << "the allocation did not fail, so this slot is not the state under test"; std::vector> disposed; - reset_shared_scratch_pool_slots( - pool, [&](int, void* buffer, cudaEvent_t slot_event) { disposed.emplace_back(buffer, slot_event); }); + reset_shared_scratch_pool_slots(pool, [&](int, void* buffer, cudaEvent_t slot_event, cudaStream_t) { + disposed.emplace_back(buffer, slot_event); + }); ASSERT_EQ(disposed.size(), 1u) << "the reset skipped a slot holding an event and no buffer, so nothing ever destroys " "that event"; @@ -548,7 +632,7 @@ TEST(SharedScratchPoolRegistry, ResetLeavesEveryEntryWhereItWas) { } int disposed = 0; - reset_shared_scratch_pool_slots(pool, [&disposed](int, void*, cudaEvent_t) { ++disposed; }); + reset_shared_scratch_pool_slots(pool, [&disposed](int, void*, cudaEvent_t, cudaStream_t) { ++disposed; }); ASSERT_EQ(disposed, kDevices); int moved = 0; @@ -605,7 +689,7 @@ TEST(SharedScratchPoolRegistry, ResetDisposesWithNoLockHeld) { // Read inside the disposer: the claim completes once the reset returns either // way, so only what was true while the disposer ran tells the two apart. bool claimed_during_dispose = false; - reset_shared_scratch_pool_slots(pool, [&](int, void*, cudaEvent_t) { + reset_shared_scratch_pool_slots(pool, [&](int, void*, cudaEvent_t, cudaStream_t) { disposing.store(true); const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2); while (!claimed.load() && std::chrono::steady_clock::now() < deadline) { @@ -667,7 +751,8 @@ TEST(SharedScratchPoolRegistry, ResetReportsALockedSlotRatherThanWaitingForIt) { int disposed = 0; const auto started = std::chrono::steady_clock::now(); - const std::size_t still_locked = reset_shared_scratch_pool_slots(pool, [&](int, void*, cudaEvent_t) { ++disposed; }); + const std::size_t still_locked = + reset_shared_scratch_pool_slots(pool, [&](int, void*, cudaEvent_t, cudaStream_t) { ++disposed; }); const auto took = std::chrono::steady_clock::now() - started; release_it.store(true); holder.join(); @@ -714,7 +799,10 @@ TEST(SharedScratchPoolRegistry, ALeakedLockDoesNotCostTheResetTheSlotsItSharesTh ASSERT_NE(call(pool.get(device_id), alloc, 1024), nullptr); } std::vector walk_order; - ASSERT_EQ(reset_shared_scratch_pool_slots(pool, [&](int id, void*, cudaEvent_t) { walk_order.push_back(id); }), 0u) + ASSERT_EQ( + reset_shared_scratch_pool_slots( + pool, [&](int id, void*, cudaEvent_t, cudaStream_t) { walk_order.push_back(id); }), + 0u) << "a reset of three unlocked slots reported one it could not take"; ASSERT_EQ(walk_order.size(), 3u); @@ -750,7 +838,8 @@ TEST(SharedScratchPoolRegistry, ALeakedLockDoesNotCostTheResetTheSlotsItSharesTh } int disposed = 0; - const std::size_t still_locked = reset_shared_scratch_pool_slots(pool, [&](int, void*, cudaEvent_t) { ++disposed; }); + const std::size_t still_locked = + reset_shared_scratch_pool_slots(pool, [&](int, void*, cudaEvent_t, cudaStream_t) { ++disposed; }); release_it.store(true); holder.join(); @@ -814,8 +903,7 @@ TEST(SharedScratchPoolRegistry, AGrowthOnOneDeviceDoesNotBlockAClaimOnAnother) { SharedScratchDevice& dev1 = pool.get(1); std::lock_guard lk(dev1.mu); RetiredScratch retired; - return shared_scratch_get_or_grow( - dev1, 2048, [&](std::size_t bytes) { return one.alloc(bytes); }, retired); + return shared_scratch_get_or_grow(dev1, 2048, [&](std::size_t bytes) { return one.alloc(bytes); }, retired); }); const bool served = claim.wait_for(std::chrono::seconds(10)) == std::future_status::ready;