From de4a79119fd5a063d6e1acccfa87f9db01462a30 Mon Sep 17 00:00:00 2001 From: wanghuancoder Date: Mon, 31 Oct 2022 17:14:01 +0800 Subject: [PATCH] fix predictor memory write overflow (#47485) * fix predictor memory write overflow --- paddle/fluid/inference/api/analysis_predictor.cc | 5 +++++ paddle/fluid/inference/api/api_impl.cc | 6 ++++++ .../tests/api/analyzer_transformer_tester_helper.h | 2 +- 3 files changed, 12 insertions(+), 1 deletion(-) diff --git a/paddle/fluid/inference/api/analysis_predictor.cc b/paddle/fluid/inference/api/analysis_predictor.cc index 0190d9c291..8663ec7d1f 100644 --- a/paddle/fluid/inference/api/analysis_predictor.cc +++ b/paddle/fluid/inference/api/analysis_predictor.cc @@ -182,6 +182,11 @@ bool PaddleTensorToLoDTensor(const PaddleTensor &pt, pt.data.data(), paddle::platform::errors::InvalidArgument( "The data contained in the input PaddleTensor is illegal.")); + PADDLE_ENFORCE_EQ( + pt.data.length(), + t->numel() * paddle::experimental::SizeOf(t->dtype()), + paddle::platform::errors::InvalidArgument( + "The data contained in the input PaddleTensor had wrong length.")); } if (platform::is_cpu_place(place)) { diff --git a/paddle/fluid/inference/api/api_impl.cc b/paddle/fluid/inference/api/api_impl.cc index 865c08fe2f..1b6779ac9f 100644 --- a/paddle/fluid/inference/api/api_impl.cc +++ b/paddle/fluid/inference/api/api_impl.cc @@ -234,6 +234,12 @@ bool NativePaddlePredictor::SetFeed(const std::vector &inputs, inputs[i].data.data(), platform::errors::InvalidArgument( "The data of input tensor should not be null.")); + PADDLE_ENFORCE_EQ( + inputs[i].data.length(), + input.numel() * paddle::experimental::SizeOf(input.dtype()), + paddle::platform::errors::InvalidArgument( + "The data contained in the input PaddleTensor had wrong length.")); + if (platform::is_cpu_place(place_)) { // TODO(panyx0718): Init LoDTensor from existing memcpy to save a copy. std::memcpy(static_cast(input_ptr), diff --git a/paddle/fluid/inference/tests/api/analyzer_transformer_tester_helper.h b/paddle/fluid/inference/tests/api/analyzer_transformer_tester_helper.h index 68167c2a31..569d62f637 100644 --- a/paddle/fluid/inference/tests/api/analyzer_transformer_tester_helper.h +++ b/paddle/fluid/inference/tests/api/analyzer_transformer_tester_helper.h @@ -154,7 +154,7 @@ void PrepareInputs(std::vector *input_slots, init_idx.name = "init_idx"; init_idx.shape.assign({batch_size}); - init_idx.dtype = PaddleDType::INT32; + init_idx.dtype = PaddleDType::INT64; TensorAssignData(&init_idx, one_batch.init_idx); trg_src_attn_bias.name = "trg_src_attn_bias"; -- GitLab