diff options
Diffstat (limited to 'src/armnn/Layer.hpp')
-rw-r--r-- | src/armnn/Layer.hpp | 19 |
1 files changed, 19 insertions, 0 deletions
diff --git a/src/armnn/Layer.hpp b/src/armnn/Layer.hpp index fff57732aa..6cb3b8f5bd 100644 --- a/src/armnn/Layer.hpp +++ b/src/armnn/Layer.hpp @@ -28,6 +28,7 @@ #include <memory> #include <string> #include <vector> +#include <backendsCommon/WorkloadData.hpp> namespace armnn { @@ -204,6 +205,7 @@ class ScopedCpuTensorHandle; // Base layer class using LayerPriority = unsigned int; +using AdditionalInfoObjectPtr = std::shared_ptr<void>; class Layer : public IConnectableLayer { @@ -333,6 +335,17 @@ public: m_ShapeInferenceMethod = shapeInferenceMethod; } + template<typename T> + std::shared_ptr<T> GetAdditionalInformation() + { + return std::static_pointer_cast<T>(m_AdditionalInfoObject); + } + + void SetAdditionalInfoForObject(const AdditionalInfoObjectPtr& additionalInfo) + { + m_AdditionalInfoObject = additionalInfo; + } + protected: // Graph needs access to the virtual destructor. friend class Graph; @@ -377,6 +390,12 @@ protected: using ConstantTensors = std::vector<std::reference_wrapper<std::unique_ptr<ScopedCpuTensorHandle>>>; virtual ConstantTensors GetConstantTensorsByRef() {return ConstantTensors(); }; + // "Blob" + AdditionalInfoObjectPtr m_AdditionalInfoObject; + + // Utility method to set a pointer in the queueDescriptor to the "blob" location in the layer + void SetAdditionalInfo(QueueDescriptor& descriptor) const; + private: void CollectWorkloadInputs(WorkloadDataCollector& dataCollector) const; void CollectWorkloadOutputs(WorkloadDataCollector& dataCollector) const; |