Function MKLTransposeND

tensorflow/core/kernels/mkl_transpose_op.cc:65–104

// Transpose of N-dimensional tensor using OneDNN
template <typename T>
Status MKLTransposeND(OpKernelContext* context, const Tensor& in_tensor,
                      Tensor* out_tensor, const gtl::ArraySlice<int32>& perm) {
  try {
    engine cpu_engine = engine(ENGINE_CPU, 0);
    MklDnnData<T> in(&cpu_engine);
    MklDnnData<T> out(&cpu_engine);

    memory::dims in_dims = TFShapeToMklDnnDims(in_tensor.shape());
    memory::dims out_dims = TFShapeToMklDnnDims(out_tensor->shape());
    memory::dims in_strides = CalculateTFStrides(in_dims);
    // Reorder output strides based on permutation requested.
    memory::dims out_strides =
        ReorderStrides(CalculateTFStrides(out_dims), perm);

    std::shared_ptr<stream> transpose_stream;
    in.SetUsrMem(in_dims, in_strides, &in_tensor);
    // Output dimensions are same as input dimensions. We adjust the layout
    // using strides.
    out.SetUsrMem(in_dims, out_strides, out_tensor);

    std::vector<primitive> net;
    MklDnnThreadPool eigen_tp(context);
    auto* prim = FindOrCreateReorder<T>(in.GetUsrMem(), out.GetUsrMem());
    transpose_stream.reset(CreateStream(&eigen_tp, prim->GetEngine()));
    in.SetUsrMemDataHandle(&in_tensor, transpose_stream);
    out.SetUsrMemDataHandle(out_tensor, transpose_stream);
    net.push_back(*(prim->GetPrimitive()));
    std::vector<MemoryArgsMap> net_args;
    net_args.push_back({{DNNL_ARG_FROM, *in.GetUsrMem()},
                        {DNNL_ARG_TO, *out.GetUsrMem()}});
    execute_primitives(net, transpose_stream, net_args);

    return Status::OK();
  } catch (dnnl::error& e) {
    string error_msg = "Status: " + std::to_string(e.status) +
                       ", message: " + std::string(e.message) + ", in file " +
                       std::string(__FILE__) + ":" + std::to_string(__LINE__);
    return errors::Aborted("Operation received an exception:", error_msg);
  }
}

}  // namespace