src/backends/CpuTensorHandle.hpp

   1 //
   2 // Copyright © 2017 Arm Ltd. All rights reserved.
   3 // SPDX-License-Identifier: MIT
   4 //
   5 #pragma once
   6
   7 #include "CpuTensorHandleFwd.hpp"
   8 #include <armnn/TypesUtils.hpp>
   9 #include <backends/OutputHandler.hpp>
  10
  11 #include <algorithm>
  12
  13 namespace armnn
  14 {
  15
  16 // Abstract tensor handles wrapping a CPU-readable region of memory, interpreting it as tensor data.
  17 class ConstCpuTensorHandle : public ITensorHandle
  18 {
  19 public:
  20     template <typename T>
  21     const T* GetConstTensor() const
  22     {
  23         BOOST_ASSERT(GetTensorInfo().GetDataType() == GetDataType<T>());
  24         return reinterpret_cast<const T*>(m_Memory);
  25     }
  26
  27     const TensorInfo& GetTensorInfo() const
  28     {
  29         return m_TensorInfo;
  30     }
  31
  32     virtual ITensorHandle::Type GetType() const override
  33     {
  34         return ITensorHandle::Cpu;
  35     }
  36
  37     virtual void Manage() override {}
  38
  39     virtual ITensorHandle* GetParent() const override { return nullptr; }
  40
  41     virtual const void* Map(bool /* blocking = true */) const override { return m_Memory; }
  42     virtual void Unmap() const override {}
  43
  44     TensorShape GetStrides() const override
  45     {
  46         TensorShape shape(m_TensorInfo.GetShape());
  47         auto size = GetDataTypeSize(m_TensorInfo.GetDataType());
  48         auto runningSize = size;
  49         std::vector<unsigned int> strides(shape.GetNumDimensions());
  50         auto lastIdx = shape.GetNumDimensions()-1;
  51         for (unsigned int i=0; i < lastIdx ; i++)
  52         {
  53             strides[lastIdx-i] = runningSize;
  54             runningSize *= shape[lastIdx-i];
  55         }
  56         strides[0] = runningSize;
  57         return TensorShape(shape.GetNumDimensions(), strides.data());
  58     }
  59     TensorShape GetShape() const override { return m_TensorInfo.GetShape(); }
  60
  61 protected:
  62     ConstCpuTensorHandle(const TensorInfo& tensorInfo);
  63
  64     void SetConstMemory(const void* mem) { m_Memory = mem; }
  65
  66 private:
  67     ConstCpuTensorHandle(const ConstCpuTensorHandle& other) = delete;
  68     ConstCpuTensorHandle& operator=(const ConstCpuTensorHandle& other) = delete;
  69
  70     TensorInfo m_TensorInfo;
  71     const void* m_Memory;
  72 };
  73
  74 // Abstract specialization of ConstCpuTensorHandle that allows write access to the same data.
  75 class CpuTensorHandle : public ConstCpuTensorHandle
  76 {
  77 public:
  78     template <typename T>
  79     T* GetTensor() const
  80     {
  81         BOOST_ASSERT(GetTensorInfo().GetDataType() == GetDataType<T>());
  82         return reinterpret_cast<T*>(m_MutableMemory);
  83     }
  84
  85 protected:
  86     CpuTensorHandle(const TensorInfo& tensorInfo);
  87
  88     void SetMemory(void* mem)
  89     {
  90         m_MutableMemory = mem;
  91         SetConstMemory(m_MutableMemory);
  92     }
  93
  94 private:
  95
  96     CpuTensorHandle(const CpuTensorHandle& other) = delete;
  97     CpuTensorHandle& operator=(const CpuTensorHandle& other) = delete;
  98     void* m_MutableMemory;
  99 };
 100
 101 // A CpuTensorHandle that owns the wrapped memory region.
 102 class ScopedCpuTensorHandle : public CpuTensorHandle
 103 {
 104 public:
 105     explicit ScopedCpuTensorHandle(const TensorInfo& tensorInfo);
 106
 107     // Copies contents from Tensor.
 108     explicit ScopedCpuTensorHandle(const ConstTensor& tensor);
 109
 110     // Copies contents from ConstCpuTensorHandle
 111     explicit ScopedCpuTensorHandle(const ConstCpuTensorHandle& tensorHandle);
 112
 113     ScopedCpuTensorHandle(const ScopedCpuTensorHandle& other);
 114     ScopedCpuTensorHandle& operator=(const ScopedCpuTensorHandle& other);
 115     ~ScopedCpuTensorHandle();
 116
 117     virtual void Allocate() override;
 118
 119 private:
 120     void CopyFrom(const ScopedCpuTensorHandle& other);
 121     void CopyFrom(const void* srcMemory, unsigned int numBytes);
 122 };
 123
 124 // A CpuTensorHandle that wraps an already allocated memory region.
 125 //
 126 // Clients must make sure the passed in memory region stays alive for the lifetime of
 127 // the PassthroughCpuTensorHandle instance.
 128 //
 129 // Note there is no polymorphism to/from ConstPassthroughCpuTensorHandle.
 130 class PassthroughCpuTensorHandle : public CpuTensorHandle
 131 {
 132 public:
 133     PassthroughCpuTensorHandle(const TensorInfo& tensorInfo, void* mem)
 134     :   CpuTensorHandle(tensorInfo)
 135     {
 136         SetMemory(mem);
 137     }
 138
 139     virtual void Allocate() override;
 140 };
 141
 142 // A ConstCpuTensorHandle that wraps an already allocated memory region.
 143 //
 144 // This allows users to pass in const memory to a network.
 145 // Clients must make sure the passed in memory region stays alive for the lifetime of
 146 // the PassthroughCpuTensorHandle instance.
 147 //
 148 // Note there is no polymorphism to/from PassthroughCpuTensorHandle.
 149 class ConstPassthroughCpuTensorHandle : public ConstCpuTensorHandle
 150 {
 151 public:
 152     ConstPassthroughCpuTensorHandle(const TensorInfo& tensorInfo, const void* mem)
 153     :   ConstCpuTensorHandle(tensorInfo)
 154     {
 155         SetConstMemory(mem);
 156     }
 157
 158     virtual void Allocate() override;
 159 };
 160
 161
 162 // Template specializations.
 163
 164 template <>
 165 const void* ConstCpuTensorHandle::GetConstTensor() const;
 166
 167 template <>
 168 void* CpuTensorHandle::GetTensor() const;
 169
 170 } // namespace armnn