xref: /aosp_15_r20/external/pytorch/test/cpp/api/tensor_options.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #include <gtest/gtest.h>
2 #include <test/cpp/api/support.h>
3 
4 #include <torch/torch.h>
5 
6 #include <string>
7 #include <vector>
8 
9 using namespace at;
10 using namespace torch::test;
11 
12 // A macro so we don't lose location information when an assertion fails.
13 #define REQUIRE_OPTIONS(device_, index_, type_, layout_)                  \
14   ASSERT_EQ(options.device().type(), Device((device_), (index_)).type()); \
15   ASSERT_TRUE(                                                            \
16       options.device().index() == Device((device_), (index_)).index());   \
17   ASSERT_EQ(options.dtype(), (type_));                                    \
18   ASSERT_TRUE(options.layout() == (layout_))
19 
20 #define REQUIRE_TENSOR_OPTIONS(device_, index_, type_, layout_)            \
21   ASSERT_EQ(tensor.device().type(), Device((device_), (index_)).type());   \
22   ASSERT_EQ(tensor.device().index(), Device((device_), (index_)).index()); \
23   ASSERT_EQ(tensor.scalar_type(), (type_));                                \
24   ASSERT_TRUE(tensor.options().layout() == (layout_))
25 
TEST(TensorOptionsTest,DefaultsToTheRightValues)26 TEST(TensorOptionsTest, DefaultsToTheRightValues) {
27   TensorOptions options;
28   REQUIRE_OPTIONS(kCPU, -1, kFloat, kStrided);
29 }
30 
TEST(TensorOptionsTest,UtilityFunctionsReturnTheRightTensorOptions)31 TEST(TensorOptionsTest, UtilityFunctionsReturnTheRightTensorOptions) {
32   auto options = dtype(kInt);
33   REQUIRE_OPTIONS(kCPU, -1, kInt, kStrided);
34 
35   options = layout(kSparse);
36   REQUIRE_OPTIONS(kCPU, -1, kFloat, kSparse);
37 
38   options = device({kCUDA, 1});
39   REQUIRE_OPTIONS(kCUDA, 1, kFloat, kStrided);
40 
41   options = device_index(1);
42   REQUIRE_OPTIONS(kCUDA, 1, kFloat, kStrided);
43 
44   options = dtype(kByte).layout(kSparse).device(kCUDA, 2).device_index(3);
45   REQUIRE_OPTIONS(kCUDA, 3, kByte, kSparse);
46 }
47 
TEST(TensorOptionsTest,ConstructsWellFromCPUTypes)48 TEST(TensorOptionsTest, ConstructsWellFromCPUTypes) {
49   TensorOptions options;
50   REQUIRE_OPTIONS(kCPU, -1, kFloat, kStrided);
51 
52   options = TensorOptions({kCPU, 0});
53   REQUIRE_OPTIONS(kCPU, 0, kFloat, kStrided);
54 
55   options = TensorOptions("cpu:0");
56   REQUIRE_OPTIONS(kCPU, 0, kFloat, kStrided);
57 
58   options = TensorOptions(kInt);
59   REQUIRE_OPTIONS(kCPU, -1, kInt, kStrided);
60 
61   options =
62       TensorOptions(getDeprecatedTypeProperties(Backend::SparseCPU, kFloat));
63   REQUIRE_OPTIONS(kCPU, -1, kFloat, kSparse);
64 
65   options =
66       TensorOptions(getDeprecatedTypeProperties(Backend::SparseCPU, kByte));
67   REQUIRE_OPTIONS(kCPU, -1, kByte, kSparse);
68 }
69 
TEST(TensorOptionsTest,ConstructsWellFromCPUTensors)70 TEST(TensorOptionsTest, ConstructsWellFromCPUTensors) {
71   auto options = empty(5, kDouble).options();
72   REQUIRE_OPTIONS(kCPU, -1, kDouble, kStrided);
73 
74   options = empty(5, getDeprecatedTypeProperties(Backend::SparseCPU, kByte))
75                 .options();
76   REQUIRE_OPTIONS(kCPU, -1, kByte, kSparse);
77 }
78 
TEST(TensorOptionsTest,ConstructsWellFromVariables)79 TEST(TensorOptionsTest, ConstructsWellFromVariables) {
80   auto options = torch::empty(5).options();
81   REQUIRE_OPTIONS(kCPU, -1, kFloat, kStrided);
82   ASSERT_FALSE(options.requires_grad());
83 
84   options = torch::empty(5, at::requires_grad()).options();
85   REQUIRE_OPTIONS(kCPU, -1, kFloat, kStrided);
86   ASSERT_FALSE(options.requires_grad());
87 }
88 
TEST(DeviceTest,ParsesCorrectlyFromString)89 TEST(DeviceTest, ParsesCorrectlyFromString) {
90   Device device("cpu:0");
91   ASSERT_EQ(device, Device(DeviceType::CPU, 0));
92 
93   device = Device("cpu");
94   ASSERT_EQ(device, Device(DeviceType::CPU));
95 
96   device = Device("cuda:123");
97   ASSERT_EQ(device, Device(DeviceType::CUDA, 123));
98 
99   device = Device("cuda");
100   ASSERT_EQ(device, Device(DeviceType::CUDA));
101 
102   device = Device("mkldnn");
103   ASSERT_EQ(device, Device(DeviceType::MKLDNN));
104 
105   device = Device("opengl");
106   ASSERT_EQ(device, Device(DeviceType::OPENGL));
107 
108   device = Device("opencl");
109   ASSERT_EQ(device, Device(DeviceType::OPENCL));
110 
111   device = Device("ideep");
112   ASSERT_EQ(device, Device(DeviceType::IDEEP));
113 
114   device = Device("hip");
115   ASSERT_EQ(device, Device(DeviceType::HIP));
116 
117   device = Device("hip:123");
118   ASSERT_EQ(device, Device(DeviceType::HIP, 123));
119 
120   std::vector<std::string> badnesses = {
121       "", "cud:1", "cuda:", "cpu::1", ":1", "3", "tpu:4", "??"};
122   for (const auto& badness : badnesses) {
123     // NOLINTNEXTLINE(cppcoreguidelines-avoid-goto,hicpp-avoid-goto)
124     ASSERT_ANY_THROW({ Device d(badness); });
125   }
126 }
127 
TEST(DefaultDtypeTest,CanSetAndGetDefaultDtype)128 TEST(DefaultDtypeTest, CanSetAndGetDefaultDtype) {
129   AutoDefaultDtypeMode dtype_mode(kFloat);
130 
131   ASSERT_EQ(at::get_default_dtype(), kFloat);
132   set_default_dtype(caffe2::TypeMeta::Make<int>());
133   ASSERT_EQ(at::get_default_dtype(), kInt);
134 }
135 
TEST(DefaultDtypeTest,NewTensorOptionsHasCorrectDefault)136 TEST(DefaultDtypeTest, NewTensorOptionsHasCorrectDefault) {
137   AutoDefaultDtypeMode dtype_mode(kFloat);
138 
139   set_default_dtype(caffe2::TypeMeta::Make<int>());
140   ASSERT_EQ(at::get_default_dtype(), kInt);
141   TensorOptions options;
142   ASSERT_EQ(options.dtype(), kInt);
143 }
144 
TEST(DefaultDtypeTest,NewTensorsHaveCorrectDefaultDtype)145 TEST(DefaultDtypeTest, NewTensorsHaveCorrectDefaultDtype) {
146   AutoDefaultDtypeMode dtype_mode(kFloat);
147   set_default_dtype(caffe2::TypeMeta::Make<int>());
148   {
149     auto tensor = torch::ones(5);
150     ASSERT_EQ(tensor.dtype(), kInt);
151   }
152   set_default_dtype(caffe2::TypeMeta::Make<double>());
153   {
154     auto tensor = torch::ones(5);
155     ASSERT_EQ(tensor.dtype(), kDouble);
156   }
157   {
158     auto tensor = torch::ones(5, kFloat);
159     ASSERT_EQ(tensor.dtype(), kFloat);
160   }
161 }
162