inference-engine/src/mkldnn_plugin/mkldnn/cpu_prim_tensor.h

   1 // Copyright (C) 2018-2019 Intel Corporation
   2 // SPDX-License-Identifier: Apache-2.0
   3 //
   4
   5 #pragma once
   6
   7 #include "inference_engine.hpp"
   8 #include "prim_tensor.h"
   9 #include "mkldnn.hpp"
  10 #include <memory>
  11
  12 namespace MKLDNNPlugin {
  13
  14 class CpuPrimTensor;
  15
  16 using CpuPrimTensorPtr = std::shared_ptr<CpuPrimTensor>;
  17
  18 class CpuPrimTensor : public PrimTensor {
  19 public:
  20     using Memory = std::shared_ptr<mkldnn::memory>;
  21     using PrimitiveDesc = std::shared_ptr<mkldnn::memory::primitive_desc>;
  22
  23     explicit CpuPrimTensor(mkldnn::memory::desc desc) :
  24             desc(desc) {}
  25
  26
  27     mkldnn::memory getPrimitive() { return *(memory.get()); }
  28
  29 private:
  30     Memory memory;
  31     mkldnn::memory::desc desc;
  32
  33     friend class CpuEngine;
  34 };
  35 }  // namespace MKLDNNPlugin