WorkingMemHandle.hpp

Go to the documentation of this file.

//
// Copyright © 2022 Arm Ltd and Contributors. All rights reserved.
// SPDX-License-Identifier: MIT
//
 
#pragma once
 
#include "ExecutionData.hpp"
#include "Layer.hpp"
#include "Network.hpp"
#include "WorkingMemDescriptor.hpp"
 
#include <armnn/IWorkingMemHandle.hpp>
#include <armnn/Tensor.hpp>
#include <armnn/backends/IBackendInternal.hpp>
 
#include <unordered_map>
#include <mutex>
#include <backendsCommon/MemoryManager.hpp>
 
namespace armnn
{
 
namespace experimental
{
 
using BackendPtrMap = std::unordered_map<BackendId, IBackendInternalUniquePtr>;
 
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::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,
                     std::vector<std::pair<BackendId, ExecutionData>> executionDataVec,
                     BackendPtrMap* backends);
 
    ~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 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];
    }
 
    /// Get the ExecutionData at an index.
    /// The ExecutionData is paired with a BackendId to be able to call backend specific functions upon it.
    /// The ExecutionData are stored in the same order as the Workloads in a topologically sorted graph.
    std::pair<BackendId, ExecutionData>& GetExecutionDataAt(unsigned int id) override
    {
        return m_ExecutionDataVec[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::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;
 
    std::vector<std::pair<BackendId, ExecutionData>> m_ExecutionDataVec;
 
    BackendPtrMap* m_Backends;
};
 
} // end experimental namespace
 
} // end armnn namespace