diff --git a/.github/workflows/executorch-build-linux.yml b/.github/workflows/executorch-build-linux.yml index f3eab9c237..33fb1071e8 100644 --- a/.github/workflows/executorch-build-linux.yml +++ b/.github/workflows/executorch-build-linux.yml @@ -101,12 +101,33 @@ jobs: python -m pip wheel --no-build-isolation --no-deps --wheel-dir dist py/torch-tensorrt-executorch-runtime # this is to build the libtorchtrt.tar.gz bazel build //:libtorchtrt --compilation_mode opt --config=linux + executorch_cmake_location="$(bazel query @executorch//:executorch/CMakeLists.txt --output=location)" + export EXECUTORCH_SOURCE_DIR="$(dirname "${executorch_cmake_location%%:*}")" + export EXECUTORCH_ROOT="${EXECUTORCH_SOURCE_DIR}" + # this is to verify the end user's workflow + python -m pip install pyyaml "executorch>=1.4.1,<1.5" + python examples/torchtrt_executorch_example/export_static_shape.py \ + --model_path="${RUNNER_TEMP}/torchtrt-reference-runner.pte" + .github/scripts/verify-executorch-reference-runner.sh \ + "${RUNNER_TEMP}/torchtrt-reference-runner.pte" # Run the ExecuTorch backend C++ unit tests, which are otherwise only - # built, never executed. These tests link TensorRT and CUDA, but Bazel's - # cc_import ships the unversioned libnvinfer.so and libcudart.so while the - # loader asks for the versioned sonames, so add the directories holding - # those to LD_LIBRARY_PATH. Append rather than replace: the toolchain - # entries already present are still needed. + # built, never executed. + # + # Last in this script, and not next to the Bazel build above, because + # verify-executorch-reference-runner.sh between the two is the only check + # anywhere in CI that the release archive still carries every header the + # packaged CMake backend needs, and that the backend still builds out of it. + # `set -e` ends the script at the first failure, so a GPU suite placed above + # that check trades it away whenever a device is missing or flaky. Run last, + # a failure here no longer costs that check. It still costs the wheel: + # linux-test.yml guards its upload step on `inputs.upload-artifact` alone, + # with no always(), so any failure in this script skips the upload. + # + # These tests link TensorRT and CUDA, but Bazel's cc_import ships the + # unversioned libnvinfer.so and libcudart.so while the loader asks for the + # versioned sonames, so add the directories holding those to + # LD_LIBRARY_PATH. Append rather than replace: the toolchain entries + # already present are still needed. bazel_external="$(bazel info output_base)/external" for _soname in libnvinfer.so libcudart.so; do _dir="$( @@ -119,15 +140,23 @@ jobs: fi export LD_LIBRARY_PATH="${_dir}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" done + # --test_summary=detailed names every case that skipped, in Bazel's own + # summary rather than in the test log, which --test_output=errors prints + # nothing of for a target that passes. Cases in that suite skip for reasons + # the variable below does not cover -- its own coverage note lists them -- + # and without those names a run in which the pool's growth and + # allocation-failure paths never executed looks like one in which they did. + # + # TORCHTRT_EXECUTORCH_REQUIRE_CUDA turns a case that skips for want of a + # CUDA device into a failure. Every row of this job's matrix is a CUDA row + # -- filter-matrix.py drops any whose desired_cuda is not a TensorRT CUDA + # version -- and linux-test.yml starts the container with `--gpus all` for + # those, so a skip here means the device went missing rather than that the + # runner never had one. Without the variable the backend suite would skip + # every case, exit zero and be reported as a passing target that covered + # nothing. bazel test //tests/cpp/executorch:executorch_backend_tests \ --compilation_mode opt --config=linux --test_output=errors \ - --test_env=LD_LIBRARY_PATH - executorch_cmake_location="$(bazel query @executorch//:executorch/CMakeLists.txt --output=location)" - export EXECUTORCH_SOURCE_DIR="$(dirname "${executorch_cmake_location%%:*}")" - export EXECUTORCH_ROOT="${EXECUTORCH_SOURCE_DIR}" - # this is to verify the end user's workflow - python -m pip install pyyaml "executorch>=1.4.1,<1.5" - python examples/torchtrt_executorch_example/export_static_shape.py \ - --model_path="${RUNNER_TEMP}/torchtrt-reference-runner.pte" - .github/scripts/verify-executorch-reference-runner.sh \ - "${RUNNER_TEMP}/torchtrt-reference-runner.pte" + --test_summary=detailed \ + --test_env=LD_LIBRARY_PATH \ + --test_env=TORCHTRT_EXECUTORCH_REQUIRE_CUDA=1 diff --git a/cpp/BUILD b/cpp/BUILD index 30619cda92..1335303cfc 100644 --- a/cpp/BUILD +++ b/cpp/BUILD @@ -191,6 +191,113 @@ cc_library( ], ) +# Implementation detail of :tensorrt_executorch_backend, so the header sits under +# src/ and ships with the sources rather than in the installed API set. A caller +# turning the pool on needs only the option key, which is in TensorRTBackend.h. +cc_library( + name = "tensorrt_executorch_shared_scratch_pool", + hdrs = [ + "src/torch_tensorrt/executorch/SharedScratchPool.h", + ], + strip_include_prefix = "src", + 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": [], + }), +) + +# The declaration of the one internal entry point of :tensorrt_executorch_backend +# that its test calls directly. Under src/ for the same reason as the pool header +# above: it is an implementation detail of execute(), and an installed header +# declaring it would make it API from the next release. +cc_library( + name = "tensorrt_executorch_pooled_scratch_install", + hdrs = [ + "src/torch_tensorrt/executorch/PooledScratchInstall.h", + ], + strip_include_prefix = "src", + target_compatible_with = select({ + ":linux_x86_64": [], + ":sbsa": [], + "//conditions:default": ["@platforms//:incompatible"], + }), + deps = select({ + ":linux_x86_64": [ + "@tensorrt//:nvinfer", + ], + ":sbsa": [ + "@tensorrt_sbsa//:nvinfer", + ], + "//conditions:default": [], + }), +) + +# The body of the pool's test-only reset. testonly and under src/, so it is in +# neither the packaged include tree nor executorch_backend_source_files below: a +# release build has no definition of it anywhere. SharedScratchPool.h keeps only +# the friend declaration that lets this reach the registry's internals. +cc_library( + name = "tensorrt_executorch_shared_scratch_pool_reset", + testonly = True, + hdrs = [ + "src/torch_tensorrt/executorch/SharedScratchPoolReset.h", + ], + strip_include_prefix = "src", + target_compatible_with = select({ + ":linux_x86_64": [], + ":sbsa": [], + "//conditions:default": ["@platforms//:incompatible"], + }), + deps = [ + ":tensorrt_executorch_shared_scratch_pool", + ], +) + +# The pool's test-only entry points. A separate target, and testonly, so the +# released archive -- Bazel's :tensorrt_executorch_backend and CMake's +# executorch_trt_backend, neither of which compiles this source -- carries no +# definition of them and exports no symbol for them. One of them frees the live +# pool with no wait for work in flight, which is not something a release build +# should offer a caller. +cc_library( + name = "tensorrt_executorch_shared_scratch_pool_test_hooks", + testonly = True, + srcs = [ + "src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp", + ], + hdrs = [ + "src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.h", + ], + strip_include_prefix = "src", + target_compatible_with = select({ + ":linux_x86_64": [], + ":sbsa": [], + "//conditions:default": ["@platforms//:incompatible"], + }), + deps = [ + ":tensorrt_executorch_shared_scratch_pool", + ":tensorrt_executorch_shared_scratch_pool_reset", + ] + select({ + ":linux_x86_64": [ + "@cuda//:cudart", + ], + ":sbsa": [ + "@cuda//:cudart", + ], + "//conditions:default": [], + }), +) + cc_library( name = "tensorrt_executorch_backend", srcs = [ @@ -211,6 +318,8 @@ cc_library( deps = [ ":tensorrt_executorch_binding_names", ":tensorrt_executorch_blob_header", + ":tensorrt_executorch_pooled_scratch_install", + ":tensorrt_executorch_shared_scratch_pool", ":tensorrt_executorch_weight_streaming_budget", ] + select({ ":linux_x86_64": [ @@ -234,7 +343,9 @@ filegroup( name = "executorch_backend_source_files", srcs = [ "src/torch_tensorrt/executorch/CMakeLists.txt", + "src/torch_tensorrt/executorch/PooledScratchInstall.h", "src/torch_tensorrt/executorch/README.md", + "src/torch_tensorrt/executorch/SharedScratchPool.h", "src/torch_tensorrt/executorch/TensorRTBackend.cpp", "src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp", "src/torch_tensorrt/executorch/WeightStreamingBudget.cpp", diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index b33d712d40..c9bbe6ec57 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -17,6 +17,7 @@ #include +#include #include #include #include @@ -75,6 +76,20 @@ struct EngineHandle { size_t num_aliased_outputs = 0; int device_id = 0; bool unified_memory = false; + // Whether exec_ctx was created kUSER_MANAGED, because + // kSharedActivationScratchKey was on at this engine's load. Such a context takes + // its activation scratch from the shared per-device pool on every call, unless + // engine_needs_scratch below is false, in which case it needs none and never + // claims from the pool at all. + bool shared_scratch = false; + // Whether the engine reports needing activation scratch under any shape, read + // at init when shared_scratch is set. It is a predicate and not a size because + // the size is never the right amount to install: enqueueV3 refuses a + // kUSER_MANAGED context with no device memory once this is true, however little + // the shapes bound to a given call need, so execute() must hand such a call + // some buffer -- but the engine's own figure covers every shape in the profile, + // and installing it would size the pool for the largest of them. + bool engine_needs_scratch = false; std::mutex mu; // Makes the skip-sync fast path safe to reuse: TensorRT forbids reconfiguring or // destroying an execution context while one of its enqueues is in flight, so when @@ -87,6 +102,15 @@ struct EngineHandle { ~EngineHandle(); }; +// Runtime backend option that backs execution-context activation scratch with a +// shared per-device pool instead of giving every context its own. Boolean, +// default false. Read by TensorRTBackend::set_option below, and delivered as +// executorch::runtime::set_option("TensorRTBackend", options.view()) +// A context's allocation strategy is fixed when the context is created, so a +// later call governs only the engines loaded after it, and a pooled context and +// a private-scratch one coexist in one process. +inline constexpr char kSharedActivationScratchKey[] = "use_shared_activation_scratch"; + class TensorRTBackend final : public ::executorch::runtime::BackendInterface { public: bool is_available() const override; @@ -102,13 +126,118 @@ 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. - // Note that other CUDA delegates sharing the same guard may instead synchronize before + // Other CUDA delegates sharing the same guard may instead synchronize before // returning, so do not assume results are ready on return from this one. + // With the shared activation scratch pool (kSharedActivationScratchKey) one + // buffer per device backs every context created while the option was on. Calls + // on two such handles on one device may overlap: a per-device lock held across + // the enqueue serializes them, so they do not run concurrently on the device. A + // handle whose context was created while the option was off keeps its own + // scratch and is not subject to this, nor is one whose engine reports needing no + // activation scratch under any shape, which is left out of the pool. For the + // handles that do draw on it, five further consequences: + // - A call needing more scratch than the pool holds grows it, and the growth + // has to get rid of the buffer it replaces before it returns, or the bytes + // of every size the pool ever grew to stay resident at once. What that costs + // depends on whether this call is one that synchronizes the stream. + // A call that does -- anything staging through host memory, aliasing an + // output, or running with no caller stream -- queues a stream-ordered + // cudaFreeAsync and its own synchronization returns the bytes, waiting for + // nothing it did not submit. A call that does not -- the caller-stream, + // device-resident case the rest of this contract is about -- makes a host + // wait on the previous enqueue and a device-wide cudaFree instead, because a + // queued free on a stream nothing ever synchronizes never returns the bytes + // at all. So that one call blocks until the device is idle however + // asynchronous the rest of this contract makes it -- an unbounded wait on + // work this call did not submit. Unbounded is meant literally: if any of + // that work is itself waiting on something only this thread supplies once + // execute() returns -- a host function it will release, a copy it will + // enqueue next -- the call does not return, and the thread that would + // unblock it is the one inside cudaFree. The same wait is also the fallback + // for the first case, on a device with no stream-ordered allocator, where + // cudaFreeAsync reports cudaErrorNotSupported rather than freeing, and on + // any other code it reports, since a queued free that did not happen is not + // a free. It is also what a call that fails after the growth takes -- a + // refused enqueue, or a failed record of one already submitted -- whichever + // kind of call it is, because by the time such a call disposes of the + // buffer it has no synchronization of that stream left to make, so a queued + // free's bytes would never come back. Which calls grow the pool is not + // knowable from here; see the README. + // - A call that synchronizes the stream waits, through the handoff, for every + // pooled enqueue submitted before it on that device. That is what one buffer + // costs, and it is another unbounded wait: park a host function ahead of a + // pooled enqueue and release it only after a later execute() returns, and + // the later call is the one that does not return. + // - Capturing a CUDA graph around this delegate is not supported, with the + // option on or off. With it on, a call whose selected stream is capturing is + // refused with Error::NotSupported, ahead of every CUDA call it makes that a + // capture cannot take -- only the device query and the device switch run + // first. The pool's event handoff waits on an event recorded outside the + // capture, which invalidates it under every capture mode. A growth's + // allocation invalidates it under every mode but cudaStreamCaptureModeRelaxed, + // which permits it but does not record it, leaving a replay pointed at a + // buffer the pool may since have freed; and a growth's disposal of the buffer + // it replaces is prohibited too -- a call that synchronizes the stream is one + // no capture mode permits, so the only disposal a capture could reach is the + // host wait and the device-wide cudaFree, which a capture cannot take outside + // Relaxed either. The alternative to refusing is a capture that silently + // comes back invalidated. Only a capture on the selected stream is caught. A + // capture running on any other stream under cudaStreamCaptureModeGlobal, or + // under cudaStreamCaptureModeThreadLocal from this thread, is invalidated by + // the same calls and is not refused, because CUDA offers no query for it: do + // not run a pooled engine while capturing anywhere in the process. + // With the option off nothing refuses, because the check is inside the + // pooled path, and a call can then be captured: one submitted on the + // capturing stream through a CallerStreamGuard that binds only non-empty + // device-resident tensors, aliases no output and follows no call that left + // an enqueue in flight makes no call a capture cannot take, and the graph + // replays the engine. It costs the handle. Such a call ends by recording the + // completion event this contract describes, and under capture that record is + // a node of the graph rather than an event the host can wait on, so the next + // execute() on the handle fails Error::InvalidProgram on the wait for it, + // and so does the destructor's wait. That failing call clears the flag + // before it returns, so the call after it records the event again outside + // the capture and runs -- but nothing waits for a replay of the graph, which + // enqueues on the same context, and reconfiguring a context under a live + // enqueue is what the event was there to prevent. A call that ends by + // synchronizing the stream instead is one no capture mode permits, so there + // is no call shape that captures and leaves the handle safe to go on using. + // The README states this in full. + // - cudaDeviceReset() invalidates the pool without emptying it. The buffer + // and the handoff event it still holds are destroyed with the primary + // context, and the next call on that device uses both. There is no guard: + // do not reset a device this backend has run a pooled engine on. + // - The pool adds six failure points to a call, on top of Error::NotSupported + // for the capture above. Only one of them has a code of its own: + // Error::Internal, which a pooled call returns when the device's handoff + // event cannot be created and an unpooled call never returns at all. The + // other five reuse codes an unpooled call already returns from elsewhere in + // execute(), so the code alone does not say the pool was involved -- the log + // line does, and every one of them logs at Error first. Claiming the + // device's scratch answers Error::MemoryAllocationFailed if the buffer + // cannot be allocated or grown and Error::InvalidState if the wait ordering + // this call behind the previous enqueue fails; installing the buffer answers + // Error::InvalidState if TensorRT refuses it, and recording this call's + // enqueue for the next claimant answers the same. Error::InvalidProgram + // means the disposal's host wait failed, which leaks the retired buffer + // rather than freeing it under an enqueue that may still be reading it. The + // last two are the only ones reached after this call's enqueue is submitted, + // and both synchronize the stream before returning, so neither leaves engine + // work in flight -- and the record failure makes that wait with the device's + // pool lock still held, which is the one place a pooled call holds it across + // a host wait. ::executorch::runtime::Error execute( ::executorch::runtime::BackendExecutionContext& context, ::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, 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/CMakeLists.txt b/cpp/src/torch_tensorrt/executorch/CMakeLists.txt index 1b503c567d..ec75e2e18c 100644 --- a/cpp/src/torch_tensorrt/executorch/CMakeLists.txt +++ b/cpp/src/torch_tensorrt/executorch/CMakeLists.txt @@ -36,6 +36,13 @@ set_target_properties(executorch_trt_backend target_include_directories(executorch_trt_backend PUBLIC "${CMAKE_CURRENT_LIST_DIR}/../../../include" + PRIVATE + # SharedScratchPool.h and PooledScratchInstall.h are implementation + # details, so they sit beside the sources instead of in the installed + # include tree. This directory is + # .../torch_tensorrt/executorch in both the repository and the released + # package, so two levels up is the root its #include path is relative to. + "${CMAKE_CURRENT_LIST_DIR}/../.." ) get_filename_component(_torchtrt_repo_root "${CMAKE_CURRENT_LIST_DIR}/../../../.." ABSOLUTE) diff --git a/cpp/src/torch_tensorrt/executorch/PooledScratchInstall.h b/cpp/src/torch_tensorrt/executorch/PooledScratchInstall.h new file mode 100644 index 0000000000..8370801218 --- /dev/null +++ b/cpp/src/torch_tensorrt/executorch/PooledScratchInstall.h @@ -0,0 +1,42 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +// The one step of execute()'s pooled path that the backend's test drives +// directly, declared beside the source that defines it rather than in the +// installed header set: it is an implementation detail of execute(), and a +// declaration in an installed header is API the next release has to keep. +// SharedScratchPool.h sits here for the same reason. + +#include + +#include + +namespace torch_tensorrt { +namespace executorch_backend { + +// Installs `bytes` of `buffer` as the activation scratch of `ctx`, a +// kUSER_MANAGED context, and reports whether TensorRT accepted it. Logs the +// refusal, naming `device_id`, when it did not. +// +// The check is the point. setDeviceMemoryV2 returns void and refuses a buffer +// smaller than the bound shapes need, so a caller that does not ask cannot tell an +// accepted install from a refused one -- and a refused one leaves the context +// pointed at the buffer it was last given, which a shared-scratch pool growth may +// since have freed. The engine then reads and writes freed memory with enqueueV3 +// reporting success. The refusal is read back through an IErrorRecorder scoped to +// this one call, the only channel that hands it to the caller: with no recorder +// attached TensorRT writes it to the runtime's ILogger and this function has +// nothing to return. +// +// The backend's own test calls it over the same TensorRT call to cover a refusal +// that cannot be induced through execute(). +bool install_pooled_scratch(nvinfer1::IExecutionContext& ctx, void* buffer, std::size_t bytes, int device_id); + +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/cpp/src/torch_tensorrt/executorch/README.md b/cpp/src/torch_tensorrt/executorch/README.md index e7367f8706..7e40c5efef 100644 --- a/cpp/src/torch_tensorrt/executorch/README.md +++ b/cpp/src/torch_tensorrt/executorch/README.md @@ -68,8 +68,7 @@ type. The upstream `CallerStreamGuard` documents the generic contract (per-thread, nested scoping; the caller owns the stream for the guard's lifetime; the caller manages host-data lifetime for async work). The TensorRT backend adds these -requirements, which previously lived on -the removed `CudaStreamGuard`: +requirements, which previously lived on the removed `CudaStreamGuard`: - The selected stream must be on the TensorRT engine's device. - Calls using one delegate handle must not overlap, and must not overlap with @@ -85,6 +84,18 @@ 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 no two enqueues against it may overlap. The backend + enforces this itself: it holds a per-device lock from the claim on the buffer + through the enqueue and the completion event recorded on it, so two + `execute()` calls on one device are serialized at submission and the second's + stream waits on the first's enqueue. They may run on one stream or on two, and + they may be submitted concurrently from two threads, but they will not run + concurrently on the device, so the pool costs the parallelism between them. + Contexts created while the option was off keep their own scratch and are + unaffected, and so is a context whose engine needs no activation scratch under + any shape: that one is left out of the pool and serializes against nothing. - The reference-runner smoke test runs inference inside a caller-stream guard on the discrete-GPU CI configuration, where all inputs and outputs are host-backed and therefore take the synchronized staging path. CI separately asserts that the @@ -104,6 +115,301 @@ 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 a device's contexts from one buffer, grown to the largest +requirement any call on that device has asked for. Which contexts: those created +while the option was on, less any whose engine needs no scratch under any shape, +which are left out of the pool entirely. The two exceptions are spelled out +below. + +```cpp +#include + +executorch::runtime::Error enable_shared_activation_scratch() { + executorch::runtime::BackendOptions<1> options; + // Returns Error::InvalidArgument when the object has no room left. One key in a + // BackendOptions<1> always fits, so this cannot fail as written; the check is + // what keeps it honest when a second key is added beside it. + const executorch::runtime::Error stored = + options.set_option("use_shared_activation_scratch", true); + if (stored != executorch::runtime::Error::Ok) { + return stored; + } + // Nothing was set if this fails: every context still allocates its own scratch. + return executorch::runtime::set_option("TensorRTBackend", options.view()); +} +``` + +`Error::NotFound` means no backend is registered under that name, which is what a +binary that has not linked the backend archive gets. Nothing forces the check: the +free `executorch::runtime::set_option` is not `ET_NODISCARD`, so dropping its +return compiles. + +N per-engine copies collapse to one, so the reclaimed memory is the sum of the N +requirements less the largest of them. Set the option before loading the methods +whose engines should use the pool, and read the `use_shared_activation_scratch` +bullet of the caller-stream contract above: engines sharing a buffer do not run +concurrently on the device. The pool never shrinks, so the largest scratch it was +ever asked for stays allocated until the process exits. + +A call that asks for nothing is handed whatever buffer the pool already holds and +never grows it. It is not handed the pool's size with it: what the context is +told it owns is the figure this call asked for, so a zero becomes the one-byte +minimum below and TensorRT is told the buffer is one byte, whatever the pool's +capacity is. Two of the things a zero from the per-shape +query can mean reach this point, and are indistinguishable where it is read: the +shapes bound to this call need none -- an empty batch inside a profile that admits +one -- or the query failed. Neither wants a buffer of its own. + +Such a call cannot be left with no buffer at all. `enqueueV3` refuses a +`kUSER_MANAGED` context with no device memory installed as soon as the engine +reports needing some under *any* shape, however little the shapes actually bound +need. So where the pool holds nothing yet, a zero allocates the smallest buffer +that satisfies that check, and the first call with a real requirement grows it. +Standing the engine's own figure in for the zero would satisfy the check too, but +that figure covers the whole profile: on the dynamic engine the tests use it is +4 MiB against the 128 KiB the next call needs, and the pool never shrinks. + +Only one of the two causes is safe to hand a buffer that small, and the install is +what separates them. `setDeviceMemoryV2` refuses a buffer smaller than the bound +shapes need, and it returns `void`: where the shapes genuinely need nothing the +expected size is zero and the one-byte buffer is accepted, but where the query +failed the engine expects what it always did, the install is refused, and the +context keeps the buffer it was last given -- which a growth may already have +freed, so the enqueue would read and write memory the pool no longer owns while +`enqueueV3` reports success. The backend reads the refusal back through a TensorRT +`IErrorRecorder` scoped to that one call and fails with `Error::InvalidState` +rather than enqueueing. Nothing reaches `enqueueV3` on an install TensorRT +rejected. + +What each call installs is its own requirement, not the capacity the pool holds, +which is larger whenever an earlier call asked for more. Handing a context the +larger figure would say it owns bytes that belong to whatever ran before it, and +the pool never clears the buffer between users -- activation scratch is written +before it is read, so a context that stays inside its own requirement never sees +those bytes, and one that runs past it sees another engine's activations rather +than zeros. The smaller figure also keeps the refusal above sharp: after a growth +the capacity may well cover what a failed query concealed, and an install sized +from it would be accepted. + +The third thing a zero can mean is an engine that needs no scratch under *any* +shape, and that one is settled before a call is ever made: such an engine is left +out of the pool entirely, so it takes no per-device lock and does not serialize +against the engines that do. + +The buffer grows when a call asks for more than every call before it did, and the +growth has to get rid of the buffer it replaces *before `execute()` returns*. +Not disposing of it inside the call is not an option: the pool grows monotonically, +so a run that let each retired buffer outlive its growth would end up holding the +sum of every size the pool was ever grown to rather than the largest of them. + +There are two ways to get rid of it, and which one a call uses is decided by +whether that call synchronizes the stream before it returns -- that is, whether it +stages anything through host memory, aliases an output, or runs with no caller +stream. Both are made with the per-device lock dropped, so neither is serialized +behind the pool's own lock and neither holds the next claimant off while it runs: +two pooled calls serialize at submission, as the caller-stream contract above says, +and a growth is not an exception to it. That is not the same as holding nobody up, +and the paragraph below on the device-wide free says what it does hold up. + +Which of the two a call uses is settled where the buffer is disposed of and not +where it is claimed, because it is a claim about what the call has still to do. A +call that fails after the growth -- a refused `enqueueV3`, or a failed record of an +enqueue already submitted -- takes the host wait whichever kind of call it is. The +failed record does drain the stream, but it drains and then returns, so the +disposal falls after the last synchronization that call makes; the refused enqueue +makes none at all. A free queued on a stream that is never synchronized again does +not come back: not at `cudaStreamDestroy`, not at `cudaDeviceSynchronize`, not from +another stream, and not at `cudaMemPoolTrimTo`. Measured on an A100 with CUDA 13.0, +a 4 GiB block queued that way and then abandoned with the stream is gone for the +life of the process, on a device with tens of gigabytes still free. + +**A call that synchronizes** queues the free: `cudaFreeAsync` on the stream it +enqueued on. Ordering is what makes that safe, not a wait -- every pooled call +makes its stream wait on the device's handoff event before it enqueues and records +on that event afterwards, so a free queued on this call's stream sits behind every +enqueue that ever used the retired buffer, and behind this call's own, which uses +the new one. The call's own `cudaStreamSynchronize` is then what returns the +bytes, and nothing waits for work the call did not submit. Measured on an A100 +with CUDA 13.0 and driver 13.0, with a host function parked for 1500 ms on a stream +neither the pool nor the growing engine had ever used: around 10 ms, against the +1505 ms the disposal below took, which returned exactly when that host function was +released. + +**A call that does not synchronize** -- the caller-stream, device-resident case, +which is the one this option exists for -- makes a host wait on the handoff event +and then a device-wide `cudaFree`. It cannot queue the free, because a queued free +does not return the bytes when the stream reaches it. Measured on the same +machine: `cudaFreeAsync` on a `cudaMalloc`'d pointer returns `cudaSuccess` and +defers the free, but the bytes come back at the next `cudaStreamSynchronize` of +that stream and at no point before it. A stream `cudaStreamQuery` reports as +drained still holds them, `cudaMemGetInfo` does not move, and a `cudaMalloc` of the +same size fails with out of memory until the synchronize. On this path nothing +this backend does will ever synchronize that stream, so replaying the pool's own +growth sequence on one stream -- 8, 16, 24 and 32 GiB, no synchronize between +them -- four growths that all succeed with `cudaFree` fail on the fourth with out +of memory with `cudaFreeAsync`. + +The same host wait and `cudaFree` are also the fallback for the first case: where a +device has no stream-ordered allocator (`cudaDevAttrMemoryPoolsSupported`), +`cudaFreeAsync` reports `cudaErrorNotSupported` instead of freeing -- and on any +other code it reports, since a queued free that did not happen is not a free. + +**That `cudaFree` is the thing to know about**, because it is what every growth on +the asynchronous path pays. It waits for everything queued on the device, not only +for the enqueues that used the buffer, so on such a growth the asynchronous return +this backend otherwise promises does not happen: the call blocks until the device +is idle. The wait has no upper bound, and a caller can +turn it into one that never ends. Anything queued anywhere on the device is enough +to hold it, so if some stream is waiting on work only this thread will submit -- a +host function it will release after `execute()` returns, a copy it will enqueue +next -- the growing call does not return and the thread that would unblock it is +the one inside `cudaFree`. This is not hypothetical: the backend's own tests +deadlocked on it once, when a change of test order turned a case that parks a host +function on its own stream into the one that grew the pool. A program that parks +work like that has to keep the pool's growths away from it, and the paragraph +below on how often the pool grows is what says whether a run order can do that. + +Sharing one buffer has an unbounded wait of its own, growth or no growth. A call +that synchronizes waits, through the handoff, for every pooled enqueue submitted +before it on that device. Park a host function ahead of one pooled call's enqueue +and release it only once a later pooled call has returned, and the later call is +the one that never returns. + +What an engine answers when asked how much it needs is decided when it is built, +not when it runs. The builder's `kRUNTIME_ACTIVATION_RESIZE_10_10` preview feature +makes an engine report what the shapes just bound need; without it, whether an +engine does that or reports its profile maximum depends on how TensorRT planned +it. Either way the pool can settle well above the live data, and nothing the +runtime does changes it. + +How often the pool grows follows from that. The backend asks afresh on every +`execute()`, after the input shapes are bound, so an engine whose answer does not +vary with the bound shapes grows the pool at most once, on its first run. For a +program built only from such engines, running the largest one first leaves the +pool with a single allocation. An engine whose answer does vary can grow it on any +call whose shapes need more than every call before them, so with one of those in +the program no run order bounds the number of allocations. One engine over a +`[1..4, 512, 512]` profile is either, depending on how it was built. + +### CUDA graph capture is not supported + +A pooled call whose selected stream is capturing is refused with +`Error::NotSupported`. Capture with the option off is not a way round that, for a +different reason given at the end of this section, so the refusal is a diagnosis +and not a redirection. + +The handoff between one enqueue and the next waits on an event recorded outside +the capture, and `cudaStreamWaitEvent` on such an event fails with +`cudaErrorStreamCaptureIsolation` and invalidates the capture under +every capture mode, `cudaStreamCaptureModeRelaxed` included; the caller learns of +that only when `cudaStreamEndCapture` hands back `cudaErrorStreamCaptureInvalidated` +and a null graph, long after the call that caused it. A growth adds more of the +same. Its `cudaMalloc` returns `cudaErrorStreamCaptureUnsupported` and invalidates +the capture under `Global` and `ThreadLocal`, and under `Relaxed` is permitted but +runs uncaptured, so a replayed graph would use whatever buffer was installed when +it was captured -- which by then the pool may have freed. So does its disposal of +the buffer it replaces. A call that synchronizes the stream is one no capture mode +permits, so the only disposal a capture could reach is the host wait on the handoff +event and the `cudaFree`, and those two return `cudaErrorStreamCaptureUnsupported` +and invalidate the capture under `Global` and `ThreadLocal`, and under `Relaxed` +are permitted but run uncaptured. The queued free would be no way round it either: +a `cudaFreeAsync` made while capturing is refused under all three modes, with +`cudaErrorInvalidValue` and without invalidating the capture -- measured; a free +made while capturing has to be given a graph allocation, and the pool's buffers +come from `cudaMalloc`. + +The backend checks the selected stream and returns instead, and it +checks ahead of everything a capture cannot take -- not merely ahead of the pool's +own calls, since the wait on a previous enqueue and the `cudaMalloc` that grows a +host-input staging buffer come before those and, outside `Relaxed`, would +invalidate the capture first. Only the device query and the device switch run +before the check, and a capture takes both. + +What that check does not cover is a capture running on some *other* stream. Under +`Global`, and under `ThreadLocal` from the thread that calls `execute()`, the same +calls are prohibited and invalidate that capture, and the call is not refused: +`cudaStreamIsCapturing` answers about one stream, and asking it about +`cudaStreamLegacy` instead does not close the gap. Measured on CUDA 13.0, that +reports the capture only when the capturing stream is a blocking one, and then +reports it under `Relaxed` and under another thread's `ThreadLocal` as well, where +the pool's calls are permitted and refusing would be wrong. CUDA has no query for +"is a capture live in this process", so the rule is the caller's to keep: do not +run a pooled engine while any capture is open anywhere in the process. + +**With the option off nothing refuses, and capture still does not work.** +Turning the option off does take the pool's calls out of the way, and the refusal +with them, because the check lives inside the pooled path. What it leaves is a +capture that succeeds, a handle whose next call fails, and a graph whose replays +nothing orders against the delegate's own calls. + +A call captures cleanly with the option off when it runs on the capturing stream +under a `CallerStreamGuard`, binds only non-empty device-resident tensors, has no +aliased outputs, and follows no call that left an enqueue in flight. A first call +on a freshly loaded handle over device-resident inputs and outputs meets all four: +`execute()` then makes only the device query, the device switch, the shape and +address binding, `enqueueV3` and the completion record, and none of those is a +call a capture cannot take. Measured, not argued -- the graph instantiates, +replays, and reproduces the bytes the same engine produces outside a capture. + +The completion record is what costs the handle. `execute()` records that event so +the next call and `~EngineHandle` can wait for an enqueue that outlived the +return; under capture the record becomes a node of the graph rather than an event +the host can wait on. The next `execute()` on that handle fails +`Error::InvalidProgram` on the wait at its head -- measured, logging +`cudaEventSynchronize failed: invalid argument` -- and so does the destructor's +wait. No call shape avoids this: the alternative to recording the event is the +`cudaStreamSynchronize` that ends any call staging through host memory, aliasing +an output or running with no caller stream, and a capturing stream cannot be +synchronized under any mode, `Relaxed` included. So the only call a capture can +take is the one that arms the event. + +It is one call that fails, not every call after the capture. The failing call +clears the in-flight flag before returning the error, so the call after it skips +the wait, records the event again outside any capture, and runs -- measured, +producing what the same engine produces uncaptured. That recovery is not a reason +to capture. What the handle does not get back is the ordering the event carried +over the captured work: a `cudaGraphLaunch` of that graph enqueues on the same +execution context, no `execute()` waits for it, and TensorRT forbids +reconfiguring a context while one of its enqueues is in flight -- so a caller who +replays the graph and calls `execute()` again is reconfiguring a live context with +nothing to say when it stops being live. +`AnUnpooledEngineCaptureReplaysAndCostsTheHandleItsNextCall` in +`tests/cpp/executorch/test_shared_scratch_backend.cpp` is that whole sequence. + +Two further calls invalidate a capture with the option off, before the handle +question arises, and nothing refuses them either: the host wait on a previous +enqueue, and a staging buffer's `cudaMalloc` or the `cudaFree` of a cached buffer +a larger tensor has outgrown. Both are prohibited under `Global` and +`ThreadLocal` and permitted under `Relaxed`, where they run for real and outside +the graph. +`AnUnpooledEngineInvalidatesACaptureOnceAnEnqueueIsInFlight` covers the first. + +Two changes would be needed to make capture work rather than documenting it, and +neither is local. The completion event would have to be reworked, and the +asynchronous return described in the caller-stream contract above rests on it. +And moving the pool's check out of the pooled path so that capture were refused +both ways would refuse the calls above, which do capture, for contexts that never +touch the pool -- a behaviour change for callers who have not turned the option +on. So the option-off path stays as it is, and this section, the `execute()` +contract in `TensorRTBackend.h` and the refusal message all say that capture is +unsupported rather than pointing a caller at it. + +### cudaDeviceReset() is not survivable + +The pool is not guarded against it: it holds its buffer and its handoff event for +the process lifetime, and a reset destroys the primary context under both. The +next pooled `execute()` on that device then waits on a destroyed event and hands +the engine a pointer that is no longer a device allocation. Guarding this would +mean revalidating both on every call, and the check that would catch it is as +expensive as the work it protects. Treat a device the backend has run a pooled +engine on as one that must not be reset. + ## Standalone Backend Archive Use this path only when you need `libexecutorch_trt_backend.a` without building diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h new file mode 100644 index 0000000000..9fb538036b --- /dev/null +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPool.h @@ -0,0 +1,251 @@ +/* + * 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 +#include + +namespace torch_tensorrt { +namespace executorch_backend { + +// Per-device handoff marker for the shared scratch buffer: the pool-owned CUDA +// event that the last enqueue against the buffer was recorded on. +struct SharedScratchMarker { + // Nothing in the normal path destroys this event; only the test-only reset + // hands it to a disposer that does. + cudaEvent_t event = nullptr; + bool pending = false; // an enqueue against the buffer has been recorded on `event` +}; + +// 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 its own +// enqueue and the record of that enqueue on the marker. Holding it that far is +// what makes the marker a complete account of who is using the buffer. Anything +// less leaves an enqueue live in a window the marker does not cover, and a +// claimant entering that window is handed the same buffer with nothing ordering +// the two. `mu` covers one device, so a growth holds no lock a claim on +// another device has to acquire. +struct SharedScratchDevice { + std::mutex mu; + void* buffer = nullptr; + std::size_t capacity = 0; + SharedScratchMarker marker; +}; + +class SharedScratchPool; + +// Defined in SharedScratchPoolReset.h; see the friend declaration inside +// SharedScratchPool for why the body is not here. +template +std::size_t reset_shared_scratch_pool_slots(SharedScratchPool& pool, Dispose dispose); + +// Holds one SharedScratchDevice per device id. +// +// `get` locks only long enough to find or create the entry, and the reference it +// returns stays usable once that lock is dropped: std::unordered_map keeps +// references to elements valid across rehashing, and entries are never erased. +// This one lock is shared by every device, so holding it couples every device to +// whoever holds it: it covers the lookup, and in the test-only reset the snapshot +// that empties the slots, and nothing else. No CUDA call is made under it. +class SharedScratchPool { + public: + SharedScratchDevice& get(int device_id) { + std::lock_guard lk(mu_); + return devices_[device_id]; + } + + // How long a reset waits, in total, for the device locks. A test calls the reset + // between cases with nothing claimed, so anything held is a leak the run has + // already failed over; the wait only has to outlast a claim still winding down. + static constexpr std::chrono::seconds kResetLockWait{5}; + + private: + // The reset is test-only and its body is not here. It frees everything the live + // pool holds without waiting for work in flight against it, which is not a thing + // this header should hand a consumer of the released source package -- and this + // header does ship with those sources. Only the grant of access is here; the + // body is in SharedScratchPoolReset.h, which is testonly and ships nowhere. What + // is left reachable from a release build is the grant, so reaching the pool that + // way means writing the traversal, the locks and the deadline again rather than + // calling something. + // + // It is an open door and not a lock: the grant is on a template, so it admits + // any definition of that name a consumer writes in this namespace, not only the + // one shipped alongside it, and what such a definition reaches is both members + // below. That includes erasing entries a live `get` reference points into, which + // the stability rule above forbids. C++ offers no narrower grant here -- a + // non-template friend admits a definition of its own just the same -- so the + // alternatives are putting the body back in this header, which is what moving it + // out was for, or making the members public, which is worse. + template + friend std::size_t reset_shared_scratch_pool_slots(SharedScratchPool& pool, Dispose dispose); + + std::mutex mu_; + std::unordered_map devices_; +}; + +// The process-wide per-device pool the TensorRT backend runs on. One buffer per +// device, grown to the largest requirement any call on that device has asked +// for, serves every kUSER_MANAGED context on it that needs activation scratch -- +// a context whose engine needs none under any shape claims nothing here -- instead +// of each of N layer-engines pinning its own scratch, which makes device memory +// scale with the layer count and OOMs multi-layer models. +// +// ORDERING: a context reads and writes its scratch for the whole enqueue, which +// can still be in flight when execute() returns, so two enqueues must never hold +// one buffer at the same time. A device's lock is what enforces that -- see the +// backend's SharedScratchClaim -- and it is held from the claim through the +// enqueue and the record of it, so two execute() calls on one device are +// serialized at submission. The lock does not couple two devices: each carries +// its own, and no CUDA call is made under the one lock the registry itself holds. +// +// The buffers and the events are intentionally never freed at teardown. Nothing +// here runs a CUDA call at process exit, which keeps the pool clear of +// teardown-order hazards against anything else holding device memory. +// +// The C++ object is never destroyed either, and that is not the same claim. +// Static destruction would destroy the registry's mutex, every device's mutex +// and the map nodes a live reference points into, while a thread anywhere +// between the lookup and the release of its device lock still holds them. +// Leaking it costs one allocation and removes that race. +// +// Inline, so the pool's test hooks can reach the same instance from the +// translation unit that defines them without the backend having to export an +// accessor for them; see SharedScratchPoolTestHooks.h. One instance still, for +// the usual reason a function-local static in an inline function is one. +inline SharedScratchPool& scratch_pool() { + static SharedScratchPool* const pool = new SharedScratchPool(); + return *pool; +} + +// Claims a device's handoff for a caller about to enqueue against its shared +// scratch, creating the marker's event on first use. Call with `dev.mu` held. +// +// `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; +} + +// A buffer a growth replaced, handed back for the caller to dispose of. +// +// A non-null `wait_for` is the marker's event, on which an enqueue that may still +// be reading and writing `buffer` has been recorded; the caller must not free the +// buffer ahead of that enqueue. A null `wait_for` means nothing was ever recorded +// against it, so nothing is using it. +// +// One event covers every enqueue the buffer ever served, but only because each of +// them claims the handoff before enqueueing -- which orders its stream after the +// event -- and records on the event afterwards, so the latest recording completes +// only once all the earlier ones have. An enqueue that reaches the buffer without +// doing both is ordered against nothing here. +// +// The caller has two ways to honour that, and both belong outside `dev.mu`: a +// free queued on a stream already ordered after the event, or a host wait on the +// event followed by a device-wide free. Neither is the one the backend always +// takes. A queued free returns the bytes only at the next synchronize of that +// stream, so the backend queues it on a call that synchronizes before returning +// and makes the host wait everywhere else: on a call that does not synchronize, +// on a call that fails after the growth and so has no synchronization left to +// make, and wherever the queued free is refused -- a device with no +// stream-ordered allocator to queue onto, or any other code cudaFreeAsync +// reports. Under the lock either one makes an unrelated claim on this device wait +// for work it has nothing to do with. +struct RetiredScratch { + void* buffer = nullptr; + cudaEvent_t wait_for = nullptr; +}; + +// Bookkeeping for a device's scratch buffer, which grows monotonically to the +// largest requested size. Call with `dev.mu` held. +// +// `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. +// +// A growth reports the buffer it displaced through `out_retired`; see +// RetiredScratch for what the caller owes it. Nothing is freed here, so a caller +// that ignores `out_retired` leaks rather than frees a buffer an enqueue may +// still be using. Every path clears `out_retired` first, so a caller reusing one +// across calls is not handed a buffer an earlier call already disposed of. +// +// What the buffer ended up sized at is `dev.capacity`, which the caller holds the +// lock for anyway. It is not reported separately, because a caller reading it +// would be reading what an earlier call grew the pool to and not its own request: +// the only figure that belongs on a caller's execution context is the one it +// asked for here. +template +void* shared_scratch_get_or_grow(SharedScratchDevice& dev, std::size_t need, Alloc alloc, RetiredScratch& out_retired) { + out_retired = RetiredScratch{}; + if (dev.buffer != nullptr && dev.capacity >= need) { + return dev.buffer; + } + void* p = alloc(need); + if (p == nullptr) { + return nullptr; + } + if (dev.buffer != nullptr) { + out_retired.buffer = dev.buffer; + out_retired.wait_for = dev.marker.pending ? dev.marker.event : nullptr; + } + dev.buffer = p; + dev.capacity = need; + return p; +} + +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolReset.h b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolReset.h new file mode 100644 index 0000000000..8b96d0ce2d --- /dev/null +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolReset.h @@ -0,0 +1,106 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +// The body of the shared scratch pool's test-only reset. +// +// It lives here rather than in SharedScratchPool.h because that header ships in +// the released source package and this operation frees everything the live pool +// holds without waiting for work in flight against it. This file is carried by a +// testonly Bazel target, is in neither the packaged include tree nor the packaged +// source set, and is compiled by nothing a release build produces. What +// SharedScratchPool leaves reachable is the friend declaration alone, so a +// consumer of the released sources who wants this has to write it. +// +// Safe only with no claim outstanding and no enqueue in flight against a pooled +// buffer; the caller earns that by synchronizing every stream it submitted on. + +#include "torch_tensorrt/executorch/SharedScratchPool.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace torch_tensorrt { +namespace executorch_backend { + +// Returns every device's slot to the state it had before anything claimed it, +// handing what the slot held to `dispose(device_id, buffer, event)` so the caller +// can release it. Entries stay in the map, so a reference `get` handed out +// remains valid. +// +// Answers with the number of devices whose lock it could not take within +// SharedScratchPool::kResetLockWait, whose slots it left alone. That case is a +// leaked claim, which is a defect a case in the pool's own test exists to catch -- +// and taking the locks unconditionally would hang on exactly that defect, in a +// fixture that resets both before and after every case, so the run would end in a +// target timeout with nothing said about the cause. The caller reports it instead. +// +// The deadline covers the whole reset rather than each device, so a leaked lock +// costs the same however many devices the pool has seen. What a leak must not do +// is spend the budget on behalf of the slots after it: the retry below sweeps +// every slot it has not taken yet on each pass, so a slot whose claim is merely +// winding down is tried again inside the window it is free, rather than once at +// whatever instant a leaked neighbour leaves in the budget. +// +// Every slot is emptied under the locks, and what came out of it is disposed of +// afterwards with neither held, because the disposer frees device memory and a +// device-wide free blocks on everything queued on that device -- a parked host +// function included. Under the device's lock that would hold off that device's +// next claimant; under the registry's, every device's. +// +// Nothing here waits for an enqueue: the buffer it frees may still be in use by +// one, and the caller is responsible for there being none. +template +std::size_t reset_shared_scratch_pool_slots(SharedScratchPool& pool, Dispose dispose) { + std::vector> taken; + std::size_t still_locked = 0; + { + std::lock_guard lk(pool.mu_); + taken.reserve(pool.devices_.size()); + const auto deadline = std::chrono::steady_clock::now() + SharedScratchPool::kResetLockWait; + std::vector> not_taken_yet; + not_taken_yet.reserve(pool.devices_.size()); + for (auto& entry : pool.devices_) { + not_taken_yet.emplace_back(entry.first, &entry.second); + } + while (!not_taken_yet.empty()) { + std::vector> still_to_try; + for (const auto& slot : not_taken_yet) { + SharedScratchDevice& dev = *slot.second; + if (!dev.mu.try_lock()) { + still_to_try.push_back(slot); + continue; + } + std::lock_guard dev_lk(dev.mu, std::adopt_lock); + taken.emplace_back(slot.first, dev.buffer, dev.marker.event); + dev.buffer = nullptr; + dev.capacity = 0; + dev.marker.event = nullptr; + dev.marker.pending = false; + } + not_taken_yet.swap(still_to_try); + if (not_taken_yet.empty() || std::chrono::steady_clock::now() >= deadline) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + still_locked = not_taken_yet.size(); + } + for (const auto& slot : taken) { + dispose(std::get<0>(slot), std::get<1>(slot), std::get<2>(slot)); + } + return still_locked; +} + +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp new file mode 100644 index 0000000000..bbeaea9837 --- /dev/null +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.cpp @@ -0,0 +1,56 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include "torch_tensorrt/executorch/SharedScratchPoolTestHooks.h" + +#include "torch_tensorrt/executorch/SharedScratchPool.h" +#include "torch_tensorrt/executorch/SharedScratchPoolReset.h" + +#include + +#include + +namespace torch_tensorrt { +namespace executorch_backend { + +std::size_t shared_scratch_capacity_for_testing(int device_id) { + SharedScratchDevice& dev = scratch_pool().get(device_id); + std::lock_guard lk(dev.mu); + return dev.capacity; +} + +bool reset_shared_scratch_pool_for_testing() { + int restore_to = 0; + const bool have_current = cudaGetDevice(&restore_to) == cudaSuccess; + const std::size_t still_locked = + reset_shared_scratch_pool_slots(scratch_pool(), [](int device_id, void* buffer, cudaEvent_t event) { + if (buffer == nullptr && event == nullptr) { + return; + } + // cudaFree and cudaEventDestroy both act on the current device, and a slot is + // keyed by the device its buffer came from. + if (cudaSetDevice(device_id) != cudaSuccess) { + return; + } + if (buffer != nullptr) { + (void)cudaFree(buffer); + } + if (event != nullptr) { + (void)cudaEventDestroy(event); + } + // Clears a non-sticky error so a reset does not leave one for the next call to + // report. A sticky one survives the clear, and no cleanup here recovers it. + (void)cudaGetLastError(); + }); + if (have_current) { + (void)cudaSetDevice(restore_to); + } + return still_locked == 0; +} + +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.h b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.h new file mode 100644 index 0000000000..01ad317227 --- /dev/null +++ b/cpp/src/torch_tensorrt/executorch/SharedScratchPoolTestHooks.h @@ -0,0 +1,44 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +// Views of the shared activation-scratch pool that only a test has any business +// calling: one reads the pool's capacity, the other frees everything it holds +// with no wait for work in flight against it. +// +// They live in this header and SharedScratchPoolTestHooks.cpp rather than beside +// the pool, and neither file is compiled into libexecutorch_trt_backend or +// shipped in the source package, so a released build contains no definition of +// either and exports no symbol for them. The Bazel target that carries them is +// testonly. Reaching the pool from a separate translation unit is what +// scratch_pool() is inline in SharedScratchPool.h for. +// +// Both are safe only with no claim outstanding and no enqueue in flight against a +// pooled buffer; a test earns that by synchronizing every stream it submitted on. +// The reset also frees what the pool holds, so a test should destroy its delegate +// handles before calling it. + +#include + +namespace torch_tensorrt { +namespace executorch_backend { + +// The bytes the pool holds for `device_id` right now; zero if it holds nothing. +std::size_t shared_scratch_capacity_for_testing(int device_id); + +// Frees every device's buffer, destroys its handoff event, and clears the marker, +// so one test does not inherit a pool an earlier one grew. +// +// False when a device's lock was still held after the pool's reset deadline: that +// slot is left as it was. A caller should report it rather than carry on, and +// must not wait for the lock itself -- the state it describes is a claim that was +// never released, so nothing is going to release it. +bool reset_shared_scratch_pool_for_testing(); + +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index 8408c13e88..fa1c4f6946 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -6,10 +6,13 @@ */ #include "torch_tensorrt/executorch/TensorRTBackend.h" +#include "torch_tensorrt/executorch/PooledScratchInstall.h" +#include "torch_tensorrt/executorch/SharedScratchPool.h" #include "torch_tensorrt/executorch/TensorRTBindingNames.h" #include "torch_tensorrt/executorch/TensorRTBlobHeader.h" #include "torch_tensorrt/executorch/WeightStreamingBudget.h" +#include #include #include #include @@ -17,6 +20,7 @@ #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; @@ -85,7 +91,9 @@ EngineHandle::~EngineHandle() { cudaError_t err = cudaEventSynchronize(inflight_event); if (err != cudaSuccess) { ET_LOG(Error, "EngineHandle::~EngineHandle: cudaEventSynchronize failed: %s", cudaGetErrorString(err)); - cudaGetLastError(); // clear sticky error; tear down regardless + // Clears a non-sticky error so it does not resurface under the name of the + // next call here; a sticky one survives the clear. Tear down regardless. + cudaGetLastError(); } inflight_pending = false; } @@ -151,6 +159,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 +176,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_needs_scratch = handle.engine->getDeviceMemorySizeV2() > 0; + } + return Error::Ok; } @@ -204,8 +232,484 @@ bool is_cuda_accessible_ptr(const void* ptr) { return attrs.type == cudaMemoryTypeDevice || attrs.type == cudaMemoryTypeManaged; } +// A caller's hold on one device's shared scratch: the device lock, plus the +// buffer a growth displaced, freed once that lock is dropped. +// +// The lock spans the enqueue, not just the choice of buffer. A claimant that +// released it as soon as it had a buffer would leave its enqueue live for a +// window the marker's event does not yet cover, and a second claimant entering +// that window is handed the same buffer and told to wait for the enqueue before +// it -- so nothing orders the two and both write the same scratch. The failure +// is silent: wrong output, no CUDA error, no TensorRT error. +// +// This lock nests inside the per-handle EngineHandle::mu, which already spans +// the enqueue, and is never taken in the other order. +class SharedScratchClaim { + public: + SharedScratchClaim() = default; + SharedScratchClaim(const SharedScratchClaim&) = delete; + SharedScratchClaim& operator=(const SharedScratchClaim&) = delete; + ~SharedScratchClaim() { + // A claim that still holds a retired buffer here was never released, so + // execute() returned somewhere between the claim and the release below -- a + // refused enqueue, or a failed record of one already submitted -- and no + // synchronization of this stream is still to come. The record failure does + // drain the stream, but it drains and then returns, so a free queued here + // would be queued after the last synchronization the call makes; the refused + // enqueue makes none at all. Such a free would come back at nothing: measured + // on an A100 with CUDA 13.0 and driver 13.0, a cudaFreeAsync'd 4 GiB block is + // still resident after the stream is destroyed, after cudaDeviceSynchronize + // and after another stream's synchronize -- and a caller stream is the + // caller's to destroy, so those bytes are gone for the life of the process. A + // bail-out therefore takes the disposal that has the bytes back before it + // returns, whatever the call would have done had it run to the end. Passing + // that in here rather than storing it at retire() is what keeps the two in + // step: a return added between the claim and the release cannot leave a + // promise to synchronize behind it. + // + // The return is the caller's to report, and on this path there is no caller + // left to report it to: an execute() that returned early has already failed. + (void)release(/*the_call_synchronizes_the_stream=*/false); + } + + SharedScratchDevice& hold(int device_id) { + dev_ = &scratch_pool().get(device_id); + device_id_ = device_id; + lock_ = std::unique_lock(dev_->mu); + return *dev_; + } + + // Null until hold() runs and null again after release(): non-null exactly while + // this claim holds the device's lock. + SharedScratchDevice* device() const { + return dev_; + } + + // Takes ownership of a buffer a growth displaced, to be freed by release(). + // `wait_for` is the marker event the enqueues that used it were recorded on, or + // null if none were; `stream` is the one this claim is about to enqueue on. + void retire(void* buffer, cudaEvent_t wait_for, cudaStream_t stream) { + retired_ = buffer; + retired_wait_ = wait_for; + stream_ = stream; + } + + // Drops the lock, and the device pointer with it so device() cannot hand out a + // pointer this claim no longer holds the lock for. Then disposes of whatever a + // growth displaced, after the unlock, because either way of doing that blocks or + // queues on device work another pooled engine on this device has nothing to do + // with. + // + // Which disposal, because a stream-ordered free does not return the bytes when + // the stream reaches it. Measured on an A100 with CUDA 13.0 and driver 13.0: + // cudaFreeAsync accepts a pointer cudaMalloc returned and defers the free, but + // the bytes come back at the next cudaStreamSynchronize of that stream and at + // no point before it. A stream cudaStreamQuery reports as drained still holds + // them -- a cudaMalloc of the same size fails with out of memory until the + // synchronize, and cudaMemGetInfo does not move. + // + // So the free is queued only where this call synchronizes the stream itself, + // which `the_call_synchronizes_the_stream` says. It is passed in here rather + // than recorded at retire() because it is a claim about what the caller has + // still to do, and only the caller, at the point it releases, knows that it is + // going to: the destructor answers it false for exactly that reason. + // + // - the_call_synchronizes_the_stream. cudaFreeAsync on the stream this + // claim enqueued on, and execute()'s own cudaStreamSynchronize returns the + // bytes before the call ends. Nothing waits for work this call did not + // submit -- measured with a host function parked for 1500 ms on a stream + // neither the pool nor this engine had ever used, around 10 ms against the + // 1505 ms the synchronous disposal took, which came back when that host + // function was released. + // + // - Otherwise -- the caller-stream, device-resident path that returns with its + // enqueue still running -- nothing this backend does will ever synchronize + // that stream, so a queued free would hold the bytes until the caller + // happened to. Every buffer a run retired would stay resident and peak + // memory would be the sum of every size the pool grew to rather than the + // largest of them: measured, replaying this pool's growth sequence on one + // stream with no synchronize, four growths that all succeed under the + // synchronous disposal fail on the fourth with out of memory. That path + // waits on the marker event and makes a device-wide cudaFree, which has the + // bytes back before execute() returns and costs a host wait for everything + // queued on the device. The execute() contract and the README say so. + // + // Ordering, not a wait, is what makes the queued free safe. Every claimant makes + // its stream wait on the marker event before it enqueues, and this claim did so + // while holding the lock, so this stream is already behind every enqueue that + // ever used the retired buffer -- the marker covers all of them, because each + // records on it afterwards. Queuing the free on this stream therefore puts it + // after the last of them, and after this claim's own enqueue, which uses the new + // buffer. + // + // cudaFreeAsync also needs the device's stream-ordered allocator, which not + // every platform has; where it is missing the call fails rather than freeing and + // the synchronous disposal runs instead, so that one is both the other path's + // disposal and this path's fallback. + // + // Returns false when the synchronous disposal's wait failed, which leaves the + // buffer leaked rather than freed under a live enqueue; the caller reports it. + // Frees on the current device, which must still be the buffer's. + bool release(bool the_call_synchronizes_the_stream) { + if (lock_.owns_lock()) { + lock_.unlock(); + } + dev_ = nullptr; + if (retired_ == nullptr) { + return true; + } + void* const retired = retired_; + const cudaEvent_t wait_for = retired_wait_; + const cudaStream_t stream = stream_; + // Clearing the pointer is what stops the destructor's release from disposing + // of the same buffer twice. + retired_ = nullptr; + + if (the_call_synchronizes_the_stream) { + const cudaError_t async_err = cudaFreeAsync(retired, stream); + if (async_err == cudaSuccess) { + return true; + } + // The error is this call's own, so it is cleared here rather than left for + // the next CUDA call in execute() to report under its own name. + cudaGetLastError(); + if (async_err == cudaErrorNotSupported) { + ET_LOG( + Info, + "TensorRTBackend::execute: device %d has no stream-ordered allocator, so the free of the shared activation scratch buffer a pool growth replaced falls back to a host wait and a device-wide free, which blocks this call until the device is idle", + device_id_); + } else { + // Any other code is a fault on a device that does have the allocator, so + // this does not say the platform lacks one. + ET_LOG( + Info, + "TensorRTBackend::execute: the stream-ordered free of the shared activation scratch buffer a pool growth replaced on device %d returned %s, so it falls back to a host wait and a device-wide free, which blocks this call until the device is idle", + device_id_, + cudaGetErrorString(async_err)); + } + } + return dispose_with_a_host_wait(retired, wait_for); + } + + private: + // The disposal that has the bytes back before it returns: wait on the host for + // the enqueue that last used the buffer, then free it device-wide. Both waits + // are unbounded -- see the execute() contract -- which is why the queued free is + // preferred wherever it can return the bytes. + bool dispose_with_a_host_wait(void* retired, cudaEvent_t wait_for) { + if (wait_for != nullptr) { + const cudaError_t wait_err = cudaEventSynchronize(wait_for); + if (wait_err != cudaSuccess) { + // This wait is the only thing keeping the free off a buffer an enqueue may + // still be reading, so a failed wait leaks it instead. What it reports is + // usually an asynchronous fault raised by earlier work on this device. + ET_LOG( + Error, + "TensorRTBackend::execute: waiting for the enqueue on the replaced shared activation scratch on device %d failed (%s), which for a wait on device work is usually an earlier asynchronous fault on this device surfacing here; leaking that buffer rather than freeing it under a live enqueue", + device_id_, + cudaGetErrorString(wait_err)); + cudaGetLastError(); + return false; + } + } + + const cudaError_t err = cudaFree(retired); + if (err != cudaSuccess) { + // cudaFree synchronizes, so what it reports is more often an earlier + // asynchronous fault on this device than a fault in the free -- which is + // why the message does not call it one. + ET_LOG( + Error, + "TensorRTBackend::execute: freeing the shared activation scratch buffer that a pool growth replaced on device %d reported %s; a device-wide free reports whatever fault this device is already in, so this need not be the pool's", + device_id_, + cudaGetErrorString(err)); + // Clears a non-sticky error so it does not resurface under the name of the + // next CUDA call in execute(). A sticky one survives the clear and will + // resurface anyway; the caller learns of it from that call. + cudaGetLastError(); + } + return true; + } + + SharedScratchDevice* dev_ = nullptr; + int device_id_ = -1; + std::unique_lock lock_; + void* retired_ = nullptr; + cudaEvent_t retired_wait_ = nullptr; + cudaStream_t stream_ = nullptr; +}; + +// What a call needing no activation scratch is given when the pool holds nothing +// yet. It cannot be given nothing: enqueueV3 refuses a kUSER_MANAGED context with +// no device memory installed as soon as the engine reports needing any under some +// shape, whatever the shapes actually bound need. So the pool starts at the +// smallest allocation that satisfies that check and the first call with a real +// requirement grows it -- as against standing the engine's profile-wide figure in +// for the zero, which would pin the pool at the largest shape the engine admits +// on the strength of a call that uses none of it. +constexpr size_t kMinPooledScratchBytes = 1; + +// Refuses a pooled call whose stream is capturing a CUDA graph. +// +// The handoff's event wait is on an event recorded outside the capture: +// cudaStreamWaitEvent fails it with cudaErrorStreamCaptureIsolation and +// invalidates the capture under every capture mode. A growth adds more. Its +// cudaMalloc invalidates the capture outside cudaStreamCaptureModeRelaxed, and +// under Relaxed is permitted but runs uncaptured, leaving a replay pointed at a +// buffer the pool may since have freed. Its disposal of the buffer it replaces +// behaves the same way: a call that synchronizes the stream is one no capture mode +// permits, so the only disposal a capture could reach is the host wait and the +// device-wide cudaFree, and both are prohibited outside Relaxed. Even the queued +// free would be no way round it -- a cudaFreeAsync made while capturing has to be +// given a graph allocation and the pool's buffers come from cudaMalloc, so it is +// refused under all three modes, though without invalidating the capture. The +// invalidations do not fail cleanly: +// the caller learns of one only when cudaStreamEndCapture hands back an error and +// a null graph. Refusing names the cause instead. +// +// execute() calls this ahead of every CUDA call it makes that a capture cannot +// take, not just the pool's own: the cudaEventSynchronize on a previous enqueue +// and the cudaMalloc that grows a host-input staging buffer each invalidate a +// capture under every mode but Relaxed, and the cudaStreamSynchronize that ends +// a call this backend does not let return early does so under every mode. Any of +// them leaves a refusal made after it nothing to save. Only the device query and +// the device switch run earlier, and a capture takes both. +// +// It sees only `stream`. A capture live on some *other* stream is invalidated by +// the pool's calls just the same under the Global mode, and under ThreadLocal +// from the same thread, and this does not catch it: cudaStreamIsCapturing +// answers per stream, and asking it about cudaStreamLegacy instead does not +// close the gap -- measured on CUDA 13.0, that reports the capture only when the +// capturing stream is a blocking one, and then reports it under Relaxed and +// under another thread's ThreadLocal too, where the pool's calls are permitted +// and refusing would be wrong. There is no query for "is any capture live in +// this process", so the gap stays; the header and the README say so. +// +// A failed query is refused with the rest. Where `stream` is the legacy stream +// and some other stream is capturing, what it reports is +// cudaErrorStreamCaptureImplicit -- a capture somewhere, not an unrelated fault. +// Refusing on that is over-conservative where the capture is one the pool's calls +// would have survived, and never the other way round. +Error refuse_pooled_call_on_a_capturing_stream(cudaStream_t stream, int device_id) { + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + const cudaError_t capture_err = cudaStreamIsCapturing(stream, &capture); + if (capture_err == cudaSuccess && capture == cudaStreamCaptureStatusNone) { + return Error::Ok; + } + ET_LOG( + Error, + "TensorRTBackend::execute: the selected stream is capturing a CUDA graph (%s), which the shared activation scratch pool on device %d does not support. Neither does the rest of this delegate: loading the engine with '%s' off removes this refusal, but a call captured that way records the handle's completion event into the graph, the next call on that handle then fails, and nothing afterwards waits for a replay of the graph. The capture section of the backend README has the detail.", + capture_err == cudaSuccess ? "capture in progress" : cudaGetErrorString(capture_err), + device_id, + kSharedActivationScratchKey); + if (capture_err != cudaSuccess) { + // The query's own failure is this call's, not the next one's. Only here: on + // the ordinary refusal the query succeeded, so anything pending on this thread + // was left by earlier work and belongs to whoever calls CUDA next. + cudaGetLastError(); + } + return Error::NotSupported; +} + +// Sets out_ptr to a buffer of at least `need` bytes on `device_id`, with `stream` +// ordered after the enqueue that last used the buffer. Returns with `claim` +// holding the device's lock: the caller must submit its enqueue, call +// record_shared_scratch_enqueue, and only then release the claim. +// +// The buffer's capacity is not reported, because no caller has any use for it: +// what a call installs on its context is its own requirement, not whatever the +// pool grew to for someone else. +// +// `need` is a real request and never zero -- execute() substitutes +// kMinPooledScratchBytes for a zero before calling, so that the size the pool +// guarantees and the size the context is told it owns are one figure and not two. +// +// The caller must already have refused a capturing `stream`: the handoff's wait +// invalidates a capture under every mode, a growth's allocation under every mode +// but Relaxed, and a growth's disposal of the buffer it replaces under every mode +// but Relaxed as well. +// +// Must be called with `device_id` already current: cudaEventCreateWithFlags and +// cudaMalloc both act on the *current* device and nothing in here sets it. +Error claim_shared_scratch(SharedScratchClaim& claim, int device_id, size_t need, cudaStream_t stream, void*& out_ptr) { + SharedScratchDevice& dev = claim.hold(device_id); + + const SharedScratchHandoff handoff = shared_scratch_claim_event(dev, []() -> cudaEvent_t { + cudaEvent_t event = nullptr; + // Blocking-sync so the host yields instead of busy-spinning. The only host + // wait ever made on this event is SharedScratchClaim::release()'s synchronous + // disposal, and it waits for a whole inference; spinning would burn a core for + // that time and be no faster, since it is followed by a device-wide cudaFree. + // Where the free is queued on the stream instead nothing waits on this event + // from the host at all, and the flag costs nothing. + if (cudaEventCreateWithFlags(&event, cudaEventDisableTiming | cudaEventBlockingSync) != cudaSuccess) { + // The pool's own failure, cleared where it is made: the caller is told by + // the return, and leaving it pending would surface it under the name of + // whatever this thread calls next. + cudaGetLastError(); + return nullptr; + } + return event; + }); + 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)); + cudaGetLastError(); + return Error::InvalidState; + } + } + + const bool first_buffer = dev.buffer == nullptr; + RetiredScratch retired; + void* const buffer = shared_scratch_get_or_grow( + dev, + need, + [device_id, first_buffer](size_t bytes) -> void* { + void* p = nullptr; + if (cudaMalloc(&p, bytes) != cudaSuccess) { + cudaGetLastError(); + 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; + }, + retired); + 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; + } + + // The retired buffer is disposed of at release(), with the device's lock + // dropped, either by a free queued on `stream` or by a host wait and a + // device-wide free; see SharedScratchClaim::release() for which and why. + // Nothing here makes a CUDA call that blocks on device work under that lock. + claim.retire(retired.buffer, retired.wait_for, stream); + + out_ptr = buffer; + return Error::Ok; +} + +// Records the enqueue now in flight on `stream` against the claimed device's +// shared scratch, so the next call to claim_shared_scratch waits for it. +// +// Call on a `claim` that claim_shared_scratch returned Error::Ok on and that +// still holds the device's lock. Both halves matter, and together they are why +// neither the device nor the event below is checked: a claim's device is non-null +// exactly while it holds the lock, and claim_shared_scratch has already failed the +// call with Error::Internal if the marker had no event, which only the test-only +// reset clears again and that needs the lock this claim is holding. +Error record_shared_scratch_enqueue(SharedScratchClaim& claim, cudaStream_t stream) { + const cudaEvent_t event = shared_scratch_mark_in_flight(*claim.device()); + const cudaError_t err = cudaEventRecord(event, stream); + if (err != cudaSuccess) { + ET_LOG( + Error, + "TensorRTBackend::execute: recording the completion event for the shared activation scratch enqueue failed: %s", + cudaGetErrorString(err)); + cudaGetLastError(); + return Error::InvalidState; + } + return Error::Ok; +} + +// Collects what TensorRT reports while it is attached. +// +// install_pooled_scratch attaches one across a single API call and detaches it +// again, so anything in it came from that call. The detach is not tidiness: while +// a recorder is attached TensorRT reports to it *instead of* to the ILogger, so +// one left in place would divert the diagnostics for setInputShape, +// setTensorAddress and enqueueV3 away from the backend's logger. +class ScopedErrorRecorder final : public nvinfer1::IErrorRecorder { + public: + int32_t getNbErrors() const noexcept override { + return static_cast(errors_.size()); + } + nvinfer1::ErrorCode getErrorCode(int32_t index) const noexcept override { + return in_range(index) ? errors_[static_cast(index)].first : nvinfer1::ErrorCode::kSUCCESS; + } + ErrorDesc getErrorDesc(int32_t index) const noexcept override { + return in_range(index) ? errors_[static_cast(index)].second.c_str() : ""; + } + bool hasOverflowed() const noexcept override { + return overflowed_; + } + void clear() noexcept override { + errors_.clear(); + } + bool reportError(nvinfer1::ErrorCode code, ErrorDesc desc) noexcept override { + // noexcept, and a throwing push_back here would terminate: an allocation + // failure while reporting an error costs the description, not the process. + try { + errors_.emplace_back(code, std::string(desc == nullptr ? "" : desc)); + } catch (...) { + overflowed_ = true; + } + return false; // false asks TensorRT to keep going; the caller decides + } + RefCount incRefCount() noexcept override { + return ++refs_; + } + RefCount decRefCount() noexcept override { + return --refs_; + } + + bool anything_reported() const { + return !errors_.empty() || overflowed_; + } + + private: + bool in_range(int32_t index) const { + return index >= 0 && static_cast(index) < errors_.size(); + } + + std::vector> errors_; + bool overflowed_ = false; + RefCount refs_ = 1; +}; + } // namespace +// Declared in PooledScratchInstall.h. See that header for why the install is +// checked rather than made and trusted. +bool install_pooled_scratch(nvinfer1::IExecutionContext& ctx, void* buffer, size_t bytes, int device_id) { + ScopedErrorRecorder recorder; + nvinfer1::IErrorRecorder* const previous = ctx.getErrorRecorder(); + ctx.setErrorRecorder(&recorder); + ctx.setDeviceMemoryV2(buffer, static_cast(bytes)); + ctx.setErrorRecorder(previous); + + if (!recorder.anything_reported()) { + return true; + } + ET_LOG( + Error, + "TensorRTBackend::execute: TensorRT refused the %zu-byte shared activation scratch buffer for device %d: %s", + bytes, + device_id, + recorder.getNbErrors() > 0 ? recorder.getErrorDesc(0) : "the refusal could not be recorded"); + return false; +} + // --------------------------------------------------------------------------- // is_available // --------------------------------------------------------------------------- @@ -611,6 +1115,29 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* nvinfer1::IExecutionContext* ctx = engine->exec_ctx.get(); TORCHTRT_ET_CHECK_NOT_NULL(ctx, Error::InvalidState, "TensorRTBackend::execute: backend is not initialized"); + const auto caller_stream = ::executorch::extension::cuda::getCallerStream(); + const bool caller_stream_set = caller_stream.has_value(); + cudaStream_t stream = caller_stream.value_or(cudaStreamPerThread); + + // An engine that needs no scratch under any shape is left out of the pool + // entirely: enqueueV3 accepts it with no device memory installed, so it need not + // claim the device and does not serialize against the engines that do. + const bool pooled_scratch = engine->shared_scratch && engine->engine_needs_scratch; + + // Refused here, ahead of every call this function makes that a capture cannot + // take. The pool's own are not the only ones: the wait on a previous enqueue just + // below, and the cudaMalloc that grows a host-input staging buffer further down, + // are prohibited under every mode but cudaStreamCaptureModeRelaxed, so outside + // that mode they would invalidate the capture before the pooled path was ever + // reached and a refusal any later would arrive after the thing it exists to + // protect was gone. + if (pooled_scratch) { + const Error capture_err = refuse_pooled_call_on_a_capturing_stream(stream, engine->device_id); + if (capture_err != Error::Ok) { + return capture_err; + } + } + // A prior fast-path execute() may have returned with its enqueue still in flight // on the shared exec_ctx. Wait for it before reconfiguring the context below: // TensorRT forbids mutating a context while one of its enqueues is in flight, and @@ -623,12 +1150,37 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* return Error::InvalidProgram; } } - const auto caller_stream = ::executorch::extension::cuda::getCallerStream(); - const bool caller_stream_set = caller_stream.has_value(); - cudaStream_t stream = caller_stream.value_or(cudaStreamPerThread); bool output_staged_to_host = false; bool input_staged_from_host = false; + // A host-resident input is staged with an asynchronous copy that reads the + // caller's own memory, and that copy must not still be running when execute() + // returns: the caller is free to write that buffer again, and a copy from + // pinned memory has not read it yet -- measured, every byte the device + // received was the value written after the return. The success path already + // synchronizes whenever anything was staged (must_sync below); this covers the + // returns between the copy and that point, which would otherwise leave it live. + // + // The wait it makes is a whole-stream one and not a wait on the copy alone: + // cudaStreamSynchronize returns when everything queued on the caller's stream + // has run, including work the caller queued before calling. So an error return + // that staged a host input can block on work this delegate never submitted. + // Waiting for the copy alone would mean recording an event after the last + // staging copy of every call that makes one and waiting on that instead, which + // costs a per-call event record on the success path to narrow a wait only error + // returns make. Staging already implies must_sync, so no successful call reaches + // the destructor with this pending. + struct StagedInputDrain { + cudaStream_t stream; + const bool& staged; + bool done = false; + ~StagedInputDrain() { + if (staged && !done) { + (void)cudaStreamSynchronize(stream); + } + } + } staged_input_drain{stream, input_staged_from_host}; + if (engine->cached_input_ptrs.empty()) { engine->cached_input_ptrs.resize(num_inputs, nullptr); engine->cached_input_sizes.resize(num_inputs, 0); @@ -905,8 +1457,91 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } } + // Whether this call ends by waiting for its own enqueue. Read twice below -- at + // the release of the scratch claim, whose disposal of a retired buffer turns on + // it, and at the wait itself -- so it is settled once here, which is the first + // point everything it reads is final: the two staging flags and the aliased + // reflects are set while the bindings are built above, and whether a caller + // stream is active was read at the top. + // + // must_sync = an output is staged to host (the caller reads the D2H result on + // return), an input was staged from host (its async H2D read the caller's host + // buffer, which the caller may reuse once we return), an aliased reflect is + // queued (ExecuTorch's buffer-mutation copy_ reads that EValue after execute() + // returns, so the reflect must complete first), or no caller stream is active + // (preserve the historical "results ready on return" behavior). + // + // Otherwise -- caller stream, all I/O device-resident, no alias -- the engine + // work is left enqueued so it composes with the caller's later GPU work, and a + // completion event is recorded instead so the next execute() and the destructor + // wait before reusing or freeing exec_ctx. + const bool aliased_reflect_pending = !aliased_reflects.empty(); + const bool must_sync = + output_staged_to_host || input_staged_from_host || aliased_reflect_pending || !caller_stream_set; + // ------------------------------------------------------------------ - // 4. 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: an engine backed by less + // than it asked for writes past the end. + // + // The buffer is installed on every call, not once, because any call needing more + // than the pool holds -- this engine on other shapes, or another one -- grows it + // and moves it. A kSTATIC context owns its private scratch, so the install must + // not be made on one. + // + // A reported zero has two causes that reach here and nothing distinguishes them: + // a call whose bound shapes need none -- an empty input inside a profile that + // admits one -- and a query that failed. Neither is asked to size the pool: a + // zero asks for kMinPooledScratchBytes, which any live buffer already covers, so + // only against an empty pool does it allocate anything, and then the minimum + // TensorRT will accept rather than the engine's profile-wide figure. + // + // The two causes are told apart by the install rather than by the query, which + // is the only place the difference is observable. Where the shapes genuinely + // need nothing the engine expects nothing and the minimum is accepted; where the + // query failed the engine still expects what it always did, is offered that same + // minimum, and refuses it -- and this call ends here rather than enqueueing + // against whatever pointer the context was last given, which a growth may since + // have freed. + // + // What is installed is this call's own requirement, not the capacity the pool + // happens to hold. They differ whenever an earlier call asked for more, and + // handing this context the larger figure would tell TensorRT it owns bytes + // holding another engine's activations -- the pool never clears the buffer -- + // and would leave an engine overrunning its own requirement inside the region it + // was told it owns, so nothing catches it. It would also blunt the refusal + // above: after any growth, the capacity may cover what a failed query concealed, + // and the install that is meant to catch it succeeds. + // + // The claim holds the device's pool lock from here through the record of the + // enqueue below; see SharedScratchClaim for why it spans that far. Every return + // in between drops it through the destructor, which runs ahead of the device + // restore above, so its free lands on the right device. + SharedScratchClaim scratch_claim; + if (pooled_scratch) { + const size_t need = ctx->updateDeviceMemorySizeForShapes(); + // The substitution for a zero is made here and once: the same figure has to + // be the size the pool guarantees and the size the context is told it owns, + // and nothing downstream checks that two copies of it still agree -- + // setDeviceMemoryV2 refuses an install smaller than the bound shapes need and + // says nothing about one that is larger. A context whose engine needs scratch + // under some shape has to be given a buffer whatever this call's shapes need, + // or enqueueV3 refuses it. + const size_t scratch_bytes = need == 0 ? kMinPooledScratchBytes : need; + void* pool = nullptr; + const Error scratch_err = claim_shared_scratch(scratch_claim, engine->device_id, scratch_bytes, stream, pool); + if (scratch_err != Error::Ok) { + return scratch_err; + } + if (!install_pooled_scratch(*ctx, pool, scratch_bytes, engine->device_id)) { + return Error::InvalidState; + } + } + + // ------------------------------------------------------------------ + // 5. Enqueue inference on the current CUDA stream // ------------------------------------------------------------------ if (!ctx->enqueueV3(stream)) { ET_LOG( @@ -918,6 +1553,54 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* return Error::InvalidState; } + // What every early return past this point owes. The enqueue is submitted and + // this handle's completion marker is not armed, so a return that left it running + // would have a later execute() or ~EngineHandle reconfigure or free exec_ctx + // underneath it -- which TensorRT forbids. Where the pool is in play the same + // wait is also what keeps the next claimant from overwriting scratch this + // enqueue is still using. The drain guard's flag is cleared with it, since this + // wait covers a staged input's copy as well. + const auto drain_the_enqueue = [&]() { + (void)cudaStreamSynchronize(stream); + engine->inflight_pending = false; + staged_input_drain.done = true; + }; + + // Pairs with claim_shared_scratch: the next claimant waits on this event. + if (pooled_scratch) { + const Error mark_err = record_shared_scratch_enqueue(scratch_claim, stream); + if (mark_err != Error::Ok) { + // Nothing will wait for this enqueue otherwise: the record that would have + // put it on the marker is the call that just failed. + // + // The wait is made with the device's pool lock still held -- the claim is + // released below -- which is the one place a pooled call holds it across a + // host wait, so another pooled engine on this device waits out this + // inference. Releasing first would be worse: it would hand the buffer to a + // claimant with nothing ordering it against the enqueue this call just + // submitted, which is the silent-corruption case the lock exists for. + drain_the_enqueue(); + return mark_err; + } + } + // The enqueue is now on the marker's event, so the device's pool is safe to + // hand to the next claimant. Released here rather than at the end of the + // function so the rest of execute() -- the aliased reflects, the D2H copies and + // their synchronizations -- does not hold up another engine on this device. + // The release is also where a growth disposes of the buffer it replaced, with + // the lock dropped. This is the only place that release is made rather than + // left to the claim's destructor, and it is the only place must_sync is a + // promise this function can still keep: every return past it synchronizes the + // stream when must_sync is set -- the branch below does, and the aliased + // reflects drain the enqueue before returning. + if (!scratch_claim.release(must_sync)) { + // Only the synchronous disposal reports a failure, and only from its host + // wait, which for a wait on device work means this device is already in a + // faulted state. + drain_the_enqueue(); + return Error::InvalidProgram; + } + // Caller-owned KV: reflect each engine in-place update into its delegate output // EValue (D2D on the same stream, after the engine work). for (const auto& r : aliased_reflects) { @@ -925,34 +1608,19 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* if (cuda_err != cudaSuccess) { ET_LOG( Error, "TensorRTBackend::execute: aliased-output reflect D2D copy failed: %s", cudaGetErrorString(cuda_err)); - // enqueueV3 already submitted engine work to `stream`, and inflight_pending - // is not armed until the end of the happy path -- drain now so a later - // execute() or the destructor never reconfigures/frees exec_ctx while this - // enqueue is still running. - (void)cudaStreamSynchronize(stream); - engine->inflight_pending = false; + drain_the_enqueue(); return Error::InvalidProgram; } } - // The engine work is now in flight on `stream`. Decide whether to wait for it: - // must_sync = an output is staged to host (the caller reads the D2H result on - // return), an input was staged from host (its async H2D read the caller's host - // buffer, which the caller may reuse once we return), or no caller stream is - // active (preserve the historical "results ready on return" behavior). - // Otherwise (caller stream + all I/O device-resident) leave the work enqueued so - // it composes with the caller's later GPU work, and record inflight_event so the - // next execute() and the destructor wait before reusing/freeing exec_ctx. The D2H - // copies live in the must_sync branch: an output staged to host always sets + // The engine work is now in flight on `stream`, and must_sync -- decided above, + // where the scratch claim needed it -- says whether to wait for it. The D2H + // copies live in this branch: an output staged to host always sets // output_staged_to_host, so outputs_needing_copy is empty on the skip path. - // An aliased reflect enqueues the engine's in-place update into the delegate - // output EValue on `stream`; ExecuTorch's buffer-mutation copy_ reads that EValue - // after execute() returns, so the reflect must complete first. A model with - // aliased outputs therefore always syncs here. - const bool aliased_reflect_pending = !aliased_reflects.empty(); - const bool must_sync = - output_staged_to_host || input_staged_from_host || aliased_reflect_pending || !caller_stream_set; if (must_sync) { + // Every return from here on is behind the cudaStreamSynchronize below, so + // the staging drain has nothing left to do. + staged_input_drain.done = true; Error copy_err = Error::Ok; for (auto& output : outputs_needing_copy) { exec_aten::Tensor et_out = args[output.first]->toTensor(); @@ -983,11 +1651,9 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } else { cuda_err = cudaEventRecord(engine->inflight_event, stream); if (cuda_err != cudaSuccess) { - // Could not arm the completion marker; drain now so a later execute() or the - // destructor never reconfigures or frees exec_ctx while this enqueue runs. + // Could not arm the completion marker, so nothing downstream would wait. ET_LOG(Error, "TensorRTBackend::execute: cudaEventRecord failed: %s", cudaGetErrorString(cuda_err)); - (void)cudaStreamSynchronize(stream); - engine->inflight_pending = false; + drain_the_enqueue(); return Error::InvalidProgram; } engine->inflight_pending = true; @@ -995,6 +1661,37 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* return Error::Ok; } +// --------------------------------------------------------------------------- +// set_option +// --------------------------------------------------------------------------- +Error TensorRTBackend::set_option(ET_UNUSED BackendOptionContext& context, const Span& backend_options) { + // The whole span is read before anything is stored. A span is one request, so a + // caller told it was refused must not find part of it applied -- and the part + // that would be applied here is process-wide and governs every engine loaded + // after it. Where a span names this key more than once the last one wins, which + // is what applying each in turn did. + bool requested = false; + bool have_request = false; + for (const auto& option : backend_options) { + // A caller may address one option span to several backends, so a key this + // backend does not read is skipped rather than refused. + if (std::strcmp(option.key, kSharedActivationScratchKey) == 0) { + const bool* const val = std::get_if(&option.value); + if (val == nullptr) { + ET_LOG(Error, "TensorRTBackend::set_option: option '%s' must be a boolean", kSharedActivationScratchKey); + return Error::InvalidArgument; + } + requested = *val; + have_request = true; + } + } + + if (have_request) { + scratch_enabled.store(requested, std::memory_order_relaxed); + } + return Error::Ok; +} + // --------------------------------------------------------------------------- // destroy // diff --git a/py/torch-tensorrt-executorch-runtime/native/CMakeLists.txt b/py/torch-tensorrt-executorch-runtime/native/CMakeLists.txt index 4dadbf56f5..287ec853d5 100644 --- a/py/torch-tensorrt-executorch-runtime/native/CMakeLists.txt +++ b/py/torch-tensorrt-executorch-runtime/native/CMakeLists.txt @@ -173,6 +173,9 @@ target_compile_definitions(torch_tensorrt_executorch_backend target_include_directories(torch_tensorrt_executorch_backend PRIVATE "${TORCH_TENSORRT_SOURCE_DIR}" "${TORCH_TENSORRT_SOURCE_DIR}/cpp/include" + # Headers the backend sources include but that are not part of the installed + # API, such as SharedScratchPool.h. + "${TORCH_TENSORRT_SOURCE_DIR}/cpp/src" "${EXECUTORCH_SOURCE_DIR}/.." "${EXECUTORCH_SOURCE_DIR}/runtime/core/portable_type/c10") if(NOT TARGET extension_cuda) diff --git a/tests/cpp/BUILD b/tests/cpp/BUILD index b5c0c15138..b8d18994af 100644 --- a/tests/cpp/BUILD +++ b/tests/cpp/BUILD @@ -62,13 +62,13 @@ test_suite( ], ) +# The suite in the executorch package is the one the workflow runs; this is the +# label for it from here. Listing the tests again would be a second list to keep +# in step, for a name nothing refers to. test_suite( name = "executorch_backend_tests", tests = [ - "//tests/cpp/executorch:test_caller_stream", - "//tests/cpp/executorch:test_executorch_binding_names", - "//tests/cpp/executorch:test_executorch_blob_header", - "//tests/cpp/executorch:test_executorch_weight_streaming_budget", + "//tests/cpp/executorch:executorch_backend_tests", ], ) diff --git a/tests/cpp/executorch/BUILD b/tests/cpp/executorch/BUILD index 17d2820bf2..1b8f123b59 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,57 @@ cc_test( "@googletest//:gtest_main", ], ) + +# The long timeout is what makes the fork case's per-child deadline mean anything: +# it reaps 24 children at 30 s each, so the path where they wedge budgets 720 s +# against Bazel's 300 s default for a target that names neither size nor timeout. +# Two source files because one of them is how the target sees that scratch_pool() +# is one registry across translation units; see the second file's comment. +cc_test( + name = "test_shared_scratch_pool", + timeout = "long", + srcs = [ + "test_shared_scratch_pool.cpp", + "test_shared_scratch_pool_other_tu.cpp", + ], + deps = [ + "//cpp:tensorrt_executorch_shared_scratch_pool", + "//cpp:tensorrt_executorch_shared_scratch_pool_reset", + "@googletest//:gtest_main", + ], +) + +# 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", + "//cpp:tensorrt_executorch_pooled_scratch_install", + "//cpp:tensorrt_executorch_shared_scratch_pool", + "//cpp:tensorrt_executorch_shared_scratch_pool_test_hooks", + "@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..da3ea2ef7f --- /dev/null +++ b/tests/cpp/executorch/test_shared_scratch_backend.cpp @@ -0,0 +1,2714 @@ +/* + * 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 pooled execute() path -- the kUSER_MANAGED context, the +// updateDeviceMemorySizeForShapes/setDeviceMemoryV2 pair, the enqueue handoff +// between two caller streams, the growth a larger engine forces on a pool a +// smaller one already allocated, two of the three things that make the per-shape +// query answer zero -- the third, a failed query, cannot be induced from inside +// this process -- and two threads submitting against one pooled buffer at once. +// +// The TensorRT engine is built here rather than loaded from a .pte so the target +// carries no exported artifact, at the cost of a few seconds of builder time. +// +// 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 -- and a skipped +// gtest case exits zero, which Bazel reports as a passing target. A run that is +// meant to have a device says so through TORCHTRT_EXECUTORCH_REQUIRE_CUDA, and +// then the missing device is a failure instead, which each failing case says for +// itself. Where that variable is unset, a count of the cases that skipped for that +// reason is printed at the end of the suite. +// +// That variable covers the missing device and nothing else. Five cases skip for a +// second reason: two need the device full, two need its memory quiet, and one +// needs the device's stream-ordered allocator. None of those is a state a test can +// insist on while sharing the device with other processes. Those skips stand +// whether the variable is set or not, so a green required-CUDA run says every case +// ran, not that every case covered what it is named for. So the suite counts these +// separately from the missing device and lists them by name and reason at the end +// of the file, whether or not the variable is set. That list is in the test log, +// which --test_output=errors prints nothing of for a target that passes; the CI +// invocation therefore also passes --test_summary=detailed, which names every +// skipped case in Bazel's own summary. +// +// SHARED-DEVICE WARNING: the two cases that need the device full take it to +// essentially zero free bytes for as long as they hold their DeviceMemoryHog -- +// measured, 16 MiB free of 81151 MiB, with a neighbour process getting +// out-of-memory on 256 MiB allocations throughout the window. That is the only way +// to reach the pool's allocation-failure path from inside this process: nothing in +// the delegate takes an allocator a test could substitute, and the other early +// returns on that path need TensorRT or CUDA to fail a call that is correct as +// made. Bazel's exclusive tag keeps other actions in the same build off this +// device; it cannot keep anything else off it. +// +// The CI invocation in .github/workflows/executorch-build-linux.yml passes that +// variable, and the job it sits in asks for a GPU runner and starts its container +// with every GPU attached. The ExecuTorch job's own gate in ci-linux-x86_64.yml +// does not keep these cases off an ordinary pull-request push: it sits out only +// when the lane is skip or the backend is RTX, and _decide.yml resolves an +// unlabelled pull_request to lane=fast and backend=standard. What does drop the +// job is the standard channel having been cancelled, which the same workflow +// already carries a comment about. + +#include "torch_tensorrt/executorch/PooledScratchInstall.h" +#include "torch_tensorrt/executorch/SharedScratchPool.h" +#include "torch_tensorrt/executorch/SharedScratchPoolTestHooks.h" +#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 +#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 TensorRTBackend.h: a test that reads the +// key through the production constant cannot pin the key's value. +constexpr char kOptionKey[] = "use_shared_activation_scratch"; + +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); + +// A second scratch-needing engine, four times the elements of the one above, so +// loading it after that one drives the pool's growth path. Every other engine in +// this file asks for the same size, which is why nothing else reaches it. +constexpr int kBigRows = 4096; +constexpr int kBigCols = 4096; + +// A dynamic-batch engine whose profile admits an empty batch. Binding one is a +// valid call that needs no activation scratch, from an engine that needs some, +// which is one of the three things a zero from updateDeviceMemorySizeForShapes +// can mean. The dimensions are small because no test compares this engine's +// memory against a figure; only whether each answer it gives is zero matters. +constexpr int kDynRows = 64; +constexpr int kDynCols = 64; +constexpr int kDynMaxBatch = 128; +// Inside the same profile and not empty, so one engine covers both answers. +constexpr int kDynBatch = 4; + +// Engines loaded together in the memory test. Four is enough for the private +// case to cost 4x the scratch and the pooled case 1x. +constexpr int kEngineCount = 4; + +// 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; +} + +// Wraps a serialized engine in the delegate's blob wire format. +std::vector wrap_engine_plan(const nvinfer1::IHostMemory& plan) { + const std::string metadata = + R"({"io_bindings":[{"name":"input_0","is_input":true},{"name":"output_0","is_input":false}],)" + R"("hardware_compatible":false,"device_id":0})"; + const auto metadata_offset = static_cast(kHeaderSize); + const auto metadata_size = static_cast(metadata.size()); + const auto engine_offset = static_cast(align_up(metadata_offset + metadata_size, kEngineAlignment)); + + std::vector blob(static_cast(engine_offset) + plan.size(), 0); + std::memcpy(blob.data(), kMagic, sizeof(kMagic)); + write_field(blob, kMetadataOffsetField, metadata_offset); + write_field(blob, kMetadataSizeField, metadata_size); + write_field(blob, kEngineOffsetField, engine_offset); + write_field(blob, kEngineSizeField, static_cast(plan.size())); + std::memcpy(blob.data() + metadata_offset, metadata.data(), metadata.size()); + std::memcpy(blob.data() + engine_offset, plan.data(), plan.size()); + return blob; +} + +std::vector build_engine_blob(bool needs_scratch, int rows = kRows, int cols = kCols) { + static BuilderLogger logger; + + 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, rows, cols}); + 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, rows, cols}; + profile->setDimensions("input_0", nvinfer1::OptProfileSelector::kMIN, shape); + profile->setDimensions("input_0", nvinfer1::OptProfileSelector::kOPT, shape); + profile->setDimensions("input_0", nvinfer1::OptProfileSelector::kMAX, shape); + config->addOptimizationProfile(profile); + + TRTUniquePtr plan(builder->buildSerializedNetwork(*network, *config)); + if (plan == nullptr) { + return {}; + } + return wrap_engine_plan(*plan); +} + +// The scratch-needing network again, over a leading dimension the caller chooses +// per call, with an empty batch as the profile minimum. +std::vector build_dynamic_batch_engine_blob() { + static BuilderLogger logger; + + TRTUniquePtr builder(nvinfer1::createInferBuilder(logger)); + if (builder == nullptr) { + return {}; + } + TRTUniquePtr network(builder->createNetworkV2(0)); + if (network == nullptr) { + return {}; + } + + nvinfer1::ITensor* input = + network->addInput("input_0", nvinfer1::DataType::kFLOAT, nvinfer1::Dims3{-1, kDynRows, kDynCols}); + if (input == nullptr || !add_scratch_needing_net(*network, *input)) { + return {}; + } + + TRTUniquePtr config(builder->createBuilderConfig()); + if (config == nullptr) { + return {}; + } + nvinfer1::IOptimizationProfile* profile = builder->createOptimizationProfile(); + profile->setDimensions("input_0", nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims3{0, kDynRows, kDynCols}); + profile->setDimensions("input_0", nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims3{kDynBatch, kDynRows, kDynCols}); + profile->setDimensions( + "input_0", nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims3{kDynMaxBatch, kDynRows, kDynCols}); + config->addOptimizationProfile(profile); + + TRTUniquePtr plan(builder->buildSerializedNetwork(*network, *config)); + if (plan == nullptr) { + return {}; + } + return wrap_engine_plan(*plan); +} + +// The activation scratch one context of `blob`'s engine needs for the given +// shape, read the way execute() reads it. Zero if the engine could not be +// measured, which is also what an empty batch answers, so a caller reading a zero +// as meaningful has to rule the failure out by some other measurement. +std::size_t measure_engine_scratch( + const std::vector& blob, + int rows = kRows, + int cols = kCols, + int batch = 1) { + static BuilderLogger logger; + TensorRTBlobHeader header; + if (!TensorRTBlobHeader::parse(blob.data(), blob.size(), header)) { + 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{batch, rows, cols})) { + 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. `batch`/`rows`/`cols` must be a shape the + // blob's profile admits; `batch` may be 0, which allocates nothing and leaves + // both device pointers null, a state an empty ExecuTorch tensor can arrive in. + Error load( + const std::vector& blob, + std::uint32_t seed, + int rows = kRows, + int cols = kCols, + int batch = 1) { + batch_ = static_cast(batch); + rows_ = static_cast(rows); + cols_ = static_cast(cols); + std::vector host_in(elems()); + for (std::size_t i = 0; i < elems(); ++i) { + host_in[i] = pattern(i, seed); + } + if (bytes() > 0) { + if (cudaMalloc(&device_in_, bytes()) != cudaSuccess || cudaMalloc(&device_out_, bytes()) != cudaSuccess) { + return Error::MemoryAllocationFailed; + } + if (cudaMemcpy(device_in_, host_in.data(), bytes(), cudaMemcpyHostToDevice) != cudaSuccess) { + return Error::Internal; + } + } + + arena_storage_.resize(kArenaBytes); + 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) { + if (bytes() == 0) { + return true; + } + const std::vector host(elems(), value); + return cudaMemcpy(device_out_, host.data(), bytes(), cudaMemcpyHostToDevice) == cudaSuccess; + } + + // Runs one inference on `stream`. Returns without waiting for the enqueue, + // which is the state the pool's handoff exists to order. + Error run(cudaStream_t stream) { + return run_with_input(stream, device_in_); + } + + // The same run with no CallerStreamGuard, which is one of the four things that + // make execute() synchronize the stream before returning. It is the one a memory + // measurement can use: the other three each need a buffer this fixture does not + // otherwise allocate -- a host-backed input or output, which execute() stages + // through a device buffer of its own, or an extra output EValue for an aliased + // output -- and the growth measurement would read those alongside the pool. The + // enqueue goes on cudaStreamPerThread, which is not ordered against any stream + // created cudaStreamNonBlocking. + Error run_on_the_synchronized_path() { + return run_with_input(nullptr, device_in_, /*scope_a_caller_stream=*/false); + } + + // The same run with the input bound to caller-owned host memory, which + // execute() stages through a device buffer of its own instead of binding + // directly. The output stays device-resident, so the only staging is the + // input's. `host_in` must hold bytes() bytes. + Error run_from_host_input(cudaStream_t stream, void* host_in) { + return run_with_input(stream, host_in); + } + + // Where execute() staged this handle's input, or null if it never had to. + void* staging_buffer_for_input_0() const { + const EngineHandle* const h = handle(); + return h->cached_input_ptrs.empty() ? nullptr : h->cached_input_ptrs[0]; + } + + std::vector read_output() const { + std::vector host_out(elems()); + if (bytes() > 0 && cudaMemcpy(host_out.data(), device_out_, bytes(), cudaMemcpyDeviceToHost) != cudaSuccess) { + host_out.clear(); + } + return host_out; + } + + const EngineHandle* handle() const { + return static_cast(handle_); + } + + // The execution context execute() configures and enqueues on, so a test can + // watch what the delegate does to it. + nvinfer1::IExecutionContext* context() const { + return static_cast(handle_)->exec_ctx.get(); + } + + std::size_t elems() const { + return static_cast(batch_) * static_cast(rows_) * static_cast(cols_); + } + + std::size_t bytes() const { + return elems() * sizeof(float); + } + + private: + // With `scope_a_caller_stream` false no CallerStreamGuard is scoped and `stream` + // is unread: execute() then finds no caller stream, enqueues on + // cudaStreamPerThread and synchronizes before returning. + Error run_with_input(cudaStream_t stream, void* in_ptr, bool scope_a_caller_stream = true) { + // Separate arrays: execute() resizes the output tensor to the shape TensorRT + // inferred, which writes through whichever array that tensor was given. + SizesType in_sizes[3] = {batch_, rows_, cols_}; + SizesType out_sizes[3] = {batch_, rows_, cols_}; + ::executorch::aten::TensorImpl in_impl(ScalarType::Float, 3, in_sizes, in_ptr); + ::executorch::aten::TensorImpl out_impl(ScalarType::Float, 3, out_sizes, device_out_); + ::executorch::aten::Tensor in_tensor(&in_impl); + ::executorch::aten::Tensor out_tensor(&out_impl); + EValue in_value(in_tensor); + EValue out_value(out_tensor); + EValue* args[2] = {&in_value, &out_value}; + + BackendExecutionContext exec_context; + if (!scope_a_caller_stream) { + return backend_.execute(exec_context, handle_, Span(args, 2)); + } + ::executorch::extension::cuda::CallerStreamGuard guard(stream); + return backend_.execute(exec_context, handle_, Span(args, 2)); + } + + // 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; + SizesType batch_ = 1; + SizesType rows_ = kRows; + SizesType cols_ = kCols; +}; + +// Sets `out` to the device-wide bytes in use, or returns false leaving it alone. +// It reports the failure rather than substituting a figure because both callers +// subtract two of these: a zero for the first reading of a pair makes the second +// look like the whole cost of what was measured between them, which is a pass. +bool device_bytes_in_use(std::size_t& out) { + std::size_t free_bytes = 0; + std::size_t total_bytes = 0; + if (cudaMemGetInfo(&free_bytes, &total_bytes) != cudaSuccess) { + return false; + } + out = total_bytes - free_bytes; + return true; +} + +// Takes device memory until an allocation of the size a caller names can no longer +// succeed, and gives it all back when it goes out of scope. The only way to reach +// the pool's allocation-failure path from inside the process, and the reason it is +// held for as short a window as possible: while it is up, every other process on +// this device is out of memory too. +class DeviceMemoryHog { + public: + DeviceMemoryHog() = default; + DeviceMemoryHog(const DeviceMemoryHog&) = delete; + DeviceMemoryHog& operator=(const DeviceMemoryHog&) = delete; + + ~DeviceMemoryHog() { + release(); + } + + // Chunked from large to small, so a device with tens of gigabytes free is filled + // in a few dozen allocations rather than thousands. + bool leave_less_free_than(std::size_t bytes) { + for (std::size_t chunk = std::size_t{1} << 30; chunk >= (std::size_t{1} << 20); chunk /= 4) { + while (free_bytes_above(bytes)) { + void* block = nullptr; + if (cudaMalloc(&block, chunk) != cudaSuccess) { + cudaGetLastError(); + break; + } + blocks_.push_back(block); + } + } + return !free_bytes_above(bytes); + } + + void release() { + for (void* block : blocks_) { + cudaFree(block); + } + blocks_.clear(); + } + + private: + static bool free_bytes_above(std::size_t bytes) { + std::size_t free_bytes = 0; + std::size_t total_bytes = 0; + return cudaMemGetInfo(&free_bytes, &total_bytes) == cudaSuccess && free_bytes > bytes; + } + + std::vector blocks_; +}; + +// Opens a live kUSER_MANAGED context over `blob`'s engine with one shape bound, +// the state execute() installs a scratch buffer into. Held by the caller, unlike +// measure_engine_scratch's, which is gone by the time it returns its figure. +struct LiveContext { + TRTUniquePtr runtime; + TRTUniquePtr engine; + TRTUniquePtr ctx; +}; + +bool open_user_managed_context(const std::vector& blob, int rows, int cols, int batch, LiveContext& out) { + static BuilderLogger logger; + TensorRTBlobHeader header; + if (!TensorRTBlobHeader::parse(blob.data(), blob.size(), header)) { + return false; + } + out.runtime.reset(nvinfer1::createInferRuntime(logger)); + if (out.runtime == nullptr) { + return false; + } + out.engine.reset( + out.runtime->deserializeCudaEngine(TensorRTBlobHeader::engine_data(blob.data(), header), header.engine_size)); + if (out.engine == nullptr) { + return false; + } + out.ctx.reset(out.engine->createExecutionContext(nvinfer1::ExecutionContextAllocationStrategy::kUSER_MANAGED)); + if (out.ctx == nullptr) { + return false; + } + return out.ctx->setInputShape("input_0", nvinfer1::Dims3{batch, rows, cols}); +} + +// Counts the reference-count calls TensorRT makes on a recorder as it attaches +// and detaches one, which is how a test attached to the delegate's own context +// sees install_pooled_scratch scope its recorder over the install. +// +// TensorRT documents setErrorRecorder as calling incRefCount on the recorder it +// takes and decRefCount on the one it replaces, and it does: measured, attaching +// a second recorder over this one and then restoring it moves detached to 1 and +// reattached to 1, and nothing else in a pooled execute() touches either. +class CountingErrorRecorder final : public nvinfer1::IErrorRecorder { + public: + int32_t getNbErrors() const noexcept override { + return 0; + } + nvinfer1::ErrorCode getErrorCode(int32_t) const noexcept override { + return nvinfer1::ErrorCode::kSUCCESS; + } + ErrorDesc getErrorDesc(int32_t) const noexcept override { + return ""; + } + bool hasOverflowed() const noexcept override { + return false; + } + void clear() noexcept override {} + bool reportError(nvinfer1::ErrorCode, ErrorDesc) noexcept override { + reported.fetch_add(1); + return false; + } + RefCount incRefCount() noexcept override { + reattached.fetch_add(1); + return ++refs_; + } + RefCount decRefCount() noexcept override { + detached.fetch_add(1); + return --refs_; + } + + void forget() { + reattached.store(0); + detached.store(0); + reported.store(0); + } + + std::atomic reattached{0}; + std::atomic detached{0}; + std::atomic reported{0}; + + private: + RefCount refs_ = 1; +}; + +// Set by a run that is supposed to have a GPU. A skip is then a failure, rather +// than a green target that covered nothing. +constexpr char kRequireCudaEnvVar[] = "TORCHTRT_EXECUTORCH_REQUIRE_CUDA"; + +bool cuda_device_is_required() { + const char* const value = std::getenv(kRequireCudaEnvVar); + return value != nullptr && value[0] != '\0' && std::strcmp(value, "0") != 0; +} + +Error set_shared_scratch(TensorRTBackend& backend, bool enabled) { + BackendOption option; + std::strncpy(option.key, kOptionKey, sizeof(option.key) - 1); + 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); + big_blob_ = build_engine_blob(true, kBigRows, kBigCols); + dynamic_blob_ = build_dynamic_batch_engine_blob(); + if (blob_.empty() || scratch_free_blob_.empty() || big_blob_.empty() || dynamic_blob_.empty()) { + return; + } + scratch_bytes_ = measure_engine_scratch(blob_); + big_scratch_bytes_ = measure_engine_scratch(big_blob_, kBigRows, kBigCols); + empty_batch_scratch_bytes_ = measure_engine_scratch(dynamic_blob_, kDynRows, kDynCols, 0); + dynamic_batch_scratch_bytes_ = measure_engine_scratch(dynamic_blob_, kDynRows, kDynCols, kDynBatch); + scratch_free_engine_bytes_ = engine_scratch_requirement(scratch_free_blob_); + dynamic_engine_bytes_ = engine_scratch_requirement(dynamic_blob_); + } + + static void TearDownTestSuite() { + report_the_second_reason_skips(); + if (skipped_for_no_device_ == 0) { + return; + } + // Nothing skipped when the requirement is on -- SetUp counts the missing device + // and then fails the case. Printing here would put a skip banner under those + // failures and tell the reader to set the very variable that produced them. + if (cuda_device_is_required()) { + return; + } + const ::testing::TestSuite* const suite = ::testing::UnitTest::GetInstance()->current_test_suite(); + std::fprintf( + stderr, + "[ SKIPPED ] %d of %d cases in this file: no CUDA device. Nothing here ran, so this target passing says " + "nothing about the shared activation scratch pool. Set %s=1 on a run that is meant to have a device and a " + "skip becomes a failure.\n", + skipped_for_no_device_, + suite == nullptr ? skipped_for_no_device_ : suite->total_test_count(), + kRequireCudaEnvVar); + } + + void SetUp() override { + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count == 0) { + ++skipped_for_no_device_; + this_case_skipped_for_no_device_ = true; + if (cuda_device_is_required()) { + FAIL() << "no CUDA device, and " << kRequireCudaEnvVar + << " says this run must have one. Every case in this file needs a device, so without one the binary " + "would exit zero having covered nothing."; + } + GTEST_SKIP() << "no CUDA device: the shared-scratch backend path is not covered by this run"; + } + ASSERT_FALSE(blob_.empty()) << "TensorRT could not build the fixture engine"; + ASSERT_FALSE(scratch_free_blob_.empty()) << "TensorRT could not build the scratch-free fixture engine"; + ASSERT_FALSE(big_blob_.empty()) << "TensorRT could not build the larger fixture engine"; + ASSERT_FALSE(dynamic_blob_.empty()) << "TensorRT could not build the dynamic-batch fixture engine"; + ASSERT_EQ(set_shared_scratch(backend_, false), Error::Ok); + // The pool outlives every test in this file, and most of them depend on what + // it holds when they start: one expects a growth, one expects none, one + // expects a first allocation. Leaving it alone makes each of those depend on + // which cases ran before it, so the suite passes in declaration order and in + // no other -- running the empty-input case first turns the handoff test's + // first run into a growth, which is not the state that case is written for. + // + // Reported rather than waited out: a lock still held here was leaked by the + // case before, and this case cannot run without the reset, so it fails now + // and says which state it found. + ASSERT_TRUE(reset_shared_scratch_pool_for_testing()) + << "a device's pool lock was still held when this case started, so an earlier one returned without releasing " + "its claim and the pool could not be reset"; + } + + void TearDown() override { + record_a_second_reason_skip(); + set_shared_scratch(backend_, false); + // Also here, so a buffer this test grew is not still resident while the next + // one measures device-wide memory. + ASSERT_TRUE(reset_shared_scratch_pool_for_testing()) + << "this case left a device's pool lock held, so the pool could not be reset for the next one"; + } + + const std::vector& blob() const { + return blob_; + } + + const std::vector& scratch_free_blob() const { + return scratch_free_blob_; + } + + const std::vector& big_blob() const { + return big_blob_; + } + + const std::vector& dynamic_blob() const { + return dynamic_blob_; + } + + // Defined below, next to the two cases that call it. + void run_a_growth_and_bound_what_it_cost(bool synchronized_path); + + // A case can also skip for a reason the device being present says nothing about + // -- the device would not fill, its memory would not stay still, it has no + // stream-ordered allocator. Those skips are the difference between a green + // target and a run that covered what the cases are named for, and with + // --test_output=errors a passing target prints not one of them. So they are + // collected here and printed at the end whether or not the CUDA requirement is + // armed, which is the case the missing-device banner cannot cover. + void record_a_second_reason_skip() { + if (this_case_skipped_for_no_device_ || !::testing::Test::IsSkipped()) { + return; + } + const ::testing::TestInfo* const info = ::testing::UnitTest::GetInstance()->current_test_info(); + if (info == nullptr) { + return; + } + std::string reason; + const ::testing::TestResult* const result = info->result(); + for (int i = 0; result != nullptr && i < result->total_part_count(); ++i) { + const ::testing::TestPartResult& part = result->GetTestPartResult(i); + if (part.type() == ::testing::TestPartResult::kSkip) { + reason = part.summary(); + break; + } + } + // The reason a case gives usually wraps, and one line per skip is what makes + // the report below readable. + for (char& c : reason) { + if (c == '\n') { + c = ' '; + } + } + while (!reason.empty() && reason.back() == ' ') { + reason.pop_back(); + } + if (reason.empty()) { + reason = "no reason recorded"; + } + second_reason_skips_.emplace_back(info->name(), reason); + } + + static void report_the_second_reason_skips() { + if (second_reason_skips_.empty()) { + return; + } + const ::testing::TestSuite* const suite = ::testing::UnitTest::GetInstance()->current_test_suite(); + std::fprintf( + stderr, + "[ SKIPPED ] %zu of %d cases in this file ran on a device but did not cover what they are named for, so this " + "target passing does not say the pool's growth, allocation-failure and queued-free paths were exercised:\n", + second_reason_skips_.size(), + suite == nullptr ? static_cast(second_reason_skips_.size()) : suite->total_test_count()); + for (const auto& skip : second_reason_skips_) { + std::fprintf(stderr, "[ SKIPPED ] %s: %s\n", skip.first.c_str(), skip.second.c_str()); + } + } + + TensorRTBackend backend_; + static std::vector blob_; + static std::vector scratch_free_blob_; + static std::vector big_blob_; + static std::vector dynamic_blob_; + static std::size_t scratch_bytes_; + static std::size_t big_scratch_bytes_; + static std::size_t empty_batch_scratch_bytes_; + static std::size_t dynamic_batch_scratch_bytes_; + static std::int64_t scratch_free_engine_bytes_; + static std::int64_t dynamic_engine_bytes_; + static int skipped_for_no_device_; + static std::vector> second_reason_skips_; + bool this_case_skipped_for_no_device_ = false; +}; + +std::vector SharedScratchBackendTest::blob_; +std::vector SharedScratchBackendTest::scratch_free_blob_; +std::vector SharedScratchBackendTest::big_blob_; +std::vector SharedScratchBackendTest::dynamic_blob_; +std::size_t SharedScratchBackendTest::scratch_bytes_ = 0; +std::size_t SharedScratchBackendTest::big_scratch_bytes_ = 0; +std::size_t SharedScratchBackendTest::empty_batch_scratch_bytes_ = 0; +std::size_t SharedScratchBackendTest::dynamic_batch_scratch_bytes_ = 0; +std::int64_t SharedScratchBackendTest::scratch_free_engine_bytes_ = -1; +std::int64_t SharedScratchBackendTest::dynamic_engine_bytes_ = -1; +int SharedScratchBackendTest::skipped_for_no_device_ = 0; +std::vector> SharedScratchBackendTest::second_reason_skips_; + +// --------------------------------------------------------------------------- +// set_option +// --------------------------------------------------------------------------- + +// 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"; +} + +// The case above sends the bad entry on its own, where returning before the store +// and storing before the return look the same. A span carrying a good entry ahead +// of a bad one tells them apart, and it is the span a caller writes when it sets +// several options at once. +TEST_F(SharedScratchBackendTest, SetOptionRejectsASpanWithoutApplyingTheEntriesBeforeTheBadOne) { + ASSERT_EQ(set_shared_scratch(backend_, false), Error::Ok); + + BackendOption turn_on; + std::strncpy(turn_on.key, kOptionKey, sizeof(turn_on.key) - 1); + turn_on.value = true; + // The same key again, so what the span asks for is unambiguous: this entry + // cannot be read as addressed to some other backend. + BackendOption wrong_type; + std::strncpy(wrong_type.key, kOptionKey, sizeof(wrong_type.key) - 1); + wrong_type.value = 1; + BackendOption options[2] = {turn_on, wrong_type}; + BackendOptionContext context; + + EXPECT_EQ(backend_.set_option(context, Span(options, 2)), Error::InvalidArgument); + + LoadedEngine engine; + ASSERT_EQ(engine.load(blob(), 28), Error::Ok); + EXPECT_FALSE(engine.handle()->shared_scratch) + << "a span that was refused still turned the pool on for every engine loaded after it"; +} + +// A context's allocation strategy is fixed when the context is created, so the +// option cannot be re-read per call. +TEST_F(SharedScratchBackendTest, EachEngineCapturesTheSettingInEffectAtItsOwnLoad) { + 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; + { + std::size_t before = 0; + ASSERT_TRUE(device_bytes_in_use(before)) << "cudaMemGetInfo failed, so this test measured nothing"; + std::vector> engines; + for (int i = 0; i < kEngineCount; ++i) { + engines.push_back(std::make_unique()); + ASSERT_EQ(engines.back()->load(blob(), 9), Error::Ok); + ASSERT_EQ(engines.back()->run(stream), Error::Ok); + } + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + std::size_t after = 0; + ASSERT_TRUE(device_bytes_in_use(after)) << "cudaMemGetInfo failed, so this test measured nothing"; + // 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; + { + std::size_t before = 0; + ASSERT_TRUE(device_bytes_in_use(before)) << "cudaMemGetInfo failed, so this test measured nothing"; + std::vector> engines; + for (int i = 0; i < kEngineCount; ++i) { + engines.push_back(std::make_unique()); + ASSERT_EQ(engines.back()->load(blob(), 9), Error::Ok); + ASSERT_EQ(engines.back()->run(stream), Error::Ok); + } + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + std::size_t after = 0; + ASSERT_TRUE(device_bytes_in_use(after)) << "cudaMemGetInfo failed, so this test measured nothing"; + 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"; +} + +// --------------------------------------------------------------------------- +// Growing the pool +// --------------------------------------------------------------------------- + +// Runs a four-times-larger engine after a smaller one to reach the growth path, +// which nothing else in this file does, and bounds what the growth cost the +// device. +// +// The reading is taken the moment execute() returns and *before* any synchronize +// of this test's, which is what makes the upper bound mean something. A growth has +// two disposals to choose from and only one of them has the bytes back by then; a +// synchronize in between hides the difference, because it is exactly what releases +// a queued free. Measured on an A100 with CUDA 13.0: cudaFreeAsync on a +// cudaMalloc'd pointer returns cudaSuccess and defers the free, and the bytes come +// back at the next cudaStreamSynchronize of that stream and at no point before it +// -- a stream cudaStreamQuery reports as drained still holds them. So a growth +// that queued the free on a stream nothing synchronizes costs the whole new buffer +// here rather than the difference, and the upper bound below fails. +// +// The lower bound also fails if the pool were already large enough for the second +// engine, which is how this test could otherwise pass vacuously. The fixture +// empties the pool before each test, so it is the smaller engine's run below that +// establishes the size the growth has to exceed. +// +// `synchronized_path` picks which of the two disposals runs, since the choice is +// execute()'s must_sync and not the pool's: a call with no caller stream ends by +// synchronizing and queues the free, a call on a caller stream with device-resident +// I/O does not and makes a host wait and a device-wide free instead. Both have to +// have the bytes back before the call returns, and this is the same measurement +// pointed at each. +void SharedScratchBackendTest::run_a_growth_and_bound_what_it_cost(bool synchronized_path) { + ASSERT_GE(big_scratch_bytes_, scratch_bytes_ + kMinMeasurableScratch) + << "the two fixture engines ask for " << scratch_bytes_ << " and " << big_scratch_bytes_ + << " bytes of activation scratch, too close for the growth to be measurable"; + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + + const auto run_one = [synchronized_path, stream](LoadedEngine& engine) { + return synchronized_path ? engine.run_on_the_synchronized_path() : engine.run(stream); + }; + + // Private-scratch references for both engines, and the run that pays the larger + // engine's one-time TensorRT and CUDA module costs so they land outside the + // measurement below. + std::vector small_expected; + std::vector big_expected; + { + LoadedEngine small_priv; + LoadedEngine big_priv; + ASSERT_EQ(small_priv.load(blob(), 15), Error::Ok); + ASSERT_EQ(big_priv.load(big_blob(), 16, kBigRows, kBigCols), Error::Ok); + ASSERT_FALSE(small_priv.handle()->shared_scratch); + ASSERT_FALSE(big_priv.handle()->shared_scratch); + ASSERT_EQ(run_one(small_priv), Error::Ok); + ASSERT_EQ(run_one(big_priv), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + small_expected = small_priv.read_output(); + big_expected = big_priv.read_output(); + } + + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + LoadedEngine small; + ASSERT_EQ(small.load(blob(), 15), Error::Ok); + ASSERT_TRUE(small.handle()->shared_scratch); + ASSERT_EQ(run_one(small), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + // Loaded before the measurement starts: its weights and its I/O are not part of + // what the growth costs. + LoadedEngine big; + ASSERT_EQ(big.load(big_blob(), 16, kBigRows, kBigCols), Error::Ok); + ASSERT_TRUE(big.handle()->shared_scratch); + + std::size_t before = 0; + ASSERT_TRUE(device_bytes_in_use(before)) << "cudaMemGetInfo failed, so this test measured nothing"; + const auto measurement_began = std::chrono::steady_clock::now(); + ASSERT_EQ(run_one(big), Error::Ok); + // No synchronize between the growth and this reading: see the note above. Both + // the allocation and both disposals are host calls execute() makes and returns + // from, so what this reads is settled whether or not the enqueue has finished. + std::size_t after = 0; + ASSERT_TRUE(device_bytes_in_use(after)) << "cudaMemGetInfo failed, so this test measured nothing"; + // A control window of the same length as the measurement, with nothing of this + // test's running in it. It is sampled before the synchronize below for the same + // reason the reading above is: a synchronize releases a free the growth only + // queued, so one taken inside this window would move `settled` away from `after` + // and turn the failure the bounds are meant to report into a skip. On a device + // this test has to itself, device-wide usage does not move across it; anything + // else means another process is moving memory on the same timescale as the + // growth, which is what makes the bounds below report a figure the pool did not + // produce. The same length matters: two readings taken back to back would sample + // microseconds against the growth's milliseconds and would miss almost + // everything. It is still a sample of a different window, so it narrows the + // misdiagnosis rather than removing it, and the upper bound names the cause as + // well. The exclusive tag keeps other Bazel actions off this device, not other + // processes. + // + // What counts as quiet is what the bounds below can absorb, rather than exact + // equality, which threw the whole measurement away for a byte. The tighter of + // the two bounds allows scratch_bytes_/2 of slack, so a control window that + // moved by at most half of that cannot turn what they report into something + // else, while anything large enough to matter still skips. The tolerance is for + // allocator granularity and not for a busy device: on a device this suite had to + // itself the window was measured moving 0 bytes against a tolerance of 8 MiB for + // these fixture engines, and a neighbour allocating at the scale they do moves it + // by far more than that. + std::this_thread::sleep_for(std::chrono::steady_clock::now() - measurement_began); + std::size_t settled = 0; + ASSERT_TRUE(device_bytes_in_use(settled)) << "cudaMemGetInfo failed, so this test measured nothing"; + const std::size_t quiet_tolerance = scratch_bytes_ / 4; + const std::size_t control_window_moved_by = settled > after ? settled - after : after - settled; + const bool device_was_quiet = control_window_moved_by <= quiet_tolerance; + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + // Only something outside this test can make device-wide usage fall across a + // growth: the growth allocates the larger buffer before it disposes of the + // smaller one. It belongs with the skip below rather than being a failure of its + // own, for the reason that skip exists -- and it has to be answered before the + // subtraction, which is unsigned. + const bool device_gave_memory_back = after < before; + const std::size_t difference = big_scratch_bytes_ - scratch_bytes_; + + // The pool must still serve the smaller engine after the growth moved the + // buffer: its context holds the address it was given on its previous call, and + // that address has been freed. This run is also the one place the suite installs + // a size below the pool's capacity, since that is what this engine's shapes + // need, so a size TensorRT refuses fails here. + ASSERT_TRUE(small.fill_output(kSentinel)); + ASSERT_EQ(run_one(small), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::vector small_actual = small.read_output(); + const std::vector big_actual = big.read_output(); + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); + + ASSERT_EQ(big_expected.size(), big_actual.size()); + ASSERT_FALSE(big_expected.empty()); + EXPECT_EQ(std::memcmp(big_expected.data(), big_actual.data(), big_expected.size() * sizeof(float)), 0) + << "the engine that grew the pool did not produce what it produces with its own scratch"; + + ASSERT_EQ(small_expected.size(), kElems); + ASSERT_EQ(small_actual.size(), kElems); + EXPECT_NE(small_expected[0], kSentinel) << "the reference output is the sentinel, so a skipped enqueue would pass"; + EXPECT_EQ(std::memcmp(small_expected.data(), small_actual.data(), kBytes), 0) + << "the smaller engine stopped producing its own output once the growth moved the shared buffer"; + + // Last, so the output comparisons above are made either way. + if (device_gave_memory_back) { + GTEST_SKIP() << "device-wide memory in use fell from " << before << " to " << after + << " bytes across the growth, which a growth cannot do, so another process released memory on this " + "device and the two bounds below would measure it rather than the growth"; + } + if (!device_was_quiet) { + GTEST_SKIP() << "device-wide memory in use moved from " << after << " to " << settled << " bytes -- " + << control_window_moved_by << " against a " << quiet_tolerance + << "-byte tolerance -- with nothing of this test's running in between, so another process is " + "allocating on this device and the two bounds below would measure it rather than the growth"; + } + const std::size_t growth_cost = after - before; + EXPECT_GE(growth_cost, difference / 2) << "the larger engine cost " << growth_cost << " bytes against a " + << difference + << "-byte difference in requirement, so the pool did not grow for it"; + EXPECT_LE(growth_cost, difference + scratch_bytes_ / 2) + << "the larger engine cost " << growth_cost << " bytes, about the whole " << big_scratch_bytes_ + << "-byte buffer rather than the " << difference + << "-byte difference, so the buffer it replaced was still resident when execute() returned -- a free the growth " + "only queued, or none at all -- unless another process allocated on this device across the measurement, " + "which the control window above samples for and cannot rule out"; +} + +// The path this option exists for: a caller stream with device-resident I/O, so +// execute() returns with the enqueue still running and never synchronizes the +// stream. A growth here cannot queue its free, because nothing would ever make the +// bytes come back. +TEST_F(SharedScratchBackendTest, ALargerEngineGrowsThePoolAndFreesTheBufferItReplaces) { + run_a_growth_and_bound_what_it_cost(/*synchronized_path=*/false); +} + +// The other path, where execute() ends by synchronizing the stream and the growth +// queues its free on it. The bytes have to be back by the time the call returns +// here too -- the same bound, and it is the call's own synchronize that satisfies +// it rather than anything this test does. +TEST_F(SharedScratchBackendTest, AGrowthOnASynchronizedCallFreesTheBufferItReplaces) { + run_a_growth_and_bound_what_it_cost(/*synchronized_path=*/true); +} + +// --------------------------------------------------------------------------- +// An engine that needs no activation scratch +// --------------------------------------------------------------------------- + +// The pooled path branches on this flag, and it must separate the two fixture +// networks or the branch is only ever taken one way below. +TEST_F(SharedScratchBackendTest, EachEngineRecordsWhetherItNeedsActivationScratch) { + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + LoadedEngine needing; + LoadedEngine scratch_free; + ASSERT_EQ(needing.load(blob(), 12), Error::Ok); + ASSERT_EQ(scratch_free.load(scratch_free_blob(), 13), Error::Ok); + + EXPECT_TRUE(needing.handle()->engine_needs_scratch) + << "the two-softmax network reports no activation scratch of its own, so no test here reaches the pooled path"; + EXPECT_FALSE(scratch_free.handle()->engine_needs_scratch) + << "the pointwise chain reports activation scratch, so it no longer covers the scratch-free case"; +} + +// Turning the pool on must not turn an engine that legitimately needs no +// activation scratch into a failure. Such an engine skips the pool altogether, +// which is what the capacity check below pins. +TEST_F(SharedScratchBackendTest, AnEngineNeedingNoActivationScratchRunsWithThePoolEnabled) { + ASSERT_EQ(scratch_free_engine_bytes_, 0) << "the fixture engine needs scratch, so this test covers nothing"; + + 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); + EXPECT_EQ(shared_scratch_capacity_for_testing(pooled.handle()->device_id), 0u) + << "an engine that needs no activation scratch under any shape claimed the pool anyway"; + const std::vector actual = pooled.read_output(); + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); + + 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); +} + +// --------------------------------------------------------------------------- +// An empty input to an engine that does need activation scratch +// --------------------------------------------------------------------------- + +// An empty batch inside the profile is a valid call, and the per-shape query +// answers it with the same zero it gives for a failed query, from an engine whose +// own requirement is not zero. The three assertions at the top are what make this +// the empty-input case rather than either of the others: an engine that needs +// nothing would fail the second, and a measurement that could not read the engine +// at all would fail the third. +// +// What the empty call must not do is size the pool for itself. The engine's own +// profile-wide requirement, `dynamic_engine_bytes_`, is the figure a zero invites +// installing, and it stands far above the `dynamic_batch_scratch_bytes_` the next +// call actually needs; the pool never shrinks, so installing it would pin the pool +// at many times that for the rest of the process. The empty call is therefore +// checked twice, once against an empty pool and once against a pool holding a real +// requirement, and neither may leave the pool above what a non-empty call asked +// for. The capacity is read directly rather than inferred from device-wide +// memory, which cannot resolve the difference reliably. +TEST_F(SharedScratchBackendTest, AnEmptyInputRunsWithThePoolEnabled) { + ASSERT_EQ(empty_batch_scratch_bytes_, 0u) + << "an empty batch reports " << empty_batch_scratch_bytes_ + << " bytes of activation scratch, so this test no longer covers a call whose per-shape query answers zero"; + ASSERT_GT(dynamic_engine_bytes_, 0) + << "the dynamic fixture engine reports no activation scratch of its own, so the zero above is the " + "scratch-free case rather than the empty-input one"; + ASSERT_GT(dynamic_batch_scratch_bytes_, 0u) << "the same engine reports no activation scratch at batch " << kDynBatch + << " either, so the zero above says nothing about the batch being empty"; + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + + // The non-empty reference is taken with private scratch, so the pooled run + // below has something to be compared against that the pool did not produce. + LoadedEngine priv; + ASSERT_EQ(priv.load(dynamic_blob(), 19, kDynRows, kDynCols, kDynBatch), Error::Ok); + ASSERT_FALSE(priv.handle()->shared_scratch); + ASSERT_TRUE(priv.fill_output(kSentinel)); + ASSERT_EQ(priv.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::vector expected = priv.read_output(); + + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + LoadedEngine empty; + ASSERT_EQ(empty.load(dynamic_blob(), 19, kDynRows, kDynCols, 0), Error::Ok); + ASSERT_TRUE(empty.handle()->shared_scratch); + const int device = empty.handle()->device_id; + ASSERT_EQ(shared_scratch_capacity_for_testing(device), 0u) + << "the fixture left the pool holding something, so an empty call adding nothing to it proves nothing"; + EXPECT_EQ(empty.run(stream), Error::Ok) << "the pool rejected an empty input to an engine that needs scratch"; + EXPECT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + // Not zero: enqueueV3 refuses this context with no device memory installed + // whatever the bound shapes need, so the pool has to allocate something. What it + // must not do is take the engine's word for how much. + const std::size_t after_empty = shared_scratch_capacity_for_testing(device); + // Pinned to the exact figure rather than bounded below the 131072 bytes the next + // call needs, because a bound that loose is met by any minimum up to 128 KiB and + // the point of the minimum is that it is the smallest thing enqueueV3 accepts. + EXPECT_EQ(after_empty, 1u) << "an empty call against an empty pool took " << after_empty + << " bytes rather than the one-byte kMinPooledScratchBytes; a zero means no buffer was " + "installed, which TensorRT refuses for this engine, and anything larger means the " + "call sized the pool from a figure of its own"; + + // The same engine on a shape that does need scratch, after the empty call, so a + // pool the empty call left in an unusable state is not left untested. + LoadedEngine pooled; + ASSERT_EQ(pooled.load(dynamic_blob(), 19, kDynRows, kDynCols, kDynBatch), Error::Ok); + ASSERT_TRUE(pooled.handle()->shared_scratch); + ASSERT_TRUE(pooled.fill_output(kSentinel)); + EXPECT_EQ(pooled.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::vector actual = pooled.read_output(); + const std::size_t after_non_empty = shared_scratch_capacity_for_testing(device); + EXPECT_EQ(after_non_empty, dynamic_batch_scratch_bytes_) + << "the pool holds " << after_non_empty << " bytes after a call needing " << dynamic_batch_scratch_bytes_ + << ", so it was not sized to what that call asked for"; + + // A second empty call, now that the pool holds a real requirement: it must be + // handed that buffer rather than grow the pool to the engine's profile-wide one. + ASSERT_EQ(empty.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + EXPECT_EQ(shared_scratch_capacity_for_testing(device), after_non_empty) + << "the empty call grew the pool from " << after_non_empty << " bytes, against the " << dynamic_engine_bytes_ + << " bytes this engine reports over its whole profile"; + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); + + const std::size_t dyn_elems = + static_cast(kDynBatch) * static_cast(kDynRows) * static_cast(kDynCols); + ASSERT_EQ(expected.size(), dyn_elems); + ASSERT_EQ(actual.size(), dyn_elems); + // Without these two the comparison would be satisfied by an execute() that + // wrote nothing, and by a network whose output does not depend on its input. + EXPECT_NE(expected[0], kSentinel) << "the reference output is the sentinel, so a skipped enqueue would pass"; + bool varies = false; + for (std::size_t i = 1; i < dyn_elems && !varies; ++i) { + varies = expected[i] != expected[0]; + } + EXPECT_TRUE(varies) << "the reference output is constant, so the comparison proves nothing"; + EXPECT_EQ(std::memcmp(expected.data(), actual.data(), dyn_elems * sizeof(float)), 0); +} + +// The install is what separates the two causes of a zero from the per-shape +// query. Both reach execute() as a zero and neither is asked to size the pool, so +// both can be handed a buffer as small as one byte; only one of them is safe with +// it. Where the bound shapes genuinely need nothing the engine expects nothing +// and the one-byte buffer is accepted. Where the query failed the engine expects what it +// always did, TensorRT refuses the install through a call that returns void, and +// the context keeps the buffer it was last given -- which a growth may already +// have freed, and enqueueV3 then reports success while the engine reads and +// writes memory the pool no longer owns. +// +// The failed query itself cannot be induced from inside this process, so what is +// pinned here is the check that catches its consequence, over the same TensorRT +// call execute() makes. +TEST_F(SharedScratchBackendTest, AnUndersizedScratchInstallIsSeenAsRefused) { + LiveContext needs_scratch; + ASSERT_TRUE(open_user_managed_context(blob(), kRows, kCols, 1, needs_scratch)) + << "could not open a user-managed context over the fixture engine"; + const std::size_t need = needs_scratch.ctx->updateDeviceMemorySizeForShapes(); + ASSERT_GT(need, 1u) << "the fixture engine needs no more than one byte of activation scratch for this shape, so an " + "undersized install cannot be built out of it"; + + void* enough = nullptr; + ASSERT_EQ(cudaMalloc(&enough, need), cudaSuccess); + EXPECT_TRUE(install_pooled_scratch(*needs_scratch.ctx, enough, need, 0)) + << "a buffer of exactly what the bound shapes need was reported as refused, which would fail every pooled call"; + + void* one_byte = nullptr; + ASSERT_EQ(cudaMalloc(&one_byte, 1), cudaSuccess); + EXPECT_FALSE(install_pooled_scratch(*needs_scratch.ctx, one_byte, 1, 0)) + << "TensorRT refused a one-byte buffer for a context expecting " << need + << " bytes and the backend read the install as accepted, so the enqueue would run on the buffer installed " + "before it"; + + // The safe cause, which has to keep working: an empty batch expects nothing, so + // the one-byte buffer is accepted and no refusal may be invented for it. + LiveContext empty_batch; + ASSERT_TRUE(open_user_managed_context(dynamic_blob(), kDynRows, kDynCols, 0, empty_batch)); + ASSERT_EQ(empty_batch.ctx->updateDeviceMemorySizeForShapes(), 0u) + << "an empty batch does not answer zero for this engine, so this half covers nothing"; + EXPECT_TRUE(install_pooled_scratch(*empty_batch.ctx, one_byte, 1, 0)) + << "the one-byte buffer an empty call is handed was reported as refused, which would fail a call that is safe"; + + EXPECT_EQ(cudaFree(enough), cudaSuccess); + EXPECT_EQ(cudaFree(one_byte), cudaSuccess); +} + +// The case above drives install_pooled_scratch directly, which leaves execute() +// free to stop calling it: replacing the checked call with a bare +// setDeviceMemoryV2 installs the same buffer, produces the same output and drops +// only the refusal, so every other case here stays green. +// +// Forcing the refusal itself through execute() is not available. The size +// execute() installs is what the per-shape query just returned, over a buffer the +// pool has grown to at least that, so the installed size is never short of what +// the context expects. Only a failed query makes the two disagree, and that +// cannot be induced from inside this process. +// +// Nor is the choice between that size and the pool's capacity observable through +// TensorRT: an install larger than the context expects is accepted in silence, so +// no black-box case here can tell the two apart, and the argument for the smaller +// one is the hazard it removes rather than anything a test can read back. What the +// suite does see is the other direction -- ALargerEngineGrowsThePoolAndFreesTheBufferItReplaces +// runs the smaller engine again after the growth, so its install is smaller than +// the capacity, and a size TensorRT would not accept fails that case. +// +// What is observable is the recorder. install_pooled_scratch attaches one for +// the duration of the install and restores the previous one after, because that +// is the only channel TensorRT reports a refusal on. A recorder this test leaves +// on the delegate's context therefore sees itself replaced and put back exactly +// once per pooled run, and sees nothing at all if the install stops going +// through the helper. +TEST_F(SharedScratchBackendTest, ExecuteInstallsPooledScratchThroughTheCheckedHelper) { + // Declared before the engine so it outlives the context TensorRT attaches it + // to: the context decrements this recorder's count as it is destroyed. + CountingErrorRecorder recorder; + + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + + LoadedEngine pooled; + ASSERT_EQ(pooled.load(blob(), 25), Error::Ok); + ASSERT_TRUE(pooled.handle()->shared_scratch); + ASSERT_TRUE(pooled.handle()->engine_needs_scratch) + << "this engine skips the pool, so its execute() makes no install for this case to watch"; + pooled.context()->setErrorRecorder(&recorder); + + recorder.forget(); + ASSERT_EQ(pooled.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + EXPECT_EQ(recorder.detached.load(), 1) + << "a pooled execute() did not replace the context's error recorder, so its activation-scratch install was not " + "the checked one and a refusal by TensorRT would go unread"; + EXPECT_EQ(recorder.reattached.load(), 1) + << "the recorder the caller had attached was not put back after the install, so TensorRT's diagnostics for the " + "rest of the call go somewhere the caller did not ask for"; + EXPECT_EQ(pooled.context()->getErrorRecorder(), &recorder) + << "execute() left an error recorder of its own attached to the context"; + EXPECT_EQ(recorder.reported.load(), 0) + << "TensorRT reported an error to the caller's recorder during a run that succeeded"; + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +// Every pooled path out of execute() that is not the happy one returns while the +// claim still holds the device's lock, and leaves the claim's destructor to drop +// it. Nothing else here takes one of those paths, and the failure a regression +// there produces is not a wrong answer: it is every later pooled call on the +// device blocking forever with nothing logged. +// +// Only one of them can be reached from inside the process. A refused install needs +// the pool to hold less than the context expects, which it never does; a failed +// enqueue and a failed completion record need TensorRT or CUDA to fail a call that +// is correct as made. The pool's own allocation can be made to fail by taking the +// device's memory first, and it returns through the same destructor as the others. +// +// What it reaches is the lock the destructor drops, not the buffer it disposes +// of: an allocation that failed retired nothing, so this case leaves the +// destructor with none to free. Nothing in this suite reaches a bail-out that +// still holds a retired buffer, since the two returns that can are the two named +// above. What keeps that path right is the shape of release() rather than a case: +// it takes "does this call still synchronize the stream?" as an argument, so a +// bail-out answers it where the bail-out happens and there is no recorded promise +// for a later return to leave stale. +TEST_F(SharedScratchBackendTest, AFailedPooledAllocationLeavesTheDeviceLockFreeAndNothingPending) { + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + + LoadedEngine pooled; + ASSERT_EQ(pooled.load(blob(), 31), Error::Ok); + ASSERT_TRUE(pooled.handle()->shared_scratch); + ASSERT_TRUE(pooled.handle()->engine_needs_scratch) << "this engine skips the pool, so it never allocates from it"; + const int device_id = pooled.handle()->device_id; + ASSERT_EQ(shared_scratch_capacity_for_testing(device_id), 0u) + << "the pool already holds a buffer, so this call would reuse it rather than allocate"; + + Error failed_run = Error::Ok; + cudaError_t pending_after_the_failure = cudaSuccess; + { + DeviceMemoryHog hog; + if (!hog.leave_less_free_than(scratch_bytes_)) { + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); + GTEST_SKIP() << "could not take the device below the " << scratch_bytes_ + << " bytes this engine's scratch needs, so its allocation would have succeeded"; + } + // The hog clears the last failed allocation of its own, and nothing else has + // made a CUDA call since, so whatever is pending after the run below was left + // by the run. + ASSERT_EQ(cudaGetLastError(), cudaSuccess) << "something before the call under test left a CUDA error pending"; + failed_run = pooled.run(stream); + pending_after_the_failure = cudaPeekAtLastError(); + } + + if (failed_run == Error::Ok) { + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); + GTEST_SKIP() << "the pooled allocation succeeded with the device full, so this run did not take the path under " + "test"; + } + ASSERT_EQ(failed_run, Error::MemoryAllocationFailed) + << "the run failed somewhere other than the pool's allocation, so it says nothing about that path"; + + SharedScratchDevice& dev = scratch_pool().get(device_id); + const bool lock_free = dev.mu.try_lock(); + if (lock_free) { + dev.mu.unlock(); + } + ASSERT_TRUE(lock_free) << "a pooled call that returned early left the device's pool lock held, so every later pooled " + "call on this device blocks forever"; + + // The failure is reported by the return value, so the CUDA error the pool's own + // cudaMalloc left has to be cleared where it was made. Left pending it is + // collected by whichever caller makes the next CUDA call on this thread and + // reported under that call's name -- which is what the capture refusal has its + // own case for, from the other direction. + EXPECT_EQ(pending_after_the_failure, cudaSuccess) + << "the pooled call left '" << cudaGetErrorString(pending_after_the_failure) + << "' pending on this thread after reporting the failure through its return value, so the next CUDA call any " + "caller makes collects the pool's error as its own"; + cudaGetLastError(); + + // And the slot is still usable, not merely unlocked: with the memory back, the + // next call allocates and runs. + ASSERT_TRUE(pooled.fill_output(kSentinel)); + EXPECT_EQ(pooled.run(stream), Error::Ok) << "the pool was left unusable by a call whose allocation failed"; + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::vector actual = pooled.read_output(); + ASSERT_EQ(actual.size(), kElems); + EXPECT_NE(actual[0], kSentinel) << "the engine did not write its output, so the run above did not reach it"; + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +// --------------------------------------------------------------------------- +// Stream capture +// --------------------------------------------------------------------------- + +// The pool's handoff waits on an event recorded outside the capture, which +// cudaStreamWaitEvent refuses with cudaErrorStreamCaptureIsolation and which +// invalidates the capture under every capture mode. Left to run, the caller learns +// about it only when cudaStreamEndCapture returns an error and a null graph, far +// from the cause. execute() refuses instead. A growth adds more of the same: its +// allocation invalidates a capture outside cudaStreamCaptureModeRelaxed, and its +// disposal of the buffer it replaces does too -- a call that synchronizes the +// stream is one no capture mode permits, so the only disposal a capture could +// reach is the host wait and the device-wide free. This test starts from a pool +// already large enough and runs in relaxed mode, so it pins the wait rather than +// either of them. +// +// The capture is ended either way: a capture left open belongs to the stream, and +// destroying that stream would abandon it. +TEST_F(SharedScratchBackendTest, APooledEngineRefusesToRunWhileItsStreamIsCapturing) { + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + + LoadedEngine pooled; + ASSERT_EQ(pooled.load(blob(), 20), Error::Ok); + ASSERT_TRUE(pooled.handle()->shared_scratch); + + // One run outside the capture, so the pool holds a buffer and an event with a + // recording against it. Against an empty pool the handoff has nothing to wait + // for, and that wait is the call here that invalidates a capture, so the refusal + // below would be refusing a run that was never a hazard. + ASSERT_EQ(pooled.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + ASSERT_GT(shared_scratch_capacity_for_testing(pooled.handle()->device_id), 0u) + << "the pool holds nothing, so the run below has no handoff to wait on and the capture it is refused for was " + "never at risk"; + + // Relaxed, so the host-side waits execute() makes before it reaches the pool -- + // the one on the previous enqueue's completion event -- are permitted. This test + // is about what the pool does under capture, not about which of the surrounding + // calls a stricter mode would reject first. + ASSERT_EQ(cudaStreamBeginCapture(stream, cudaStreamCaptureModeRelaxed), cudaSuccess); + const Error captured = pooled.run(stream); + cudaGraph_t graph = nullptr; + const cudaError_t end_err = cudaStreamEndCapture(stream, &graph); + if (graph != nullptr) { + cudaGraphDestroy(graph); + } + + EXPECT_EQ(captured, Error::NotSupported) << "execute() did not refuse a pooled run on a capturing stream"; + // The refusal is worth nothing if it still spoiled the capture on the way out: + // a caller told no is expected to end the capture and get its graph. + EXPECT_EQ(end_err, cudaSuccess) << "the refused run left the capture invalidated: " << cudaGetErrorString(end_err); + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +// The refusal clears the CUDA error state only when the capture query itself +// failed, because then the error is the query's own. On the ordinary refusal -- +// the query succeeded and reported a capture -- whatever is pending was left by +// earlier work on this thread, and clearing it takes it away from the caller, +// which learns of it from no later call either: the error is not sticky. +TEST_F(SharedScratchBackendTest, ARefusedPooledCallLeavesTheCallersPendingCudaErrorAlone) { + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + + LoadedEngine pooled; + ASSERT_EQ(pooled.load(blob(), 32), Error::Ok); + ASSERT_TRUE(pooled.handle()->shared_scratch); + + // An error of the caller's own, unread. Nothing but cudaGetLastError clears one, + // so it is still there when the refused call returns unless that call took it. + void* impossible = nullptr; + ASSERT_EQ(cudaMalloc(&impossible, static_cast(1) << 62), cudaErrorMemoryAllocation); + + ASSERT_EQ(cudaStreamBeginCapture(stream, cudaStreamCaptureModeRelaxed), cudaSuccess); + const Error refused = pooled.run(stream); + cudaGraph_t graph = nullptr; + const cudaError_t end_err = cudaStreamEndCapture(stream, &graph); + if (graph != nullptr) { + cudaGraphDestroy(graph); + } + // Read here, before any assertion can return early with it still pending and + // hand it to the next case. + const cudaError_t pending = cudaGetLastError(); + + ASSERT_EQ(refused, Error::NotSupported) << "the run was not refused, so it did not take the path under test"; + EXPECT_EQ(end_err, cudaSuccess) << "the refused run invalidated the capture: " << cudaGetErrorString(end_err); + EXPECT_EQ(pending, cudaErrorMemoryAllocation) + << "the refusal cleared an error this backend did not cause, so the caller never sees it: got " + << cudaGetErrorString(pending); + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +// Where in execute() the refusal comes, which the relaxed-mode case above cannot +// see. Under cudaStreamCaptureModeRelaxed everything execute() does before it +// reaches the pool is permitted, so a guard sitting anywhere ahead of the pool +// looks the same. Under the default Global mode it does not: the host wait on the +// previous enqueue, and further down the cudaMalloc that grows a host-input +// staging buffer, are prohibited and invalidate the capture where they stand. A +// refusal after either is a clean error handed to a caller whose capture is +// already dead, which is the outcome the guard exists to prevent. +// +// The first run is what arms that wait -- it returns with its enqueue still in +// flight, and the next call waits for it before touching the context. +TEST_F(SharedScratchBackendTest, APooledEngineRefusesACaptureBeforeAnythingCanInvalidateIt) { + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + // A node of the caller's own, so a null graph below means the capture was + // invalidated rather than that nothing was ever captured. + void* captured_target = nullptr; + ASSERT_EQ(cudaMalloc(&captured_target, 16), cudaSuccess); + + LoadedEngine pooled; + ASSERT_EQ(pooled.load(blob(), 21), Error::Ok); + ASSERT_TRUE(pooled.handle()->shared_scratch); + ASSERT_EQ(pooled.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + ASSERT_TRUE(pooled.handle()->inflight_pending) + << "the first run left no enqueue in flight, so the next call makes no host wait and this case pins nothing"; + + ASSERT_EQ(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal), cudaSuccess); + ASSERT_EQ(cudaMemsetAsync(captured_target, 0, 16, stream), cudaSuccess); + const Error captured = pooled.run(stream); + cudaGraph_t graph = nullptr; + const cudaError_t end_err = cudaStreamEndCapture(stream, &graph); + const bool have_graph = graph != nullptr; + if (have_graph) { + cudaGraphDestroy(graph); + } + + EXPECT_EQ(captured, Error::NotSupported) + << "execute() did not refuse a pooled run on a stream capturing in the default mode"; + EXPECT_EQ(end_err, cudaSuccess) << "the refused run had already made a call the capture could not take: " + << cudaGetErrorString(end_err); + EXPECT_TRUE(have_graph) << "the capture ended with no graph, so the run invalidated it before refusing"; + + EXPECT_EQ(cudaFree(captured_target), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +// The two cases above are the pool refusing. These two are what a caller gets +// with the option off, where nothing refuses because the guard sits inside the +// pooled branch. Between them they are why the README and the refusal message +// say this delegate does not support capture at all, rather than offering the +// option-off path as the way to capture. +// +// This one takes the capture as far as it goes. A first call over non-empty +// device-resident tensors, submitted on the capturing stream through a +// CallerStreamGuard, makes no call a capture cannot take: the device query, the +// device switch, the shape and address binding, enqueueV3 and the completion +// record, none of which allocates, waits or synchronizes. The graph that comes +// out of it instantiates and replays the engine. +// +// The completion record is the problem. execute() records it so that the next +// call and the destructor can wait for an enqueue that outlived the return, and +// under capture that record becomes a node of the graph rather than an event the +// host can wait on. It costs the handle its next call, and the destructor's wait: +// the call after the failing one records the event again outside the capture and +// works. The assertions below describe what the delegate does, not what it should +// do; they are the measurement the capture documentation rests on. +TEST_F(SharedScratchBackendTest, AnUnpooledEngineCaptureReplaysAndCostsTheHandleItsNextCall) { + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + // A node of the caller's own, so a null graph below means the capture was + // invalidated rather than that nothing was ever captured. + void* captured_target = nullptr; + ASSERT_EQ(cudaMalloc(&captured_target, 16), cudaSuccess); + + // The same engine on the same input, run outside any capture, so the replay + // below has something to be compared against. A separate handle, because a run + // on the capturing handle would leave an enqueue in flight and the capture + // would then fail on the wait for it, which is the other case. + LoadedEngine reference; + ASSERT_EQ(reference.load(blob(), 22), Error::Ok); + ASSERT_EQ(reference.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::vector expected = reference.read_output(); + ASSERT_EQ(expected.size(), kElems); + EXPECT_NE(expected[0], kSentinel) << "the reference output is the sentinel, so a graph that ran nothing would pass"; + + LoadedEngine unpooled; + ASSERT_EQ(unpooled.load(blob(), 22), Error::Ok); + ASSERT_FALSE(unpooled.handle()->shared_scratch) + << "this handle took the pool, so it says nothing about what capturing with the option off gets"; + ASSERT_FALSE(unpooled.handle()->inflight_pending) + << "a freshly loaded handle already reports an enqueue in flight, so the host wait this case is built to avoid " + "would run anyway"; + + ASSERT_EQ(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal), cudaSuccess); + ASSERT_EQ(cudaMemsetAsync(captured_target, 0, 16, stream), cudaSuccess); + const Error captured = unpooled.run(stream); + cudaGraph_t graph = nullptr; + const cudaError_t end_err = cudaStreamEndCapture(stream, &graph); + + EXPECT_EQ(captured, Error::Ok) << "a call with the option off was refused or failed under capture"; + EXPECT_EQ(end_err, cudaSuccess) << "a call with the option off invalidated the capture: " + << cudaGetErrorString(end_err); + ASSERT_NE(graph, nullptr) << "the capture ended with no graph, so the call invalidated it"; + + // The launch is not optional here. A graph destroyed unexecuted says only that + // the capture was not invalidated; it shows neither what the graph computes nor + // what capturing it cost the handle, which is the whole of what this case is for. + cudaGraphExec_t graph_exec = nullptr; + ASSERT_EQ(cudaGraphInstantiate(&graph_exec, graph, 0), cudaSuccess); + ASSERT_TRUE(unpooled.fill_output(kSentinel)); + ASSERT_EQ(cudaGraphLaunch(graph_exec, stream), cudaSuccess); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::vector replayed = unpooled.read_output(); + ASSERT_EQ(replayed.size(), kElems); + EXPECT_EQ(std::memcmp(expected.data(), replayed.data(), kBytes), 0) + << "replaying the captured graph did not reproduce what the same engine produces outside a capture"; + + // And what the capture cost. The completion event was recorded into the graph, + // so the host wait the next call makes on it fails, and the call with it. + const Error after_replay = unpooled.run(stream); + EXPECT_EQ(after_replay, Error::InvalidProgram) + << "a handle whose completion event was recorded inside a captured graph ran again; if that is now supported, " + "the capture section of the README and the refusal message are stale"; + // Not sticky, but left pending it would surface under the name of whatever the + // next case calls first. + (void)cudaGetLastError(); + + // And what it does not cost. The failing call clears the in-flight flag before + // it returns the error, so the call after it skips the wait, records the + // completion event again outside any capture and runs normally -- measured + // here, byte-for-byte against the uncaptured reference. The handle recovers; + // what does not is the ordering the event carried, because a replay of the + // graph enqueues on this same context and no execute() waits for it. + ASSERT_TRUE(unpooled.fill_output(kSentinel)); + const Error after_the_failure = unpooled.run(stream); + EXPECT_EQ(after_the_failure, Error::Ok) + << "the call after the failing one did not recover; the capture section of the README and the option key's " + "documentation say it does, so one of the two is now wrong"; + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::vector recovered = unpooled.read_output(); + ASSERT_EQ(recovered.size(), kElems); + EXPECT_EQ(std::memcmp(expected.data(), recovered.data(), kBytes), 0) + << "the recovered call returned Ok without producing what the same engine produces outside a capture"; + + EXPECT_EQ(cudaGraphExecDestroy(graph_exec), cudaSuccess); + EXPECT_EQ(cudaGraphDestroy(graph), cudaSuccess); + EXPECT_EQ(cudaFree(captured_target), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +// The other half. The case above ends with a handle whose next call the completion +// event has cost it; this one is the failure a caller meets first if the handle +// already ran once. A call that returned with its enqueue in flight leaves the next call on +// that handle waiting for it on the host before it may touch the context, and that +// wait is prohibited under the default capture mode. Nothing refuses it: the +// caller gets an error from execute() and a capture that ends with no graph. +// +// Between the two, a handle with the option off has no state in which capturing it +// is useful: a capture that reaches the completion record succeeds and costs the +// handle its next call, and one made behind a call that left an enqueue in flight +// is destroyed before it gets that far. +TEST_F(SharedScratchBackendTest, AnUnpooledEngineInvalidatesACaptureOnceAnEnqueueIsInFlight) { + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + void* captured_target = nullptr; + ASSERT_EQ(cudaMalloc(&captured_target, 16), cudaSuccess); + + LoadedEngine unpooled; + ASSERT_EQ(unpooled.load(blob(), 23), Error::Ok); + ASSERT_FALSE(unpooled.handle()->shared_scratch); + ASSERT_EQ(unpooled.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + ASSERT_TRUE(unpooled.handle()->inflight_pending) + << "the first run left no enqueue in flight, so the next call makes no host wait and this case pins nothing"; + + ASSERT_EQ(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal), cudaSuccess); + ASSERT_EQ(cudaMemsetAsync(captured_target, 0, 16, stream), cudaSuccess); + const Error captured = unpooled.run(stream); + cudaGraph_t graph = nullptr; + const cudaError_t end_err = cudaStreamEndCapture(stream, &graph); + const bool have_graph = graph != nullptr; + if (have_graph) { + cudaGraphDestroy(graph); + } + // The wait's own cudaErrorStreamCaptureUnsupported is not sticky and execute() + // does not clear it, so it would otherwise surface under the name of whatever + // the next case calls first. + (void)cudaGetLastError(); + + EXPECT_EQ(captured, Error::InvalidProgram) + << "a call following one that left an enqueue in flight did not fail on the host wait for it, so the capture " + "documentation is stale in saying that wait is what a caller meets on a handle that has already run"; + EXPECT_NE(end_err, cudaSuccess) + << "that wait did not invalidate the capture, so the capture documentation is stale in listing it among the " + "calls a capture cannot take"; + EXPECT_FALSE(have_graph) << "the capture ended with a graph despite the host wait, so the capture documentation is " + "stale in listing that wait among the calls a capture cannot take"; + + EXPECT_EQ(cudaFree(captured_target), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +// --------------------------------------------------------------------------- +// 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; +}; + +// A thread that is joined however the case leaves its scope, so an assertion that +// returns early does not leave one joinable and take the process down with it. +class JoinAtScopeExit { + public: + explicit JoinAtScopeExit(std::thread t) : t_(std::move(t)) {} + JoinAtScopeExit(const JoinAtScopeExit&) = delete; + JoinAtScopeExit& operator=(const JoinAtScopeExit&) = delete; + + ~JoinAtScopeExit() { + join(); + } + + void join() { + if (t_.joinable()) { + t_.join(); + } + } + + private: + std::thread t_; +}; + +// Two engines on one device share one scratch buffer, so the second engine's +// enqueue must not start before the first one's has finished with it. The two +// run on different streams, which is what the README permits and what the event +// 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); +} + +// --------------------------------------------------------------------------- +// What a growth waits for +// --------------------------------------------------------------------------- + +// A growth's disposal of the buffer it replaces is one of two things, and each of +// these two cases parks the work one of them waits for. +// +// A call that ends by synchronizing the stream queues the free on that stream, so +// it waits for nothing it did not submit -- in particular not for work parked on +// some unrelated stream, which is what the first case pins. A call that does not +// synchronize cannot queue the free, because the bytes would never come back, so +// it makes a host wait on the marker event and a device-wide cudaFree: it does +// wait for the enqueue against the buffer it retires, and the second case pins +// that it waits rather than freeing under a live enqueue. +// +// Each measures what the call came back before through the gate's watchdog, which +// is the only other thing that can open a gate: if a call waited when it should +// not have, the watchdog would have had to open the gate to end the test, and the +// case says so. + +// Nothing the growing engine or the pool ever submitted to runs on the parked +// stream, so only a device-wide synchronization has any reason to wait for it. +// This is the queued free's whole point, so the case runs on the path that queues +// it: with no caller stream, execute() ends by synchronizing its own stream, which +// is what returns the bytes. +// +// Gated on the device's stream-ordered allocator, and gated before the gate is +// parked so a device without one costs milliseconds rather than the watchdog +// interval. Where there is none, cudaFreeAsync reports cudaErrorNotSupported, the +// backend correctly falls back to the host wait and the device-wide free, and this +// assertion would report a defect that is not there -- the target is built for +// sbsa as well as x86_64, and nothing in the delegate asks the device about memory +// pools, so the fallback is what a whole platform would take. +// +// That fallback is not left untested by the skip. It is the same host wait and +// device-wide free that every growth on a caller stream makes on this device -- +// AGrowthOnACallerStreamWaitsForTheEnqueueOnTheBufferItRetires and +// ALargerEngineGrowsThePoolAndFreesTheBufferItReplaces both run it and check what +// it costs and what it frees. What a device with a memory pool cannot exercise is +// either branch that enters it from a call that does synchronize: cudaFreeAsync +// answering cudaErrorNotSupported, or answering any other code. +TEST_F(SharedScratchBackendTest, AGrowthDoesNotWaitForUnrelatedWorkQueuedOnTheDevice) { + ASSERT_GE(big_scratch_bytes_, scratch_bytes_ + kMinMeasurableScratch) + << "the two fixture engines ask for " << scratch_bytes_ << " and " << big_scratch_bytes_ + << " bytes of activation scratch, too close for the second to be sure of growing the pool"; + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + + cudaStream_t unrelated = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&unrelated, cudaStreamNonBlocking), cudaSuccess); + + LoadedEngine small; + LoadedEngine big; + ASSERT_EQ(small.load(blob(), 26), Error::Ok); + ASSERT_EQ(big.load(big_blob(), 27, kBigRows, kBigCols), Error::Ok); + ASSERT_TRUE(small.handle()->shared_scratch); + ASSERT_TRUE(big.handle()->shared_scratch); + const int device_id = big.handle()->device_id; + + int memory_pools = 0; + ASSERT_EQ(cudaDeviceGetAttribute(&memory_pools, cudaDevAttrMemoryPoolsSupported, device_id), cudaSuccess); + if (memory_pools == 0) { + GTEST_SKIP() << "device " << device_id + << " reports no stream-ordered allocator (cudaDevAttrMemoryPoolsSupported = 0), so a growth here " + "correctly takes the host wait and the device-wide free, which does wait for work parked " + "anywhere on the device"; + } + + // Without this the larger engine allocates rather than grows, and a growth that + // retires nothing disposes of nothing. + ASSERT_EQ(small.run_on_the_synchronized_path(), Error::Ok); + const std::size_t before = shared_scratch_capacity_for_testing(device_id); + ASSERT_GT(before, 0u) << "the smaller engine left the pool empty, so the larger one has nothing to retire"; + + StreamGate gate; + ASSERT_EQ(cudaLaunchHostFunc(unrelated, hold_stream, &gate), cudaSuccess); + GateRelease gate_release(gate, unrelated); + + ASSERT_TRUE(big.fill_output(kSentinel)); + const Error growth_error = big.run_on_the_synchronized_path(); + const bool came_back_with_the_device_held = !gate.forced_open.load(); + + gate_release.release(); + + EXPECT_EQ(growth_error, Error::Ok); + EXPECT_GT(shared_scratch_capacity_for_testing(device_id), before) + << "the pool did not grow, so no buffer was retired and nothing was disposed of"; + const std::vector grown_output = big.read_output(); + ASSERT_FALSE(grown_output.empty()); + EXPECT_NE(grown_output[0], kSentinel) + << "the growing engine did not write its output, so it never reached the engine"; + + EXPECT_TRUE(came_back_with_the_device_held) + << "the growing call did not return until the watchdog released a host function parked on a stream it never " + "submitted to, so its disposal of the retired buffer waits for the whole device"; + + ASSERT_EQ(cudaStreamDestroy(unrelated), cudaSuccess); +} + +// The other disposal. On a caller stream with device-resident I/O nothing will +// ever synchronize the stream, so the growth cannot queue the free -- the bytes +// would stay resident until the caller happened to synchronize, and every buffer a +// run retired would be resident at once. It waits on the marker event and frees +// device-wide instead, so it *does* wait for the enqueue against the buffer it +// retires, and what has to hold is that it waits rather than freeing under it. +// +// The smaller engine's own enqueue is parked behind a host function, and a thread +// releases it a short while after the growing call starts. A growth that waited +// comes back after that release; one that freed the retired buffer under the live +// enqueue comes back before it. +// +// The smaller engine's output is checked against the same engine run with private +// scratch, because it is the one whose buffer was retired underneath it: a free +// that landed early would take its activations with it. +TEST_F(SharedScratchBackendTest, AGrowthOnACallerStreamWaitsForTheEnqueueOnTheBufferItRetires) { + ASSERT_GE(big_scratch_bytes_, scratch_bytes_ + kMinMeasurableScratch) + << "the two fixture engines ask for " << scratch_bytes_ << " and " << big_scratch_bytes_ + << " bytes of activation scratch, too close for the second to be sure of growing the pool"; + + // Long enough that the growing call cannot come back after it by accident, short + // enough not to lengthen the suite. + constexpr std::chrono::milliseconds kHoldTheEnqueueFor{500}; + + cudaStream_t held = nullptr; + cudaStream_t growing = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&held, cudaStreamNonBlocking), cudaSuccess); + ASSERT_EQ(cudaStreamCreateWithFlags(&growing, cudaStreamNonBlocking), cudaSuccess); + + // Loaded with the option off, so this one keeps its own scratch and its output + // is what the pooled run below has to reproduce. + LoadedEngine reference; + ASSERT_EQ(set_shared_scratch(backend_, false), Error::Ok); + ASSERT_EQ(reference.load(blob(), 29), Error::Ok); + ASSERT_FALSE(reference.handle()->shared_scratch); + ASSERT_EQ(reference.run(held), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(held), cudaSuccess); + const std::vector expected = reference.read_output(); + ASSERT_EQ(expected.size(), kElems); + + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + LoadedEngine small; + LoadedEngine big; + ASSERT_EQ(small.load(blob(), 29), Error::Ok); + ASSERT_EQ(big.load(big_blob(), 30, kBigRows, kBigCols), Error::Ok); + ASSERT_TRUE(small.handle()->shared_scratch); + ASSERT_TRUE(big.handle()->shared_scratch); + const int device_id = big.handle()->device_id; + ASSERT_TRUE(small.fill_output(kSentinel)); + + // Parked ahead of the smaller engine's enqueue, so that enqueue and the event + // recorded after it both stay pending for as long as this test wants them to. + StreamGate gate; + ASSERT_EQ(cudaLaunchHostFunc(held, hold_stream, &gate), cudaSuccess); + GateRelease gate_release(gate, held); + + ASSERT_EQ(small.run(held), Error::Ok); + const std::size_t before = shared_scratch_capacity_for_testing(device_id); + ASSERT_GT(before, 0u) << "the smaller engine left the pool empty, so the larger one has nothing to retire"; + + // Opens the gate, without the stream synchronization GateRelease::release() also + // makes: that one runs on this thread at scope exit, after the join below, so the + // two never touch GateRelease at once. + std::atomic the_enqueue_was_released{false}; + JoinAtScopeExit opener{std::thread([&] { + std::this_thread::sleep_for(kHoldTheEnqueueFor); + the_enqueue_was_released.store(true); + { + std::lock_guard lock(gate.mu); + gate.open = true; + } + gate.cv.notify_all(); + })}; + + const Error growth_error = big.run(growing); + const bool waited_for_the_parked_enqueue = the_enqueue_was_released.load(); + opener.join(); + + gate_release.release(); + ASSERT_EQ(cudaStreamSynchronize(held), cudaSuccess); + ASSERT_EQ(cudaStreamSynchronize(growing), cudaSuccess); + + EXPECT_EQ(growth_error, Error::Ok); + ASSERT_FALSE(gate.forced_open.load()) + << "the watchdog had to open the gate, so the growing call blocked for the whole watchdog interval rather than " + "for the enqueue it had to wait for"; + EXPECT_TRUE(waited_for_the_parked_enqueue) + << "the growing call returned while the enqueue against the buffer it retired was still parked, so it did not " + "wait for that enqueue -- either it freed the buffer under it or it only queued the free, which on this path " + "never gives the bytes back"; + EXPECT_GT(shared_scratch_capacity_for_testing(device_id), before) + << "the pool did not grow, so no buffer was retired and there was nothing to order the free against"; + const std::vector small_output = small.read_output(); + ASSERT_EQ(small_output.size(), kElems); + EXPECT_EQ(std::memcmp(expected.data(), small_output.data(), kBytes), 0) + << "the engine whose scratch buffer the growth retired did not produce what the same engine produces with " + "private scratch, so the buffer went away while its enqueue was still using it"; + + ASSERT_EQ(cudaStreamDestroy(held), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(growing), cudaSuccess); +} + +// Neither growth case above can see whether the disposal is made with the device's +// lock held: both run their growth alone, so a lock held across it holds nothing +// up and both stay green with the unlock moved below the free. This one is what +// covers that. +// +// It matters because the disposal on this path is a host wait and a cudaFree, and +// cudaFree waits for every stream on the device rather than for the ones that +// touched the buffer -- the 1505 ms against a host function parked for 1500 ms on +// an unrelated stream that the execute() contract cites. Under the lock, that wait +// is one every other pooled engine on the device has to sit through, which is the +// serialization the contract says a pooled call does not impose past its own +// submission. +// +// So: park work on an unrelated stream to stall the free, then check the pool lock +// while the growth is still inside it. The capacity read under that same try_lock +// is what stops the case passing on a lock that is free because the growth has not +// started -- past a growth it is above what the smaller engine left, and the growth +// has not returned. +TEST_F(SharedScratchBackendTest, AGrowthOnACallerStreamDisposesWithTheDeviceLockDropped) { + ASSERT_GE(big_scratch_bytes_, scratch_bytes_ + kMinMeasurableScratch) + << "the two fixture engines ask for " << scratch_bytes_ << " and " << big_scratch_bytes_ + << " bytes of activation scratch, too close for the second to be sure of growing the pool"; + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + + cudaStream_t stream = nullptr; + cudaStream_t unrelated = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + ASSERT_EQ(cudaStreamCreateWithFlags(&unrelated, cudaStreamNonBlocking), cudaSuccess); + + LoadedEngine small; + LoadedEngine big; + ASSERT_EQ(small.load(blob(), 32), Error::Ok); + ASSERT_EQ(big.load(big_blob(), 33, kBigRows, kBigCols), Error::Ok); + ASSERT_TRUE(small.handle()->shared_scratch); + ASSERT_TRUE(big.handle()->shared_scratch); + const int device_id = big.handle()->device_id; + + // Without this the larger engine allocates rather than grows, and a growth that + // retires nothing frees nothing. Synchronized, so the growth's wait on the + // marker returns at once and the free is the only thing left to stall it. + ASSERT_EQ(small.run(stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + const std::size_t before = shared_scratch_capacity_for_testing(device_id); + ASSERT_GT(before, 0u) << "the smaller engine left the pool empty, so the larger one has nothing to retire"; + + // Neither engine submits to this stream, so nothing on the growth path but the + // device-wide free has any reason to wait for it. + StreamGate gate; + ASSERT_EQ(cudaLaunchHostFunc(unrelated, hold_stream, &gate), cudaSuccess); + GateRelease gate_release(gate, unrelated); + + std::atomic growth_returned{false}; + Error growth_error = Error::Internal; + JoinAtScopeExit grower{std::thread([&] { + // A caller stream with device-resident I/O, which is the path whose disposal + // is the host wait and the device-wide free. + growth_error = big.run(stream); + growth_returned.store(true); + })}; + + SharedScratchDevice& dev = scratch_pool().get(device_id); + bool lock_free_during_the_disposal = false; + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (!growth_returned.load() && std::chrono::steady_clock::now() < deadline) { + if (dev.mu.try_lock()) { + const std::size_t capacity_now = dev.capacity; + dev.mu.unlock(); + if (capacity_now > before && !growth_returned.load()) { + lock_free_during_the_disposal = true; + break; + } + } + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + const bool still_inside_execute = !growth_returned.load(); + + gate_release.release(); + grower.join(); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(unrelated), cudaSuccess); + + ASSERT_FALSE(gate.forced_open.load()) + << "the watchdog had to open the gate: a call blocked on the held stream rather than returning, so nothing " + "below was measured under the conditions it describes"; + ASSERT_EQ(growth_error, Error::Ok); + ASSERT_TRUE(still_inside_execute) + << "the growing call returned before the gate opened, so its free never stalled and there was no window in " + "which to observe the lock"; + ASSERT_GT(shared_scratch_capacity_for_testing(device_id), before) + << "the pool did not grow, so no buffer was retired and no free was made"; + EXPECT_TRUE(lock_free_during_the_disposal) + << "the device's pool lock stayed held for the whole of a stalled growth free, so every other pooled engine on " + "this device waits out a device-wide free it has nothing to do with"; +} + +// --------------------------------------------------------------------------- +// A failing call and the host-input copy it queued +// --------------------------------------------------------------------------- + +// An input that is not device-resident is staged with cudaMemcpyAsync from the +// caller's own memory, and from pinned memory that copy reads the caller's bytes +// when the stream reaches it, not when it is queued. So a call that fails after +// queueing one must not return while it is still pending: the caller owns that +// buffer again the moment execute() returns, and what it writes there is what +// the device then receives. +// +// The failure driven here is the pool's allocation, the one early return in the +// pooled scratch block that can be reached from inside this process. The stream +// is parked ahead of the copy so it cannot complete on its own, and the caller +// overwrites its buffer as soon as the call comes back -- which is exactly what a +// caller may do. +// +// Three outcomes are separated at the end, by giving the failing call a +// different input from the successful one before it: the sentinel means the copy +// took the caller's post-return write, the first call's pattern means it never +// ran at all, and the second call's pattern is the only pass. Which of the three +// happens is decided by where the caller's write falls relative to the copy and +// not by any interval this case picks, so a slow return cannot turn the first +// outcome into the third. +TEST_F(SharedScratchBackendTest, AFailedPooledCallDrainsTheHostInputCopyItQueued) { + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + + cudaStream_t stream = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + + LoadedEngine pooled; + ASSERT_EQ(pooled.load(blob(), 41), Error::Ok); + ASSERT_TRUE(pooled.handle()->shared_scratch); + const int device_id = pooled.handle()->device_id; + + // Pinned: cudaMemcpyAsync from pageable memory does not return until the + // driver has taken the bytes, so a later write cannot reach the device and + // there is no hazard to pin. + float* host_in = nullptr; + ASSERT_EQ(cudaHostAlloc(reinterpret_cast(&host_in), pooled.bytes(), cudaHostAllocDefault), cudaSuccess); + for (std::size_t i = 0; i < pooled.elems(); ++i) { + host_in[i] = pattern(i, 41); + } + + // One good run first, so the failing one reaches the copy: the staging buffer + // is allocated on the call that first needs it, and with the device full that + // allocation would fail ahead of everything this case is about. + ASSERT_EQ(pooled.run_from_host_input(stream, host_in), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + void* const staging = pooled.staging_buffer_for_input_0(); + ASSERT_NE(staging, nullptr) << "the run bound the caller's host memory directly, so nothing was staged and there is " + "no asynchronous copy to leave running"; + + // The second call's input, so what the device holds at the end says which of + // the three outcomes happened. + for (std::size_t i = 0; i < pooled.elems(); ++i) { + host_in[i] = pattern(i, 42); + } + + // Empties the pool, so the failing run allocates instead of reusing what the + // run above left in it. + ASSERT_TRUE(reset_shared_scratch_pool_for_testing()); + ASSERT_EQ(shared_scratch_capacity_for_testing(device_id), 0u); + + DeviceMemoryHog hog; + if (!hog.leave_less_free_than(scratch_bytes_)) { + ASSERT_EQ(cudaFreeHost(host_in), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); + GTEST_SKIP() << "could not take the device below the " << scratch_bytes_ + << " bytes this engine's scratch needs, so its allocation would have succeeded"; + } + + StreamGate gate; + ASSERT_EQ(cudaLaunchHostFunc(stream, hold_stream, &gate), cudaSuccess); + GateRelease gate_release(gate, stream); + + // The caller overwrites its own buffer the moment the call comes back, in the + // thread that made the call, which is what a caller is entitled to do and is + // what decides this case. A build that returns with the copy still queued does + // that write before the gate opens, so the copy reads the sentinel and the + // comparison at the end sees it. A build that drains cannot reach the write + // until the copy has run. Neither outcome depends on how long anything takes: + // the wait below is a bound on this case, not the thing that discriminates. + std::atomic call_returned{false}; + std::atomic caller_overwrote_its_buffer{false}; + Error failed_run = Error::Ok; + std::thread caller([&] { + failed_run = pooled.run_from_host_input(stream, host_in); + call_returned.store(true); + for (std::size_t i = 0; i < pooled.elems(); ++i) { + host_in[i] = kSentinel; + } + caller_overwrote_its_buffer.store(true); + }); + + // Waited for rather than slept through, so a build that returns early is not + // handed the rest of the interval to finish its write in. Reaching the deadline + // is the passing outcome: the stream cannot move until the gate opens, so a call + // that waits for its own copy is still inside execute() here. + const auto give_up_waiting_at = std::chrono::steady_clock::now() + std::chrono::milliseconds(500); + while (!caller_overwrote_its_buffer.load() && std::chrono::steady_clock::now() < give_up_waiting_at) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + const bool returned_with_the_copy_still_queued = call_returned.load(); + gate_release.release(); + caller.join(); + // After the gate, because a free is device-wide and would otherwise block on it. + hog.release(); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + std::vector staged(pooled.elems(), 0.0f); + ASSERT_EQ(cudaMemcpy(staged.data(), staging, pooled.bytes(), cudaMemcpyDeviceToHost), cudaSuccess); + ASSERT_EQ(cudaFreeHost(host_in), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); + + ASSERT_FALSE(gate.forced_open.load()) + << "the watchdog had to open the gate, so the timings below were not measured under the conditions they describe"; + if (failed_run == Error::Ok) { + GTEST_SKIP() << "the pooled allocation succeeded with the device full, so this run did not fail where the case " + "needs it to"; + } + ASSERT_EQ(failed_run, Error::MemoryAllocationFailed) + << "the run failed somewhere other than the pool's allocation, so it says nothing about that early return"; + + EXPECT_FALSE(returned_with_the_copy_still_queued) + << "the failing call returned while the copy reading the caller's host input was still queued behind a parked " + "stream, so whatever the caller writes next is what the device receives"; + std::size_t bytes_the_device_did_not_get = 0; + for (std::size_t i = 0; i < pooled.elems(); ++i) { + if (staged[i] != pattern(i, 42)) { + ++bytes_the_device_did_not_get; + } + } + EXPECT_EQ(bytes_the_device_did_not_get, 0u) + << "the device does not hold the input this call was made with. It holds " << staged[0] + << " at element 0: " << kSentinel << " is what the caller wrote after execute() returned, " << pattern(0, 41) + << " is the previous call's input and means the copy never ran, and " << pattern(0, 42) << " is a pass"; +} + +// --------------------------------------------------------------------------- +// The pooled path under two concurrent callers +// --------------------------------------------------------------------------- + +// The window a dropped lock would leave open is a few microseconds wide, so one +// pair of runs would find it only by luck. At this count a build that leaves the +// window open loses most of the runs, and the test costs about two seconds. +constexpr int kConcurrentRunsPerThread = 60; + +// Two pooled engines on one device, submitted from two threads on two streams, +// with nothing but the backend ordering them. Both are backed by the same buffer, +// so a claim that ends before the enqueue is recorded hands a second caller the +// same scratch with nothing ordering the two -- silently wrong output, no CUDA +// error, no TensorRT error. Each thread compares byte-for-byte against what its own +// engine produces with private scratch. +TEST_F(SharedScratchBackendTest, TwoThreadsRunningPooledEnginesOnOneDeviceKeepTheirOwnOutputs) { + cudaStream_t reference_stream = nullptr; + ASSERT_EQ(cudaStreamCreate(&reference_stream), cudaSuccess); + + std::vector first_expected; + std::vector second_expected; + { + LoadedEngine first_priv; + LoadedEngine second_priv; + ASSERT_EQ(first_priv.load(blob(), 17), Error::Ok); + ASSERT_EQ(second_priv.load(blob(), 18), Error::Ok); + ASSERT_FALSE(first_priv.handle()->shared_scratch); + ASSERT_FALSE(second_priv.handle()->shared_scratch); + ASSERT_EQ(first_priv.run(reference_stream), Error::Ok); + ASSERT_EQ(second_priv.run(reference_stream), Error::Ok); + ASSERT_EQ(cudaStreamSynchronize(reference_stream), cudaSuccess); + first_expected = first_priv.read_output(); + second_expected = second_priv.read_output(); + } + ASSERT_EQ(cudaStreamDestroy(reference_stream), cudaSuccess); + + ASSERT_EQ(first_expected.size(), kElems); + ASSERT_EQ(second_expected.size(), kElems); + // Two engines producing the same bytes would let each thread pass on the other + // one's output, which is the outcome this test exists to catch. + ASSERT_NE(std::memcmp(first_expected.data(), second_expected.data(), kBytes), 0) + << "the two engines were given different inputs but produced the same output"; + ASSERT_NE(first_expected[0], kSentinel) << "the reference output is the sentinel, so a skipped enqueue would pass"; + ASSERT_NE(second_expected[0], kSentinel) << "the reference output is the sentinel, so a skipped enqueue would pass"; + + ASSERT_EQ(set_shared_scratch(backend_, true), Error::Ok); + LoadedEngine first; + LoadedEngine second; + ASSERT_EQ(first.load(blob(), 17), Error::Ok); + ASSERT_EQ(second.load(blob(), 18), Error::Ok); + ASSERT_TRUE(first.handle()->shared_scratch); + ASSERT_TRUE(second.handle()->shared_scratch); + + cudaStream_t first_stream = nullptr; + cudaStream_t second_stream = nullptr; + ASSERT_EQ(cudaStreamCreateWithFlags(&first_stream, cudaStreamNonBlocking), cudaSuccess); + ASSERT_EQ(cudaStreamCreateWithFlags(&second_stream, cudaStreamNonBlocking), cudaSuccess); + + // The host copies that bracket each run synchronize the whole device, so two + // threads left to themselves take turns rather than overlap. Against a build + // that leaves the window open, taking turns caught it in 2 of the 120 runs + // below; releasing both threads together caught nearly all of them. + // + // The wait has a deadline because the loop it sits in contains a delegate call, + // and a pooled call that blocks is one of the failures this case exists to + // catch. Without one the thread whose partner is stuck spins here until the + // target's own timeout, which reports as a killed binary rather than as this + // case; with one, the run ends and the assertion below names it. + // + // The loop below stops at the first rendezvous that expires, which is what makes + // the deadline affordable. It is armed per rendezvous -- a whole-case wall clock + // would have to cover 60 real runs and would fail a slow machine instead -- and + // paid at most once per thread, so the worst case is a bounded 30 s rather than + // the 60 x 30 s a loop that carried on would cost against this target's 900 s + // timeout. Stopping also keeps the surviving thread out of the delegate call + // that follows the rendezvous, which its wedged partner may be holding the + // device's pool lock across. + constexpr std::chrono::seconds kRendezvousDeadline{30}; + std::atomic arrived{0}; + std::atomic partner_never_arrived{false}; + auto submit_together = [&](int iteration) { + arrived.fetch_add(1); + const auto deadline = std::chrono::steady_clock::now() + kRendezvousDeadline; + while (arrived.load() < 2 * (iteration + 1)) { + if (partner_never_arrived.load() || std::chrono::steady_clock::now() >= deadline) { + partner_never_arrived.store(true); + return; + } + std::this_thread::yield(); + } + }; + + std::atomic wrong_outputs{0}; + std::atomic failures{0}; + auto run_repeatedly = [&](LoadedEngine& engine, const std::vector& expected, cudaStream_t stream) { + for (int i = 0; i < kConcurrentRunsPerThread; ++i) { + // Rewritten every iteration, so a run whose enqueue never reached the engine + // leaves the sentinel behind rather than the previous iteration's output. + if (!engine.fill_output(kSentinel)) { + failures.fetch_add(1); + } + submit_together(i); + if (partner_never_arrived.load()) { + break; + } + if (engine.run(stream) != Error::Ok || cudaStreamSynchronize(stream) != cudaSuccess) { + failures.fetch_add(1); + continue; + } + const std::vector actual = engine.read_output(); + if (actual.size() != expected.size() || std::memcmp(actual.data(), expected.data(), kBytes) != 0) { + wrong_outputs.fetch_add(1); + } + } + }; + + std::thread first_thread([&] { run_repeatedly(first, first_expected, first_stream); }); + std::thread second_thread([&] { run_repeatedly(second, second_expected, second_stream); }); + first_thread.join(); + second_thread.join(); + + ASSERT_EQ(cudaStreamDestroy(first_stream), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(second_stream), cudaSuccess); + + EXPECT_FALSE(partner_never_arrived.load()) + << "a thread waited out the " << kRendezvousDeadline.count() + << "-second rendezvous deadline without its partner arriving, so a pooled call blocked instead of returning and " + "both threads stopped short of " + << kConcurrentRunsPerThread << " runs"; + EXPECT_EQ(failures.load(), 0) << "a run failed outright, so fewer than " << (2 * kConcurrentRunsPerThread) + << " runs reached the comparison below"; + EXPECT_EQ(wrong_outputs.load(), 0) << wrong_outputs.load() << " of " << (2 * kConcurrentRunsPerThread) + << " concurrent pooled runs did not produce what the same engine produces with " + "its own scratch"; +} + +} // namespace +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/tests/cpp/executorch/test_shared_scratch_pool.cpp b/tests/cpp/executorch/test_shared_scratch_pool.cpp new file mode 100644 index 0000000000..acd446b302 --- /dev/null +++ b/tests/cpp/executorch/test_shared_scratch_pool.cpp @@ -0,0 +1,1023 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Pins the shared scratch pool helper: its grow, reuse and per-device policy and +// its enqueue-handoff rule, driven over fakes so no CUDA device is needed. +// +// This exercises the helper, not the backend: it does not link the delegate, so +// it cannot catch the delegate calling the helper wrongly or ceasing to call it. +// test_shared_scratch_backend covers that, and needs a GPU to do it. +// +// One case at the end is the exception to "over fakes": the lifetime of the +// process-wide registry scratch_pool() hands out is a property of that object and +// not of any fake, and it is observable only after main returns. It forks to see +// it, and still makes no CUDA call. + +#include "torch_tensorrt/executorch/SharedScratchPool.h" +#include "torch_tensorrt/executorch/SharedScratchPoolReset.h" + +#include "gtest/gtest.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(__unix__) +#include +#include + +#include +#include +#endif + +namespace torch_tensorrt { +namespace executorch_backend { + +// Defined in test_shared_scratch_pool_other_tu.cpp; see +// EveryTranslationUnitSeesOneProcessPool below. +const SharedScratchPool* scratch_pool_seen_from_another_translation_unit(); + +namespace { + +// Fake device allocator: hands out distinct non-null pointers and records every +// allocation size and every buffer a growth retired, so tests can assert the +// pool's grow/reuse policy and what each retirement has to wait for, without a +// CUDA device. +struct FakeAllocator { + std::vector alloc_sizes; + std::vector> retirements; + 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 retire(void* p, cudaEvent_t wait_for) { + retirements.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; + } +}; + +// Stands in for the backend: passes the allocator through and records whatever +// the call retired, the way execute() hands a retired buffer to its claim. +void* call(SharedScratchDevice& dev, FakeAllocator& a, std::size_t need) { + RetiredScratch retired; + void* const p = shared_scratch_get_or_grow( + dev, need, [&a](std::size_t bytes) { return a.alloc(bytes); }, retired); + if (retired.buffer != nullptr) { + a.retire(retired.buffer, retired.wait_for); + } + return p; +} + +TEST(SharedScratchPool, FirstRequestAllocatesExactSize) { + SharedScratchDevice dev; + FakeAllocator a; + + void* p = call(dev, a, /*need=*/1024); + + EXPECT_NE(p, nullptr); + EXPECT_EQ(dev.capacity, 1024u); + ASSERT_EQ(a.alloc_count(), 1); + EXPECT_EQ(a.alloc_sizes[0], 1024u); + EXPECT_TRUE(a.retirements.empty()); +} + +TEST(SharedScratchPool, ReusesWhenExistingBufferIsLargeEnough) { + SharedScratchDevice dev; + FakeAllocator a; + + void* first = call(dev, a, 4096); + // A smaller and an equal request must both reuse the same buffer (no realloc). + void* second = call(dev, a, 1000); + void* third = call(dev, a, 4096); + + EXPECT_EQ(second, first); + EXPECT_EQ(third, first); + // The smaller request did not shrink the pool to what it asked for. + EXPECT_EQ(dev.capacity, 4096u); + EXPECT_EQ(a.alloc_count(), 1); + EXPECT_TRUE(a.retirements.empty()); +} + +TEST(SharedScratchPool, GrowsMonotonicallyToMaxAndRetiresOldBuffer) { + SharedScratchDevice dev; + FakeAllocator a; + + void* small = call(dev, a, 1024); + void* big = call(dev, a, 8192); + + EXPECT_NE(big, small); + EXPECT_EQ(dev.capacity, 8192u); + ASSERT_EQ(a.alloc_count(), 2); + EXPECT_EQ(a.alloc_sizes[1], 8192u); + ASSERT_EQ(a.retirements.size(), 1u); + EXPECT_EQ(a.retirements[0].first, small); + + // A subsequent smaller request reuses the grown buffer -- pool never shrinks. + void* reuse = call(dev, a, 512); + EXPECT_EQ(reuse, big); + EXPECT_EQ(dev.capacity, 8192u); + EXPECT_EQ(a.alloc_count(), 2); +} + +TEST(SharedScratchPool, GrowRetiresTheOldBufferWithTheEventToWaitOn) { + SharedScratchDevice dev; + FakeAllocator a; + FakeEventFactory events; + + void* small = call(dev, a, 1024); + ASSERT_NE(small, nullptr); + + // An enqueue against `small` has been submitted and recorded, so its + // retirement has something specific to outlive. + const SharedScratchHandoff handoff = shared_scratch_claim_event(dev, std::ref(events)); + ASSERT_EQ(shared_scratch_mark_in_flight(dev), handoff.event); + + ASSERT_NE(call(dev, a, 8192), nullptr); + + ASSERT_EQ(a.retirements.size(), 1u); + EXPECT_EQ(a.retirements[0].first, small); + // The retirement carries the event that enqueue was recorded on, so the caller + // has one specific enqueue to wait for, rather than needing a device-wide + // synchronize to be correct. + EXPECT_EQ(a.retirements[0].second, handoff.event); +} + +TEST(SharedScratchPool, GrowHasNothingToWaitForWhenNoEnqueueWasRecorded) { + SharedScratchDevice dev; + FakeAllocator a; + FakeEventFactory events; + + void* small = call(dev, a, 1024); + ASSERT_NE(small, nullptr); + // The slot has an event, but nothing has been recorded on it: claiming the + // handoff is not the same as enqueueing against the buffer. + ASSERT_NE(shared_scratch_claim_event(dev, std::ref(events)).event, nullptr); + + ASSERT_NE(call(dev, a, 8192), nullptr); + + ASSERT_EQ(a.retirements.size(), 1u); + EXPECT_EQ(a.retirements[0].first, small); + EXPECT_EQ(a.retirements[0].second, nullptr); +} + +TEST(SharedScratchPool, AllocationFailureLeavesExistingBufferUntouched) { + SharedScratchDevice dev; + FakeAllocator a; + + void* first = call(dev, a, 1024); + ASSERT_NE(first, nullptr); + + // A growth whose allocation fails must return nullptr and keep the old buffer, + // so the caller can surface the error without corrupting the pool. + a.fail_next = true; + void* failed = call(dev, a, 8192); + EXPECT_EQ(failed, nullptr); + EXPECT_TRUE(a.retirements.empty()); + + // The device still holds the original buffer and serves it on the next request. + void* again = call(dev, a, 1024); + EXPECT_EQ(again, first); + EXPECT_EQ(dev.capacity, 1024u); +} + +TEST(SharedScratchPool, FirstAllocationFailureReturnsNullAndStoresNothing) { + SharedScratchDevice dev; + FakeAllocator a; + + a.fail_next = true; + void* p = call(dev, a, 1024); + EXPECT_EQ(p, nullptr); + EXPECT_EQ(dev.buffer, nullptr); + EXPECT_EQ(dev.capacity, 0u); + + // Nothing stored: a later successful request allocates fresh. + void* q = call(dev, a, 1024); + EXPECT_NE(q, nullptr); + EXPECT_EQ(a.alloc_count(), 1); +} + +// The backend passes one RetiredScratch through a call and acts on what it holds, +// so a path that leaves the previous growth's buffer in it hands that buffer to a +// second caller to free. +TEST(SharedScratchPool, EveryPathClearsTheRetirementItReports) { + SharedScratchDevice dev; + FakeAllocator a; + // Deliberately reused across the calls below, which is the state the clearing is + // for; `call` above gives each of its calls a fresh one. + RetiredScratch retired; + const auto request = [&](std::size_t need) { + return shared_scratch_get_or_grow( + dev, need, [&a](std::size_t bytes) { return a.alloc(bytes); }, retired); + }; + + void* const first = request(1024); + ASSERT_NE(first, nullptr); + ASSERT_EQ(retired.buffer, nullptr); + ASSERT_NE(request(8192), nullptr); + ASSERT_EQ(retired.buffer, first) << "the growth did not report the buffer it replaced"; + + EXPECT_NE(request(512), nullptr); + EXPECT_EQ(retired.buffer, nullptr) << "a reuse left the previous growth's buffer in the result, so a caller acting " + "on it frees a buffer that was already handed back once"; + + void* const second = dev.buffer; + ASSERT_NE(request(16384), nullptr); + ASSERT_EQ(retired.buffer, second); + a.fail_next = true; + EXPECT_EQ(request(65536), nullptr); + EXPECT_EQ(retired.buffer, nullptr) << "a failed allocation left the previous growth's buffer in the result"; +} + +// --------------------------------------------------------------------------- +// Ordering the shared buffer's handoff from one enqueue to the next. +// --------------------------------------------------------------------------- + +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; + + void* dev0 = call(pool.get(0), a, 2048); + void* dev1 = call(pool.get(1), a, 2048); + + EXPECT_NE(dev0, dev1); + EXPECT_EQ(a.alloc_count(), 2); + EXPECT_TRUE(a.retirements.empty()); + + // Growing device 1 must not touch device 0's buffer. + void* dev1_big = call(pool.get(1), a, 9000); + void* dev0_again = call(pool.get(0), a, 2048); + EXPECT_NE(dev1_big, dev1); + EXPECT_EQ(dev0_again, dev0); + ASSERT_EQ(a.retirements.size(), 1u); + EXPECT_EQ(a.retirements[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); +} + +// scratch_pool() is inline so that the backend and the pool's test hooks, which +// are separate translation units, reach one registry. Nothing else here can see +// that: every case above builds its own SharedScratchPool. Give the accessor +// internal linkage instead and each translation unit gets a registry of its own, +// at which point the reset hook clears one nobody uses, the capacity hook always +// answers zero, and several of the backend suite's pool assertions pass without +// testing anything -- on a GPU host, which is the only place they run. +// +// The second translation unit is test_shared_scratch_pool_other_tu.cpp. +TEST(SharedScratchPoolRegistry, EveryTranslationUnitSeesOneProcessPool) { + EXPECT_EQ(scratch_pool_seen_from_another_translation_unit(), &scratch_pool()) + << "two translation units got different process pools, so scratch_pool() no longer has external linkage and " + "the pool's test hooks operate on a registry the backend never uses"; +} + +TEST(SharedScratchPoolRegistry, ResetHandsBackEverySlotAndLeavesItEmpty) { + // The backend test fixture runs this between cases, so every case that reads + // what the pool holds depends on it. The two things it has to get right are + // handing the caller both the buffer and the event to release -- neither the + // pool nor the disposer can find them afterwards -- and clearing the slot, so + // the next claim allocates instead of reusing a pointer that has been freed. + SharedScratchPool pool; + FakeAllocator zero; + FakeAllocator one; + FakeEventFactory events; + + SharedScratchDevice& dev0 = pool.get(0); + SharedScratchDevice& dev1 = pool.get(1); + void* const dev0_buffer = call(dev0, zero, 4096); + void* const dev1_buffer = call(dev1, one, 2048); + const cudaEvent_t dev0_event = shared_scratch_claim_event(dev0, std::ref(events)).event; + shared_scratch_mark_in_flight(dev0); + // Device 1 is left with a buffer and no event, which is the state of a slot + // whose event creation failed: the reset has to cope with a null there. + ASSERT_NE(dev0_buffer, nullptr); + ASSERT_NE(dev1_buffer, nullptr); + ASSERT_NE(dev0_event, nullptr); + ASSERT_EQ(dev1.marker.event, nullptr); + + struct Disposal { + int device_id; + void* buffer; + cudaEvent_t event; + }; + std::vector disposed; + reset_shared_scratch_pool_slots(pool, [&](int device_id, void* buffer, cudaEvent_t event) { + disposed.push_back({device_id, buffer, event}); + }); + + // The registry iterates in unspecified order, so each device is looked up. + const auto disposal_for = [&disposed](int device_id) -> const Disposal* { + for (const Disposal& d : disposed) { + if (d.device_id == device_id) { + return &d; + } + } + return nullptr; + }; + ASSERT_EQ(disposed.size(), 2u) << "the reset skipped a device's slot, whose buffer is then never freed"; + ASSERT_NE(disposal_for(0), nullptr); + ASSERT_NE(disposal_for(1), nullptr); + EXPECT_EQ(disposal_for(0)->buffer, dev0_buffer); + EXPECT_EQ(disposal_for(0)->event, dev0_event); + EXPECT_EQ(disposal_for(1)->buffer, dev1_buffer); + EXPECT_EQ(disposal_for(1)->event, nullptr); + + // The slot the reset cleared is the one a later lookup finds, so the reads below + // are of the entry the backend would go on using. One address cannot tell a + // surviving entry from a recycled allocation, so the invariant itself is pinned + // by ResetLeavesEveryEntryWhereItWas rather than here. + EXPECT_EQ(&pool.get(0), &dev0); + EXPECT_EQ(dev0.buffer, nullptr); + EXPECT_EQ(dev0.capacity, 0u); + EXPECT_EQ(dev0.marker.event, nullptr); + EXPECT_FALSE(dev0.marker.pending); + + // A smaller request than the freed buffer served: reuse would satisfy it from + // the stale capacity and allocate nothing, so this is what distinguishes a + // cleared slot from one the reset only emptied of its event. + void* const fresh = call(dev0, zero, 1024); + EXPECT_NE(fresh, nullptr); + EXPECT_EQ(dev0.capacity, 1024u); + EXPECT_EQ(zero.alloc_count(), 2) << "the slot was not cleared, so the request reused the freed buffer"; + EXPECT_TRUE(zero.retirements.empty()) << "the cleared slot retired a buffer the reset had already handed back"; +} + +// The mirror of device 1 in the case above: an event and no buffer. That state is +// reachable, because the handoff event is created before the allocation, so a +// claim whose allocation fails leaves one behind. A reset that skipped slots with +// no buffer would leak that event and pass every other case in this file. +TEST(SharedScratchPoolRegistry, ResetHandsBackTheEventOfASlotThatNeverGotABuffer) { + SharedScratchPool pool; + FakeAllocator a; + FakeEventFactory events; + + SharedScratchDevice& dev = pool.get(3); + const cudaEvent_t event = shared_scratch_claim_event(dev, std::ref(events)).event; + ASSERT_NE(event, nullptr); + a.fail_next = true; + ASSERT_EQ(call(dev, a, 1024), nullptr); + ASSERT_EQ(dev.buffer, nullptr) << "the allocation did not fail, so this slot is not the state under test"; + + std::vector> disposed; + reset_shared_scratch_pool_slots( + pool, [&](int, void* buffer, cudaEvent_t slot_event) { disposed.emplace_back(buffer, slot_event); }); + + ASSERT_EQ(disposed.size(), 1u) << "the reset skipped a slot holding an event and no buffer, so nothing ever destroys " + "that event"; + EXPECT_EQ(disposed[0].first, nullptr); + EXPECT_EQ(disposed[0].second, event); + EXPECT_EQ(dev.marker.event, nullptr) << "the slot kept the event the reset handed to the disposer"; +} + +// The registry hands out a reference and drops its lock, so callers go on using +// that reference; the reset has to empty a slot without moving it. Checking one +// address cannot pin that. An erasing reset returns the node to the allocator, +// which hands the very same address back to the next lookup, so the comparison +// still holds -- and every read through the old reference between the two is of +// freed memory. Enough entries and the allocator cannot reproduce them all. +TEST(SharedScratchPoolRegistry, ResetLeavesEveryEntryWhereItWas) { + constexpr int kDevices = 512; + SharedScratchPool pool; + FakeAllocator alloc; + FakeEventFactory events; + + std::vector before; + before.reserve(kDevices); + for (int id = 0; id < kDevices; ++id) { + SharedScratchDevice& dev = pool.get(id); + // Give each slot something, so the reset has work to do on all of them rather + // than skipping past empty ones. + ASSERT_NE(call(dev, alloc, 1024), nullptr); + ASSERT_NE(shared_scratch_claim_event(dev, std::ref(events)).event, nullptr); + before.push_back(&dev); + } + + int disposed = 0; + reset_shared_scratch_pool_slots(pool, [&disposed](int, void*, cudaEvent_t) { ++disposed; }); + ASSERT_EQ(disposed, kDevices); + + int moved = 0; + for (int id = 0; id < kDevices; ++id) { + if (&pool.get(id) != before[static_cast(id)]) { + ++moved; + } + } + EXPECT_EQ(moved, 0) << moved << " of " << kDevices + << " entries moved, so a reference the registry handed out before the reset names freed memory " + "or another device's slot"; +} + +// The reset's disposer frees device memory, and a device-wide free waits on +// everything queued on that device -- a parked host function included. Under +// either lock that would hold up whatever the lock covers: the device's own next +// claimant under the device lock, and every device's claimants under the +// registry's. So each slot is emptied under the locks and what came out of it is +// disposed of with neither held. +TEST(SharedScratchPoolRegistry, ResetDisposesWithNoLockHeld) { + constexpr int kUntouchedDevice = 99; + SharedScratchPool pool; + FakeAllocator alloc; + + SharedScratchDevice& dev0 = pool.get(0); + ASSERT_NE(call(dev0, alloc, 1024), nullptr); + + std::atomic disposing{false}; + std::atomic reset_returned{false}; + std::atomic nothing_disposed{false}; + std::atomic claimed{false}; + std::atomic device_lock_free{false}; + std::thread claimer([&] { + // The regression this case exists to catch is a reset that disposes of + // nothing, and this wait is where that regression arrives. Waiting for + // `disposing` alone would turn it into a target that hangs until Bazel kills + // it, with no assertion to say what broke. + const auto rendezvous_deadline = std::chrono::steady_clock::now() + std::chrono::seconds(10); + while (!disposing.load() && !reset_returned.load() && std::chrono::steady_clock::now() < rendezvous_deadline) { + std::this_thread::yield(); + } + if (!disposing.load()) { + nothing_disposed.store(true); + return; + } + if (dev0.mu.try_lock()) { + device_lock_free.store(true); + dev0.mu.unlock(); + } + pool.get(kUntouchedDevice); + claimed.store(true); + }); + + // Read inside the disposer: the claim completes once the reset returns either + // way, so only what was true while the disposer ran tells the two apart. + bool claimed_during_dispose = false; + reset_shared_scratch_pool_slots(pool, [&](int, void*, cudaEvent_t) { + disposing.store(true); + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (!claimed.load() && std::chrono::steady_clock::now() < deadline) { + std::this_thread::yield(); + } + claimed_during_dispose = claimed.load(); + }); + reset_returned.store(true); + claimer.join(); + + ASSERT_FALSE(nothing_disposed.load()) << "the reset returned without disposing of the slot that holds a buffer, so " + "there was no window in which to observe either lock"; + EXPECT_TRUE(device_lock_free.load()) << "the disposer ran holding the slot's own device lock, so a claim on that " + "device waits for a free that waits on the whole device"; + EXPECT_TRUE(claimed_during_dispose) << "a lookup for a device the reset never touched could not complete while the " + "disposer ran, so the registry's lock was held across it and one device's " + "teardown blocks every other device"; +} + +// A device lock still held between cases is a claim that was never released -- +// the defect the backend suite's own case hunts. Waiting for it would hang the +// fixture that resets before and after every case, so the run would end in a +// target timeout and the message naming the leak would never be printed. The +// reset reports the slot and leaves it instead. +// +// The holder below releases on its own deadline, so a reset that waits for the +// lock rather than reporting it ends this case in a failure rather than a hang. +TEST(SharedScratchPoolRegistry, ResetReportsALockedSlotRatherThanWaitingForIt) { + constexpr std::chrono::seconds kHolderDeadline{30}; + + SharedScratchPool pool; + FakeAllocator alloc; + + // Two unlocked slots against one locked one, so the number the reset answers + // with -- devices it could not take -- differs from the number it cleared. With + // one of each the two are both 1 and a reset that reported the wrong one of them + // would pass. + SharedScratchDevice& leaked = pool.get(0); + SharedScratchDevice& ordinary = pool.get(1); + SharedScratchDevice& also_ordinary = pool.get(2); + ASSERT_NE(call(leaked, alloc, 1024), nullptr); + ASSERT_NE(call(ordinary, alloc, 2048), nullptr); + ASSERT_NE(call(also_ordinary, alloc, 4096), nullptr); + + std::atomic holding{false}; + std::atomic release_it{false}; + std::thread holder([&] { + leaked.mu.lock(); + holding.store(true); + const auto deadline = std::chrono::steady_clock::now() + kHolderDeadline; + while (!release_it.load() && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + leaked.mu.unlock(); + }); + while (!holding.load()) { + std::this_thread::yield(); + } + + int disposed = 0; + const auto started = std::chrono::steady_clock::now(); + const std::size_t still_locked = reset_shared_scratch_pool_slots(pool, [&](int, void*, cudaEvent_t) { ++disposed; }); + const auto took = std::chrono::steady_clock::now() - started; + release_it.store(true); + holder.join(); + + EXPECT_EQ(still_locked, 1u) << "the reset did not report the one device whose lock it could not take"; + EXPECT_LT(took, kHolderDeadline) << "the reset returned only once the lock was released, so it waited for a claim " + "that a leak would never release"; + EXPECT_EQ(disposed, 2) << "the reset disposed of " << disposed + << " slots: it should hand back the two unlocked ones and leave the locked one alone"; + EXPECT_NE(leaked.buffer, nullptr) << "the reset emptied a slot whose lock it never held, so it handed the disposer a " + "buffer a live claimant is still using"; + EXPECT_EQ(ordinary.buffer, nullptr) << "one locked slot stopped the reset clearing the others"; + EXPECT_EQ(also_ordinary.buffer, nullptr) << "one locked slot stopped the reset clearing the others"; +} + +// The budget covers the whole reset, so a leaked lock has to be prevented from +// spending it on behalf of the devices it shares the pool with. A device that is +// merely busy -- claims arriving and finishing, which is every device the pool is +// for -- must still be taken, and a single try_lock made once the leak has run the +// clock down only takes it if it happens to be idle at that one instant. +// +// So the neighbour below is locked when the reset starts, free for half a second +// in the middle, and locked again long before the deadline. A reset that retries +// takes it inside that window; one that spends its budget on the leak first and +// then tries once reports it as locked and never hands its buffer to the disposer +// -- which, in the fixture that resets between cases, means the next case starts +// on a device still holding what an earlier one grew. +// +// Which slot is which is decided by the map rather than by this case, and only a +// leak the reset reaches *before* the busy slot can cost it anything. So the order +// is read off a reset of the same three entries first -- entries are never erased, +// so the second walk takes them in the same order -- and the roles are handed out +// from that. +TEST(SharedScratchPoolRegistry, ALeakedLockDoesNotCostTheResetTheSlotsItSharesThePoolWith) { + constexpr std::chrono::milliseconds kNeighbourFreeFrom{200}; + constexpr std::chrono::milliseconds kNeighbourBusyAgainFrom{700}; + const auto holder_deadline = SharedScratchPool::kResetLockWait + std::chrono::seconds{25}; + + SharedScratchPool pool; + FakeAllocator alloc; + // Three slots, so the number the reset answers with and the number it cleared + // stay different figures: one leaked, one busy, one free throughout. + for (const int device_id : {0, 1, 2}) { + ASSERT_NE(call(pool.get(device_id), alloc, 1024), nullptr); + } + std::vector walk_order; + ASSERT_EQ(reset_shared_scratch_pool_slots(pool, [&](int id, void*, cudaEvent_t) { walk_order.push_back(id); }), 0u) + << "a reset of three unlocked slots reported one it could not take"; + ASSERT_EQ(walk_order.size(), 3u); + + SharedScratchDevice& leaked = pool.get(walk_order[0]); + SharedScratchDevice& busy = pool.get(walk_order[1]); + SharedScratchDevice& ordinary = pool.get(walk_order[2]); + ASSERT_NE(call(leaked, alloc, 1024), nullptr); + ASSERT_NE(call(busy, alloc, 2048), nullptr); + ASSERT_NE(call(ordinary, alloc, 4096), nullptr); + + std::atomic both_held{false}; + std::atomic release_it{false}; + const auto started = std::chrono::steady_clock::now(); + std::thread holder([&] { + leaked.mu.lock(); + busy.mu.lock(); + both_held.store(true); + std::this_thread::sleep_until(started + kNeighbourFreeFrom); + busy.mu.unlock(); + std::this_thread::sleep_until(started + kNeighbourBusyAgainFrom); + // Locked again well inside the deadline, so only a reset that looked during + // the window above ever had it. + busy.mu.lock(); + const auto deadline = std::chrono::steady_clock::now() + holder_deadline; + while (!release_it.load() && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + busy.mu.unlock(); + leaked.mu.unlock(); + }); + while (!both_held.load()) { + std::this_thread::yield(); + } + + int disposed = 0; + const std::size_t still_locked = reset_shared_scratch_pool_slots(pool, [&](int, void*, cudaEvent_t) { ++disposed; }); + release_it.store(true); + holder.join(); + + EXPECT_EQ(still_locked, 1u) + << "the reset reported " << still_locked + << " locked slots: only the leaked one is unreachable, and the busy one was free for " + << (kNeighbourBusyAgainFrom - kNeighbourFreeFrom).count() << " ms of the " + << std::chrono::duration_cast(SharedScratchPool::kResetLockWait).count() + << " ms it had to look"; + EXPECT_EQ(disposed, 2) << "the reset disposed of " << disposed + << " slots: the leaked one it must leave alone, and the busy one it must come back to rather " + "than give up on because another device's lock spent the budget"; + EXPECT_EQ(busy.buffer, nullptr) << "one leaked lock left a busy device's buffer resident, so the pool a case is " + "reset for still holds what an earlier one grew"; + EXPECT_EQ(ordinary.buffer, nullptr) << "the reset left a slot nothing ever locked"; + EXPECT_NE(leaked.buffer, nullptr) << "the reset emptied a slot whose lock it never held"; +} + +TEST(SharedScratchPoolRegistry, AGrowthOnOneDeviceDoesNotBlockAClaimOnAnother) { + SharedScratchPool pool; + // One allocator per thread: the two claims share the registry and nothing else. + 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); + RetiredScratch retired; + shared_scratch_get_or_grow( + dev0, + 4096, + [&](std::size_t bytes) { + entered_alloc.set_value(); + leave.wait(); + return zero.alloc(bytes); + }, + retired); + }); + // The cap matters as much as the wait: a growth that takes the reuse path never + // reaches its allocation, so nothing fires this promise and an uncapped wait + // 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); + RetiredScratch retired; + return shared_scratch_get_or_grow( + dev1, 2048, [&](std::size_t bytes) { return one.alloc(bytes); }, retired); + }); + const bool served = claim.wait_for(std::chrono::seconds(10)) == std::future_status::ready; + + 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; + } +} + +#if defined(__unix__) + +// scratch_pool() hands back a registry that is allocated once and never +// destroyed. The alternative -- an object with static storage duration -- is +// destroyed after main returns, which destroys the registry's mutex, every +// device's mutex and the map nodes a live reference points into, while a thread +// anywhere between the lookup and the release of its device lock still holds +// them. +// +// Nothing gtest asserts inside this process can see that: it happens after the +// last test has finished. So these run in forked children that put a thread +// inside the pooled path and then exit, and read the status the child died with. +// +// Probabilistic in one direction only. A child that is torn down under a live +// claimant does not have to crash, so the count below is a floor on how often it +// would; a child that is not torn down under one cannot crash at all, so a +// single non-zero status is a real failure and not a flake. +constexpr int kTeardownChildren = 24; + +// What a child exits with when it could not get as far as the state under test -- +// starting a thread, two dozen forks in, is the way that happens. Distinct from +// zero and from the teardown crash this case counts, so a run that could not set +// itself up is not read as evidence either way. +constexpr int kChildCouldNotStart = 121; + +// Long enough that the children below -- a thread, a 30 ms sleep and exit() -- +// are nowhere near it. +constexpr std::chrono::seconds kChildDeadline{30}; + +// Reaps `child`, killing it if it overruns `kChildDeadline`; returns false when +// it had to. Without a deadline a wedged child parks the parent here until the +// whole target times out, which reports as a timeout on the binary rather than as +// this case failing. The deadline only delivers that if the whole reaping loop +// fits inside the target's timeout -- kTeardownChildren children at kChildDeadline +// each is 720 s -- which is why the BUILD file gives this target the long (900 s) +// timeout rather than leaving it Bazel's 300 s default. An interrupted wait is +// retried, since a signal can arrive at any point and says nothing about the +// child. +bool reap_child(pid_t child, int& status) { + const auto deadline = std::chrono::steady_clock::now() + kChildDeadline; + for (;;) { + const pid_t reaped = waitpid(child, &status, WNOHANG); + if (reaped == child) { + return true; + } + if (reaped == -1 && errno != EINTR) { + return false; + } + if (std::chrono::steady_clock::now() >= deadline) { + kill(child, SIGKILL); + // Reaped even so. A zombie is not inherited by anything -- it stays this + // process's child until this process reaps it or exits -- so what leaving it + // costs is a process-table entry per overrun, held for the rest of the run. + while (waitpid(child, &status, 0) == -1 && errno == EINTR) { + } + return false; + } + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } +} + +[[noreturn]] void hold_the_pool_open_then_exit() { + static std::atomic inside{false}; + std::thread claimant([] { + for (unsigned i = 0;; ++i) { + SharedScratchDevice& dev = scratch_pool().get(static_cast(i % 8)); + std::lock_guard lk(dev.mu); + dev.capacity += 1; + inside.store(true); + } + }); + // Never joined: the point is for it to still be in there at teardown. + claimant.detach(); + while (!inside.load()) { + std::this_thread::yield(); + } + std::this_thread::sleep_for(std::chrono::milliseconds(30)); + // exit() and not _exit(): static destruction is the event under test. + std::exit(0); +} + +TEST(SharedScratchPoolRegistry, TheProcessPoolOutlivesStaticDestructionUnderALiveClaimant) { + int killed_by_signal = 0; + int exited_nonzero = 0; + int could_not_start = 0; + int overran = 0; + for (int i = 0; i < kTeardownChildren; ++i) { + // Anything gtest has buffered would otherwise be written twice, once by each + // side of the fork. + std::fflush(nullptr); + const pid_t child = fork(); + ASSERT_NE(child, -1) << "fork failed: " << std::strerror(errno); + if (child == 0) { + // Nothing may leave the child by any route but exit. Left to unwind, a + // failed thread construction returns into the test body and the child + // carries on through the rest of the binary as a second gtest process, + // forking children of its own and writing over the parent's output. + try { + hold_the_pool_open_then_exit(); + } catch (...) { + _exit(kChildCouldNotStart); + } + } + int status = 0; + if (!reap_child(child, status)) { + ++overran; + continue; + } + if (WIFSIGNALED(status)) { + ++killed_by_signal; + } else if (WEXITSTATUS(status) == kChildCouldNotStart) { + ++could_not_start; + } else if (WEXITSTATUS(status) != 0) { + ++exited_nonzero; + } + } + + EXPECT_EQ(overran, 0) << overran << " of " << kTeardownChildren << " children were still running after " + << kChildDeadline.count() + << " seconds and were killed, so they neither reached teardown nor reported anything about it"; + EXPECT_EQ(could_not_start, 0) << could_not_start << " of " << kTeardownChildren + << " children could not start their claimant thread, so those runs never put anything " + "inside the pool and say nothing about what teardown does to a live claimant"; + EXPECT_EQ(killed_by_signal, 0) << killed_by_signal << " of " << kTeardownChildren + << " children died on a signal at exit with a thread still inside the pool, so the " + "registry is being destroyed out from under a live claimant"; + EXPECT_EQ(exited_nonzero, 0) << exited_nonzero << " of " << kTeardownChildren + << " children exited non-zero at teardown with a thread still inside the pool"; +} + +#endif // defined(__unix__) + +} // namespace +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/tests/cpp/executorch/test_shared_scratch_pool_other_tu.cpp b/tests/cpp/executorch/test_shared_scratch_pool_other_tu.cpp new file mode 100644 index 0000000000..18cebba619 --- /dev/null +++ b/tests/cpp/executorch/test_shared_scratch_pool_other_tu.cpp @@ -0,0 +1,35 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// A second translation unit, so the pool's test target can see whether +// scratch_pool() hands the same registry to every one of them. +// +// It has to be a separate file: the property is a linkage one, and a single-file +// target cannot observe it. scratch_pool() is inline for exactly this reason -- +// the backend and the pool's test hooks are different translation units and both +// have to reach the one registry -- and with internal linkage instead each would +// get a registry of its own, the reset hook would clear one nobody uses and the +// capacity hook would always answer zero. +// +// The backend suite does not let that through: six of its assertions across five +// cases read the capacity back and expect a figure above zero, and every one of +// them would fail. What it cannot do is say why, since those cases are named for +// the empty input, the capture refusal and the three growths rather than for +// linkage. That is what this file buys -- one named case in the pool's own target, +// which needs no GPU. + +#include "torch_tensorrt/executorch/SharedScratchPool.h" + +namespace torch_tensorrt { +namespace executorch_backend { + +const SharedScratchPool* scratch_pool_seen_from_another_translation_unit() { + return &scratch_pool(); +} + +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/third_party/cuda/BUILD b/third_party/cuda/BUILD index 204b9cee23..49869f4f0f 100644 --- a/third_party/cuda/BUILD +++ b/third_party/cuda/BUILD @@ -17,6 +17,17 @@ config_setting( ], ) +# Every header in the toolkit, and the only place the glob is written: the library +# targets below take their headers from here. One catch-all pattern rather than a +# list of extensions, because `include/**/*` already matches every .h, .hpp and +# .inl under it. Keeping the headers in a target of their own is the shape +# third_party/cublas and third_party/cudnn/archive already use. +cc_library( + name = "cuda_headers", + hdrs = glob(["include/**/*"]), + includes = ["include/"], +) + cc_library( name = "cudart", srcs = select({ @@ -30,13 +41,7 @@ cc_library( "lib64/libcudart.so", ], }), - hdrs = glob([ - "include/**/*.h", - "include/**/*.hpp", - "include/**/*.inl", - "include/**/*", - ]), - includes = ["include/"], + deps = [":cuda_headers"], ) cc_library( @@ -71,14 +76,8 @@ cc_library( allow_empty = True, ), }), - hdrs = glob([ - "include/**/*.h", - "include/**/*.hpp", - "include/**/*.inl", - "include/**/*", - ]), - includes = ["include/"], linkopts = ["-Wl,-rpath,lib/"], + deps = [":cuda_headers"], ) cc_library( @@ -97,12 +96,6 @@ cc_library( allow_empty = True, ), }), - hdrs = glob([ - "include/**/*cublas*.h", - "include/**/*.hpp", - "include/**/*.inl", - "include/**/*", - ]), - includes = ["include/"], linkopts = ["-Wl,-rpath,lib/"], + deps = [":cuda_headers"], )