1 #include <gtest/gtest.h>
2
3 #include <cuda.h>
4 #include <cuda_runtime.h>
5
6 #include <ATen/cuda/detail/TensorInfo.cuh>
7 #include <ATen/cuda/CUDAContext.h>
8 #define ASSERT_EQ_CUDA(X, Y) \
9 { \
10 bool _isEQ = X == Y; \
11 ASSERT_TRUE(_isEQ); \
12 }
13 /*
14 Tests related to tensor indexing and applying operations.
15 */
16 #ifndef _WIN32
17
18 // CATCH_TEST_CASE("2D Contiguous", "Collapses a 2D contiguous tensor to 1D
19 // contiguous") {
TEST(ApplyTest,Contiguous2D)20 TEST(ApplyTest, Contiguous2D) {
21 if (!at::cuda::is_available()) return;
22 int sizes[] = {4, 4};
23 int strides[] = {4, 1};
24 ::at::cuda::detail::TensorInfo<void, int> ti{nullptr, 2, sizes, strides};
25 ti.collapseDims();
26 ASSERT_EQ_CUDA(ti.dims, 1);
27 ASSERT_EQ_CUDA(ti.sizes[0], (4 * 4));
28 }
29
30 // CATCH_TEST_CASE("3D Contiguous", "Collapses a 3D contiguous tensor to a 1D
31 // contiguous") {
TEST(ApplyTest,Contiguous3D)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};
37 ti.collapseDims();
38 ASSERT_EQ_CUDA(ti.dims, 1);
39 ASSERT_EQ_CUDA(ti.sizes[0], (6 * 3 * 7));
40 }
41 // CATCH_TEST_CASE("3D Partial Collapse", "Collapses a 3D noncontiguous tensor
42 // to a 2D tensor") {
TEST(ApplyTest,PartialCollapse3D)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};
48 ti.collapseDims();
49 ASSERT_EQ_CUDA(ti.dims, 2);
50 ASSERT_EQ_CUDA(ti.sizes[0], (4 * 3));
51 ASSERT_EQ_CUDA(ti.sizes[1], 2);
52 }
53
54 // Collapses a 2D skip contiguous tensor to a 1D skip contiguous tensor
TEST(ApplyTest,StridedCollapse2D)55 TEST(ApplyTest, StridedCollapse2D) {
56 if (!at::cuda::is_available()) return;
57 int sizes[] = {3, 2};
58 int strides[] = {2 * 2, 2};
59 ::at::cuda::detail::TensorInfo<void, int> ti{nullptr, 2, sizes, strides};
60 ti.collapseDims();
61 ASSERT_EQ_CUDA(ti.dims, 1);
62 ASSERT_EQ_CUDA(ti.sizes[0], (3 * 2));
63 ASSERT_EQ_CUDA(ti.strides[0], 2);
64 }
65
66 // Collapses a 4D tensor to a 2D tensor
TEST(ApplyTest,PartialStridedCollapse4D)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};
72 ti.collapseDims();
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);
78 }
79
80 // Collapses a 5D tensor to a 1D tensor
TEST(ApplyTest,CollapsesZerosAndOnes)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};
86 ti.collapseDims();
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);
92 }
93
94 // Collapses a 3D tensor to a point tensor
TEST(ApplyTest,CollapseToPointTensor)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);
104 }
105
106 // Collapses a 4D tensor to a 3D tensor
TEST(ApplyTest,ExcludingInContiguous4D)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);
120 }
121
122 // Collapses a 4D tensor to a 3D tensor
TEST(ApplyTest,RovingExclusion)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);
136 }
137
138 // Attempts to exclude a nonexisting dimension
TEST(ApplyTest,InvalidExclusion)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));
145 }
146 #endif
147