From 7f8add7ab88c5ccda4b2983e88cbf5a31b8262c9 Mon Sep 17 00:00:00 2001 From: Raman Sarokin Date: Tue, 21 Apr 2020 10:45:01 -0700 Subject: [PATCH] Fixed reading out of bounds in Pooling when no output_indices. PiperOrigin-RevId: 307635625 Change-Id: I1294acd8ab58df6ad181e50f3743704f002294a8 --- tensorflow/lite/delegates/gpu/cl/kernels/pooling.cc | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/pooling.cc b/tensorflow/lite/delegates/gpu/cl/kernels/pooling.cc index fb985461c02..e292f2dad7d 100644 --- a/tensorflow/lite/delegates/gpu/cl/kernels/pooling.cc +++ b/tensorflow/lite/delegates/gpu/cl/kernels/pooling.cc @@ -185,9 +185,11 @@ std::string GetMaxPoolingKernelCode( TensorCodeGenerator dst_tensor( "dst_data", WHSPoint{"dst_size.x", "dst_size.y", "dst_size.z"}, op_def.dst_tensors[0]); + const auto dst_ind_def = + output_indices ? op_def.dst_tensors[1] : op_def.dst_tensors[0]; TensorCodeGenerator indices_tensor( "dst_indices", WHSPoint{"dst_size.x", "dst_size.y", "dst_size.z"}, - op_def.dst_tensors[1]); + dst_ind_def); std::string c = GetCommonDefines(op_def.precision); @@ -281,10 +283,12 @@ std::string GetMaxPooling3DKernelCode( "dst_data", WHDSPoint{"dst_size.x", "dst_size.y", "dst_size.z", "dst_size.w"}, op_def.dst_tensors[0]); + const auto dst_ind_def = + output_indices ? op_def.dst_tensors[1] : op_def.dst_tensors[0]; TensorCodeGenerator indices_tensor( "dst_indices", WHDSPoint{"dst_size.x", "dst_size.y", "dst_size.z", "dst_size.w"}, - op_def.dst_tensors[1]); + dst_ind_def); std::string c = GetCommonDefines(op_def.precision);