From d49e923a3b0f0247fea1aa0d5c82b2d0c77cc0db Mon Sep 17 00:00:00 2001
From: Olivier BICHLER <olivier.bichler@cea.fr>
Date: Tue, 4 Jun 2024 09:33:49 +0200
Subject: [PATCH] Added capacity() function and fix issue
 eclipse/aidge/aidge#112

---
 include/aidge/backend/TensorImpl.hpp          | 11 +++++++++++
 include/aidge/backend/cpu/data/TensorImpl.hpp |  2 ++
 src/backend/TensorImpl.cpp                    |  6 ++++--
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/include/aidge/backend/TensorImpl.hpp b/include/aidge/backend/TensorImpl.hpp
index f3fa4ef51..e11a6d26f 100644
--- a/include/aidge/backend/TensorImpl.hpp
+++ b/include/aidge/backend/TensorImpl.hpp
@@ -182,6 +182,17 @@ public:
     */
     inline std::size_t size() const noexcept { return mNbElts; }
 
+    /**
+     * @brief Return the current capacity of the tensor, i.e. the actual memory
+     * currently being allocated. It can be different from the size:
+     * - Capacity can be 0 if the tensor memory was not yet initialized (because
+     *   of lazy initialization, memory is allocated only when it needs to be
+     *   accessed the first time).
+     * - Capacity can be > size if the tensor was downsized but memory was not
+     *   reallocated.
+    */
+    virtual std::size_t capacity() const noexcept = 0;
+
     /**
      * @brief Return the size (in bytes) of one element (scalar).
     */
diff --git a/include/aidge/backend/cpu/data/TensorImpl.hpp b/include/aidge/backend/cpu/data/TensorImpl.hpp
index e9ad19291..ff5a4fc4b 100644
--- a/include/aidge/backend/cpu/data/TensorImpl.hpp
+++ b/include/aidge/backend/cpu/data/TensorImpl.hpp
@@ -43,6 +43,8 @@ public:
         return std::make_shared<TensorImpl_cpu<T>>(device, dims);
     }
 
+    inline std::size_t capacity() const noexcept override final { return mData.size(); }
+
     inline std::size_t scalarSize() const noexcept override final { return sizeof(T); }
 
     void zeros() override final;
diff --git a/src/backend/TensorImpl.cpp b/src/backend/TensorImpl.cpp
index ee2f82a9c..335122d05 100644
--- a/src/backend/TensorImpl.cpp
+++ b/src/backend/TensorImpl.cpp
@@ -15,7 +15,9 @@
 #include "aidge/utils/ErrorHandling.hpp"
 
 void Aidge::TensorImpl::copyFrom(const TensorImpl& srcImpl, NbElts_t length, NbElts_t srcOffset, NbElts_t dstOffset) {
-    if (&srcImpl == this && srcOffset == dstOffset) {
+    // Returns if src and dst are the same
+    // OR if src capacity is 0 (no valid data must be copied)
+    if ((&srcImpl == this && srcOffset == dstOffset) || srcImpl.capacity() == 0) {
         return;
     }
 
@@ -24,7 +26,7 @@ void Aidge::TensorImpl::copyFrom(const TensorImpl& srcImpl, NbElts_t length, NbE
             // Same backend, but different device
             copyFromDevice(srcImpl.rawPtr(srcOffset), srcImpl.device(), length, dstOffset);
         }
-        else if (srcImpl.hostPtr() != nullptr) {
+        else if (srcImpl.hostPtr() != nullptr) { // capacity() is > 0 so hostPtr() will not assert
             // Different backend, but input is valid on host
             copyFromHost(srcImpl.hostPtr(srcOffset), length, dstOffset);
         }
-- 
GitLab