Create DNN handles before command buffers capture convolutions

xla/backends/gpu/runtime/BUILD

@@ -793,8 +793,14 @@ xla_cc_test(
         ":thunk",
         ":thunk_proto_cc",
         "//xla/service:buffer_assignment",
+        "//xla/stream_executor:mock_stream_executor",
         "//xla/tsl/util/proto:parse_text_proto",
         "//xla/tsl/util/proto:proto_matchers",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest_main",
     ],
 )

xla/backends/gpu/runtime/convolution_thunk.cc

@@ -94,6 +94,16 @@ absl::StatusOr<RunConvOptions> ConvolutionThunk::GetOrCreate(
   return RunConvOptions{nullptr, &**runner_states_.Find(device_ordinal)};
 }
 
+absl::Status ConvolutionThunk::Initialize(const InitializeParams& params) {
+  // The first ExecuteOnStream may run inside a command buffer's stream
+  // capture, and creating the DNN handle there fails on ROCm (it allocates
+  // device memory).
+  if (!params.executor->AsDnn()) {
+    return absl::InternalError("Failed to initialize DNN support");
+  }
+  return absl::OkStatus();
+}
+
 absl::Status ConvolutionThunk::ExecuteOnStream(const ExecuteParams& params) {
   VLOG(5) << "ConvolutionThunk";
   const auto& buffer_allocations = *params.buffer_allocations;

xla/backends/gpu/runtime/convolution_thunk.h

@@ -55,6 +55,7 @@ class ConvolutionThunk : public TracedCommand {
   ConvolutionThunk(const ConvolutionThunk&) = delete;
   ConvolutionThunk& operator=(const ConvolutionThunk&) = delete;
 
+  absl::Status Initialize(const InitializeParams& params) override;
   absl::Status ExecuteOnStream(const ExecuteParams& params) override;
 
   BufferUses buffer_uses() const override;

xla/backends/gpu/runtime/convolution_thunk_test.cc

@@ -15,26 +15,33 @@ limitations under the License.
 
 #include "xla/backends/gpu/runtime/convolution_thunk.h"
 
-#include <gmock/gmock.h>
-#include <gtest/gtest.h>
-
 #include <memory>
 #include <vector>
 
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/status/status_macros.h"
+#include "absl/status/status_matchers.h"
+#include "absl/status/statusor.h"
+#include "absl/types/span.h"
 #include "xla/backends/gpu/runtime/thunk.h"
 #include "xla/backends/gpu/runtime/thunk.pb.h"
 #include "xla/service/buffer_assignment.h"
+#include "xla/stream_executor/mock_stream_executor.h"
 #include "xla/tsl/util/proto/parse_text_proto.h"
 #include "xla/tsl/util/proto/proto_matchers.h"
 
 namespace xla::gpu {
 namespace {
 
+using ::absl_testing::StatusIs;
+using ::testing::Return;
 using ::tsl::proto_testing::EqualsProto;
 using ::tsl::proto_testing::ParseTextProtoOrDie;
 
-TEST(ConvolutionThunkTest, ProtoRoundTrip) {
-  auto proto = ParseTextProtoOrDie<ThunkProto>(R"pb(
+ThunkProto ConvolutionThunkProto() {
+  return ParseTextProtoOrDie<ThunkProto>(R"pb(
     thunk_info { profile_annotation: "conv_thunk_profile" }
     convolution_thunk {
       conv_descriptor {
@@ -144,23 +151,48 @@ TEST(ConvolutionThunkTest, ProtoRoundTrip) {
       scratch_buffer { offset: 0 size: 1024 buffer_allocation_index: 3 }
     }
   )pb");
+}
 
+std::vector<BufferAllocation> CreateBufferAllocations() {
   std::vector<BufferAllocation> buffer_allocations;
   buffer_allocations.emplace_back(/*index=*/0, /*size=*/4, /*color=*/0);
   buffer_allocations.emplace_back(/*index=*/1, /*size=*/4, /*color=*/0);
   buffer_allocations.emplace_back(/*index=*/2, /*size=*/4, /*color=*/0);
   buffer_allocations.emplace_back(/*index=*/3, /*size=*/1024, /*color=*/0);
+  return buffer_allocations;
+}
 
-  ASSERT_OK_AND_ASSIGN(Thunk::ThunkInfo thunk_info,
-                       Thunk::ThunkInfo::FromProto(proto.thunk_info()));
-  ASSERT_OK_AND_ASSIGN(
-      std::unique_ptr<ConvolutionThunk> thunk,
-      ConvolutionThunk::FromProto(thunk_info, proto.convolution_thunk(),
-                                  buffer_allocations,
-                                  /*devices_per_host=*/1));
+absl::StatusOr<std::unique_ptr<ConvolutionThunk>> CreateThunk(
+    const ThunkProto& proto,
+    absl::Span<const BufferAllocation> buffer_allocations) {
+  ABSL_ASSIGN_OR_RETURN(Thunk::ThunkInfo thunk_info,
+                        Thunk::ThunkInfo::FromProto(proto.thunk_info()));
+  return ConvolutionThunk::FromProto(thunk_info, proto.convolution_thunk(),
+                                     buffer_allocations,
+                                     /*devices_per_host=*/1);
+}
+
+TEST(ConvolutionThunkTest, ProtoRoundTrip) {
+  ThunkProto proto = ConvolutionThunkProto();
+  std::vector<BufferAllocation> buffer_allocations = CreateBufferAllocations();
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<ConvolutionThunk> thunk,
+                       CreateThunk(proto, buffer_allocations));
   ASSERT_OK_AND_ASSIGN(ThunkProto round_trip_proto, thunk->ToProto());
   EXPECT_THAT(round_trip_proto, EqualsProto(proto));
 }
 
+TEST(ConvolutionThunkTest, InitializeRequestsDnnSupport) {
+  std::vector<BufferAllocation> buffer_allocations = CreateBufferAllocations();
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<ConvolutionThunk> thunk,
+      CreateThunk(ConvolutionThunkProto(), buffer_allocations));
+  se::MockStreamExecutor executor;
+  EXPECT_CALL(executor, AsDnn()).WillOnce(Return(nullptr));
+
+  Thunk::InitializeParams params;
+  params.executor = &executor;
+  EXPECT_THAT(thunk->Initialize(params), StatusIs(absl::StatusCode::kInternal));
+}
+
 }  // namespace
 }  // namespace xla::gpu

xla/service/BUILD

@@ -1046,7 +1046,9 @@ cc_library(
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
-    ],
+    ] + if_rocm([
+        "//xla/stream_executor/rocm:rocm_capture_invalidation",
+    ]),
 )
 
 xla_cc_test(
@@ -1065,6 +1067,29 @@ xla_cc_test(
     ],
 )
 
+xla_test(
+    name = "platform_util_multigpu_test",
+    srcs = ["platform_util_multigpu_test.cc"],
+    backend_tags = {
+        "gpu": ["multi_gpu"],
+    },
+    backends = ["gpu"],
+    tags = ["rocm-only"],
+    deps = [
+        ":platform_util",
+        "//xla/stream_executor:command_buffer",
+        "//xla/stream_executor:dnn",
+        "//xla/stream_executor:platform",
+        "//xla/stream_executor:stream",
+        "//xla/stream_executor:stream_executor_h",
+        "//xla/stream_executor:trace_command_buffer_factory",
+        "//xla/tsl/platform:env",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "backend",
     srcs = ["backend.cc"],

xla/service/platform_util.cc

@@ -50,6 +50,10 @@ limitations under the License.
 #include "xla/tsl/platform/threadpool.h"
 #include "xla/util.h"
 
+#if TENSORFLOW_USE_ROCM
+#include "xla/stream_executor/rocm/rocm_capture_invalidation.h"
+#endif  // TENSORFLOW_USE_ROCM
+
 namespace xla {
 
 namespace {
@@ -289,6 +293,12 @@ PlatformUtil::GetStreamExecutors(
     return Internal("no supported devices found for platform %s: %s",
                     platform->Name(), absl::StrJoin(device_errors, ", "));
   }
+
+#if TENSORFLOW_USE_ROCM
+  if (platform->id() == se::rocm::kROCmPlatformId) {
+    se::gpu::CreateDnnSupportUnlessCaptureInvalidationIsDeviceScoped(out);
+  }
+#endif  // TENSORFLOW_USE_ROCM
   return out;
 }
 

xla/service/platform_util.h

@@ -68,6 +68,10 @@ class PlatformUtil {
   // their StreamExecutors initialized, otherwise all StreamExecutors will be
   // initialized and returned.
   //
+  // On ROCm, with more than one device, this can also create every device's
+  // DNN support (MIOpen handle); see
+  // CreateDnnSupportUnlessCaptureInvalidationIsDeviceScoped.
+  //
   // If the platform has no visible devices, a not-found error is returned.
   static absl::StatusOr<std::vector<se::StreamExecutor*>> GetStreamExecutors(
       se::Platform* platform,

xla/service/platform_util_multigpu_test.cc

@@ -0,0 +1,70 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+#include <set>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "xla/service/platform_util.h"
+#include "xla/stream_executor/command_buffer.h"
+#include "xla/stream_executor/dnn.h"
+#include "xla/stream_executor/platform.h"
+#include "xla/stream_executor/stream.h"
+#include "xla/stream_executor/stream_executor.h"
+#include "xla/stream_executor/trace_command_buffer_factory.h"
+#include "xla/tsl/platform/env.h"
+
+namespace xla {
+namespace {
+
+using ::absl_testing::IsOk;
+
+// As when device 1's first convolution initializes while device 0 records a
+// command buffer. Only meaningful while nothing earlier in this process creates
+// device 1's DNN support.
+TEST(PlatformUtilMultiGpuTest,
+     CreatingDnnSupportLeavesAnotherDevicesCaptureIntact) {
+  ASSERT_OK_AND_ASSIGN(se::Platform * platform,
+                       PlatformUtil::GetPlatform("gpu"));
+  if (platform->VisibleDeviceCount() < 2) {
+    GTEST_SKIP() << "Test requires at least 2 GPUs";
+  }
+  ASSERT_OK_AND_ASSIGN(
+      std::vector<se::StreamExecutor*> executors,
+      PlatformUtil::GetStreamExecutors(platform, std::set<int>{0, 1}));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<se::Stream> stream,
+                       executors[0]->CreateStream());
+
+  se::dnn::DnnSupport* dnn = nullptr;
+  EXPECT_THAT(se::TraceCommandBufferFactory::Create(
+                  executors[0], stream.get(),
+                  [&](se::Stream*) {
+                    std::unique_ptr<tsl::Thread> thread(
+                        tsl::Env::Default()->StartThread(
+                            tsl::ThreadOptions(), "dnn_support",
+                            [&] { dnn = executors[1]->AsDnn(); }));
+                    return absl::OkStatus();
+                  },
+                  se::CommandBuffer::Mode::kPrimary),
+              IsOk());
+  EXPECT_NE(dnn, nullptr);
+}
+
+}  // namespace
+}  // namespace xla

xla/stream_executor/rocm/BUILD

@@ -1274,6 +1274,51 @@ xla_cc_test(
     ],
 )
 
+cc_library(
+    name = "rocm_capture_invalidation",
+    srcs = ["rocm_capture_invalidation.cc"],
+    hdrs = ["rocm_capture_invalidation.h"],
+    tags = [
+        "gpu",
+        "rocm-only",
+    ],
+    deps = [
+        ":rocm_status",
+        "//xla/stream_executor:stream_executor_h",
+        "//xla/tsl/platform:env",
+        "@com_google_absl//absl/base:no_destructor",
+        "@com_google_absl//absl/cleanup",
+        "@com_google_absl//absl/functional:function_ref",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/types:span",
+        "@local_config_rocm//rocm:hip",  # buildcleaner: keep
+        "@local_config_rocm//rocm:rocm_headers",
+    ],
+)
+
+xla_test(
+    name = "rocm_capture_invalidation_test",
+    srcs = ["rocm_capture_invalidation_test.cc"],
+    backends = ["gpu"],
+    tags = ["rocm-only"],
+    deps = [
+        ":rocm_capture_invalidation",
+        "//xla/stream_executor:mock_stream_executor",
+        "//xla/stream_executor:stream_executor_h",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_googletest//:gtest_main",
+        "@local_config_rocm//rocm:hip",  # buildcleaner: keep
+        "@local_config_rocm//rocm:rocm_headers",
+    ],
+)
+
 cc_library(
     name = "rocm_command_buffer",
     srcs = [