Refactor: refactor module_container (deepmodeling#2876)

* refactor module_container * set include_directory for container * refactor CMakeLists.txt for container * fix logic error of the resize method of class Tensor. --------- Co-authored-by: Qianrui <[email protected]>
maki49 · Sep 1, 2023 · d37e0f6 · d37e0f6
1 parent 71110a8
commit d37e0f6
Show file tree

Hide file tree

Showing 75 changed files with 7,385 additions and 2,586 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -99,6 +99,7 @@ set(ABACUS_SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/source)
 set(ABACUS_TEST_DIR ${CMAKE_CURRENT_SOURCE_DIR}/tests)
 set(ABACUS_BIN_PATH ${CMAKE_CURRENT_BINARY_DIR}/${ABACUS_BIN_NAME})
 include_directories(${ABACUS_SOURCE_DIR})
+include_directories(${ABACUS_SOURCE_DIR}/module_base/module_container)
 
 add_executable(${ABACUS_BIN_NAME} source/main.cpp)
 if(ENABLE_COVERAGE)
@@ -603,7 +604,6 @@ target_link_libraries(${ABACUS_BIN_NAME}
     vdw
     device
     container
-    kernels
 )
 
 if(ENABLE_LCAO)

diff --git a/source/module_base/module_container/ATen/core/CMakeLists.txt b/source/module_base/module_container/ATen/core/CMakeLists.txt
@@ -0,0 +1,2 @@
+file(GLOB ATen_CORE_SRCS "*.cpp")
+set(ATen_CPU_SRCS ${ATen_CPU_SRCS} ${ATen_CORE_SRCS} PARENT_SCOPE)
diff --git a/.../module_base/module_container/allocator.h → ...se/module_container/ATen/core/allocator.h b/.../module_base/module_container/allocator.h → ...se/module_container/ATen/core/allocator.h
@@ -1,9 +1,7 @@
-#ifndef CONTAINER_ALLOCATOR_H
-#define CONTAINER_ALLOCATOR_H
+#ifndef ATEN_CORE_ALLOCATOR_H_
+#define ATEN_CORE_ALLOCATOR_H_
 
-#include <cstddef>
-
-#include "tensor_types.h"
+#include <ATen/core/tensor_types.h>
 
 namespace container {
 
@@ -50,16 +48,24 @@ class Allocator {
      * @param ptr The pointer to get the allocated size of.
      * @return size_t The size of the allocated block of memory, in bytes.
      */
-    virtual size_t AllocatedSize(void* ptr) = 0;
+    virtual size_t AllocatedSize(void* ptr) {
+        return allocated_size_;
+    }
 
     /**
      * @brief Get the type of memory used by the TensorBuffer.
      *
      * @return MemoryType The type of memory used by the TensorBuffer.
      */
     virtual DeviceType GetDeviceType() = 0;
+
+  protected:
+    /**
+     * @brief The total number of bytes allocated by this allocator.
+     */
+    size_t allocated_size_ = 0;
 };
 
 } // namespace ABACUS
 
-#endif // CONTAINER_ALLOCATOR_H
+#endif // ATEN_CORE_ALLOCATOR_H_
diff --git a/...e_base/module_container/cpu_allocator.cpp → ...ule_container/ATen/core/cpu_allocator.cpp b/...e_base/module_container/cpu_allocator.cpp → ...ule_container/ATen/core/cpu_allocator.cpp
@@ -1,17 +1,16 @@
-#include <cstdlib> // for ::operator new, ::operator delete
-#include <cassert> // for assert
-
-#include "cpu_allocator.h"
+#include <ATen/core/cpu_allocator.h>
 
 namespace container {
 
 // Allocate a block of CPU memory with the given size and default alignment.
 void* CPUAllocator::allocate(size_t size) {
+    this->allocated_size_ = size;
     return ::operator new(size);
 }
 
 // Allocate a block of CPU memory with the given size and alignment.
 void* CPUAllocator::allocate(size_t size, size_t alignment) {
+    this->allocated_size_ = size;
     void* ptr = nullptr;
     if (posix_memalign(&ptr, alignment, size) != 0) {
         ptr = nullptr;
@@ -21,15 +20,10 @@ void* CPUAllocator::allocate(size_t size, size_t alignment) {
 
 // Free a block of CPU memory that was previously allocated by this allocator.
 void CPUAllocator::free(void* ptr) {
+    this->allocated_size_ = 0;
     ::operator delete(ptr);
 }
 
-// Get the allocated size of a given pointer.
-size_t CPUAllocator::AllocatedSize(void* ptr) {
-    std::cerr << "Method `AllocatedSize` not implemented!" << std::endl;
-    exit(EXIT_FAILURE);
-}
-
 //  Get the type of device used by the TensorBuffer.
 DeviceType CPUAllocator::GetDeviceType() {
     return DeviceType::CpuDevice;

diff --git a/...ule_base/module_container/cpu_allocator.h → ...odule_container/ATen/core/cpu_allocator.h b/...ule_base/module_container/cpu_allocator.h → ...odule_container/ATen/core/cpu_allocator.h
@@ -1,7 +1,7 @@
-#ifndef CONTAINER_CPU_ALLOCATOR_H
-#define CONTAINER_CPU_ALLOCATOR_H
+#ifndef ATEN_CORE_CPU_ALLOCATOR_H_
+#define ATEN_CORE_CPU_ALLOCATOR_H_
 
-#include "allocator.h"
+#include <ATen/core/allocator.h>
 
 namespace container {
 
@@ -41,16 +41,6 @@ class CPUAllocator : public Allocator {
      */
     void free(void* ptr) override;
 
-    /**
-     * @brief Get the allocated size of a given pointer.
-     *
-     * @param ptr The pointer to get the allocated size of.
-     * @return size_t The size of the allocated block of memory, in bytes.
-     *
-     * @note This function is not implemented for CPUAllocator and always returns 0.
-     */
-    size_t AllocatedSize(void* ptr) override;
-
     /**
      * @brief Get the type of device used by the TensorBuffer.
      *
@@ -62,4 +52,4 @@ class CPUAllocator : public Allocator {
 
 } // namespace container
 
-#endif // CONTAINER_CPU_ALLOCATOR_H
+#endif // ATEN_CORE_CPU_ALLOCATOR_H_
diff --git a/...e_base/module_container/gpu_allocator.cpp → ...ule_container/ATen/core/gpu_allocator.cpp b/...e_base/module_container/gpu_allocator.cpp → ...ule_container/ATen/core/gpu_allocator.cpp
@@ -1,7 +1,7 @@
-#include <cassert> // for assert
-#include <cuda_runtime.h> // for CUDA APIs
+#ifdef __CUDA || __ROCM
 
-#include "gpu_allocator.h"
+#include <cuda_runtime.h> // for CUDA APIs
+#include <ATen/core/gpu_allocator.h>
 
 namespace container {
 
@@ -12,6 +12,7 @@ void* GPUAllocator::allocate(size_t size) {
     if (result != cudaSuccess) {
         return nullptr;
     }
+    this->allocated_size_ = size;
     return ptr;
 }
 
@@ -22,18 +23,14 @@ void* GPUAllocator::allocate(size_t size, size_t alignment) {
     if (result != cudaSuccess) {
         return nullptr;
     }
+    this->allocated_size_ = size;
     return ptr;
 }
 
 // Free a block of CPU memory that was previously allocated by this allocator.
 void GPUAllocator::free(void* ptr) {
     cudaFree(ptr);
-}
-
-// Get the allocated size of a given pointer.
-size_t GPUAllocator::AllocatedSize(void* ptr) {
-    assert(false && "not implemented");
-    return 0;
+    this->allocated_size_ = 0;
 }
 
 // Get the type of device used by the TensorBuffer.
@@ -42,3 +39,5 @@ DeviceType GPUAllocator::GetDeviceType() {
 }
 
 } // namespace container
+
+#endif // __CUDA || __ROCM
diff --git a/...ule_base/module_container/gpu_allocator.h → ...odule_container/ATen/core/gpu_allocator.h b/...ule_base/module_container/gpu_allocator.h → ...odule_container/ATen/core/gpu_allocator.h
@@ -1,7 +1,7 @@
-#ifndef CONTAINER_GPU_ALLOCATOR_H
-#define CONTAINER_GPU_ALLOCATOR_H
+#ifndef ATEN_CORE_GPU_ALLOCATOR_H_
+#define ATEN_CORE_GPU_ALLOCATOR_H_
 
-#include "allocator.h"
+#include <ATen/core/allocator.h>
 
 namespace container {
 
@@ -39,16 +39,6 @@ class GPUAllocator : public Allocator {
      */
     void free(void* ptr) override;
 
-    /**
-     * @brief Get the allocated size of a given pointer.
-     *
-     * @param ptr The pointer to get the allocated size of.
-     * @return size_t The size of the allocated block of memory, in bytes.
-     *
-     * @note This function is not implemented for CPUAllocator and always returns 0.
-     */
-    size_t AllocatedSize(void* ptr) override;
-
     /**
      * @brief Get the type of memory used by the TensorBuffer.
      *
@@ -59,4 +49,4 @@ class GPUAllocator : public Allocator {
 
 } // namespace ABACUS
 
-#endif // CONTAINER_GPU_ALLOCATOR_H
+#endif // ATEN_CORE_GPU_ALLOCATOR_H_