chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,238 @@
|
||||
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
|
||||
//
|
||||
// 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 <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/core/Stream.h>
|
||||
#include <c10/cuda/CUDAFunctions.h>
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#endif
|
||||
|
||||
#include <atomic>
|
||||
#include <chrono>
|
||||
#include <thread>
|
||||
|
||||
#include "gtest/gtest.h"
|
||||
#include "paddle/phi/api/include/context_pool.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
namespace {
|
||||
|
||||
using StreamCallbackGate = std::atomic<bool>;
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
void BlockingStreamCallback(hipStream_t /*stream*/,
|
||||
hipError_t /*status*/,
|
||||
void* user_data) {
|
||||
auto* gate = static_cast<StreamCallbackGate*>(user_data);
|
||||
while (!gate->load(std::memory_order_acquire)) {
|
||||
std::this_thread::sleep_for(std::chrono::milliseconds(1));
|
||||
}
|
||||
}
|
||||
|
||||
void CreateRawStream(hipStream_t* stream) {
|
||||
C10_CUDA_CHECK(hipStreamCreate(stream));
|
||||
}
|
||||
|
||||
void DestroyRawStream(hipStream_t stream) {
|
||||
C10_CUDA_CHECK(hipStreamDestroy(stream));
|
||||
}
|
||||
|
||||
void ClearLastStreamError() { (void)hipGetLastError(); }
|
||||
#else
|
||||
void CUDART_CB BlockingStreamCallback(void* user_data) {
|
||||
auto* gate = static_cast<StreamCallbackGate*>(user_data);
|
||||
while (!gate->load(std::memory_order_acquire)) {
|
||||
std::this_thread::sleep_for(std::chrono::milliseconds(1));
|
||||
}
|
||||
}
|
||||
|
||||
void ClearLastStreamError() { (void)cudaGetLastError(); }
|
||||
#endif
|
||||
|
||||
} // namespace
|
||||
#endif
|
||||
|
||||
// Test device_count() works in both CPU and CUDA builds
|
||||
TEST(StreamTest, DeviceCount) {
|
||||
c10::DeviceIndex count = c10::cuda::device_count();
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
// In CUDA builds, should return actual device count (>= 0)
|
||||
EXPECT_GE(count, 0);
|
||||
#else
|
||||
// In CPU-only builds, should return 0
|
||||
EXPECT_EQ(count, 0);
|
||||
#endif
|
||||
}
|
||||
|
||||
// ==================== native_handle ====================
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
// CUDA stream: native_handle() should return the underlying cudaStream_t
|
||||
// encoded as void*. For the default (null) stream the id is 0, so the
|
||||
// pointer is nullptr; for a real stream it must be non-null.
|
||||
TEST(StreamTest, NativeHandleCudaDefaultStream) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
c10::Stream s = c10::cuda::getDefaultCUDAStream().unwrap();
|
||||
// Default stream encodes nullptr (id == 0), so native_handle() == nullptr.
|
||||
EXPECT_EQ(s.native_handle(), nullptr);
|
||||
}
|
||||
|
||||
TEST(StreamTest, NativeHandleCudaCurrentStream) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
auto cuda_stream = c10::cuda::getCurrentCUDAStream();
|
||||
c10::Stream s = cuda_stream.unwrap();
|
||||
// getCurrentCUDAStream wraps the real phi stream handle; calling
|
||||
// native_handle() must not throw.
|
||||
EXPECT_NO_THROW({ (void)s.native_handle(); });
|
||||
}
|
||||
#endif // PADDLE_WITH_CUDA || PADDLE_WITH_HIP
|
||||
|
||||
// CPU stream: native_handle() is not supported and must throw.
|
||||
TEST(StreamTest, NativeHandleCpuStreamThrows) {
|
||||
c10::Stream cpu_stream(c10::Stream::DEFAULT,
|
||||
c10::Device(c10::DeviceType::CPU, 0));
|
||||
EXPECT_THROW({ (void)cpu_stream.native_handle(); }, std::exception);
|
||||
}
|
||||
|
||||
// ==================== query ====================
|
||||
|
||||
// CPU stream is always ready.
|
||||
TEST(StreamTest, QueryCpuStreamReturnsTrue) {
|
||||
c10::Stream cpu_stream(c10::Stream::DEFAULT,
|
||||
c10::Device(c10::DeviceType::CPU, 0));
|
||||
EXPECT_TRUE(cpu_stream.query());
|
||||
}
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
// A freshly-obtained CUDA stream with no pending work must report ready.
|
||||
TEST(StreamTest, QueryCudaStreamReady) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
auto cuda_stream = c10::cuda::getCurrentCUDAStream();
|
||||
c10::Stream s = cuda_stream.unwrap();
|
||||
// synchronize first to ensure no pending work, then query should be true.
|
||||
EXPECT_NO_THROW(s.synchronize());
|
||||
EXPECT_TRUE(s.query());
|
||||
}
|
||||
|
||||
#endif // PADDLE_WITH_CUDA || PADDLE_WITH_HIP
|
||||
|
||||
// ==================== synchronize ====================
|
||||
|
||||
// CPU stream: synchronize() is a no-op and must not throw.
|
||||
TEST(StreamTest, SynchronizeCpuStream) {
|
||||
c10::Stream cpu_stream(c10::Stream::DEFAULT,
|
||||
c10::Device(c10::DeviceType::CPU, 0));
|
||||
EXPECT_NO_THROW(cpu_stream.synchronize());
|
||||
}
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
// CUDA stream: synchronize() must complete without error.
|
||||
TEST(StreamTest, SynchronizeCudaStream) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
auto cuda_stream = c10::cuda::getCurrentCUDAStream();
|
||||
c10::Stream s = cuda_stream.unwrap();
|
||||
EXPECT_NO_THROW(s.synchronize());
|
||||
}
|
||||
#endif // PADDLE_WITH_CUDA || PADDLE_WITH_HIP
|
||||
|
||||
// ==================== getDefaultCUDAStream ====================
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
// getDefaultCUDAStream must always return the null stream (id == 0),
|
||||
// which corresponds to cudaStreamDefault on the device.
|
||||
TEST(CUDAStreamTest, DefaultStreamIsNullStream) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
auto default_stream = c10::cuda::getDefaultCUDAStream();
|
||||
// id == 0 encodes cudaStreamDefault (the null stream, handle nullptr).
|
||||
EXPECT_EQ(default_stream.id(), static_cast<c10::StreamId>(0));
|
||||
}
|
||||
|
||||
// getDefaultCUDAStream must be stable: calling it twice returns equal streams.
|
||||
TEST(CUDAStreamTest, DefaultStreamIsStable) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
auto s1 = c10::cuda::getDefaultCUDAStream();
|
||||
auto s2 = c10::cuda::getDefaultCUDAStream();
|
||||
EXPECT_EQ(s1, s2);
|
||||
}
|
||||
|
||||
TEST(CUDAStreamTest, GetStreamFromPoolBoolOverloadPreservesHighPriority) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
auto low_priority_stream =
|
||||
c10::cuda::getStreamFromPool(/*isHighPriority=*/false);
|
||||
auto high_priority_stream =
|
||||
c10::cuda::getStreamFromPool(/*isHighPriority=*/true);
|
||||
auto explicit_high_priority_stream = c10::cuda::getStreamFromPool(-1);
|
||||
|
||||
const int low_priority = low_priority_stream.priority();
|
||||
const int high_priority = high_priority_stream.priority();
|
||||
const int explicit_high_priority = explicit_high_priority_stream.priority();
|
||||
|
||||
if (low_priority == explicit_high_priority) {
|
||||
return;
|
||||
}
|
||||
|
||||
EXPECT_EQ(high_priority, explicit_high_priority);
|
||||
EXPECT_NE(high_priority, low_priority);
|
||||
}
|
||||
|
||||
// After setCurrentCUDAStream redirects the current stream,
|
||||
// getDefaultCUDAStream must still return the null stream.
|
||||
TEST(CUDAStreamTest, DefaultStreamUnaffectedBySetCurrentCUDAStream) {
|
||||
if (!at::cuda::is_available()) {
|
||||
return;
|
||||
}
|
||||
// Snapshot the current stream before we touch it so we can
|
||||
// restore it afterward and avoid polluting subsequent tests.
|
||||
auto original_stream = c10::cuda::getCurrentCUDAStream();
|
||||
|
||||
// Obtain a non-default stream from the pool.
|
||||
auto pool_stream = c10::cuda::getStreamFromPool(/*isHighPriority=*/false);
|
||||
|
||||
// Redirect the current stream.
|
||||
c10::cuda::setCurrentCUDAStream(pool_stream);
|
||||
|
||||
auto default_stream = c10::cuda::getDefaultCUDAStream();
|
||||
auto current_stream = c10::cuda::getCurrentCUDAStream();
|
||||
auto place = phi::GPUPlace(current_stream.device_index());
|
||||
|
||||
// Default stream is still null; current stream has changed.
|
||||
EXPECT_EQ(default_stream.id(), static_cast<c10::StreamId>(0));
|
||||
EXPECT_NE(default_stream, current_stream);
|
||||
EXPECT_EQ(paddle::GetCurrentCUDAStream(place)->raw_stream(),
|
||||
current_stream.stream());
|
||||
|
||||
// Restore the original current stream.
|
||||
c10::cuda::setCurrentCUDAStream(original_stream);
|
||||
EXPECT_EQ(paddle::GetCurrentCUDAStream(place)->raw_stream(),
|
||||
original_stream.stream());
|
||||
}
|
||||
|
||||
#endif // PADDLE_WITH_CUDA || PADDLE_WITH_HIP
|
||||
Reference in New Issue
Block a user