<commit_before>/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

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.
==============================================================================*/

#ifdef INTEL_MKL

#include <limits>
#include <vector>

#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/framework/bounds_check.h"
#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/kernels/mkl_pooling_ops_common.h"

namespace tensorflow {
using mkldnn::prop_kind;

template <typename T>
void MklPoolingFwdPrimitive<T>::Setup(const MklPoolingParams& fwdParams) {
  DCHECK(fwdParams.alg_kind == ALGORITHM::pooling_max ||
         fwdParams.alg_kind == ALGORITHM::pooling_avg ||
         fwdParams.alg_kind == ALGORITHM::pooling_avg_include_padding ||
         fwdParams.alg_kind == ALGORITHM::pooling_avg_exclude_padding)
      << "Pooling algorithm kind is not supported";

  context_.alg_kind = fwdParams.alg_kind;
  context_.prop_kind = fwdParams.prop_kind;

  // Create memory descriptor
  // FIXME: Pooling doesn't expose to get the src_primitive_desc,
  //        so src format is currently hard-coded.
  //        A utility function is used to do this,
  //        which may be broken with future CPU architectures
  bool is_2d = (fwdParams.src_dims.size() == 4);
  if (std::is_same<T, qint8>::value || std::is_same<T, quint8>::value)
    context_.src_fmt = is_2d ? MEMORY_FORMAT::nhwc : MEMORY_FORMAT::ndhwc;
  else
    context_.src_fmt = fwdParams.src_format;

  context_.src_md.reset(new memory::desc({fwdParams.src_dims}, MklDnnType<T>(),
                                         context_.src_fmt));
  context_.dst_md.reset(new memory::desc({fwdParams.dst_dims}, MklDnnType<T>(),
                                         MEMORY_FORMAT::any));

  // Create a pooling descriptor.
  context_.fwd_desc.reset(new pooling_forward::desc(
      fwdParams.prop_kind, fwdParams.alg_kind, *context_.src_md,
      *context_.dst_md, fwdParams.strides, fwdParams.filter_dims,
      fwdParams.padding_left, fwdParams.padding_right, padding_kind::zero));
  context_.fwd_pd.reset(
      new pooling_forward::primitive_desc(*context_.fwd_desc, cpu_engine_));

#ifndef ENABLE_MKLDNN_V1
  context_.dst_fmt = static_cast<MEMORY_FORMAT>(
      context_.fwd_pd.get()->PRIMITIVE_DESC_DST.desc().data.format);
#endif  // ENABLE_MKLDNN_V1

  // Create MKL-DNN internal memory object with dummy data.
  context_.src_mem.reset(new MEMORY_CONSTRUCTOR(
      context_.fwd_pd.get()->PRIMITIVE_DESC_SRC, cpu_engine_, DummyData));
  context_.dst_mem.reset(new MEMORY_CONSTRUCTOR(
      context_.fwd_pd.get()->PRIMITIVE_DESC_DST, cpu_engine_, DummyData));

  // For max pooling, need to return workspace (ws) for backward computing.
  if (fwdParams.alg_kind == ALGORITHM::pooling_max &&
      fwdParams.prop_kind == prop_kind::forward_training) {
#ifdef ENABLE_MKLDNN_V1
    context_.ws_mem.reset(
        new MEMORY_CONSTRUCTOR(context_.fwd_pd.get()->PRIMITIVE_DESC_WORKSPACE,
                               cpu_engine_, DummyData));
    context_.net_args.push_back({{MKLDNN_ARG_SRC, *context_.src_mem},
                                 {MKLDNN_ARG_DST, *context_.dst_mem},
                                 {MKLDNN_ARG_WORKSPACE, *context_.ws_mem}});
    context_.fwd.reset(new pooling_forward(*context_.fwd_pd));
#else
    auto ws_pd = context_.fwd_pd.get()->PRIMITIVE_DESC_WORKSPACE.desc().data;
    // Store workspace's dims and format to create workspace tensor.
    context_.ws_fmt = static_cast<MEMORY_FORMAT>(ws_pd.format);
    context_.ws_dims.assign(ws_pd.dims, ws_pd.dims + ws_pd.ndims);
    context_.ws_dt = static_cast<mkldnn::memory::data_type>(ws_pd.data_type);
    context_.ws_size =
        context_.fwd_pd.get()->PRIMITIVE_DESC_WORKSPACE.get_size();

    context_.ws_mem.reset(
        new memory(context_.fwd_pd.get()->PRIMITIVE_DESC_WORKSPACE, DummyData));
    context_.fwd.reset(new pooling_forward(*context_.fwd_pd, *context_.src_mem,
                                           *context_.dst_mem,
                                           *context_.ws_mem));
#endif  // ENABLE_MKLDNN_V1
  } else {
#ifdef ENABLE_MKLDNN_V1
    context_.net_args.push_back({{MKLDNN_ARG_SRC, *context_.src_mem},
                                 {MKLDNN_ARG_DST, *context_.dst_mem}});
    context_.fwd.reset(new pooling_forward(*context_.fwd_pd));
#else
    context_.fwd.reset(new pooling_forward(*context_.fwd_pd, *context_.src_mem,
                                           *context_.dst_mem));
#endif  // ENABLE_MKLDNN_V1
  }

  context_.fwd_primitives.push_back(*context_.fwd);
}

template <typename T>
void MklPoolingFwdPrimitive<T>::Execute(const T* src_data, T* dst_data,
                                        void* ws_data) {
  context_.src_mem->set_data_handle(
      static_cast<void*>(const_cast<T*>(src_data)));
  context_.dst_mem->set_data_handle(static_cast<void*>(dst_data));
  if (context_.alg_kind == ALGORITHM::pooling_max &&
      context_.prop_kind ==
          prop_kind::forward_training) {  // Max pooling must have workspace.
    DCHECK(ws_data != nullptr);
    context_.ws_mem->set_data_handle(ws_data);
  }

#ifdef ENABLE_MKLDNN_V1
  //ExecuteFwdPrimitives(context_);
  DCHECK_EQ(context_.fwd_primitives.size(), context_.net_args.size());
  for (size_t i = 0; i < context_.fwd_primitives.size(); ++i) {<commit_msg>reviews comments 5<commit_after>/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

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.
==============================================================================*/

#ifdef INTEL_MKL

#include <limits>
#include <vector>

#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/framework/bounds_check.h"
#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/kernels/mkl_pooling_ops_common.h"

namespace tensorflow {
using mkldnn::prop_kind;

template <typename T>
void MklPoolingFwdPrimitive<T>::Setup(const MklPoolingParams& fwdParams) {
  DCHECK(fwdParams.alg_kind == ALGORITHM::pooling_max ||
         fwdParams.alg_kind == ALGORITHM::pooling_avg ||
         fwdParams.alg_kind == ALGORITHM::pooling_avg_include_padding ||
         fwdParams.alg_kind == ALGORITHM::pooling_avg_exclude_padding)
      << "Pooling algorithm kind is not supported";

  context_.alg_kind = fwdParams.alg_kind;
  context_.prop_kind = fwdParams.prop_kind;

  // Create memory descriptor
  // FIXME: Pooling doesn't expose to get the src_primitive_desc,
  //        so src format is currently hard-coded.
  //        A utility function is used to do this,
  //        which may be broken with future CPU architectures
  bool is_2d = (fwdParams.src_dims.size() == 4);
  if (std::is_same<T, qint8>::value || std::is_same<T, quint8>::value)
    context_.src_fmt = is_2d ? MEMORY_FORMAT::nhwc : MEMORY_FORMAT::ndhwc;
  else
    context_.src_fmt = fwdParams.src_format;

  context_.src_md.reset(new memory::desc({fwdParams.src_dims}, MklDnnType<T>(),
                                         context_.src_fmt));
  context_.dst_md.reset(new memory::desc({fwdParams.dst_dims}, MklDnnType<T>(),
                                         MEMORY_FORMAT::any));

  // Create a pooling descriptor.
  context_.fwd_desc.reset(new pooling_forward::desc(
      fwdParams.prop_kind, fwdParams.alg_kind, *context_.src_md,
      *context_.dst_md, fwdParams.strides, fwdParams.filter_dims,
      fwdParams.padding_left, fwdParams.padding_right, padding_kind::zero));
  context_.fwd_pd.reset(
      new pooling_forward::primitive_desc(*context_.fwd_desc, cpu_engine_));

#ifndef ENABLE_MKLDNN_V1
  context_.dst_fmt = static_cast<MEMORY_FORMAT>(
      context_.fwd_pd.get()->PRIMITIVE_DESC_DST.desc().data.format);
#endif  // ENABLE_MKLDNN_V1

  // Create MKL-DNN internal memory object with dummy data.
  context_.src_mem.reset(new MEMORY_CONSTRUCTOR(
      context_.fwd_pd.get()->PRIMITIVE_DESC_SRC, cpu_engine_, DummyData));
  context_.dst_mem.reset(new MEMORY_CONSTRUCTOR(
      context_.fwd_pd.get()->PRIMITIVE_DESC_DST, cpu_engine_, DummyData));

  // For max pooling, need to return workspace (ws) for backward computing.
  if (fwdParams.alg_kind == ALGORITHM::pooling_max &&
      fwdParams.prop_kind == prop_kind::forward_training) {
#ifdef ENABLE_MKLDNN_V1
    context_.ws_mem.reset(
        new MEMORY_CONSTRUCTOR(context_.fwd_pd.get()->PRIMITIVE_DESC_WORKSPACE,
                               cpu_engine_, DummyData));
    context_.net_args.push_back({{MKLDNN_ARG_SRC, *context_.src_mem},
                                 {MKLDNN_ARG_DST, *context_.dst_mem},
                                 {MKLDNN_ARG_WORKSPACE, *context_.ws_mem}});
    context_.fwd.reset(new pooling_forward(*context_.fwd_pd));
#else
    auto ws_pd = context_.fwd_pd.get()->PRIMITIVE_DESC_WORKSPACE.desc().data;
    // Store workspace's dims and format to create workspace tensor.
    context_.ws_fmt = static_cast<MEMORY_FORMAT>(ws_pd.format);
    context_.ws_dims.assign(ws_pd.dims, ws_pd.dims + ws_pd.ndims);
    context_.ws_dt = static_cast<mkldnn::memory::data_type>(ws_pd.data_type);
    context_.ws_size =
        context_.fwd_pd.get()->PRIMITIVE_DESC_WORKSPACE.get_size();

    context_.ws_mem.reset(
        new memory(context_.fwd_pd.get()->PRIMITIVE_DESC_WORKSPACE, DummyData));
    context_.fwd.reset(new pooling_forward(*context_.fwd_pd, *context_.src_mem,
                                           *context_.dst_mem,
                                           *context_.ws_mem));
#endif  // ENABLE_MKLDNN_V1
  } else {
#ifdef ENABLE_MKLDNN_V1
    context_.net_args.push_back({{MKLDNN_ARG_SRC, *context_.src_mem},
                                 {MKLDNN_ARG_DST, *context_.dst_mem}});
    context_.fwd.reset(new pooling_forward(*context_.fwd_pd));
#else
    context_.fwd.reset(new pooling_forward(*context_.fwd_pd, *context_.src_mem,
                                           *context_.dst_mem));
#endif  // ENABLE_MKLDNN_V1
  }

  context_.fwd_primitives.push_back(*context_.fwd);
}

template <typename T>
void MklPoolingFwdPrimitive<T>::Execute(const T* src_data, T* dst_data,
                                        void* ws_data) {
  context_.src_mem->set_data_handle(
      static_cast<void*>(const_cast<T*>(src_data)));
  context_.dst_mem->set_data_handle(static_cast<void*>(dst_data));
  if (context_.alg_kind == ALGORITHM::pooling_max &&
      context_.prop_kind ==
          prop_kind::forward_training) {  // Max pooling must have workspace.
    DCHECK(ws_data != nullptr);
    context_.ws_mem->set_data_handle(ws_data);
  }

#ifdef ENABLE_MKLDNN_V1
  DCHECK_EQ(context_.fwd_primitives.size(), context_.net_args.size());
  for (size_t i = 0; i < context_.fwd_primitives.size(); ++i) {
{'fraction_non_alphanumeric': 0.12285201333675301, 'fraction_numerical': 0.0036761562793878773, 'mean_word_length': 3.534108527131783, 'pattern_counts': {'":': 0, '<': 41, '<?xml version=': 0, '>': 55, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 28, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1685403', 'n_tokens_mistral': 3873, 'n_tokens_neox': 3802, 'n_words': 703}