From 911170b5547f2b564bb768d43ac6ff94f37e9d22 Mon Sep 17 00:00:00 2001
From: Olivier BICHLER <olivier.bichler@cea.fr>
Date: Thu, 9 Jan 2025 10:27:39 +0100
Subject: [PATCH] Fixed concat block size for NCHW
---
src/scheduler/Scheduler.cpp | 19 +++++++++++++++++--
1 file changed, 17 insertions(+), 2 deletions(-)
diff --git a/src/scheduler/Scheduler.cpp b/src/scheduler/Scheduler.cpp
index d456c8af2..b928334f3 100644
--- a/src/scheduler/Scheduler.cpp
+++ b/src/scheduler/Scheduler.cpp
@@ -707,6 +707,7 @@ Aidge::MemoryManager Aidge::Scheduler::generateMemoryAutoConcat(bool incProducer
}
}
+ size_t concatSize = size;
size_t concatOffset = 0;
if (concat) {
@@ -730,6 +731,9 @@ Aidge::MemoryManager Aidge::Scheduler::generateMemoryAutoConcat(bool incProducer
if (parentOutputFormat == DataFormat::NHWC) {
parentSize = parentOutputDims.end()[-3];
}
+ else {
+ concatSize = parentSize;
+ }
concatOffset += parentSize;
}
@@ -747,16 +751,27 @@ Aidge::MemoryManager Aidge::Scheduler::generateMemoryAutoConcat(bool incProducer
if (wrapAroundBuffer && wrapAroundSize > 0) {
memManager.reallocate(memPlane,
node, concatOffset,
- size, true, wrapAroundExtra, childs, stride, length, count);
+ concatSize, true, wrapAroundExtra, childs, stride, length, count);
}
else {
memManager.reallocate(memPlane.memSpace,
node, memPlane.offset + concatOffset,
- size, false, 0, childs, stride, length, count);
+ concatSize, false, 0, childs, stride, length, count);
}
if (concat && itConcat == concatMemPlane.end()) {
concatMemPlane.emplace(concat, memPlane);
+
+ if (wrapAroundBuffer && wrapAroundSize > 0) {
+ memManager.reallocate(memPlane,
+ concat, 0,
+ size, true, wrapAroundExtra, childs, stride, length, count);
+ }
+ else {
+ memManager.reallocate(memPlane.memSpace,
+ concat, memPlane.offset,
+ size, false, 0, childs, stride, length, count);
+ }
}
}
--
GitLab