From aa9a7b2e3f8b4978db3e8f60677d4ffedbd28ba2 Mon Sep 17 00:00:00 2001 From: SlavikMIPT Date: Tue, 3 Dec 2024 17:14:55 +0300 Subject: [PATCH] [onert-micro] Fix training configure tool UB (#14406) This commit fixes undefined behavior in training configure tool. ONE-DCO-1.0-Signed-off-by: Vyacheslav Bazhenov Co-authored-by: Vyacheslav Bazhenov --- onert-micro/onert-micro/src/train/kernels/FullyConnected.cpp | 5 +++-- .../src/TrainingConfigureFileHandler.cpp | 3 ++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/onert-micro/onert-micro/src/train/kernels/FullyConnected.cpp b/onert-micro/onert-micro/src/train/kernels/FullyConnected.cpp index 266d5c4b51b..a8248a7fc91 100644 --- a/onert-micro/onert-micro/src/train/kernels/FullyConnected.cpp +++ b/onert-micro/onert-micro/src/train/kernels/FullyConnected.cpp @@ -153,8 +153,9 @@ OMStatus onert_micro::train::train_kernel_CircleFullyConnected(const OMBackpropE // 2. Calculate weight gradient // Init weight grads with zeros - std::memset(dloss_dweight_data, 0, - output_shape.dims(1) * input_shape.dims(1) * sizeof(float)); + for (int i = 0; i < output_shape.dims(1) * input_shape.dims(1); i += sizeof(float)) + *static_cast(static_cast(dloss_dweight_data + i)) = 0; + pal::FullyConnectedWeightGrad( core::utils::castInputData(dloss_doutput_data), output_shape, core::utils::castInputData(input_data), input_shape, diff --git a/onert-micro/training-configure-tool/src/TrainingConfigureFileHandler.cpp b/onert-micro/training-configure-tool/src/TrainingConfigureFileHandler.cpp index 2d99113488f..efb116f809a 100644 --- a/onert-micro/training-configure-tool/src/TrainingConfigureFileHandler.cpp +++ b/onert-micro/training-configure-tool/src/TrainingConfigureFileHandler.cpp @@ -32,7 +32,8 @@ void writeTrainConfigFileDataIntoBuffer( const auto &train_op_indexes_with_ranks = train_data.trainable_op_indexes_with_ranks; // Resize to calculated size - auto buffer_size = 8 + train_op_indexes_with_ranks.size() * 2; + // 8 Bytes - handler + U16 indexes + U8 ranks + auto buffer_size = 8 + train_op_indexes_with_ranks.size() * (sizeof(uint16_t) + sizeof(uint8_t)); buffer.resize(buffer_size); // Point to start of the buffer