1 #include <gtest/gtest.h> 4 #include <cuda_runtime.h> 6 #include <ATen/cuda/detail/TensorInfo.cuh> 7 #include <ATen/cuda/CUDAContext.h> 8 #define ASSERT_EQ_CUDA(X, Y) \ 10 bool _isEQ = X == Y; \ 20 TEST(ApplyTest, Contiguous2D) {
21 if (!at::cuda::is_available())
return;
23 int strides[] = {4, 1};
24 ::at::cuda::detail::TensorInfo<void, int> ti{
nullptr, 2, sizes, strides};
26 ASSERT_EQ_CUDA(ti.dims, 1);
27 ASSERT_EQ_CUDA(ti.sizes[0], (4 * 4));
32 TEST(ApplyTest, Contiguous3D) {
33 if (!at::cuda::is_available())
return;
34 int sizes[] = {6, 3, 7};
35 int strides[] = {3 * 7, 7, 1};
36 ::at::cuda::detail::TensorInfo<void, int> ti{
nullptr, 3, sizes, strides};
38 ASSERT_EQ_CUDA(ti.dims, 1);
39 ASSERT_EQ_CUDA(ti.sizes[0], (6 * 3 * 7));
43 TEST(ApplyTest, PartialCollapse3D) {
44 if (!at::cuda::is_available())
return;
45 int sizes[] = {4, 3, 2};
46 int strides[] = {3 * 3, 3, 1};
47 ::at::cuda::detail::TensorInfo<void, int> ti{
nullptr, 3, sizes, strides};
49 ASSERT_EQ_CUDA(ti.dims, 2);
50 ASSERT_EQ_CUDA(ti.sizes[0], (4 * 3));
51 ASSERT_EQ_CUDA(ti.sizes[1], 2);
55 TEST(ApplyTest, StridedCollapse2D) {
56 if (!at::cuda::is_available())
return;
58 int strides[] = {2 * 2, 2};
59 ::at::cuda::detail::TensorInfo<void, int> ti{
nullptr, 2, sizes, strides};
61 ASSERT_EQ_CUDA(ti.dims, 1);
62 ASSERT_EQ_CUDA(ti.sizes[0], (3 * 2));
63 ASSERT_EQ_CUDA(ti.strides[0], 2);
67 TEST(ApplyTest, PartialStridedCollapse4D) {
68 if (!at::cuda::is_available())
return;
69 int sizes[] = {3, 6, 5, 2};
70 int strides[] = {6 * 22, 22, 2 * 2, 2};
71 ::at::cuda::detail::TensorInfo<void, int> ti{
nullptr, 4, sizes, strides};
73 ASSERT_EQ_CUDA(ti.dims, 2);
74 ASSERT_EQ_CUDA(ti.sizes[0], (3 * 6));
75 ASSERT_EQ_CUDA(ti.strides[0], 22);
76 ASSERT_EQ_CUDA(ti.sizes[1], (5 * 2));
77 ASSERT_EQ_CUDA(ti.strides[1], 2);
81 TEST(ApplyTest, CollapsesZerosAndOnes) {
82 if (!at::cuda::is_available())
return;
83 int sizes[] = {1, 10, 1, 5, 4};
84 int strides[] = {4, 0, 16, 0, 1};
85 ::at::cuda::detail::TensorInfo<void, int> ti{
nullptr, 5, sizes, strides};
87 ASSERT_EQ_CUDA(ti.dims, 2);
88 ASSERT_EQ_CUDA(ti.sizes[0], (10 * 5));
89 ASSERT_EQ_CUDA(ti.strides[0], 0);
90 ASSERT_EQ_CUDA(ti.sizes[1], 4);
91 ASSERT_EQ_CUDA(ti.strides[1], 1);
95 TEST(ApplyTest, CollapseToPointTensor) {
96 if (!at::cuda::is_available())
return;
97 int sizes[] = {1, 1, 1};
98 int strides[] = {17, 12, 3};
99 ::at::cuda::detail::TensorInfo<void, int> ti{
nullptr, 3, sizes, strides};
100 ASSERT_EQ_CUDA(ti.collapseDims(), 0);
101 ASSERT_EQ_CUDA(ti.dims, 1);
102 ASSERT_EQ_CUDA(ti.sizes[0], 1);
103 ASSERT_EQ_CUDA(ti.strides[0], 1);
107 TEST(ApplyTest, ExcludingInContiguous4D) {
108 if (!at::cuda::is_available())
return;
109 int sizes[] = {3, 6, 5, 2};
110 int strides[] = {6 * 22, 22, 2 * 2, 2};
111 ::at::cuda::detail::TensorInfo<void, int> ti{
nullptr, 4, sizes, strides};
112 ASSERT_EQ_CUDA(ti.collapseDims(1), 1);
113 ASSERT_EQ_CUDA(ti.dims, 3);
114 ASSERT_EQ_CUDA(ti.sizes[0], 3);
115 ASSERT_EQ_CUDA(ti.strides[0], (6 * 22));
116 ASSERT_EQ_CUDA(ti.sizes[1], 6);
117 ASSERT_EQ_CUDA(ti.strides[1], 22);
118 ASSERT_EQ_CUDA(ti.sizes[2], (5 * 2));
119 ASSERT_EQ_CUDA(ti.strides[2], 2);
123 TEST(ApplyTest, RovingExclusion) {
124 if (!at::cuda::is_available())
return;
125 int sizes[] = {3, 6, 5, 2};
126 int strides[] = {6 * 22, 22, 2 * 2, 2};
127 ::at::cuda::detail::TensorInfo<void, int> ti{
nullptr, 4, sizes, strides};
128 ASSERT_EQ_CUDA(ti.collapseDims(2), 1);
129 ASSERT_EQ_CUDA(ti.dims, 3);
130 ASSERT_EQ_CUDA(ti.sizes[0], (3 * 6));
131 ASSERT_EQ_CUDA(ti.strides[0], 22);
132 ASSERT_EQ_CUDA(ti.sizes[1], 5);
133 ASSERT_EQ_CUDA(ti.strides[1], 4);
134 ASSERT_EQ_CUDA(ti.sizes[2], 2);
135 ASSERT_EQ_CUDA(ti.strides[2], 2);
139 TEST(ApplyTest, InvalidExclusion) {
140 if (!at::cuda::is_available())
return;
141 int sizes[] = {1, 1, 1};
142 int strides[] = {17, 12, 3};
143 ::at::cuda::detail::TensorInfo<void, int> ti{
nullptr, 3, sizes, strides};
144 ASSERT_ANY_THROW(ti.collapseDims(5));