From 821fa6f64dd0e44025f021dfa79350e62a9d301b Mon Sep 17 00:00:00 2001 From: bolunz Date: Mon, 12 Oct 2026 15:53:40 +0800 Subject: [PATCH 1/3] feat: init stream guard --- .../include/core/runtime/device_guard.h | 6 + .../include/core/runtime/stream_guard.h | 30 +++ .../src/core/runtime/cpu/cpu_guard_impl.cc | 7 + .../src/core/runtime/cpu/cpu_guard_impl.h | 2 + .../src/core/runtime/cuda/cuda_guard_impl.cc | 25 ++- .../src/core/runtime/cuda/cuda_guard_impl.h | 2 + infini_train/src/core/runtime/device_guard.cc | 5 + infini_train/src/core/runtime/stream_guard.cc | 16 ++ tests/backend/CMakeLists.txt | 10 + tests/backend/test_stream_guard.cc | 199 ++++++++++++++++++ 10 files changed, 300 insertions(+), 2 deletions(-) create mode 100644 infini_train/include/core/runtime/stream_guard.h create mode 100644 infini_train/src/core/runtime/stream_guard.cc create mode 100644 tests/backend/test_stream_guard.cc diff --git a/infini_train/include/core/runtime/device_guard.h b/infini_train/include/core/runtime/device_guard.h index ced3aeabf..ebb0a6fde 100644 --- a/infini_train/include/core/runtime/device_guard.h +++ b/infini_train/include/core/runtime/device_guard.h @@ -86,6 +86,12 @@ class DeviceGuardImpl { virtual Stream *GetStream(Device) const; + // Set the thread-local current stream for device and return its previous + // stream. Does not change the current device or synchronize. Stream ownership + // is unchanged. The caller must provide a stream belonging to device and keep + // it alive while selected. + virtual Stream *ExchangeStream(Device device, Stream *stream) const; + virtual Stream *CreateStream(Device) const; virtual Stream *CreateStreamWithPriority(Device, int priority) const; diff --git a/infini_train/include/core/runtime/stream_guard.h b/infini_train/include/core/runtime/stream_guard.h new file mode 100644 index 000000000..993fc5dba --- /dev/null +++ b/infini_train/include/core/runtime/stream_guard.h @@ -0,0 +1,30 @@ +#pragma once + +#include "infini_train/include/core/runtime/device_guard.h" + +namespace infini_train::core { + +// Select device and its current stream for this scope. On destruction, restore +// the target device's previous stream, then restore the original device within +// the same backend. The stream is borrowed, must belong to device, and must +// outlive this guard. No synchronization is performed; callers remain +// responsible for cross-stream dependencies. +class StreamGuard { + public: + explicit StreamGuard(Device device, Stream *stream); + ~StreamGuard(); + + StreamGuard(const StreamGuard &) = delete; + StreamGuard &operator=(const StreamGuard &) = delete; + StreamGuard(StreamGuard &&) = delete; + StreamGuard &operator=(StreamGuard &&) = delete; + + private: + // Declared first so device restoration also runs if stream exchange throws + // during construction. + DeviceGuard device_guard_; + DeviceGuardImpl *impl_; + Stream *previous_stream_; +}; + +} // namespace infini_train::core diff --git a/infini_train/src/core/runtime/cpu/cpu_guard_impl.cc b/infini_train/src/core/runtime/cpu/cpu_guard_impl.cc index 65217ba9b..5c7673358 100644 --- a/infini_train/src/core/runtime/cpu/cpu_guard_impl.cc +++ b/infini_train/src/core/runtime/cpu/cpu_guard_impl.cc @@ -39,6 +39,13 @@ Stream *CpuGuardImpl::GetStream(Device device) const { return nullptr; } +Stream *CpuGuardImpl::ExchangeStream(Device device, Stream *stream) const { + CHECK(device.type() == Device::DeviceType::kCPU); + VLOG(3) << "CpuGuardImpl::ExchangeStream is not supported. " + "Return nullptr."; + return nullptr; +} + Stream *CpuGuardImpl::CreateStream(Device device) const { CHECK(device.type() == Device::DeviceType::kCPU); VLOG(3) << "CpuGuardImpl::CreateStream is not supported. " diff --git a/infini_train/src/core/runtime/cpu/cpu_guard_impl.h b/infini_train/src/core/runtime/cpu/cpu_guard_impl.h index b4a211fbb..bb0c66650 100644 --- a/infini_train/src/core/runtime/cpu/cpu_guard_impl.h +++ b/infini_train/src/core/runtime/cpu/cpu_guard_impl.h @@ -27,6 +27,8 @@ class CpuGuardImpl final : public DeviceGuardImpl { // Stream management (explicitly unsupported for now) Stream *GetStream(Device device) const override; + Stream *ExchangeStream(Device device, Stream *stream) const override; + Stream *CreateStream(Device device) const override; Stream *CreateStreamWithPriority(Device device, int priority) const override; diff --git a/infini_train/src/core/runtime/cuda/cuda_guard_impl.cc b/infini_train/src/core/runtime/cuda/cuda_guard_impl.cc index 33974ad49..69c2c0b87 100644 --- a/infini_train/src/core/runtime/cuda/cuda_guard_impl.cc +++ b/infini_train/src/core/runtime/cuda/cuda_guard_impl.cc @@ -16,6 +16,7 @@ constexpr int kMaxGpus = 8; constexpr size_t kBytesPerMB = 1024ULL * 1024ULL; static std::array, kMaxGpus> cuda_streams; +static thread_local std::array current_streams = {}; static std::array, kMaxGpus> cuda_blas_handles; static std::array device_stream_flags; @@ -51,6 +52,7 @@ void CudaGuardImpl::InitSingleStream(Device device) { CUDA_CHECK(cudaSetDevice(device.index())); cuda_streams[device.index()] = std::make_unique(); + current_streams[device.index()] = cuda_streams[device.index()].get(); CUDA_CHECK(cudaSetDevice(current_device)); } @@ -103,7 +105,23 @@ Stream *CudaGuardImpl::GetStream(Device device) const { // times within the same process (e.g. in unit tests). std::call_once(device_stream_flags.at(device.index()), InitSingleStream, device); - return cuda_streams.at(device.index()).get(); + if (current_streams.at(device.index()) == nullptr) { + current_streams.at(device.index()) = cuda_streams.at(device.index()).get(); + } + return current_streams.at(device.index()); +} + +Stream *CudaGuardImpl::ExchangeStream(Device device, Stream *stream) const { + CheckCudaDevice(device); + std::call_once(device_stream_flags.at(device.index()), InitSingleStream, + device); + CHECK_NOTNULL(stream); + if (current_streams.at(device.index()) == nullptr) { + current_streams.at(device.index()) = cuda_streams.at(device.index()).get(); + } + auto *previous_stream = current_streams.at(device.index()); + current_streams[device.index()] = stream; + return previous_stream; } Stream *CudaGuardImpl::CreateStream(Device device) const { @@ -231,7 +249,10 @@ BlasHandle *CudaGuardImpl::GetBlasHandle(Device device) const { CheckCudaDevice(device); std::call_once(device_handle_flags.at(device.index()), InitSingleHandle, device); - return cuda_blas_handles.at(device.index()).get(); + auto *handle = cuda_blas_handles.at(device.index()).get(); + CUBLAS_CHECK(cublasSetStream(handle->cublas_handle(), + GetCudaStream(GetStream(device)))); + return handle; } // memory diff --git a/infini_train/src/core/runtime/cuda/cuda_guard_impl.h b/infini_train/src/core/runtime/cuda/cuda_guard_impl.h index 9f821bba4..287ace253 100644 --- a/infini_train/src/core/runtime/cuda/cuda_guard_impl.h +++ b/infini_train/src/core/runtime/cuda/cuda_guard_impl.h @@ -32,6 +32,8 @@ class CudaGuardImpl final : public DeviceGuardImpl { // constructor/destructor of class Stream Stream *GetStream(Device device) const override; + Stream *ExchangeStream(Device device, Stream *stream) const override; + Stream *CreateStream(Device device) const override; Stream *CreateStreamWithPriority(Device device, int priority) const override; diff --git a/infini_train/src/core/runtime/device_guard.cc b/infini_train/src/core/runtime/device_guard.cc index a965fa94f..398579763 100644 --- a/infini_train/src/core/runtime/device_guard.cc +++ b/infini_train/src/core/runtime/device_guard.cc @@ -29,6 +29,11 @@ Stream *DeviceGuardImpl::GetStream(Device) const { return nullptr; // unreachable } +Stream *DeviceGuardImpl::ExchangeStream(Device, Stream *) const { + LOG(FATAL) << "DeviceGuardImpl::ExchangeStream is not implemented."; + return nullptr; // unreachable +} + Stream *DeviceGuardImpl::CreateStream(Device) const { LOG(FATAL) << "DeviceGuardImpl::CreateStream is not implemented."; return nullptr; // unreachable diff --git a/infini_train/src/core/runtime/stream_guard.cc b/infini_train/src/core/runtime/stream_guard.cc new file mode 100644 index 000000000..ad05593bf --- /dev/null +++ b/infini_train/src/core/runtime/stream_guard.cc @@ -0,0 +1,16 @@ +#include "infini_train/include/core/runtime/stream_guard.h" + +namespace infini_train::core { + +StreamGuard::StreamGuard(Device device, Stream *stream) + : device_guard_(device), + impl_(GetDeviceGuardImpl(device.type())), + previous_stream_(impl_->ExchangeStream(device, stream)) {} + +StreamGuard::~StreamGuard() { + impl_->ExchangeStream(device_guard_.current_device(), previous_stream_); + // device_guard_ restores the original device after the stream has been + // restored. +} + +} // namespace infini_train::core diff --git a/tests/backend/CMakeLists.txt b/tests/backend/CMakeLists.txt index 8ddd45755..6bb3affb4 100644 --- a/tests/backend/CMakeLists.txt +++ b/tests/backend/CMakeLists.txt @@ -7,3 +7,13 @@ infini_train_add_test(test_privateuse1_backend_validation SOURCES test_privateuse1_backend_validation.cc LABELS cpu ) + +if(USE_CUDA) + infini_train_add_test(test_stream_guard_cuda + SOURCES test_stream_guard.cc + LABELS cuda accelerator hardware + BINARY_REGISTRATION + RUN_SERIAL + TEST_TIMEOUT 60 + ) +endif() diff --git a/tests/backend/test_stream_guard.cc b/tests/backend/test_stream_guard.cc new file mode 100644 index 000000000..3ca20d96e --- /dev/null +++ b/tests/backend/test_stream_guard.cc @@ -0,0 +1,199 @@ +#include +#include +#include +#include +#include + +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/core/runtime/stream_guard.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h" + +namespace infini_train::test { +namespace { + +static_assert(!std::is_default_constructible_v); +static_assert(!std::is_copy_constructible_v); +static_assert(!std::is_copy_assignable_v); +static_assert(!std::is_move_constructible_v); +static_assert(!std::is_move_assignable_v); + +class StreamGuardTest : public ::testing::Test { + protected: + void SetUp() override { + if (cudaGetDeviceCount(&device_count_) != cudaSuccess || + device_count_ == 0) { + GTEST_SKIP() << "Requires CUDA"; + } + impl_ = core::GetDeviceGuardImpl(Device::DeviceType::kCUDA); + original_device_ = impl_->GetDevice(); + impl_->SetDevice(device0_); + original_stream0_ = impl_->GetStream(device0_); + if (device_count_ > 1) { + original_stream1_ = impl_->GetStream(device1_); + } + } + + void TearDown() override { + if (!impl_) { + return; + } + impl_->ExchangeStream(device0_, original_stream0_); + if (device_count_ > 1) { + impl_->ExchangeStream(device1_, original_stream1_); + } + impl_->SetDevice(original_device_); + } + + const Device device0_{Device::DeviceType::kCUDA, 0}; + const Device device1_{Device::DeviceType::kCUDA, 1}; + int device_count_ = 0; + core::DeviceGuardImpl *impl_ = nullptr; + Device original_device_; + core::Stream *original_stream0_ = nullptr; + core::Stream *original_stream1_ = nullptr; +}; + +TEST_F(StreamGuardTest, NestedSameDeviceRestoresPreviousStream) { + std::unique_ptr first(impl_->CreateStream(device0_)); + std::unique_ptr second(impl_->CreateStream(device0_)); + { + core::StreamGuard outer(device0_, first.get()); + EXPECT_EQ(impl_->GetDevice(), device0_); + EXPECT_EQ(impl_->GetStream(device0_), first.get()); + { + core::StreamGuard inner(device0_, second.get()); + EXPECT_EQ(impl_->GetStream(device0_), second.get()); + } + EXPECT_EQ(impl_->GetStream(device0_), first.get()); + } + EXPECT_EQ(impl_->GetStream(device0_), original_stream0_); + EXPECT_EQ(impl_->GetDevice(), device0_); +} + +TEST_F(StreamGuardTest, CrossDeviceRestoresTargetStreamBeforeOriginalDevice) { + if (device_count_ < 2) { + GTEST_SKIP() << "Requires two GPUs"; + } + std::unique_ptr first(impl_->CreateStream(device0_)); + std::unique_ptr second(impl_->CreateStream(device1_)); + std::unique_ptr third(impl_->CreateStream(device0_)); + { + core::StreamGuard outer(device0_, first.get()); + { + core::StreamGuard inner(device1_, second.get()); + EXPECT_EQ(impl_->GetDevice(), device1_); + EXPECT_EQ(impl_->GetStream(device0_), first.get()); + EXPECT_EQ(impl_->GetStream(device1_), second.get()); + { + core::StreamGuard nested(device0_, third.get()); + EXPECT_EQ(impl_->GetDevice(), device0_); + EXPECT_EQ(impl_->GetStream(device0_), third.get()); + } + EXPECT_EQ(impl_->GetDevice(), device1_); + EXPECT_EQ(impl_->GetStream(device0_), first.get()); + } + EXPECT_EQ(impl_->GetDevice(), device0_); + EXPECT_EQ(impl_->GetStream(device0_), first.get()); + EXPECT_EQ(impl_->GetStream(device1_), original_stream1_); + } + EXPECT_EQ(impl_->GetStream(device0_), original_stream0_); +} + +TEST_F(StreamGuardTest, RestoresAfterException) { + if (device_count_ < 2) { + GTEST_SKIP() << "Requires two GPUs"; + } + std::unique_ptr stream(impl_->CreateStream(device1_)); + EXPECT_THROW( + { + core::StreamGuard guard(device1_, stream.get()); + EXPECT_EQ(impl_->GetDevice(), device1_); + throw std::runtime_error("leave guarded scope"); + }, + std::runtime_error); + EXPECT_EQ(impl_->GetDevice(), device0_); + EXPECT_EQ(impl_->GetStream(device1_), original_stream1_); +} + +TEST_F(StreamGuardTest, RestoresEvenIfCurrentDeviceChangesInsideScope) { + if (device_count_ < 2) { + GTEST_SKIP() << "Requires two GPUs"; + } + std::unique_ptr stream(impl_->CreateStream(device1_)); + { + core::StreamGuard guard(device1_, stream.get()); + impl_->SetDevice(device0_); + } + EXPECT_EQ(impl_->GetStream(device1_), original_stream1_); + EXPECT_EQ(impl_->GetStream(device0_), original_stream0_); + EXPECT_EQ(impl_->GetDevice(), device0_); +} + +TEST_F(StreamGuardTest, ExchangeStreamAloneDoesNotSwitchDevices) { + if (device_count_ < 2) { + GTEST_SKIP() << "Requires two GPUs"; + } + std::unique_ptr stream(impl_->CreateStream(device1_)); + auto *previous = impl_->ExchangeStream(device1_, stream.get()); + EXPECT_EQ(previous, original_stream1_); + EXPECT_EQ(impl_->GetDevice(), device0_); + EXPECT_EQ(impl_->GetStream(device1_), stream.get()); + impl_->ExchangeStream(device1_, previous); +} + +TEST_F(StreamGuardTest, CurrentStreamSelectionIsThreadLocal) { + std::unique_ptr stream(impl_->CreateStream(device0_)); + core::StreamGuard guard(device0_, stream.get()); + std::atomic independent{false}; + std::thread worker([&] { + core::DeviceGuard device_guard(device0_); + independent = impl_->GetStream(device0_) == original_stream0_; + }); + worker.join(); + EXPECT_TRUE(independent.load()); + EXPECT_EQ(impl_->GetStream(device0_), stream.get()); +} + +TEST_F(StreamGuardTest, AllocationAndGemmFollowGuardedStream) { + if (device_count_ < 2) { + GTEST_SKIP() << "Requires two GPUs"; + } + std::unique_ptr stream(impl_->CreateStream(device1_)); + std::shared_ptr output; + { + core::StreamGuard guard(device1_, stream.get()); + EXPECT_EQ(impl_->GetDevice(), device1_); + auto *handle = dynamic_cast( + impl_->GetBlasHandle(device1_)); + ASSERT_NE(handle, nullptr); + cudaStream_t blas_stream = nullptr; + ASSERT_EQ(cublasGetStream(handle->cublas_handle(), &blas_stream), + CUBLAS_STATUS_SUCCESS); + EXPECT_EQ( + blas_stream, + dynamic_cast(stream.get())->cuda_stream()); + auto input = std::make_shared(std::vector{2, 2}, + DataType::kFLOAT32, device1_); + input->Fill(2.0f); + output = input->Matmul(input); + // The guard does not synchronize: join explicitly before consuming output + // on the original stream. + impl_->SynchronizeStream(stream.get()); + } + EXPECT_EQ(impl_->GetDevice(), device0_); + EXPECT_EQ(impl_->GetStream(device1_), original_stream1_); + auto cpu = output->To(Device()); + const auto *values = static_cast(cpu.DataPtr()); + for (size_t i = 0; i < cpu.NumElements(); ++i) { + EXPECT_FLOAT_EQ(values[i], 8.0f); + } +} + +} // namespace +} // namespace infini_train::test From f05f12185e2eb695c7faa6ce43d711721589c47e Mon Sep 17 00:00:00 2001 From: bolunz Date: Wed, 14 Oct 2026 10:51:55 +0800 Subject: [PATCH 2/3] fix: add lightweight BlasHandlePool to avoid multiple use of a handle from multiple streams --- .../src/core/runtime/cuda/cuda_guard_impl.cc | 36 ++++++-- tests/backend/test_stream_guard.cc | 85 +++++++++++++++++++ 2 files changed, 115 insertions(+), 6 deletions(-) diff --git a/infini_train/src/core/runtime/cuda/cuda_guard_impl.cc b/infini_train/src/core/runtime/cuda/cuda_guard_impl.cc index 69c2c0b87..8174b89d6 100644 --- a/infini_train/src/core/runtime/cuda/cuda_guard_impl.cc +++ b/infini_train/src/core/runtime/cuda/cuda_guard_impl.cc @@ -3,6 +3,7 @@ #include #include #include +#include #include "infini_train/include/common/cuda/common_cuda.h" #include "infini_train/include/core/runtime/runtime_common.h" @@ -17,10 +18,18 @@ constexpr size_t kBytesPerMB = 1024ULL * 1024ULL; static std::array, kMaxGpus> cuda_streams; static thread_local std::array current_streams = {}; -static std::array, kMaxGpus> cuda_blas_handles; +// Each thread exclusively reserves one handle per device. The pool retains +// handles after thread exit for reuse, since CudaBlasHandle deliberately does +// not destroy its CUDA resources during teardown. +struct BlasHandlePool { + std::mutex mutex; + std::array>, kMaxGpus> handles; +}; +static BlasHandlePool cuda_blas_handle_pool; +static thread_local std::array, kMaxGpus> + cuda_blas_handles; static std::array device_stream_flags; -static std::array device_handle_flags; inline void CheckCudaDevice(Device device) { CHECK(device.type() == Device::DeviceType::kCUDA) << std::format( @@ -60,6 +69,20 @@ void CudaGuardImpl::InitSingleStream(Device device) { void CudaGuardImpl::InitSingleHandle(Device device) { CheckCudaDevice(device); + auto &handle = cuda_blas_handles.at(device.index()); + if (handle) { + return; + } + std::lock_guard lock(cuda_blas_handle_pool.mutex); + auto &pooled_handles = cuda_blas_handle_pool.handles.at(device.index()); + for (const auto &pooled_handle : pooled_handles) { + // A pool-only reference means no thread currently reserves this handle. + if (pooled_handle.use_count() == 1) { + handle = pooled_handle; + return; + } + } + int current_device = -1; CUDA_CHECK(cudaGetDevice(¤t_device)); CUDA_CHECK(cudaSetDevice(device.index())); @@ -67,8 +90,8 @@ void CudaGuardImpl::InitSingleHandle(Device device) { std::call_once(device_stream_flags.at(device.index()), InitSingleStream, device); - cuda_blas_handles[device.index()] = - std::make_unique(cuda_streams[device.index()].get()); + handle = std::make_shared(cuda_streams[device.index()].get()); + pooled_handles.push_back(handle); CUDA_CHECK(cudaSetDevice(current_device)); } @@ -247,8 +270,9 @@ void CudaGuardImpl::SynchronizeStream(Stream *stream) const { // blas BlasHandle *CudaGuardImpl::GetBlasHandle(Device device) const { CheckCudaDevice(device); - std::call_once(device_handle_flags.at(device.index()), InitSingleHandle, - device); + if (!cuda_blas_handles.at(device.index())) { + InitSingleHandle(device); + } auto *handle = cuda_blas_handles.at(device.index()).get(); CUBLAS_CHECK(cublasSetStream(handle->cublas_handle(), GetCudaStream(GetStream(device)))); diff --git a/tests/backend/test_stream_guard.cc b/tests/backend/test_stream_guard.cc index 3ca20d96e..cb4505d02 100644 --- a/tests/backend/test_stream_guard.cc +++ b/tests/backend/test_stream_guard.cc @@ -160,6 +160,91 @@ TEST_F(StreamGuardTest, CurrentStreamSelectionIsThreadLocal) { EXPECT_EQ(impl_->GetStream(device0_), stream.get()); } +TEST_F(StreamGuardTest, BlasHandleBindingIsThreadLocal) { + std::unique_ptr first(impl_->CreateStream(device0_)); + std::unique_ptr second(impl_->CreateStream(device0_)); + core::StreamGuard guard(device0_, first.get()); + auto *handle = dynamic_cast( + impl_->GetBlasHandle(device0_)); + ASSERT_NE(handle, nullptr); + + cublasHandle_t worker_handle = nullptr; + std::thread worker([&] { + core::StreamGuard worker_guard(device0_, second.get()); + auto *worker_blas = dynamic_cast( + impl_->GetBlasHandle(device0_)); + ASSERT_NE(worker_blas, nullptr); + worker_handle = worker_blas->cublas_handle(); + EXPECT_EQ(impl_->GetBlasHandle(device0_), worker_blas); + cudaStream_t stream = nullptr; + ASSERT_EQ(cublasGetStream(worker_handle, &stream), CUBLAS_STATUS_SUCCESS); + EXPECT_EQ( + stream, + dynamic_cast(second.get())->cuda_stream()); + }); + worker.join(); + + EXPECT_NE(handle->cublas_handle(), worker_handle); + // Do not refresh the main thread's handle: another thread must not change its + // binding. + cudaStream_t stream = nullptr; + ASSERT_EQ(cublasGetStream(handle->cublas_handle(), &stream), + CUBLAS_STATUS_SUCCESS); + EXPECT_EQ(stream, + dynamic_cast(first.get())->cuda_stream()); +} + +TEST_F(StreamGuardTest, BlasHandleIsReusedAfterThreadExit) { + std::unique_ptr first(impl_->CreateStream(device0_)); + std::unique_ptr second(impl_->CreateStream(device0_)); + auto get_handle = [&](core::Stream *selected_stream, cublasHandle_t &result) { + core::StreamGuard guard(device0_, selected_stream); + auto *handle = dynamic_cast( + impl_->GetBlasHandle(device0_)); + ASSERT_NE(handle, nullptr); + result = handle->cublas_handle(); + cudaStream_t stream = nullptr; + ASSERT_EQ(cublasGetStream(result, &stream), CUBLAS_STATUS_SUCCESS); + EXPECT_EQ( + stream, + dynamic_cast(selected_stream)->cuda_stream()); + }; + cublasHandle_t first_handle = nullptr; + std::thread first_worker([&] { get_handle(first.get(), first_handle); }); + first_worker.join(); + cublasHandle_t second_handle = nullptr; + std::thread second_worker([&] { get_handle(second.get(), second_handle); }); + second_worker.join(); + ASSERT_NE(first_handle, nullptr); + EXPECT_EQ(first_handle, second_handle); +} + +TEST_F(StreamGuardTest, BlasHandleFollowsNestedStreamGuards) { + std::unique_ptr first(impl_->CreateStream(device0_)); + std::unique_ptr second(impl_->CreateStream(device0_)); + auto expect_blas_stream = [&](core::Stream *expected_stream) { + auto *handle = dynamic_cast( + impl_->GetBlasHandle(device0_)); + ASSERT_NE(handle, nullptr); + cudaStream_t stream = nullptr; + ASSERT_EQ(cublasGetStream(handle->cublas_handle(), &stream), + CUBLAS_STATUS_SUCCESS); + EXPECT_EQ( + stream, + dynamic_cast(expected_stream)->cuda_stream()); + }; + { + core::StreamGuard outer(device0_, first.get()); + expect_blas_stream(first.get()); + { + core::StreamGuard inner(device0_, second.get()); + expect_blas_stream(second.get()); + } + expect_blas_stream(first.get()); + } + expect_blas_stream(original_stream0_); +} + TEST_F(StreamGuardTest, AllocationAndGemmFollowGuardedStream) { if (device_count_ < 2) { GTEST_SKIP() << "Requires two GPUs"; From 62f978007cebf93a5254ae050c39531c9ed9e7bb Mon Sep 17 00:00:00 2001 From: bolunz Date: Thu, 22 Oct 2026 14:34:00 +0800 Subject: [PATCH 3/3] style: clang-tidy check --- infini_train/include/core/runtime/stream_guard.h | 6 ++++++ infini_train/src/core/runtime/stream_guard.cc | 2 ++ 2 files changed, 8 insertions(+) diff --git a/infini_train/include/core/runtime/stream_guard.h b/infini_train/include/core/runtime/stream_guard.h index 993fc5dba..0d20ce612 100644 --- a/infini_train/include/core/runtime/stream_guard.h +++ b/infini_train/include/core/runtime/stream_guard.h @@ -2,8 +2,14 @@ #include "infini_train/include/core/runtime/device_guard.h" +namespace infini_train { +class Device; +} + namespace infini_train::core { +class Stream; + // Select device and its current stream for this scope. On destruction, restore // the target device's previous stream, then restore the original device within // the same backend. The stream is borrowed, must belong to device, and must diff --git a/infini_train/src/core/runtime/stream_guard.cc b/infini_train/src/core/runtime/stream_guard.cc index ad05593bf..86d7c072f 100644 --- a/infini_train/src/core/runtime/stream_guard.cc +++ b/infini_train/src/core/runtime/stream_guard.cc @@ -1,5 +1,7 @@ #include "infini_train/include/core/runtime/stream_guard.h" +#include "infini_train/include/device.h" + namespace infini_train::core { StreamGuard::StreamGuard(Device device, Stream *stream)