Method ComputeMklSlice

tensorflow/core/kernels/mkl_slice_op.cc:357–482

Slice op implemented using MKL-DNN APIs.

 private:
  // Slice op implemented using MKL-DNN APIs.
  void ComputeMklSlice(OpKernelContext* context,
                       const gtl::InlinedVector<int64, 4>& begin,
                       const gtl::InlinedVector<int64, 4>& size) {
    try {
      // OneDNN API usage below is guided by description at:
      //  https://github.com/01org/mkl-dnn/issues/69
      //
      // Relevant part of the description is copied below:
      //
      // Let's say you want to copy a part of memory into another buffer (and
      // probably change the format). Then your steps are:
      //
      // 1. create memory primitive descriptor in_mem_pd and memory primitive
      //    in_mem_p for the entire source data. create view primitive
      //    descriptor in_submem_pd based on in_mem_pd, initial offsets,
      //    and sub-sizes
      // 2. create memory primitive descriptor out_mem_pd and memory primitive
      //    out_mem_p for the output (the logical sizes should match sub-sizes
      //    used in step 1, but the format might be arbitrary)
      // 3. create reorder primitive descriptor reorder_pd based on in_submem_pd
      //    and out_mem_pd. create reorder primitive itself based on reorder_pd,
      //    in_mem_p, and out_mem_p.
      //
      // Please notice that there is no view primitive. There is only view
      // primitive descriptor. And the reorder uses source memory as input but
      // traverses it according to a view in_submem_pd.

      auto cpu_engine = engine(ENGINE_CPU, 0);
      MklDnnData<T> src(&cpu_engine);
      MklDnnData<T> output(&cpu_engine);

      // Populate offsets and sizes in memory::dims format based on vector.
      memory::dims begin_dims = {};
      begin_dims.resize(begin.size());
      for (size_t i = 0; i < begin.size(); ++i) begin_dims[i] = begin[i];
      memory::dims size_dims = {};
      bool empty = false;
      size_dims.resize(size.size());
      for (size_t i = 0; i < size.size(); ++i) {
        size_dims[i] = size[i];
        if (size_dims[i] == 0) empty = true;
      }

      Tensor* output_tensor = nullptr;
      MklDnnShape output_mkl_shape;

      // If no dimension is selected in slice, the result should be empty.
      // Just return an empty output tensor, and a dummy OneDNN-shape tensor.
      if (empty) {  // for empty dims
        auto shape_to = MklDnnDimsToTFShape(size_dims);
        AllocateOutputSetMklShape(context, 0, &output_tensor, shape_to,
                                  output_mkl_shape);
        return;
      }

      // Step 1 (as per above description) - Create memory for user data.
      // We use blocked format here to describe input tensor.
      const Tensor& input_tensor = MklGetInput(context, 0);