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..5f43fdf3ce --- /dev/null +++ b/cpp/include/torch_tensorrt/executorch/SharedScratchPool.h @@ -0,0 +1,162 @@ +/* + * 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 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 + +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; +}; + +// 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. 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. +// +// 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(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 {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(SharedScratchDevice& dev) { + if (dev.marker.event != nullptr) { + dev.marker.pending = true; + } + return dev.marker.event; +} + +// 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 buffer is then left untouched. +// Allocating before releasing is what makes that true, and it costs peak +// residency: while the buffer grows, the old and the new one are both resident. +// +// `release(old, wait_for)` frees `old`. A non-null `wait_for` is the marker's +// event, on which an enqueue that may still be reading and writing `old` has been +// recorded; the release must wait for that event on the host before freeing. One +// event covers every enqueue the buffer ever served, but only because each of +// them claims the handoff before enqueueing -- which orders its stream after the +// event -- and records on the event afterwards, so the latest recording completes +// only once all the earlier ones have. An enqueue that reaches the buffer without +// doing both is covered by no wait here. A null `wait_for` means nothing was ever +// recorded against this buffer, so there is nothing to wait for. +template +void* shared_scratch_get_or_grow( + SharedScratchDevice& dev, + std::size_t need, + std::size_t& out_size, + Alloc alloc, + Release release) { + if (dev.buffer != nullptr && dev.capacity >= need) { + out_size = dev.capacity; + return dev.buffer; + } + void* p = alloc(need); + if (p == nullptr) { + return nullptr; + } + if (dev.buffer != nullptr) { + release(dev.buffer, dev.marker.pending ? dev.marker.event : nullptr); + } + dev.buffer = p; + dev.capacity = 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..e164bd01f9 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -75,6 +75,14 @@ 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; + // 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 @@ -102,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( @@ -109,6 +122,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..b5bd5cc526 100644 --- a/cpp/src/torch_tensorrt/executorch/README.md +++ b/cpp/src/torch_tensorrt/executorch/README.md @@ -85,6 +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. 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 @@ -104,6 +114,41 @@ 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 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 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..a362ea7efc 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -6,10 +6,12 @@ */ #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 @@ -17,6 +19,7 @@ #include #include #include +#include #include #include @@ -34,6 +37,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 +156,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,10 +173,23 @@ 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"); + 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; } @@ -204,6 +229,127 @@ 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 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: 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. +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. +// 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 and cudaFree all act on the *current* device and nothing in here +// sets it. +Error get_or_grow_shared_scratch(int device_id, size_t need, cudaStream_t stream, void*& out_ptr, size_t& out_size) { + SharedScratchDevice& dev = scratch_pool.get(device_id); + std::lock_guard lk(dev.mu); + + const SharedScratchHandoff handoff = shared_scratch_claim_event(dev, []() -> 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 bool first_buffer = dev.buffer == nullptr; + void* const buffer = shared_scratch_get_or_grow( + dev, + 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, cudaEvent_t wait_for) { + if (wait_for != nullptr) { + cudaEventSynchronize(wait_for); + } + 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) { + SharedScratchDevice& dev = scratch_pool.get(device_id); + std::lock_guard lk(dev.mu); + + 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; + } + 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 +1052,47 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } // ------------------------------------------------------------------ - // 4. Enqueue inference on the current CUDA stream + // 4. Back activation scratch with the shared per-device pool + // ------------------------------------------------------------------ + // The query requires every input shape to be bound, which they are by here. + // Whatever it answers is binding rather than advisory: setDeviceMemoryV2 refuses + // a smaller buffer, and an engine backed by less than it asked for writes past + // the end. + // + // 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(); + 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; + } + } + + // ------------------------------------------------------------------ + // 5. Enqueue inference on the current CUDA stream // ------------------------------------------------------------------ if (!ctx->enqueueV3(stream)) { ET_LOG( @@ -918,6 +1104,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 +1193,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/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 17d2820bf2..224e270295 100644 --- a/tests/cpp/executorch/BUILD +++ b/tests/cpp/executorch/BUILD @@ -9,6 +9,8 @@ test_suite( ":test_executorch_binding_names", ":test_executorch_blob_header", ":test_executorch_weight_streaming_budget", + ":test_shared_scratch_backend", + ":test_shared_scratch_pool", ], ) @@ -47,3 +49,44 @@ 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", + ], +) + +# 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 new file mode 100644 index 0000000000..0532d0fe40 --- /dev/null +++ b/tests/cpp/executorch/test_shared_scratch_pool.cpp @@ -0,0 +1,485 @@ +// 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 + +namespace torch_tensorrt { +namespace executorch_backend { +namespace { + +// Fake device allocator: hands out distinct non-null pointers and records every +// allocation size and every release, so tests can assert the pool's grow/reuse +// policy and what each release was told to wait for, without a CUDA device. +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, cudaEvent_t wait_for) { + released.emplace_back(p, wait_for); + } + + int alloc_count() const { + return static_cast(alloc_sizes.size()); + } +}; + +// 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(SharedScratchDevice& dev, FakeAllocator& a, std::size_t need, std::size_t& out_size) { + return shared_scratch_get_or_grow( + dev, + need, + out_size, + [&a](std::size_t bytes) { return a.alloc(bytes); }, + [&a](void* p, cudaEvent_t wait_for) { a.release(p, wait_for); }); +} + +TEST(SharedScratchPool, FirstRequestAllocatesExactSize) { + SharedScratchDevice dev; + FakeAllocator a; + std::size_t out = 0; + + void* p = call(dev, a, /*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) { + SharedScratchDevice dev; + FakeAllocator a; + std::size_t out = 0; + + 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(dev, a, 1000, out2); + void* third = call(dev, a, 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) { + SharedScratchDevice dev; + FakeAllocator a; + std::size_t out = 0; + + 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].first, small); + + // A subsequent smaller request reuses the grown buffer -- pool never shrinks. + void* reuse = call(dev, a, 512, out); + EXPECT_EQ(reuse, big); + EXPECT_EQ(out, 8192u); + EXPECT_EQ(a.alloc_count(), 2); +} + +TEST(SharedScratchPool, GrowWaitsOnTheRecordedEnqueueBeforeReleasing) { + SharedScratchDevice dev; + FakeAllocator a; + FakeEventFactory events; + std::size_t out = 0; + + void* small = call(dev, a, 1024, out); + ASSERT_NE(small, nullptr); + + // An enqueue against `small` has been submitted and recorded, so the release + // has something specific to outlive. + const SharedScratchHandoff handoff = shared_scratch_claim_event(dev, std::ref(events)); + ASSERT_EQ(shared_scratch_mark_in_flight(dev), handoff.event); + + ASSERT_NE(call(dev, a, 8192, out), nullptr); + + ASSERT_EQ(a.released.size(), 1u); + EXPECT_EQ(a.released[0].first, small); + // The release is handed the event that enqueue was recorded on, so it waits for + // that enqueue rather than for everything queued on the device. + EXPECT_EQ(a.released[0].second, handoff.event); +} + +TEST(SharedScratchPool, GrowHasNothingToWaitForWhenNoEnqueueWasRecorded) { + SharedScratchDevice dev; + FakeAllocator a; + FakeEventFactory events; + std::size_t out = 0; + + 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(dev, a, 8192, out2); + EXPECT_EQ(failed, nullptr); + EXPECT_TRUE(a.released.empty()); + + // 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) { + SharedScratchDevice dev; + FakeAllocator a; + std::size_t out = 0; + + a.fail_next = true; + 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(dev, a, 1024, out); + EXPECT_NE(q, nullptr); + EXPECT_EQ(a.alloc_count(), 1); +} + +// --------------------------------------------------------------------------- +// Ordering the shared buffer's handoff from one enqueue to the next. +// --------------------------------------------------------------------------- + +TEST(SharedScratchHandoffTest, FirstUseCreatesTheSlotsEventAndWaitsForNothing) { + SharedScratchDevice dev; + FakeEventFactory events; + + const SharedScratchHandoff handoff = shared_scratch_claim_event(dev, std::ref(events)); + + EXPECT_NE(handoff.event, nullptr); + EXPECT_FALSE(handoff.needs_wait); + EXPECT_EQ(events.created, 1); +} + +TEST(SharedScratchHandoffTest, EveryUseAfterAnEnqueueWaitsOnTheSameEvent) { + SharedScratchDevice dev; + FakeEventFactory events; + const SharedScratchHandoff first = shared_scratch_claim_event(dev, std::ref(events)); + ASSERT_FALSE(first.needs_wait); + + 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(dev, std::ref(events)); + EXPECT_TRUE(second.needs_wait); + EXPECT_EQ(second.event, first.event); + + const SharedScratchHandoff third = shared_scratch_claim_event(dev, 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) { + SharedScratchPool pool; + FakeEventFactory events; + 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 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(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) { + SharedScratchDevice dev; + FakeEventFactory events; + + events.fail_next = true; + 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(dev, std::ref(events)); + EXPECT_NE(retried.event, nullptr); + EXPECT_FALSE(retried.needs_wait); + EXPECT_EQ(events.created, 1); +} + +TEST(SharedScratchHandoffTest, ASlotWithNoEventIsNotMarked) { + 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(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(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 +} // 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({