未验证 提交 48937020 编写于 作者: F From00 提交者: GitHub

Skip zero-size Allocation in RecordStream (#38264)

上级 1f445bf3
...@@ -311,6 +311,10 @@ class AllocatorFacadePrivate { ...@@ -311,6 +311,10 @@ class AllocatorFacadePrivate {
void RecordStream(std::shared_ptr<Allocation> allocation, void RecordStream(std::shared_ptr<Allocation> allocation,
const gpuStream_t& stream) { const gpuStream_t& stream) {
if (allocation->size() == 0) {
return;
}
StreamSafeCUDAAllocation* stream_safe_cuda_allocation = StreamSafeCUDAAllocation* stream_safe_cuda_allocation =
dynamic_cast<StreamSafeCUDAAllocation*>(allocation.get()); dynamic_cast<StreamSafeCUDAAllocation*>(allocation.get());
PADDLE_ENFORCE_NOT_NULL(stream_safe_cuda_allocation, PADDLE_ENFORCE_NOT_NULL(stream_safe_cuda_allocation,
......
...@@ -243,6 +243,27 @@ TEST(StreamSafeCUDAAllocInterfaceTest, GetAllocatorInterfaceTest) { ...@@ -243,6 +243,27 @@ TEST(StreamSafeCUDAAllocInterfaceTest, GetAllocatorInterfaceTest) {
CheckMemLeak(place); CheckMemLeak(place);
} }
TEST(StreamSafeCUDAAllocInterfaceTest, ZeroSizeRecordStreamTest) {
platform::CUDAPlace place = platform::CUDAPlace();
std::shared_ptr<Allocation> zero_size_allocation = AllocShared(place, 0);
EXPECT_EQ(zero_size_allocation->ptr(), nullptr);
gpuStream_t stream;
#ifdef PADDLE_WITH_CUDA
PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamCreate(&stream));
#else
PADDLE_ENFORCE_GPU_SUCCESS(hipStreamCreate(&stream));
#endif
EXPECT_NO_THROW(RecordStream(zero_size_allocation, stream));
#ifdef PADDLE_WITH_CUDA
PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamDestroy(stream));
#else
PADDLE_ENFORCE_GPU_SUCCESS(hipStreamDestroy(stream));
#endif
}
TEST(StreamSafeCUDAAllocInterfaceTest, GetStreamInterfaceTest) { TEST(StreamSafeCUDAAllocInterfaceTest, GetStreamInterfaceTest) {
platform::CUDAPlace place = platform::CUDAPlace(); platform::CUDAPlace place = platform::CUDAPlace();
size_t alloc_size = 256; size_t alloc_size = 256;
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册