xref: /aosp_15_r20/external/tensorflow/tensorflow/core/kernels/tensor_cord_test.cc (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7     http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 
16 #include "tensorflow/core/kernels/tensor_cord.h"
17 
18 #include "tensorflow/core/framework/tensor.h"
19 #include "tensorflow/core/framework/tensor_reference.h"
20 #include "tensorflow/core/framework/tensor_util.h"
21 #include "tensorflow/core/framework/variant.h"
22 #include "tensorflow/core/framework/variant_encode_decode.h"
23 #include "tensorflow/core/platform/cord.h"
24 #include "tensorflow/core/platform/test.h"
25 #include "tensorflow/core/platform/test_benchmark.h"
26 
27 namespace tensorflow {
28 
29 namespace {
30 
DoNothingReleaser(void *)31 void DoNothingReleaser(void*) {}
32 
TEST(TensorCordTest,Empty)33 TEST(TensorCordTest, Empty) {
34   TensorCord tc;
35   EXPECT_EQ(tc.size(), 0);
36   EXPECT_EQ(tc.chunk_begin(), tc.chunk_end());
37   auto chunks = tc.Chunks();
38   EXPECT_EQ(chunks.begin(), chunks.end());
39 }
40 
TEST(TensorCordTest,ViewOfValue)41 TEST(TensorCordTest, ViewOfValue) {
42   TensorCord tc("abc", &DoNothingReleaser, nullptr);
43   EXPECT_EQ(*tc.chunk_begin(), "abc");
44   auto it = tc.chunk_begin();
45   EXPECT_EQ(*it, "abc");
46   ++it;
47   EXPECT_EQ(it, tc.chunk_end());
48 }
49 
TEST(TensorCordTest,Chunks)50 TEST(TensorCordTest, Chunks) {
51   TensorCord tc("abc", &DoNothingReleaser, nullptr);
52   int counter = 0;
53   for (auto string_piece : tc.Chunks()) {
54     EXPECT_EQ(string_piece, "abc");
55     ++counter;
56   }
57   EXPECT_EQ(counter, 1);
58 }
59 
60 // This function  takes an arg-less std::function that may have a closure, and
61 // creates a std::function with no closure: one that can be cast
62 // directly to a (*)(void*) function pointer that takes the original function.
63 // Use it this way:
64 //
65 //  void callback_with_arg((* fn)(void*), void* arg) { fn(arg); }
66 //
67 //  auto fn = [&]() { ... }
68 //  auto thunk = CreateThunkFor(fn);
69 //  callback_with_arg(thunk, &fn);
70 //
71 // Idea from:
72 //   http://bannalia.blogspot.com/2016/07/passing-capturing-c-lambda-functions-as.html
73 template <typename T>
CreateThunkFor(const T & fn)74 CordRepReleaser CreateThunkFor(const T& fn) {
75   return [](void* ptr) { (*static_cast<T*>(ptr))(); };
76 }
77 
TEST(TensorCordTest,Copy)78 TEST(TensorCordTest, Copy) {
79   int cleaned = 0;
80   auto cleaner = [&cleaned]() { ++cleaned; };
81   auto thunk = CreateThunkFor(cleaner);
82   TensorCord tc_copy;
83   string a = "abc";
84   {
85     TensorCord tc(a, thunk, &cleaner);
86     tc_copy = tc;
87   }
88   auto it = tc_copy.chunk_begin();
89   EXPECT_EQ(*it, "abc");
90   ++it;
91   EXPECT_EQ(it, tc_copy.chunk_end());
92   EXPECT_EQ(cleaned, 0);
93   tc_copy = TensorCord();
94   EXPECT_EQ(cleaned, 1);
95 }
96 
TEST(TensorCordTest,AppendCord)97 TEST(TensorCordTest, AppendCord) {
98   int cleaned_0 = 0;
99   int cleaned_1 = 0;
100   auto cleaner_0 = [&cleaned_0]() { ++cleaned_0; };
101   auto cleaner_1 = [&cleaned_1]() { ++cleaned_1; };
102   auto thunk_0 = CreateThunkFor(cleaner_0);
103   auto thunk_1 = CreateThunkFor(cleaner_1);
104   TensorCord tc_0("abc", thunk_0, &cleaner_0);
105   TensorCord tc_1("cba", thunk_1, &cleaner_1);
106   tc_0.Append(tc_1);
107   EXPECT_EQ(string(tc_0), "abccba");
108   auto it = tc_0.chunk_begin();
109   EXPECT_EQ(*it, "abc");
110   ++it;
111   EXPECT_EQ(*it, "cba");
112   ++it;
113   EXPECT_EQ(it, tc_0.chunk_end());
114   tc_1 = TensorCord();
115   EXPECT_EQ(cleaned_0, 0);
116   EXPECT_EQ(cleaned_1, 0);
117   tc_0 = TensorCord();
118   EXPECT_EQ(cleaned_0, 1);
119   EXPECT_EQ(cleaned_1, 1);
120 }
121 
TEST(TensorCordTest,AppendView)122 TEST(TensorCordTest, AppendView) {
123   int cleaned_0 = 0;
124   int cleaned_1 = 0;
125   auto cleaner_0 = [&cleaned_0]() { ++cleaned_0; };
126   auto cleaner_1 = [&cleaned_1]() { ++cleaned_1; };
127   auto thunk_0 = CreateThunkFor(cleaner_0);
128   auto thunk_1 = CreateThunkFor(cleaner_1);
129   TensorCord tc_0("abc", thunk_0, &cleaner_0);
130   tc_0.Append("cba", thunk_1, &cleaner_1);
131   EXPECT_EQ(string(tc_0), "abccba");
132   auto it = tc_0.chunk_begin();
133   EXPECT_EQ(*it, "abc");
134   ++it;
135   EXPECT_EQ(*it, "cba");
136   ++it;
137   EXPECT_EQ(it, tc_0.chunk_end());
138   EXPECT_EQ(cleaned_0, 0);
139   EXPECT_EQ(cleaned_1, 0);
140   tc_0 = TensorCord();
141   EXPECT_EQ(cleaned_0, 1);
142   EXPECT_EQ(cleaned_1, 1);
143 }
144 
TEST(TensorCordTest,Move)145 TEST(TensorCordTest, Move) {
146   int cleaned = 0;
147   auto cleaner = [&cleaned]() { ++cleaned; };
148   auto thunk = CreateThunkFor(cleaner);
149   TensorCord tc_copy;
150   string a = "abc";
151   {
152     TensorCord tc(a, thunk, &cleaner);
153     tc_copy = std::move(tc);
154   }
155   EXPECT_EQ(tc_copy.size(), 3);
156   auto it = tc_copy.chunk_begin();
157   EXPECT_EQ(*it, "abc");
158   ++it;
159   EXPECT_EQ(it, tc_copy.chunk_end());
160   EXPECT_EQ(cleaned, 0);
161   tc_copy = TensorCord();
162   EXPECT_EQ(tc_copy.size(), 0);
163   EXPECT_EQ(cleaned, 1);
164 }
165 
TEST(TensorCordTest,CopyConstructor)166 TEST(TensorCordTest, CopyConstructor) {
167   int cleaned = 0;
168   auto cleaner = [&cleaned]() { ++cleaned; };
169   auto thunk = CreateThunkFor(cleaner);
170   string a = "abc";
171   TensorCord tc(a, thunk, &cleaner);
172   TensorCord tc_copy(tc);
173   EXPECT_EQ(tc.size(), 3);
174   EXPECT_EQ(tc_copy.size(), 3);
175   auto it = tc_copy.chunk_begin();
176   EXPECT_EQ(*it, "abc");
177   ++it;
178   EXPECT_EQ(it, tc_copy.chunk_end());
179   EXPECT_EQ(cleaned, 0);
180   tc = TensorCord();
181   EXPECT_EQ(cleaned, 0);
182   tc_copy = TensorCord();
183   EXPECT_EQ(cleaned, 1);
184 }
185 
TEST(TensorCordTest,MoveConstructor)186 TEST(TensorCordTest, MoveConstructor) {
187   int cleaned = 0;
188   auto cleaner = [&cleaned]() { ++cleaned; };
189   auto thunk = CreateThunkFor(cleaner);
190   string a = "abc";
191   TensorCord tc(a, thunk, &cleaner);
192   TensorCord tc_copy(std::move(tc));
193   EXPECT_EQ(tc_copy.size(), 3);
194   auto it = tc_copy.chunk_begin();
195   EXPECT_EQ(*it, "abc");
196   ++it;
197   EXPECT_EQ(it, tc_copy.chunk_end());
198   EXPECT_EQ(cleaned, 0);
199   tc_copy = TensorCord();
200   EXPECT_EQ(cleaned, 1);
201 }
202 
203 #ifdef PLATFORM_GOOGLE
204 
TensorCopyFromTensorBenchmark(benchmark::State & state,int num_elem,int string_size)205 void TensorCopyFromTensorBenchmark(benchmark::State& state, int num_elem,
206                                    int string_size) {
207   Tensor strings(DT_STRING, {num_elem});
208   auto t = strings.flat<tstring>();
209   for (int i = 0; i < num_elem; ++i) {
210     t(i).insert(0, string_size, 'a');
211   }
212   for (auto _ : state) {
213     benchmark::DoNotOptimize(tensor::DeepCopy(strings));
214   }
215 }
216 
TensorCordFromTensorBenchmark(benchmark::State & state,int num_elem,int string_size)217 void TensorCordFromTensorBenchmark(benchmark::State& state, int num_elem,
218                                    int string_size) {
219   Tensor strings(DT_STRING, {num_elem});
220   auto t = strings.flat<tstring>();
221   for (int i = 0; i < num_elem; ++i) {
222     t(i).insert(0, string_size, 'a');
223   }
224   for (auto _ : state) {
225     Tensor copy(DT_VARIANT, {num_elem});
226     auto t_copy = copy.flat<Variant>();
227     for (int i = 0; i < num_elem; ++i) {
228       t_copy(i) = TensorCord(t(i), &strings);
229     }
230   }
231 }
232 
CordReleaser(void * cord_ptr)233 void CordReleaser(void* cord_ptr) { delete static_cast<absl::Cord*>(cord_ptr); }
234 
TensorCordFromAbslCordBenchmark(benchmark::State & state,int num_elem,int string_size)235 void TensorCordFromAbslCordBenchmark(benchmark::State& state, int num_elem,
236                                      int string_size) {
237   std::vector<absl::Cord> cords(num_elem);
238   for (int i = 0; i < num_elem; ++i) {
239     string s(string_size, 'a');
240     cords[i] = s;
241   }
242 
243   for (auto _ : state) {
244     Tensor copy(DT_VARIANT, {num_elem});
245     auto t_copy = copy.flat<Variant>();
246     for (int i = 0; i < num_elem; ++i) {
247       auto my_cord = new absl::Cord(cords[i]);
248       t_copy(i) = TensorCord(*my_cord->chunk_begin(), CordReleaser, my_cord);
249     }
250   }
251 }
252 
253 #define CreateBM(NUM_ELEM, STRING_SIZE)                                        \
254   void BM_TensorCopyFromTensor_NumElem_##NUM_ELEM##_StringSize_##STRING_SIZE(  \
255       benchmark::State& state) {                                               \
256     TensorCopyFromTensorBenchmark(state, NUM_ELEM, STRING_SIZE);               \
257   }                                                                            \
258   BENCHMARK(                                                                   \
259       BM_TensorCopyFromTensor_NumElem_##NUM_ELEM##_StringSize_##STRING_SIZE);  \
260   void BM_TensorCordFromTensor_NumElem_##NUM_ELEM##_StringSize_##STRING_SIZE(  \
261       benchmark::State& state) {                                               \
262     TensorCordFromTensorBenchmark(state, NUM_ELEM, STRING_SIZE);               \
263   }                                                                            \
264   BENCHMARK(                                                                   \
265       BM_TensorCordFromTensor_NumElem_##NUM_ELEM##_StringSize_##STRING_SIZE);  \
266   void                                                                         \
267       BM_TensorCordFromAbslCord_NumElem_##NUM_ELEM##_StringSize_##STRING_SIZE( \
268           benchmark::State& state) {                                           \
269     TensorCordFromAbslCordBenchmark(state, NUM_ELEM, STRING_SIZE);             \
270   }                                                                            \
271   BENCHMARK(                                                                   \
272       BM_TensorCordFromAbslCord_NumElem_##NUM_ELEM##_StringSize_##STRING_SIZE);
273 
274 #define CreateStringBMs(NUM_ELEM)           \
275   CreateBM(NUM_ELEM, /*STRING_SIZE=*/16);   \
276   CreateBM(NUM_ELEM, /*STRING_SIZE=*/32);   \
277   CreateBM(NUM_ELEM, /*STRING_SIZE=*/128);  \
278   CreateBM(NUM_ELEM, /*STRING_SIZE=*/1024); \
279   CreateBM(NUM_ELEM, /*STRING_SIZE=*/4096);
280 
281 CreateStringBMs(/*NUM_ELEM=*/1);
282 CreateStringBMs(/*NUM_ELEM=*/16);
283 CreateStringBMs(/*NUM_ELEM=*/32);
284 CreateStringBMs(/*NUM_ELEM=*/64);
285 CreateStringBMs(/*NUM_ELEM=*/128);
286 
287 #endif  // PLATFORM_GOOGLE
288 
289 }  // namespace
290 
291 }  // namespace tensorflow
292