src/armnn/backends/MakeWorkloadHelper.hpp

   1 //
   2 // Copyright © 2017 Arm Ltd. All rights reserved.
   3 // See LICENSE file in the project root for full license information.
   4 //
   5 #pragma once
   6
   7 namespace armnn
   8 {
   9 namespace
  10 {
  11
  12 // Make a workload of the specified WorkloadType.
  13 template<typename WorkloadType>
  14 struct MakeWorkloadForType
  15 {
  16     template<typename QueueDescriptorType, typename... Args>
  17     static std::unique_ptr<WorkloadType> Func(const QueueDescriptorType& descriptor,
  18                                               const WorkloadInfo& info,
  19                                               Args&&... args)
  20     {
  21         return std::make_unique<WorkloadType>(descriptor, info, std::forward<Args>(args)...);
  22     }
  23 };
  24
  25 // Specialization for void workload type used for unsupported workloads.
  26 template<>
  27 struct MakeWorkloadForType<NullWorkload>
  28 {
  29     template<typename QueueDescriptorType, typename... Args>
  30     static std::unique_ptr<NullWorkload> Func(const QueueDescriptorType& descriptor,
  31                                               const WorkloadInfo& info,
  32                                               Args&&... args)
  33     {
  34         return nullptr;
  35     }
  36 };
  37
  38 // Makes a workload for one the specified types based on the data type requirements of the tensorinfo.
  39 // Specify type void as the WorkloadType for unsupported DataType/WorkloadType combos.
  40 template <typename Float16Workload, typename Float32Workload, typename Uint8Workload, typename QueueDescriptorType,
  41     typename... Args>
  42 std::unique_ptr<IWorkload> MakeWorkload(const QueueDescriptorType& descriptor, const WorkloadInfo& info, Args&&... args)
  43 {
  44     const DataType dataType = !info.m_InputTensorInfos.empty() ?
  45         info.m_InputTensorInfos[0].GetDataType()
  46         : info.m_OutputTensorInfos[0].GetDataType();
  47
  48     BOOST_ASSERT(info.m_InputTensorInfos.empty() || info.m_OutputTensorInfos.empty()
  49         || info.m_InputTensorInfos[0].GetDataType() == info.m_OutputTensorInfos[0].GetDataType());
  50
  51     switch (dataType)
  52     {
  53         case DataType::Float16:
  54             return MakeWorkloadForType<Float16Workload>::Func(descriptor, info, std::forward<Args>(args)...);
  55         case DataType::Float32:
  56             return MakeWorkloadForType<Float32Workload>::Func(descriptor, info, std::forward<Args>(args)...);
  57         case DataType::QuantisedAsymm8:
  58             return MakeWorkloadForType<Uint8Workload>::Func(descriptor, info, std::forward<Args>(args)...);
  59         default:
  60             BOOST_ASSERT_MSG(false, "Unknown DataType.");
  61             return nullptr;
  62     }
  63 }
  64
  65 // Makes a workload for one the specified types based on the data type requirements of the tensorinfo.
  66 // Calling this method is the equivalent of calling the three typed MakeWorkload method with <FloatWorkload,
  67 // FloatWorkload, Uint8Workload>.
  68 // Specify type void as the WorkloadType for unsupported DataType/WorkloadType combos.
  69 template <typename FloatWorkload, typename Uint8Workload, typename QueueDescriptorType, typename... Args>
  70 std::unique_ptr<IWorkload> MakeWorkload(const QueueDescriptorType& descriptor, const WorkloadInfo& info, Args&&... args)
  71 {
  72     return MakeWorkload<FloatWorkload, FloatWorkload, Uint8Workload>(descriptor, info,
  73        std::forward<Args>(args)...);
  74 }
  75
  76
  77 } //namespace
  78 } //namespace armnn