fix: fix bugs on security

N/A
Signed-off-by: NLuxuhui <luxuhui@xiaomi.com>

include/mace/public/mace.h

@@ -416,7 +416,8 @@ class MACE_API MaceEngine {
   MaceStatus Init(const NetDef *net_def,
                   const std::vector<std::string> &input_nodes,
                   const std::vector<std::string> &output_nodes,
-                  const unsigned char *model_data);
+                  const unsigned char *model_data,
+                  const int64_t model_data_size);
 
   MaceStatus Init(const NetDef *net_def,
                   const std::vector<std::string> &input_nodes,

mace/core/arg_helper.cc

@@ -74,11 +74,8 @@ MACE_GET_OPTIONAL_ARGUMENT_FUNC(std::string, s, false)
 #define MACE_GET_REPEATED_ARGUMENT_FUNC(T, fieldname, lossless_conversion) \
   template <>                                                              \
   std::vector<T> ProtoArgHelper::GetRepeatedArgs<T>(                       \
-      const std::string &arg_name, const std::vector<T> &default_value)    \
-      const {                                                              \
-    if (arg_map_.count(arg_name) == 0) {                                   \
-      return default_value;                                                \
-    }                                                                      \
+      const std::string &arg_name) const {                                 \
+    MACE_CHECK(arg_map_.count(arg_name) > 0, arg_name, "not exist.");      \
     std::vector<T> values;                                                 \
     for (const auto &v : arg_map_.at(arg_name).fieldname()) {              \
       if (lossless_conversion) {                                           \
@@ -89,8 +86,19 @@ MACE_GET_OPTIONAL_ARGUMENT_FUNC(std::string, s, false)
       values.push_back(v);                                                 \
     }                                                                      \
     return values;                                                         \
+  }                                                                        \
+  template <>                                                              \
+  std::vector<T> ProtoArgHelper::GetRepeatedArgs<T>(                       \
+      const std::string &arg_name, const std::vector<T> &default_value)    \
+      const {                                                              \
+    if (arg_map_.count(arg_name) == 0) {                                   \
+      return default_value;                                                \
+    } else {                                                               \
+      return GetRepeatedArgs<T>(arg_name);                                 \
+    }                                                                      \
   }
 
+
 MACE_GET_REPEATED_ARGUMENT_FUNC(float, floats, false)
 MACE_GET_REPEATED_ARGUMENT_FUNC(int, ints, true)
 MACE_GET_REPEATED_ARGUMENT_FUNC(int64_t, ints, true)

mace/core/arg_helper.h

@@ -46,10 +46,14 @@ class ProtoArgHelper {
 
   template <typename T>
   T GetOptionalArg(const std::string &arg_name, const T &default_value) const;
+
   template <typename T>
   std::vector<T> GetRepeatedArgs(
       const std::string &arg_name,
-      const std::vector<T> &default_value = std::vector<T>()) const;
+      const std::vector<T> &default_value) const;
+
+  template <typename T>
+  std::vector<T> GetRepeatedArgs(const std::string &arg_name) const;
 
  private:
   std::map<std::string, Argument> arg_map_;

mace/core/memory_optimizer.cc

@@ -114,12 +114,9 @@ void MemoryOptimizer::Optimize(
     const mace::OperatorDef *op_def,
     const std::unordered_map<std::string, MemoryType> *mem_types) {
   MACE_LATENCY_LOGGER(2, "Optimize memory");
-  if (op_def->output_size() != op_def->output_shape_size()) {
-    VLOG(1) << op_def->name()
-            << ": the number of output shape "
-            << "is not equal to the number of output";
-    return;
-  }
+  MACE_CHECK(op_def->output_size() == op_def->output_shape_size(),
+             op_def->name(), "The number of output shapes is"
+                             " not equal to the number of outputs");
 
   auto device = static_cast<DeviceType>(op_def->device_type());
   DataType op_dtype = static_cast<DataType>(ProtoArgHelper::GetOptionalArg(
@@ -220,8 +217,8 @@ void MemoryOptimizer::Optimize(
       } else {
         mem_ref_count_[best_mem_id] = 1;
       }
-      tensor_mem_map_.emplace(op_def->output(i), TensorMemInfo(best_mem_id,
-          dt, data_format));
+      tensor_mem_map_.emplace(op_def->output(i),
+                              TensorMemInfo(best_mem_id, dt, data_format));
     }
   }
 
@@ -249,7 +246,7 @@ void MemoryOptimizer::Optimize(
   }
 }
 
-const std::vector<MemoryBlock>& MemoryOptimizer::mem_blocks() const {
+const std::vector<MemoryBlock> &MemoryOptimizer::mem_blocks() const {
   return mem_blocks_;
 }
 
@@ -277,8 +274,8 @@ std::string MemoryOptimizer::DebugInfo() const {
     sstream << i << " " << memory_type_to_str(mem_blocks_[i].mem_type())
             << " ";
     if (mem_blocks_[i].mem_type() == MemoryType::GPU_IMAGE) {
-      sstream << DataTypeToString(mem_blocks_[i].data_type()) << " "
-              "[" << mem_blocks_[i].x() << ", " << mem_blocks_[i].y() << "]";
+      sstream << DataTypeToString(mem_blocks_[i].data_type()) << " ["
+              << mem_blocks_[i].x() << ", " << mem_blocks_[i].y() << "]";
     } else {
       sstream << "[" << mem_blocks_[i].x() << "]";
     }

mace/core/memory_optimizer.h

@@ -94,9 +94,9 @@ class MemoryOptimizer {
                               const MemoryType mem_type);
   void UpdateTensorRef(const std::string &tensor_name);
   void UpdateTensorRef(const OperatorDef *op_def);
-  void Optimize(
+  virtual void Optimize(
       const OperatorDef *op_def,
-      const std::unordered_map<std::string, MemoryType> *mem_types = nullptr);
+      const std::unordered_map<std::string, MemoryType> *mem_types);
 
   const std::vector<MemoryBlock> &mem_blocks() const;
 

mace/core/net.cc

@@ -54,6 +54,9 @@ SerialNet::SerialNet(const OpRegistry *op_registry,
 #ifdef MACE_ENABLE_OPENCL
   // used for memory optimization
   std::unordered_map<std::string, MemoryType> output_mem_map;
+  std::unordered_map<std::string, MemoryType> *mem_types = &output_mem_map;
+#else
+  std::unordered_map<std::string, MemoryType> *mem_types = nullptr;
 #endif  // MACE_ENABLE_OPENCL
 
   OpConstructContext construct_context(ws_);
@@ -100,11 +103,7 @@ SerialNet::SerialNet(const OpRegistry *op_registry,
   for (auto &op : operators_) {
     VLOG(2) << "Operator " << op->debug_def().name() << "<" << op->device_type()
             << ", " << op->debug_def().type() << ">";
-#ifdef MACE_ENABLE_OPENCL
-    mem_optimizer->Optimize(op->operator_def().get(), &output_mem_map);
-#else
-    mem_optimizer->Optimize(op->operator_def().get());
-#endif  // MACE_ENABLE_OPENCL
+    mem_optimizer->Optimize(op->operator_def().get(), mem_types);
   }
   VLOG(1) << mem_optimizer->DebugInfo();
 }
@@ -128,8 +127,7 @@ MaceStatus SerialNet::Init() {
 
 MaceStatus SerialNet::Run(RunMetadata *run_metadata) {
   const char *profiling = getenv("MACE_OPENCL_PROFILING");
-  bool
-  enable_opencl_profiling =
+  bool enable_opencl_profiling =
       profiling != nullptr && strlen(profiling) == 1 && profiling[0] == '1';
 
   MACE_MEMORY_LOGGING_GUARD();

mace/core/ops/operator.h

@@ -48,14 +48,22 @@ class Operation {
     return ProtoArgHelper::GetOptionalArg<OperatorDef, T>(
         *operator_def_, name, default_value);
   }
+
   template<typename T>
   std::vector<T> GetRepeatedArgs(
-      const std::string &name, const std::vector<T> &default_value = {}) const {
+      const std::string &name, const std::vector<T> &default_value) const {
     MACE_CHECK(operator_def_, "operator_def was null!");
     return ProtoArgHelper::GetRepeatedArgs<OperatorDef, T>(
         *operator_def_, name, default_value);
   }
 
+  template<typename T>
+  std::vector<T> GetRepeatedArgs(const std::string &name) const {
+    MACE_CHECK(operator_def_, "operator_def was null!");
+    return ProtoArgHelper::GetRepeatedArgs<OperatorDef, T>(
+        *operator_def_, name);
+  }
+
   DeviceType device_type() const {
     return static_cast<DeviceType>(operator_def_->device_type());
   }
@@ -65,7 +73,10 @@ class Operation {
     return inputs_[idx];
   }
 
-  Tensor *Output(int idx) { return outputs_[idx]; }
+  Tensor *Output(unsigned int idx) {
+    MACE_CHECK(idx < outputs_.size());
+    return outputs_[idx];
+  }
 
   int InputSize() { return inputs_.size(); }
   int OutputSize() { return outputs_.size(); }

mace/core/runtime/apu/apu_wrapper.cc

@@ -74,7 +74,8 @@ apu_eltwise_mode ApuWrapper::MapToApuEltwiseMode(int mace_mode) {
 }
 
 bool ApuWrapper::Init(const NetDef &net_def,
-                      unsigned const char *model_data) {
+                      unsigned const char *model_data,
+                      const index_t model_data_size) {
   frontend = new ApuFrontend();
 
   // parse model argument
@@ -100,11 +101,15 @@ bool ApuWrapper::Init(const NetDef &net_def,
     tensor.dim_size = const_tensor.dims_size();
     MACE_CHECK(tensor.dim_size <= APU_TENSOR_MAX_DIMS,
                "tensor dimension size not supported");
-    for (auto i = 0 ; i < tensor.dim_size ; i++) {
+    for (auto i = 0; i < tensor.dim_size; i++) {
       tensor.dims[i] = const_tensor.dims(i);
     }
+    const auto tensor_end = const_tensor.offset() +
+        const_tensor->data_size() * GetEnumTypeSize(const_tensor.data_type());
+    MACE_CHECK(tensor_end <= model_data_size, "tensor_end (", tensor_end,
+               ") should <= ", model_data_size);
     tensor.data_buf =
-        const_cast<unsigned char*>(model_data + const_tensor.offset());
+        const_cast<unsigned char *>(model_data + const_tensor.offset());
     const_tensors.push_back(tensor);
   }
 
@@ -121,10 +126,10 @@ bool ApuWrapper::Init(const NetDef &net_def,
     tensor.dim_size = input_info.dims_size();
     MACE_CHECK(tensor.dim_size <= APU_TENSOR_MAX_DIMS,
                "tensor dimension size not supported");
-    tensor_info info;
+    ApuTensorInfo info;
     info.name = input_info.name();
     info.size = 1;
-    for (auto i = 0 ; i < tensor.dim_size ; i++) {
+    for (auto i = 0; i < tensor.dim_size; i++) {
       tensor.dims[i] = input_info.dims(i);
       info.size *= input_info.dims(i);
       info.shape.push_back(input_info.dims(i));
@@ -140,13 +145,13 @@ bool ApuWrapper::Init(const NetDef &net_def,
 
   // output tensors
   std::vector<int> output_tensor_ids;
-  std::vector<void*> output_buffers;
+  std::vector<void *> output_buffers;
   for (auto output_info : net_def.output_info()) {
     output_tensor_ids.push_back(output_info.node_id());
-    tensor_info info;
+    ApuTensorInfo info;
     info.name = output_info.name();
     info.size = 1;
-    for (auto i = 0 ; i < output_info.dims().size() ; i++) {
+    for (auto i = 0; i < output_info.dims().size(); i++) {
       info.size *= output_info.dims(i);
       info.shape.push_back(output_info.dims(i));
     }
@@ -170,7 +175,7 @@ bool ApuWrapper::Init(const NetDef &net_def,
     strncpy(op.type, op_def.type().c_str(), APU_OP_TYPE_MAX_SIZE);
     op.input_size = op_def.node_input_size();
     std::vector<int> input_ids;
-    for (auto i = 0 ; i < op.input_size ; i++) {
+    for (auto i = 0; i < op.input_size; i++) {
       input_ids.push_back(op_def.node_input(i).node_id());
     }
     cached_op_inputs.push_back(input_ids);
@@ -188,7 +193,7 @@ bool ApuWrapper::Init(const NetDef &net_def,
     op.output.dim_size = op_def.output_shape(0).dims_size();
     MACE_CHECK(op.output.dim_size <= APU_TENSOR_MAX_DIMS,
                "tensor dimension size not supported");
-    for (auto i = 0 ; i < op.output.dim_size ; i++) {
+    for (auto i = 0; i < op.output.dim_size; i++) {
       op.output.dims[i] = op_def.output_shape(0).dims(i);
     }
     op.output.data_buf = nullptr;
@@ -243,7 +248,7 @@ bool ApuWrapper::Run(const std::map<std::string, Tensor *> &input_tensors,
 
   // prepare input
   for (int i = 0 ; i < static_cast<int>(input_tensors.size()) ; i++) {
-    Tensor* tensor = input_tensors.at(input_infos[i].name);
+    Tensor *tensor = input_tensors.at(input_infos[i].name);
 
     // check size
     int size = input_infos[i].size;
@@ -251,7 +256,7 @@ bool ApuWrapper::Run(const std::map<std::string, Tensor *> &input_tensors,
 
     // quantize
     quantize_util_.QuantizeWithScaleAndZeropoint(
-        (const float*)tensor->raw_data(),
+        tensor->data<float>(),
         size,
         input_infos[i].scale,
         input_infos[i].zero_point,
@@ -263,8 +268,8 @@ bool ApuWrapper::Run(const std::map<std::string, Tensor *> &input_tensors,
   MACE_CHECK(ret == true, "neuron run model failed");
 
   // process output
-  for (int i = 0 ; i < static_cast<int>(output_tensors->size()) ; i++) {
-    Tensor* tensor = output_tensors->at(output_infos[i].name);
+  for (int i = 0; i < static_cast<int>(output_tensors->size()); i++) {
+    Tensor *tensor = output_tensors->at(output_infos[i].name);
 
     // prepare out buffer
     tensor->SetDtype(DT_FLOAT);
@@ -278,7 +283,7 @@ bool ApuWrapper::Run(const std::map<std::string, Tensor *> &input_tensors,
         size,
         output_infos[i].scale,
         output_infos[i].zero_point,
-        reinterpret_cast<float*>(tensor->raw_mutable_data()));
+        tensor->mutable_data<float>());
   }
 
   return true;

mace/core/runtime/apu/apu_wrapper.h

@@ -21,27 +21,28 @@
 #include <memory>
 
 #include "mace/proto/mace.pb.h"
-#include "mace/core/tensor.h"
 #include "mace/core/device.h"
 #include "mace/core/quantize.h"
+#include "mace/core/tensor.h"
 
 #include "third_party/apu/ApuFrontend.h"
 
 namespace mace {
 
 class ApuWrapper {
-struct tensor_info {
+  struct ApuTensorInfo {
     std::string name;
     std::shared_ptr<uint8_t> buf;
     std::vector<index_t> shape;
     int size;
     float scale;
     int zero_point;
-};
+  };
 
  public:
   explicit ApuWrapper(Device *device);
-    bool Init(const NetDef& net_def, unsigned const char *model_data);
+  bool Init(const NetDef &net_def, unsigned const char *model_data,
+            const index_t model_data_size);
   bool Run(const std::map<std::string, Tensor *> &input_tensors,
            std::map<std::string, Tensor *> *output_tensors);
   bool Uninit();
@@ -52,9 +53,9 @@ struct tensor_info {
   apu_eltwise_mode MapToApuEltwiseMode(int mace_mode);
 
  private:
-    ApuFrontend* frontend;
-    std::vector<tensor_info> input_infos;
-    std::vector<tensor_info> output_infos;
+  ApuFrontend *frontend;
+  std::vector<ApuTensorInfo> input_infos;
+  std::vector<ApuTensorInfo> output_infos;
   QuantizeUtil<float, uint8_t> quantize_util_;
 };
 

mace/core/runtime/hexagon/hexagon_control_wrapper.h

@@ -46,7 +46,8 @@ class HexagonControlWrapper {
   virtual bool Init() = 0;
   virtual bool Finalize() = 0;
   virtual bool SetupGraph(const NetDef &net_def,
-                          const unsigned char *model_data) = 0;
+                          const unsigned char *model_data,
+                          const index_t model_data_size) = 0;
   virtual bool ExecuteGraph(const Tensor &input_tensor,
                             Tensor *output_tensor) = 0;
   virtual bool ExecuteGraphNew(

mace/core/runtime/hexagon/hexagon_dsp_wrapper.cc

@@ -165,7 +165,8 @@ bool HexagonDSPWrapper::Finalize() {
 }
 
 bool HexagonDSPWrapper::SetupGraph(const NetDef &net_def,
-                                   unsigned const char *model_data) {
+                                   unsigned const char *model_data,
+                                   const index_t model_data_size) {
   LOG(INFO) << "Hexagon setup graph";
 
   int64_t t0 = NowMicros();
@@ -188,6 +189,9 @@ bool HexagonDSPWrapper::SetupGraph(const NetDef &net_def,
       data = const_cast<unsigned char *>(model_data + const_tensor.offset());
       data_len =
           const_tensor.data_size() * GetEnumTypeSize(const_tensor.data_type());
+      MACE_CHECK(const_tensor.offset() + data_len <= model_data_size,
+                 "tensor end (", const_tensor.offset() + data_len,
+                 ") should <= ", model_data_size);
     }
     MACE_CHECK(
         hexagon_nn_append_const_node(nn_id_,

mace/core/runtime/hexagon/hexagon_dsp_wrapper.h

@@ -35,7 +35,8 @@ class HexagonDSPWrapper : public HexagonControlWrapper {
   bool Init() override;
   bool Finalize() override;
   bool SetupGraph(const NetDef &net_def,
-                  const unsigned char *model_data) override;
+                  const unsigned char *model_data,
+                  const index_t model_data_size) override;
   bool ExecuteGraph(const Tensor &input_tensor,
                     Tensor *output_tensor) override;
   bool ExecuteGraphNew(const std::map<std::string, Tensor*> &input_tensors,

mace/core/runtime/hexagon/hexagon_hta_wrapper.cc

@@ -138,7 +138,8 @@ bool HexagonHTAWrapper::Finalize() {
 }
 
 bool HexagonHTAWrapper::SetupGraph(const NetDef &net_def,
-                                   unsigned const char *model_data) {
+                                   unsigned const char *model_data,
+                                   const index_t model_data_size) {
   LOG(INFO) << "Hexagon setup graph";
 
   int64_t t0 = NowMicros();
@@ -159,6 +160,9 @@ bool HexagonHTAWrapper::SetupGraph(const NetDef &net_def,
           const_cast<unsigned char *>(model_data + const_tensor.offset());
       const_node_data_len = const_tensor.data_size() *
           GetEnumTypeSize(const_tensor.data_type());
+      MACE_CHECK(const_tensor.offset() + const_node_data_len <= model_data_size,
+                 "tensor end (", const_tensor.offset() + const_node_data_len,
+                 ") should <= ", model_data_size);
     }
 
     hexagon_hta_nn_append_const_node(nn_id_,

mace/core/runtime/hexagon/hexagon_hta_wrapper.h

@@ -59,7 +59,8 @@ class HexagonHTAWrapper : public HexagonControlWrapper {
   bool Init() override;
   bool Finalize() override;
   bool SetupGraph(const NetDef &net_def,
-                  const unsigned char *model_data) override;
+                  const unsigned char *model_data,
+                  const index_t model_data_size) override;
   bool ExecuteGraph(const Tensor &input_tensor,
                     Tensor *output_tensor) override;
   bool ExecuteGraphNew(const std::map<std::string, Tensor*> &input_tensors,

mace/core/tensor.h

@@ -276,6 +276,7 @@ class Tensor {
 
   inline void *raw_mutable_data() {
     MACE_CHECK_NOTNULL(buffer_);
+
     return buffer_->raw_mutable_data();
   }
 

mace/core/workspace.cc

@@ -64,6 +64,18 @@ void DequantizeTensor(Device *device,
                            dequantized_data);
 }
 
+
+index_t GetModelValidSize(const NetDef &net_def) {
+  index_t valid_data_size = 0;
+  for (auto &const_tensor : net_def.tensors()) {
+    valid_data_size = std::max<index_t>(
+        valid_data_size, const_tensor.offset() +
+            const_tensor.data_size()
+                * GetEnumTypeSize(const_tensor.data_type()));
+  }
+  return valid_data_size;
+}
+
 }  // namespace
 
 Workspace::Workspace(const OpDelegatorRegistry *registry) :
@@ -105,23 +117,17 @@ std::vector<std::string> Workspace::Tensors() const {
   return names;
 }
 
-MaceStatus Workspace::LoadModelTensor(const NetDef &net_def,
-                                      Device *device,
-                                      const unsigned char *model_data) {
+MaceStatus Workspace::LoadModelTensor(const NetDef &net_def, Device *device,
+                                      const unsigned char *model_data,
+                                      const index_t model_data_size) {
   MACE_LATENCY_LOGGER(1, "Load model tensors");
-  index_t model_data_size = 0;
-  for (auto &const_tensor : net_def.tensors()) {
-    model_data_size = std::max(
-        model_data_size,
-        static_cast<index_t>(const_tensor.offset() +
-            const_tensor.data_size() *
-                GetEnumTypeSize(const_tensor.data_type())));
-  }
-  VLOG(3) << "Model data size: " << model_data_size;
+  index_t valid_data_size = GetModelValidSize(net_def);
+  VLOG(3) << "Model valid data size: " << valid_data_size;
+  MACE_CHECK(valid_data_size <= model_data_size,
+             valid_data_size, "should be smaller than", model_data_size);
 
   const DeviceType device_type = device->device_type();
-
-  if (model_data_size > 0) {
+  if (valid_data_size > 0) {
     bool is_quantize_model = IsQuantizedModel(net_def);
     diffused_buffer_ =
         (device_type == DeviceType::CPU && HasHalfTensor(net_def)) ||
@@ -129,7 +135,7 @@ MaceStatus Workspace::LoadModelTensor(const NetDef &net_def,
 #ifdef MACE_ENABLE_OPENCL
     diffused_buffer_ = diffused_buffer_ || (device_type == DeviceType::GPU &&
         device->gpu_runtime()->opencl_runtime()->GetDeviceMaxMemAllocSize() <=
-            static_cast<uint64_t>(model_data_size));
+            static_cast<uint64_t>(valid_data_size));
 #endif
     if (diffused_buffer_) {
       for (auto &const_tensor : net_def.tensors()) {
@@ -162,15 +168,10 @@ MaceStatus Workspace::LoadModelTensor(const NetDef &net_def,
 
         MACE_CHECK(tensor->size() == const_tensor.data_size(),
                    "Tensor's data_size not equal with the shape");
-        MACE_CHECK(static_cast<index_t>(const_tensor.offset() +
-            tensor->size() * GetEnumTypeSize(const_tensor.data_type())) <=
-            model_data_size,
-                   "buffer offset + length (",
-                   const_tensor.offset(),
-                   " + ",
-                   tensor->size() * GetEnumTypeSize(const_tensor.data_type()),
-                   ") should <= ",
-                   model_data_size);
+        const index_t tensor_end = const_tensor.offset() +
+            tensor->size() * GetEnumTypeSize(const_tensor.data_type());
+        MACE_CHECK(tensor_end <= model_data_size, "tensor_end (", tensor_end,
+                   ") should <= ", model_data_size);
 
         if (device_type == DeviceType::CPU &&
             const_tensor.data_type() == DataType::DT_HALF) {
@@ -207,14 +208,14 @@ MaceStatus Workspace::LoadModelTensor(const NetDef &net_def,
         tensor_buffer_ = std::unique_ptr<Buffer>(
             new Buffer(device->allocator(),
                        const_cast<unsigned char *>(model_data),
-                       model_data_size));
+                       valid_data_size));
       } else {
         tensor_buffer_ = std::unique_ptr<Buffer>(
             new Buffer(device->allocator()));
-        MACE_RETURN_IF_ERROR(tensor_buffer_->Allocate(model_data_size));
+        MACE_RETURN_IF_ERROR(tensor_buffer_->Allocate(valid_data_size));
         tensor_buffer_->Map(nullptr);
         tensor_buffer_->Copy(const_cast<unsigned char *>(model_data),
-                             0, model_data_size);
+                             0, valid_data_size);
         tensor_buffer_->UnMap();
       }
       for (auto &const_tensor : net_def.tensors()) {

mace/core/workspace.h

@@ -56,9 +56,9 @@ class Workspace {
 
   std::vector<std::string> Tensors() const;
 
-  MaceStatus LoadModelTensor(const NetDef &net_def,
-                             Device *device,
-                             const unsigned char *model_data);
+  MaceStatus LoadModelTensor(const NetDef &net_def, Device *device,
+                             const unsigned char *model_data,
+                             const index_t model_data_size);
 
   MaceStatus PreallocateOutputTensor(const NetDef &net_def,
                                      const MemoryOptimizer *mem_optimizer,

mace/libmace/capability.cc

@@ -350,7 +350,8 @@ MaceStatus BMNet::Run(DeviceType device,
   }
   MaceEngine engine(config);
 
-  status = engine.Init(&net_, input_names_, output_names_, weight_.data());
+  status = engine.Init(&net_, input_names_, output_names_,
+                       weight_.data(), weight_.size());
   if (status != MaceStatus::MACE_SUCCESS) {
     return status;
   }

mace/libmace/mace.cc

@@ -433,7 +433,8 @@ class MaceEngine::Impl {
   MaceStatus Init(const NetDef *net_def,
                   const std::vector<std::string> &input_nodes,
                   const std::vector<std::string> &output_nodes,
-                  const unsigned char *model_data);
+                  const unsigned char *model_data,
+                  const int64_t model_data_size);
 
   MaceStatus Init(const NetDef *net_def,
                   const std::vector<std::string> &input_nodes,
@@ -557,7 +558,7 @@ MaceStatus MaceEngine::Impl::Init(
     const NetDef *net_def,
     const std::vector<std::string> &input_nodes,
     const std::vector<std::string> &output_nodes,
-    const unsigned char *model_data) {
+    const unsigned char *model_data, const int64_t model_data_size) {
   LOG(INFO) << "Initializing MaceEngine";
   // Check avalibility
 #ifdef MACE_ENABLE_OPENCL
@@ -634,8 +635,8 @@ MaceStatus MaceEngine::Impl::Init(
     MACE_CHECK(hexagon_controller_->Init(), "hexagon init error");
     hexagon_controller_->SetDebugLevel(
         static_cast<int>(mace::port::MinVLogLevelFromEnv()));
-    MACE_CHECK(hexagon_controller_->SetupGraph(*net_def, model_data),
-               "hexagon setup graph error");
+    MACE_CHECK(hexagon_controller_->SetupGraph(
+        *net_def, model_data, model_data_size), "hexagon setup graph error");
     if (VLOG_IS_ON(2)) {
       hexagon_controller_->PrintGraph();
     }
@@ -644,12 +645,12 @@ MaceStatus MaceEngine::Impl::Init(
 #ifdef MACE_ENABLE_APU
   if (device_type_ == APU) {
     apu_controller_.reset(new ApuWrapper(device_.get()));
-    MACE_CHECK(apu_controller_->Init(*net_def, model_data), "apu init error");
+    MACE_CHECK(apu_controller_->Init(
+        *net_def, model_data, model_data_size), "apu init error");
   } else {
 #endif
-    MACE_RETURN_IF_ERROR(ws_->LoadModelTensor(*net_def,
-                                              device_.get(),
-                                              model_data));
+    MACE_RETURN_IF_ERROR(ws_->LoadModelTensor(
+        *net_def, device_.get(), model_data, model_data_size));
 
     NetDef adapted_net_def;
     NetDefAdapter net_def_adapter(op_registry_.get(), ws_.get());
@@ -696,7 +697,8 @@ MaceStatus MaceEngine::Impl::Init(
 
   MACE_RETURN_IF_ERROR(Init(
       net_def, input_nodes, output_nodes,
-      reinterpret_cast<const unsigned char *>(model_data_->data())));
+      reinterpret_cast<const unsigned char *>(model_data_->data()),
+      model_data_->length()));
 
   if (device_type_ == DeviceType::GPU || device_type_ == DeviceType::HEXAGON ||
       device_type_ == DeviceType::HTA ||
@@ -945,6 +947,7 @@ MaceStatus MaceEngine::Impl::Run(
                  << MakeString(MapKeys(output_info_map_));
     }
     Tensor *output_tensor = ws_->GetTensor(output.first);
+    MACE_CHECK(output_tensor != nullptr);
     output_tensors[output.first] = output_tensor;
   }
 #if defined(MACE_ENABLE_HEXAGON) || defined(MACE_ENABLE_HTA)
@@ -988,8 +991,10 @@ MaceEngine::~MaceEngine() = default;
 MaceStatus MaceEngine::Init(const NetDef *net_def,
                             const std::vector<std::string> &input_nodes,
                             const std::vector<std::string> &output_nodes,
-                            const unsigned char *model_data) {
-  return impl_->Init(net_def, input_nodes, output_nodes, model_data);
+                            const unsigned char *model_data,
+                            int64_t model_data_size) {
+  return impl_->Init(net_def, input_nodes, output_nodes,
+                     model_data, model_data_size);
 }
 
 MaceStatus MaceEngine::Init(const NetDef *net_def,
@@ -1032,7 +1037,8 @@ MaceStatus CreateMaceEngineFromProto(
 
   engine->reset(new mace::MaceEngine(config));
   MaceStatus status = (*engine)->Init(
-      net_def.get(), input_nodes, output_nodes, model_weights_data);
+      net_def.get(), input_nodes, output_nodes,
+      model_weights_data, model_weights_data_size);
 
   return status;
 }

mace/ops/common/conv_pool_2d_util.cc

@@ -31,7 +31,9 @@ void CalcPaddingAndOutputSize(const index_t *input_shape,
                               index_t *output_shape,
                               int *padding_size) {
   MACE_CHECK(dilations[0] > 0 && dilations[1] > 0,
-             "Invalid dilations, must >= 1");
+             "Invalid dilations, must > 0");
+  MACE_CHECK(strides[0] > 0 && strides[1] > 0,
+             "Invalid strides, must > 0");
   MACE_CHECK((dilations[0] == 1 || strides[0] == 1) &&
       (dilations[1] == 1 || strides[1] == 1),
              "If dilations > 1, strides should be 1");
@@ -145,7 +147,9 @@ void CalcOutputSize(const index_t *input_shape,
                     const RoundType round_type,
                     index_t *output_shape) {
   MACE_CHECK(dilations[0] > 0 && dilations[1] > 0,
-             "Invalid dilations, must >= 1");
+             "Invalid dilations, must > 0");
+  MACE_CHECK(strides[0] > 0 && strides[1] > 0,
+             "Invalid strides, must > 0");
   MACE_CHECK((dilations[0] == 1 || strides[0] == 1) &&
       (dilations[1] == 1 || strides[1] == 1),
              "If dilations > 1, strides should be 1");

mace/ops/pooling.cc

@@ -94,6 +94,7 @@ class PoolingOp<DeviceType::CPU, T> : public PoolingOpBase {
     Tensor::MappingGuard input_guard(input_tensor);
     Tensor::MappingGuard output_guard(output_tensor);
     const T *input = input_tensor->data<T>();
+    MACE_CHECK(output_tensor->dtype() == DataTypeToEnum<T>::value);
     T *output = output_tensor->mutable_data<T>();
     const index_t *input_shape = input_tensor->shape().data();
     int pad_hw[2] = {paddings[0] / 2, paddings[1] / 2};
@@ -299,6 +300,7 @@ class PoolingOp<DeviceType::CPU, uint8_t> : public PoolingOpBase {
     Tensor::MappingGuard input_guard(input_tensor);
     Tensor::MappingGuard output_guard(output_tensor);
     const uint8_t *input = input_tensor->data<uint8_t>();
+    MACE_CHECK(output_tensor->dtype() == DataTypeToEnum<uint8_t>::value);
     uint8_t *output = output_tensor->mutable_data<uint8_t>();
     int pad_hw[2] = {paddings[0] / 2, paddings[1] / 2};
 

mace/ops/ref/conv_2d.cc

@@ -39,9 +39,9 @@ MaceStatus Conv2d<T>::Compute(const OpContext *context,
                               const Tensor *filter,
                               Tensor *output) {
   MACE_UNUSED(context);
-
   const std::vector<index_t> in_shape = input->shape();
   const std::vector<index_t> filter_shape = filter->shape();
+  MACE_CHECK(in_shape[1] == filter_shape[1]);
   std::vector<index_t> out_shape(4);
 
   std::vector<int> paddings(2);

mace/python/tools/mace_engine_factory.h.jinja2

@@ -77,6 +77,7 @@ __attribute__((deprecated)) MaceStatus CreateMaceEngineFromCode(
 {% if embed_model_data %}
   (void)model_data_file;
   const unsigned char * model_data;
+  const int64_t model_size;
 {% endif %}
   MaceStatus status = MaceStatus::MACE_SUCCESS;
   switch (model_name_map[model_name]) {
@@ -86,8 +87,9 @@ __attribute__((deprecated)) MaceStatus CreateMaceEngineFromCode(
       engine->reset(new mace::MaceEngine(config));
 {% if embed_model_data %}
       model_data = mace::{{model_tags[i]}}::LoadModelData();
+      model_size = mace::{{model_tags[i]}}::GetModelSize();
       status = (*engine)->Init(net_def.get(), input_nodes, output_nodes,
-                               model_data);
+                               model_data, model_size);
 {% else %}
       status = (*engine)->Init(net_def.get(), input_nodes, output_nodes,
                                model_data_file);
@@ -116,10 +118,9 @@ MaceStatus CreateMaceEngineFromCode(
   std::shared_ptr<NetDef> net_def;
 {% if embed_model_data %}
   const unsigned char * model_data;
+  const int64_t model_size;
   (void)model_weights_data;
 {% endif %}
-  // TODO(yejianwu) Add buffer range checking
-  (void)model_weights_data_size;
 
   MaceStatus status = MaceStatus::MACE_SUCCESS;
   switch (model_name_map[model_name]) {
@@ -129,11 +130,12 @@ MaceStatus CreateMaceEngineFromCode(
       engine->reset(new mace::MaceEngine(config));
 {% if embed_model_data %}
       model_data = mace::{{model_tags[i]}}::LoadModelData();
+      model_size = mace::{{model_tags[i]}}::GetModelSize();
       status = (*engine)->Init(net_def.get(), input_nodes, output_nodes,
-                               model_data);
+                               model_data, model_size);
 {% else %}
       status = (*engine)->Init(net_def.get(), input_nodes, output_nodes,
-                               model_weights_data);
+                               model_weights_data, model_weights_data_size);
 {% endif %}
       break;
 {% endfor %}

test/ccunit/mace/libmace/mace_api_exception_test.cc

@@ -36,7 +36,7 @@ TEST(MaceAPIExceptionTest, WrongInputTest) {
   }
 
   MaceEngine engine(config);
-  ASSERT_DEATH(engine.Init(net_def.get(), {"input"}, output_names, nullptr),
+  ASSERT_DEATH(engine.Init(net_def.get(), {"input"}, output_names, nullptr, 0),
                "");
 }
 

test/ccunit/mace/libmace/mace_api_mt_test.cc

@@ -68,8 +68,10 @@ void MaceRunFunc(const int in_out_size) {
   config.SetGPUContext(mace::ops::test::OpTestContext::Get()->gpu_context());
 
   MaceEngine engine(config);
-  MaceStatus status = engine.Init(net_def.get(), input_names, output_names,
-      reinterpret_cast<unsigned char *>(data.data()));
+  MaceStatus status = engine.Init(
+      net_def.get(), input_names, output_names,
+      reinterpret_cast<unsigned char *>(data.data()),
+      data.size() * sizeof(half));
   EXPECT_EQ(status, MaceStatus::MACE_SUCCESS);
 
   std::map<std::string, mace::MaceTensor> inputs;

test/ccunit/mace/libmace/mace_api_test.cc

@@ -65,7 +65,7 @@ void MaceRun(const int in_out_size,
 
   MaceEngine engine(config);
   MaceStatus status = engine.Init(net_def.get(), input_names, output_names,
-      reinterpret_cast<unsigned char *>(data.data()));
+      reinterpret_cast<unsigned char *>(data.data()), data.size() * sizeof(T));
   EXPECT_EQ(status, MaceStatus::MACE_SUCCESS);
 
   std::map<std::string, mace::MaceTensor> inputs;

test/ccutils/mace/ops/memory_optimizer_for_test.cc

+// Copyright 2018 The MACE Authors. All Rights Reserved.
+//
+// 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 to 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.
+
+#include "mace/ops/memory_optimizer_for_test.h"
+
+namespace mace {
+
+void MemoryOptimizerForTest::Optimize(
+    const OperatorDef *op_def,
+    const std::unordered_map<std::string, MemoryType> *mem_types) {
+  if (op_def->output_size() == op_def->output_shape_size()) {
+    MemoryOptimizer::Optimize(op_def, mem_types);
+  }
+}
+
+}  // namespace mace

test/ccutils/mace/ops/memory_optimizer_for_test.h

+// Copyright 2020 The MACE Authors. All Rights Reserved.
+//
+// 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 to 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 MACE_OPS_MEMORY_OPTIMIZER_FOR_TEST_H_
+#define MACE_OPS_MEMORY_OPTIMIZER_FOR_TEST_H_
+
+#include <string>
+#include <unordered_map>
+
+#include "mace/core/memory_optimizer.h"
+
+namespace mace {
+
+class MemoryOptimizerForTest : public MemoryOptimizer {
+ public:
+  MemoryOptimizerForTest() {}
+  ~MemoryOptimizerForTest() {}
+
+  void Optimize(
+      const OperatorDef *op_def,
+      const std::unordered_map<std::string, MemoryType> *mem_types) override;
+};
+
+}  // namespace mace
+
+#endif  // MACE_OPS_MEMORY_OPTIMIZER_FOR_TEST_H_

test/ccutils/mace/ops/ops_test_util.cc

@@ -13,9 +13,11 @@
 // limitations under the License.
 
 #include "mace/ops/ops_test_util.h"
+
 #include "mace/core/memory_optimizer.h"
-#include "mace/utils/memory.h"
 #include "mace/core/net_def_adapter.h"
+#include "mace/ops/memory_optimizer_for_test.h"
+#include "mace/utils/memory.h"
 
 namespace mace {
 namespace ops {
@@ -228,7 +230,7 @@ bool OpsTestNet::Setup(mace::DeviceType device) {
                               OpTestContext::Get()->GetDevice(device),
                               &adapted_net_def);
 
-  MemoryOptimizer mem_optimizer;
+  MemoryOptimizerForTest mem_optimizer;
   net_ = make_unique<SerialNet>(
       op_registry_.get(),
       &adapted_net_def,
@@ -284,7 +286,7 @@ MaceStatus OpsTestNet::RunNet(const mace::NetDef &net_def,
   net_def_adapter.AdaptNetDef(&net_def,
                               OpTestContext::Get()->GetDevice(device),
                               &adapted_net_def);
-  MemoryOptimizer mem_optimizer;
+  MemoryOptimizerForTest mem_optimizer;
   net_ = make_unique<SerialNet>(
       op_registry_.get(),
       &adapted_net_def,

tools/python/template/mace_engine_factory.h.jinja2

@@ -77,6 +77,7 @@ __attribute__((deprecated)) MaceStatus CreateMaceEngineFromCode(
 {% if embed_model_data %}
   (void)model_data_file;
   const unsigned char * model_data;
+  const int64_t model_size;
 {% endif %}
   MaceStatus status = MaceStatus::MACE_SUCCESS;
   switch (model_name_map[model_name]) {
@@ -86,8 +87,9 @@ __attribute__((deprecated)) MaceStatus CreateMaceEngineFromCode(
       engine->reset(new mace::MaceEngine(config));
 {% if embed_model_data %}
       model_data = mace::{{model_tags[i]}}::LoadModelData();
+      model_size = mace::{{model_tags[i]}}::GetModelSize();
       status = (*engine)->Init(net_def.get(), input_nodes, output_nodes,
-                               model_data);
+                               model_data, model_size);
 {% else %}
       status = (*engine)->Init(net_def.get(), input_nodes, output_nodes,
                                model_data_file);
@@ -116,10 +118,9 @@ MACE_API MaceStatus CreateMaceEngineFromCode(
   std::shared_ptr<NetDef> net_def;
 {% if embed_model_data %}
   const unsigned char * model_data;
+  const int64_t model_size;
   (void)model_weights_data;
 {% endif %}
-  // TODO(yejianwu) Add buffer range checking
-  (void)model_weights_data_size;
 
   MaceStatus status = MaceStatus::MACE_SUCCESS;
   switch (model_name_map[model_name]) {
@@ -129,11 +130,12 @@ MACE_API MaceStatus CreateMaceEngineFromCode(
       engine->reset(new mace::MaceEngine(config));
 {% if embed_model_data %}
       model_data = mace::{{model_tags[i]}}::LoadModelData();
+      model_size = mace::{{model_tags[i]}}::GetModelSize();
       status = (*engine)->Init(net_def.get(), input_nodes, output_nodes,
-                               model_data);
+                               model_data, model_size);
 {% else %}
       status = (*engine)->Init(net_def.get(), input_nodes, output_nodes,
-                               model_weights_data);
+                               model_weights_data, model_weights_data_size);
 {% endif %}
       break;
 {% endfor %}

tools/python/template/tensor_data.jinja2

@@ -27,6 +27,10 @@ MACE_API const unsigned char *LoadModelData() {
   return model_data;
 }
 
+MACE_API int64_t GetModelSize() {
+  return {{ model_data_size }};
+}
+
 }  // namespace {{tag}}
 }  // namespace mace