!17248 [lite]open allocator

From: @xu_anyue
Reviewed-by: 
Signed-off-by:

mindspore/lite/include/allocator.h

@@ -0,0 +1,84 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef MINDSPORE_LITE_INCLUDE_ALLOCATOR_H_
+#define MINDSPORE_LITE_INCLUDE_ALLOCATOR_H_
+
+#include <memory>
+#include "include/lite_utils.h"
+
+namespace mindspore {
+/// \brief Allocator defined a memory pool for malloc memory and free memory dynamically.
+class MS_API Allocator {
+ public:
+  /// \brief Destructor of MindSpore Allocator.
+  virtual ~Allocator() = default;
+
+  /// \brief Method to request memory.
+  ///
+  /// \param[in] size Define the memory size to request.
+  virtual void *Malloc(size_t size) = 0;
+
+  /// \brief Method to free memory.
+  ///
+  /// \param[in] ptr Define the pointer of a certain memory.
+  virtual void Free(void *ptr) = 0;
+
+  /// \brief Reference count of a certain memory.
+  ///
+  /// \param[in] ptr Define the pointer of a certain memory.
+  ///
+  /// \return Reference count of a certain memory currently.
+  virtual int RefCount(void *ptr) = 0;
+
+  /// \brief Set reference count of a certain memory.
+  ///
+  /// \param[in] ptr Define the pointer of a certain memory.
+  /// \param[in] ref_count Define the reference count to set.
+  ///
+  /// \return Reference count of a certain memory after setting.
+  virtual int SetRefCount(void *ptr, int ref_count) = 0;
+
+  /// \brief Decrease the reference count of a certain memory.
+  ///
+  /// \param[in] ptr Define the pointer of a certain memory.
+  /// \param[in] ref_count Define the reference count to reduce.
+  ///
+  /// \return Reference count of a certain memory after decreating.
+  virtual int DecRefCount(void *ptr, int ref_count) = 0;
+
+  /// \brief Increase the reference count of a certain memory.
+  ///
+  /// \param[in] ptr Define the pointer of a certain memory.
+  /// \param[in] ref_count Define the reference count to increase.
+  ///
+  /// \return Reference count of a certain memory after increasing.
+  virtual int IncRefCount(void *ptr, int ref_count) = 0;
+
+  /// \brief Static method to create an allocator.
+  ///
+  /// \return Smart pointer of an allocator.
+  static std::shared_ptr<Allocator> Create();
+
+  /// \brief Prepare a certain memory.
+  ///
+  /// \param[in] ptr Define the pointer of a certain memory to prepare.
+  ///
+  /// \return Pointer of ready memory.
+  virtual void *Prepare(void *ptr) { return ptr; }
+};
+}  // namespace mindspore
+#endif  // MINDSPORE_LITE_INCLUDE_ALLOCATOR_H_

mindspore/lite/include/lite_utils.h

@@ -602,6 +602,7 @@ class Vector {
 };
 using TensorPtrVector = Vector<mindspore::schema::Tensor *>;
 using Uint32Vector = Vector<uint32_t>;
+class Allocator;
 using AllocatorPtr = void *;
 using DeviceContextVector = Vector<lite::DeviceContext>;
 using KernelCallBack = void (*)(void *, void *);

mindspore/lite/include/ms_tensor.h

