Skip to content

Commit e222396

Browse files
youge325youge325
committed
add RecordStreamTest.cpp
1 parent 4a2ec44 commit e222396

File tree

2 files changed

+269
-0
lines changed

2 files changed

+269
-0
lines changed

‎CMakeLists.txt‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,8 +87,12 @@ set(TORCH_LIBRARIES "")
8787
file(GLOB_RECURSE TORCH_LIBRARIES "${TORCH_DIR}/lib/*.so"
8888
"${TORCH_DIR}/lib/*.a")
8989

90+
find_package(CUDAToolkit QUIET)
9091
set(TORCH_INCLUDE_DIR "${TORCH_DIR}/include"
9192
"${TORCH_DIR}/include/torch/csrc/api/include/")
93+
if(CUDAToolkit_FOUND)
94+
list(APPEND TORCH_INCLUDE_DIR "${CUDAToolkit_INCLUDE_DIRS}")
95+
endif()
9296

9397
set(TORCH_TARGET_FOLDER ${CMAKE_BINARY_DIR}/torch)
9498
set(BIN_PREFIX "torch_")
@@ -120,6 +124,9 @@ set(PADDLE_INCLUDE_DIR
120124
"${PADDLE_DIR}/include/paddle/phi/api/include/compat/"
121125
"${PADDLE_DIR}/include/paddle/phi/api/include/compat/torch/csrc/api/include/"
122126
)
127+
if(CUDAToolkit_FOUND)
128+
list(APPEND PADDLE_INCLUDE_DIR "${CUDAToolkit_INCLUDE_DIRS}")
129+
endif()
123130

