diff --git a/mindspore/lite/nnacl/fp32/reverse_fp32.c b/mindspore/lite/nnacl/fp32/reverse_fp32.c index aa15d32016..45aa7179d6 100644 --- a/mindspore/lite/nnacl/fp32/reverse_fp32.c +++ b/mindspore/lite/nnacl/fp32/reverse_fp32.c @@ -16,10 +16,13 @@ #include "nnacl/fp32/reverse_fp32.h" #include -#include "nnacl/op_base.h" #include "nnacl/errorcode.h" +#include "nnacl/nnacl_utils.h" int Reverse(const float *input, float *output, size_t elem_size, int *index) { + for (int i = 0; i < elem_size; i++) { + NNACL_ASSERT(index[i] >= 0); + } for (int i = 0; i < elem_size; i++) { output[index[i]] = input[i]; } diff --git a/mindspore/lite/nnacl/fp32/transpose_fp32.c b/mindspore/lite/nnacl/fp32/transpose_fp32.c index 5b4f587953..011594c630 100644 --- a/mindspore/lite/nnacl/fp32/transpose_fp32.c +++ b/mindspore/lite/nnacl/fp32/transpose_fp32.c @@ -176,6 +176,11 @@ int DoTransposeFp32(const float *in_data, float *out_data, int *input_shape, con (void)memcpy(out_data, in_data, data_size); return NNACL_OK; } + for (int i = 0; i < num_axes; ++i) { + if (perm[i] < 0) { + return NNACL_PARAM_INVALID; + } + } if (num_axes == 2) { TransposeDim2Fp32(in_data, out_data, strides, out_strides, perm, output_shape, h_start, h_end); } else if (num_axes == 3) { diff --git a/mindspore/lite/nnacl/nnacl_utils.h b/mindspore/lite/nnacl/nnacl_utils.h index b8dd2fbf44..735425d3bb 100644 --- a/mindspore/lite/nnacl/nnacl_utils.h +++ b/mindspore/lite/nnacl/nnacl_utils.h @@ -21,9 +21,18 @@ #ifdef __cplusplus extern "C" { #endif + #if defined(__arm__) || defined(__aarch64__) uint32_t getHwCap(int hwcap_type); #endif + +#ifdef DEBUG +#include +#define NNACL_ASSERT(f) assert(f) +#else +#define NNACL_ASSERT(f) ((void)0) +#endif + #ifdef __cplusplus } #endif diff --git a/mindspore/lite/nnacl/squeeze.c b/mindspore/lite/nnacl/squeeze.c index 3ec5199ce0..f038631202 100644 --- a/mindspore/lite/nnacl/squeeze.c +++ b/mindspore/lite/nnacl/squeeze.c @@ -20,7 +20,7 @@ int DoSqueeze(const float *in_data, float *out_data, size_t data_size) { if (in_data == NULL || out_data == NULL) { - return -1; + return NNACL_ERR; } (void)memcpy(out_data, in_data, data_size); return NNACL_OK; @@ -28,7 +28,7 @@ int DoSqueeze(const float *in_data, float *out_data, size_t data_size) { int DoSqueezeInt32(int32_t *in_data, int32_t *out_data, size_t data_size) { if (in_data == NULL || out_data == NULL) { - return -1; + return NNACL_ERR; } (void)memcpy(out_data, in_data, data_size); return NNACL_OK;