From 16946d05558ee831663e253f66eb4717855abe7a Mon Sep 17 00:00:00 2001
From: Stanislav Minakov <sminakov@tenstorrent.com>
Date: Wed, 23 Oct 2024 22:36:01 +0000
Subject: [PATCH] #0: Fixup

---
 tt_metal/impl/buffers/buffer.cpp | 32 +++++++++++++++-----------------
 tt_metal/impl/buffers/buffer.hpp |  1 +
 2 files changed, 16 insertions(+), 17 deletions(-)

diff --git a/tt_metal/impl/buffers/buffer.cpp b/tt_metal/impl/buffers/buffer.cpp
index adcf2d315dd..446bf930242 100644
--- a/tt_metal/impl/buffers/buffer.cpp
+++ b/tt_metal/impl/buffers/buffer.cpp
@@ -237,31 +237,29 @@ std::shared_ptr<Buffer> Buffer::create(
 void Buffer::deallocate() {
     deallocation_requested_.store(true, std::memory_order::relaxed);
     device_->push_work([self = weak_self.lock()] {
-        if (self->allocation_status_.load(std::memory_order::relaxed) == AllocationStatus::ALLOCATED) {
-            return;
-        }
-
-        if (self->device_->initialized_ && self->size_ != 0) {
-            // address_ is only modified from this thread, no sync required
-            detail::BUFFER_MAP.erase({self->device_->id(), self->address_});
-            detail::DeallocateBuffer(self.get());
-        }
-
-        self->allocation_status_.store(AllocationStatus::DEALLOCATED, std::memory_order::relaxed);
+        self->deallocate_impl();
     });
 }
 
 void Buffer::deleter(Buffer* buffer) {
     buffer->device_->push_work([buffer] {
         std::unique_ptr<Buffer> unique_buffer = std::unique_ptr<Buffer>(buffer);
+        buffer->deallocate_impl();
+    });
+}
 
-        if (!buffer->device_->initialized_ || buffer->size_ == 0) {
-            return;
-        }
+void Buffer::deallocate_impl() {
+    if (allocation_status_.load(std::memory_order::relaxed) == AllocationStatus::DEALLOCATED) {
+        return;
+    }
 
-        detail::BUFFER_MAP.erase({buffer->device_->id(), buffer->address_});
-        detail::DeallocateBuffer(buffer);
-    });
+    if (device_->initialized_ && size_ != 0) {
+        // address_ is only modified from this thread, no sync required
+        detail::BUFFER_MAP.erase({device_->id(), address_});
+        detail::DeallocateBuffer(this);
+    }
+
+    allocation_status_.store(AllocationStatus::DEALLOCATED, std::memory_order::relaxed);
 }
 
 bool Buffer::is_allocated() const {
diff --git a/tt_metal/impl/buffers/buffer.hpp b/tt_metal/impl/buffers/buffer.hpp
index eea60d7db10..c11f4701b84 100644
--- a/tt_metal/impl/buffers/buffer.hpp
+++ b/tt_metal/impl/buffers/buffer.hpp
@@ -221,6 +221,7 @@ class Buffer final {
     // Deallocate is allowed to be called multiple times on the same buffer
     void deallocate();
     static void deleter(Buffer* buffer);
+    void deallocate_impl();
     friend void DeallocateBuffer(Buffer &buffer);
 
     DeviceAddr translate_page_address(uint64_t offset, uint32_t bank_id) const;