test(base): cover IO and container ownership

Author: zxy.monado

source/source_base/module_container/ATen/core/tensor_buffer.cpp

@@ -1,26 +1,33 @@
 #include <ATen/core/tensor_buffer.h>
-
 #include <base/core/cpu_allocator.h>
 #include <base/macros/macros.h>
 
 #if defined(__CUDA) || defined(__ROCM)
 #include <base/core/gpu_allocator.h>
 #endif
 
-namespace container {
+namespace container
+{
 
 // Construct a new TensorBuffer object.
-TensorBuffer::TensorBuffer(base::core::Allocator* alloc, void* data_ptr) : alloc_(alloc), data_(data_ptr), owns_memory_(true) {}
+TensorBuffer::TensorBuffer(base::core::Allocator* alloc, void* data_ptr)
+    : alloc_(alloc), data_(data_ptr), owns_memory_(true)
+{
+}
 
 // Construct a new TensorBuffer object.
 // Note, this is a reference TensorBuffer, does not own memory itself.
-TensorBuffer::TensorBuffer(void* data_ptr) : alloc_(), data_(data_ptr), owns_memory_(false) {}
+TensorBuffer::TensorBuffer(void* data_ptr) : alloc_(), data_(data_ptr), owns_memory_(false)
+{
+}
 
-// Class members are initialized in the order of their declaration, 
+// Class members are initialized in the order of their declaration,
 // rather than the order they appear in the initialization list!
-TensorBuffer::TensorBuffer(base::core::Allocator* alloc, size_t size) {
-    alloc_ = alloc; 
-    if (size > 0) {
+TensorBuffer::TensorBuffer(base::core::Allocator* alloc, size_t size)
+{
+    alloc_ = alloc;
+    if (size > 0)
+    {
         data_ = alloc_->allocate(size);
         owns_memory_ = true;
         allocated_bytes_ = size;
@@ -29,106 +36,142 @@ TensorBuffer::TensorBuffer(base::core::Allocator* alloc, size_t size) {
 
 // Move constructor.
 TensorBuffer::TensorBuffer(TensorBuffer&& other) noexcept
-        : alloc_(other.alloc_),
-          data_(other.data_), 
-          owns_memory_(other.owns_memory_),
-          allocated_bytes_(other.allocated_bytes_)
+    : alloc_(other.alloc_), data_(other.data_), owns_memory_(other.owns_memory_),
+      allocated_bytes_(other.allocated_bytes_)
 {
     // Reset the other TensorBuffer.
+    other.alloc_ = nullptr;
     other.data_ = nullptr;
     other.owns_memory_ = false;
     other.allocated_bytes_ = 0;
 }
 
 // Destroy the TensorBuffer object.
-TensorBuffer::~TensorBuffer() {
-    if (this->OwnsMemory() && data_ != nullptr) {
+TensorBuffer::~TensorBuffer()
+{
+    if (this->OwnsMemory() && data_ != nullptr)
+    {
         alloc_->free(data_);
     }
-    if (alloc_ != nullptr) {
+    if (alloc_ != nullptr)
+    {
         delete alloc_;
     }
 }
 
 // Get the raw data pointer.
-void* TensorBuffer::data() const { return data_; }
+void* TensorBuffer::data() const
+{
+    return data_;
+}
 
 // Get the total number of bytes allocated for the buffer.
 // This method returns the total number of bytes allocated for the buffer by the allocator
 // associated with the TensorBuffer. If the buffer is not yet allocated, the function returns 0.
-size_t TensorBuffer::GetAllocatedBytes() const {
+size_t TensorBuffer::GetAllocatedBytes() const
+{
     return allocated_bytes_;
 }
 
 // Get the root TensorBuffer object.
 // If this TensorBuffer is a sub-buffer of another TensorBuffer, returns that
 // TensorBuffer. Otherwise, returns this.
-TensorBuffer* TensorBuffer::root_buffer() { return this; } // Implementation goes here.
+TensorBuffer* TensorBuffer::root_buffer()
+{
+    return this;
+} // Implementation goes here.
 
 // Get the Allocator object used in this class.
-base::core::Allocator* TensorBuffer::allocator() const {
+base::core::Allocator* TensorBuffer::allocator() const
+{
     return alloc_;
 }
 
 // Check whether this TensorBuffer owns the underlying memory.
-bool TensorBuffer::OwnsMemory() const { return this->owns_memory_; }
+bool TensorBuffer::OwnsMemory() const
+{
+    return this->owns_memory_;
+}
 
 // Get the type of device used by the TensorBuffer.
-DeviceType TensorBuffer::GetDeviceType() const {
-    if (alloc_ != nullptr) {
+DeviceType TensorBuffer::GetDeviceType() const
+{
+    if (alloc_ != nullptr)
+    {
         return alloc_->GetDeviceType();
     }
     return DeviceType::UnKnown;
 }
 
-void TensorBuffer::resize(size_t size) {
+void TensorBuffer::resize(size_t size)
+{
     // Allocate a new buffer.
     void* new_data = this->alloc_->allocate(size);
 
     // Free the old buffer.
-    if (this->OwnsMemory()) {
+    if (this->OwnsMemory())
+    {
         this->alloc_->free(data_);
     }
 
     // Update the internal state.
     this->data_ = new_data;
     this->owns_memory_ = true;
+    this->allocated_bytes_ = size;
 }
+TensorBuffer& TensorBuffer::operator=(const TensorBuffer& other)
+{
+    if (this == &other)
+    {
+        return *this;
+    }
 
-TensorBuffer& TensorBuffer::operator=(const TensorBuffer& other) {
-    if (this->OwnsMemory()) {
+    if (this->OwnsMemory())
+    {
         this->alloc_->free(data_);
     }
 
     delete this->alloc_;
-    if (other.GetDeviceType() == DeviceType::CpuDevice) {
+    if (other.GetDeviceType() == DeviceType::CpuDevice)
+    {
         this->alloc_ = new base::core::CPUAllocator();
     }
-    #if defined(__CUDA) || defined(__ROCM)
-    else if (other.GetDeviceType() == DeviceType::GpuDevice) {
+#if defined(__CUDA) || defined(__ROCM)
+    else if (other.GetDeviceType() == DeviceType::GpuDevice)
+    {
         this->alloc_ = new base::core::GPUAllocator();
     }
-    #endif // __CUDA || __ROCM
-
+#endif // __CUDA || __ROCM
 
     this->data_ = this->alloc_->allocate(other.GetAllocatedBytes());
     this->owns_memory_ = true;
+    this->allocated_bytes_ = other.GetAllocatedBytes();
     return *this;
 }
 
-TensorBuffer& TensorBuffer::operator=(TensorBuffer&& other) noexcept {
-    if (this->OwnsMemory()) {
+TensorBuffer& TensorBuffer::operator=(TensorBuffer&& other) noexcept
+{
+    if (this == &other)
+    {
+        return *this;
+    }
+
+    if (this->OwnsMemory())
+    {
         this->alloc_->free(data_);
     }
     delete this->alloc_;
     this->alloc_ = other.alloc_;
     this->data_ = other.data_;
     this->owns_memory_ = other.owns_memory_;
+    this->allocated_bytes_ = other.allocated_bytes_;
 
     // Reset the other TensorBuffer.
+    other.alloc_ = nullptr;
     other.data_ = nullptr;
     other.owns_memory_ = false;
+    other.allocated_bytes_ = 0;
     return *this;
 }
 
-}  // namespace container
+} // namespace container

source/source_base/module_container/test/allocator_test.cpp

@@ -25,11 +25,14 @@ TEST(CPUAllocator, AllocateAndFree) {
 }
 
 TEST(CPUAllocator, AllocatedSize) {
-  base::core::CPUAllocator alloc;
+  base::core::Allocator* alloc = new base::core::CPUAllocator();
   // Allocate memory of size 100 and check its size.
-  void* ptr = alloc.allocate(100);
+  void* ptr = alloc->allocate(100);
   EXPECT_NE(nullptr, ptr);
-  alloc.free(ptr);
+  EXPECT_EQ(alloc->AllocatedSize(ptr), 100);
+  alloc->free(ptr);
+  EXPECT_EQ(alloc->AllocatedSize(ptr), 0);
+  delete alloc;
 }
 
 TEST(CPUAllocator, GetDeviceType) {
@@ -49,3 +52,32 @@ TEST(LoggingDeathTest, AbortsWithContext)
   EXPECT_DEATH(base::utils::check_exit_impl("allocate", "allocator_test.cpp", 42, "allocation failed"),
                "Fatal error.*allocation failed");
 }
+
+namespace
+{
+class TestCounted : public base::core::counted_base
+{
+  public:
+    explicit TestCounted(bool* destroyed) : destroyed_(destroyed) {}
+
+    ~TestCounted() override
+    {
+        *destroyed_ = true;
+    }
+
+  private:
+    bool* destroyed_;
+};
+} // namespace
+
+TEST(RefCount, TracksReferencesAndDeletes)
+{
+    bool destroyed = false;
+    TestCounted* counted = new TestCounted(&destroyed);
+    EXPECT_TRUE(counted->ref_count_is_one());
+    counted->ref();
+    EXPECT_EQ(counted->ref_count(), 2);
+    EXPECT_FALSE(counted->unref());
+    EXPECT_TRUE(counted->unref());
+    EXPECT_TRUE(destroyed);
+}

source/source_base/module_container/test/tensor_buffer_test.cpp

@@ -40,6 +40,7 @@ TEST(TensorBuffer, resize) {
     // Resize the buffer.
     const size_t new_buffer_size = 200;
     tensor_buffer.resize(new_buffer_size);
+    EXPECT_EQ(tensor_buffer.GetAllocatedBytes(), new_buffer_size);
 
     // Free the memory.
     // auto free by the destructor
@@ -82,4 +83,55 @@ TEST(TensorBuffer, empty_allocator) {
 
     // Free the memory.
     alloc.free(buffer);
-}
+}
+
+TEST(TensorBuffer, OwningPointerConstructor)
+{
+    base::core::Allocator* alloc = new base::core::CPUAllocator();
+    void* data = alloc->allocate(16);
+    container::TensorBuffer buffer(alloc, data);
+
+    EXPECT_EQ(buffer.allocator(), alloc);
+    EXPECT_EQ(buffer.data(), data);
+    EXPECT_TRUE(buffer.OwnsMemory());
+}
+
+TEST(TensorBuffer, MoveConstructor)
+{
+    container::TensorBuffer source(new base::core::CPUAllocator(), 16);
+    void* data = source.data();
+
+    container::TensorBuffer destination(std::move(source));
+
+    EXPECT_EQ(destination.data(), data);
+    EXPECT_EQ(destination.GetAllocatedBytes(), 16);
+    EXPECT_TRUE(destination.OwnsMemory());
+    EXPECT_EQ(source.data(), nullptr);
+    EXPECT_FALSE(source.OwnsMemory());
+}
+
+TEST(TensorBuffer, CopyAssignment)
+{
+    container::TensorBuffer source(new base::core::CPUAllocator(), 16);
+    container::TensorBuffer destination(new base::core::CPUAllocator(), 8);
+
+    destination = source;
+
+    EXPECT_NE(destination.data(), source.data());
+    EXPECT_EQ(destination.GetDeviceType(), container::DeviceType::CpuDevice);
+    EXPECT_TRUE(destination.OwnsMemory());
+}
+
+TEST(TensorBuffer, MoveAssignment)
+{
+    container::TensorBuffer source(new base::core::CPUAllocator(), 16);
+    void* data = source.data();
+    container::TensorBuffer destination(new base::core::CPUAllocator(), 8);
+
+    destination = std::move(source);
+
+    EXPECT_EQ(destination.data(), data);
+    EXPECT_TRUE(destination.OwnsMemory());
+    EXPECT_EQ(source.data(), nullptr);
+    EXPECT_FALSE(source.OwnsMemory());
+}

source/source_base/module_container/test/tensor_test.cpp

@@ -43,6 +43,18 @@ TEST(Tensor, Constructor) {
     EXPECT_EQ(t4.data(), vec.data());
 }
 
+TEST(Tensor, CopyAssignment)
+{
+    container::Tensor source({1, 2, 3, 4});
+    source.reshape({2, 2});
+    container::Tensor destination;
+
+    destination = source;
+
+    EXPECT_EQ(destination, source);
+    EXPECT_NE(destination.data(), source.data());
+}
+
 
 TEST(Tensor, GetDataPointer) {
     // Create a 1x1 float tensor with data [1.0, 2.0, 3.0, 4.0].