1 #include <gtest/gtest.h>
2
3 #include <ATen/ATen.h>
4 #include <ATen/CPUApplyUtils.h>
5 #include <ATen/test/test_assert.h>
6 #include <c10/util/irange.h>
7
8 #include <iostream>
9 using namespace std;
10 using namespace at;
11
fill_tensor(int64_t scalar,Tensor & t_)12 void fill_tensor(int64_t scalar, Tensor& t_) {
13 auto t = t_.view(-1);
14 for (const auto i : c10::irange(t.numel())) {
15 t[i] = (i + 1) * scalar;
16 }
17 }
18
19 // This test exercises all sequential applyX functions. Given a shape and two
20 // transpose dimensions we create 5 tensors (a0, ..., a4) of the given shape and
21 // transpose the dimension a with b for each tensor. Then we call the applyX
22 // function on each floating type. a4 is allocated in doubles only, whereas a0,
23 // ..., a3 are allocated in the given type. For each applyX function we once
24 // write the same type as we read (using a0, ..., aX-1) and we once write to
25 // double (using a4 as a target). We also exercise on a zero_dim and empty
26 // tensor.
test(DeprecatedTypeProperties & type,IntArrayRef shape,int64_t a=0,int64_t b=1)27 void test(DeprecatedTypeProperties& type, IntArrayRef shape, int64_t a = 0, int64_t b = 1) {
28 auto zero_dim = at::empty({}, type);
29 zero_dim.fill_(2);
30 zero_dim.exp_();
31 AT_DISPATCH_FLOATING_TYPES(zero_dim.scalar_type(), "test0", [&] {
32 ASSERT(zero_dim.const_data_ptr<scalar_t>()[0] == std::exp(2));
33 });
34
35 auto empty_t = at::empty({0}, type);
36 empty_t.fill_(3);
37 empty_t.exp_();
38
39 auto a0 = at::empty({0}, type.options());
40 auto a1 = at::empty({0}, type.options());
41 auto a2 = at::empty({0}, type.options());
42 auto a3 = at::empty({0}, type.options());
43 auto a4 = at::empty({0}, at::TensorOptions(kCPU).dtype(kDouble));
44
45 std::vector<Tensor> tensors({a0, a1, a2, a3, a4});
46 for (const auto i : c10::irange(tensors.size())) {
47 tensors[i].resize_(shape);
48 fill_tensor(i + 1, tensors[i]);
49 if (a >= 0 && b >= 0) {
50 tensors[i].transpose_(a, b);
51 }
52 }
53
54 AT_DISPATCH_FLOATING_TYPES(a0.scalar_type(), "test1", [&] {
55 CPU_tensor_apply2<scalar_t, scalar_t>(
56 a0, a1, [](scalar_t& y, const scalar_t& x) { y = x * x; });
57 CPU_tensor_apply2<double, scalar_t>(
58 a4, a1, [](double& y, scalar_t x) { y = (double)(x * x); });
59 for (const auto i : c10::irange(a0.numel())) {
60 auto target = a1.const_data_ptr<scalar_t>()[i] * a1.const_data_ptr<scalar_t>()[i];
61 ASSERT(a0.const_data_ptr<scalar_t>()[i] == target);
62 ASSERT(a4.const_data_ptr<double>()[i] == target);
63 }
64 });
65
66 AT_DISPATCH_FLOATING_TYPES(a0.scalar_type(), "test2", [&] {
67 CPU_tensor_apply3<scalar_t, scalar_t, scalar_t>(
68 a0, a1, a2, [](scalar_t& y, const scalar_t& x, const scalar_t& z) {
69 y = x * x + z;
70 });
71 CPU_tensor_apply3<double, scalar_t, scalar_t>(
72 a4, a1, a2, [](double& y, const scalar_t& x, const scalar_t& z) {
73 y = (double)(x * x + z);
74 });
75 for (const auto i : c10::irange(a0.numel())) {
76 auto target = a1.const_data_ptr<scalar_t>()[i] * a1.const_data_ptr<scalar_t>()[i];
77 target = target + a2.const_data_ptr<scalar_t>()[i];
78 ASSERT(a0.const_data_ptr<scalar_t>()[i] == target);
79 ASSERT(a4.const_data_ptr<double>()[i] == target);
80 }
81 });
82
83 AT_DISPATCH_FLOATING_TYPES(a0.scalar_type(), "test3", [&] {
84 CPU_tensor_apply4<scalar_t, scalar_t, scalar_t, scalar_t>(
85 a0,
86 a1,
87 a2,
88 a3,
89 [](scalar_t& y,
90 const scalar_t& x,
91 const scalar_t& z,
92 const scalar_t& a) { y = x * x + z * a; });
93 CPU_tensor_apply4<double, scalar_t, scalar_t, scalar_t>(
94 a4,
95 a1,
96 a2,
97 a3,
98 [](double& y, const scalar_t& x, const scalar_t& z, const scalar_t& a) {
99 y = (double)(x * x + z * a);
100 });
101 for (const auto i : c10::irange(a0.numel())) {
102 auto target = a1.const_data_ptr<scalar_t>()[i] * a1.const_data_ptr<scalar_t>()[i];
103 target = target + a2.const_data_ptr<scalar_t>()[i] * a3.const_data_ptr<scalar_t>()[i];
104 ASSERT(a0.const_data_ptr<scalar_t>()[i] == target);
105 ASSERT(a4.const_data_ptr<double>()[i] == target);
106 }
107 });
108 }
109
110 // apply utils test 2-dim small contiguous
TEST(ApplyUtilsTest,Contiguous2D)111 TEST(ApplyUtilsTest, Contiguous2D) {
112 manual_seed(123);
113 test(CPU(kDouble), {2, 1}, -1, -1);
114 }
115
116 // apply utils test 2-dim small
TEST(ApplyUtilsTest,Small2D)117 TEST(ApplyUtilsTest, Small2D) {
118 manual_seed(123);
119 test(CPU(kDouble), {2, 1});
120 }
121
122 // apply utils test 2-dim
TEST(ApplyUtilsTest,_2D)123 TEST(ApplyUtilsTest, _2D) {
124 manual_seed(123);
125 test(CPU(kDouble), {20, 10});
126 }
127
128 // apply utils test 3-dim
TEST(ApplyUtilsTest,_3D)129 TEST(ApplyUtilsTest, _3D) {
130 manual_seed(123);
131 test(CPU(kDouble), {3, 4, 2});
132 }
133
134 // apply utils test 3-dim medium
TEST(ApplyUtilsTest,Medium3D)135 TEST(ApplyUtilsTest, Medium3D) {
136 manual_seed(123);
137 test(CPU(kDouble), {3, 40, 2});
138 }
139
140 // apply utils test 10-dim
TEST(ApplyUtilsTest,_10D)141 TEST(ApplyUtilsTest, _10D) {
142 manual_seed(123);
143 test(CPU(kDouble), {3, 4, 2, 5, 2, 1, 3, 4, 2, 3});
144 }
145