[onert-micro] Fix training configure tool UB

This commit fixes undefined behavior in training configure tool. ONE-DCO-1.0-Signed-off-by: Vyacheslav Bazhenov <[email protected]>
Samsung · Dec 3, 2024 · 71aa892 · 71aa892
1 parent 5d8026f
commit 71aa892
Show file tree

Hide file tree

Showing 2 changed files with 5 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,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