Changed the way CPUBuffers are implemented

datasig-ac-uk · Oct 19, 2023 · 831d1e8 · 831d1e8
1 parent b2b0fab
commit 831d1e8
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 53 deletions.
diff --git a/device/src/cpu/cpu_buffer.cpp b/device/src/cpu/cpu_buffer.cpp
@@ -42,60 +42,52 @@
 using namespace rpy;
 using namespace rpy::devices;
 
-CPUBuffer::CPUBuffer(cl_mem buffer, CPUDevice dev)
-    : ocl_buffer(buffer, dev->ocl_device()), is_ocl(true)
-{}
+CPUBuffer::CPUBuffer(CPUBuffer::RawBuffer raw, CPUBuffer::Flags arg_flags)
+    : raw_buffer(std::move(raw)),
+      flags(arg_flags)
+{
+    RPY_DBG_ASSERT(raw_ref_count(std::memory_order_acq_rel) > 0);
+    inc_ref();
+}
 
-CPUBuffer::CPUBuffer(void* raw_ptr, dimn_t size)
-    : raw_buffer{ raw_ptr, size, IsOwned}, is_ocl(false)
+CPUBuffer::CPUBuffer(void* raw_ptr, dimn_t size, atomic_t rc)
+    : raw_buffer{raw_ptr, size, rc},
+      flags()
 {}
 
-CPUBuffer::CPUBuffer(const void* raw_ptr, dimn_t size)
-    : raw_buffer { const_cast<void*>(raw_ptr), size, IsConst }, is_ocl(false)
+CPUBuffer::CPUBuffer(const void* raw_ptr, dimn_t size, atomic_t rc)
+    : raw_buffer{const_cast<void*>(raw_ptr), size, rc},
+      flags(IsConst)
 {}
 
-CPUBuffer::~CPUBuffer() {
-    if (is_ocl) {
-        ocl_buffer.~OCLBuffer();
-    }
-
-    if (raw_buffer.flags & IsOwned) {
-
+CPUBuffer::~CPUBuffer()
+{
+    if (dec_ref() == 1) {
+        RPY_DBG_ASSERT(raw_ref_count(std::memory_order_acq_rel) == 0);
+        CPUDeviceHandle::get()->raw_free(raw_buffer.ptr, raw_buffer.size);
+        raw_buffer.ptr = nullptr;
+        raw_buffer.size = 0;
     }
-
 }
 
 BufferMode CPUBuffer::mode() const
 {
-    if (is_ocl) {
-        return ocl_buffer.mode();
-    }
-
-    if (raw_buffer.flags & IsConst) {
-        return BufferMode::Read;
-    }
+    if (flags & IsConst) { return BufferMode::Read; }
     return BufferMode::ReadWrite;
 }
-dimn_t CPUBuffer::size() const {
-    if (is_ocl) {
-        return ocl_buffer.size();
-    }
+dimn_t CPUBuffer::size() const
+{
     return raw_buffer.size;
 }
-void* CPUBuffer::ptr() {
-    if (is_ocl) {
-        return ocl_buffer.ptr();
-    }
-    return raw_buffer.ptr;
-}
-std::unique_ptr<rpy::devices::dtl::InterfaceBase> CPUBuffer::clone() const
+void* CPUBuffer::ptr()
 {
-    return nullptr;
+    return raw_buffer.ptr;
 }
-Device CPUBuffer::device() const noexcept {
-    return CPUDeviceHandle::get();
+std::unique_ptr<rpy::devices::dtl::InterfaceBase> CPUBuffer::clone() const {
+    return std::unique_ptr<CPUBuffer>(new CPUBuffer(raw_buffer, flags));
 }
-bool CPUBuffer::owning() const noexcept {
-    return is_ocl ? false : (raw_buffer.flags & IsOwned != 0);
+Device CPUBuffer::device() const noexcept { return CPUDeviceHandle::get(); }
+dimn_t CPUBuffer::ref_count() const noexcept
+{
+    return raw_ref_count();
 }
-
diff --git a/device/src/cpu/cpu_buffer.h b/device/src/cpu/cpu_buffer.h
@@ -35,52 +35,74 @@
 #include <roughpy/device/buffer.h>
 
 #include "cpu_decls.h"
-#include "opencl/ocl_headers.h"
 #include "opencl/ocl_buffer.h"
+#include "opencl/ocl_headers.h"
+
+#include <atomic>
 
 namespace rpy {
 namespace devices {
 
 class CPUBuffer : public BufferInterface
 {
+public:
+    using atomic_t = std::atomic_size_t*;
+private:
+
+
     enum Flags
     {
         IsConst = 1,
-        IsOwned = 2
     };
 
     struct RawBuffer {
         void* ptr = nullptr;
         dimn_t size = 0;
-        Flags flags;
+        atomic_t ref_count;
     };
 
-    union {
-        OCLBuffer ocl_buffer;
-        RawBuffer raw_buffer;
-    };
+    RawBuffer raw_buffer;
+    Flags flags;
+
 
-    bool is_ocl = false;
+    inline dimn_t inc_ref() noexcept
+    {
+        return raw_buffer.ref_count->fetch_add(1, std::memory_order_relaxed);
+    }
+
+    inline dimn_t dec_ref() noexcept
+    {
+        RPY_DBG_ASSERT(raw_ref_count(std::memory_order_acq_rel) > 0);
+        return raw_buffer.ref_count->fetch_sub(1, std::memory_order_acq_rel);
+    }
+
+    inline dimn_t raw_ref_count(std::memory_order order = std::memory_order_relaxed
+    )
+            const noexcept
+    {
+        return raw_buffer.ref_count->load(order);
+    }
+
+    CPUBuffer(RawBuffer raw, Flags arg_flags);
 
 public:
 
-    CPUBuffer(cl_mem buffer, CPUDevice dev);
-    CPUBuffer(void* raw_ptr, dimn_t size);
-    CPUBuffer(const void* raw_ptr, dimn_t size);
+
+    CPUBuffer(void* raw_ptr, dimn_t size, atomic_t rc);
+    CPUBuffer(const void* raw_ptr, dimn_t size, atomic_t rc);
 
     ~CPUBuffer();
 
+    dimn_t ref_count() const noexcept override;
     std::unique_ptr<dtl::InterfaceBase> clone() const override;
     Device device() const noexcept override;
 
-    virtual bool owning() const noexcept;
     BufferMode mode() const override;
     dimn_t size() const override;
     void* ptr() override;
-
 };
 
-}// namespace device
+}// namespace devices
 }// namespace rpy
 
 #endif// ROUGHPY_DEVICE_SRC_CPUDEVICE_CPU_BUFFER_H_