inference-engine/thirdparty/clDNN/kernel_selector/core/cl_kernels/index_select_gpu_ref.cl

   1 // Copyright (c) 2018 Intel Corporation
   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 #include "include/include_all.cl"
  16
  17
  18 KERNEL(index_select_gpu_ref)(
  19     const __global UNIT_TYPE* input,
  20 #ifndef REVERSE
  21     const __global int* indices,
  22 #endif
  23     __global UNIT_TYPE* output)
  24 {
  25     // [CONSTEXPR]:
  26     const uint input_sx  = INPUT0_SIZE_X;
  27     const uint input_sy  = INPUT0_SIZE_Y;
  28     const uint input_sf  = INPUT0_FEATURE_NUM;
  29     const uint input_sb  = INPUT0_BATCH_NUM;
  30
  31     const uint out_b         = (uint) get_global_id(0);
  32     const uint indices_idx   = (uint) get_global_id(1);
  33     const uint feature_idx   = (uint) get_global_id(2);
  34
  35     #if AXES_NUMBER == 1
  36         #ifdef REVERSE
  37             const uint indices_value = REVERSE_AXIS_SIZE - 1 - indices_idx;
  38         #else
  39             const uint indices_value = indices[indices_idx];
  40         #endif
  41     #elif AXES_NUMBER > 1
  42         #ifdef REVERSE
  43             uint indices_value[4] =  {
  44                 #ifdef REVERSE_INDEX_SELECT_AXIS_BATCH_SIZE
  45                     REVERSE_INDEX_SELECT_AXIS_BATCH_SIZE - 1 - out_b,
  46                 #else
  47                     out_b,
  48                 #endif
  49                 #ifdef REVERSE_INDEX_SELECT_AXIS_FEATURE_SIZE
  50                     REVERSE_INDEX_SELECT_AXIS_FEATURE_SIZE - 1 - feature_idx,
  51                 #else
  52                     feature_idx,
  53                 #endif
  54                 #ifdef REVERSE_INDEX_SELECT_AXIS_Y_SIZE
  55                     REVERSE_INDEX_SELECT_AXIS_Y_SIZE - 1 - indices_idx,
  56                 #else
  57                     indices_idx,
  58                 #endif
  59                     0
  60              };
  61         #endif
  62     #endif
  63
  64     // [LOGIC]:
  65     #if AXES_NUMBER > 1
  66         for(uint x = 0; x < input_sx; x++)
  67         {
  68             #ifdef REVERSE_INDEX_SELECT_AXIS_X_SIZE
  69                 indices_value[3] = REVERSE_INDEX_SELECT_AXIS_X_SIZE - 1 - x;
  70             #else
  71                 indices_value[3] = x;
  72             #endif
  73             output[GET_DATA_INDEX(OUTPUT, out_b, feature_idx, indices_idx, x)] = input[GET_DATA_INDEX(INPUT0, indices_value[0], indices_value[1], indices_value[2], indices_value[3])];
  74         }
  75
  76     #else
  77         #ifdef INDEX_SELECT_AXIS_BATCH
  78             for(uint x = 0; x < input_sx; x++)
  79             {
  80                 for(uint y = 0; y < input_sy; y++)
  81                 {
  82                     output[GET_DATA_INDEX(OUTPUT, indices_idx, feature_idx, y, x)] = input[GET_DATA_INDEX(INPUT0, indices_value, feature_idx, y, x)];
  83                 }
  84             }
  85         #elif defined INDEX_SELECT_AXIS_FEATURE
  86             for(uint x = 0; x < input_sx; x++)
  87             {
  88                 output[GET_DATA_INDEX(OUTPUT, out_b, indices_idx, feature_idx, x)] = input[GET_DATA_INDEX(INPUT0, out_b, indices_value, feature_idx, x)];
  89             }
  90         #elif defined INDEX_SELECT_AXIS_X
  91             for(uint i = 0; i < input_sy; i++)
  92             {
  93                 output[GET_DATA_INDEX(OUTPUT, out_b, feature_idx, i, indices_idx)] = input[GET_DATA_INDEX(INPUT0, out_b, feature_idx, i, indices_value)];
  94             }
  95         #elif defined INDEX_SELECT_AXIS_Y
  96
  97             for(uint i = 0; i < input_sx; i++)
  98             {
  99                 output[GET_DATA_INDEX(OUTPUT, out_b, feature_idx, indices_idx, i)] = input[GET_DATA_INDEX(INPUT0, out_b, feature_idx, indices_value, i)];
 100             }
 101         #endif
 102     #endif
 103 }