Unify memory block API.

spacetelescope · Dec 23, 2024 · af420b2 · af420b2
1 parent cb496f8
commit af420b2
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 6 deletions.
diff --git a/catkit_core/CudaSharedMemory.h b/catkit_core/CudaSharedMemory.h
@@ -1,6 +1,8 @@
 #ifndef CUDA_SHARED_MEMORY_H
 #define CUDA_SHARED_MEMORY_H
 
+#include "Memory.h"
+
 #include <memory>
 
 #ifdef HAVE_CUDA
@@ -11,7 +13,7 @@ typedef cudaIpcMemHandle_t CudaIpcHandle;
 typedef char CudaIpcHandle[64];
 #endif
 
-class CudaSharedMemory
+class CudaSharedMemory : public Memory
 {
 private:
     CudaSharedMemory(const CudaIpcHandle &ipc_handle, void *device_pointer=nullptr);
@@ -22,7 +24,7 @@ class CudaSharedMemory
     static std::shared_ptr<CudaSharedMemory> Create(size_t num_bytes_in_buffer);
     static std::shared_ptr<CudaSharedMemory> Open(const CudaIpcHandle &ipc_handle);
 
-    void *GetAddress();
+    void *GetAddress(std::size_t offset = 0) override;
 };
 
 #endif // CUDA_SHARED_MEMORY_H
diff --git a/catkit_core/Memory.h b/catkit_core/Memory.h
@@ -0,0 +1,16 @@
+#ifndef MEMORY_H
+#define MEMORY_H
+
+#include <cstddef>
+
+class Memory
+{
+public:
+    virtual ~Memory()
+    {
+    }
+
+    virtual void *GetAddress(std::size_t offset = 0) = 0;
+};
+
+#endif // MEMORY_H
diff --git a/catkit_core/SharedMemory.cpp b/catkit_core/SharedMemory.cpp
@@ -82,7 +82,7 @@ SharedMemory::SharedMemory(const std::string &id, FileObject file, bool is_owner
 		throw std::runtime_error("Something went wrong while mapping shared memory file.");
 }
 
-void *SharedMemory::GetAddress()
+void *SharedMemory::GetAddress(std::size_t offset)
 {
-	return m_Buffer;
+	return static_cast<char *>(m_Buffer) + offset;
 }
diff --git a/catkit_core/SharedMemory.h b/catkit_core/SharedMemory.h
@@ -1,6 +1,8 @@
 #ifndef SHARED_MEMORY_H
 #define SHARED_MEMORY_H
 
+#include "Memory.h"
+
 #include <memory>
 #include <string>
 
@@ -21,7 +23,7 @@
 	typedef int FileObject;
 #endif
 
-class SharedMemory
+class SharedMemory : public Memory
 {
 private:
 	SharedMemory(const std::string &id, FileObject file, bool is_owner);
@@ -32,7 +34,7 @@ class SharedMemory
 	static std::shared_ptr<SharedMemory> Create(const std::string &id, size_t num_bytes_in_buffer);
 	static std::shared_ptr<SharedMemory> Open(const std::string &id);
 
-	void *GetAddress();
+	void *GetAddress(std::size_t offset = 0) override;
 
 private:
 	std::string m_Id;