1 | #include <gtest/gtest.h> |
2 | |
3 | #include <ATen/ATen.h> |
4 | |
5 | #include <cmath> |
6 | |
7 | #define REQUIRE_TENSOR_OPTIONS(device_, index_, type_, layout_) \ |
8 | ASSERT_TRUE( \ |
9 | tensor.device().type() == at::Device((device_), (index_)).type()); \ |
10 | ASSERT_TRUE( \ |
11 | tensor.device().index() == at::Device((device_), (index_)).index()); \ |
12 | ASSERT_EQ(tensor.dtype(), (type_)); \ |
13 | ASSERT_TRUE(tensor.layout() == (layout_)) |
14 | |
15 | TEST(TensorTest, AllocatesTensorOnTheCorrectDevice_MultiCUDA) { |
16 | auto tensor = at::tensor({1, 2, 3}, at::device({at::kCUDA, 1})); |
17 | ASSERT_EQ(tensor.device().type(), at::Device::Type::CUDA); |
18 | ASSERT_EQ(tensor.device().index(), 1); |
19 | } |
20 | |
21 | TEST(TensorTest, ToDevice_MultiCUDA) { |
22 | auto tensor = at::empty({3, 4}); |
23 | REQUIRE_TENSOR_OPTIONS(at::kCPU, -1, at::kFloat, at::kStrided); |
24 | |
25 | tensor = tensor.to({at::kCUDA, 1}); |
26 | REQUIRE_TENSOR_OPTIONS(at::kCUDA, 1, at::kFloat, at::kStrided); |
27 | |
28 | tensor = tensor.to({at::kCUDA, 0}); |
29 | REQUIRE_TENSOR_OPTIONS(at::kCUDA, 0, at::kFloat, at::kStrided); |
30 | |
31 | tensor = tensor.to({at::kCUDA, 1}); |
32 | REQUIRE_TENSOR_OPTIONS(at::kCUDA, 1, at::kFloat, at::kStrided); |
33 | |
34 | tensor = tensor.to(at::Device(at::kCPU)); |
35 | REQUIRE_TENSOR_OPTIONS(at::kCPU, -1, at::kFloat, at::kStrided); |
36 | |
37 | tensor = tensor.to(at::kCUDA); |
38 | REQUIRE_TENSOR_OPTIONS(at::kCUDA, 0, at::kFloat, at::kStrided); |
39 | |
40 | tensor = tensor.to(at::TensorOptions({at::kCUDA, 1})); |
41 | REQUIRE_TENSOR_OPTIONS(at::kCUDA, 1, at::kFloat, at::kStrided); |
42 | |
43 | tensor = tensor.to(at::TensorOptions({at::kCUDA, 0})); |
44 | REQUIRE_TENSOR_OPTIONS(at::kCUDA, 0, at::kFloat, at::kStrided); |
45 | |
46 | tensor = tensor.to(at::TensorOptions(at::kDouble)); |
47 | REQUIRE_TENSOR_OPTIONS(at::kCUDA, 0, at::kDouble, at::kStrided); |
48 | |
49 | tensor = tensor.to(at::TensorOptions({at::kCUDA, 1})); |
50 | REQUIRE_TENSOR_OPTIONS(at::kCUDA, 1, at::kDouble, at::kStrided); |
51 | |
52 | tensor = tensor.to(at::TensorOptions(at::kInt)); |
53 | REQUIRE_TENSOR_OPTIONS(at::kCUDA, 1, at::kInt, at::kStrided); |
54 | |
55 | tensor = tensor.to(at::TensorOptions(at::Device(at::kCPU))); |
56 | REQUIRE_TENSOR_OPTIONS(at::kCPU, -1, at::kInt, at::kStrided); |
57 | |
58 | tensor = tensor.to(at::TensorOptions(at::kCUDA)); |
59 | REQUIRE_TENSOR_OPTIONS(at::kCUDA, 0, at::kInt, at::kStrided); |
60 | } |
61 | |
62 | TEST(TensorTest, ToTensorAndTensorAttributes_MultiCUDA) { |
63 | auto tensor = at::empty({3, 4}); |
64 | REQUIRE_TENSOR_OPTIONS(at::kCPU, -1, at::kFloat, at::kStrided); |
65 | |
66 | auto other = at::empty({3, 4}, at::kFloat); |
67 | tensor = tensor.to(other); |
68 | REQUIRE_TENSOR_OPTIONS(at::kCPU, -1, at::kFloat, at::kStrided); |
69 | |
70 | other = at::empty({3, 4}, at::TensorOptions(at::kCUDA).dtype(at::kDouble)); |
71 | tensor = tensor.to(other.dtype()); |
72 | REQUIRE_TENSOR_OPTIONS(at::kCPU, -1, at::kDouble, at::kStrided); |
73 | tensor = tensor.to(other.device()); |
74 | REQUIRE_TENSOR_OPTIONS(at::kCUDA, 0, at::kDouble, at::kStrided); |
75 | |
76 | other = at::empty({3, 4}, at::TensorOptions({at::kCUDA, 1}).dtype(at::kLong)); |
77 | tensor = tensor.to(other.device(), other.dtype()); |
78 | REQUIRE_TENSOR_OPTIONS(at::kCUDA, 1, at::kLong, at::kStrided); |
79 | |
80 | other = at::empty({3, 4}, at::kFloat); |
81 | tensor = tensor.to(other.options()); |
82 | REQUIRE_TENSOR_OPTIONS(at::kCPU, -1, at::kFloat, at::kStrided); |
83 | } |
84 | |
85 | TEST(TensorTest, ToDoesNotCopyWhenOptionsAreAllTheSame_CUDA) { |
86 | auto tensor = at::empty( |
87 | {3, 4}, at::TensorOptions(at::kFloat).device(at::Device("cuda" ))); |
88 | auto hopefully_not_copy = tensor.to(tensor.options()); |
89 | ASSERT_EQ(hopefully_not_copy.data_ptr<float>(), tensor.data_ptr<float>()); |
90 | hopefully_not_copy = tensor.to(at::kFloat); |
91 | ASSERT_EQ(hopefully_not_copy.data_ptr<float>(), tensor.data_ptr<float>()); |
92 | hopefully_not_copy = tensor.to("cuda" ); |
93 | ASSERT_EQ(hopefully_not_copy.data_ptr<float>(), tensor.data_ptr<float>()); |
94 | hopefully_not_copy = tensor.to(at::TensorOptions("cuda" )); |
95 | ASSERT_EQ(hopefully_not_copy.data_ptr<float>(), tensor.data_ptr<float>()); |
96 | hopefully_not_copy = tensor.to(at::TensorOptions(at::kFloat)); |
97 | ASSERT_EQ(hopefully_not_copy.data_ptr<float>(), tensor.data_ptr<float>()); |
98 | } |
99 | |
100 | TEST(TensorTest, ToDeviceAndDtype_MultiCUDA) { |
101 | auto tensor = at::empty({3, 4}); |
102 | REQUIRE_TENSOR_OPTIONS(at::kCPU, -1, at::kFloat, at::kStrided); |
103 | |
104 | tensor = tensor.to({at::kCUDA, 1}, at::kInt); |
105 | REQUIRE_TENSOR_OPTIONS(at::kCUDA, 1, at::kInt, at::kStrided); |
106 | |
107 | tensor = tensor.to(at::TensorOptions({at::kCUDA, 0}).dtype(at::kLong)); |
108 | REQUIRE_TENSOR_OPTIONS(at::kCUDA, 0, at::kLong, at::kStrided); |
109 | |
110 | tensor = tensor.to(at::TensorOptions({at::kCUDA, 1}).dtype(at::kDouble)); |
111 | REQUIRE_TENSOR_OPTIONS(at::kCUDA, 1, at::kDouble, at::kStrided); |
112 | |
113 | tensor = tensor.to(at::kCPU, at::kInt); |
114 | REQUIRE_TENSOR_OPTIONS(at::kCPU, -1, at::kInt, at::kStrided); |
115 | } |
116 | |
117 | TEST(TensorTest, MagmaInitializesCorrectly_CUDA) { |
118 | // Any tensor will work here as long as it's invertible |
119 | // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays) |
120 | float data[] = {1, 1, 1, 0, 0, 3, 1, 2, 2, 3, 1, 0, 1, 0, 2, 1}; |
121 | auto tensor = |
122 | at::from_blob(data, {4, 4}, at::TensorOptions(at::kFloat)).cuda(); |
123 | if (at::hasMAGMA()) { |
124 | at::inverse(tensor); |
125 | } |
126 | } |
127 | |