blob: 91c586b4a1b3c16f5f21b3ba3895a269aa9f95f6 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
|
// Copyright (C) 2018 Intel Corporation
//
// SPDX-License-Identifier: Apache-2.0
//
#pragma once
#include <ie_api.h>
#include <memory>
#include "mkldnn_memory.h"
#include "mkldnn_dims.h"
#include <map>
#include <vector>
namespace MKLDNNPlugin {
class MKLDNNNode;
class MKLDNNEdge;
using MKLDNNEdgePtr = std::shared_ptr<MKLDNNEdge>;
using MKLDNNEdgeWeakPtr = std::weak_ptr<MKLDNNEdge>;
class MKLDNNEdge : public InferenceEngine::details::no_copy {
public:
enum class Status {
Uninitialized,
NeedAllocation,
NotAllocated,
Allocated,
Validated
};
MKLDNNEdge(const std::shared_ptr<MKLDNNNode>& parent, const std::shared_ptr<MKLDNNNode>& child);
inline Status getStatus() noexcept {
return status;
}
void changeStatus(Status state);
virtual void init();
virtual void allocate(const void* mem_ptr = nullptr);
virtual void validate();
const std::shared_ptr<MKLDNNNode> getParent() const;
const std::shared_ptr<MKLDNNNode> getChild() const;
bool needReorder();
InferenceEngine::Blob::Ptr getBlob();
const MKLDNNMemory& getMemory();
MKLDNNMemoryPtr& getMemoryPtr();
bool isDropped();
InferenceEngine::TensorDesc getDesc();
int getInputNum();
int getOutputNum();
std::vector<int> getAllOutputNums();
std::vector<int> getAllInputNums();
MKLDNNDims &getDims();
void setDims(MKLDNNDims &dims);
void sharedMemFrom(const MKLDNNEdgePtr& edge);
MKLDNNEdgePtr getSharedEdge() const;
private:
std::weak_ptr<MKLDNNNode> parent;
std::weak_ptr<MKLDNNNode> child;
MKLDNNEdgeWeakPtr memoryFromEdge;
MKLDNNDims dims;
MKLDNNMemoryPtr memoryPtr;
Status status = Status::Uninitialized;
InferenceEngine::TensorDesc getInputDesc();
InferenceEngine::TensorDesc getOutputDesc();
InferenceEngine::TensorDesc getSpecifiedInputDesc(std::map<mkldnn::memory::format, size_t> formats);
InferenceEngine::TensorDesc getSpecifiedOutputDesc(std::map<mkldnn::memory::format, size_t> formats);
InferenceEngine::TensorDesc inputDesc;
InferenceEngine::TensorDesc outputDesc;
bool nodeCanChangeDesc(const std::shared_ptr<MKLDNNPlugin::MKLDNNNode>& node) const;
enum LOOK { LOOK_UP = 1, LOOK_DOWN = 2, LOOK_BOTH = LOOK_UP | LOOK_DOWN };
MKLDNNEdgePtr getBaseEdge(LOOK look = LOOK_BOTH);
bool inPlace(LOOK look = LOOK_BOTH);
friend class MKLDNNGraph;
};
} // namespace MKLDNNPlugin
|
