Exposing memory view object for buffers / mapping

kecho · Dec 30, 2023 · e428fbb · e428fbb
1 parent 27a4c97
commit e428fbb
Show file tree

Hide file tree

Showing 4 changed files with 63 additions and 1 deletion.
diff --git a/Source/pymodules/gpu/ResourceDownloadRequest.cpp b/Source/pymodules/gpu/ResourceDownloadRequest.cpp
@@ -236,7 +236,7 @@ PyObject* dataAsByteArray(PyObject* self, PyObject* vargs, PyObject* kwds)
         return nullptr;
 
     request.rowBytesPitchObject = PyLong_FromLongLong((long)status.rowPitch);
-    request.dataAsByteArray = PyBytes_FromStringAndSize((const char*)status.downloadPtr, status.downloadByteSize);
+    request.dataAsByteArray = PyMemoryView_FromMemory((char*)(status.downloadPtr), status.downloadByteSize, 0);
 
     Py_INCREF(request.dataAsByteArray);
     return request.dataAsByteArray;

diff --git a/Source/pymodules/gpu/Resources.cpp b/Source/pymodules/gpu/Resources.cpp
@@ -5,12 +5,14 @@
 #include "structmember.h"
 #include <coalpy.render/IDevice.h>
 #include <coalpy.texture/ITextureLoader.h>
+#include "HelperMacros.h"
 
 namespace coalpy
 {
 namespace gpu
 {
 
+
 bool validateEnum(ModuleState& state, int value, int count, const char* name, const char* typeName)
 {
     if (value > count || value < 0)
@@ -24,6 +26,18 @@ bool validateEnum(ModuleState& state, int value, int count, const char* name, co
     return true;
 }
 
+namespace methods
+{
+    #include "bindings/MethodDecl.h"
+    #include "bindings/Buffer.inl"
+}
+
+static PyMethodDef g_bufferMethods[] = {
+    #include "bindings/MethodDef.h"
+    #include "bindings/Buffer.inl"
+    FN_END
+};
+
 void Buffer::constructType(CoalpyTypeObject& o)
 {
     auto& t = o.pyObj;
@@ -45,6 +59,7 @@ void Buffer::constructType(CoalpyTypeObject& o)
     t.tp_flags = Py_TPFLAGS_DEFAULT;
     t.tp_new = PyType_GenericNew;
     t.tp_init = Buffer::init;
+    t.tp_methods = g_bufferMethods;
     t.tp_dealloc = Buffer::destroy;
 }
 
@@ -102,10 +117,49 @@ void Buffer::destroy(PyObject* self)
     if (buffer->owned)
         moduleState.device().release(buffer->buffer);
 
+    Py_XDECREF(buffer->mappedMemory);
     buffer->~Buffer();
     Py_TYPE(self)->tp_free(self);
 }
 
+namespace methods
+{
+    PyObject* bufferMappedMemory(PyObject* self, PyObject* vargs, PyObject* kwds)
+    {
+        Buffer* bufferObj = (Buffer*)self;
+        ModuleState& state = parentModule(self);
+        if (!state.checkValidDevice())
+        {
+            PyErr_SetString(state.exObj(), "Current device is invalid.");
+            return nullptr;
+        }
+
+        if (!bufferObj->buffer.valid())
+        {
+            PyErr_SetString(state.exObj(), "Current buffer is invalid.");
+            return nullptr;
+        }
+
+        if (bufferObj->mappedMemory != nullptr)
+        {
+            Py_INCREF(bufferObj->mappedMemory);
+            return bufferObj->mappedMemory;
+        }
+
+        void* mappedMemory = state.device().mappedMemory(bufferObj->buffer);
+        if (mappedMemory == nullptr)
+            Py_RETURN_NONE;
+
+        render::ResourceMemoryInfo memInfo = {};
+        state.device().getResourceMemoryInfo(bufferObj->buffer, memInfo);
+
+        PyObject* memoryView = PyMemoryView_FromMemory((char*)mappedMemory, memInfo.byteSize, 0);
+        bufferObj->mappedMemory = memoryView;
+        Py_INCREF(bufferObj->mappedMemory);
+        return bufferObj->mappedMemory;
+    }
+}
+
 void Texture::constructType(CoalpyTypeObject& o)
 {
     auto& t = o.pyObj;

diff --git a/Source/pymodules/gpu/Resources.h b/Source/pymodules/gpu/Resources.h
@@ -35,6 +35,7 @@ struct Buffer
     bool isAppendConsume = false;
     bool owned = true;
     render::Buffer buffer;
+    PyObject* mappedMemory = nullptr;
 
     //Functions
     static const TypeId s_typeId = TypeId::Buffer;

diff --git a/Source/pymodules/gpu/bindings/Buffer.inl b/Source/pymodules/gpu/bindings/Buffer.inl
@@ -0,0 +1,7 @@
+COALPY_FN(mappedMemory, bufferMappedMemory, R"(
+    Returns:
+        Returns a memory view object with the mapped memory.
+        NOTE: only works if the buffer is created with BufferUsage.Upload flag.
+)")
+
+#undef COALPY_FN