[XLA:GPU] Defer CUDA VMM memory deallocations while a cuDNN handle is active.

cuDNN v9 internally captures temporary CUDA graphs during heuristic queries and execution plan creation (e.g. via `cudnn::backend::cudnnBackendGetMetadataFromGraph` inside `cudnnBackendFinalize`), which does not go through `CudaStream::CaptureHandle`.

During parallel autotuning (`ConfigAssigner::GetConfigsForAll`), a worker thread running `GpuProfiler::Profile` can deallocate VMM buffers via `CudaDeviceMemoryAllocation::~CudaDeviceMemoryAllocation` -> `DrainPendingVmmDeallocations` -> `cuCtxSynchronize()` concurrently with the main thread querying cuDNN heuristics in `CudnnSupport::GetConvolveRunners`. Calling `cuCtxSynchronize()` while cuDNN is capturing a CUDA graph on the same CUDA context either fails with `CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED` / `cudaErrorStreamCaptureInvalidated` or segfaults inside `libcuda.so` (`cuCtxSynchronize`).

Call `CudaDeviceAllocator::EnterStreamCapture` and `CudaDeviceAllocator::ExitStreamCapture` in `CudnnHandle` (`xla/stream_executor/cuda/cuda_dnn.cc`) so that VMM deallocations are deferred until the active cuDNN operation completes.

PiperOrigin-RevId: 1003883162

third_party/xla/xla/stream_executor/cuda/BUILD

@@ -563,6 +563,7 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         ":cuda_compute_capability",
+        ":cuda_device_allocator",
         ":cuda_diagnostics",
         ":cuda_platform_id",
         ":cuda_status",
@@ -2145,6 +2146,9 @@ xla_test(
         "//xla/stream_executor/gpu:gpu_test_kernels",
         "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/base",
+        "@com_google_absl//absl/log:log_entry",
+        "@com_google_absl//absl/log:log_sink",
+        "@com_google_absl//absl/log:log_sink_registry",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",  # build_cleaner: keep
         "@com_google_absl//absl/strings",

third_party/xla/xla/stream_executor/cuda/cuda_command_buffer_test.cc

@@ -16,18 +16,24 @@ limitations under the License.
 #include "xla/stream_executor/cuda/cuda_command_buffer.h"
 
 #include <array>
+#include <atomic>
 #include <cstddef>
 #include <cstdint>
 #include <memory>
 #include <optional>
+#include <thread>  // NOLINT
 #include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/base/casts.h"
+#include "absl/log/log_entry.h"
+#include "absl/log/log_sink.h"
+#include "absl/log/log_sink_registry.h"
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"
 #include "absl/strings/ascii.h"
+#include "absl/strings/match.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "third_party/cudnn_frontend/include/cudnn_frontend.h"  // IWYU pragma: keep - cudnn frontend headers are not hermetic
@@ -508,5 +514,77 @@ TEST_F(CudaCommandBufferTest,
   ASSERT_NE(cuda_cmd_buffer, nullptr);
   EXPECT_TRUE(cuda_cmd_buffer->has_node_priority());
 }
+
+TEST_F(CudaCommandBufferTest,
+       ConcurrentVmmDeallocationDuringGetConvolveRunners) {
+  Platform* platform = CudaPlatform();
+  ASSERT_OK_AND_ASSIGN(StreamExecutor * executor,
+                       platform->ExecutorForDevice(0));
+  dnn::DnnSupport* dnn = executor->AsDnn();
+  ASSERT_NE(dnn, nullptr);
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<Stream> stream,
+                       executor->CreateStream());
+
+  class CaptureErrorSink : public absl::LogSink {
+   public:
+    void Send(const absl::LogEntry& entry) override {
+      if (absl::StrContains(entry.text_message(),
+                            "CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED")) {
+        capture_errors.fetch_add(1, std::memory_order_relaxed);
+      }
+    }
+    std::atomic<int> capture_errors{0};
+  } sink;
+  absl::AddLogSink(&sink);
+
+  std::atomic<bool> stop{false};
+  std::thread dealloc_thread([&]() {
+    while (!stop.load(std::memory_order_relaxed)) {
+      DeviceAddressBase buf = executor->Allocate(4096);
+      ASSERT_FALSE(buf.is_null());
+      executor->Deallocate(&buf);
+    }
+  });
+
+  for (int iter = 0; iter < 5; ++iter) {
+    dnn::BatchDescriptor input_desc;
+    input_desc.set_count(2)
+        .set_feature_map_count(3)
+        .set_height(9 + iter)
+        .set_width(10 + iter)
+        .set_layout(dnn::DataLayout::kBatchDepthYX);
+    dnn::FilterDescriptor filter_desc;
+    filter_desc.set_input_feature_map_count(3)
+        .set_output_feature_map_count(4)
+        .set_input_filter_height(3)
+        .set_input_filter_width(3)
+        .set_layout(dnn::FilterLayout::kOutputInputYX);
+    dnn::ConvolutionDescriptor conv_desc;
+    conv_desc.set_vertical_filter_stride(1)
+        .set_horizontal_filter_stride(1)
+        .set_zero_padding_height(1)
+        .set_zero_padding_width(1);
+    dnn::BatchDescriptor output_desc;
+    output_desc.set_count(2)
+        .set_feature_map_count(4)
+        .set_height(9 + iter)
+        .set_width(10 + iter)
+        .set_layout(dnn::DataLayout::kBatchDepthYX);
+
+    std::vector<std::unique_ptr<const dnn::ConvRunner>> runners;
+    EXPECT_OK(dnn->GetConvolveRunners(
+        dnn::ConvolutionKind::FORWARD, dnn::DataType::kHalf,
+        dnn::DataType::kHalf, stream.get(), input_desc,
+        DeviceAddressBase(nullptr, 0), filter_desc,
+        DeviceAddressBase(nullptr, 0), output_desc,
+        DeviceAddressBase(nullptr, 0), conv_desc, /*use_fallback=*/false,
+        /*scratch_allocator=*/nullptr, EngineOptions{}, &runners));
+  }
+
+  stop.store(true, std::memory_order_relaxed);
+  dealloc_thread.join();
+  absl::RemoveLogSink(&sink);
+  EXPECT_EQ(sink.capture_errors.load(), 0);
+}
 }  // namespace
 }  // namespace stream_executor::cuda

