diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/common_infer.c b/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/common_infer.c index 0bee95b8467..032143b02c2 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/common_infer.c +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/common_infer.c @@ -273,6 +273,9 @@ int GetChannel(const TensorC *tensor) { } int GetElementNum(const TensorC *tensor) { + if (tensor == NULL) { + return -1; + } if (tensor->shape_size_ == 0) { return 1; // scalar mode } diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/range_infer.c b/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/range_infer.c index 4064800b6fb..c039ddf380e 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/range_infer.c +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/range_infer.c @@ -20,21 +20,14 @@ int RangeInferShape(const TensorC *const *inputs, size_t inputs_size, TensorC **outputs, size_t outputs_size, OpParameter *parameter) { - int check_ret = CheckAugmentWithMinSize(inputs, inputs_size, outputs, outputs_size, parameter, 1, 1); + int check_ret = CheckAugmentNullSizeInputTwo(inputs, inputs_size, outputs, outputs_size, parameter, 1, C3NUM, 1); if (check_ret != NNACL_OK) { return check_ret; } const TensorC *input = inputs[0]; TensorC *output = outputs[0]; - NNACL_CHECK_NULL_RETURN_ERR(input); - NNACL_CHECK_NULL_RETURN_ERR(output); - - if (inputs_size == C3NUM) { - output->data_type_ = input->data_type_; - } else { - output->data_type_ = kNumberTypeInt32; - } + output->data_type_ = inputs_size == C3NUM ? input->data_type_ : kNumberTypeInt32; output->format_ = input->format_; if (!InferFlag(inputs, inputs_size)) { return NNACL_INFER_INVALID; diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/slice_infer.c b/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/slice_infer.c index b63521a6fec..ce12667bc71 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/slice_infer.c +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/slice_infer.c @@ -95,7 +95,11 @@ int SliceInferShape(const TensorC *const *inputs, size_t inputs_size, TensorC ** int begin[MAX_SHAPE_SIZE]; int size[MAX_SHAPE_SIZE]; for (int32_t i = 0; i < param->param_length_; ++i) { - MS_CHECK_TRUE_RET(param->axis_[i] < param->param_length_, NNACL_PARAM_INVALID); + if (param->axis_[i] < 0) { + MS_CHECK_INT_ADD_NOT_OVERFLOW(param->axis_[i], (int)input->shape_size_, NNACL_PARAM_INVALID); + param->axis_[i] += (int)input->shape_size_; + } + MS_CHECK_TRUE_RET(param->axis_[i] >= 0 && param->axis_[i] < param->param_length_, NNACL_PARAM_INVALID); begin[param->axis_[i]] = param->begin_[i]; size[param->axis_[i]] = param->size_[i]; } diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/strided_slice_infer.c b/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/strided_slice_infer.c index 81543607f51..cffedbf2f39 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/strided_slice_infer.c +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/nnacl/infer/strided_slice_infer.c @@ -16,6 +16,7 @@ #include "nnacl/infer/strided_slice_infer.h" #include "nnacl/infer/infer_register.h" +#include "nnacl/op_base.h" const size_t kStridedSliceOutputNum = 1; const size_t kStridedSliceInputNum = 1; @@ -124,10 +125,9 @@ int HandleAxesInputExist(const TensorC *const *inputs, int *ndim, int *in_shape, int *stride_data = NULL; const TensorC *stride_tensor = inputs[4]; - if (GetElementNum(stride_tensor) != 0) { - if (GetElementNum(stride_tensor) != begin_ndim) { - return NNACL_ERR; - } + int stride_data_num = GetElementNum(stride_tensor); + if (stride_data_num != 0) { + MS_CHECK_TRUE_RET(stride_data_num == begin_ndim, NNACL_ERR); stride_data = (int *)(stride_tensor->data_); } @@ -271,9 +271,10 @@ int ApplyEllipsisMask(StridedSliceTransferBuffer *transfer_buffer, const int *in return NNACL_OK; } -int TransIndexToPositive(StridedSliceTransferBuffer *transfer_buffer, const int *in_shape, size_t in_shape_size) { +int TransIndexToPositive(StridedSliceTransferBuffer *transfer_buffer, const int *in_shape, size_t max_shape_size, + size_t in_shape_size) { for (size_t i = 0; i < transfer_buffer->begins_size_; i++) { - if (i >= in_shape_size) { + if (i >= max_shape_size) { return NNACL_ERR; } if (transfer_buffer->begins_[i] < 0) { @@ -282,6 +283,15 @@ int TransIndexToPositive(StridedSliceTransferBuffer *transfer_buffer, const int if (transfer_buffer->ends_[i] < 0) { transfer_buffer->ends_[i] += in_shape[i]; } + if (i < in_shape_size) { + if (transfer_buffer->begins_[i] < 0 || transfer_buffer->begins_[i] > in_shape[i]) { + return NNACL_ERR; + } + if ((transfer_buffer->ends_[i] < 0 && transfer_buffer->ends_[i] != -1) || + transfer_buffer->ends_[i] > in_shape[i]) { + return NNACL_ERR; + } + } } return NNACL_OK; } @@ -426,7 +436,7 @@ int StridedSliceInferShape(const TensorC *const *inputs, size_t inputs_size, Ten int output_shape[MAX_SHAPE_SIZE]; size_t output_shape_size = 0; ShapeSet(output_shape, &output_shape_size, in_shape, in_shape_size); - ret = TransIndexToPositive(&transfer_buffer, in_shape, MAX_SHAPE_SIZE); + ret = TransIndexToPositive(&transfer_buffer, in_shape, MAX_SHAPE_SIZE, input->shape_size_); if (ret != NNACL_OK) { return ret; } diff --git a/mindspore/lite/src/ops/compat/attr_transfer_common.cc b/mindspore/lite/src/ops/compat/attr_transfer_common.cc index 3789a10c739..d222d38716d 100644 --- a/mindspore/lite/src/ops/compat/attr_transfer_common.cc +++ b/mindspore/lite/src/ops/compat/attr_transfer_common.cc @@ -21,24 +21,26 @@ namespace mindspore { namespace lite { -schema::Tensor *AttrToTensor(const void *data, int data_size, bool is_array, TypeId type_id, +schema::Tensor *AttrToTensor(const void *data, size_t data_size, bool is_array, TypeId type_id, std::vector *const tensor_bufs) { if (data == nullptr || tensor_bufs == nullptr) { MS_LOG(ERROR) << "the parameter of this function is nullptr."; return nullptr; } - auto dst_tensor = (is_array ? new (std::nothrow) Tensor(type_id, {data_size}, mindspore::NHWC, Category::CONST_TENSOR) - : new (std::nothrow) Tensor(type_id, {}, mindspore::NHWC, Category::CONST_SCALAR)); - auto dst_data = dst_tensor->MutableData(); - if (dst_data == nullptr) { - MS_LOG(ERROR) << "Data from tensor is nullptr"; - delete dst_tensor; + if (data_size > static_cast(INT32_MAX)) { + MS_LOG(ERROR) << "the amount of data exceeds the INT32_MAX."; + return nullptr; + } + auto shape = is_array ? std::vector{static_cast(data_size)} : std::vector{}; + auto dst_tensor = (is_array ? new (std::nothrow) Tensor(type_id, shape, mindspore::NHWC, Category::CONST_TENSOR) + : new (std::nothrow) Tensor(type_id, shape, mindspore::NHWC, Category::CONST_SCALAR)); + if (dst_tensor == nullptr) { + MS_LOG(ERROR) << "w a tensor failed."; return nullptr; } std::vector uint8_data; uint8_data.resize(dst_tensor->Size()); memcpy(uint8_data.data(), data, dst_tensor->Size()); - auto shape = dst_tensor->shape(); flatbuffers::FlatBufferBuilder fbb(1024); auto tensor_offset = schema::CreateTensorDirect(fbb, NodeType_ValueNode, type_id, &shape, schema::Format_NHWC, 0, 0, &uint8_data); diff --git a/mindspore/lite/src/ops/compat/attr_transfer_common.h b/mindspore/lite/src/ops/compat/attr_transfer_common.h index 2ba6cdc3cc2..c9040cd4c93 100644 --- a/mindspore/lite/src/ops/compat/attr_transfer_common.h +++ b/mindspore/lite/src/ops/compat/attr_transfer_common.h @@ -26,7 +26,7 @@ namespace mindspore { namespace lite { -schema::Tensor *AttrToTensor(const void *data, int data_size, bool is_array, TypeId type_id, +schema::Tensor *AttrToTensor(const void *data, size_t data_size, bool is_array, TypeId type_id, std::vector *const tensor_bufs); } // namespace lite } // namespace mindspore diff --git a/mindspore/lite/src/runtime/infer_manager.cc b/mindspore/lite/src/runtime/infer_manager.cc index 652bfcdf16c..01be12724e9 100644 --- a/mindspore/lite/src/runtime/infer_manager.cc +++ b/mindspore/lite/src/runtime/infer_manager.cc @@ -73,6 +73,9 @@ int KernelInferShape(const std::vector &inputs, const std::vecto auto ret = kernel_interface->Infer(&in_tensors, &out_tensors, static_cast(primitive), kernel); if (ret == kLiteInferInvalid) { + for (auto output : outputs) { + output->set_shape({-1}); + } return RET_INFER_INVALID; } if (ret != kSuccess) { diff --git a/mindspore/lite/src/runtime/kernel/arm/int8/concat_int8.h b/mindspore/lite/src/runtime/kernel/arm/int8/concat_int8.h index fb163cfa56d..79e88c12723 100644 --- a/mindspore/lite/src/runtime/kernel/arm/int8/concat_int8.h +++ b/mindspore/lite/src/runtime/kernel/arm/int8/concat_int8.h @@ -40,13 +40,13 @@ class ConcatInt8CPUKernel : public InnerKernel { if (output_shape != nullptr) { free(output_shape); } - for (std::size_t i = 0; i < in_tensors().size(); i++) { - int *input_shape = concat_param_->input_shapes_[i]; - if (input_shape != nullptr) { - free(input_shape); - } - } if (concat_param_->input_shapes_ != nullptr) { + for (std::size_t i = 0; i < in_tensors().size(); i++) { + int *input_shape = concat_param_->input_shapes_[i]; + if (input_shape != nullptr) { + free(input_shape); + } + } free(concat_param_->input_shapes_); } if (concat_param_->quant_arg_.in_args_ != nullptr) {