aboutsummaryrefslogtreecommitdiff
path: root/src/armnn/WorkingMemHandle.hpp
blob: bca1d2d80c7bc3b052cc9ecf84bef87361aef094 (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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
//
// Copyright © 2021 Arm Ltd and Contributors. All rights reserved.
// SPDX-License-Identifier: MIT
//

#pragma once

#include "Layer.hpp"
#include "Network.hpp"
#include "WorkingMemDescriptor.hpp"

#include <armnn/IWorkingMemHandle.hpp>
#include <armnn/Tensor.hpp>

#include <unordered_map>
#include <mutex>
#include <backendsCommon/MemoryManager.hpp>

namespace armnn
{

namespace experimental
{


class WorkingMemHandle final : public IWorkingMemHandle
{

public:
    struct InputMemDescriptorCoords
    {
        LayerBindingId m_LayerBindingId;

        std::vector<std::pair<unsigned int, unsigned int>> m_InputSlotCoords;
    };

    struct OutputMemDescriptorCoords
    {
        std::vector<LayerBindingId> m_LayerBindingIds;

        std::pair<unsigned int, unsigned int> m_OutputSlotCoords;
        std::vector<std::pair<unsigned int, unsigned int>> m_InputSlotCoords;
    };

    WorkingMemHandle(NetworkId networkId) : m_NetworkId(networkId){}

    WorkingMemHandle(NetworkId networkId,
                     std::vector<InputMemDescriptorCoords> inputLayerInfo,
                     std::vector<OutputMemDescriptorCoords> outputLayerInfo,
                     std::vector<WorkingMemDescriptor> workingMemDescriptors,
                     std::unordered_map<LayerGuid, WorkingMemDescriptor> workingMemDescriptorMap,
                     std::unique_ptr<MemoryManager> memoryManager,
                     std::vector<std::pair<std::shared_ptr<TensorMemory>, MemorySource>> tensorMemory,
                     std::vector<std::unique_ptr<ITensorHandle>> managedTensorHandles,
                     std::vector<std::unique_ptr<ITensorHandle>> unmanagedTensorHandles);

    ~WorkingMemHandle()
    { Free(); }

    NetworkId GetNetworkId() override
    {
        return m_NetworkId;
    }

    /// Allocate the backing memory required for execution. If this is not called, then allocation will be
    /// deferred to execution time.
    void Allocate() override;

    /// Free the backing memory required for execution.
    void Free() override;

    /// IsAllocated returns true if the backing memory is currently allocated.
    bool IsAllocated() override
    {
        return m_IsAllocated;
    }

    /// Get the WorkingMemDescriptor for a Layer.
    WorkingMemDescriptor& GetWorkingMemDescriptor(LayerGuid id) override
    {
        auto result = m_WorkingMemDescriptorMap.find(id);
        ARMNN_ASSERT(result != m_WorkingMemDescriptorMap.end());
        return result->second;
    }

    /// Get the WorkingMemDescriptor at an index. The WorkingMemDescriptors are stored in the same order as
    /// the Workloads in a topologically sorted graph.
    WorkingMemDescriptor& GetWorkingMemDescriptorAt(unsigned int id) override
    {
        return m_WorkingMemDescriptors[id];
    }

    ITensorHandle* GetInputHandle(LayerBindingId layerBindingId) const
    {
        return m_InputHandleMap.at(layerBindingId);
    };

    ITensorHandle* GetOutputHandle(LayerBindingId layerBindingId) const
    {
        return m_OutputHandleMap.at(layerBindingId);
    };

    const std::vector<std::vector<ITensorHandle*>::iterator>& GetInputConnections(LayerBindingId layerBindingId) const
    {
        return m_InputConnectionMap.at(layerBindingId);
    };

    const std::vector<std::vector<ITensorHandle*>::iterator>& GetOutputConnection(LayerBindingId layerBindingId) const
    {
        return m_OutputConnectionMap.at(layerBindingId);
    };

    void MemSyncOutputs();

    std::vector<LayerBindingId>& GetBindingIdVector()
    {
        return m_BindingIdVec;
    };

    void ValidateBindingIds();

private:
    using DifferenceType = std::vector<ITensorHandle*>::difference_type;
    NetworkId m_NetworkId;

    std::unordered_map<LayerBindingId, ITensorHandle*> m_InputHandleMap;
    std::unordered_map<LayerBindingId, ITensorHandle*> m_OutputHandleMap;
    std::unordered_map<LayerBindingId, std::vector<std::vector<ITensorHandle*>::iterator>> m_InputConnectionMap;
    std::unordered_map<LayerBindingId, std::vector<std::vector<ITensorHandle*>::iterator>> m_OutputConnectionMap;

    std::vector<WorkingMemDescriptor> m_WorkingMemDescriptors;
    std::unordered_map<LayerGuid, WorkingMemDescriptor> m_WorkingMemDescriptorMap;

    std::unique_ptr<MemoryManager> m_MemoryManager;

    // Memory to be imported into the tensorHandles after allocation
    std::vector<std::pair<std::shared_ptr<TensorMemory>, MemorySource>> m_TensorMemory;


    // Tensors that will need to be allocated internally within armnn
    std::vector<std::unique_ptr<ITensorHandle>> m_ManagedTensorHandles;

    // Tensors that will be allocated externally by the user
    std::vector<std::unique_ptr<ITensorHandle>> m_UnmanagedTensorHandles;

    std::unordered_map<LayerBindingId, bool> m_InputValidationMap;
    std::unordered_map<LayerBindingId, bool> m_OutputValidationMap;

    std::vector<LayerBindingId> m_BindingIdVec;

    DifferenceType m_InputSize;

    bool m_IsAllocated;
};

} // end experimental namespace

} // end armnn namespace