Capture and serialise MTLBuffer contents

MTLBuffer initial state chunks.
CPU modified buffers chunks (MTLBuffer_InternalModifyCPUContents) stored in the command buffer record.

renderdoc/driver/metal/metal_buffer.cpp

@@ -110,4 +110,43 @@ void WrappedMTLBuffer::didModifyRange(NS::Range &range)
   }
 }
 
+template <typename SerialiserType>
+bool WrappedMTLBuffer::Serialise_InternalModifyCPUContents(SerialiserType &ser, uint64_t start,
+                                                           uint64_t end, MetalBufferInfo *bufInfo)
+{
+  SERIALISE_ELEMENT_LOCAL(Buffer, this).Important();
+  SERIALISE_ELEMENT(start).Important();
+  uint64_t size = end - start;
+  SERIALISE_ELEMENT(size).Important();
+  byte *pData = NULL;
+  if(ser.IsWriting())
+  {
+    pData = (byte *)Unwrap(this)->contents() + start;
+  }
+  if(IsReplayingAndReading())
+  {
+    pData = (byte *)Unwrap(Buffer)->contents() + start;
+  }
+
+  // serialise directly using buffer memory
+  ser.Serialise("data"_lit, pData, size, SerialiserFlags::NoFlags);
+
+  if(IsCaptureMode(m_State))
+  {
+    // update the base snapshot from the serialised data
+    size_t offs = size_t(ser.GetWriter()->GetOffset() - size);
+    const byte *serialisedData = ser.GetWriter()->GetData() + offs;
+    if(bufInfo->baseSnapshot.isEmpty())
+      bufInfo->baseSnapshot.resize(bufInfo->length);
+    RDCASSERTEQUAL(bufInfo->baseSnapshot.size(), bufInfo->length);
+    memcpy(bufInfo->baseSnapshot.data() + start, serialisedData, size);
+  }
+
+  SERIALISE_CHECK_READ_ERRORS();
+
+  return true;
+}
+
 INSTANTIATE_FUNCTION_SERIALISED(WrappedMTLBuffer, void, didModifyRange, NS::Range &);
+INSTANTIATE_FUNCTION_SERIALISED(WrappedMTLBuffer, void, InternalModifyCPUContents, uint64_t,
+                                uint64_t, MetalBufferInfo *);

renderdoc/driver/metal/metal_buffer.h

@@ -36,6 +36,9 @@ public:
   void *contents();
 
   DECLARE_FUNCTION_SERIALISED(void, didModifyRange, NS::Range &range);
