src/ml/ml_tensors_data_manager.h

   1 /*
   2  * Copyright (c) 2021 Samsung Electronics Co., Ltd All Rights Reserved
   3  *
   4  *    Licensed under the Apache License, Version 2.0 (the "License");
   5  *    you may not use this file except in compliance with the License.
   6  *    You may obtain a copy of the License at
   7  *
   8  *        http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  *    Unless required by applicable law or agreed to in writing, software
  11  *    distributed under the License is distributed on an "AS IS" BASIS,
  12  *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  *    See the License for the specific language governing permissions and
  14  *    limitations under the License.
  15  */
  16
  17 #ifndef __ML_TENSORS_DATA_MANAGER_H__
  18 #define __ML_TENSORS_DATA_MANAGER_H__
  19
  20 #include <unordered_map>
  21
  22 #include "common/logger.h"
  23 #include "common/platform_result.h"
  24 #include "ml_utils.h"
  25
  26 using common::PlatformResult;
  27 using common::ErrorCode;
  28
  29 namespace extension {
  30 namespace ml {
  31
  32 class TensorsInfo;
  33
  34 struct TensorRawData {
  35   // TensorRawData does not take ownership of data, remember to handle it outside
  36   uint8_t* data;
  37   size_t size;
  38   std::string type_str;
  39   unsigned int shape[ML_TENSOR_RANK_LIMIT];
  40 };
  41
  42 class TensorsData {
  43  public:
  44   TensorsData(ml_tensors_data_h handle, int id, TensorsInfo* tensors_info);
  45   ~TensorsData();
  46
  47   ml_tensors_data_h Handle();
  48   int Id();
  49   int TensorsInfoId();
  50   int Count();
  51   ml_tensor_type_e GetTensorType(int index);
  52   PlatformResult GetTensorRawData(int index, int location[ML_TENSOR_RANK_LIMIT],
  53                                   int size[ML_TENSOR_RANK_LIMIT], TensorRawData* tensor_raw_data);
  54   PlatformResult SetTensorRawData(int index, int location[ML_TENSOR_RANK_LIMIT],
  55                                   int size[ML_TENSOR_RANK_LIMIT], TensorRawData& tensor_raw_data);
  56
  57   PlatformResult NativeDestroy();
  58
  59  private:
  60   TensorsData(TensorsData const&) = delete;
  61   TensorsData& operator=(TensorsData const&) = delete;
  62
  63   ml_tensors_data_h handle_;
  64   int id_;
  65   TensorsInfo* tensors_info_;
  66 };
  67
  68 class TensorsDataManager {
  69  public:
  70   TensorsDataManager();
  71   ~TensorsDataManager();
  72
  73   TensorsData* CreateTensorsData(TensorsInfo* tensors_info);
  74   TensorsData* GetTensorsData(int id);
  75
  76   PlatformResult DisposeTensorsData(int id);
  77   PlatformResult DisposeTensorsData(TensorsData* t);
  78
  79  private:
  80   TensorsDataManager(TensorsDataManager const&) = delete;
  81   TensorsDataManager& operator=(TensorsDataManager const&) = delete;
  82
  83   std::unordered_map<int, std::unique_ptr<TensorsData>> map_;
  84   int nextId_;
  85 };
  86
  87 }  // ml
  88 }  // extension
  89 #endif  // __ML_TENSORS_DATA_MANAGER_H__