perf(executorch): opt-in shared per-device activation-scratch pool

.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

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",