From 71de31e05b6e7223b6248fb038b33ae287e4cdab Mon Sep 17 00:00:00 2001 From: luxuhui Date: Tue, 19 May 2020 20:01:52 +0800 Subject: [PATCH] fix: fix bugs on security N/A Signed-off-by: Luxuhui --- include/mace/public/mace.h | 3 +- mace/core/arg_helper.cc | 18 ++++-- mace/core/arg_helper.h | 6 +- mace/core/memory_optimizer.cc | 19 +++---- mace/core/memory_optimizer.h | 4 +- mace/core/net.cc | 12 ++-- mace/core/ops/operator.h | 15 ++++- mace/core/runtime/apu/apu_wrapper.cc | 57 ++++++++++--------- mace/core/runtime/apu/apu_wrapper.h | 31 +++++----- .../runtime/hexagon/hexagon_control_wrapper.h | 3 +- .../runtime/hexagon/hexagon_dsp_wrapper.cc | 6 +- .../runtime/hexagon/hexagon_dsp_wrapper.h | 3 +- .../runtime/hexagon/hexagon_hta_wrapper.cc | 6 +- .../runtime/hexagon/hexagon_hta_wrapper.h | 3 +- mace/core/tensor.h | 1 + mace/core/workspace.cc | 55 +++++++++--------- mace/core/workspace.h | 6 +- mace/libmace/capability.cc | 3 +- mace/libmace/mace.cc | 30 ++++++---- mace/ops/common/conv_pool_2d_util.cc | 8 ++- mace/ops/pooling.cc | 2 + mace/ops/ref/conv_2d.cc | 2 +- .../python/tools/mace_engine_factory.h.jinja2 | 12 ++-- .../mace/libmace/mace_api_exception_test.cc | 2 +- test/ccunit/mace/libmace/mace_api_mt_test.cc | 6 +- test/ccunit/mace/libmace/mace_api_test.cc | 2 +- .../mace/ops/memory_optimizer_for_test.cc | 27 +++++++++ .../mace/ops/memory_optimizer_for_test.h | 37 ++++++++++++ test/ccutils/mace/ops/ops_test_util.cc | 8 ++- .../template/mace_engine_factory.h.jinja2 | 12 ++-- tools/python/template/tensor_data.jinja2 | 4 ++ 31 files changed, 265 insertions(+), 138 deletions(-) create mode 100644 test/ccutils/mace/ops/memory_optimizer_for_test.cc create mode 100644 test/ccutils/mace/ops/memory_optimizer_for_test.h diff --git a/include/mace/public/mace.h b/include/mace/public/mace.h index ca723c41..6417fe98 100644 --- a/include/mace/public/mace.h +++ b/include/mace/public/mace.h @@ -416,7 +416,8 @@ class MACE_API MaceEngine { MaceStatus Init(const NetDef *net_def, const std::vector &input_nodes, const std::vector &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 &input_nodes, diff --git a/mace/core/arg_helper.cc b/mace/core/arg_helper.cc index 2cb1379b..0b078bc7 100644 --- a/mace/core/arg_helper.cc +++ b/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 ProtoArgHelper::GetRepeatedArgs( \ - const std::string &arg_name, const std::vector &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 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 ProtoArgHelper::GetRepeatedArgs( \ + const std::string &arg_name, const std::vector &default_value) \ + const { \ + if (arg_map_.count(arg_name) == 0) { \ + return default_value; \ + } else { \ + return GetRepeatedArgs(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) diff --git a/mace/core/arg_helper.h b/mace/core/arg_helper.h index e3a6319a..34a0c1f5 100644 --- a/mace/core/arg_helper.h +++ b/mace/core/arg_helper.h @@ -46,10 +46,14 @@ class ProtoArgHelper { template T GetOptionalArg(const std::string &arg_name, const T &default_value) const; + template std::vector GetRepeatedArgs( const std::string &arg_name, - const std::vector &default_value = std::vector()) const; + const std::vector &default_value) const; + + template + std::vector GetRepeatedArgs(const std::string &arg_name) const; private: std::map arg_map_; diff --git a/mace/core/memory_optimizer.cc b/mace/core/memory_optimizer.cc index 12581af6..a6aacc2c 100644 --- a/mace/core/memory_optimizer.cc +++ b/mace/core/memory_optimizer.cc @@ -114,12 +114,9 @@ void MemoryOptimizer::Optimize( const mace::OperatorDef *op_def, const std::unordered_map *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(op_def->device_type()); DataType op_dtype = static_cast(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& MemoryOptimizer::mem_blocks() const { +const std::vector &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() << "]"; } diff --git a/mace/core/memory_optimizer.h b/mace/core/memory_optimizer.h index 8cbfc5dc..6d7fdfb1 100644 --- a/mace/core/memory_optimizer.h +++ b/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 *mem_types = nullptr); + const std::unordered_map *mem_types); const std::vector &mem_blocks() const; diff --git a/mace/core/net.cc b/mace/core/net.cc index 1e116549..0a2f2f7a 100644 --- a/mace/core/net.cc +++ b/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 output_mem_map; + std::unordered_map *mem_types = &output_mem_map; +#else + std::unordered_map *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(); diff --git a/mace/core/ops/operator.h b/mace/core/ops/operator.h index bb4a20d5..bed56720 100644 --- a/mace/core/ops/operator.h +++ b/mace/core/ops/operator.h @@ -48,14 +48,22 @@ class Operation { return ProtoArgHelper::GetOptionalArg( *operator_def_, name, default_value); } + template std::vector GetRepeatedArgs( - const std::string &name, const std::vector &default_value = {}) const { + const std::string &name, const std::vector &default_value) const { MACE_CHECK(operator_def_, "operator_def was null!"); return ProtoArgHelper::GetRepeatedArgs( *operator_def_, name, default_value); } + template + std::vector GetRepeatedArgs(const std::string &name) const { + MACE_CHECK(operator_def_, "operator_def was null!"); + return ProtoArgHelper::GetRepeatedArgs( + *operator_def_, name); + } + DeviceType device_type() const { return static_cast(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(); } diff --git a/mace/core/runtime/apu/apu_wrapper.cc b/mace/core/runtime/apu/apu_wrapper.cc index 7c69e69c..08740364 100644 --- a/mace/core/runtime/apu/apu_wrapper.cc +++ b/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(model_data + const_tensor.offset()); + const_cast(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 output_tensor_ids; - std::vector output_buffers; + std::vector 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 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; @@ -223,12 +228,12 @@ bool ApuWrapper::Init(const NetDef &net_def, bool print_model = false; bool ret = frontend->InitGraph( - const_tensors.size(), const_tensors.data(), - input_tensors.size(), input_tensors.data(), - output_tensor_ids.size(), output_tensor_ids.data(), - output_buffers.data(), - ops.size(), ops.data(), - print_model); + const_tensors.size(), const_tensors.data(), + input_tensors.size(), input_tensors.data(), + output_tensor_ids.size(), output_tensor_ids.data(), + output_buffers.data(), + ops.size(), ops.data(), + print_model); cached_op_inputs.clear(); MACE_CHECK(ret == true, "apu init graph failed"); @@ -243,7 +248,7 @@ bool ApuWrapper::Run(const std::map &input_tensors, // prepare input for (int i = 0 ; i < static_cast(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 &input_tensors, // quantize quantize_util_.QuantizeWithScaleAndZeropoint( - (const float*)tensor->raw_data(), + tensor->data(), size, input_infos[i].scale, input_infos[i].zero_point, @@ -263,8 +268,8 @@ bool ApuWrapper::Run(const std::map &input_tensors, MACE_CHECK(ret == true, "neuron run model failed"); // process output - for (int i = 0 ; i < static_cast(output_tensors->size()) ; i++) { - Tensor* tensor = output_tensors->at(output_infos[i].name); + for (int i = 0; i < static_cast(output_tensors->size()); i++) { + Tensor *tensor = output_tensors->at(output_infos[i].name); // prepare out buffer tensor->SetDtype(DT_FLOAT); @@ -278,18 +283,18 @@ bool ApuWrapper::Run(const std::map &input_tensors, size, output_infos[i].scale, output_infos[i].zero_point, - reinterpret_cast(tensor->raw_mutable_data())); + tensor->mutable_data()); } return true; } bool ApuWrapper::Uninit() { - bool ret = frontend->UninitGraph(); - frontend = nullptr; - input_infos.clear(); - output_infos.clear(); - return ret; + bool ret = frontend->UninitGraph(); + frontend = nullptr; + input_infos.clear(); + output_infos.clear(); + return ret; } } // namespace mace diff --git a/mace/core/runtime/apu/apu_wrapper.h b/mace/core/runtime/apu/apu_wrapper.h index ea0fb012..7b87e56c 100755 --- a/mace/core/runtime/apu/apu_wrapper.h +++ b/mace/core/runtime/apu/apu_wrapper.h @@ -21,41 +21,42 @@ #include #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 buf; std::vector 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 Run(const std::map &input_tensors, - std::map *output_tensors); - bool Uninit(); + explicit ApuWrapper(Device *device); + bool Init(const NetDef &net_def, unsigned const char *model_data, + const index_t model_data_size); + bool Run(const std::map &input_tensors, + std::map *output_tensors); + bool Uninit(); private: - apu_data_type MapToApuDataType(DataType mace_type); - apu_pooling_mode MapToApuPoolingMode(int mace_mode); - apu_eltwise_mode MapToApuEltwiseMode(int mace_mode); + apu_data_type MapToApuDataType(DataType mace_type); + apu_pooling_mode MapToApuPoolingMode(int mace_mode); + apu_eltwise_mode MapToApuEltwiseMode(int mace_mode); private: - ApuFrontend* frontend; - std::vector input_infos; - std::vector output_infos; - QuantizeUtil quantize_util_; + ApuFrontend *frontend; + std::vector input_infos; + std::vector output_infos; + QuantizeUtil quantize_util_; }; } // namespace mace diff --git a/mace/core/runtime/hexagon/hexagon_control_wrapper.h b/mace/core/runtime/hexagon/hexagon_control_wrapper.h index 801c095d..8fa3bd26 100644 --- a/mace/core/runtime/hexagon/hexagon_control_wrapper.h +++ b/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( diff --git a/mace/core/runtime/hexagon/hexagon_dsp_wrapper.cc b/mace/core/runtime/hexagon/hexagon_dsp_wrapper.cc index dedff7f9..4519e19c 100644 --- a/mace/core/runtime/hexagon/hexagon_dsp_wrapper.cc +++ b/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(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_, diff --git a/mace/core/runtime/hexagon/hexagon_dsp_wrapper.h b/mace/core/runtime/hexagon/hexagon_dsp_wrapper.h index 2ec98e7a..204d5c49 100644 --- a/mace/core/runtime/hexagon/hexagon_dsp_wrapper.h +++ b/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 &input_tensors, diff --git a/mace/core/runtime/hexagon/hexagon_hta_wrapper.cc b/mace/core/runtime/hexagon/hexagon_hta_wrapper.cc index 24a5eb82..a6b25939 100644 --- a/mace/core/runtime/hexagon/hexagon_hta_wrapper.cc +++ b/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(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_, diff --git a/mace/core/runtime/hexagon/hexagon_hta_wrapper.h b/mace/core/runtime/hexagon/hexagon_hta_wrapper.h index aabcdc61..b774ef80 100644 --- a/mace/core/runtime/hexagon/hexagon_hta_wrapper.h +++ b/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 &input_tensors, diff --git a/mace/core/tensor.h b/mace/core/tensor.h index 33ce0c22..1340f188 100644 --- a/mace/core/tensor.h +++ b/mace/core/tensor.h @@ -276,6 +276,7 @@ class Tensor { inline void *raw_mutable_data() { MACE_CHECK_NOTNULL(buffer_); + return buffer_->raw_mutable_data(); } diff --git a/mace/core/workspace.cc b/mace/core/workspace.cc index 08bf59b0..c735b2fb 100644 --- a/mace/core/workspace.cc +++ b/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( + 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 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(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(model_data_size)); + static_cast(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(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( new Buffer(device->allocator(), const_cast(model_data), - model_data_size)); + valid_data_size)); } else { tensor_buffer_ = std::unique_ptr( 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(model_data), - 0, model_data_size); + 0, valid_data_size); tensor_buffer_->UnMap(); } for (auto &const_tensor : net_def.tensors()) { diff --git a/mace/core/workspace.h b/mace/core/workspace.h index eae7ebd5..af27adf5 100644 --- a/mace/core/workspace.h +++ b/mace/core/workspace.h @@ -56,9 +56,9 @@ class Workspace { std::vector 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, diff --git a/mace/libmace/capability.cc b/mace/libmace/capability.cc index 46896fcd..40a342c8 100644 --- a/mace/libmace/capability.cc +++ b/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; } diff --git a/mace/libmace/mace.cc b/mace/libmace/mace.cc index 732b688f..dcd61557 100644 --- a/mace/libmace/mace.cc +++ b/mace/libmace/mace.cc @@ -433,7 +433,8 @@ class MaceEngine::Impl { MaceStatus Init(const NetDef *net_def, const std::vector &input_nodes, const std::vector &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 &input_nodes, @@ -557,7 +558,7 @@ MaceStatus MaceEngine::Impl::Init( const NetDef *net_def, const std::vector &input_nodes, const std::vector &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(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(model_data_->data()))); + reinterpret_cast(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 &input_nodes, const std::vector &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; } diff --git a/mace/ops/common/conv_pool_2d_util.cc b/mace/ops/common/conv_pool_2d_util.cc index 072efaa7..d679e9ef 100644 --- a/mace/ops/common/conv_pool_2d_util.cc +++ b/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"); diff --git a/mace/ops/pooling.cc b/mace/ops/pooling.cc index e74232a3..96c4e85f 100644 --- a/mace/ops/pooling.cc +++ b/mace/ops/pooling.cc @@ -94,6 +94,7 @@ class PoolingOp : public PoolingOpBase { Tensor::MappingGuard input_guard(input_tensor); Tensor::MappingGuard output_guard(output_tensor); const T *input = input_tensor->data(); + MACE_CHECK(output_tensor->dtype() == DataTypeToEnum::value); T *output = output_tensor->mutable_data(); 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 : public PoolingOpBase { Tensor::MappingGuard input_guard(input_tensor); Tensor::MappingGuard output_guard(output_tensor); const uint8_t *input = input_tensor->data(); + MACE_CHECK(output_tensor->dtype() == DataTypeToEnum::value); uint8_t *output = output_tensor->mutable_data(); int pad_hw[2] = {paddings[0] / 2, paddings[1] / 2}; diff --git a/mace/ops/ref/conv_2d.cc b/mace/ops/ref/conv_2d.cc index a97ddb1f..0bfec1ed 100644 --- a/mace/ops/ref/conv_2d.cc +++ b/mace/ops/ref/conv_2d.cc @@ -39,9 +39,9 @@ MaceStatus Conv2d::Compute(const OpContext *context, const Tensor *filter, Tensor *output) { MACE_UNUSED(context); - const std::vector in_shape = input->shape(); const std::vector filter_shape = filter->shape(); + MACE_CHECK(in_shape[1] == filter_shape[1]); std::vector out_shape(4); std::vector paddings(2); diff --git a/mace/python/tools/mace_engine_factory.h.jinja2 b/mace/python/tools/mace_engine_factory.h.jinja2 index e1502389..dc71b25d 100644 --- a/mace/python/tools/mace_engine_factory.h.jinja2 +++ b/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 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 %} diff --git a/test/ccunit/mace/libmace/mace_api_exception_test.cc b/test/ccunit/mace/libmace/mace_api_exception_test.cc index 7bfb39d0..3163d205 100644 --- a/test/ccunit/mace/libmace/mace_api_exception_test.cc +++ b/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), ""); } diff --git a/test/ccunit/mace/libmace/mace_api_mt_test.cc b/test/ccunit/mace/libmace/mace_api_mt_test.cc index 77a5977b..687a279b 100644 --- a/test/ccunit/mace/libmace/mace_api_mt_test.cc +++ b/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(data.data())); + MaceStatus status = engine.Init( + net_def.get(), input_names, output_names, + reinterpret_cast(data.data()), + data.size() * sizeof(half)); EXPECT_EQ(status, MaceStatus::MACE_SUCCESS); std::map inputs; diff --git a/test/ccunit/mace/libmace/mace_api_test.cc b/test/ccunit/mace/libmace/mace_api_test.cc index 50e897f0..01f7a64d 100644 --- a/test/ccunit/mace/libmace/mace_api_test.cc +++ b/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(data.data())); + reinterpret_cast(data.data()), data.size() * sizeof(T)); EXPECT_EQ(status, MaceStatus::MACE_SUCCESS); std::map inputs; diff --git a/test/ccutils/mace/ops/memory_optimizer_for_test.cc b/test/ccutils/mace/ops/memory_optimizer_for_test.cc new file mode 100644 index 00000000..d16abadb --- /dev/null +++ b/test/ccutils/mace/ops/memory_optimizer_for_test.cc @@ -0,0 +1,27 @@ +// 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 *mem_types) { + if (op_def->output_size() == op_def->output_shape_size()) { + MemoryOptimizer::Optimize(op_def, mem_types); + } +} + +} // namespace mace diff --git a/test/ccutils/mace/ops/memory_optimizer_for_test.h b/test/ccutils/mace/ops/memory_optimizer_for_test.h new file mode 100644 index 00000000..1b8e2491 --- /dev/null +++ b/test/ccutils/mace/ops/memory_optimizer_for_test.h @@ -0,0 +1,37 @@ +// 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 +#include + +#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 *mem_types) override; +}; + +} // namespace mace + +#endif // MACE_OPS_MEMORY_OPTIMIZER_FOR_TEST_H_ diff --git a/test/ccutils/mace/ops/ops_test_util.cc b/test/ccutils/mace/ops/ops_test_util.cc index a0761101..4de2d272 100644 --- a/test/ccutils/mace/ops/ops_test_util.cc +++ b/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( 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( op_registry_.get(), &adapted_net_def, diff --git a/tools/python/template/mace_engine_factory.h.jinja2 b/tools/python/template/mace_engine_factory.h.jinja2 index d9557be3..e84e7777 100644 --- a/tools/python/template/mace_engine_factory.h.jinja2 +++ b/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 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 %} diff --git a/tools/python/template/tensor_data.jinja2 b/tools/python/template/tensor_data.jinja2 index 337bf2b6..9dee6c64 100644 --- a/tools/python/template/tensor_data.jinja2 +++ b/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 -- GitLab