Fix unused size calculation when requested size is greater than the m…

…ax size of the buffer

Reviewed By: EricGriffith

Differential Revision: D50226988

fbshipit-source-id: 9577718c8525c9c47aa5da8efd6b21c87848a1ca

src/igl/vulkan/VulkanStagingDevice.cpp

@@ -168,15 +168,19 @@ VulkanStagingDevice::MemoryRegion VulkanStagingDevice::nextFreeBlock(VkDeviceSiz
   regions_.clear();
 
   // Store the unused size in the deque first...
-  const uint32_t unusedSize = stagingBufferSize_ - requestedAlignedSize;
+  const VkDeviceSize unusedSize =
+      stagingBufferSize_ > requestedAlignedSize ? stagingBufferSize_ - requestedAlignedSize : 0;
   if (unusedSize > 0) {
     const uint32_t unusedOffset = requestedAlignedSize;
     regions_.push_front(
         {unusedOffset, unusedSize, unusedSize, VulkanImmediateCommands::SubmitHandle()});
   }
 
   //... and then return the smallest free region that can hold the requested size
-  return {0, size, requestedAlignedSize, VulkanImmediateCommands::SubmitHandle()};
+  return {0,
+          std::min(size, stagingBufferSize_),
+          std::min(requestedAlignedSize, stagingBufferSize_),
+          VulkanImmediateCommands::SubmitHandle()};
 }
 
 void VulkanStagingDevice::getBufferSubData(VulkanBuffer& buffer,
@@ -195,6 +199,7 @@ void VulkanStagingDevice::getBufferSubData(VulkanBuffer& buffer,
 
   size_t chunkSrcOffset = srcOffset;
   auto* dstData = static_cast<uint8_t*>(data);
+  const size_t bufferSize = size;
 
   while (size) {
     MemoryRegion memoryChunk = nextFreeBlock(size);
@@ -213,7 +218,7 @@ void VulkanStagingDevice::getBufferSubData(VulkanBuffer& buffer,
 
     // Copy data into data
     const uint8_t* src = stagingBuffer_->getMappedPtr() + memoryChunk.offset;
-    checked_memcpy(dstData, size - chunkSrcOffset, src, memoryChunk.alignedSize);
+    checked_memcpy(dstData, bufferSize - chunkSrcOffset, src, memoryChunk.size);
 
     size -= copySize;
     dstData = (uint8_t*)dstData + copySize;