+  template <typename SerialiserType>
+  bool Serialise_InternalModifyCPUContents(SerialiserType &ser, uint64_t start, uint64_t end,
+                                           MetalBufferInfo *bufInfo);
 
   enum
   {

renderdoc/driver/metal/metal_core.cpp

@@ -24,6 +24,7 @@
 
 #include "metal_core.h"
 #include "serialise/rdcfile.h"
+#include "metal_buffer.h"
 #include "metal_command_buffer.h"
 #include "metal_device.h"
 #include "metal_texture.h"
@@ -395,7 +396,6 @@ void WrappedMTLDevice::CaptureCmdBufSubmit(MetalResourceRecord *record)
   WrappedMTLCommandBuffer *commandBuffer = (WrappedMTLCommandBuffer *)(record->m_Resource);
   if(IsActiveCapturing(m_State))
   {
-    Chunk *chunk = NULL;
     std::unordered_set<ResourceId> refIDs;
     // The record will get deleted at the end of active frame capture
     record->AddRef();
@@ -408,12 +408,44 @@ void WrappedMTLDevice::CaptureCmdBufSubmit(MetalResourceRecord *record)
       MetalResourceRecord *refRecord = GetResourceManager()->GetResourceRecord(id);
       if(refRecord->m_Type == eResBuffer)
       {
-        // TODO: capture CPU modified buffers
+        MetalBufferInfo *bufInfo = refRecord->bufInfo;
+        if(bufInfo->storageMode == MTL::StorageModeShared)
+        {
+          size_t diffStart = 0;
+          size_t diffEnd = bufInfo->length;
+          bool foundDifference = true;
+          if(!bufInfo->baseSnapshot.isEmpty())
+          {
+            foundDifference = FindDiffRange(bufInfo->data, bufInfo->baseSnapshot.data(),
+                                            bufInfo->length, diffStart, diffEnd);
+            if(diffEnd <= diffStart)
+              foundDifference = false;
+          }
+
+          if(foundDifference)
+          {
+            if(bufInfo->data == NULL)
+            {
+              RDCERR("Writing buffer memory %s that is NULL", ToStr(id).c_str());
+              continue;
+            }
+            Chunk *chunk = NULL;
+            {
+              CACHE_THREAD_SERIALISER();
+              SCOPED_SERIALISE_CHUNK(MetalChunk::MTLBuffer_InternalModifyCPUContents);
+              ((WrappedMTLBuffer *)refRecord->m_Resource)
+                  ->Serialise_InternalModifyCPUContents(ser, diffStart, diffEnd, bufInfo);
+              chunk = scope.Get();
+            }
+            record->AddChunk(chunk);
+          }
+        }
       }
     }
     record->MarkResourceFrameReferenced(GetResID(commandBuffer->GetCommandQueue()), eFrameRef_Read);
     // pull in frame refs from this command buffer
     record->AddResourceReferences(GetResourceManager());
+    Chunk *chunk = NULL;
     {
       CACHE_THREAD_SERIALISER();
       SCOPED_SERIALISE_CHUNK(MetalChunk::MTLCommandBuffer_commit);

renderdoc/driver/metal/metal_device.cpp

@@ -643,6 +643,21 @@ WrappedMTLBuffer *WrappedMTLDevice::Common_NewBuffer(bool withBytes, const void
 
     MetalResourceRecord *record = GetResourceManager()->AddResourceRecord(wrappedMTLBuffer);
     record->AddChunk(chunk);
+
+    MTL::StorageMode mode = realMTLBuffer->storageMode();
+    record->bufInfo = new MetalBufferInfo(mode);
+
+    // Create CPU side tracking info for CPU shared buffers
+    if(mode == MTL::StorageModeShared)
+    {
+      record->bufInfo->data = (byte *)realMTLBuffer->contents();
+      record->bufInfo->length = realMTLBuffer->length();
+    }
+    // Snapshot GPU only buffers
+    else if(mode == MTL::StorageModePrivate)
+    {
+      GetResourceManager()->MarkDirtyResource(id);
+    }
   }
   else
   {

renderdoc/driver/metal/metal_init_state.cpp

@@ -22,14 +22,88 @@
  * THE SOFTWARE.
  ******************************************************************************/
 
+#include "metal_buffer.h"
 #include "metal_common.h"
 #include "metal_device.h"
 
+static rdcliteral NameOfType(MetalResourceType type)
+{
+  switch(type)
+  {
+    case eResBuffer: return "MTLBuffer"_lit;
+    default: break;
+  }
+  return "MTLResource"_lit;
+}
+
 bool WrappedMTLDevice::Prepare_InitialState(WrappedMTLObject *res)
 {
   ResourceId id = GetResourceManager()->GetID(res);
 
   MetalResourceType type = res->m_Record->m_Type;
+
+  if(type == eResBuffer)
+  {
+    WrappedMTLBuffer *buffer = (WrappedMTLBuffer *)res;
+    MTL::Buffer *mtlBuffer = Unwrap(buffer);
+    MTL::Buffer *mtlSharedBuffer = NULL;
+    MTL::StorageMode storageMode = mtlBuffer->storageMode();
+    size_t len = mtlBuffer->length();
+    byte *data = NULL;
+    if(storageMode == MTL::StorageModeShared)
+    {
+      // MTLStorageModeShared buffers are automatically synchronized
+      data = (byte *)mtlBuffer->contents();
+    }
+    else if(storageMode == MTL::StorageModeManaged)
+    {
+      // MTLStorageModeManaged buffers need to call MTLBlitCommandEncoder::synchronizeResource
+      MTL::CommandBuffer *mtlCommandBuffer = m_mtlCommandQueue->commandBuffer();
+      MTL::BlitCommandEncoder *mtlBlitEncoder = mtlCommandBuffer->blitCommandEncoder();
+      mtlBlitEncoder->synchronizeResource(mtlBuffer);
+      mtlBlitEncoder->endEncoding();
+      mtlCommandBuffer->commit();
+      mtlCommandBuffer->waitUntilCompleted();
+      data = (byte *)mtlBuffer->contents();
+    }
+    else if(storageMode == MTL::StorageModePrivate)
+    {
+      // TODO: postpone readback until data is required
+      // TODO: batch readback for multiple resources to avoid sync per resource
+      // MTLStorageModePrivate buffer need to copy into a temporary MTLStorageModeShared buffer
+      mtlSharedBuffer = Unwrap(this)->newBuffer(len, MTL::ResourceStorageModeShared);
+      MTL::CommandBuffer *mtlCommandBuffer = m_mtlCommandQueue->commandBuffer();
+      MTL::BlitCommandEncoder *mtlBlitEncoder = mtlCommandBuffer->blitCommandEncoder();
+      mtlBlitEncoder->copyFromBuffer(mtlBuffer, 0, mtlSharedBuffer, 0, len);
+      mtlBlitEncoder->endEncoding();
+      mtlCommandBuffer->commit();
+      mtlCommandBuffer->waitUntilCompleted();
+      data = (byte *)mtlSharedBuffer->contents();
+    }
+    else
+    {
+      RDCERR("Unhandled buffer storage mode 0x%X", storageMode);
+    }
+
+    bytebuf bufferContents(data, len);
+    MetalInitialContents initialContents(type, bufferContents);
+    GetResourceManager()->SetInitialContents(id, initialContents);
+    if(mtlSharedBuffer)
+    {
+      mtlSharedBuffer->release();
+    }
+    if(storageMode == MTL::StorageModeShared)
+    {
+      // Set the base snapshot to match the initial contents
+      MetalBufferInfo *bufInfo = res->m_Record->bufInfo;
+      if(bufInfo->baseSnapshot.isEmpty())
+        bufInfo->baseSnapshot.resize(len);
+      RDCASSERTEQUAL(bufInfo->baseSnapshot.size(), len);
+      memcpy(bufInfo->baseSnapshot.data(), bufferContents.data(), len);
+    }
+    return true;
+  }
+  else
   {
     RDCERR("Unhandled resource type %d", type);
   }
@@ -39,8 +113,16 @@ bool WrappedMTLDevice::Prepare_InitialState(WrappedMTLObject *res)
 
 uint64_t WrappedMTLDevice::GetSize_InitialState(ResourceId id, const MetalInitialContents &initial)
 {
-  METAL_NOT_IMPLEMENTED();
-  return 128;
+  uint64_t ret = 128;
+
+  if(initial.type == eResBuffer)
+  {
+    ret += uint64_t(initial.resourceContents.size() + WriteSerialiser::GetChunkAlignment());
+    return ret;
+  }
+
+  RDCERR("Unhandled resource type %s", ToStr(initial.type).c_str());
+  return 0;
 }
 
 template <typename SerialiserType>
@@ -48,7 +130,29 @@ bool WrappedMTLDevice::Serialise_InitialState(SerialiserType &ser, ResourceId id
                                               MetalResourceRecord *record,
                                               const MetalInitialContents *initial)
 {
-  METAL_NOT_IMPLEMENTED();
+  SERIALISE_ELEMENT_LOCAL(type, initial->type);
+  SERIALISE_ELEMENT(id).TypedAs(NameOfType(type)).Important();
+  if(type == eResBuffer)
+  {
+    SERIALISE_CHECK_READ_ERRORS();
+
+    bytebuf contents;
+    if(ser.IsWriting())
+    {
+      ser.Serialise("Contents"_lit, initial->resourceContents);
+    }
+    else
+    {
+      ser.Serialise("Contents"_lit, contents);
+    }
+
+    if(IsReplayingAndReading())
+    {
+      // TODO: implement RD MTL replay
+    }
+    return true;
+  }
+  RDCERR("Unhandled resource type %d", type);
   return false;
 }
 

renderdoc/driver/metal/metal_manager.h

@@ -42,6 +42,13 @@ struct MetalInitialContents
     type = t;
   }
 
+  MetalInitialContents(MetalResourceType t, bytebuf data)
+  {
+    memset(this, 0, sizeof(*this));
+    type = t;
+    resourceContents = data;
+  }
+
   template <typename Configuration>
   void Free(ResourceManager<Configuration> *rm)
   {

renderdoc/driver/metal/metal_resources.cpp

@@ -56,6 +56,8 @@ MetalResourceRecord::~MetalResourceRecord()
 {
   if(m_Type == eResCommandBuffer)
     SAFE_DELETE(cmdInfo);
+  else if(m_Type == eResBuffer)
+    SAFE_DELETE(bufInfo);
 }
 
 void WrappedMTLObject::AddEvent()

renderdoc/driver/metal/metal_resources.h

@@ -150,6 +150,16 @@ struct MetalCmdBufferRecordingInfo
   bool presented = false;
 };
 
+struct MetalBufferInfo
+{
+  MetalBufferInfo() = delete;
+  MetalBufferInfo(MTL::StorageMode mode) : storageMode(mode), data(NULL), length(0) {}
+  MTL::StorageMode storageMode;
+  bytebuf baseSnapshot;
+  byte *data;
+  size_t length;
+};
+
 struct MetalResourceRecord : public ResourceRecord
 {
 public:
@@ -171,5 +181,6 @@ public:
   {
     void *ptrUnion;                          // for initialisation to NULL
     MetalCmdBufferRecordingInfo *cmdInfo;    // only for command buffers
+    MetalBufferInfo *bufInfo;                // only for buffers
   };
 };