diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/concat_infer.c b/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/concat_infer.c index 81fda9c0b65..e6b6798067c 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/concat_infer.c +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/concat_infer.c @@ -65,11 +65,7 @@ int ConcatInferShape(const TensorC *const *inputs, size_t inputs_size, TensorC * for (size_t i = 1; i < inputs_size; ++i) { size_t input_i_shape_size = inputs[i]->shape_size_; if (input_i_shape_size != input0_shape_size) { - if (input_i_shape_size != 0) { - return NNACL_PARAM_INVALID; - } else { - continue; - } + return NNACL_PARAM_INVALID; } int shape_tmp[MAX_SHAPE_SIZE] = {0}; size_t shape_tmp_size = 0; diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/crop_infer.c b/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/crop_infer.c index 2c0393c0fe6..5f2ca237df7 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/crop_infer.c +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/crop_infer.c @@ -16,6 +16,7 @@ #include "nnacl/infer/crop_infer.h" #include "nnacl/infer/infer_register.h" +#include "nnacl/crop_parameter.h" int CropInferShape(const TensorC *const *inputs, size_t inputs_size, TensorC **outputs, size_t outputs_size, OpParameter *parameter) { @@ -24,6 +25,13 @@ int CropInferShape(const TensorC *const *inputs, size_t inputs_size, TensorC **o return check_ret; } + size_t input_shape_size = inputs[0]->shape_size_; + CropParameter *param = (CropParameter *)parameter; + int64_t axis = param->axis_ < 0 ? param->axis_ + (int64_t)input_shape_size : param->axis_; + if (axis < 0 || axis >= input_shape_size) { + return NNACL_ERR; + } + SetDataTypeFormat(outputs[0], inputs[0]); if (!InferFlag(inputs, inputs_size)) { return NNACL_INFER_INVALID; diff --git a/mindspore/lite/src/runtime/kernel/arm/base/crop_base.cc b/mindspore/lite/src/runtime/kernel/arm/base/crop_base.cc index 06954fdfcb2..f63280bd7b6 100644 --- a/mindspore/lite/src/runtime/kernel/arm/base/crop_base.cc +++ b/mindspore/lite/src/runtime/kernel/arm/base/crop_base.cc @@ -49,7 +49,7 @@ int CropBaseCPUKernel::ReSize() { int CropBaseCPUKernel::PadOffset(int input_dim, CropParameter *crop_para) const { auto axis = crop_para->axis_; auto offsets_size = crop_para->offset_size_; - MS_ASSERT(axis <= input_dim); + MS_CHECK_TRUE_MSG(axis < input_dim, RET_ERROR, "The axis is invalid."); if (offsets_size > 1) { MS_CHECK_TRUE_MSG(axis + offsets_size == input_dim, RET_ERROR, "The axis and offsets is invalid"); } diff --git a/mindspore/lite/src/runtime/kernel/arm/fp32/gatherNd_fp32.cc b/mindspore/lite/src/runtime/kernel/arm/fp32/gatherNd_fp32.cc index f0239e76574..8163632f50a 100644 --- a/mindspore/lite/src/runtime/kernel/arm/fp32/gatherNd_fp32.cc +++ b/mindspore/lite/src/runtime/kernel/arm/fp32/gatherNd_fp32.cc @@ -105,7 +105,7 @@ int GatherNdCPUKernel::InitOffset() { (void)memset(in_offset_, 0, count_ * sizeof(int)); for (int j = 0; j < count_; ++j) { for (int k = 0; k < idx_lastshape; ++k) { - if (indices_ptr[j * idx_stride + k] < in_shape[k]) { + if (indices_ptr[j * idx_stride + k] >= 0 && indices_ptr[j * idx_stride + k] < in_shape[k]) { in_offset_[j] += indices_ptr[j * idx_stride + k] * in_stride.at(k); } else { MS_LOG(ERROR) << name() << " indices value invalid!"; diff --git a/mindspore/lite/src/runtime/kernel/arm/fp32/matmul_fp32_base.cc b/mindspore/lite/src/runtime/kernel/arm/fp32/matmul_fp32_base.cc index 8f131e80bb4..de8492ca9df 100644 --- a/mindspore/lite/src/runtime/kernel/arm/fp32/matmul_fp32_base.cc +++ b/mindspore/lite/src/runtime/kernel/arm/fp32/matmul_fp32_base.cc @@ -312,7 +312,11 @@ int MatmulFp32BaseCPUKernel::Prepare() { CHECK_LESS_RETURN(in_tensors_.size(), C2NUM); CHECK_LESS_RETURN(out_tensors_.size(), 1); init_global_variable(); + MS_CHECK_INT_MUL_NOT_OVERFLOW(a_batch_, params_->row_align_, RET_ERROR); + MS_CHECK_INT_MUL_NOT_OVERFLOW(a_batch_ * params_->row_align_, params_->deep_, RET_ERROR); matrix_a_pack_size_ = a_batch_ * params_->row_align_ * params_->deep_; + MS_CHECK_INT_MUL_NOT_OVERFLOW(a_batch_, params_->col_align_, RET_ERROR); + MS_CHECK_INT_MUL_NOT_OVERFLOW(a_batch_ * params_->col_align_, params_->deep_, RET_ERROR); matrix_b_pack_size_ = b_batch_ * params_->col_align_ * params_->deep_; if (matrix_a_pack_size_ < 0 || matrix_b_pack_size_ < 0) { MS_LOG(ERROR) << "Matrix pack size is negative " diff --git a/mindspore/lite/src/runtime/kernel/arm/fp32/reverse_fp32.cc b/mindspore/lite/src/runtime/kernel/arm/fp32/reverse_fp32.cc index 28f2f1163fb..f4e2e958003 100644 --- a/mindspore/lite/src/runtime/kernel/arm/fp32/reverse_fp32.cc +++ b/mindspore/lite/src/runtime/kernel/arm/fp32/reverse_fp32.cc @@ -37,7 +37,11 @@ int ReverseCPUKernel::Stride(int index) { int ReverseCPUKernel::ReSize() { // trans negative to positive axis - UpdateAxisInfo(); + auto ret = UpdateAxisInfo(); + if (ret != RET_OK) { + MS_LOG(ERROR) << "Get axis failed."; + return ret; + } data_size_ = in_tensors_.at(0)->ElementsNum(); thread_sz_count_ = MSMIN(op_parameter_->thread_num_, data_size_); @@ -146,14 +150,19 @@ int ReverseCPUKernel::Run() { return RET_OK; } -void ReverseCPUKernel::UpdateAxisInfo() { +int ReverseCPUKernel::UpdateAxisInfo() { auto reverse_param = reinterpret_cast(op_parameter_); int in_shape_len = static_cast(in_tensors_.front()->shape().size()); for (int i = 0; i < reverse_param->num_axis_; ++i) { if (reverse_param->axis_[i] < 0) { reverse_param->axis_[i] += in_shape_len; } + if (reverse_param->axis_[i] < 0 || reverse_param->axis_[i] >= in_shape_len) { + MS_LOG(ERROR) << "Invalid axis."; + return RET_ERROR; + } } + return RET_OK; } REG_KERNEL(kCPU, kNumberTypeFloat32, PrimitiveType_ReverseV2, LiteKernelCreator) diff --git a/mindspore/lite/src/runtime/kernel/arm/fp32/reverse_fp32.h b/mindspore/lite/src/runtime/kernel/arm/fp32/reverse_fp32.h index 08c81379106..571493ac5a3 100644 --- a/mindspore/lite/src/runtime/kernel/arm/fp32/reverse_fp32.h +++ b/mindspore/lite/src/runtime/kernel/arm/fp32/reverse_fp32.h @@ -42,7 +42,7 @@ class ReverseCPUKernel : public InnerKernel { int DoReverse(int task_id); private: - void UpdateAxisInfo(); + int UpdateAxisInfo(); int thread_sz_count_ = 0; int thread_sz_stride_ = 0; int data_size_ = 0; diff --git a/mindspore/lite/src/runtime/kernel/arm/int8/softmax_int8.cc b/mindspore/lite/src/runtime/kernel/arm/int8/softmax_int8.cc index 78e51c99d58..0ca2691230d 100644 --- a/mindspore/lite/src/runtime/kernel/arm/int8/softmax_int8.cc +++ b/mindspore/lite/src/runtime/kernel/arm/int8/softmax_int8.cc @@ -51,6 +51,7 @@ int SoftmaxInt8CPUKernel::Prepare() { MS_ASSERT(input_tensor != nullptr); auto in_quant_args = input_tensor->quant_params(); + CHECK_LESS_RETURN(in_quant_args.size(), 1); quant_param_->in_quant_args_.scale_ = static_cast(in_quant_args.front().scale); quant_param_->in_quant_args_.zp_ = -in_quant_args.front().zeroPoint; @@ -58,6 +59,7 @@ int SoftmaxInt8CPUKernel::Prepare() { MS_ASSERT(out_tensor != nullptr); auto out_quant_args = out_tensor->quant_params(); + CHECK_LESS_RETURN(out_quant_args.size(), 1); quant_param_->out_quant_arg_.scale_ = static_cast(out_quant_args.front().scale); quant_param_->out_quant_arg_.zp_ = -out_quant_args.front().zeroPoint; quant_param_->output_activation_min_ = std::numeric_limits::min();