diff --git a/xla/backends/gpu/runtime/BUILD b/xla/backends/gpu/runtime/BUILD index abd87e3353b1ac..da6e97dc54e3c2 100644 --- a/xla/backends/gpu/runtime/BUILD +++ b/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", ], ) diff --git a/xla/backends/gpu/runtime/convolution_thunk.cc b/xla/backends/gpu/runtime/convolution_thunk.cc index 9f02754d2dcd3f..b7a7d551f19625 100644 --- a/xla/backends/gpu/runtime/convolution_thunk.cc +++ b/xla/backends/gpu/runtime/convolution_thunk.cc @@ -94,6 +94,16 @@ absl::StatusOr 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; diff --git a/xla/backends/gpu/runtime/convolution_thunk.h b/xla/backends/gpu/runtime/convolution_thunk.h index 0883992022c7c9..ff29c6d1a281d3 100644 --- a/xla/backends/gpu/runtime/convolution_thunk.h +++ b/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; diff --git a/xla/backends/gpu/runtime/convolution_thunk_test.cc b/xla/backends/gpu/runtime/convolution_thunk_test.cc index a81bed6a4327f4..a075153bd63e00 100644 --- a/xla/backends/gpu/runtime/convolution_thunk_test.cc +++ b/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 -#include - #include #include +#include +#include +#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(R"pb( +ThunkProto ConvolutionThunkProto() { + return ParseTextProtoOrDie(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 CreateBufferAllocations() { std::vector 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 thunk, - ConvolutionThunk::FromProto(thunk_info, proto.convolution_thunk(), - buffer_allocations, - /*devices_per_host=*/1)); +absl::StatusOr> CreateThunk( + const ThunkProto& proto, + absl::Span 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 buffer_allocations = CreateBufferAllocations(); + ASSERT_OK_AND_ASSIGN(std::unique_ptr 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 buffer_allocations = CreateBufferAllocations(); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr 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 diff --git a/xla/service/BUILD b/xla/service/BUILD index ca6e64d28ab413..58301d7c07668a 100644 --- a/xla/service/BUILD +++ b/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"], diff --git a/xla/service/platform_util.cc b/xla/service/platform_util.cc index b463a5ea70ff44..abeda02c55a46e 100644 --- a/xla/service/platform_util.cc +++ b/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; } diff --git a/xla/service/platform_util.h b/xla/service/platform_util.h index 373954bc6302b5..b1e2ee8c5bb07f 100644 --- a/xla/service/platform_util.h +++ b/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> GetStreamExecutors( se::Platform* platform, diff --git a/xla/service/platform_util_multigpu_test.cc b/xla/service/platform_util_multigpu_test.cc new file mode 100644 index 00000000000000..f3c801a3dd6029 --- /dev/null +++ b/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 +#include +#include + +#include +#include +#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 executors, + PlatformUtil::GetStreamExecutors(platform, std::set{0, 1})); + ASSERT_OK_AND_ASSIGN(std::unique_ptr stream, + executors[0]->CreateStream()); + + se::dnn::DnnSupport* dnn = nullptr; + EXPECT_THAT(se::TraceCommandBufferFactory::Create( + executors[0], stream.get(), + [&](se::Stream*) { + std::unique_ptr 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 diff --git a/xla/stream_executor/rocm/BUILD b/xla/stream_executor/rocm/BUILD index 5f7928d3e04c5c..2200049ce0a8a8 100644 --- a/xla/stream_executor/rocm/BUILD +++ b/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 = [ diff --git a/xla/stream_executor/rocm/rocm_capture_invalidation.cc b/xla/stream_executor/rocm/rocm_capture_invalidation.cc new file mode 100644 index 00000000000000..96ea4e63a25e30 --- /dev/null +++ b/xla/stream_executor/rocm/rocm_capture_invalidation.cc @@ -0,0 +1,191 @@ +/* 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 "xla/stream_executor/rocm/rocm_capture_invalidation.h" + +#include +#include + +#include "absl/base/no_destructor.h" +#include "absl/cleanup/cleanup.h" +#include "absl/functional/function_ref.h" +#include "absl/log/log.h" +#include "absl/status/status.h" +#include "absl/status/status_macros.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "rocm/include/hip/hip_runtime.h" +#include "xla/stream_executor/rocm/rocm_status.h" +#include "xla/stream_executor/stream_executor.h" +#include "xla/tsl/platform/env.h" + +namespace stream_executor::gpu { +namespace { + +constexpr size_t kSyncCallBytes = 4; + +void LogIfError(hipError_t result, absl::string_view detail) { + if (absl::Status status = ToStatus(result, detail); !status.ok()) { + LOG(ERROR) << status; + } +} + +absl::StatusOr SyncCallInvalidatesCaptureOnThisThread(int capture_device, + int sync_device) { + void* buffer = nullptr; + ABSL_RETURN_IF_ERROR(ToStatus(hipSetDevice(sync_device), "hipSetDevice")); + ABSL_RETURN_IF_ERROR( + ToStatus(hipMalloc(&buffer, kSyncCallBytes), "hipMalloc")); + absl::Cleanup free_buffer = [buffer] { + LogIfError(hipFree(buffer), "hipFree"); + }; + + hipStream_t stream = nullptr; + ABSL_RETURN_IF_ERROR(ToStatus(hipSetDevice(capture_device), "hipSetDevice")); + ABSL_RETURN_IF_ERROR( + ToStatus(hipStreamCreateWithFlags(&stream, hipStreamDefault), + "hipStreamCreateWithFlags")); + absl::Cleanup destroy_stream = [stream] { + LogIfError(hipStreamDestroy(stream), "hipStreamDestroy"); + }; + ABSL_RETURN_IF_ERROR( + ToStatus(hipStreamBeginCapture(stream, hipStreamCaptureModeThreadLocal), + "hipStreamBeginCapture")); + + hipError_t sync_result = hipErrorUnknown; + std::unique_ptr sync_thread(tsl::Env::Default()->StartThread( + tsl::ThreadOptions(), "hip_sync_call", [&] { + sync_result = hipSetDevice(sync_device); + if (sync_result == hipSuccess) { + sync_result = hipMemset(buffer, 0, kSyncCallBytes); + } + })); + sync_thread.reset(); + + hipStreamCaptureStatus capture_status = hipStreamCaptureStatusNone; + hipError_t query_result = hipStreamIsCapturing(stream, &capture_status); + hipGraph_t graph = nullptr; + hipError_t end_result = hipStreamEndCapture(stream, &graph); + if (graph != nullptr) { + LogIfError(hipGraphDestroy(graph), "hipGraphDestroy"); + } + ABSL_RETURN_IF_ERROR(ToStatus(query_result, "hipStreamIsCapturing")); + + if (capture_status == hipStreamCaptureStatusActive && + end_result == hipSuccess && sync_result == hipSuccess) { + return false; + } + if (capture_status == hipStreamCaptureStatusInvalidated && + end_result == hipErrorStreamCaptureInvalidated && + sync_result != hipSuccess) { + return true; + } + return absl::InternalError(absl::StrCat( + "Inconsistent stream capture after a synchronous HIP call: status ", + static_cast(capture_status), ", hipStreamEndCapture ", + hipGetErrorName(end_result), ", hipMemset ", + hipGetErrorName(sync_result))); +} + +// Returns whether a synchronous HIP call invalidates only the captures on the +// calling device. +absl::StatusOr IsCaptureInvalidationDeviceScoped( + int device, int other_device, + absl::FunctionRef(int, int)> + sync_call_invalidates_capture) { + ABSL_ASSIGN_OR_RETURN(bool same_device_invalidated, + sync_call_invalidates_capture(device, device)); + if (!same_device_invalidated) { + return absl::FailedPreconditionError( + "A synchronous HIP call left a capture on its own device intact, so " + "invalidation cannot be observed"); + } + ABSL_ASSIGN_OR_RETURN(bool other_device_invalidated, + sync_call_invalidates_capture(device, other_device)); + return !other_device_invalidated; +} + +void CreateDnnSupportUnlessDeviceScoped( + absl::Span executors, + absl::FunctionRef(int device, int other_device)> + is_capture_invalidation_device_scoped) { + if (executors.size() < 2) { + return; + } + absl::StatusOr device_scoped = is_capture_invalidation_device_scoped( + executors[0]->device_ordinal(), executors[1]->device_ordinal()); + if (device_scoped.ok() && *device_scoped) { + VLOG(1) << "HIP invalidates stream captures per device; DNN support stays " + "lazy"; + return; + } + if (device_scoped.ok()) { + VLOG(1) << "HIP invalidates stream captures on every device; creating DNN " + "support for " + << executors.size() << " devices"; + } else { + LOG(WARNING) << "Creating DNN support for " << executors.size() + << " devices now: " << device_scoped.status(); + } + for (StreamExecutor* executor : executors) { + // Not fatal: programs without convolutions must still start. + if (executor->AsDnn() == nullptr) { + LOG(WARNING) << "Failed to create DNN support for device " + << executor->device_ordinal(); + } + } +} + +} // namespace + +absl::StatusOr SyncCallInvalidatesCapture(int capture_device, + int sync_device) { + absl::StatusOr result; + std::unique_ptr thread(tsl::Env::Default()->StartThread( + tsl::ThreadOptions(), "hip_capture_probe", [&] { + result = + SyncCallInvalidatesCaptureOnThisThread(capture_device, sync_device); + })); + thread.reset(); + return result; +} + +void CreateDnnSupportUnlessCaptureInvalidationIsDeviceScoped( + absl::Span executors) { + CreateDnnSupportUnlessDeviceScoped( + executors, [](int device, int other_device) { + // How far invalidation reaches is a property of the HIP runtime, so + // the first answer holds for every pair of devices. + static const absl::NoDestructor> device_scoped( + IsCaptureInvalidationDeviceScoped(device, other_device, + SyncCallInvalidatesCapture)); + return *device_scoped; + }); +} + +void CreateDnnSupportUnlessCaptureInvalidationIsDeviceScoped( + absl::Span executors, + absl::FunctionRef(int capture_device, int sync_device)> + sync_call_invalidates_capture) { + CreateDnnSupportUnlessDeviceScoped( + executors, [&](int device, int other_device) { + return IsCaptureInvalidationDeviceScoped(device, other_device, + sync_call_invalidates_capture); + }); +} + +} // namespace stream_executor::gpu diff --git a/xla/stream_executor/rocm/rocm_capture_invalidation.h b/xla/stream_executor/rocm/rocm_capture_invalidation.h new file mode 100644 index 00000000000000..82a9d0d87b28a7 --- /dev/null +++ b/xla/stream_executor/rocm/rocm_capture_invalidation.h @@ -0,0 +1,58 @@ +/* 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. +==============================================================================*/ + +#ifndef XLA_STREAM_EXECUTOR_ROCM_ROCM_CAPTURE_INVALIDATION_H_ +#define XLA_STREAM_EXECUTOR_ROCM_ROCM_CAPTURE_INVALIDATION_H_ + +#include "absl/functional/function_ref.h" +#include "absl/status/statusor.h" +#include "absl/types/span.h" +#include "xla/stream_executor/stream_executor.h" + +namespace stream_executor::gpu { + +// Returns whether a synchronous HIP call (hipMemset) that one thread makes on +// `sync_device` invalidates a stream capture that another thread has in +// progress on `capture_device`. Runs on threads of its own, so the caller's +// current device and last HIP error are left alone, and ends the capture it +// starts. +absl::StatusOr SyncCallInvalidatesCapture(int capture_device, + int sync_device); + +// Creating a device's DNN support makes synchronous HIP calls. Up to +// ROCm/rocm-systems#12274, HIP invalidates every stream capture in the process +// on such a call, not only the captures on the calling device, so DNN support +// created lazily while another device records a command buffer breaks that +// recording. +// +// If `executors` spans more than one device, checks once per process, with +// SyncCallInvalidatesCapture, that a call on one device invalidates a capture +// on that device and not one on another device. Unless both hold, creates the +// DNN support of every executor now, one at a time. Call it only after every +// executor's Init() has returned: MIOpen's first handle loads a library whose +// HIP code objects take comgr's lock under the dynamic loader's lock, and a +// device still in Init() can hold comgr's lock while waiting for the loader's. +void CreateDnnSupportUnlessCaptureInvalidationIsDeviceScoped( + absl::Span executors); + +// As above, but asks `sync_call_invalidates_capture`, on every call. +void CreateDnnSupportUnlessCaptureInvalidationIsDeviceScoped( + absl::Span executors, + absl::FunctionRef(int capture_device, int sync_device)> + sync_call_invalidates_capture); + +} // namespace stream_executor::gpu + +#endif // XLA_STREAM_EXECUTOR_ROCM_ROCM_CAPTURE_INVALIDATION_H_ diff --git a/xla/stream_executor/rocm/rocm_capture_invalidation_test.cc b/xla/stream_executor/rocm/rocm_capture_invalidation_test.cc new file mode 100644 index 00000000000000..b0a570c442cc08 --- /dev/null +++ b/xla/stream_executor/rocm/rocm_capture_invalidation_test.cc @@ -0,0 +1,88 @@ +/* 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 "xla/stream_executor/rocm/rocm_capture_invalidation.h" + +#include +#include + +#include +#include +#include "absl/status/status.h" +#include "absl/status/status_matchers.h" +#include "absl/status/statusor.h" +#include "rocm/include/hip/hip_runtime.h" +#include "xla/stream_executor/mock_stream_executor.h" +#include "xla/stream_executor/stream_executor.h" + +namespace stream_executor::gpu { +namespace { + +using ::absl_testing::IsOk; +using ::absl_testing::IsOkAndHolds; +using ::testing::MockFunction; +using ::testing::NiceMock; +using ::testing::Return; +using ::testing::StrictMock; + +TEST(CreateDnnSupportTest, CreatesDnnSupportOnEveryDeviceWhenProbeFails) { + constexpr std::array kOrdinals = {3, 5, 6}; + std::array, kOrdinals.size()> executors; + std::array executor_ptrs; + for (size_t i = 0; i < kOrdinals.size(); ++i) { + ON_CALL(executors[i], device_ordinal()).WillByDefault(Return(kOrdinals[i])); + EXPECT_CALL(executors[i], AsDnn()); + executor_ptrs[i] = &executors[i]; + } + StrictMock(int, int)>> probe; + EXPECT_CALL(probe, Call(3, 3)).WillOnce(Return(true)); + EXPECT_CALL(probe, Call(3, 5)) + .WillOnce(Return(absl::InternalError("probe failed"))); + + CreateDnnSupportUnlessCaptureInvalidationIsDeviceScoped( + executor_ptrs, probe.AsStdFunction()); +} + +TEST(SyncCallInvalidatesCaptureTest, + InvalidatesASameDeviceCaptureAndLeavesNoTrace) { + int device_count = 0; + if (hipGetDeviceCount(&device_count) != hipSuccess || device_count < 1) { + GTEST_SKIP() << "Test requires a GPU"; + } + const int last_device = device_count - 1; + ASSERT_EQ(hipSetDevice(last_device), hipSuccess); + ASSERT_EQ(hipGetLastError(), hipSuccess); + + EXPECT_THAT(SyncCallInvalidatesCapture(0, 0), IsOkAndHolds(true)); + EXPECT_THAT(SyncCallInvalidatesCapture(0, last_device), IsOk()); + + int current_device = -1; + EXPECT_EQ(hipGetDevice(¤t_device), hipSuccess); + EXPECT_EQ(current_device, last_device); + EXPECT_EQ(hipGetLastError(), hipSuccess); + hipStreamCaptureStatus status = hipStreamCaptureStatusActive; + EXPECT_EQ(hipStreamIsCapturing(nullptr, &status), hipSuccess); + EXPECT_EQ(status, hipStreamCaptureStatusNone); + + // A capture left behind, even an invalidated one, fails synchronous calls. + void* buffer = nullptr; + ASSERT_EQ(hipSetDevice(0), hipSuccess); + ASSERT_EQ(hipMalloc(&buffer, 4), hipSuccess); + EXPECT_EQ(hipMemset(buffer, 0, 4), hipSuccess); + EXPECT_EQ(hipFree(buffer), hipSuccess); +} + +} // namespace +} // namespace stream_executor::gpu