inference-engine/thirdparty/clDNN/tests/test_cases/index_select_gpu_test.cpp

   1 /*
   2 // Copyright (c) 2016 Intel Corporation
   3 //
   4 // Licensed under the Apache License, Version 2.0 (the "License");
   5 // you may not use this file except in compliance with the License.
   6 // You may obtain a copy of the License at
   7 //
   8 //      http://www.apache.org/licenses/LICENSE-2.0
   9 //
  10 // Unless required by applicable law or agreed to in writing, software
  11 // distributed under the License is distributed on an "AS IS" BASIS,
  12 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 // See the License for the specific language governing permissions and
  14 // limitations under the License.
  15 */
  16
  17 ///////////////////////////////////////////////////////////////////////////////////////////////////
  18
  19 #include <gtest/gtest.h>
  20
  21 #include <api/CPP/engine.hpp>
  22 #include <api/CPP/input_layout.hpp>
  23 #include <api/CPP/index_select.hpp>
  24 #include <api/CPP/memory.hpp>
  25 #include <api/CPP/topology.hpp>
  26 #include <api/CPP/network.hpp>
  27
  28 #include "test_utils/test_utils.h"
  29
  30 #include <vector>
  31 #include <algorithm>
  32
  33 using namespace cldnn;
  34 using namespace tests;
  35
  36 std::vector<float> generate_reference_bfyx(const std::vector<float>& input, const std::vector<int32_t>& indices, index_select_axis_name axis, const size_t b_size, const size_t f_size,
  37     const size_t y_size, const size_t x_size)
  38 {
  39
  40     auto unique_indices = std::set<int32_t>(indices.begin(), indices.end());
  41     std::map<uint32_t, std::vector<float>> mapped_axises;
  42
  43     auto append_to_vector = [&](std::vector<float>& vec)
  44     {
  45         for (auto const& id : indices)
  46         {
  47             vec.insert(vec.end(), mapped_axises.at(id).begin(), mapped_axises.at(id).end());
  48         }
  49     };
  50
  51     std::vector<float> ret;
  52     size_t offset = 0;
  53     switch (axis)
  54     {
  55     case index_select_axis_name::along_b:
  56         offset = f_size * y_size * x_size;
  57         for (auto const& ui : unique_indices)
  58         {
  59             mapped_axises[ui] = std::vector<float>(input.begin() + ui * offset, input.begin() + ui * offset + offset);
  60         }
  61         append_to_vector(ret);
  62         return ret;
  63     case index_select_axis_name::along_f:
  64         offset = y_size * x_size;
  65         for (size_t i = 0; i < b_size; i++)
  66         {
  67             size_t batch_index = i * f_size * y_size * x_size;
  68             mapped_axises.clear();
  69             for (auto const& ui : unique_indices)
  70             {
  71                 mapped_axises[ui] = std::vector<float>(input.begin() + batch_index + ui * offset, input.begin() + batch_index + ui * offset + offset);
  72             }
  73             append_to_vector(ret);
  74         }
  75         return ret;
  76     case index_select_axis_name::along_x:
  77         offset = y_size;
  78         ret.resize(b_size * f_size * y_size * indices.size());
  79         for (size_t i = 0; i < b_size; i++)
  80         {
  81             size_t batch_index = i * f_size * y_size * x_size;
  82             for (size_t j = 0; j < f_size; j++)
  83             {
  84                 size_t feature_index = j * y_size * x_size;
  85                 size_t start_idx = batch_index + feature_index;
  86                 mapped_axises.clear();
  87                 for (auto const& ui : unique_indices)
  88                 {
  89                     std::vector<float> values = {};
  90                     for (size_t k = 0; k < offset; k++)
  91                     {
  92                         values.push_back(input.at(start_idx + k * x_size +  ui));
  93                     }
  94                     mapped_axises[ui] = values;
  95                 }
  96
  97                 for (size_t idx = 0; idx < indices.size(); idx++)
  98                 {
  99                     auto const id = indices.at(idx);
 100                     //ret.insert(ret.end(), mapped_axises.at(id).begin(), mapped_axises.at(id).end());
 101                     auto out_idx = i * f_size * y_size * indices.size() + j * y_size * indices.size() + idx;
 102                     for (size_t y = 0; y < y_size; y++)
 103                     {
 104                         ret.at(out_idx + y * indices.size()) = mapped_axises.at(id).at(y);
 105                     }
 106                 }
 107             }
 108         }
 109         return ret;
 110         break;
 111     case index_select_axis_name::along_y:
 112         offset = x_size;
 113         for (size_t i = 0; i < b_size; i++)
 114         {
 115             size_t batch_index = i * f_size * y_size * x_size;
 116             for (size_t j = 0; j < f_size; j++)
 117             {
 118                 size_t feature_index = j * y_size * x_size;
 119                 size_t start_idx = batch_index + feature_index;
 120                 mapped_axises.clear();
 121                 for (auto const& ui : unique_indices)
 122                 {
 123                     mapped_axises[ui] = std::vector<float>(input.begin() + start_idx + ui * offset, input.begin() + start_idx + ui * offset + offset);
 124                 }
 125                 append_to_vector(ret);
 126             }
 127         }
 128         return ret;
 129         break;
 130     default:
 131         throw std::runtime_error("Unknown index_select axis!");
 132         break;
 133     }
 134 }
 135
 136
 137 std::vector<float> generate_reference_yxfb(const std::vector<float>& input, const std::vector<int32_t>& indices, index_select_axis_name axis, const cldnn::layout& input_lay)
 138 {
 139     auto memory_desc_inp = generic_test::get_linear_memory_desc(input_lay);
 140
 141     std::vector<float> ret;
 142     switch (axis)
 143     {
 144     case index_select_axis_name::along_b:
 145
 146         for (auto y = 0; y < input_lay.size.spatial[1]; y++)
 147         {
 148             for (auto x = 0; x < input_lay.size.spatial[0]; x++)
 149             {
 150                 for (auto f = 0; f < input_lay.size.feature[0]; f++)
 151                 {
 152                     for (auto const& ind : indices)
 153                     {
 154
 155                         size_t index = generic_test::get_linear_index(input_lay, ind, f, y, x, memory_desc_inp);
 156                         ret.push_back(input.at(index));
 157                     }
 158
 159                 }
 160             }
 161         }
 162         return ret;
 163     case index_select_axis_name::along_f:
 164         for (auto y = 0; y < input_lay.size.spatial[1]; y++)
 165         {
 166             for (auto x = 0; x < input_lay.size.spatial[0]; x++)
 167             {
 168                 for (auto const& ind : indices)
 169                 {
 170                     for (auto b = 0; b < input_lay.size.batch[0]; b++)
 171                     {
 172                     size_t index = generic_test::get_linear_index(input_lay, b, ind, y, x, memory_desc_inp);
 173                     ret.push_back(input.at(index));
 174                     }
 175                 }
 176             }
 177         }
 178         return ret;
 179     case index_select_axis_name::along_x:
 180         for (auto y = 0; y < input_lay.size.spatial[1]; y++)
 181         {
 182             for (auto const& ind : indices)
 183             {
 184                 for (auto f = 0; f < input_lay.size.feature[0]; f++)
 185                 {
 186                     for (auto b = 0; b < input_lay.size.batch[0]; b++)
 187                     {
 188                         size_t index = generic_test::get_linear_index(input_lay, b, f, y, ind, memory_desc_inp);
 189                         ret.push_back(input.at(index));
 190                     }
 191                 }
 192             }
 193         }
 194         return ret;
 195     case index_select_axis_name::along_y:
 196
 197         for (auto const& ind : indices)
 198         {
 199             for (auto x = 0; x < input_lay.size.spatial[0]; x++)
 200             {
 201                 for (auto f = 0; f < input_lay.size.feature[0]; f++)
 202                 {
 203                     for (auto b = 0; b < input_lay.size.batch[0]; b++)
 204                     {
 205                         size_t index = generic_test::get_linear_index(input_lay, b, f, ind, x, memory_desc_inp);
 206                         ret.push_back(input.at(index));
 207                     }
 208                 }
 209             }
 210         }
 211         return ret;
 212     default:
 213         throw std::runtime_error("Unknown index_select axis!");
 214         break;
 215     }
 216 }
 217
 218
 219 TEST(index_select_gpu, basic_along_b_3_executes_bfyx)
 220 {
 221     /*
 222     input: {5, 2, 3, 4}
 223     indices: {1, 1, 4, 1}
 224     output: {4, 2, 3, 4}
 225     */
 226     const auto& engine = get_test_engine();
 227     constexpr auto in_size_b = 5;
 228     constexpr auto in_size_f = 2;
 229     constexpr auto in_size_x = 3;
 230     constexpr auto in_size_y = 4;
 231     constexpr auto count = in_size_b * in_size_f * in_size_x * in_size_y;
 232     constexpr auto new_indicies_size = 4;
 233     constexpr auto axis = index_select_axis_name::along_b;
 234     auto input = memory::allocate(engine, { data_types::f32, format::bfyx, { in_size_b, in_size_f, in_size_x, in_size_y } });
 235     auto indices = memory::allocate(engine, { data_types::i32, format::bfyx, { 1, 1, new_indicies_size, 1 } });
 236
 237     auto input_data = generate_random_1d<float>(count, 0, 10);
 238     set_values(input, input_data);
 239
 240     /*
 241     Network will be executed 3 times (for 3 different indicies_data).
 242     */
 243     std::vector<std::vector<int32_t>> indices_data =
 244     {
 245         {0, 1, 1, 1}, //for run: 0
 246         {0, 1, 2, 3}, //for run: 1
 247         {4, 3, 2, 1} // for run: 2
 248     };
 249
 250     topology topo;
 251     topo.add(
 252         input_layout("input", input.get_layout())
 253     );
 254     topo.add(
 255         input_layout("indices", indices.get_layout())
 256     );
 257     topo.add(
 258         index_select("index_select", "input", "indices", axis)
 259     );
 260
 261     network net(engine, topo);
 262     net.set_input_data("input", input);
 263     for (auto const& id : indices_data)
 264     {
 265         set_values(indices, id);
 266         net.set_input_data("indices", indices);
 267         auto outputs = net.execute();
 268
 269                 ASSERT_EQ(outputs.size(), size_t(1));
 270         EXPECT_EQ(outputs.begin()->first, "index_select");
 271
 272         auto output_mem = outputs.at("index_select").get_memory();
 273         auto output_layout = output_mem.get_layout();
 274
 275         int b_size = output_layout.size.batch[0];
 276         int f_size = output_layout.size.feature[0];
 277         int x_size = output_layout.size.spatial[0];
 278         int y_size = output_layout.size.spatial[1];
 279         EXPECT_EQ(output_layout.format, format::bfyx);
 280         EXPECT_EQ(b_size, new_indicies_size);
 281         EXPECT_EQ(f_size, in_size_f);
 282         EXPECT_EQ(x_size, in_size_x);
 283         EXPECT_EQ(y_size, in_size_y);
 284
 285         auto ref = generate_reference_bfyx(input_data, id, axis, in_size_b, in_size_f, in_size_y, in_size_x);
 286
 287         auto output_ptr = output_mem.pointer<float>();
 288         for (size_t i = 0; i < output_ptr.size(); i++)
 289         {
 290             EXPECT_EQ(output_ptr[i], ref[i]);
 291         }
 292     }
 293 }
 294
 295 TEST(index_select_gpu, basic_along_f_3_executes_bfyx)
 296 {
 297     /*
 298     input: {2, 5, 3, 3}
 299     indices: {1, 1, 10, 1}
 300     output: {2, 10, 3, 3}
 301     */
 302     const auto& engine = get_test_engine();
 303     constexpr auto in_size_b = 2;
 304     constexpr auto in_size_f = 5;
 305     constexpr auto in_size_x = 3;
 306     constexpr auto in_size_y = 3;
 307     constexpr auto count = in_size_b * in_size_f * in_size_x * in_size_y;
 308     constexpr auto new_indicies_size = 10;
 309     constexpr auto axis = index_select_axis_name::along_f;
 310     auto input = memory::allocate(engine, { data_types::f32, format::bfyx, { in_size_b, in_size_f, in_size_x, in_size_y } });
 311     auto indices = memory::allocate(engine, { data_types::i32, format::bfyx, { 1, 1, new_indicies_size, 1 } });
 312
 313     auto input_data = generate_random_1d<float>(count, 0, 10);
 314     set_values(input, input_data);
 315
 316     /*
 317     Network will be executed 3 times (for 3 different indicies_data).
 318     */
 319     std::vector<std::vector<int32_t>> indices_data =
 320     {
 321         { 0, 1, 2, 3, 4, 0, 1, 2, 3, 4 }, //for run: 0
 322         { 1, 1, 3, 3, 2, 2, 4, 4, 0, 0 }, //for run: 1
 323         { 0, 0, 0, 0, 0, 4, 3, 2, 1, 0 } // for run: 2
 324     };
 325
 326     topology topo;
 327     topo.add(
 328         input_layout("input", input.get_layout())
 329     );
 330     topo.add(
 331         input_layout("indices", indices.get_layout())
 332     );
 333     topo.add(
 334         index_select("index_select", "input", "indices", axis)
 335     );
 336
 337     network net(engine, topo);
 338     net.set_input_data("input", input);
 339     for (auto const& id : indices_data)
 340     {
 341         set_values(indices, id);
 342         net.set_input_data("indices", indices);
 343         auto outputs = net.execute();
 344
 345                 ASSERT_EQ(outputs.size(), size_t(1));
 346         EXPECT_EQ(outputs.begin()->first, "index_select");
 347
 348         auto output_mem = outputs.at("index_select").get_memory();
 349         auto output_layout = output_mem.get_layout();
 350
 351         int b_size = output_layout.size.batch[0];
 352         int f_size = output_layout.size.feature[0];
 353         int x_size = output_layout.size.spatial[0];
 354         int y_size = output_layout.size.spatial[1];
 355         EXPECT_EQ(output_layout.format, format::bfyx);
 356         EXPECT_EQ(b_size, in_size_b);
 357         EXPECT_EQ(f_size, new_indicies_size);
 358         EXPECT_EQ(x_size, in_size_x);
 359         EXPECT_EQ(y_size, in_size_y);
 360
 361         auto ref = generate_reference_bfyx(input_data, id, axis, in_size_b, in_size_f, in_size_y, in_size_x);
 362
 363         auto output_ptr = output_mem.pointer<float>();
 364         for (size_t i = 0; i < output_ptr.size(); i++)
 365         {
 366             EXPECT_EQ(output_ptr[i], ref[i]);
 367         }
 368     }
 369 }
 370
 371 TEST(index_select_gpu, basic_along_x_3_executes_bfyx)
 372 {
 373     /*
 374     input: {3, 4, 6, 5}
 375     indices: {1, 1, 3, 1}
 376     output: {3, 4, 3, 5}
 377     */
 378     const auto& engine = get_test_engine();
 379     constexpr auto in_size_b = 3;
 380     constexpr auto in_size_f = 4;
 381     constexpr auto in_size_x = 6;
 382     constexpr auto in_size_y = 5;
 383     constexpr auto count = in_size_b * in_size_f * in_size_x * in_size_y;
 384     constexpr auto new_indicies_size = 3;
 385     constexpr auto axis = index_select_axis_name::along_x;
 386     auto input = memory::allocate(engine, { data_types::f32, format::bfyx, { in_size_b, in_size_f, in_size_x, in_size_y } });
 387     auto indices = memory::allocate(engine, { data_types::i32, format::bfyx, { 1, 1, new_indicies_size, 1 } });
 388
 389     auto input_data = generate_random_1d<float>(count, 0, 10);
 390     set_values(input, input_data);
 391
 392     /*
 393     Network will be executed 3 times (for 3 different indicies_data).
 394     */
 395     std::vector<std::vector<int32_t>> indices_data =
 396     {
 397         { 2, 1, 0 }, //for run: 0
 398         { 0, 0, 0 }, //for run: 1
 399         { 1, 1, 0 } // for run: 2
 400     };
 401
 402     topology topo;
 403     topo.add(
 404         input_layout("input", input.get_layout())
 405     );
 406     topo.add(
 407         input_layout("indices", indices.get_layout())
 408     );
 409     topo.add(
 410         index_select("index_select", "input", "indices", axis)
 411     );
 412
 413     network net(engine, topo);
 414     net.set_input_data("input", input);
 415     for (auto const& id : indices_data)
 416     {
 417         set_values(indices, id);
 418         net.set_input_data("indices", indices);
 419         auto outputs = net.execute();
 420
 421                 ASSERT_EQ(outputs.size(), size_t(1));
 422         EXPECT_EQ(outputs.begin()->first, "index_select");
 423
 424         auto output_mem = outputs.at("index_select").get_memory();
 425         auto output_layout = output_mem.get_layout();
 426
 427         int b_size = output_layout.size.batch[0];
 428         int f_size = output_layout.size.feature[0];
 429         int x_size = output_layout.size.spatial[0];
 430         int y_size = output_layout.size.spatial[1];
 431         EXPECT_EQ(output_layout.format, format::bfyx);
 432         EXPECT_EQ(b_size, in_size_b);
 433         EXPECT_EQ(f_size, in_size_f);
 434         EXPECT_EQ(x_size, new_indicies_size);
 435         EXPECT_EQ(y_size, in_size_y);
 436
 437         auto ref = generate_reference_bfyx(input_data, id, axis, in_size_b, in_size_f, in_size_y, in_size_x);
 438
 439         auto output_ptr = output_mem.pointer<float>();
 440         for (size_t i = 0; i < output_ptr.size(); i++)
 441         {
 442             EXPECT_EQ(output_ptr[i], ref[i]);
 443         }
 444     }
 445 }
 446
 447 TEST(index_select_gpu, basic_along_y_3_executes_bfyx)
 448 {
 449     /*
 450     input: {2, 4, 4, 3}
 451     indices: {1, 1, 5, 1}
 452     output: {2, 4, 4, 5}
 453     */
 454     const auto& engine = get_test_engine();
 455     constexpr auto in_size_b = 2;
 456     constexpr auto in_size_f = 4;
 457     constexpr auto in_size_x = 4;
 458     constexpr auto in_size_y = 3;
 459     constexpr auto count = in_size_b * in_size_f * in_size_x * in_size_y;
 460     constexpr auto new_indicies_size = 5;
 461     constexpr auto axis = index_select_axis_name::along_y;
 462     auto input = memory::allocate(engine, { data_types::f32, format::bfyx,{ in_size_b, in_size_f, in_size_x, in_size_y } });
 463     auto indices = memory::allocate(engine, { data_types::i32, format::bfyx,{ 1, 1, new_indicies_size, 1 } });
 464
 465     auto input_data = generate_random_1d<float>(count, 0, 10);
 466     set_values(input, input_data);
 467
 468     /*
 469     Network will be executed 3 times (for 3 different indicies_data).
 470     */
 471     std::vector<std::vector<int32_t>> indices_data =
 472     {
 473         { 0, 1, 2, 2, 1 }, //for run: 0
 474         { 2, 2, 1, 0, 1 }, //for run: 1
 475         { 1, 1, 2, 1, 0 } // for run: 2
 476     };
 477
 478     topology topo;
 479     topo.add(
 480         input_layout("input", input.get_layout())
 481     );
 482     topo.add(
 483         input_layout("indices", indices.get_layout())
 484     );
 485     topo.add(
 486         index_select("index_select", "input", "indices", axis)
 487     );
 488
 489     network net(engine, topo);
 490     net.set_input_data("input", input);
 491     for (auto const& id : indices_data)
 492     {
 493         set_values(indices, id);
 494         net.set_input_data("indices", indices);
 495         auto outputs = net.execute();
 496
 497                 ASSERT_EQ(outputs.size(), size_t(1));
 498         EXPECT_EQ(outputs.begin()->first, "index_select");
 499
 500         auto output_mem = outputs.at("index_select").get_memory();
 501         auto output_layout = output_mem.get_layout();
 502
 503         int b_size = output_layout.size.batch[0];
 504         int f_size = output_layout.size.feature[0];
 505         int x_size = output_layout.size.spatial[0];
 506         int y_size = output_layout.size.spatial[1];
 507         EXPECT_EQ(output_layout.format, format::bfyx);
 508         EXPECT_EQ(b_size, in_size_b);
 509         EXPECT_EQ(f_size, in_size_f);
 510         EXPECT_EQ(x_size, x_size);
 511         EXPECT_EQ(y_size, new_indicies_size);
 512
 513         auto ref = generate_reference_bfyx(input_data, id, axis, in_size_b, in_size_f, in_size_y, in_size_x);
 514
 515         auto output_ptr = output_mem.pointer<float>();
 516         for (size_t i = 0; i < output_ptr.size(); i++)
 517         {
 518             EXPECT_EQ(output_ptr[i], ref[i]);
 519         }
 520     }
 521 }
 522
 523 TEST(index_select_gpu, basic_along_b_3_executes_yxfb)
 524 {
 525     /*
 526     input: {5, 2, 3, 4}
 527     indices: {1, 1, 4, 1}
 528     output: {4, 2, 3, 4}
 529     */
 530     const auto& engine = get_test_engine();
 531     constexpr auto in_size_b = 5;
 532     constexpr auto in_size_f = 2;
 533     constexpr auto in_size_x = 3;
 534     constexpr auto in_size_y = 4;
 535     constexpr auto count = in_size_b * in_size_f * in_size_x * in_size_y;
 536     constexpr auto new_indicies_size = 4;
 537     constexpr auto axis = index_select_axis_name::along_b;
 538     auto input_lay= cldnn::layout(data_types::f32, format::yxfb,{ in_size_b, in_size_f, in_size_x, in_size_y });
 539     auto input = memory::allocate(engine, input_lay);
 540     auto indices = memory::allocate(engine, { data_types::i32, format::yxfb, { 1, 1, new_indicies_size, 1 } });
 541
 542     auto input_data = generate_random_1d<float>(count, 0, 10);
 543     set_values(input, input_data);
 544
 545     /*
 546     Network will be executed 3 times (for 3 different indicies_data).
 547     */
 548     std::vector<std::vector<int32_t>> indices_data =
 549     {
 550         { 0, 1, 1, 1 }, //for run: 0
 551         { 0, 1, 2, 3 }, //for run: 1
 552         { 4, 3, 2, 1 } // for run: 2
 553     };
 554
 555     topology topo;
 556     topo.add(
 557         input_layout("input", input.get_layout())
 558     );
 559     topo.add(
 560         input_layout("indices", indices.get_layout())
 561     );
 562     topo.add(
 563         index_select("index_select", "input", "indices", axis)
 564     );
 565
 566     network net(engine, topo);
 567     net.set_input_data("input", input);
 568     for (auto const& id : indices_data)
 569     {
 570         set_values(indices, id);
 571         net.set_input_data("indices", indices);
 572         auto outputs = net.execute();
 573
 574                 ASSERT_EQ(outputs.size(), size_t(1));
 575         EXPECT_EQ(outputs.begin()->first, "index_select");
 576
 577         auto output_mem = outputs.at("index_select").get_memory();
 578         auto output_layout = output_mem.get_layout();
 579
 580         int b_size = output_layout.size.batch[0];
 581         int f_size = output_layout.size.feature[0];
 582         int x_size = output_layout.size.spatial[0];
 583         int y_size = output_layout.size.spatial[1];
 584         EXPECT_EQ(output_layout.format, format::yxfb);
 585         EXPECT_EQ(b_size, new_indicies_size);
 586         EXPECT_EQ(f_size, in_size_f);
 587         EXPECT_EQ(x_size, in_size_x);
 588         EXPECT_EQ(y_size, in_size_y);
 589
 590         auto ref = generate_reference_yxfb(input_data, id, axis, input_lay);
 591
 592         auto output_ptr = output_mem.pointer<float>();
 593         for (size_t i = 0; i < output_ptr.size(); i++)
 594         {
 595             EXPECT_EQ(output_ptr[i], ref[i]);
 596         }
 597     }
 598 }
 599
 600 TEST(index_select_gpu, basic_along_f_3_executes_yxfb)
 601 {
 602     /*
 603     input: {2, 5, 3, 3}
 604     indices: {1, 1, 10, 1}
 605     output: {2, 10, 3, 3}
 606     */
 607     const auto& engine = get_test_engine();
 608     constexpr auto in_size_b = 2;
 609     constexpr auto in_size_f = 5;
 610     constexpr auto in_size_x = 3;
 611     constexpr auto in_size_y = 3;
 612     constexpr auto count = in_size_b * in_size_f * in_size_x * in_size_y;
 613     constexpr auto new_indicies_size = 10;
 614     constexpr auto axis = index_select_axis_name::along_f;
 615     auto input_lay = cldnn::layout(data_types::f32, format::yxfb, { in_size_b, in_size_f, in_size_x, in_size_y });
 616     auto input = memory::allocate(engine, input_lay);
 617     auto indices = memory::allocate(engine, { data_types::i32, format::yxfb,{ 1, 1, new_indicies_size, 1 } });
 618
 619     auto input_data = generate_random_1d<float>(count, 0, 10);
 620     set_values(input, input_data);
 621
 622     /*
 623     Network will be executed 3 times (for 3 different indicies_data).
 624     */
 625     std::vector<std::vector<int32_t>> indices_data =
 626     {
 627         { 0, 1, 2, 3, 4, 0, 1, 2, 3, 4 }, //for run: 0
 628     { 1, 1, 3, 3, 2, 2, 4, 4, 0, 0 }, //for run: 1
 629     { 0, 0, 0, 0, 0, 4, 3, 2, 1, 0 } // for run: 2
 630     };
 631
 632     topology topo;
 633     topo.add(
 634         input_layout("input", input.get_layout())
 635     );
 636     topo.add(
 637         input_layout("indices", indices.get_layout())
 638     );
 639     topo.add(
 640         index_select("index_select", "input", "indices", axis)
 641     );
 642
 643     network net(engine, topo);
 644     net.set_input_data("input", input);
 645     for (auto const& id : indices_data)
 646     {
 647         set_values(indices, id);
 648         net.set_input_data("indices", indices);
 649         auto outputs = net.execute();
 650
 651                 ASSERT_EQ(outputs.size(), size_t(1));
 652         EXPECT_EQ(outputs.begin()->first, "index_select");
 653
 654         auto output_mem = outputs.at("index_select").get_memory();
 655         auto output_layout = output_mem.get_layout();
 656
 657         int b_size = output_layout.size.batch[0];
 658         int f_size = output_layout.size.feature[0];
 659         int x_size = output_layout.size.spatial[0];
 660         int y_size = output_layout.size.spatial[1];
 661         EXPECT_EQ(output_layout.format, format::yxfb);
 662         EXPECT_EQ(b_size, in_size_b);
 663         EXPECT_EQ(f_size, new_indicies_size);
 664         EXPECT_EQ(x_size, in_size_x);
 665         EXPECT_EQ(y_size, in_size_y);
 666
 667         auto ref = generate_reference_yxfb(input_data, id, axis, input_lay);
 668
 669         auto output_ptr = output_mem.pointer<float>();
 670         for (size_t i = 0; i < output_ptr.size(); i++)
 671         {
 672             EXPECT_EQ(output_ptr[i], ref[i]);
 673         }
 674     }
 675 }
 676
 677 TEST(index_select_gpu, basic_along_x_3_executes_yxfb)
 678 {
 679     /*
 680     input: {3, 4, 6, 5}
 681     indices: {1, 1, 3, 1}
 682     output: {3, 4, 3, 5}
 683     */
 684     const auto& engine = get_test_engine();
 685     constexpr auto in_size_b = 3;
 686     constexpr auto in_size_f = 4;
 687     constexpr auto in_size_x = 6;
 688     constexpr auto in_size_y = 5;
 689     constexpr auto count = in_size_b * in_size_f * in_size_x * in_size_y;
 690     constexpr auto new_indicies_size = 3;
 691     constexpr auto axis = index_select_axis_name::along_x;
 692     auto input_lay = cldnn::layout(data_types::f32, format::yxfb, { in_size_b, in_size_f, in_size_x, in_size_y });
 693     auto input = memory::allocate(engine, input_lay);
 694     auto indices = memory::allocate(engine, { data_types::i32, format::yxfb,{ 1, 1, new_indicies_size, 1 } });
 695
 696     auto input_data = generate_random_1d<float>(count, 0, 10);
 697     set_values(input, input_data);
 698
 699     /*
 700     Network will be executed 3 times (for 3 different indicies_data).
 701     */
 702     std::vector<std::vector<int32_t>> indices_data =
 703     {
 704         { 2, 1, 0 }, //for run: 0
 705         { 0, 0, 0 }, //for run: 1
 706         { 1, 1, 0 } // for run: 2
 707     };
 708
 709     topology topo;
 710     topo.add(
 711         input_layout("input", input.get_layout())
 712     );
 713     topo.add(
 714         input_layout("indices", indices.get_layout())
 715     );
 716     topo.add(
 717         index_select("index_select", "input", "indices", axis)
 718     );
 719
 720     network net(engine, topo);
 721     net.set_input_data("input", input);
 722     for (auto const& id : indices_data)
 723     {
 724         set_values(indices, id);
 725         net.set_input_data("indices", indices);
 726         auto outputs = net.execute();
 727
 728                 ASSERT_EQ(outputs.size(), size_t(1));
 729         EXPECT_EQ(outputs.begin()->first, "index_select");
 730
 731         auto output_mem = outputs.at("index_select").get_memory();
 732         auto output_layout = output_mem.get_layout();
 733
 734         int b_size = output_layout.size.batch[0];
 735         int f_size = output_layout.size.feature[0];
 736         int x_size = output_layout.size.spatial[0];
 737         int y_size = output_layout.size.spatial[1];
 738         EXPECT_EQ(output_layout.format, format::yxfb);
 739         EXPECT_EQ(b_size, in_size_b);
 740         EXPECT_EQ(f_size, in_size_f);
 741         EXPECT_EQ(x_size, new_indicies_size);
 742         EXPECT_EQ(y_size, in_size_y);
 743
 744         auto ref = generate_reference_yxfb(input_data, id, axis, input_lay);
 745
 746         auto output_ptr = output_mem.pointer<float>();
 747         for (size_t i = 0; i < output_ptr.size(); i++)
 748         {
 749             EXPECT_EQ(output_ptr[i], ref[i]);
 750         }
 751     }
 752 }
 753 TEST(index_select_gpu, basic_along_y_3_executes_yxfb)
 754 {
 755     /*
 756     input: {2, 4, 4, 3}
 757     indices: {1, 1, 5, 1}
 758     output: {2, 4, 4, 5}
 759     */
 760     const auto& engine = get_test_engine();
 761     constexpr auto in_size_b = 2;
 762     constexpr auto in_size_f = 4;
 763     constexpr auto in_size_x = 4;
 764     constexpr auto in_size_y = 3;
 765     constexpr auto count = in_size_b * in_size_f * in_size_x * in_size_y;
 766     constexpr auto new_indicies_size = 5;
 767     constexpr auto axis = index_select_axis_name::along_y;
 768     auto input_lay = cldnn::layout(data_types::f32, format::yxfb, { in_size_b, in_size_f, in_size_x, in_size_y });
 769     auto input = memory::allocate(engine, input_lay);
 770     auto indices = memory::allocate(engine, { data_types::i32, format::yxfb,{ 1, 1, new_indicies_size, 1 } });
 771
 772     auto input_data = generate_random_1d<float>(count, 0, 10);
 773     set_values(input, input_data);
 774
 775     /*
 776     Network will be executed 3 times (for 3 different indicies_data).
 777     */
 778     std::vector<std::vector<int32_t>> indices_data =
 779     {
 780         { 0, 1, 2, 2, 1 }, //for run: 0
 781         { 2, 2, 1, 0, 1 }, //for run: 1
 782         { 1, 1, 2, 1, 0 } // for run: 2
 783     };
 784
 785     topology topo;
 786     topo.add(
 787         input_layout("input", input.get_layout())
 788     );
 789     topo.add(
 790         input_layout("indices", indices.get_layout())
 791     );
 792     topo.add(
 793         index_select("index_select", "input", "indices", axis)
 794     );
 795
 796     network net(engine, topo);
 797     net.set_input_data("input", input);
 798     for (auto const& id : indices_data)
 799     {
 800         set_values(indices, id);
 801         net.set_input_data("indices", indices);
 802         auto outputs = net.execute();
 803
 804                 ASSERT_EQ(outputs.size(), size_t(1));
 805         EXPECT_EQ(outputs.begin()->first, "index_select");
 806
 807         auto output_mem = outputs.at("index_select").get_memory();
 808         auto output_layout = output_mem.get_layout();
 809
 810         int b_size = output_layout.size.batch[0];
 811         int f_size = output_layout.size.feature[0];
 812         int x_size = output_layout.size.spatial[0];
 813         int y_size = output_layout.size.spatial[1];
 814         EXPECT_EQ(output_layout.format, format::yxfb);
 815         EXPECT_EQ(b_size, in_size_b);
 816         EXPECT_EQ(f_size, in_size_f);
 817         EXPECT_EQ(x_size, x_size);
 818         EXPECT_EQ(y_size, new_indicies_size);
 819
 820         auto ref = generate_reference_yxfb(input_data, id, axis, input_lay);
 821
 822         auto output_ptr = output_mem.pointer<float>();
 823         for (size_t i = 0; i < output_ptr.size(); i++)
 824         {
 825             EXPECT_EQ(output_ptr[i], ref[i]);
 826         }
 827     }
 828 }
 829
 830 TEST(index_select_gpu, reverse_along_b_bfyx)
 831 {
 832     const auto& engine = get_test_engine();
 833     auto input = memory::allocate(engine, { data_types::f32, format::bfyx,{ 2, 2, 4, 2 } });
 834
 835     std::vector<float> input_data = {
 836         0.f,  1.f,  2.f,  3.f,
 837         4.f,  5.f,  6.f,  7.f,
 838
 839         8.f,  9.f, 10.f, 11.f,
 840         12.f, 13.f, 14.f, 15.f,
 841
 842
 843
 844         16.f, 17.f, 18.f, 19.f,
 845         20.f, 21.f, 22.f, 23.f,
 846
 847         24.f, 25.f, 26.f, 27.f,
 848         28.f, 29.f, 30.f, 31.f,
 849     };
 850
 851     std::vector<float> out_data = {
 852         16.f, 17.f, 18.f, 19.f,
 853         20.f, 21.f, 22.f, 23.f,
 854
 855         24.f, 25.f, 26.f, 27.f,
 856         28.f, 29.f, 30.f, 31.f,
 857
 858
 859
 860         0.f,  1.f,  2.f,  3.f,
 861         4.f,  5.f,  6.f,  7.f,
 862
 863         8.f,  9.f, 10.f, 11.f,
 864         12.f, 13.f, 14.f, 15.f,
 865     };
 866
 867     constexpr auto axis = index_select_axis_name::along_b;
 868
 869     topology topo;
 870     topo.add(
 871         input_layout("input", input.get_layout())
 872     );
 873     topo.add(
 874         index_select("index_select", "input", axis)
 875     );
 876
 877     network net(engine, topo);
 878
 879     set_values(input, input_data);
 880     net.set_input_data("input", input);
 881
 882     auto outputs = net.execute();
 883     auto output_mem = outputs.at("index_select").get_memory();
 884     auto output_ptr = output_mem.pointer<float>();
 885
 886     for (size_t i = 0; i < output_ptr.size(); i++)
 887     {
 888         EXPECT_EQ(output_ptr[i], out_data[i]);
 889     }
 890 }
 891
 892 TEST(index_select_gpu, reverse_along_f_bfyx)
 893 {
 894     const auto& engine = get_test_engine();
 895     auto input = memory::allocate(engine, { data_types::f32, format::bfyx,{ 1, 2, 3, 4 } });
 896
 897     std::vector<float> input_data = {
 898         0.f,  1.f,  2.f,  3.f,
 899         4.f,  5.f,  6.f,  7.f,
 900         8.f,  9.f, 10.f, 11.f,
 901
 902         12.f, 13.f, 14.f, 15.f,
 903         16.f, 17.f, 18.f, 19.f,
 904         20.f, 21.f, 22.f, 23.f
 905     };
 906
 907     std::vector<float> out_data = {
 908         12.f, 13.f, 14.f, 15.f,
 909         16.f, 17.f, 18.f, 19.f,
 910         20.f, 21.f, 22.f, 23.f,
 911
 912         0.f,  1.f,  2.f,  3.f,
 913         4.f,  5.f,  6.f,  7.f,
 914         8.f,  9.f, 10.f, 11.f
 915     };
 916
 917     constexpr auto axis = index_select_axis_name::along_f;
 918
 919     topology topo;
 920     topo.add(
 921         input_layout("input", input.get_layout())
 922     );
 923     topo.add(
 924         index_select("index_select", "input", axis)
 925     );
 926
 927     network net(engine, topo);
 928
 929     set_values(input, input_data);
 930     net.set_input_data("input", input);
 931
 932     auto outputs = net.execute();
 933     auto output_mem = outputs.at("index_select").get_memory();
 934     auto output_ptr = output_mem.pointer<float>();
 935
 936     for (size_t i = 0; i < output_ptr.size(); i++)
 937     {
 938         EXPECT_EQ(output_ptr[i], out_data[i]);
 939     }
 940 }
 941
 942 TEST(index_select_gpu, reverse_along_y_bfyx)
 943 {
 944     const auto& engine = get_test_engine();
 945     auto input = memory::allocate(engine, { data_types::f32, format::bfyx,{ 1, 2, 4, 3 } });
 946
 947     std::vector<float> input_data = {
 948         0.f,  1.f,  2.f,  3.f,
 949         4.f,  5.f,  6.f,  7.f,
 950         8.f,  9.f, 10.f, 11.f,
 951
 952         12.f, 13.f, 14.f, 15.f,
 953         16.f, 17.f, 18.f, 19.f,
 954         20.f, 21.f, 22.f, 23.f
 955     };
 956
 957     std::vector<float> out_data = {
 958         8.f,  9.f, 10.f, 11.f,
 959         4.f,  5.f,  6.f,  7.f,
 960         0.f,  1.f,  2.f,  3.f,
 961
 962         20.f, 21.f, 22.f, 23.f,
 963         16.f, 17.f, 18.f, 19.f,
 964         12.f, 13.f, 14.f, 15.f
 965     };
 966
 967     constexpr auto axis = index_select_axis_name::along_y;
 968
 969     topology topo;
 970     topo.add(
 971         input_layout("input", input.get_layout())
 972     );
 973     topo.add(
 974         index_select("index_select", "input", axis)
 975     );
 976
 977     network net(engine, topo);
 978
 979     set_values(input, input_data);
 980     net.set_input_data("input", input);
 981
 982     auto outputs = net.execute();
 983     auto output_mem = outputs.at("index_select").get_memory();
 984     auto output_ptr = output_mem.pointer<float>();
 985
 986     for (size_t i = 0; i < output_ptr.size(); i++)
 987     {
 988         EXPECT_EQ(output_ptr[i], out_data[i]);
 989     }
 990 }
 991
 992 TEST(index_select_gpu, reverse_along_x_bfyx)
 993 {
 994     const auto& engine = get_test_engine();
 995     auto input = memory::allocate(engine, { data_types::f32, format::bfyx,{ 1, 2, 4, 3 } });
 996
 997     std::vector<float> input_data = {
 998         0.f,  1.f,  2.f,  3.f,
 999         4.f,  5.f,  6.f,  7.f,
1000         8.f,  9.f, 10.f, 11.f,
1001
1002         12.f, 13.f, 14.f, 15.f,
1003         16.f, 17.f, 18.f, 19.f,
1004         20.f, 21.f, 22.f, 23.f
1005     };
1006
1007     std::vector<float> out_data = {
1008         3.f,  2.f,  1.f,  0.f,
1009         7.f,  6.f,  5.f,  4.f,
1010         11.f,  10.f, 9.f, 8.f,
1011
1012         15.f, 14.f, 13.f, 12.f,
1013         19.f, 18.f, 17.f, 16.f,
1014         23.f, 22.f, 21.f, 20.f
1015     };
1016
1017     constexpr auto axis = index_select_axis_name::along_x;
1018
1019     topology topo;
1020     topo.add(
1021         input_layout("input", input.get_layout())
1022     );
1023     topo.add(
1024         index_select("index_select", "input", axis)
1025     );
1026
1027     network net(engine, topo);
1028
1029     set_values(input, input_data);
1030     net.set_input_data("input", input);
1031
1032     auto outputs = net.execute();
1033     auto output_mem = outputs.at("index_select").get_memory();
1034     auto output_ptr = output_mem.pointer<float>();
1035
1036     for (size_t i = 0; i < output_ptr.size(); i++)
1037     {
1038         EXPECT_EQ(output_ptr[i], out_data[i]);
1039     }
1040 }
1041
1042
1043 TEST(index_select_gpu, reverse_along_y_yxfb)
1044 {
1045     const auto& engine = get_test_engine();
1046     auto input = memory::allocate(engine, { data_types::f32, format::yxfb,{ 4, 2, 2, 2 } });
1047
1048     std::vector<float> input_data = {
1049         0.f,  1.f,  2.f,  3.f,
1050         4.f,  5.f,  6.f,  7.f,
1051
1052         8.f,  9.f, 10.f, 11.f,
1053         12.f, 13.f, 14.f, 15.f,
1054
1055
1056
1057         16.f, 17.f, 18.f, 19.f,
1058         20.f, 21.f, 22.f, 23.f,
1059
1060         24.f, 25.f, 26.f, 27.f,
1061         28.f, 29.f, 30.f, 31.f,
1062     };
1063
1064     std::vector<float> out_data = {
1065         16.f, 17.f, 18.f, 19.f,
1066         20.f, 21.f, 22.f, 23.f,
1067
1068         24.f, 25.f, 26.f, 27.f,
1069         28.f, 29.f, 30.f, 31.f,
1070
1071
1072
1073         0.f,  1.f,  2.f,  3.f,
1074         4.f,  5.f,  6.f,  7.f,
1075
1076         8.f,  9.f, 10.f, 11.f,
1077         12.f, 13.f, 14.f, 15.f,
1078     };
1079
1080     constexpr auto axis = index_select_axis_name::along_y;
1081
1082     topology topo;
1083     topo.add(
1084         input_layout("input", input.get_layout())
1085     );
1086     topo.add(
1087         index_select("index_select", "input", axis)
1088     );
1089
1090     network net(engine, topo);
1091
1092     set_values(input, input_data);
1093     net.set_input_data("input", input);
1094
1095     auto outputs = net.execute();
1096     auto output_mem = outputs.at("index_select").get_memory();
1097     auto output_ptr = output_mem.pointer<float>();
1098
1099     for (size_t i = 0; i < output_ptr.size(); i++)
1100     {
1101         EXPECT_EQ(output_ptr[i], out_data[i]);
1102     }
1103 }
1104
1105 TEST(index_select_gpu, reverse_along_x_yxfb)
1106 {
1107     const auto& engine = get_test_engine();
1108     auto input = memory::allocate(engine, { data_types::f32, format::yxfb,{ 3, 4, 2, 1 } });
1109
1110     std::vector<float> input_data = {
1111         0.f,  1.f,  2.f,  3.f,
1112         4.f,  5.f,  6.f,  7.f,
1113         8.f,  9.f, 10.f, 11.f,
1114
1115         12.f, 13.f, 14.f, 15.f,
1116         16.f, 17.f, 18.f, 19.f,
1117         20.f, 21.f, 22.f, 23.f
1118     };
1119
1120     std::vector<float> out_data = {
1121         12.f, 13.f, 14.f, 15.f,
1122         16.f, 17.f, 18.f, 19.f,
1123         20.f, 21.f, 22.f, 23.f,
1124
1125         0.f,  1.f,  2.f,  3.f,
1126         4.f,  5.f,  6.f,  7.f,
1127         8.f,  9.f, 10.f, 11.f
1128     };
1129
1130     constexpr auto axis = index_select_axis_name::along_x;
1131
1132     topology topo;
1133     topo.add(
1134         input_layout("input", input.get_layout())
1135     );
1136     topo.add(
1137         index_select("index_select", "input", axis)
1138     );
1139
1140     network net(engine, topo);
1141
1142     set_values(input, input_data);
1143     net.set_input_data("input", input);
1144
1145     auto outputs = net.execute();
1146     auto output_mem = outputs.at("index_select").get_memory();
1147     auto output_ptr = output_mem.pointer<float>();
1148
1149     for (size_t i = 0; i < output_ptr.size(); i++)
1150     {
1151         EXPECT_EQ(output_ptr[i], out_data[i]);
1152     }
1153 }
1154
1155 TEST(index_select_gpu, reverse_along_f_yxfb)
1156 {
1157     const auto& engine = get_test_engine();
1158     auto input = memory::allocate(engine, { data_types::f32, format::yxfb,{ 4, 3, 2, 1 } });
1159
1160     std::vector<float> input_data = {
1161         0.f,  1.f,  2.f,  3.f,
1162         4.f,  5.f,  6.f,  7.f,
1163         8.f,  9.f, 10.f, 11.f,
1164
1165         12.f, 13.f, 14.f, 15.f,
1166         16.f, 17.f, 18.f, 19.f,
1167         20.f, 21.f, 22.f, 23.f
1168     };
1169
1170     std::vector<float> out_data = {
1171         8.f,  9.f, 10.f, 11.f,
1172         4.f,  5.f,  6.f,  7.f,
1173         0.f,  1.f,  2.f,  3.f,
1174
1175         20.f, 21.f, 22.f, 23.f,
1176         16.f, 17.f, 18.f, 19.f,
1177         12.f, 13.f, 14.f, 15.f
1178     };
1179
1180     constexpr auto axis = index_select_axis_name::along_f;
1181
1182     topology topo;
1183     topo.add(
1184         input_layout("input", input.get_layout())
1185     );
1186     topo.add(
1187         index_select("index_select", "input", axis)
1188     );
1189
1190     network net(engine, topo);
1191
1192     set_values(input, input_data);
1193     net.set_input_data("input", input);
1194
1195     auto outputs = net.execute();
1196     auto output_mem = outputs.at("index_select").get_memory();
1197     auto output_ptr = output_mem.pointer<float>();
1198
1199     for (size_t i = 0; i < output_ptr.size(); i++)
1200     {
1201         EXPECT_EQ(output_ptr[i], out_data[i]);
1202     }
1203 }
1204
1205 TEST(index_select_gpu, reverse_along_b_yxfb)
1206 {
1207     const auto& engine = get_test_engine();
1208     auto input = memory::allocate(engine, { data_types::f32, format::yxfb,{ 4, 3, 2, 1 } });
1209
1210     std::vector<float> input_data = {
1211         0.f,  1.f,  2.f,  3.f,
1212         4.f,  5.f,  6.f,  7.f,
1213         8.f,  9.f, 10.f, 11.f,
1214
1215         12.f, 13.f, 14.f, 15.f,
1216         16.f, 17.f, 18.f, 19.f,
1217         20.f, 21.f, 22.f, 23.f
1218     };
1219
1220     std::vector<float> out_data = {
1221         3.f,  2.f,  1.f,  0.f,
1222         7.f,  6.f,  5.f,  4.f,
1223         11.f,  10.f, 9.f, 8.f,
1224
1225         15.f, 14.f, 13.f, 12.f,
1226         19.f, 18.f, 17.f, 16.f,
1227         23.f, 22.f, 21.f, 20.f
1228     };
1229
1230     constexpr auto axis = index_select_axis_name::along_b;
1231
1232     topology topo;
1233     topo.add(
1234         input_layout("input", input.get_layout())
1235     );
1236     topo.add(
1237         index_select("index_select", "input", axis)
1238     );
1239
1240     network net(engine, topo);
1241
1242     set_values(input, input_data);
1243     net.set_input_data("input", input);
1244
1245     auto outputs = net.execute();
1246     auto output_mem = outputs.at("index_select").get_memory();
1247     auto output_ptr = output_mem.pointer<float>();
1248
1249     for (size_t i = 0; i < output_ptr.size(); i++)
1250     {
1251         EXPECT_EQ(output_ptr[i], out_data[i]);
1252     }
1253 }
1254
1255
1256 TEST(index_select_gpu, reverse_along_yx_bfyx)
1257 {
1258     const auto& engine = get_test_engine();
1259     auto input = memory::allocate(engine, { data_types::f32, format::bfyx,{ 1, 2, 4, 3 } });
1260
1261     std::vector<float> input_data = {
1262         0.f,  1.f,  2.f,  3.f,
1263         4.f,  5.f,  6.f,  7.f,
1264         8.f,  9.f, 10.f, 11.f,
1265
1266         12.f, 13.f, 14.f, 15.f,
1267         16.f, 17.f, 18.f, 19.f,
1268         20.f, 21.f, 22.f, 23.f
1269     };
1270
1271     std::vector<float> out_data = {
1272         11.f,  10.f, 9.f, 8.f,
1273         7.f,  6.f,  5.f,  4.f,
1274         3.f,  2.f,  1.f,  0.f,
1275
1276         23.f, 22.f, 21.f, 20.f,
1277         19.f, 18.f, 17.f, 16.f,
1278         15.f, 14.f, 13.f, 12.f
1279     };
1280
1281     std::vector<index_select_axis_name> axis = { index_select_axis_name::along_y, index_select_axis_name::along_x };
1282
1283     topology topo;
1284     topo.add(
1285         input_layout("input", input.get_layout())
1286     );
1287     topo.add(
1288         index_select("index_select", "input", axis)
1289     );
1290
1291     network net(engine, topo);
1292
1293     set_values(input, input_data);
1294     net.set_input_data("input", input);
1295
1296     auto outputs = net.execute();
1297     auto output_mem = outputs.at("index_select").get_memory();
1298     auto output_ptr = output_mem.pointer<float>();
1299
1300     for (size_t i = 0; i < output_ptr.size(); i++)
1301     {
1302         EXPECT_EQ(output_ptr[i], out_data[i]);
1303     }
1304 }
1305
1306 TEST(index_select_gpu, reverse_along_fyx_bfyx)
1307 {
1308     const auto& engine = get_test_engine();
1309     auto input = memory::allocate(engine, { data_types::f32, format::bfyx,{ 1, 2, 4, 3 } });
1310
1311     std::vector<float> input_data = {
1312         0.f,  1.f,  2.f,  3.f,
1313         4.f,  5.f,  6.f,  7.f,
1314         8.f,  9.f, 10.f, 11.f,
1315
1316         12.f, 13.f, 14.f, 15.f,
1317         16.f, 17.f, 18.f, 19.f,
1318         20.f, 21.f, 22.f, 23.f
1319     };
1320
1321     std::vector<float> out_data = {
1322         23.f, 22.f, 21.f, 20.f,
1323         19.f, 18.f, 17.f, 16.f,
1324         15.f, 14.f, 13.f, 12.f,
1325
1326         11.f,  10.f, 9.f, 8.f,
1327         7.f,  6.f,  5.f,  4.f,
1328         3.f,  2.f,  1.f,  0.f
1329     };
1330
1331     std::vector<index_select_axis_name> axis = { index_select_axis_name::along_f, index_select_axis_name::along_y, index_select_axis_name::along_x };
1332
1333     topology topo;
1334     topo.add(
1335         input_layout("input", input.get_layout())
1336     );
1337     topo.add(
1338         index_select("index_select", "input", axis)
1339     );
1340
1341     network net(engine, topo);
1342
1343     set_values(input, input_data);
1344     net.set_input_data("input", input);
1345
1346     auto outputs = net.execute();
1347     auto output_mem = outputs.at("index_select").get_memory();
1348     auto output_ptr = output_mem.pointer<float>();
1349
1350     for (size_t i = 0; i < output_ptr.size(); i++)
1351     {
1352         EXPECT_EQ(output_ptr[i], out_data[i]);
1353     }
1354 }
1355
1356 TEST(index_select_gpu, reverse_along_bfyx_bfyx)
1357 {
1358     const auto& engine = get_test_engine();
1359     auto input = memory::allocate(engine, { data_types::f32, format::bfyx,{ 3, 3, 4, 3 } });
1360
1361     std::vector<float> input_data = {
1362         // b0f0
1363         0.f,  1.f,  2.f,  3.f,
1364         4.f,  5.f,  6.f,  7.f,
1365         8.f,  9.f, 10.f, 11.f,
1366         // f1
1367         12.f, 13.f, 14.f, 15.f,
1368         16.f, 17.f, 18.f, 19.f,
1369         20.f, 21.f, 22.f, 23.f,
1370         // f2
1371         24.f, 25.f, 26.f, 27.f,
1372         28.f, 29.f, 30.f, 31.f,
1373         32.f, 33.f, 34.f, 35.f,
1374
1375         // b1f0
1376         36.f, 37.f, 38.f, 39.f,
1377         40.f, 41.f, 42.f, 43.f,
1378         44.f, 45.f, 46.f, 47.f,
1379         // f1
1380         48.f, 49.f, 50.f, 51.f,
1381         52.f, 53.f, 54.f, 55.f,
1382         56.f, 57.f, 58.f, 59.f,
1383         // f2
1384         60.f, 61.f, 62.f, 63.f,
1385         64.f, 65.f, 66.f, 67.f,
1386         68.f, 69.f, 70.f, 71.f,
1387
1388         // b2f0
1389         72.f, 73.f, 74.f, 75.f,
1390         76.f, 77.f, 78.f, 79.f,
1391         80.f, 81.f, 82.f, 83.f,
1392         // f1
1393         84.f, 85.f, 86.f, 87.f,
1394         88.f, 89.f, 90.f, 91.f,
1395         92.f, 93.f, 94.f, 95.f,
1396         // f2
1397         96.f, 97.f, 98.f, 99.f,
1398         100.f, 101.f, 102.f, 103.f,
1399         104.f, 105.f, 106.f, 107.f
1400     };
1401
1402     std::vector<float> out_data = {
1403         107.f, 106.f, 105.f, 104.f,
1404         103.f, 102.f, 101.f, 100.f,
1405         99.f, 98.f, 97.f, 96.f,
1406
1407         95.f, 94.f, 93.f, 92.f,
1408         91.f, 90.f, 89.f, 88.f,
1409         87.f, 86.f, 85.f, 84.f,
1410
1411         83.f, 82.f, 81.f, 80.f,
1412         79.f, 78.f, 77.f, 76.f,
1413         75.f, 74.f, 73.f, 72.f,
1414
1415
1416         71.f, 70.f, 69.f, 68.f,
1417         67.f, 66.f, 65.f, 64.f,
1418         63.f, 62.f, 61.f, 60.f,
1419
1420         59.f, 58.f, 57.f, 56.f,
1421         55.f, 54.f, 53.f, 52.f,
1422         51.f, 50.f, 49.f, 48.f,
1423
1424         47.f, 46.f, 45.f, 44.f,
1425         43.f, 42.f, 41.f, 40.f,
1426         39.f, 38.f, 37.f, 36.f,
1427
1428
1429         35.f, 34.f, 33.f, 32.f,
1430         31.f, 30.f, 29.f, 28.f,
1431         27.f, 26.f, 25.f, 24.f,
1432
1433         23.f, 22.f, 21.f, 20.f,
1434         19.f, 18.f, 17.f, 16.f,
1435         15.f, 14.f, 13.f, 12.f,
1436
1437         11.f,  10.f, 9.f, 8.f,
1438         7.f,  6.f,  5.f,  4.f,
1439         3.f,  2.f,  1.f,  0.f
1440     };
1441
1442     std::vector<index_select_axis_name> axis = { index_select_axis_name::along_b, index_select_axis_name::along_f, index_select_axis_name::along_y, index_select_axis_name::along_x };
1443
1444     topology topo;
1445     topo.add(
1446         input_layout("input", input.get_layout())
1447     );
1448     topo.add(
1449         index_select("index_select", "input", axis)
1450     );
1451
1452     network net(engine, topo);
1453
1454     set_values(input, input_data);
1455     net.set_input_data("input", input);
1456
1457     auto outputs = net.execute();
1458     auto output_mem = outputs.at("index_select").get_memory();
1459     auto output_ptr = output_mem.pointer<float>();
1460
1461     for (size_t i = 0; i < output_ptr.size(); i++)
1462     {
1463         EXPECT_EQ(output_ptr[i], out_data[i]);
1464     }
1465 }
1466
1467 TEST(index_select_gpu, reverse_along_bfx_yxfb)
1468 {
1469     const auto& engine = get_test_engine();
1470     auto input = memory::allocate(engine, { data_types::f32, format::yxfb,{ 4, 3, 3, 3 } });
1471
1472     std::vector<float> input_data = {
1473         // y0x0
1474         0.f,  1.f,  2.f,  3.f,
1475         4.f,  5.f,  6.f,  7.f,
1476         8.f,  9.f, 10.f, 11.f,
1477         // x1
1478         12.f, 13.f, 14.f, 15.f,
1479         16.f, 17.f, 18.f, 19.f,
1480         20.f, 21.f, 22.f, 23.f,
1481         // x2
1482         24.f, 25.f, 26.f, 27.f,
1483         28.f, 29.f, 30.f, 31.f,
1484         32.f, 33.f, 34.f, 35.f,
1485
1486         // y1x0
1487         36.f, 37.f, 38.f, 39.f,
1488         40.f, 41.f, 42.f, 43.f,
1489         44.f, 45.f, 46.f, 47.f,
1490         // x1
1491         48.f, 49.f, 50.f, 51.f,
1492         52.f, 53.f, 54.f, 55.f,
1493         56.f, 57.f, 58.f, 59.f,
1494         // x2
1495         60.f, 61.f, 62.f, 63.f,
1496         64.f, 65.f, 66.f, 67.f,
1497         68.f, 69.f, 70.f, 71.f,
1498
1499         // y2x0
1500         72.f, 73.f, 74.f, 75.f,
1501         76.f, 77.f, 78.f, 79.f,
1502         80.f, 81.f, 82.f, 83.f,
1503         // x1
1504         84.f, 85.f, 86.f, 87.f,
1505         88.f, 89.f, 90.f, 91.f,
1506         92.f, 93.f, 94.f, 95.f,
1507         // x2
1508         96.f, 97.f, 98.f, 99.f,
1509         100.f, 101.f, 102.f, 103.f,
1510         104.f, 105.f, 106.f, 107.f
1511     };
1512
1513     std::vector<float> out_data = {
1514         35.f, 34.f, 33.f, 32.f,
1515         31.f, 30.f, 29.f, 28.f,
1516         27.f, 26.f, 25.f, 24.f,
1517
1518         23.f, 22.f, 21.f, 20.f,
1519         19.f, 18.f, 17.f, 16.f,
1520         15.f, 14.f, 13.f, 12.f,
1521
1522         11.f,  10.f, 9.f, 8.f,
1523         7.f,  6.f,  5.f,  4.f,
1524         3.f,  2.f,  1.f,  0.f,
1525
1526
1527         71.f, 70.f, 69.f, 68.f,
1528         67.f, 66.f, 65.f, 64.f,
1529         63.f, 62.f, 61.f, 60.f,
1530
1531         59.f, 58.f, 57.f, 56.f,
1532         55.f, 54.f, 53.f, 52.f,
1533         51.f, 50.f, 49.f, 48.f,
1534
1535         47.f, 46.f, 45.f, 44.f,
1536         43.f, 42.f, 41.f, 40.f,
1537         39.f, 38.f, 37.f, 36.f,
1538
1539
1540         107.f, 106.f, 105.f, 104.f,
1541         103.f, 102.f, 101.f, 100.f,
1542         99.f, 98.f, 97.f, 96.f,
1543
1544         95.f, 94.f, 93.f, 92.f,
1545         91.f, 90.f, 89.f, 88.f,
1546         87.f, 86.f, 85.f, 84.f,
1547
1548         83.f, 82.f, 81.f, 80.f,
1549         79.f, 78.f, 77.f, 76.f,
1550         75.f, 74.f, 73.f, 72.f
1551     };
1552
1553     std::vector<index_select_axis_name> axis = { index_select_axis_name::along_f, index_select_axis_name::along_b, index_select_axis_name::along_x };
1554
1555     topology topo;
1556     topo.add(
1557         input_layout("input", input.get_layout())
1558     );
1559     topo.add(
1560         index_select("index_select", "input", axis)
1561     );
1562
1563     network net(engine, topo);
1564
1565     set_values(input, input_data);
1566     net.set_input_data("input", input);
1567
1568     auto outputs = net.execute();
1569     auto output_mem = outputs.at("index_select").get_memory();
1570     auto output_ptr = output_mem.pointer<float>();
1571
1572     for (size_t i = 0; i < output_ptr.size(); i++)
1573     {
1574         EXPECT_EQ(output_ptr[i], out_data[i]);
1575     }
1576 }
1577
1578 TEST(index_select_gpu, reverse_along_bfyx_yxfb)
1579 {
1580     const auto& engine = get_test_engine();
1581     auto input = memory::allocate(engine, { data_types::f32, format::yxfb,{ 4, 3, 3, 3 } });
1582
1583     std::vector<float> input_data = {
1584         // y0x0
1585         0.f,  1.f,  2.f,  3.f,
1586         4.f,  5.f,  6.f,  7.f,
1587         8.f,  9.f, 10.f, 11.f,
1588         // x1
1589         12.f, 13.f, 14.f, 15.f,
1590         16.f, 17.f, 18.f, 19.f,
1591         20.f, 21.f, 22.f, 23.f,
1592         // x2
1593         24.f, 25.f, 26.f, 27.f,
1594         28.f, 29.f, 30.f, 31.f,
1595         32.f, 33.f, 34.f, 35.f,
1596
1597         // y1x0
1598         36.f, 37.f, 38.f, 39.f,
1599         40.f, 41.f, 42.f, 43.f,
1600         44.f, 45.f, 46.f, 47.f,
1601         // x1
1602         48.f, 49.f, 50.f, 51.f,
1603         52.f, 53.f, 54.f, 55.f,
1604         56.f, 57.f, 58.f, 59.f,
1605         // x2
1606         60.f, 61.f, 62.f, 63.f,
1607         64.f, 65.f, 66.f, 67.f,
1608         68.f, 69.f, 70.f, 71.f,
1609
1610         // y2x0
1611         72.f, 73.f, 74.f, 75.f,
1612         76.f, 77.f, 78.f, 79.f,
1613         80.f, 81.f, 82.f, 83.f,
1614         // x1
1615         84.f, 85.f, 86.f, 87.f,
1616         88.f, 89.f, 90.f, 91.f,
1617         92.f, 93.f, 94.f, 95.f,
1618         // x2
1619         96.f, 97.f, 98.f, 99.f,
1620         100.f, 101.f, 102.f, 103.f,
1621         104.f, 105.f, 106.f, 107.f
1622     };
1623
1624     std::vector<float> out_data = {
1625         107.f, 106.f, 105.f, 104.f,
1626         103.f, 102.f, 101.f, 100.f,
1627         99.f, 98.f, 97.f, 96.f,
1628
1629         95.f, 94.f, 93.f, 92.f,
1630         91.f, 90.f, 89.f, 88.f,
1631         87.f, 86.f, 85.f, 84.f,
1632
1633         83.f, 82.f, 81.f, 80.f,
1634         79.f, 78.f, 77.f, 76.f,
1635         75.f, 74.f, 73.f, 72.f,
1636
1637
1638         71.f, 70.f, 69.f, 68.f,
1639         67.f, 66.f, 65.f, 64.f,
1640         63.f, 62.f, 61.f, 60.f,
1641
1642         59.f, 58.f, 57.f, 56.f,
1643         55.f, 54.f, 53.f, 52.f,
1644         51.f, 50.f, 49.f, 48.f,
1645
1646         47.f, 46.f, 45.f, 44.f,
1647         43.f, 42.f, 41.f, 40.f,
1648         39.f, 38.f, 37.f, 36.f,
1649
1650
1651         35.f, 34.f, 33.f, 32.f,
1652         31.f, 30.f, 29.f, 28.f,
1653         27.f, 26.f, 25.f, 24.f,
1654
1655         23.f, 22.f, 21.f, 20.f,
1656         19.f, 18.f, 17.f, 16.f,
1657         15.f, 14.f, 13.f, 12.f,
1658
1659         11.f,  10.f, 9.f, 8.f,
1660         7.f,  6.f,  5.f,  4.f,
1661         3.f,  2.f,  1.f,  0.f
1662     };
1663
1664     std::vector<index_select_axis_name> axis = { index_select_axis_name::along_b, index_select_axis_name::along_f, index_select_axis_name::along_y, index_select_axis_name::along_x };
1665
1666     topology topo;
1667     topo.add(
1668         input_layout("input", input.get_layout())
1669     );
1670     topo.add(
1671         index_select("index_select", "input", axis)
1672     );
1673
1674     network net(engine, topo);
1675
1676     set_values(input, input_data);
1677     net.set_input_data("input", input);
1678
1679     auto outputs = net.execute();
1680     auto output_mem = outputs.at("index_select").get_memory();
1681     auto output_ptr = output_mem.pointer<float>();
1682
1683     for (size_t i = 0; i < output_ptr.size(); i++)
1684     {
1685         EXPECT_EQ(output_ptr[i], out_data[i]);
1686     }
1687 }