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