inference-engine/src/mkldnn_plugin/mkldnn_primitive.h

   1 // Copyright (C) 2018-2019 Intel Corporation
   2 // SPDX-License-Identifier: Apache-2.0
   3 //
   4
   5 #pragma once
   6
   7 #include <mkldnn.hpp>
   8 #include <functional>
   9 #include <ie_common.h>
  10 #include <vector>
  11 #include <memory>
  12 #include <details/ie_exception.hpp>
  13
  14 namespace MKLDNNPlugin {
  15
  16 class MKLDNNPrimitive {
  17 public:
  18     MKLDNNPrimitive();
  19     operator bool();
  20     MKLDNNPrimitive& operator=(const std::shared_ptr<mkldnn::primitive>& prim);
  21     mkldnn::primitive operator*();
  22
  23     void reset(mkldnn::primitive* prim);
  24     void setBatchLimit(int batch, size_t inputNum, size_t outputNum);
  25
  26 private:
  27     std::shared_ptr<mkldnn::primitive> prim;
  28     std::vector<int> originInputBatches;
  29     std::vector<int> originOutputBatches;
  30 };
  31
  32 }  // namespace MKLDNNPlugin