124131
set(PADDLE_LIBRARIES
125132
"${PADDLE_DIR}/base/libpaddle.so" "${PADDLE_DIR}/libs/libcommon.so"

‎test/RecordStreamTest.cpp‎

Lines changed: 262 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,262 @@
1+
#include <ATen/ATen.h>
2+
#include <ATen/core/Tensor.h>
3+
#include <ATen/ops/zeros.h>
4+
#include <c10/core/Stream.h>
5+
#include <gtest/gtest.h>
6+
#include <torch/all.h>
7+
8+
// Paddle compat 的 c10/cuda/CUDAStream.h 依赖 PADDLE_WITH_CUDA 宏，
9+
// 不能在普通编译环境中直接包含。libtorch 的版本则依赖 cuda_runtime_api.h。
10+
// 两者均只在 USE_PADDLE_API=0（libtorch build）下包含。
11+
#if !USE_PADDLE_API
12+
#include <c10/cuda/CUDAStream.h>
13+
#endif
14+
15+
#include <string>
16+
#include <vector>
17+
18+
#include "../src/file_manager.h"
19+
20+
extern paddle_api_test::ThreadSafeParam g_custom_param;
21+
22+
namespace at {
23+
namespace test {
24+
25+
using paddle_api_test::FileManerger;
26+
using paddle_api_test::ThreadSafeParam;
27+
28+
class RecordStreamTest : public ::testing::Test {
29+
protected:
30+
void SetUp() override { cpu_tensor = at::zeros({2, 3}, at::kFloat); }
31+
at::Tensor cpu_tensor;
32+
};
33+
34+
// 返回一个指向 device 0 默认 CUDA stream 的 at::Stream
35+
// libtorch: 通过 CUDAStream（有 operator Stream() 隐式转换）
36+
// Paddle compat: CUDAStream 未提供隐式转换，手动以 DEFAULT stream id 0 构造
37+
static at::Stream get_default_cuda_stream() {
38+
#if USE_PADDLE_API
39+
// Paddle: 直接构造（id=0 = CUDA null/default stream）
40+
return at::Stream(at::Stream::DEFAULT, c10::Device(c10::DeviceType::CUDA, 0));
41+
#else
42+
// libtorch: CUDAStream 隐式转换为 at::Stream
43+
return c10::cuda::getCurrentCUDAStream(0);
44+
#endif
45+
}
46+
47+
// --- 基础功能测试：CUDA tensor + CUDA stream ---
48+
49+
// kFloat, shape {2,3} (small)
50+
TEST_F(RecordStreamTest, CudaFloat2x3) {
51+
auto file_name = g_custom_param.get();
52+
FileManerger file(file_name);
53+
file.createFile();
54+
file << "CudaFloat2x3 ";
55+
if (!torch::cuda::is_available()) {
56+
file << "no_cuda";
57+
file << "\n";
58+
file.saveFile();
59+
return;
60+
}
61+
try {
62+
at::Tensor t = cpu_tensor.cuda();
63+
at::Stream stream = get_default_cuda_stream();
64+
t.record_stream(stream);
65+
file << "1";
66+
} catch (const std::exception& e) {
67+
file << "exception";
68+
}
69+
file << "\n";
70+
file.saveFile();
71+
}
72+
73+
// kDouble, shape {4} (small, different dtype)
74+
TEST_F(RecordStreamTest, CudaDouble4) {
75+
auto file_name = g_custom_param.get();
76+
FileManerger file(file_name);
77+
file.openAppend();
78+
file << "CudaDouble4 ";
79+
if (!torch::cuda::is_available()) {
80+
file << "no_cuda";
81+
file << "\n";
82+
file.saveFile();
83+
return;
84+
}
85+
try {
86+
at::Tensor t = at::zeros({4}, at::kDouble).cuda();
87+
at::Stream stream = get_default_cuda_stream();
88+
t.record_stream(stream);
89+
file << "1";
90+
} catch (const std::exception& e) {
91+
file << "exception";
92+
}
93+
file << "\n";
94+
file.saveFile();
95+
}
96+
97+
// kInt, shape {100,100} (large, >= 10000 elements)
98+
TEST_F(RecordStreamTest, CudaInt100x100) {
99+
auto file_name = g_custom_param.get();
100+
FileManerger file(file_name);
101+
file.openAppend();
102+
file << "CudaInt100x100 ";
103+
if (!torch::cuda::is_available()) {
104+
file << "no_cuda";
105+
file << "\n";
106+
file.saveFile();
107+
return;
108+
}
109+
try {
110+
at::Tensor t = at::zeros({100, 100}, at::kInt).cuda();
111+
at::Stream stream = get_default_cuda_stream();
112+
t.record_stream(stream);
113+
file << "1";
114+
} catch (const std::exception& e) {
115+
file << "exception";
116+
}
117+
file << "\n";
118+
file.saveFile();
119+
}
120+
121+
// kLong, shape {} (0-d scalar tensor)
122+
TEST_F(RecordStreamTest, CudaLongScalar) {
123+
auto file_name = g_custom_param.get();
124+
FileManerger file(file_name);
125+
file.openAppend();
126+
file << "CudaLongScalar ";
127+
if (!torch::cuda::is_available()) {
128+
file << "no_cuda";
129+
file << "\n";
130+
file.saveFile();
131+
return;
132+
}
133+
try {
134+
at::Tensor t = at::zeros({}, at::kLong).cuda();
135+
at::Stream stream = get_default_cuda_stream();
136+
t.record_stream(stream);
137+
file << "1";
138+
} catch (const std::exception& e) {
139+
file << "exception";
140+
}
141+
file << "\n";
142+
file.saveFile();
143+
}
144+
145+
// kFloat, shape {0} (空 tensor，边界 shape)
146+
TEST_F(RecordStreamTest, CudaEmptyShape) {
147+
auto file_name = g_custom_param.get();
148+
FileManerger file(file_name);
149+
file.openAppend();
150+
file << "CudaEmptyShape ";
151+
if (!torch::cuda::is_available()) {
152+
file << "no_cuda";
153+
file << "\n";
154+
file.saveFile();
155+
return;
156+
}
157+
try {
158+
at::Tensor t = at::zeros({0}, at::kFloat).cuda();
159+
at::Stream stream = get_default_cuda_stream();
160+
t.record_stream(stream);
161+
file << "1";
162+
} catch (const std::exception& e) {
163+
file << "exception";
164+
}
165+
file << "\n";
166+
file.saveFile();
167+
}
168+
169+
// kFloat, shape {1,1,1} (全一维度，边界 shape)
170+
TEST_F(RecordStreamTest, CudaAllOnes) {
171+
auto file_name = g_custom_param.get();
172+
FileManerger file(file_name);
173+
file.openAppend();
174+
file << "CudaAllOnes ";
175+
if (!torch::cuda::is_available()) {
176+
file << "no_cuda";
177+
file << "\n";
178+
file.saveFile();
179+
return;
180+
}
181+
try {
182+
at::Tensor t = at::zeros({1, 1, 1}, at::kFloat).cuda();
183+
at::Stream stream = get_default_cuda_stream();
184+
t.record_stream(stream);
185+
file << "1";
186+
} catch (const std::exception& e) {
187+
file << "exception";
188+
}
189+
file << "\n";
190+
file.saveFile();
191+
}
192+
193+
// kFloat, 非连续 tensor（经 transpose）
194+
TEST_F(RecordStreamTest, CudaNonContiguous) {
195+
auto file_name = g_custom_param.get();
196+
FileManerger file(file_name);
197+
file.openAppend();
198+
file << "CudaNonContiguous ";
199+
if (!torch::cuda::is_available()) {
200+
file << "no_cuda";
201+
file << "\n";
202+
file.saveFile();
203+
return;
204+
}
205+
try {
206+
at::Tensor base = at::zeros({3, 4}, at::kFloat).cuda();
207+
at::Tensor t = base.transpose(0, 1); // 非连续
208+
at::Stream stream = get_default_cuda_stream();
209+
t.record_stream(stream);
210+
file << "1";
211+
} catch (const std::exception& e) {
212+
file << "exception";
213+
}
214+
file << "\n";
215+
file.saveFile();
216+
}
217+
218+
// --- 异常路径：CPU tensor + CUDA stream（如有 CUDA） ---
219+
// record_stream 在两个框架下对 CPU tensor 的处理行为
220+
TEST_F(RecordStreamTest, CpuTensorCudaStream) {
221+
auto file_name = g_custom_param.get();
222+
FileManerger file(file_name);
223+
file.openAppend();
224+
file << "CpuTensorCudaStream ";
225+
if (!torch::cuda::is_available()) {
226+
file << "no_cuda";
227+
file << "\n";
228+
file.saveFile();
229+
return;
230+
}
231+
try {
232+
at::Stream stream = get_default_cuda_stream();
233+
cpu_tensor.record_stream(stream);
234+
file << "1";
235+
} catch (const std::exception& e) {
236+
file << "exception";
237+
}
238+
file << "\n";
239+
file.saveFile();
240+
}
241+
242+
// --- 异常路径：CPU tensor + CPU stream（无 CUDA 依赖） ---
243+
// record_stream 是 CUDA-only API，CPU stream 应触发异常
244+
TEST_F(RecordStreamTest, CpuTensorCpuStream) {
245+
auto file_name = g_custom_param.get();
246+
FileManerger file(file_name);
247+
file.openAppend();
248+
file << "CpuTensorCpuStream ";
249+
c10::Stream stream(c10::Stream::DEFAULT,
250+
c10::Device(c10::DeviceType::CPU, 0));
251+
try {
252+
cpu_tensor.record_stream(stream);
253+
file << "1";
254+
} catch (const std::exception& e) {
255+
file << "exception";
256+
}
257+
file << "\n";
258+
file.saveFile();
259+
}
260+
261+
} // namespace test
262+
} // namespace at

0 commit comments

Comments
 (0)