Caffe2 - C++ API
A deep learning, cross platform ML framework
tensor_cuda.cpp
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 
86 TEST(TensorTest, ToDoesNotCopyWhenOptionsAreAllTheSame_CUDA) {
87  auto tensor = at::empty({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<float>(), tensor.data<float>());
90  hopefully_not_copy = tensor.to(at::kFloat);
91  ASSERT_EQ(hopefully_not_copy.data<float>(), tensor.data<float>());
92  hopefully_not_copy = tensor.to("cuda");
93  ASSERT_EQ(hopefully_not_copy.data<float>(), tensor.data<float>());
94  hopefully_not_copy = tensor.to(at::TensorOptions("cuda"));
95  ASSERT_EQ(hopefully_not_copy.data<float>(), tensor.data<float>());
96  hopefully_not_copy = tensor.to(at::TensorOptions(at::kFloat));
97  ASSERT_EQ(hopefully_not_copy.data<float>(), tensor.data<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 }
C10_NODISCARD TensorOptions device(c10::optional< Device > device) const noexcept
Return a copy of TensorOptions with device set to the given one, or cleared if device is nullopt...
Represents a a compute device on which a tensor is located.
Definition: Device.h:30
C10_NODISCARD TensorOptions dtype(c10::optional< caffe2::TypeMeta > dtype) const noexcept
Return a copy of TensorOptions with dtype set to the given one.