[onert-micro] Fix training configure tool UB

This commit fixes undefined behavior in training configure tool.
Samsung · Dec 3, 2024 · 9111592 · 9111592
1 parent 5d8026f
commit 9111592
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git 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<float *>(static_cast<void *>(dloss_dweight_data + i)) = 0;
+
       pal::FullyConnectedWeightGrad(
         core::utils::castInputData<float>(dloss_doutput_data), output_shape,
         core::utils::castInputData<float>(input_data), input_shape,

diff --git a/onert-micro/training-configure-tool/src/TrainingConfigureFileHandler.cpp b/onert-micro/training-configure-tool/src/TrainingConfigureFileHandler.cpp
@@ -32,7 +32,7 @@ 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;
+  auto buffer_size = 8 + train_op_indexes_with_ranks.size() * 3;
   buffer.resize(buffer_size);
 
   // Point to start of the buffer