third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc

@@ -57,6 +57,7 @@ limitations under the License.
 #include "xla/status_macros.h"
 #include "xla/stream_executor/activate_context.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
+#include "xla/stream_executor/cuda/cuda_device_allocator.h"
 #include "xla/stream_executor/cuda/cuda_diagnostics.h"
 #include "xla/stream_executor/cuda/cuda_platform_id.h"
 #include "xla/stream_executor/cuda/cuda_status.h"
@@ -213,15 +214,39 @@ class CudnnHandle {
   // Takes ownership of the lock to access cuDNN using handle.
   CudnnHandle(StreamExecutor* executor, std::unique_ptr<absl::MutexLock> lock,
               cudnnHandle_t handle)
-      : context_(executor->Activate()),
+      : executor_(executor),
+        context_(executor->Activate()),
         lock_(std::move(lock)),
-        handle_(handle) {}
+        handle_(handle) {
+    // cuDNN v9 internally captures temporary CUDA graphs during heuristic
+    // queries and execution plan creation (e.g. via
+    // cudnnBackendGetMetadataFromGraph), which does not go through
+    // CudaStream::CaptureHandle. Defer VMM deallocations (which call
+    // cuCtxSynchronize()) while a cuDNN handle is active on this executor.
+    CudaDeviceAllocator::EnterStreamCapture(executor_);
+  }
+
+  CudnnHandle(CudnnHandle&& other) noexcept
+      : executor_(std::exchange(other.executor_, nullptr)),
+        context_(std::move(other.context_)),
+        lock_(std::move(other.lock_)),
+        handle_(other.handle_) {}
+
+  CudnnHandle& operator=(CudnnHandle&&) = delete;
+  CudnnHandle(const CudnnHandle&) = delete;
+  CudnnHandle& operator=(const CudnnHandle&) = delete;
+
+  ~CudnnHandle() {
+    lock_.reset();
+    CudaDeviceAllocator::ExitStreamCapture(executor_);
+  }
 
   // Returns cuDNN handle. To be passed directly to cuDNN APIs, don't keep
   // a copy.
   cudnnHandle_t handle() const { return handle_; }
 
  private:
+  StreamExecutor* executor_;
   std::unique_ptr<ActivateContext> context_;
   std::unique_ptr<absl::MutexLock> lock_;
   cudnnHandle_t handle_;  // Not owned.