cxx11_tensor_block_io.cpp (15858B)
1 // This file is part of Eigen, a lightweight C++ template library 2 // for linear algebra. 3 // 4 // This Source Code Form is subject to the terms of the Mozilla 5 // Public License v. 2.0. If a copy of the MPL was not distributed 6 // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 8 // clang-format off 9 #include "main.h" 10 #include <Eigen/CXX11/Tensor> 11 // clang-format on 12 13 // -------------------------------------------------------------------------- // 14 // A set of tests for TensorBlockIO: copying data between tensor blocks. 15 16 template <int NumDims> 17 static DSizes<Index, NumDims> RandomDims(Index min, Index max) { 18 DSizes<Index, NumDims> dims; 19 for (int i = 0; i < NumDims; ++i) { 20 dims[i] = internal::random<Index>(min, max); 21 } 22 return DSizes<Index, NumDims>(dims); 23 } 24 25 static internal::TensorBlockShapeType RandomBlockShape() { 26 return internal::random<bool>() 27 ? internal::TensorBlockShapeType::kUniformAllDims 28 : internal::TensorBlockShapeType::kSkewedInnerDims; 29 } 30 31 template <int NumDims> 32 static size_t RandomTargetBlockSize(const DSizes<Index, NumDims>& dims) { 33 return internal::random<size_t>(1, dims.TotalSize()); 34 } 35 36 template <int Layout, int NumDims> 37 static Index GetInputIndex(Index output_index, 38 const array<Index, NumDims>& output_to_input_dim_map, 39 const array<Index, NumDims>& input_strides, 40 const array<Index, NumDims>& output_strides) { 41 int input_index = 0; 42 if (Layout == ColMajor) { 43 for (int i = NumDims - 1; i > 0; --i) { 44 const Index idx = output_index / output_strides[i]; 45 input_index += idx * input_strides[output_to_input_dim_map[i]]; 46 output_index -= idx * output_strides[i]; 47 } 48 return input_index + 49 output_index * input_strides[output_to_input_dim_map[0]]; 50 } else { 51 for (int i = 0; i < NumDims - 1; ++i) { 52 const Index idx = output_index / output_strides[i]; 53 input_index += idx * input_strides[output_to_input_dim_map[i]]; 54 output_index -= idx * output_strides[i]; 55 } 56 return input_index + 57 output_index * input_strides[output_to_input_dim_map[NumDims - 1]]; 58 } 59 } 60 61 template <typename T, int NumDims, int Layout> 62 static void test_block_io_copy_data_from_source_to_target() { 63 using TensorBlockIO = internal::TensorBlockIO<T, Index, NumDims, Layout>; 64 using IODst = typename TensorBlockIO::Dst; 65 using IOSrc = typename TensorBlockIO::Src; 66 67 // Generate a random input Tensor. 68 DSizes<Index, NumDims> dims = RandomDims<NumDims>(1, 30); 69 Tensor<T, NumDims, Layout> input(dims); 70 input.setRandom(); 71 72 // Write data to an output Tensor. 73 Tensor<T, NumDims, Layout> output(dims); 74 75 // Construct a tensor block mapper. 76 using TensorBlockMapper = 77 internal::TensorBlockMapper<NumDims, Layout, Index>; 78 TensorBlockMapper block_mapper( 79 dims, {RandomBlockShape(), RandomTargetBlockSize(dims), {0, 0, 0}}); 80 81 // We will copy data from input to output through this buffer. 82 Tensor<T, NumDims, Layout> block(block_mapper.blockDimensions()); 83 84 // Precompute strides for TensorBlockIO::Copy. 85 auto input_strides = internal::strides<Layout>(dims); 86 auto output_strides = internal::strides<Layout>(dims); 87 88 const T* input_data = input.data(); 89 T* output_data = output.data(); 90 T* block_data = block.data(); 91 92 for (int i = 0; i < block_mapper.blockCount(); ++i) { 93 auto desc = block_mapper.blockDescriptor(i); 94 95 auto blk_dims = desc.dimensions(); 96 auto blk_strides = internal::strides<Layout>(blk_dims); 97 98 { 99 // Read from input into a block buffer. 100 IODst dst(blk_dims, blk_strides, block_data, 0); 101 IOSrc src(input_strides, input_data, desc.offset()); 102 103 TensorBlockIO::Copy(dst, src); 104 } 105 106 { 107 // Write from block buffer to output. 108 IODst dst(blk_dims, output_strides, output_data, desc.offset()); 109 IOSrc src(blk_strides, block_data, 0); 110 111 TensorBlockIO::Copy(dst, src); 112 } 113 } 114 115 for (int i = 0; i < dims.TotalSize(); ++i) { 116 VERIFY_IS_EQUAL(input_data[i], output_data[i]); 117 } 118 } 119 120 template <typename T, int NumDims, int Layout> 121 static void test_block_io_copy_using_reordered_dimensions() { 122 // Generate a random input Tensor. 123 DSizes<Index, NumDims> dims = RandomDims<NumDims>(1, 30); 124 Tensor<T, NumDims, Layout> input(dims); 125 input.setRandom(); 126 127 // Create a random dimension re-ordering/shuffle. 128 std::vector<int> shuffle; 129 130 for (int i = 0; i < NumDims; ++i) shuffle.push_back(i); 131 std::shuffle(shuffle.begin(), shuffle.end(), std::mt19937(g_seed)); 132 133 DSizes<Index, NumDims> output_tensor_dims; 134 DSizes<Index, NumDims> input_to_output_dim_map; 135 DSizes<Index, NumDims> output_to_input_dim_map; 136 for (Index i = 0; i < NumDims; ++i) { 137 output_tensor_dims[shuffle[i]] = dims[i]; 138 input_to_output_dim_map[i] = shuffle[i]; 139 output_to_input_dim_map[shuffle[i]] = i; 140 } 141 142 // Write data to an output Tensor. 143 Tensor<T, NumDims, Layout> output(output_tensor_dims); 144 145 // Construct a tensor block mapper. 146 // NOTE: Tensor block mapper works with shuffled dimensions. 147 using TensorBlockMapper = 148 internal::TensorBlockMapper<NumDims, Layout, Index>; 149 TensorBlockMapper block_mapper(output_tensor_dims, 150 {RandomBlockShape(), 151 RandomTargetBlockSize(output_tensor_dims), 152 {0, 0, 0}}); 153 154 // We will copy data from input to output through this buffer. 155 Tensor<T, NumDims, Layout> block(block_mapper.blockDimensions()); 156 157 // Precompute strides for TensorBlockIO::Copy. 158 auto input_strides = internal::strides<Layout>(dims); 159 auto output_strides = internal::strides<Layout>(output_tensor_dims); 160 161 const T* input_data = input.data(); 162 T* output_data = output.data(); 163 T* block_data = block.data(); 164 165 for (Index i = 0; i < block_mapper.blockCount(); ++i) { 166 auto desc = block_mapper.blockDescriptor(i); 167 168 const Index first_coeff_index = GetInputIndex<Layout, NumDims>( 169 desc.offset(), output_to_input_dim_map, input_strides, 170 output_strides); 171 172 // NOTE: Block dimensions are in the same order as output dimensions. 173 174 using TensorBlockIO = internal::TensorBlockIO<T, Index, NumDims, Layout>; 175 using IODst = typename TensorBlockIO::Dst; 176 using IOSrc = typename TensorBlockIO::Src; 177 178 auto blk_dims = desc.dimensions(); 179 auto blk_strides = internal::strides<Layout>(blk_dims); 180 181 { 182 // Read from input into a block buffer. 183 IODst dst(blk_dims, blk_strides, block_data, 0); 184 IOSrc src(input_strides, input_data, first_coeff_index); 185 186 // TODO(ezhulenev): Remove when fully switched to TensorBlock. 187 DSizes<int, NumDims> dim_map; 188 for (int j = 0; j < NumDims; ++j) 189 dim_map[j] = static_cast<int>(output_to_input_dim_map[j]); 190 TensorBlockIO::Copy(dst, src, /*dst_to_src_dim_map=*/dim_map); 191 } 192 193 { 194 // We need to convert block dimensions from output to input order. 195 auto dst_dims = blk_dims; 196 for (int out_dim = 0; out_dim < NumDims; ++out_dim) { 197 dst_dims[output_to_input_dim_map[out_dim]] = blk_dims[out_dim]; 198 } 199 200 // Write from block buffer to output. 201 IODst dst(dst_dims, input_strides, output_data, first_coeff_index); 202 IOSrc src(blk_strides, block_data, 0); 203 204 // TODO(ezhulenev): Remove when fully switched to TensorBlock. 205 DSizes<int, NumDims> dim_map; 206 for (int j = 0; j < NumDims; ++j) 207 dim_map[j] = static_cast<int>(input_to_output_dim_map[j]); 208 TensorBlockIO::Copy(dst, src, /*dst_to_src_dim_map=*/dim_map); 209 } 210 } 211 212 for (Index i = 0; i < dims.TotalSize(); ++i) { 213 VERIFY_IS_EQUAL(input_data[i], output_data[i]); 214 } 215 } 216 217 // This is the special case for reading data with reordering, when dimensions 218 // before/after reordering are the same. Squeezing reads along inner dimensions 219 // in this case is illegal, because we reorder innermost dimension. 220 template <int Layout> 221 static void test_block_io_copy_using_reordered_dimensions_do_not_squeeze() { 222 DSizes<Index, 3> tensor_dims(7, 9, 7); 223 DSizes<Index, 3> block_dims = tensor_dims; 224 225 DSizes<int, 3> block_to_tensor_dim; 226 block_to_tensor_dim[0] = 2; 227 block_to_tensor_dim[1] = 1; 228 block_to_tensor_dim[2] = 0; 229 230 auto tensor_strides = internal::strides<Layout>(tensor_dims); 231 auto block_strides = internal::strides<Layout>(block_dims); 232 233 Tensor<float, 3, Layout> block(block_dims); 234 Tensor<float, 3, Layout> tensor(tensor_dims); 235 tensor.setRandom(); 236 237 float* tensor_data = tensor.data(); 238 float* block_data = block.data(); 239 240 using TensorBlockIO = internal::TensorBlockIO<float, Index, 3, Layout>; 241 using IODst = typename TensorBlockIO::Dst; 242 using IOSrc = typename TensorBlockIO::Src; 243 244 // Read from a tensor into a block. 245 IODst dst(block_dims, block_strides, block_data, 0); 246 IOSrc src(tensor_strides, tensor_data, 0); 247 248 TensorBlockIO::Copy(dst, src, /*dst_to_src_dim_map=*/block_to_tensor_dim); 249 250 TensorMap<Tensor<float, 3, Layout> > block_tensor(block_data, block_dims); 251 TensorMap<Tensor<float, 3, Layout> > tensor_tensor(tensor_data, tensor_dims); 252 253 for (Index d0 = 0; d0 < tensor_dims[0]; ++d0) { 254 for (Index d1 = 0; d1 < tensor_dims[1]; ++d1) { 255 for (Index d2 = 0; d2 < tensor_dims[2]; ++d2) { 256 float block_value = block_tensor(d2, d1, d0); 257 float tensor_value = tensor_tensor(d0, d1, d2); 258 VERIFY_IS_EQUAL(block_value, tensor_value); 259 } 260 } 261 } 262 } 263 264 // This is the special case for reading data with reordering, when dimensions 265 // before/after reordering are the same. Squeezing reads in this case is allowed 266 // because we reorder outer dimensions. 267 template <int Layout> 268 static void test_block_io_copy_using_reordered_dimensions_squeeze() { 269 DSizes<Index, 4> tensor_dims(7, 5, 9, 9); 270 DSizes<Index, 4> block_dims = tensor_dims; 271 272 DSizes<int, 4> block_to_tensor_dim; 273 block_to_tensor_dim[0] = 0; 274 block_to_tensor_dim[1] = 1; 275 block_to_tensor_dim[2] = 3; 276 block_to_tensor_dim[3] = 2; 277 278 auto tensor_strides = internal::strides<Layout>(tensor_dims); 279 auto block_strides = internal::strides<Layout>(block_dims); 280 281 Tensor<float, 4, Layout> block(block_dims); 282 Tensor<float, 4, Layout> tensor(tensor_dims); 283 tensor.setRandom(); 284 285 float* tensor_data = tensor.data(); 286 float* block_data = block.data(); 287 288 using TensorBlockIO = internal::TensorBlockIO<float, Index, 4, Layout>; 289 using IODst = typename TensorBlockIO::Dst; 290 using IOSrc = typename TensorBlockIO::Src; 291 292 // Read from a tensor into a block. 293 IODst dst(block_dims, block_strides, block_data, 0); 294 IOSrc src(tensor_strides, tensor_data, 0); 295 296 TensorBlockIO::Copy(dst, src, /*dst_to_src_dim_map=*/block_to_tensor_dim); 297 298 TensorMap<Tensor<float, 4, Layout> > block_tensor(block_data, block_dims); 299 TensorMap<Tensor<float, 4, Layout> > tensor_tensor(tensor_data, tensor_dims); 300 301 for (Index d0 = 0; d0 < tensor_dims[0]; ++d0) { 302 for (Index d1 = 0; d1 < tensor_dims[1]; ++d1) { 303 for (Index d2 = 0; d2 < tensor_dims[2]; ++d2) { 304 for (Index d3 = 0; d3 < tensor_dims[3]; ++d3) { 305 float block_value = block_tensor(d0, d1, d3, d2); 306 float tensor_value = tensor_tensor(d0, d1, d2, d3); 307 VERIFY_IS_EQUAL(block_value, tensor_value); 308 } 309 } 310 } 311 } 312 } 313 314 template <int Layout> 315 static void test_block_io_zero_stride() { 316 DSizes<Index, 5> rnd_dims = RandomDims<5>(1, 30); 317 318 DSizes<Index, 5> input_tensor_dims = rnd_dims; 319 input_tensor_dims[0] = 1; 320 input_tensor_dims[2] = 1; 321 input_tensor_dims[4] = 1; 322 323 Tensor<float, 5, Layout> input(input_tensor_dims); 324 input.setRandom(); 325 326 DSizes<Index, 5> output_tensor_dims = rnd_dims; 327 328 auto input_tensor_strides = internal::strides<Layout>(input_tensor_dims); 329 auto output_tensor_strides = internal::strides<Layout>(output_tensor_dims); 330 331 auto input_tensor_strides_with_zeros = input_tensor_strides; 332 input_tensor_strides_with_zeros[0] = 0; 333 input_tensor_strides_with_zeros[2] = 0; 334 input_tensor_strides_with_zeros[4] = 0; 335 336 Tensor<float, 5, Layout> output(output_tensor_dims); 337 output.setRandom(); 338 339 using TensorBlockIO = internal::TensorBlockIO<float, Index, 5, Layout>; 340 using IODst = typename TensorBlockIO::Dst; 341 using IOSrc = typename TensorBlockIO::Src; 342 343 // Write data from input to output with broadcasting in dims [0, 2, 4]. 344 IODst dst(output_tensor_dims, output_tensor_strides, output.data(), 0); 345 IOSrc src(input_tensor_strides_with_zeros, input.data(), 0); 346 TensorBlockIO::Copy(dst, src); 347 348 for (int i = 0; i < output_tensor_dims[0]; ++i) { 349 for (int j = 0; j < output_tensor_dims[1]; ++j) { 350 for (int k = 0; k < output_tensor_dims[2]; ++k) { 351 for (int l = 0; l < output_tensor_dims[3]; ++l) { 352 for (int m = 0; m < output_tensor_dims[4]; ++m) { 353 float input_value = input(0, j, 0, l, 0); 354 float output_value = output(i, j, k, l, m); 355 VERIFY_IS_EQUAL(input_value, output_value); 356 } 357 } 358 } 359 } 360 } 361 } 362 363 template <int Layout> 364 static void test_block_io_squeeze_ones() { 365 using TensorBlockIO = internal::TensorBlockIO<float, Index, 5, Layout>; 366 using IODst = typename TensorBlockIO::Dst; 367 using IOSrc = typename TensorBlockIO::Src; 368 369 // Total size > 1. 370 { 371 DSizes<Index, 5> block_sizes(1, 2, 1, 2, 1); 372 auto strides = internal::strides<Layout>(block_sizes); 373 374 // Create a random input tensor. 375 Tensor<float, 5> input(block_sizes); 376 input.setRandom(); 377 378 Tensor<float, 5> output(block_sizes); 379 380 IODst dst(block_sizes, strides, output.data(), 0); 381 IOSrc src(strides, input.data()); 382 TensorBlockIO::Copy(dst, src); 383 384 for (Index i = 0; i < block_sizes.TotalSize(); ++i) { 385 VERIFY_IS_EQUAL(output.data()[i], input.data()[i]); 386 } 387 } 388 389 // Total size == 1. 390 { 391 DSizes<Index, 5> block_sizes(1, 1, 1, 1, 1); 392 auto strides = internal::strides<Layout>(block_sizes); 393 394 // Create a random input tensor. 395 Tensor<float, 5> input(block_sizes); 396 input.setRandom(); 397 398 Tensor<float, 5> output(block_sizes); 399 400 IODst dst(block_sizes, strides, output.data(), 0); 401 IOSrc src(strides, input.data()); 402 TensorBlockIO::Copy(dst, src); 403 404 for (Index i = 0; i < block_sizes.TotalSize(); ++i) { 405 VERIFY_IS_EQUAL(output.data()[i], input.data()[i]); 406 } 407 } 408 } 409 410 #define CALL_SUBTESTS(NAME) \ 411 CALL_SUBTEST((NAME<float, 1, RowMajor>())); \ 412 CALL_SUBTEST((NAME<float, 2, RowMajor>())); \ 413 CALL_SUBTEST((NAME<float, 4, RowMajor>())); \ 414 CALL_SUBTEST((NAME<float, 5, RowMajor>())); \ 415 CALL_SUBTEST((NAME<float, 1, ColMajor>())); \ 416 CALL_SUBTEST((NAME<float, 2, ColMajor>())); \ 417 CALL_SUBTEST((NAME<float, 4, ColMajor>())); \ 418 CALL_SUBTEST((NAME<float, 5, ColMajor>())); \ 419 CALL_SUBTEST((NAME<bool, 1, RowMajor>())); \ 420 CALL_SUBTEST((NAME<bool, 2, RowMajor>())); \ 421 CALL_SUBTEST((NAME<bool, 4, RowMajor>())); \ 422 CALL_SUBTEST((NAME<bool, 5, RowMajor>())); \ 423 CALL_SUBTEST((NAME<bool, 1, ColMajor>())); \ 424 CALL_SUBTEST((NAME<bool, 2, ColMajor>())); \ 425 CALL_SUBTEST((NAME<bool, 4, ColMajor>())); \ 426 CALL_SUBTEST((NAME<bool, 5, ColMajor>())) 427 428 EIGEN_DECLARE_TEST(cxx11_tensor_block_io) { 429 // clang-format off 430 CALL_SUBTESTS(test_block_io_copy_data_from_source_to_target); 431 CALL_SUBTESTS(test_block_io_copy_using_reordered_dimensions); 432 433 CALL_SUBTEST(test_block_io_copy_using_reordered_dimensions_do_not_squeeze<RowMajor>()); 434 CALL_SUBTEST(test_block_io_copy_using_reordered_dimensions_do_not_squeeze<ColMajor>()); 435 436 CALL_SUBTEST(test_block_io_copy_using_reordered_dimensions_squeeze<RowMajor>()); 437 CALL_SUBTEST(test_block_io_copy_using_reordered_dimensions_squeeze<ColMajor>()); 438 439 CALL_SUBTEST(test_block_io_zero_stride<RowMajor>()); 440 CALL_SUBTEST(test_block_io_zero_stride<ColMajor>()); 441 442 CALL_SUBTEST(test_block_io_squeeze_ones<RowMajor>()); 443 CALL_SUBTEST(test_block_io_squeeze_ones<ColMajor>()); 444 // clang-format on 445 }