@@ -39,6 +39,16 @@ class MS_API MSTensor {
   static MSTensor *CreateTensor(const String &name, TypeId type, const Vector<int> &shape, const void *data,
                                 size_t data_len);
 
+  /// \brief Set memory allocator for current MSTensor.
+  ///
+  /// \param[in] allocator Define memory allocator, which is shown in allocator.h.
+  virtual void set_allocator(mindspore::Allocator *allocator) = 0;
+
+  /// \brief Get memory allocator of current MSTensor.
+  ///
+  /// \return Pointer of memory allocator class.
+  virtual mindspore::Allocator *allocator() const = 0;
+
   /// \brief Get data type of the MindSpore Lite MSTensor.
   ///
   /// \note TypeId is defined in mindspore/mindspore/include/api/type_id.h. Only number types in TypeId enum are
@@ -47,6 +57,9 @@ class MS_API MSTensor {
   /// \return MindSpore Lite TypeId of the MindSpore Lite MSTensor.
   virtual TypeId data_type() const = 0;
 
+  /// \brief Set data type of current MSTensor.
+  virtual void set_data_type(TypeId data_type) = 0;
+
   /// \brief Get shape of the MindSpore Lite MSTensor.
   ///
   /// \return A vector of int as the shape of the MindSpore Lite MSTensor.

mindspore/lite/include/registry/kernel_interface.h

@@ -14,8 +14,8 @@
  * limitations under the License.
  */
 
-#ifndef MINDSPORE_LITE_SRC_REGISTRY_KERNEL_INTERFACE_H_
-#define MINDSPORE_LITE_SRC_REGISTRY_KERNEL_INTERFACE_H_
+#ifndef MINDSPORE_LITE_INCLUDE_REGISTRY_KERNEL_INTERFACE_H_
+#define MINDSPORE_LITE_INCLUDE_REGISTRY_KERNEL_INTERFACE_H_
 
 #include <set>
 #include <string>
@@ -27,51 +27,114 @@
 
 namespace mindspore {
 namespace kernel {
+/// \brief CapabilityParam defined performance of op when running.
 struct MS_API CapabilityParam {
-  float exec_time_;
-  float power_usage_;
+  float exec_time_;   /**< op running time argument */
+  float power_usage_; /**< op power waste argument */
 };
 
+/// \brief KernelInterface defined customized op's interface, such as infershape, and so on.
 class MS_API KernelInterface {
  public:
+  /// \brief Destructor of KernelInterface.
   virtual ~KernelInterface() = default;
+
+  /// \brief Method to infer customized op's output shape.
+  ///
+  /// \param[in] inputs Define the input tensors of op.
+  /// \param[in] outputs Define the output tensors of op.
+  /// \param[in] primitive Define the attributes of op.
+  ///
+  /// \return  STATUS as an error code of inferring, STATUS is defined in errorcode.h..
   virtual int Infer(const std::vector<tensor::MSTensor *> &inputs, const std::vector<tensor::MSTensor *> &outputs,
                     const schema::Primitive *primitive) {
     return 0;
   }
 
+  /// \brief Method to get performance of an op when running.
+  ///
+  /// \param[in] tensor_in Define the input tensors of op.
+  /// \param[in] primitive Define the attributes of op.
+  /// \param[in] param Define the contr of performance.
+  ///
+  /// \return STATUS as an error code of inferring, STATUS is defined in errorcode.h.
   virtual int GetCapability(const std::vector<tensor::MSTensor *> &tensor_in, const schema::Primitive *primitive,
                             CapabilityParam *param) {
     return 0;
   }
 };
 
+/// \brief KernelInterfaceCreator defined a functor to create KernelInterface.
 using KernelInterfaceCreator MS_API = std::function<std::shared_ptr<KernelInterface>()>;
 
+/// \brief RegisterKernelInterface defined registration and acquisition of KernelInterface.
 class MS_API RegisterKernelInterface {
  public:
+  /// \brief Static method to register op whose primitive type is custom.
+  ///
+  /// \param[in] provider Define the identification of user.
+  /// \param[in] op_type Define the concrete type of a custom op.
+  /// \param[in] creator Define the KernelInterface create function.
+  ///
+  /// \return STATUS as an error code of registering, STATUS is defined in errorcode.h.
   static int CustomReg(const std::string &provider, const std::string &op_type, KernelInterfaceCreator creator);
+
+  /// \brief Static method to register op whose primitive type is ordinary.
+  ///
+  /// \param[in] provider Define the identification of user.
+  /// \param[in] op_type Define the ordinary op type.
+  /// \param[in] creator Define the KernelInterface create function.
+  ///
+  /// \return STATUS as an error code of registering, STATUS is defined in errorcode.h.
   static int Reg(const std::string &provider, int op_type, KernelInterfaceCreator creator);
+
+  /// \brief Static method to get registration of a certain op.
+  ///
+  /// \param[in] provider Define the identification of user.
+  /// \param[in] primitive Define the attributes of a certain op.
+  ///
+  /// \return Boolean value to represent registration of a certain op is existing or not.
   static std::shared_ptr<kernel::KernelInterface> GetKernelInterface(const std::string &provider,
                                                                      const schema::Primitive *primitive);
 };
 
+/// \brief KernelInterfaceReg defined registration class of KernelInterface.
 class MS_API KernelInterfaceReg {
  public:
+  /// \brief Constructor of KernelInterfaceReg to register an ordinary op.
+  ///
+  /// \param[in] provider Define the identification of user.
+  /// \param[in] op_type Define the ordinary op type.
+  /// \param[in] creator Define the KernelInterface create function.
   KernelInterfaceReg(const std::string &provider, int op_type, KernelInterfaceCreator creator) {
     RegisterKernelInterface::Reg(provider, op_type, creator);
   }
 
+  /// \brief Constructor of KernelInterfaceReg to register custom op.
+  ///
+  /// \param[in] provider Define the identification of user.
+  /// \param[in] op_type Define the concrete type of a custom op.
+  /// \param[in] creator Define the KernelInterface create function.
   KernelInterfaceReg(const std::string &provider, const std::string &op_type, KernelInterfaceCreator creator) {
     RegisterKernelInterface::CustomReg(provider, op_type, creator);
   }
 };
 
+/// \brief Defined registering macro to register ordinary op, which called by user directly.
+///
+/// \param[in] provider Define the identification of user.
+/// \param[in] op_type Define the ordinary op type.
+/// \param[in] creator Define the KernelInterface create function.
 #define REGISTER_KERNEL_INTERFACE(provider, op_type, creator)                               \
   namespace {                                                                               \
   static KernelInterfaceReg g_##provider##op_type##_inter_reg(#provider, op_type, creator); \
   }  // namespace
 
+/// \brief Defined registering macro to register custom op, which called by user directly.
+///
+/// \param[in] provider Define the identification of user.
+/// \param[in] op_type Define the concrete type of a custom op.
+/// \param[in] creator Define the KernelInterface create function.
 #define REGISTER_CUSTOM_KERNEL_INTERFACE(provider, op_type, creator)                                \
   namespace {                                                                                       \
   static KernelInterfaceReg g_##provider##op_type##_custom_inter_reg(#provider, #op_type, creator); \
@@ -79,4 +142,4 @@ class MS_API KernelInterfaceReg {
 }  // namespace kernel
 }  // namespace mindspore
 
-#endif  // MINDSPORE_LITE_SRC_REGISTRY_KERNEL_INTERFACE_H_
+#endif  // MINDSPORE_LITE_INCLUDE_REGISTRY_KERNEL_INTERFACE_H_

mindspore/lite/include/registry/model_parser_registry.h

@@ -14,38 +14,70 @@
  * limitations under the License.
  */
 
-#ifndef MINDSPORE_LITE_TOOLS_CONVERTER_MODEL_PARSER_REGISTRY_H
-#define MINDSPORE_LITE_TOOLS_CONVERTER_MODEL_PARSER_REGISTRY_H
+#ifndef MINDSPORE_LITE_INCLUDE_REGISTRY_MODEL_PARSER_REGISTRY_H
+#define MINDSPORE_LITE_INCLUDE_REGISTRY_MODEL_PARSER_REGISTRY_H
 #include <string>
 #include <memory>
 #include <unordered_map>
 #include "include/lite_utils.h"
 
 namespace mindspore::lite {
+/// \brief ModelParser defined a model parser
 class MS_API ModelParser;
+
+/// \brief ModelParserCreator defined function pointer to get a ModelParser class.
 typedef ModelParser *(*ModelParserCreator)();
 
+/// \brief ModelParserRegistry defined registration and storage of ModelParser.
 class MS_API ModelParserRegistry {
  public:
+  /// \brief Constructor of ModelParserRegistry.
   ModelParserRegistry() = default;
+
+  /// \brief Destructor of ModelParserRegistry.
   ~ModelParserRegistry() = default;
 
+  /// \brief Static method to get a single instance.
+  ///
+  /// \return Pointer of ModelParserRegistry.
   static ModelParserRegistry *GetInstance();
+
+  /// \brief Method to get a model parser.
+  ///
+  /// \param[in] fmk Define identification of a certain framework.
+  ///
+  /// \return Pointer of ModelParser.
   ModelParser *GetModelParser(const std::string &fmk);
+
+  /// \brief Method to register model parser.
+  ///
+  /// \param[in] fmk Define identification of a certain framework.
+  /// \param[in] creator Define function pointer of creating ModelParser.
   void RegParser(const std::string &fmk, ModelParserCreator creator);
 
   std::unordered_map<std::string, ModelParserCreator> parsers_;
 };
 
+/// \brief ModelRegistrar defined registration class of ModelParser.
 class MS_API ModelRegistrar {
  public:
+  /// \brief Constructor of ModelRegistrar to register ModelParser.
+  ///
+  /// \param[in] fmk Define identification of a certain framework.
+  /// \param[in] creator Define function pointer of creating ModelParser.
   ModelRegistrar(const std::string &fmk, ModelParserCreator creator) {
     ModelParserRegistry::GetInstance()->RegParser(fmk, creator);
   }
+
+  /// \brief Destructor of ModelRegistrar.
   ~ModelRegistrar() = default;
 };
 
+/// \brief Defined registering macro to register ModelParser, which called by user directly.
+///
+/// \param[in] fmk Define identification of a certain framework.
+/// \param[in] parserCreator Define function pointer of creating ModelParser.
 #define REG_MODEL_PARSER(fmk, parserCreator) static ModelRegistrar g_##type##fmk##ModelParserReg(#fmk, parserCreator);
 }  // namespace mindspore::lite
 
-#endif  // MINDSPORE_LITE_TOOLS_CONVERTER_MODEL_PARSER_REGISTRY_H
+#endif  // MINDSPORE_LITE_INCLUDE_REGISTRY_MODEL_PARSER_REGISTRY_H

mindspore/lite/include/registry/pass_registry.h

@@ -14,8 +14,8 @@
  * limitations under the License.
  */
 
-#ifndef MINDSPORE_LITE_TOOLS_CONVERTER_REGISTRY_PASS_REGISTRY_H_
-#define MINDSPORE_LITE_TOOLS_CONVERTER_REGISTRY_PASS_REGISTRY_H_
+#ifndef MINDSPORE_LITE_INCLUDE_REGISTRY_PASS_REGISTRY_H_
+#define MINDSPORE_LITE_INCLUDE_REGISTRY_PASS_REGISTRY_H_
 
 #include <vector>
 #include <string>
@@ -27,18 +27,39 @@
 
 namespace mindspore {
 namespace opt {
+/// \brief PassPosition defined where to plae user's pass.
 enum MS_API PassPosition { POSITION_BEGIN = 0, POSITION_END = 1 };
 
+/// \brief P defined a basic interface.
+///
+/// \note List public class and interface for reference.
 class MS_API Pass;
 using PassPtr = std::shared_ptr<Pass>;
+
+/// \brief PassRegistry defined registration of Pass.
 class MS_API PassRegistry {
  public:
+  /// \brief Destructor of PassRegistry.
   virtual ~PassRegistry() = default;
+
+  /// \brief Static method to get a single instance of PassRegistry.
+  ///
+  /// \return Pointer of PassRegistry.
   static PassRegistry *GetInstance();
+
+  /// \brief Method to register Pass.
+  ///
+  /// \param[in] position Define where to replace the pass.
+  /// \param[in] pass Define user's defined pass.
   void RegPass(int position, const PassPtr &pass);
+
+  /// \brief Method to get all passes user write.
+  ///
+  /// \return A map include all pass.
   const std::unordered_map<int, PassPtr> &GetPasses() const;
 
  private:
+  /// \brief Constructor of PassRegistry.
   PassRegistry() = default;
 
  private:
@@ -46,15 +67,26 @@ class MS_API PassRegistry {
   std::mutex mutex_;
 };
 
+/// \brief PassRegistrar defined registration class of Pass.
 class MS_API PassRegistrar {
  public:
+  /// \brief Constructor of PassRegistrar to register pass.
+  ///
+  /// \param[in] pos Define where to replace the pass.
+  /// \param[in] pass Define user's defined pass.
   PassRegistrar(int pos, const PassPtr &pass) { PassRegistry::GetInstance()->RegPass(pos, pass); }
+
+  /// \brief Destructor of PassRegistrar.
   ~PassRegistrar() = default;
 };
 
+/// \brief Defined registering macro to register Pass, which called by user directly.
+///
+/// \param[in] position Define where to replace the pass.
+/// \param[in] pass Define user's defined pass.
 #define REG_PASS(position, pass) static PassRegistrar g_##position##PassReg(position, std::make_shared<pass>());
 
 }  // namespace opt
 }  // namespace mindspore
 
-#endif  // MINDSPORE_LITE_TOOLS_CONVERTER_REGISTRY_PASS_REGISTRY_H_
+#endif  // MINDSPORE_LITE_INCLUDE_REGISTRY_PASS_REGISTRY_H_

mindspore/lite/include/registry/register_kernel.h

@@ -14,8 +14,8 @@
  * limitations under the License.
  */
 
-#ifndef MINDSPORE_LITE_SRC_REGISTRY_REGISTER_KERNEL_H_
-#define MINDSPORE_LITE_SRC_REGISTRY_REGISTER_KERNEL_H_
+#ifndef MINDSPORE_LITE_INCLUDE_REGISTRY_REGISTER_KERNEL_H_
+#define MINDSPORE_LITE_INCLUDE_REGISTRY_REGISTER_KERNEL_H_
 
 #include <set>
 #include <string>
@@ -31,11 +31,12 @@ namespace kernel {
 static const char *const kArchCPU __attribute__((unused)) = "CPU";
 static const char *const kArchGPU __attribute__((unused)) = "GPU";
 
+/// \brief KernelDesc defined kernel's basic attribute.
 struct MS_API KernelDesc {
-  TypeId data_type;
-  int type;
-  std::string arch;
-  std::string provider;
+  TypeId data_type;     /**< kernel data type argument */
+  int type;             /**< op type argument */
+  std::string arch;     /**< deviceType argument */
+  std::string provider; /**< user identification argument */
 
   bool operator<(const KernelDesc &dst) const {
     if (provider != dst.provider) {
@@ -50,38 +51,103 @@ struct MS_API KernelDesc {
   }
 };
 
+/// \brief CreateKernel Defined a functor to create a kernel.
+///
+/// \param[in] inputs Define input tensors of kernel.
+/// \param[in] outputs Define output tensors of kernel.
+/// \param[in] primitive Define attributes of op.
+/// \param[in] ctx Define for holding environment variables during runtime.
+///
+/// \return Smart Pointer of kernel.
 using CreateKernel MS_API = std::function<std::shared_ptr<kernel::Kernel>(
   const std::vector<tensor::MSTensor *> &inputs, const std::vector<tensor::MSTensor *> &outputs,
   const schema::Primitive *primitive, const lite::Context *ctx)>;
 
+/// \brief RegisterKernel Defined registration of kernel.
 class MS_API RegisterKernel {
  public:
+  /// \brief Static method to register kernel which is correspondng to an ordinary op.
+  ///
+  /// \param[in] arch Define deviceType, such as CPU.
+  /// \param[in] provider Define the identification of user.
+  /// \param[in] data_type Define kernel's input data type.
+  /// \param[in] type Define the ordinary op type.
+  /// \param[in] creator Define a function pointer to create a kernel.
+  ///
+  /// \return STATUS as an error code of registering, STATUS is defined in errorcode.h.
   static int RegKernel(const std::string &arch, const std::string &provider, TypeId data_type, int type,
                        CreateKernel creator);
+
+  /// \brief Static method to register kernel which is corresponding to custom op.
+  ///
+  /// \param[in] arch Define deviceType, such as CPU.
+  /// \param[in] provider Define the identification of user.
+  /// \param[in] data_type Define kernel's input data type.
+  /// \param[in] type Define the concrete type of a custom op.
+  /// \param[in] creator Define a function pointer to create a kernel.
+  ///
+  /// \return STATUS as an error code of registering, STATUS is defined in errorcode.h.
   static int RegCustomKernel(const std::string &arch, const std::string &provider, TypeId data_type,
                              const std::string &type, CreateKernel creator);
+
+  /// \brief Static methon to get a kernel's create function.
+  ///
+  /// \param[in] desc Define kernel's basic attribute.
+  /// \param[in] primitive Define the attributes of op.
+  ///
+  /// \return Function pointer to create a kernel.
   static CreateKernel GetCreator(const kernel::KernelDesc &desc, const schema::Primitive *primitive);
 };
 
+/// \brief KernelReg Defined registration class of kernel.
 class MS_API KernelReg {
  public:
+  /// \brief Destructor of KernelReg.
   ~KernelReg() = default;
 
+  /// \brief Method to register ordinary op.
+  ///
+  /// \param[in] arch Define deviceType, such as CPU.
+  /// \param[in] provider Define the identification of user.
+  /// \param[in] data_type Define kernel's input data type.
+  /// \param[in] op_type Define the ordinary op type.
+  /// \param[in] creator Define a function pointer to create a kernel.
   KernelReg(const std::string &arch, const std::string &provider, TypeId data_type, int op_type, CreateKernel creator) {
     RegisterKernel::RegKernel(arch, provider, data_type, op_type, creator);
   }
 
+  /// \brief Method to register customized op.
+  ///
+  /// \param[in] arch Define deviceType, such as CPU.
+  /// \param[in] provider Define the identification of user.
+  /// \param[in] data_type Define kernel's input data type.
+  /// \param[in] op_type Define the concrete type of a custom op.
+  /// \param[in] creator Define a function pointer to create a kernel.
   KernelReg(const std::string &arch, const std::string &provider, TypeId data_type, const std::string &op_type,
             CreateKernel creator) {
     RegisterKernel::RegCustomKernel(arch, provider, data_type, op_type, creator);
   }
 };
 
+/// \brief Defined registering macro to register ordinary op kernel, which called by user directly.
+///
+/// \param[in] arch Define deviceType, such as CPU.
+/// \param[in] provider Define the identification of user.
+/// \param[in] data_type Define kernel's input data type.
+/// \param[in] op_type Define the ordinary op type.
+/// \param[in] creator Define a function pointer to create a kernel.
 #define REGISTER_KERNEL(arch, provider, data_type, op_type, creator)                                                 \
   namespace {                                                                                                        \
   static KernelReg g_##arch##provider##data_type##op_type##kernelReg(#arch, #provider, data_type, op_type, creator); \
   }  // namespace
 
+/// \brief Defined registering macro to register custom op kernel, which called by user directly.
+///
+/// \param[in] arch Define deviceType, such as CPU.
+/// \param[in] provider Define the identification of user.
+/// \param[in] data_type Define kernel's input data type.
+/// \param[in] op_type Define the concrete type of a custom op.
+/// \param[in] creator Define a function pointer to create a kernel.
 #define REGISTER_CUSTOM_KERNEL(arch, provider, data_type, op_type, creator)                                           \
   namespace {                                                                                                         \
   static KernelReg g_##arch##provider##data_type##op_type##kernelReg(#arch, #provider, data_type, #op_type, creator); \
@@ -89,4 +155,4 @@ class MS_API KernelReg {
 }  // namespace kernel
 }  // namespace mindspore
 
-#endif  // MINDSPORE_LITE_SRC_REGISTRY_REGISTER_KERNEL_H_
+#endif  // MINDSPORE_LITE_INCLUDE_REGISTRY_REGISTER_KERNEL_H_

mindspore/lite/micro/coder/generator/component/const_blocks/mtensor.cc

@@ -61,7 +61,10 @@ class MTensor : public mindspore::tensor::MSTensor {
   MTensor(String name, TypeId type, Vector<int32_t> shape) : tensor_name_(name), data_type_(type), shape_(shape) {}
   ~MTensor() override;
 
+  void set_allocator(mindspore::Allocator *allocator) override {}
+  mindspore::Allocator *allocator() const override { return nullptr; }
   TypeId data_type() const override { return data_type_; }
+  void set_data_type(TypeId data_type) override { data_type_ = data_type; }
   Vector<int> shape() const override { return shape_; }
   void set_shape(const Vector<int> &shape) override { shape_ = shape; }
   int ElementsNum() const override;

mindspore/lite/micro/example/mnist_stm32f746/mnist_stm32f746/include/ms_tensor.h

@@ -39,6 +39,16 @@ class MS_API MSTensor {
   static MSTensor *CreateTensor(const String &name, TypeId type, const Vector<int> &shape, const void *data,
                                 size_t data_len);
 
+  /// \brief Set memory allocator for current MSTensor.
+  ///
+  /// \param[in] allocator Define memory allocator, which is shown in allocator.h.
+  virtual void set_allocator(mindspore::Allocator *allocator) = 0;
+
+  /// \brief Get memory allocator of current MSTensor.
+  ///
+  /// \return Pointer of memory allocator class.
+  virtual mindspore::Allocator *allocator() const = 0;
+
   /// \brief Get data type of the MindSpore Lite MSTensor.
   ///
   /// \note TypeId is defined in mindspore/mindspore/include/api/type_id.h. Only number types in TypeId enum are
@@ -47,6 +57,9 @@ class MS_API MSTensor {
   /// \return MindSpore Lite TypeId of the MindSpore Lite MSTensor.
   virtual TypeId data_type() const = 0;
 
+  /// \brief Set data type of current MSTensor.
+  virtual void set_data_type(TypeId data_type) = 0;
+
   /// \brief Get shape of the MindSpore Lite MSTensor.
   ///
   /// \return A vector of int as the shape of the MindSpore Lite MSTensor.

mindspore/lite/micro/example/mnist_stm32f746/mnist_stm32f746/src/tensor.h

@@ -41,7 +41,10 @@ class MTensor : public mindspore::tensor::MSTensor {
   MTensor(String name, TypeId type, Vector<int> shape) : tensor_name_(name), data_type_(type), shape_(shape) {}
   ~MTensor() override;
 
+  void set_allocator(mindspore::Allocator *allocator) override {}
+  mindspore::Allocator *allocator() const override { return nullptr; }
   TypeId data_type() const override { return data_type_; }
+  void set_data_type(TypeId data_type) override { data_type_ = data_type; }
   Vector<int> shape() const override { return shape_; }
   void set_shape(const Vector<int> &shape) override { shape_ = shape; }
   int ElementsNum() const override;

mindspore/lite/micro/example/mnist_x86/src/tensor.h

@@ -41,7 +41,10 @@ class MTensor : public mindspore::tensor::MSTensor {
   MTensor(String name, TypeId type, Vector<int> shape) : tensor_name_(name), data_type_(type), shape_(shape) {}
   ~MTensor() override;
 
+  void set_allocator(mindspore::Allocator *allocator) override {}
+  mindspore::Allocator *allocator() const override { return nullptr; }
   TypeId data_type() const override { return data_type_; }
+  void set_data_type(TypeId data_type) override { data_type_ = data_type; }
   Vector<int> shape() const override { return shape_; }
   void set_shape(const Vector<int> &shape) override { shape_ = shape; }
   int ElementsNum() const override;

mindspore/lite/src/CMakeLists.txt

@@ -53,7 +53,7 @@ set(LITE_SRC
         ${CMAKE_CURRENT_SOURCE_DIR}/common/tensor_util.cc
         ${CMAKE_CURRENT_SOURCE_DIR}/common/dynamic_library_loader.cc
         ${CMAKE_CURRENT_SOURCE_DIR}/common/quant_utils.cc
-        ${CMAKE_CURRENT_SOURCE_DIR}/runtime/allocator.cc
+        ${CMAKE_CURRENT_SOURCE_DIR}/runtime/inner_allocator.cc
         ${CMAKE_CURRENT_SOURCE_DIR}/runtime/runtime_api.cc
         ${CMAKE_CURRENT_SOURCE_DIR}/runtime/infer_manager.cc
         ${CMAKE_CURRENT_SOURCE_DIR}/tensor.cc

mindspore/lite/src/cxx_api/context.cc

@@ -20,7 +20,7 @@
 #include <any>
 #include "include/api/types.h"
 #include "include/api/data_type.h"
-#include "src/runtime/allocator.h"
+#include "src/runtime/inner_allocator.h"
 #include "src/common/log_adapter.h"
 
 namespace mindspore {

mindspore/lite/src/cxx_api/model/model_impl.cc

@@ -24,7 +24,7 @@
 #include "include/lite_session.h"
 #include "include/context.h"
 #include "src/lite_model.h"
-#include "src/runtime/allocator.h"
+#include "src/runtime/inner_allocator.h"
 #include "src/common/string_util.h"
 #include "src/cxx_api/graph/graph_data.h"
 #include "src/cxx_api/tensor/tensor_impl.h"

mindspore/lite/src/executor.h

@@ -18,7 +18,7 @@
 #define MINDSPORE_LITE_SRC_EXECUTOR_H_
 
 #include <vector>
-#include "src/runtime/allocator.h"
+#include "src/runtime/inner_allocator.h"
 #include "src/lite_kernel.h"
 #include "include/lite_session.h"
 

mindspore/lite/src/inner_context.h

@@ -20,7 +20,7 @@
 #include <string>
 #include "include/context.h"
 #include "src/runtime/runtime_api.h"
-#include "src/runtime/allocator.h"
+#include "src/runtime/inner_allocator.h"
 #include "thread/inter_threadpool.h"
 #ifdef ENABLE_ARM
 #include "src/cpu_info.h"

mindspore/lite/src/lite_session.cc

@@ -21,7 +21,7 @@
 #include "include/errorcode.h"
 #include "src/common/log_adapter.h"
 #include "src/scheduler.h"
-#include "src/runtime/allocator.h"
+#include "src/runtime/inner_allocator.h"
 #include "src/executor.h"
 #include "src/common/utils.h"
 #include "src/common/prim_util.h"

mindspore/lite/src/mindrt_executor.h

@@ -18,7 +18,7 @@
 #define MINDSPORE_LITE_SRC_MINDRT_EXECUTOR_H_
 #include <memory>
 #include <vector>
-#include "src/runtime/allocator.h"
+#include "src/runtime/inner_allocator.h"
 #include "src/lite_kernel.h"
 #include "src/lite_mindrt.h"
 #include "src/executor.h"

mindspore/lite/src/runtime/gpu/gpu_runtime.h

@@ -25,7 +25,7 @@
 #include <type_traits>
 #include "include/errorcode.h"
 #include "src/common/log_adapter.h"
-#include "src/runtime/allocator.h"
+#include "src/runtime/inner_allocator.h"
 #include "schema/gpu_cache_generated.h"
 
 namespace mindspore::lite::gpu {

mindspore/lite/src/runtime/gpu/opencl/opencl_allocator.h

@@ -24,7 +24,7 @@
 #include <map>
 #include <unordered_map>
 #include <unordered_set>
-#include "src/runtime/allocator.h"
+#include "src/runtime/inner_allocator.h"
 #include "CL/cl2.hpp"
 
 namespace mindspore::lite::opencl {
@@ -52,7 +52,7 @@ class OpenCLAllocator : public mindspore::Allocator {
  public:
   explicit OpenCLAllocator(OpenCLRuntime *ocl_runtime);
   ~OpenCLAllocator() override;
-  void SetContext(const AllocatorContext &ctx) override;
+  void SetContext(const AllocatorContext &ctx);
   void *Malloc(size_t size, MemType type) { return _Malloc(type, nullptr, size); }
 
   // malloc shared
@@ -66,7 +66,7 @@ class OpenCLAllocator : public mindspore::Allocator {
   int SetRefCount(void *ptr, int ref_count) override;
   int DecRefCount(void *ptr, int ref_count) override;
   int IncRefCount(void *ptr, int ref_count) override;
-  size_t total_size() override;
+  size_t total_size();
 
   void Clear();
   void *GetImage(void *host_ptr);

mindspore/lite/src/runtime/gpu/opencl/opencl_executor.h

@@ -19,7 +19,7 @@
 
 #include <vector>
 #include "src/runtime/gpu/opencl/opencl_runtime.h"
-#include "src/runtime/allocator.h"
+#include "src/runtime/inner_allocator.h"
 #include "src/runtime/kernel/opencl/opencl_kernel.h"
 #include "src/executor.h"
 #include "include/lite_session.h"

mindspore/lite/src/runtime/allocator.cc → mindspore/lite/src/runtime/inner_allocator.cc

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-#include "src/runtime/allocator.h"
+#include "src/runtime/inner_allocator.h"
 #include <utility>
 #include "src/common/log_adapter.h"
 

mindspore/lite/src/runtime/allocator.h → mindspore/lite/src/runtime/inner_allocator.h

@@ -25,43 +25,26 @@
 #include <unordered_map>
 #include <unordered_set>
 #include <atomic>
+#include "include/allocator.h"
 
 namespace mindspore {
-
 struct AllocatorContext {
   int shiftFactor;
   bool lockFlag;
 };
 
-class Allocator {
- public:
-  Allocator() : name("default") {}
-  virtual ~Allocator() = default;
-  virtual void *Malloc(size_t size) = 0;
-  virtual void Free(void *ptr) = 0;
-  virtual int RefCount(void *ptr) = 0;
-  virtual int SetRefCount(void *ptr, int ref_count) = 0;
-  virtual int DecRefCount(void *ptr, int ref_count) = 0;
-  virtual int IncRefCount(void *ptr, int ref_count) = 0;
-  virtual void SetContext(const AllocatorContext &ctx) {}
-  virtual size_t total_size() = 0;
-  static std::shared_ptr<Allocator> Create();
-  virtual void *Prepare(void *ptr) { return ptr; }
-  std::string name;
-};
-
 class DefaultAllocator : public Allocator {
  public:
   DefaultAllocator();
   ~DefaultAllocator() override;
-  void SetContext(const AllocatorContext &ctx) override;
+  void SetContext(const AllocatorContext &ctx);
   void *Malloc(size_t size) override;
   void Free(void *ptr) override;
   int RefCount(void *ptr) override;
   int SetRefCount(void *ptr, int ref_count) override;
   int DecRefCount(void *ptr, int ref_count) override;
   int IncRefCount(void *ptr, int ref_count) override;
-  size_t total_size() override { return this->total_size_; }
+  size_t total_size() { return this->total_size_; }
   void Clear();
 
  private:

mindspore/lite/src/runtime/parallel_executor.h

@@ -20,7 +20,7 @@
 #include <vector>
 #include <thread>
 #include <unordered_map>
-#include "src/runtime/allocator.h"
+#include "src/runtime/inner_allocator.h"
 #include "src/lite_kernel.h"
 #include "include/lite_session.h"
 #include "src/executor.h"

mindspore/lite/src/tensor.h

@@ -24,7 +24,7 @@
 #include <functional>
 #include <atomic>
 #include "include/ms_tensor.h"
-#include "src/runtime/allocator.h"
+#include "src/runtime/inner_allocator.h"
 
 #include "src/common/log_adapter.h"
 #include "schema/model_generated.h"
@@ -81,7 +81,7 @@ class Tensor : public mindspore::tensor::MSTensor {
 
   TypeId data_type() const override { return data_type_; }
 
-  void set_data_type(TypeId data_type) { data_type_ = data_type; }
+  void set_data_type(TypeId data_type) override { data_type_ = data_type; }
 
   std::vector<int> shape() const override { return shape_; }
 
@@ -103,9 +103,9 @@ class Tensor : public mindspore::tensor::MSTensor {
 
   size_t Size() const override;
 
-  void set_allocator(mindspore::Allocator *allocator) { allocator_ = allocator; }
+  void set_allocator(mindspore::Allocator *allocator) override { allocator_ = allocator; }
 
-  mindspore::Allocator *allocator() const { return this->allocator_; }
+  mindspore::Allocator *allocator() const override { return this->allocator_; }
 
   virtual int MallocData(const mindspore::Allocator *allocator = nullptr);
 

mindspore/lite/test/CMakeLists.txt

@@ -75,7 +75,7 @@ set(TEST_LITE_SRC
         ${CCSRC_SRC}
         ${OPS_SRC}
         ${KERNEL_OP_SRC}
-        ${LITE_DIR}/src/runtime/allocator.cc
+        ${LITE_DIR}/src/runtime/inner_allocator.cc
         ${LITE_DIR}/src/runtime/runtime_api.cc
         ${LITE_DIR}/src/runtime/parallel_executor.cc
         ${LITE_DIR}/src/runtime/infer_manager.cc

mindspore/lite/test/ut/src/runtime/kernel/arm/fp32/reduce_fp32_tests.cc

@@ -20,7 +20,7 @@
 #include "schema/inner/model_generated.h"
 #include "src/tensor.h"
 #include "mindspore/lite/src/kernel_registry.h"
-#include "mindspore/lite/src/runtime/allocator.h"
+#include "mindspore/lite/src/runtime/inner_allocator.h"
 
 using mindspore::lite::Tensor;
 using mindspore::schema::ReduceMode;

mindspore/lite/tools/converter/CMakeLists.txt

@@ -119,7 +119,7 @@ set(LITE_SRC
         ${SRC_DIR}/common/string_util.cc
         ${SRC_DIR}/common/prim_util.cc
         ${SRC_DIR}/common/tensor_util.cc
-        ${SRC_DIR}/runtime/allocator.cc
+        ${SRC_DIR}/runtime/inner_allocator.cc
         ${SRC_DIR}/runtime/runtime_api.cc
         ${SRC_DIR}/runtime/infer_manager.cc
         ${SRC_DIR}/inner_context.cc