[MSLITE] fix fuzz bug

This commit is contained in:
ling 2022-01-19 16:49:40 +08:00
parent 0b668d10fb
commit 027a475265
2 changed files with 18 additions and 0 deletions

View File

@ -42,6 +42,13 @@ int ReshapeBaseCPUKernel::Run() {
CHECK_NULL_RETURN(in_tensor);
auto out_tensor = out_tensors().front();
CHECK_NULL_RETURN(out_tensor);
if (in_tensor->data_type() != out_tensor->data_type() || in_tensor->data() == nullptr ||
in_tensor->Size() != out_tensor->Size()) {
MS_LOG(ERROR) << "Invalid param in reshape";
return RET_ERROR;
}
if (in_tensor->allocator() == nullptr || in_tensor->allocator() != out_tensor->allocator() ||
in_tensor->allocator() != ms_context_->allocator || /* runtime allocator */
op_parameter_->is_train_session_) {

View File

@ -180,6 +180,11 @@ int MatmulFp32BaseCPUKernel::InitBiasData() {
}
int MatmulFp32BaseCPUKernel::InitMatrixA(const float *src_ptr) const {
if (in_tensors().at(FIRST_INPUT) == nullptr || in_tensors().at(FIRST_INPUT)->data_type() != kNumberTypeFloat32) {
MS_LOG(ERROR) << "Invalid param in matmul";
return RET_ERROR;
}
CHECK_NULL_RETURN(src_ptr);
if (vec_matmul_) {
return RET_OK;
@ -197,6 +202,11 @@ int MatmulFp32BaseCPUKernel::InitMatrixA(const float *src_ptr) const {
}
int MatmulFp32BaseCPUKernel::InitMatrixB(const float *src_ptr) const {
if (in_tensors().at(SECOND_INPUT) == nullptr || in_tensors().at(SECOND_INPUT)->data_type() != kNumberTypeFloat32) {
MS_LOG(ERROR) << "Invalid param in matmul";
return RET_ERROR;
}
CHECK_NULL_RETURN(src_ptr);
for (int i = 0; i < b_batch_; i++) {
const float *src = src_ptr + i * params_->deep_ * params_->col_;
@ -402,6 +412,7 @@ int MatmulFp32BaseCPUKernel::Prepare() {
MS_LOG(ERROR) << "InitBiasData failed";
return ret;
}
if (params_->a_const_) {
auto a_tensor = in_tensors_[0];
CHECK_NULL_RETURN(a_tensor);