text: /*******************************************************************************
  * Copyright 2021-2022 Intel Corporation
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  *
  *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  *******************************************************************************/
  
  #include <cctype>
  #include <sstream>
  #include <thread>
  
  #include "gpu/jit/conv/tensor.hpp"
  
  namespace dnnl {
  namespace impl {
  namespace gpu {
  namespace jit {
  
  layout_t::layout_t(const type_t &type, const expr_t &offset,
          const std::vector<std::pair<int, dim_t>> &parts,
          const std::vector<dim_t> &dims, bool do_normalize)
      : type_(type), offset_(offset) {
      ndims_ = 0;
      for (auto &p : parts) {
          int dim_idx = p.first;
          dim_t block = p.second;
          ndims_ = std::max(ndims_, dim_idx + 1);
          if (block == 0 && dims.empty())
              ir_error_not_expected()
                      << "Dimensions are missing. Can't deduce them from "
                         "the format.";
      }
      if (!dims.empty() && ndims_ != int(dims.size())) {
          ir_error_not_expected() << "Format and dimensions do not match.";
      }
  
      dim_t stride = 1;
      // Iterate from right to left (innermost to outermost).
      for (auto it = parts.rbegin(); it != parts.rend(); ++it) {
          int dim_idx = it->first;
          dim_t block = it->second;
          if (block == 0) {
              dim_t full_block = 1;
              for (auto &b : blocks_)
                  if (b.dim_idx == dim_idx) full_block *= b.block;
  
              block = utils::div_up(dims[dim_idx], full_block);
          }
  
          blocks_.emplace_back(dim_idx, block, stride);
          stride = block * stride;
      }
  
      if (do_normalize) blocks_ = normalize_blocks(ndims_, blocks_);
      sanity_check();
  }
  
  layout_t::layout_t(const memory_desc_wrapper &mdw, bool do_normalize)
      : type_(mdw.data_type()), offset_(mdw.offset0()) {
      ir_assert(mdw.is_blocking_desc()) << "Expected blocking memory descriptor.";
  
      ndims_ = mdw.ndims();
      auto &blocking = mdw.blocking_desc();
      auto *padded_dims = mdw.padded_dims();
  
      dim_t stride = 1;
      std::vector<dim_t> full_blocks(ndims_, 1);
      for (int i = blocking.inner_nblks - 1; i >= 0; i--) {
          int dim_idx = blocking.inner_idxs[i];
          dim_t block = blocking.inner_blks[i];
          blocks_.emplace_back(dim_idx, block, stride);
          stride *= block;
          full_blocks[dim_idx] *= block;
      }
  
      for (int i = 0; i < ndims_; i++) {
          dim_t block = padded_dims[i] / full_blocks[i];
          blocks_.emplace_back(i, block, blocking.strides[i]);
      }
  
      // Sort outer blocks by their stride.
      std::sort(blocks_.begin() + blocking.inner_nblks, blocks_.end(),
              [](const block_t &a, const block_t &b) {
                  if (a.stride == b.stride) return a.dim_idx > b.dim_idx;
                  return a.stride < b.stride;
              });
  
      if (do_normalize) blocks_ = normalize_blocks(ndims_, blocks_);
      sanity_check();
  }
  
  memory_desc_t layout_t::to_dnnl(const dim_t *dims_hint) const {
      memory_desc_t md = {};
      md.ndims = ndims();
      std::copy(dims_hint, dims_hint + ndims(), md.dims);
      md.data_type = jit::to_dnnl(type_);
      md.offset0 = to_cpp<dim_t>(offset_);
      md.format_kind = format_kind::blocked;
  
      auto &blk = md.format_desc.blocking;
      bool seen[DNNL_MAX_NDIMS] = {};
  
      bool in_inner_block = false;
      dim_t prev_stride = 0;
  
      for (auto it = blocks_.rbegin(); it != blocks_.rend(); ++it) {
          auto &b = *it;
          if (!seen[b.dim_idx]) {
              // Outer block.
              ir_assert(!in_inner_block);
              MAYBE_UNUSED(in_inner_block);
              blk.strides[b.dim_idx] = b.stride;
              md.padded_dims[b.dim_idx] = b.block;
          } else {
              // Inner block.
              md.padded_dims[b.dim_idx] *= b.block;
              blk.inner_idxs[blk.inner_nblks] = b.dim_idx;
              blk.inner_blks[blk.inner_nblks] = b.block;
              blk.inner_nblks++;
              if (prev_stride > 0) {
                  // Inner block must be dense.
                  ir_assert(prev_stride == b.block * dim_t(b.stride));
              }
              prev_stride = b.stride;
              in_inner_block = true;
          }
          seen[b.dim_idx] = true;
      }
  
      return md;
  }
  
  layout_t layout_t::map(const tensor_t &tensor) const {
      if (ndims() != tensor.ndims())
          ir_error_not_expected() << "Dimensions do not match.";
  
      std::vector<dim_t> remaining_dims = tensor.dims();
      std::vector<block_t> mapped_blocks;
  
      for (auto &eb : enumerated_blocks()) {
          block_t &b = eb.second;
          bool b_is_outermost = is_outermost(eb);
  
          dim_t block = b.block;
          dim_t &rem_dim = remaining_dims[b.dim_idx];
          if (rem_dim == 1) {
              if (b_is_outermost) {
                  // This is to have similarity between the current and
                  // mapped layouts.
                  mapped_blocks.emplace_back(b.dim_idx, 1, b.stride);
              }
              continue;
          }
          if (b_is_outermost) {
              block = rem_dim;
          } else if (rem_dim % block != 0) {
              // Try to split the current block and start mapping from
              // scratch.
              if (block % rem_dim == 0)
                  return split_block(eb, rem_dim, block / rem_dim).map(tensor);
  
              ir_error_not_expected() << "Can't map tensor layout.";
          }
          rem_dim /= block;
          mapped_blocks.emplace_back(b.dim_idx, block, b.stride);
      }
  
      for (auto &d : remaining_dims) {
          ir_assert(d == 1) << "Can't map tensor layout.";
          MAYBE_UNUSED(d);
      }
  
      return layout_t(type(), ndims(), operator()(tensor.start()), mapped_blocks);
  }
  
  layout_t layout_t::reinterpret(
          const type_t &new_type, bool do_normalize) const {
      int old_size = type().size();
      int new_size = new_type.size();
      if (new_size == old_size) return *this;
  
      expr_t new_offset = 0;
      if (!has_zero_offset()) {
          ir_assert(is_const(offset_)) << "Expected constant offset.";
          int64_t off = to_cpp<int64_t>(offset_) * old_size;
          ir_assert(off % new_size == 0);
          new_offset = off / new_size;
      }
  
      if (old_size % new_size != 0 && new_size % old_size != 0) {
          ir_error_not_expected();
          return layout_t();
      }
  
      auto new_blocks = blocks_;
      if (new_blocks.empty()) {
          ir_error_not_expected() << "Can't reinterpret.";
          return layout_t();
      }
  
      auto &b0 = new_blocks.front();
      if (dim_t(b0.stride) != 1) {
          ir_error_not_expected();
          return layout_t();
      }
  
      if (new_size < old_size) {
          int factor = (old_size / new_size);
          b0.block *= factor;
          // Recompute strides.
          for (auto &b : new_blocks) {
              if (&b == &b0) continue;
              b.stride *= factor;
          }
      } else {
          int factor = (new_size / old_size);
          if (b0.block % factor != 0) {
              ir_error_not_expected();
              return layout_t();
          }
          b0.block /= factor;
          // Recompute strides.
          for (auto &b : new_blocks) {
              if (&b == &b0) continue;
              if (b.stride % factor != 0) {
                  ir_error_not_expected();
                  return layout_t();
              }
              b.stride /= factor;
          }
      }
  
      return layout_t(new_type, ndims(), new_offset, new_blocks, do_normalize);
  }
  
  layout_t layout_t::split_block(
          const std::pair<int, block_t> &eb, dim_t block0, dim_t block1) const {
      int block_idx = eb.first;
      auto &b = eb.second;
      ir_assert(b.block == block0 * block1) << "Incompatible block sizes.";
      MAYBE_UNUSED(b);
  
      auto new_blocks = blocks_;
  
      block_t &b0 = new_blocks[block_idx];
      block_t b1 = b0;
  
      b0.block = block0;
      b1.block = block1;
      b1.stride = b0.stride * block0;
  
      new_blocks.insert(new_blocks.begin() + block_idx + 1, b1);
  
      return layout_t(
              type(), ndims(), offset(), new_blocks, /*do_normalize=*/false);
  }
  
  layout_t layout_t::split_into_multi_blocks(
          const std::vector<dim_t> &multi_blocks) const {
      if (is_empty()) return *this;
  
      layout_t tmp(*this);
      std::vector<dim_t> rem_elems = multi_blocks;
      std::vector<dim_t> cur_elems(rem_elems.size(), 1);
      for (auto &eb : tmp.enumerated_blocks()) {
          auto &b = eb.second;
          for (int i = 0; i < int(rem_elems.size()); i++) {
              auto &e = rem_elems[i];
              if (e == 1) continue;
              if (b.block > e) {
                  // Try to split this block.
                  int next_block = utils::max_div(b.block, e);
                  if (next_block == 1) return layout_t();
                  return tmp.split_block(eb, next_block, b.block / next_block)
                          .split_into_multi_blocks(multi_blocks);
              }
              if (e % b.block != 0) return layout_t();
              e /= b.block;
              cur_elems[i] *= b.block;
              break;
          }
      }
      for (int i = 0; i < int(cur_elems.size()); i++) {
          if (cur_elems[i] != multi_blocks[i]) { return layout_t(); }
      }
      return tmp;
  }
  
  tensor_t layout_t::split_into_max_tile(
          dim_t max_tile_elems, bool is_dense_tile) const {
      stride_t dense_stride = 1;
      std::vector<dim_t> tile_dims(ndims(), 1);
      dim_t cur_elems = 1;
      for (auto &eb : enumerated_blocks()) {
          auto &b = eb.second;
          if (b.block == 1) continue;
          if (b.block * cur_elems <= max_tile_elems) {
              if (is_dense_tile) {
                  if (b.stride.is_unknown()) break;
                  if (dense_stride != b.stride) break;
                  dense_stride = b.block * b.stride;
              }
              cur_elems *= b.block;
              tile_dims[b.dim_idx] *= b.block;
              continue;
          }
          dim_t max_block = utils::max_div(b.block, max_tile_elems / cur_elems);
          if (max_block == 1) break;
          auto tmp_layout = split_block(eb, max_block, b.block / max_block);
          return tmp_layout.split_into_max_tile(max_tile_elems, is_dense_tile);
      }
      return tensor_t(tile_dims);
  }
  
  void layout_t::align_layouts(layout_t &a, layout_t &b) {
      for (int i = 0; i < a.ndims(); i++) {
          auto a_blocks = a.blocks();
          auto b_blocks = b.blocks();
  
          int a_max = int(a_blocks.size());
          int b_max = int(b_blocks.size());
          int a_idx = 0;
          int b_idx = 0;
  
          for (;;) {
              while (a_idx < a_max && a_blocks[a_idx].dim_idx != i)
                  a_idx++;
              while (b_idx < b_max && b_blocks[b_idx].dim_idx != i)
                  b_idx++;
  
              if (a_idx >= a_max || b_idx >= b_max) break;
  
              auto &ab = a_blocks[a_idx];
              auto &bb = b_blocks[b_idx];
              dim_t common_block = math::gcd(ab.block, bb.block);
              if (ab.block == common_block && bb.block == common_block) {
                  a_idx++;
                  b_idx++;
                  continue;
              }
  
              if (ab.block != common_block) {
                  a = a.split_block(
                          {a_idx, ab}, common_block, ab.block / common_block);
              }
              if (bb.block != common_block) {
                  b = b.split_block(
                          {b_idx, bb}, common_block, bb.block / common_block);
              }
              break;
          }
      }
  }
  
  std::vector<std::pair<char, dim_t>> layout_t::parse_letter_blocks(
          const std::string &format) {
      std::vector<std::pair<char, dim_t>> ret;
  
      std::stringstream ss(format);
      while (!ss.eof()) {
          int next = ss.peek();
          if (ss.eof()) break;
          dim_t block = 0;
          while (std::isdigit(next)) {
              block = 10 * block + (next - '0');
              ss.ignore(1);
              next = ss.peek();
          }
          char letter = char(ss.peek());
          ir_assert(!ss.eof()) << "EOF is unexpected.";
          ss.ignore(1);
          ret.emplace_back(letter, block);
      }
      return ret;
  }
  
  std::vector<std::pair<int, dim_t>> layout_t::parse_format(
          const std::string &format, int ndims_hint) {
      bool seen_letters[DNNL_MAX_NDIMS] = {};
      int letter_ndims = 0;
      for (char c = 'a'; c < 'a' + DNNL_MAX_NDIMS; c++) {
          if (format.find(c) != std::string::npos) {
              seen_letters[c - 'a'] = true;
              MAYBE_UNUSED(seen_letters);
              letter_ndims++;
          }
      }
  
      for (int i = 0; i < DNNL_MAX_NDIMS; i++) {
          ir_assert(seen_letters[i] == (i < letter_ndims));
      }
  
      auto letter_blocks = parse_letter_blocks(format);
  
      std::vector<std::pair<int, dim_t>> parts;
      for (auto &p : letter_blocks) {
          char letter = p.first;
          dim_t block = p.second;
          if (letter != 'x') {
              int dim_idx = std::tolower(letter) - 'a';
              parts.emplace_back(dim_idx, block);
          } else {
              ir_assert(ndims_hint >= letter_ndims);
              for (int i = letter_ndims; i < ndims_hint; i++) {
                  parts.emplace_back(i, 0);
              }
          }
      }
  
      return parts;
  }
  
  void layout_t::sanity_check() const {
  #ifdef NDEBUG
      return;
  #endif
      if (is_empty()) return;
  
      for (auto &b : blocks_) {
          ir_assert(b.block > 0) << "Incorrect block size.";
          MAYBE_UNUSED(b);
      }
      ir_assert(ndims_ <= max_ndims);
  }
  
  expr_t grid_splitter_t::pop_block(int size) {
      ir_assert(size > 1);
      ir_assert(can_pop_block(size));
  
      int new_stride = cur_stride_ * size;
  
      auto idx_expr = grid_.idx(cur_idx_);
      if (cur_stride_ != 1) idx_expr /= cur_stride_;
      if (new_stride != grid_.dim(cur_idx_)) idx_expr %= size;
  
      cur_stride_ = new_stride;
      if (cur_stride_ == grid_.dim(cur_idx_)) {
          // Move to the next dimension.
          cur_idx_--;
          skip_size_1_dims();
          cur_stride_ = 1;
      }
      return idx_expr;
  }
  
  stride_t tdim_info_t::compute_stride(
          const expr_t &e, int idx, const expr_t &var) {
      // e == var -> fixed stride.
      if (e.is_same(var)) return stride_t(1);
  
      auto vars = find_objects<var_t>(e);
  
      auto e0 = e;
      auto e1 = substitute(e, var, var + 1);
      auto e_stride = simplify(e1 - e0);
  
      if (is_const(e_stride)) return stride_t(to_cpp<dim_t>(e_stride));
  
      // Stride is not a constant.
      return stride_t::unknown();
  }
  
  view_t view_t::create_sub_view(const tensor_t &sub_tensor) const {
      ir_assert(sub_tensor.ndims() == nvdims()) << "Dimensions don't match.";
  
      auto ret = *this;
      ret.vdims_ = sub_tensor.dims();
      for (int i = 0; i < nvdims(); i++) {
          auto &i_start = sub_tensor.start()[i];
          if (is_zero(i_start)) continue;
          auto &s = ret.vstart_[i];
          s += i_start;
          s = simplify(s);
      }
      return ret;
  }
  
  view_t view_t::substitute(const expr_t &from, const expr_t &to) const {
      view_t ret = *this;
      for (int i = 0; i < nvdims(); i++) {
          ret.vstart_[i] = jit::substitute(ret.vstart_[i], from, to);
          ret.vstart_[i] = simplify(ret.vstart_[i]);
      }
      return ret;
  }
  
  std::vector<expr_t> view_t::create_vvars(int nvdims) {
      static const int max_nvdims = 128;
      static thread_local std::vector<expr_t> _vvars([] {
          std::vector<expr_t> ret;
          ret.reserve(max_nvdims);
          for (int i = 0; i < max_nvdims; i++)
              ret.push_back(var_t::make(type_t::s32(), "_" + std::to_string(i)));
          return ret;
      }());
  
      ir_assert(nvdims <= max_nvdims) << "Too many dimensions: " << nvdims;
      return std::vector<expr_t>(_vvars.begin(), _vvars.begin() + nvdims);
  }
  
  layout_t view_t::create_pseudo_vlayout(const layout_t &tlayout) const {
      ir_assert(!tlayout.is_empty());
  
      std::vector<dim_t> rem_vdims = vdims_;
      std::vector<block_t> blocks;
  
      for (auto &teb : tlayout.enumerated_blocks()) {
          block_t &tb = teb.second;
          bool tb_is_outermost = tlayout.is_outermost(teb);
          dim_t tblock = tb.block;
  
          auto &tinfo = tdims_[tb.dim_idx];
          if (tb_is_outermost) {
              bool is_first = true;
              for (int i = tinfo.nvargs() - 1; i >= 0; i--) {
                  int vidx = tinfo.vidx(i);
                  if (rem_vdims[vidx] == 1) continue;
  
                  // When expression contains 2+ variables, use unknown
                  // stride unless the view variable is the innermost.
                  stride_t stride
                          = (is_first ? tinfo.vstride(i) : stride_t::unknown());
                  blocks.emplace_back(
                          vidx, rem_vdims[vidx], stride * stride_t(tb.stride));
                  rem_vdims[vidx] = 1;
                  is_first = false;
              }
              continue;
          }
  
          ir_assert(tinfo.is_identity()) << "Can't create pseudo-layout.";
  
          int vidx = tinfo.vidx(0);
          dim_t &rem_vdim = rem_vdims[vidx];
          if (rem_vdim == 1) continue;
  
          if (tb_is_outermost) {
              tblock = rem_vdim;
              rem_vdim = 1;
          } else if (rem_vdim % tblock == 0) {
              rem_vdim /= tblock;
          } else if (rem_vdim % tblock != 0) {
              // Try to split the current block and start from scratch.
              if (tblock % rem_vdim == 0) {
                  auto tmp_layout
                          = tlayout.split_block(teb, rem_vdim, tblock / rem_vdim);
                  return create_pseudo_vlayout(tmp_layout);
              }
  
              ir_error_not_expected() << "Can't create pseudo-layout.";
          }
          blocks.emplace_back(tb.dim_idx, tblock, tb.stride);
      }
  
      for (auto &d : rem_vdims) {
          ir_assert(d == 1) << "Can't create pseudo-layout.";
          MAYBE_UNUSED(d);
      }
  
      return layout_t(tlayout.type(), nvdims(), 0, blocks);
  }
  
  layout_t dim_assignment_t::map(const layout_t &layout) const {
      std::vector<block_t> new_blocks;
      for (auto &b : layout.blocks()) {
          int new_idx = assignments_[b.dim_idx];
          if (new_idx == -1) continue; // Drop this block.
          auto new_b = b;
          new_b.dim_idx = new_idx;
          new_blocks.push_back(new_b);
      }
      new_blocks = layout_t::normalize_blocks(new_ndims(), new_blocks,
              /*remove_size_1_blocks=*/false);
      auto ret = layout_t(layout.type(), new_ndims(), layout.offset(), new_blocks,
              /*do_normalize=*/false);
      ir_assert(layout.elems() == ret.elems())
              << "Assignment doesn't preserve number of elements.";
      return ret;
  }
  
  // Adds size one spatial dimensions according to input parameters. Spatial
  // dimensions are assumed to be the last dimensions.
  layout_t normalize_conv_spatial(const layout_t &layout, int old_sp_ndims,
          int reduced_dim, bool fuse_spatial) {
      int old_ndims = layout.ndims();
      int new_ndims = old_ndims - old_sp_ndims + (fuse_spatial ? 1 : 3);
  
      dim_assignment_t to_3d(old_ndims, new_ndims);
      for (int i = 0; i < old_ndims; i++) {
          if (i < old_ndims - old_sp_ndims) {
              // Non-spatial dimensions.
              to_3d.assign(i, i);
          } else if (fuse_spatial) {
              // Spatial dimensions, fuse to one.
              to_3d.assign(i, new_ndims - 1);
          } else {
              // Spatial dimensions.
              int sp_idx = 3 - (old_ndims - i);
              if (reduced_dim == 3) {
                  sp_idx = 2;
              } else if (sp_idx < reduced_dim) {
                  sp_idx += 1;
              }
              to_3d.assign(i, new_ndims - (3 - sp_idx));
          }
      }
      return to_3d.map(layout);
  }
  
  layout_t insert_dimension(const layout_t &layout, int dim_idx) {
      auto new_blocks = layout.blocks();
      for (auto &b : new_blocks) {
          if (b.dim_idx >= dim_idx) b.dim_idx++;
      }
      return layout_t(layout.type(), layout.ndims() + 1, layout.offset(),
              new_blocks,
              /*do_normalize=*/false);
  }
  
  layout_t remove_size_1_dimension(const layout_t &layout, int dim_idx) {
      ir_assert(0 <= dim_idx && dim_idx < layout.ndims());
      ir_assert(layout.dim(dim_idx) == 1);
      dim_assignment_t a(layout.ndims(), layout.ndims() - 1);
      for (int i = 0; i < layout.ndims(); i++) {
          if (i == dim_idx) continue;
          a.assign(i, i < dim_idx ? i : i - 1);
      }
      return a.map(layout);
  }
  
  layout_t split_dimension(
          const layout_t &_layout, int dim_idx, int outer_block) {
      int rem_inner_block
              = ir_utils::safe_divide(_layout.dim(dim_idx), outer_block);
      auto layout = insert_dimension(_layout, dim_idx);
      std::vector<block_t> new_blocks;
      for (auto &eb : layout.enumerated_blocks()) {
          auto &b = eb.second;
          if (b.dim_idx != dim_idx + 1) {
              new_blocks.push_back(b);
              continue;
          }
          if (b.block % rem_inner_block == 0) {
              new_blocks.emplace_back(dim_idx + 1, rem_inner_block, b.stride);
              new_blocks.emplace_back(dim_idx, b.block / rem_inner_block,
                      dim_t(b.stride) * rem_inner_block);
              rem_inner_block = 1;
          } else {
              new_blocks.push_back(b);
              rem_inner_block = ir_utils::safe_divide(rem_inner_block, b.block);
          }
      }
  
      // Remove inner blocks with size one.
      std::vector<block_t> _new_blocks;
      std::vector<bool> seen(layout.ndims());
      for (auto it = new_blocks.rbegin(); it != new_blocks.rend(); ++it) {
          if (it->block == 1 && seen[it->dim_idx]) continue;
          _new_blocks.push_back(*it);
          seen[it->dim_idx] = true;
      }
      std::reverse(_new_blocks.begin(), _new_blocks.end());
      return layout_t(layout.type(), layout.ndims(), layout.offset(), _new_blocks,
              /*do_normalize=*/false);
  }
  
  layout_t normalize_conv_groups(const layout_t &layout, bool with_groups,
          int groups, bool is_dw, bool add_groups, bool is_wei) {
      if (with_groups == add_groups) return layout;
      if (is_wei) {
          ir_assert(groups == 1)
                  << "Adding/removing groups can be done only for single group.";
          if (add_groups) return insert_dimension(layout, 0);
          return remove_size_1_dimension(layout, 0);
      }
  
      ir_assert(!with_groups) << "Unexpected groups in source/destination.";
      if (is_dw) groups = layout.dim(1);
      if (layout.dim(1) == 1) groups = 1;
      return split_dimension(layout, /*dim_idx=*/1, groups);
  }
  
  layout_t normalize_conv_layout(const layout_t &_layout, bool with_groups,
          int groups, bool is_dw, int reduced_dim, bool fuse_spatial,
          bool add_groups, bool is_wei) {
      int old_sp_ndims = _layout.ndims() - (with_groups ? 3 : 2);
  
      layout_t layout = _layout;
      layout = normalize_conv_spatial(
              layout, old_sp_ndims, reduced_dim, fuse_spatial);
      layout = normalize_conv_groups(
              layout, with_groups, groups, is_dw, add_groups, is_wei);
  
      return layout;
  }
  
  std::vector<dim_t> normalize_conv_dims(std::vector<dim_t> &dims,
          bool with_groups, int groups, bool is_dw, int reduced_dim,
          bool fuse_spatial, bool add_groups, bool is_wei) {
      layout_t dummy_layout(type_t::u8(), 0, dims);
      return normalize_conv_layout(dummy_layout, with_groups, groups, is_dw,
              reduced_dim, fuse_spatial, add_groups, is_wei)
              .dims();
  }
  
  void normalize_conv_layouts(layout_t &src_layout, layout_t &wei_layout,
          layout_t &dst_layout, layout_t &bia_layout, bool with_groups,
          int groups, bool is_dw, int reduced_dim, bool fuse_spatial,
          bool add_groups) {
      src_layout = normalize_conv_layout(src_layout, /*with_groups=*/false,
              groups, is_dw, reduced_dim, fuse_spatial, add_groups,
              /*is_wei=*/false);
      wei_layout = normalize_conv_layout(wei_layout, with_groups, groups, is_dw,
              reduced_dim, /*fuse_spatial=*/false, add_groups, /*is_wei=*/true);
      dst_layout = normalize_conv_layout(dst_layout, /*with_groups=*/false,
              groups, is_dw, reduced_dim, fuse_spatial, add_groups,
              /*is_wei=*/false);
      if (add_groups && !bia_layout.is_empty()) {
          ir_assert(bia_layout.ndims() == 1) << bia_layout;
          bia_layout = split_dimension(bia_layout, 0, groups);
      }
  }
  
  } // namespace jit
  } // namespace gpu
  } // namespace impl
  } // namespace dnnl
  