Feature/add freeOnDevice() method (#298)

XL64 · Xavier Lacoste · acitrain · web-flow · commit 7aace428c3ea · 2024-06-24T14:31:13.000-07:00
* Add freeOnDevice method to free device buffer

* Add logging for ACTION_ALLOC/ACTION_FREE when user callback activated

* add freeOnDevice for arrays of arrays

---------

Co-authored-by: Xavier Lacoste &lt;xavier.lacoste@totalenergies.com&gt;
Co-authored-by: Aurelien &lt;aurelien.citrain@inria.fr&gt;
Co-authored-by: Jie Meng &lt;jie.meng@total.com&gt;
diff --git a/src/ArrayOfArraysView.hpp b/src/ArrayOfArraysView.hpp
@@ -596,6 +596,16 @@ class ArrayOfArraysView
   }
 
 
+  /**
+   * Move data to host if not valid and register host as the last valid space and free data on device.
+   */
+  void freeOnDevice() const
+  {
+    m_values.freeOnDevice();
+    m_sizes.freeOnDevice();
+    m_offsets.freeOnDevice();
+  }
+
   ///@}
 
 protected:
diff --git a/src/ArrayView.hpp b/src/ArrayView.hpp
@@ -675,6 +675,14 @@ class ArrayView
   void move( MemorySpace const space, bool const touch=true ) const
   { m_dataBuffer.moveNested( space, size(), touch ); }
 
+  /**
+   * Move data to host if not valid and register host as the last valid space and free data on device.
+   */
+  void freeOnDevice() const
+  {
+    m_dataBuffer.freeOnDevice();
+  }
+
   ///@}
 
 #if defined(LVARRAY_USE_TOTALVIEW_OUTPUT) && !defined(__CUDA_ARCH__)
diff --git a/src/ChaiBuffer.hpp b/src/ChaiBuffer.hpp
@@ -373,6 +373,21 @@ class ChaiBuffer
   #endif
   }
 
+  /**
+   * @brief Delete the buffer on device (no-op here)
+   */
+  inline
+  void freeOnDevice() const
+  {
+  #if defined(LVARRAY_USE_CUDA) || defined(LVARRAY_USE_HIP )
+    chai::ExecutionSpace const chaiSpace = chai::CPU;
+    move( internal::toMemorySpace( chaiSpace ), true );
+    // if T is const the touch == true is ignored in move(), we force it.
+    m_pointerRecord->m_touched[ chaiSpace ] = true;
+    internal::getArrayManager().free( m_pointerRecord, chai::GPU );
+  #endif
+  }
+
   /**
    * @return Return the capacity of the buffer.
    */
@@ -499,6 +514,40 @@ class ChaiBuffer
         char const * const spaceStr = ( s == chai::CPU ) ? "HOST  " : "DEVICE";
         LVARRAY_LOG( "Moved " << paddedSize << " to the " << spaceStr << ": " << typeString << " " << name );
       }
+
+      if( act == chai::ACTION_ALLOC )
+      {
+        std::string const size = system::calculateSize( record->m_size );
+        std::string const paddedSize = std::string( 9 - size.size(), ' ' ) + size;
+        #if defined(LVARRAY_USE_CUDA)
+        size_t free, total;
+        cudaMemGetInfo( &free, &total );
+        std::string const size2 = system::calculateSize( free );
+        #endif
+        char const * const spaceStr = ( s == chai::CPU ) ? "HOST  " : "DEVICE";
+        #if defined(LVARRAY_USE_CUDA)
+        LVARRAY_LOG( "Allocated " << paddedSize << " to the " << spaceStr << ": " << typeString << " " << name << " Free memory on device: " << size2 );
+        #else
+        LVARRAY_LOG( "Allocated " << paddedSize << " to the " << spaceStr << ": " << typeString << " " << name );
+        #endif
+      }
+
+      if( act == chai::ACTION_FREE )
+      {
+        std::string const size = system::calculateSize( record->m_size );
+        std::string const paddedSize = std::string( 9 - size.size(), ' ' ) + size;
+        #if defined(LVARRAY_USE_CUDA)
+        size_t free, total;
+        cudaMemGetInfo( &free, &total );
+        std::string const size2 = system::calculateSize( free );
+        #endif
+        char const * const spaceStr = ( s == chai::CPU ) ? "HOST  " : "DEVICE";
+        #if defined(LVARRAY_USE_CUDA)
+        LVARRAY_LOG( "Freed " << paddedSize << " to the " << spaceStr << ": " << typeString << " " << name  << " Free memory on device: " << size2 );
+        #else
+        LVARRAY_LOG( "Freed " << paddedSize << " to the " << spaceStr << ": " << typeString << " " << name );
+        #endif
+      }
     };
   }
 
diff --git a/src/MallocBuffer.hpp b/src/MallocBuffer.hpp
@@ -163,6 +163,12 @@ class MallocBuffer : public bufferManipulation::VoidBuffer
     m_data = nullptr;
   }
 
+  /**
+   * @brief Delete the buffer on device (no-op here)
+   */
+  LVARRAY_HOST_DEVICE inline
+  void freeOnDevice() const {}
+
   /**
    * @return Return the capacity of the buffer.
    */
