text: ﻿// Copyright (C) 2018-2022 Intel Corporation
  // SPDX-License-Identifier: Apache-2.0
  //
  
  #pragma once
  
  #include "common_types.h"
  #include "common_tools.h"
  #include <vector>
  #include <assert.h>
  #include <numeric>
  #include <cstddef>
  #include <algorithm>
  #include <array>
  #include <string>
  #include <utility>
  #include <stdexcept>
  
  namespace kernel_selector {
  #define KERNEL_SELECTOR_TENSOR_DIM_MAX 9
  
  namespace Tensor {
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  // DataLayout
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  enum DataLayout {
      f = 0,                  // 1D
      bf,                     // 1D+batch
      fb,                     // 1D+batch
      bfyx,                   // 3D+batch
      yxfb,                   // 3D+batch
      byxf,                   // 3D+batch
      fyxb,                   // 3D+batch
      bfxy,                   // 3D+batch
      b_fs_yx_fsv2,
      b_fs_yx_fsv4,           // reordering format for swizzled input for convolution using IMAD
      b_fs_yx_fsv16,          // 3D+batch
      b_fs_zyx_fsv16,         // batch, feature, 3D spatial. Blocks of 16 input channels
      b_fs_yx_fsv32,          // 3D+batch
      b_fs_zyx_fsv32,         // 4D+batch
      bs_fs_yx_bsv16_fsv16,   // batch, feature, 2D spatial. Blocks of 16 batch and channels
      bs_fs_zyx_bsv16_fsv16,  // batch, feature, 3D spatial. Blocks of 16 batch and channels
      bs_fs_yx_bsv4_fsv4,     // batch, feature, 2D spatial. Blocks of 4 batch and 4 channels
      bs_fs_yx_bsv8_fsv4,     // batch, feature, 2D spatial. Blocks of 8 batch and 4 channels
      bs_fs_yx_bsv8_fsv2,     // batch, feature, 2D spatial. Blocks of 8 batch and 2 channels
      bs_fs_yx_bsv4_fsv2,     // batch, feature, 2D spatial. Blocks of 4 batch and 2 channels
      bs_fs_yx_bsv32_fsv32,   // batch, feature, 2D spatial. Blocks of 32 batch and 32 channels
      bs_fs_yx_bsv32_fsv16,   // batch, feature, 2D spatial. Blocks of 32 batch and 16 channels
      bs_f_bsv8__af8,         // for optimized FC
      bs_f_bsv16__af8,        // for optimized FC
      winograd_2x3_s1_data,   // winograd convolution input, F(2,3) -- filter 3x3 with stride 1
      bfzyx,                  // batch+feature+3D spatial
      fs_b_yx_fsv32,          // for FP16 kernels, 32 features to avoid partial writes
      b_fs_yx_32fp,           // bfyx with blocks of 16 packed binary input channels
      bfwzyx,                 // batch, feature, 4D spatial
      nv12,                   // media nv12 layout
      image_2d_rgba,          // image2d RGBA
      DataLayoutCount         // NUMBER OF ELEMENTS IN ENUM
  };
  
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  // WeightsLayout
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  enum WeightsLayout {
      oi = 0,
      io,
      oiyx,
      ioyx,
      oyxi,
      iyxo,
      yxio,
      o_is_yx_isv16,
      os_yxi_osv16,
      os_iyx_osv16,
      os_iyx_osv32,
      os_iyx_osv32__ai32,
      os_iyx_osv64,
      os_is_zyx_isv16_osv16,
      is_os_zyx_isv16_osv16,
      is_os_yx_isv16_osv16,
      os_is_zyx_isv8_osv16_isv2,
      os_is_yx_isv8_osv16_isv2,
      os_is_yx_isv16_osv16,
      os_zyxi_osv16,
      os_iyx_osv16_rotate_180,
      os_i_osv8__ai8,  // TODO can we drop the alignment form layout name?
      os_i_osv16__ai8,
      os_i_osv16,
      os_is_yx_osv16_isv16,           // wieghts for int8 blocked conv
      os_is_zyx_osv16_isv16,
      os_is_zyx_osv32_isv16,
      os_is_zyx_osv64_isv16,
      i_yxs_os_yxsv2_osv16,
      iy_xs_os_xsv2_osv16__ao32,
      iy_xs_os_xsv2_osv8__ao32,
      image_2d_weights_c4_fyx_b,      // image type fyx_b
      image_2d_weights_c1_b_fyx,      // image type 2d b_fyx single channel
      winograd_2x3_s1_weights,        // winograd convolution weights, F(2, 3) --filter 3x3 with stride 1
      winograd_2x3_s1_fused_weights,  // winograd convolution weights for fused kernel, F(2, 3) --filter 3x3 with stride 1
      winograd_6x3_s1_fused_weights,  // winograd convolution weights for fused kernel, F(6, 3) --filter 3x3 with stride 1
      image_2d_weights_winograd_6x3_s1_fbxyb,  // image 2d winograd convolution weights for fused kernel, F(2, 3) --filter
                                               // 3x3 with stride 1
      image_2d_weights_winograd_6x3_s1_xfbyb,  // image 2d winograd convolution weights for fused kernel, F(2, 3) --filter
                                               // 3x3 with stride 1
      dlstm_dir_io,                            // dlstm weights layout direction, input_size, 4* hiden_size
      os_is_yx_isa8_osv8_isv4,                 // for MMAD convolution
      os_is_zyx_isa8_osv8_isv4,                // for MMAD convolution
      os_is_yx_isa8_osv16_isv4,                // for fully connected MMAD
      os_is_zyx_isa8_osv16_isv4,               // for fully connected MMAD
      os_is_yx_osa4_isa8_osv8_isv4,            // for MMAD convolution swizzled from ofm 0..7 to 0,4,8,12,16,20,24,28,
      is_os_yx_isa4_osa8_isv8_osv4,            // for onednn deconvolution
      is_os_yx_osa4_isa8_osv8_isv4,
      g_os_is_yx_osa2_isa8_osv8_isv2,          // for MMAD convolution swizzled from ofm 0..7 to 0,4,8,12,16,20,24,28,
      g_os_is_yx_osa4_isa8_osv8_isv4,          // for MMAD convolution swizzled from ofm 0..7 to 0,4,8,12,16,20,24,28,
      g_os_is_zyx_osa4_isa8_osv8_isv4,         // for MMAD convolution swizzled from ofm 0..7 to 0,4,8,12,16,20,24,28,
      os_is_yx_osa4_isa8_osv8_isv2,            // for MMAD convolution swizzled from ofm 0..7 to 0,4,8,12,16,20,24,28,
      os_is_zyx_osa4_isa8_osv8_isv2,           // for MMAD convolution swizzled from ofm 0..7 to 0,4,8,12,16,20,24,28,
      os_is_zyx_osa4_isa8_osv8_isv4,           // for MMAD convolution swizzled from ofm 0..7 to 0,4,8,12,16,20,24,28,
      g_os_is_yx_osa4_isa8_osv8_isv2,          // for MMAD convolution swizzled from ofm 0..7 to 0,4,8,12,16,20,24,28,
      g_os_is_zyx_osa4_isa8_osv8_isv2,         // for MMAD convolution swizzled from ofm 0..7 to 0,4,8,12,16,20,24,28,
      os_is_yx_osa2_isa8_osv8_isv2,
      os_is_yx_osa2_isa8_osv16_isv4,
      os_is_yx_osa2_isa8_osv16_isv2,
      os_is_yx_isa8_osv8_isv2,
      is_os_yx_isa8_osv8_isv2,
      is_os_yx_isa2_osa8_isv8_osv2,
      g_os_is_yx_osa2_isa8_osv16_isv4,
      g_os_is_yx_osa2_isa8_osv16_isv2,
      os_is_yx_osa4_isa8_osv8_isv4_swizzled_by_4,  // for MMAD convolution swizzled from ofm 0..7 to 0,4,8,12,16,20,24,28,
                                                   // 1,5...
      os_is_zyx_osa4_isa8_osv8_isv4_swizzled_by_4,  // for MMAD convolution swizzled from ofm 0..7 to 0,4,8,12,16,20,24,28,
                                                    // 1,5...
      os_is_yx_isa8_osv8_isv4_swizzled_by_4,   // for MMAD convolution swizzled from ofm 0..7 to 0,4,8,12,16,20,24,28,
                                               // 1,5...
      is_o_yx_isv32,                           // for MMAD 1x1 convolutions
      is_o32_yx_isv32_swizzled_by_4,  // for MMAD 1x1 convolutions swizzled from ofm 0..7 to 0,4,8,12,16,20,24,28, 1,5...
      os_is_y_x8_osv8_isv4,           // for MMAD convolutions
      os_is_y_x8_osv8_isv4_swizzled_by_4,  // for MMAD 1x1 convolutions swizzled from ofm 0..7 to 0,4,8,12,16,20,24,28,
                                           // 1,5...
      os_is_yx_osv16_isv4,                 // swizzled weights for convolution using IMAD
      os_is_yx_osv8_isv4,                      // weights for int8 blocked conv
      os_is_yx_osv8_isv2,                      // weights for int8 blocked conv
      os_is_yx_osv32_isv4_swizzled_by_2,   //  weights for bfyx -> b_fs_yx_fsv32 convolution using IMAD with swizzled ofm (0, 2, 4..), (1, 3, 5...)
      os_is_yx_osv32_isv4,                 //  weights for bfyx -> b_fs_yx_fsv{32,16} convolution using IMAD
      os_is_zyx_osv32_isv4,                //  weights for bfzyx -> b_fs_zyx_fsv16 convolution using IMAD
      oizyx,
      iozyx,
      os_is_yx_osv32_isv32p,  // 2 blocks: 32 packed binary in channels and 32 output channels
      os_is_osv32_isv32_swizzled_by_4,     // for weights for 1x1 IMAD convolution
      os_i_yxs_osv4_yxsv4,                 // for weights for depthwise IMAD convolution
      goiyx,
      gioyx,
      goizyx,
      giozyx,
      gyxio,
      g_os_iyx_osv16,
      g_os_iyx_osv32,
      gs_oiyx_gsv16,
      gs_oizyx_gsv16,
      gs_oiyx_gsv32,
      g_os_iyx_osv16_rotate_180,
      gi_yxs_os_yxsv2_osv16,
      g_is_os_zyx_isv16_osv16,
      g_is_os_yx_isv16_osv16,
      g_os_is_yx_isa8_osv8_isv2,
      g_os_is_zyx_isv8_osv16_isv2,
      g_os_is_yx_isv8_osv16_isv2,
      g_os_is_zyx_isv16_osv16,
      g_os_is_zyx_osv16_isv16,
      giy_xs_os_xsv2_osv16__ao32,
      giy_xs_os_xsv2_osv8__ao32,
      g_os_is_yx_isv16_osv16,
      gs_oi_yxs_gsv4_yxsv4,                // grouped weights for depthwise IMAD convolution (b_fs_yx_fsv4 format)
      gs_oi_yxs_gsv16_yxsv4,               // grouped weights for depthwise IMAD convolution (b_fs_yx_fsv16 format)
      gs_oi_yxs_gsv32_yxsv4,               // grouped weights for depthwise IMAD convolution (b_fs_yx_fsv32 format)
      g_os_is_yx_osv8_isv2,
      g_os_is_yx_osv8_isv4,
      g_os_is_yx_osv16_isv4,
  
      g_os_zyx_is_osv16_isv4,
      g_os_zyx_is_osv16_isv16,
      g_os_zyx_is_osv16_isv32,
      g_os_zyx_is_osv32_isv4,
      g_os_zyx_is_osv32_isv16,
      g_os_zyx_is_osv32_isv32,
  
      WeightsLayoutCount                   // NUMBER OF ELEMENTS IN ENUM
  };
  
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  // Pad
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  struct Pad {
      size_t before;
      size_t after;
  
      size_t Total() const { return before + after; }
  };
  
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  // Dim
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  struct Dim {
      size_t v;
      size_t pitch;
      Pad pad;
  
      size_t LogicalDimPadded() const { return v + pad.Total(); }
  };
  
  using NDims = std::vector<Dim>;
  
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  // extract code
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  enum class DataChannelName { X = 0, Y = 1, Z = 2, W = 3, FEATURE = 4, BATCH = 5, COUNT = 6 };
  
  enum class WeightsChannelName { X = 0, Y = 1, Z = 2, IFM = 3, OFM = 4, G = 5, COUNT = 6 };
  
  inline bool SimpleLayout(WeightsLayout l) {
      switch (l) {
          case WeightsLayout::oi:
          case WeightsLayout::io:
          case WeightsLayout::oiyx:
          case WeightsLayout::ioyx:
          case WeightsLayout::oyxi:
          case WeightsLayout::iyxo:
          case WeightsLayout::yxio:
          case WeightsLayout::oizyx:
          case WeightsLayout::iozyx:
          case WeightsLayout::dlstm_dir_io:
              return true;
          default:
              return false;
      }
  }
  
  inline bool SimpleLayout(DataLayout l) {
      switch (l) {
          case DataLayout::bf:
          case DataLayout::fb:
          case DataLayout::bfyx:
          case DataLayout::yxfb:
          case DataLayout::byxf:
          case DataLayout::fyxb:
          case DataLayout::bfxy:
          case DataLayout::bfzyx:
          case DataLayout::bfwzyx:
              return true;
          default:
              return false;
      }
  }
  
  inline bool GroupedLayout(WeightsLayout l);
  
  inline bool GroupedLayout(DataLayout) {
      return false;
  }
  
  inline bool IsImageType(WeightsLayout l) {
      switch (l) {
          case WeightsLayout::image_2d_weights_c4_fyx_b:
          case WeightsLayout::image_2d_weights_c1_b_fyx:
          case WeightsLayout::image_2d_weights_winograd_6x3_s1_fbxyb:
          case WeightsLayout::image_2d_weights_winograd_6x3_s1_xfbyb:
              return true;
          default:
              return false;
      }
  }
  
  inline bool IsDynamicLSTMType(WeightsLayout l) {
      switch (l) {
      case WeightsLayout::dlstm_dir_io:
          return true;
      default:
          return false;
      }
  }
  
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  // Tensor Explanation
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  //
  // resource     - 80x80
  //      totalSize   - 6400
  //      x pitch     - 1
  //      y pitch     - 80
  //
  // view         - 60x60
  // viewOffset   - (20,20) => 20*80+20 = 1620
  //
  // padding (contains "paddedVal"):
  //      before  - x=20, y=20
  //      after   - x=20, y=20.
  //
  // logical data - 40x40 (contains the actual data).
  //
  // firstElementOffset:
  //      (viewOffset_x + padBefore_x) + (viewOffset_y + padBefore_y)*y_pitch =
  //      viewOffset + padBefore_x + padBefore_y*y_pitch =
  //      1620 + 20 + 20*80 = 3240
  //
  //
  //                                      whole resource (80x80)
  // +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
  // +                                                                                               +
  // +                                                                                               +
  // +                                view inside resource (60x60)                                   +
  // +       +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++       +
  // +       + start of padded part(20,20) = viewOffset                                      +       +
  // +       +                                                                               +       +
  // +       +                             logical data (40x40)                              +       +
  // +       +       +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++       +       +
  // +       +       + first element (40,40)                                         +       +       +
  // +       +       +                                                               +       +       +
  // +       +       +                                                               +       +       +
  // +       +       +                                                               +       +       +
  // +       +       +                                                               +       +       +
  // +       +       +                                                               +       +       +
  // +       +       +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++       +       +
  // +       +                                                                               +       +
  // +       +                                                                               +       +
  // +       +                                                                               +       +
  // +       +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++       +
  // +                                                                                               +
  // +                                                                                               +
  // +                                                                                               +
  // +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
  //
  //
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  // TensorBase
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  struct TensorBase {
  protected:
      NDims dims;
      size_t viewOffset = 0;  // in elements
      size_t firstElementOffset = 0;
      size_t totalSize = 0;  // in elements
      float paddedVal = 0.f;
  
  public:
      TensorBase() = default;
      TensorBase(const TensorBase&) = default;
      TensorBase& operator=(const TensorBase&) = default;
  
      TensorBase(const NDims& nd, size_t viewOf, size_t sz, float pv)
          : dims(nd),
            viewOffset(viewOf),
            firstElementOffset(std::accumulate(nd.cbegin(),
                                               nd.cend(),
                                               viewOf,
                                               [](size_t val, const Dim& d) { return val + d.pitch * d.pad.before; })),
            totalSize(sz),
            paddedVal(pv) {
          if (totalSize == 0) {
              for (const auto& d : dims) {
                  totalSize = std::max(totalSize, d.pitch * (d.LogicalDimPadded()));
              }
  
              totalSize += viewOffset;
          }
  
          size_t minimalPitch = 1;
  
          for (const auto& d : dims) {
              if (d.pitch < minimalPitch) {
                  throw std::runtime_error("Tensor pitches didn't set correctly");
              }
  
              minimalPitch *= d.LogicalDimPadded();
          }
  
          if (totalSize < (minimalPitch + viewOffset)) {
              throw std::runtime_error("Tensor total Size didn't set correctly");
          }
      }
  
      float GetPaddedVal() const { return paddedVal; }
      size_t GetFirstElementOffset() const { return firstElementOffset; }
      size_t GetViewOffset() const { return viewOffset; }
      const NDims& GetDims() const { return dims; }
  
      virtual uint32_t ElementSize() const = 0;
  
      // Size of the actual data (without padded part)
      size_t LogicalSize() const {
          return std::accumulate(dims.cbegin(), dims.cend(), (size_t)1, [](size_t val, const Dim& d) {
              return val * d.v;
          });
      }
  
      // Dimensions of the actual data (without padded part)
      std::vector<size_t> LogicalDims() const {
          std::vector<size_t> res(dims.size());
          std::transform(dims.begin(), dims.end(), res.begin(), [](const Dim& d) { return d.v; });
          return res;
      }
  
      // Whole buffer size (in elements)
      size_t PhysicalSize() const { return totalSize; }
  
      // Whole buffer size (in bytes)
      size_t PhysicalSizeInBytes() const { return totalSize * ElementSize(); }
  
      // if padded/view exists between logical dimensions.
      // in other words, if we can consider the data as a 1Dim resource.
      bool PitchesDifferFromLogicalDims() const {
          bool differ = false;
  
          size_t calc_pitch = 1;
          for (const auto& d : dims) {
              differ |= (d.pitch != calc_pitch);
              calc_pitch *= d.v;
          }
  
          return differ;
      }
  
      virtual ~TensorBase() = default;
  };
  
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  // TensorBaseT
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  template <typename DType, typename Layout>
  struct TensorBaseT : public TensorBase {
  protected:
      DType dtype;
      Layout layout;
  
      template <typename ArrayT, typename ChannelName>
      static inline int ChannelIndex(const ArrayT& channelArr, Layout l, ChannelName channelName) {
          size_t channel = static_cast<size_t>(channelName);
  
          for (auto& entry : channelArr) {
              if (entry.first == l)
                  return entry.second[channel];
          }
  
          return -1;
      }
  
      template <typename ArrayT, typename ChannelName>
      static inline Dim Extract(const ArrayT& channelArr, Layout l, ChannelName channelName, const NDims& dims) {
          const int i = ChannelIndex(channelArr, l, channelName);
          return ((i < 0) || (i >= static_cast<int>(dims.size()))) ? Dim{1, 1, {0, 0}} : dims[i];
      }
  
      template <typename ArrayT>
      static inline uint32_t ChannelsCount(const ArrayT& channelArr, Layout l) {
          const auto& entry =
              std::find_if(std::begin(channelArr),
                           std::end(channelArr),
                           [&](typename std::tuple_element<0, ArrayT>::type entry) { return entry.first == l; });
  
          if (entry == channelArr.end())
              throw std::invalid_argument("Failed to get channels count for layout " +
                                          std::to_string(static_cast<uint32_t>(l)));
  
          return std::accumulate(entry->second.begin(), entry->second.end(), 0U, [](uint32_t count, int v) {
              return count + ((v != -1) ? 1 : 0);
          });
      }
  
  public:
      TensorBaseT() = default;
      TensorBaseT(const TensorBaseT&) = default;
      TensorBaseT& operator=(const TensorBaseT&) = default;
  
      TensorBaseT(const NDims& nd, DType dt, Layout l, size_t of = 0, size_t sz = 0, float pv = 0.f)
          : TensorBase(nd, of, sz, pv), dtype(dt), layout(l) {}
  
      DType GetDType() const { return dtype; }
      Layout GetLayout() const { return layout; }
      uint32_t ElementSize() const override { return BytesPerElement(dtype); }
      size_t Dimentions() const { return dims.size(); }
      bool SimpleLayout() const { return Tensor::SimpleLayout(layout); }
      bool GroupedLayout() const { return Tensor::GroupedLayout(layout); }
  
      bool operator==(const TensorBaseT& t) const {
          bool same = dtype == t.dtype && layout == t.layout && paddedVal == t.paddedVal && viewOffset == t.viewOffset &&
                      dims.size() == t.dims.size();
          if (same) {
              for (size_t i = 0; i < dims.size(); i++) {
                  same &= dims[i].v == t.dims[i].v && dims[i].pad.before == t.dims[i].pad.before &&
                          dims[i].pad.after == t.dims[i].pad.after && dims[i].pitch == t.dims[i].pitch;
              }
          }
  
          return same;
      }
  
      bool SameDims(const TensorBaseT& t) const {
          bool same = dtype == t.dtype && layout == t.layout && dims.size() == t.dims.size();
          if (same) {
              for (size_t i = 0; i < dims.size(); i++) {
                  same &= dims[i].v == t.dims[i].v;
              }
          }
  
          return same;
      }
  
      bool SameDimsSizes(const TensorBaseT& t) const {
          bool same = dims.size() == t.dims.size();
          if (same) {
              for (size_t i = 0; i < dims.size(); i++) {
                  same &= dims[i].v == t.dims[i].v;
              }
          }
  
          return same;
      }
  };
  
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  // DataTensor
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  struct DataTensor : public TensorBaseT<Datatype, DataLayout> {
      DataTensor() = default;
      DataTensor(const DataTensor&) = default;
      DataTensor& operator=(const DataTensor&) = default;
  
      DataTensor(const NDims& nd, Datatype dt, DataLayout l, size_t of = 0, size_t sz = 0, float pv = 0.f)
          : TensorBaseT(nd, dt, l, of, sz, pv) {}
  
      DataTensor(const std::vector<size_t>& d, Datatype dt, DataLayout l)
          : TensorBaseT<Datatype, DataLayout>(GetSimpleDims(d, l), dt, l) {}
  
      Dim X() const { return Extract(layout, DataChannelName::X, dims); }
      Dim Y() const { return Extract(layout, DataChannelName::Y, dims); }
      Dim Z() const { return Extract(layout, DataChannelName::Z, dims); }
      Dim W() const { return Extract(layout, DataChannelName::W, dims); }
      Dim Feature() const { return Extract(layout, DataChannelName::FEATURE, dims); }
      Dim Batch() const { return Extract(layout, DataChannelName::BATCH, dims); }
  
      DataTensor TransformIgnorePadding(DataLayout l) const;
      DataTensor FlattenFeatureAndSpatials() const;
      DataTensor FlattenEverything() const;
  
      static inline Dim Extract(DataLayout l, DataChannelName channel, const NDims& d) {
          return TensorBaseT::Extract(dataChannelArray, l, channel, d);
      }
  
      static inline int Channelndex(DataLayout l, DataChannelName channel) {
          return TensorBaseT::ChannelIndex(dataChannelArray, l, channel);
      }
  
      static inline uint32_t ChannelsCount(DataLayout l) { return TensorBaseT::ChannelsCount(dataChannelArray, l); }
  
  private:
      using DataChannelDesc = std::pair<DataLayout, std::array<int, static_cast<size_t>(DataChannelName::COUNT)>>;
      using DataChannelArray = std::array<DataChannelDesc, DataLayout::DataLayoutCount>;
      static DataChannelArray dataChannelArray;
      static NDims GetSimpleDims(const std::vector<size_t>& d, DataLayout l);
  };
  
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  // WeightsTensor
  ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
  struct WeightsTensor : TensorBaseT<WeightsType, WeightsLayout> {
      WeightsTensor() = default;
      WeightsTensor(const WeightsTensor&) = default;
      WeightsTensor& operator=(const WeightsTensor&) = default;
  
      WeightsTensor(const NDims& nd, WeightsType dt, WeightsLayout l, size_t of = 0, size_t sz = 0, float pv = 0.f)
          : TensorBaseT(nd, dt, l, of, sz, pv) {}
  
      WeightsTensor(const std::vector<size_t>& d, WeightsType dt, WeightsLayout l)
          : TensorBaseT<WeightsType, WeightsLayout>(GetSimpleDims(d, l), dt, l) {}
  
      WeightsTensor TransformIgnorePadding(WeightsLayout l) const { return TransformIgnorePadding(l, dtype); }
      WeightsTensor TransformIgnorePadding(WeightsLayout l, WeightsType t, size_t g = 1, bool should_split = true) const;
  
      Dim X() const { return Extract(layout, WeightsChannelName::X, dims); }
      Dim Y() const { return Extract(layout, WeightsChannelName::Y, dims); }
      Dim Z() const { return Extract(layout, WeightsChannelName::Z, dims); }
      Dim IFM() const { return Extract(layout, WeightsChannelName::IFM, dims); }
      Dim OFM() const { return Extract(layout, WeightsChannelName::OFM, dims); }
      Dim G() const { return Extract(layout, WeightsChannelName::G, dims); }
  
      static inline Dim Extract(WeightsLayout l, WeightsChannelName channel, const NDims& d) {
          return TensorBaseT::Extract(weightsChannelArray, l, channel, d);
      }
  
      static inline int Channelndex(WeightsLayout l, WeightsChannelName channel) {
          return TensorBaseT::ChannelIndex(weightsChannelArray, l, channel);
      }
  
      static inline bool DoesGroupDimExist(WeightsLayout l) {
          return TensorBaseT::ChannelIndex(weightsChannelArray, l, WeightsChannelName::G) != -1;
      }
  
      static inline uint32_t ChannelsCount(WeightsLayout l) { return TensorBaseT::ChannelsCount(weightsChannelArray, l); }
  
  private:
      using WeightsChannelDesc =
          std::pair<WeightsLayout, std::array<int, static_cast<size_t>(WeightsChannelName::COUNT)>>;
      using WeightsChannelArray = std::array<WeightsChannelDesc, WeightsLayout::WeightsLayoutCount>;
      static WeightsChannelArray weightsChannelArray;
      static NDims GetSimpleDims(const std::vector<size_t>& d, WeightsLayout l);
  };
  
  inline bool GroupedLayout(WeightsLayout l) {
      return WeightsTensor::DoesGroupDimExist(l);
  }
  
  }  // namespace Tensor
  }  // namespace kernel_selector
  