From 8d7f0c64ddf217a68abef8aedc7425bcf578b403 Mon Sep 17 00:00:00 2001
From: NAUD Maxence <maxence.naud@cea.fr>
Date: Wed, 7 Feb 2024 13:49:53 +0000
Subject: [PATCH] [WIP][NF] Start fixing ReduceMean forward kernel memory leak
---
.../ReduceMeanImpl_forward_kernels.hpp | 27 +++++++++----------
1 file changed, 13 insertions(+), 14 deletions(-)
diff --git a/include/aidge/backend/cpu/operator/ReduceMeanImpl_forward_kernels.hpp b/include/aidge/backend/cpu/operator/ReduceMeanImpl_forward_kernels.hpp
index 810f8457..d9202615 100644
--- a/include/aidge/backend/cpu/operator/ReduceMeanImpl_forward_kernels.hpp
+++ b/include/aidge/backend/cpu/operator/ReduceMeanImpl_forward_kernels.hpp
@@ -12,13 +12,15 @@
#ifndef AIDGE_CPU_OPERATOR_REDUCEMEANIMPL_FORWARD_KERNEL_H_
#define AIDGE_CPU_OPERATOR_REDUCEMEANIMPL_FORWARD_KERNEL_H_
-#include "aidge/utils/Registrar.hpp"
-#include "aidge/operator/ReduceMean.hpp"
-#include "aidge/backend/cpu/operator/ReduceMeanImpl.hpp"
-#include <array>
#include <cstddef>
#include <algorithm>
+#include <numeric> //std::accumulate
+#include <functional> //std::multiplies
+
+#include "aidge/backend/cpu/operator/ReduceMeanImpl.hpp"
#include "aidge/data/Data.hpp"
+#include "aidge/operator/ReduceMean.hpp"
+#include "aidge/utils/Registrar.hpp"
namespace Aidge {
template <class I, class O, DimSize_t DIM>
@@ -32,22 +34,19 @@ void ReduceMeanImpl_cpu_forward_kernel(const typename ReduceMean_Op<DIM>::Attrs&
const DimSize_t keepDims = std::get<1>(attrs);
// Calculate the total number of elements in the input array
- size_t totalElements = 1;
- for (size_t dimSize : inputDims) {
- totalElements *= dimSize;
- }
+ const std::size_t totalElements = std::accumulate(inputDims.cbegin(), inputDims.cend(), 1, std::multiplies<std::size_t>());
+
// Create a temporary arrays to store intermediate input/output for each Reduce op
std::vector<I> tempInArray(input, input + totalElements);
std::vector<I> tempOutArray(input, input + totalElements);
- std::vector<size_t> currentDims = inputDims;
+ std::vector<std::size_t> currentDims = inputDims;
std::size_t addedElems = 0;
- for(std::size_t i=0; i<DIM ; ++i)
+ for(std::size_t i = 0; i < DIM ; ++i)
{
addedElems = 0;
- int axis_ = std::get<0>(attrs)[i];
- std::size_t axis = axis_>=0? axis_: axis_ + inputDims.size();
+ const std::size_t axis = static_cast<std::size_t>(std::get<0>(attrs)[i]);
I* tempOutArrayPtr = tempOutArray.data();
@@ -68,7 +67,7 @@ void ReduceMeanImpl_cpu_forward_kernel(const typename ReduceMean_Op<DIM>::Attrs&
I mean = 0;
for(std::size_t l=0; l<currentDims[axis];l++)
{
- size_t idx = j * (postAxisElems * currentDims[axis]) + l * postAxisElems + k;
+ std::size_t idx = j * (postAxisElems * currentDims[axis]) + l * postAxisElems + k;
mean += tempInArray[idx];
}
tempOutArrayPtr[addedElems] = mean / currentDims[axis];
@@ -83,7 +82,7 @@ void ReduceMeanImpl_cpu_forward_kernel(const typename ReduceMean_Op<DIM>::Attrs&
else if (currentDims.size()>1)
currentDims.erase(currentDims.begin()+axis);
}
- std::copy_n(tempInArray.data(), addedElems, output);
+ std::copy_n(tempInArray.cbegin(), addedElems, output);
}
namespace {
// DIM = 1
--
GitLab