inference-engine/thirdparty/mkl-dnn/src/cpu/ref_shuffle.hpp

   1 /*******************************************************************************
   2 * Copyright 2018 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 #ifndef CPU_REF_SHUFFLE_HPP
  18 #define CPU_REF_SHUFFLE_HPP
  19
  20 #include <assert.h>
  21
  22 #include "c_types_map.hpp"
  23 #include "cpu_shuffle_pd.hpp"
  24 #include "cpu_engine.hpp"
  25 #include "type_helpers.hpp"
  26 #include "utils.hpp"
  27
  28 namespace mkldnn {
  29 namespace impl {
  30 namespace cpu {
  31
  32 template<int data_type_size>
  33 struct ref_shuffle_t : public cpu_primitive_t {
  34     using shuffle_class = ref_shuffle_t<data_type_size>;
  35
  36     struct pd_t: public cpu_shuffle_pd_t {
  37         pd_t(engine_t *engine, const shuffle_desc_t *adesc,
  38                 const primitive_attr_t *attr,
  39                 const shuffle_pd_t *hint_fwd_pd)
  40             : cpu_shuffle_pd_t(engine, adesc, attr, hint_fwd_pd) {}
  41
  42         DECLARE_COMMON_PD_T("ref:any",shuffle_class);
  43
  44         virtual status_t init() override {
  45             assert(this->engine()->kind() == engine_kind::cpu);
  46
  47             bool ok = true
  48                  && data_type_size ==
  49                      types::data_type_size(this->desc()->data_desc.data_type);
  50             if (!ok)
  51                 return status::unimplemented;
  52             return status::success;
  53         }
  54     };
  55
  56     ref_shuffle_t(const pd_t *apd, const input_vector &inputs,
  57             const output_vector &outputs)
  58         : cpu_primitive_t(apd, inputs, outputs)
  59     {
  60         const int axis_size = pd()->axis_size();
  61         const int group_size = pd()->group_size();
  62         const int transpose_row = pd()->is_fwd() ? group_size
  63                                                  : axis_size / group_size;
  64         const int transpose_col = pd()->is_fwd() ? axis_size / group_size
  65                                                  : group_size;
  66         rev_transposed_ = (int *)malloc(axis_size * sizeof(int), 64);
  67         parallel_nd(transpose_col, transpose_row, [&](int i, int j) {
  68             rev_transposed_[j * transpose_col + i] = i * transpose_row + j;
  69         });
  70     }
  71
  72     ~ref_shuffle_t() { free(rev_transposed_); }
  73
  74     typedef typename typesize_traits<data_type_size>::type data_t;
  75
  76     virtual void execute(event_t *e) const {
  77         using namespace memory_format;
  78         switch (pd()->data_pd()->desc()->format) {
  79         case nCdhw16c: execute_<nCdhw16c>(); break;
  80         case nChw16c:  execute_<nChw16c>(); break;
  81         case nCdhw8c:  execute_<nCdhw8c>(); break;
  82         case nChw8c:   execute_<nChw8c>(); break;
  83         case ncdhw:    execute_<ncdhw>(); break;
  84         case nchw:     execute_<nchw>(); break;
  85         case ndhwc:    execute_<ndhwc>(); break;
  86         case nhwc:     execute_<nhwc>(); break;
  87         default:       execute_<mkldnn_any>(); break;
  88         }
  89
  90         e->set_state(event_t::ready);
  91     }
  92
  93 private:
  94     template<memory_format_t fmt>void execute_() const;
  95     const pd_t *pd() const { return (const pd_t *)primitive_t::pd(); }
  96     int *rev_transposed_;
  97 };
  98
  99 }
 100 }
 101 }
 102
 103 #endif
 104
 105 // vim: et ts=4 sw=4 cindent cino^=l0,\:0,N-s