From 2198a9e001c0c1fe55895e5e60f8455f3f1390cd Mon Sep 17 00:00:00 2001
From: leegao <lg342@cornell.edu>
Date: Thu, 14 Aug 2025 18:42:28 +1000
Subject: [PATCH] Fix BC6h staging buffer size calculation to use the correct
 texel stride #131

---
 src/vulkan/wrapper/artifacts.cpp    | 17 ++++++++++++++---
 src/vulkan/wrapper/wrapper_device.c | 12 ++++++------
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/src/vulkan/wrapper/artifacts.cpp b/src/vulkan/wrapper/artifacts.cpp
index 2d24e714b1a..def2528176b 100644
--- a/src/vulkan/wrapper/artifacts.cpp
+++ b/src/vulkan/wrapper/artifacts.cpp
@@ -55,7 +55,7 @@ void RecordBCnArtifacts(struct wrapper_device* device, VkFormat original_format,
         };
         VkResult result = WCHECK(MapMemory2KHR((VkDevice) device, &mapInfoSrc, &srcData));
         if (result != VK_SUCCESS) {
-            WLOGE("ERROR: Failed to map srcBuffer memory: %d", result);
+            WLOGE("Failed to map srcBuffer memory: %d", result);
         } else {
             int w = region->bufferRowLength ? region->bufferRowLength : region->imageExtent.width;
             int h = region->bufferImageHeight ? region->bufferImageHeight : region->imageExtent.height;
@@ -75,6 +75,17 @@ void RecordBCnArtifacts(struct wrapper_device* device, VkFormat original_format,
     } else if (wbuf->memory == VK_NULL_HANDLE) {
         WLOGE("dstBuffer not bound, skipping (decode_id=%d)", decode_id);
     } else {
+        // Invalidate the memory to be visible to the host in case this was decoded via the compute shader
+        VkMappedMemoryRange flushRange = {
+            .sType = VK_STRUCTURE_TYPE_MAPPED_MEMORY_RANGE,
+            .memory = wbuf->memory,
+            .offset = 0,
+            .size = VK_WHOLE_SIZE,
+        };
+        VkResult result = WCHECK(InvalidateMappedMemoryRanges((VkDevice) device, 1, &flushRange));
+        if (result != VK_SUCCESS) {
+            WLOGE("Failed to flush dstBuffer memory: %d", result);
+        }
         void* dstData;
         VkMemoryMapInfoKHR mapInfoSrc = {
             .sType = VK_STRUCTURE_TYPE_MEMORY_MAP_INFO_KHR,
@@ -82,9 +93,9 @@ void RecordBCnArtifacts(struct wrapper_device* device, VkFormat original_format,
             .offset = 0,
             .size = VK_WHOLE_SIZE,
         };
-        VkResult result = WCHECK(MapMemory2KHR((VkDevice) device, &mapInfoSrc, &dstData));
+        result = WCHECK(MapMemory2KHR((VkDevice) device, &mapInfoSrc, &dstData));
         if (result != VK_SUCCESS) {
-            WLOGE("ERROR: Failed to map dstBuffer memory: %d", result);
+            WLOGE("Failed to map dstBuffer memory: %d", result);
         } else {
             int block_size = get_bc_target_size(device->physical, original_format);
             auto fd = open_log_file("_dst.dat", original_format, decode_id, "wb");
diff --git a/src/vulkan/wrapper/wrapper_device.c b/src/vulkan/wrapper/wrapper_device.c
index 398da6e8148..e24a6c6c1ea 100644
--- a/src/vulkan/wrapper/wrapper_device.c
+++ b/src/vulkan/wrapper/wrapper_device.c
@@ -1392,7 +1392,7 @@ static VkResult HostSideDecompression(
       wrapper_buffer* srcBuffer,
       VkDeviceMemory dstMemory,
       const VkBufferImageCopy* region,
-      VkFormat dstFormat
+      VkFormat in_format
 ) {
    // Map the source buffer
    void* srcData;
@@ -1413,9 +1413,8 @@ static VkResult HostSideDecompression(
    VkMemoryMapInfoKHR mapInfoDst = {
       .sType = VK_STRUCTURE_TYPE_MEMORY_MAP_INFO_KHR,
       .memory = dstMemory,
-      .offset = 0, // Assuming dstMemory is large enough to hold the decompressed
-      // TODO: FIX THIS for non rgba8 data
-      .size = region->imageExtent.width * region->imageExtent.height * 4, // 4 bytes per pixel for RGBA
+      .offset = 0,
+      .size = region->imageExtent.width * region->imageExtent.height * get_bc_target_size(_device->physical, in_format),
    };
    result = WCHECK(MapMemory2KHR((VkDevice) _device, &mapInfoDst, &dstData));
    if (result != VK_SUCCESS) {
@@ -1424,7 +1423,7 @@ static VkResult HostSideDecompression(
 
    // Decompress the data from srcData to dstData
    BCnDecompression(
-      dstFormat,
+      in_format,
       srcData,
       dstData,
       region
@@ -1838,7 +1837,7 @@ WRAPPER_CmdCopyBufferToImage(VkCommandBuffer commandBuffer,
       if (!use_image_view) {
          result = CreateStagingBuffer(
             _device, 
-            region->imageExtent.width * region->imageExtent.height * region->imageExtent.depth * 4,
+            region->imageExtent.width * region->imageExtent.height * get_bc_target_size(_device->physical, wimg->original_format),
             VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
             VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
             &stagingBuffer, &stagingBufferMemory);
@@ -1848,6 +1847,7 @@ WRAPPER_CmdCopyBufferToImage(VkCommandBuffer commandBuffer,
          }
 
          // TODO: track these using the fence/semaphore for the submitted queue
+         // This should be safe to do now with the queue synchronization, check with vvl
          WLOGD("Tracking new staging buffer for image %d", wimg->obj_id);
          TRACK_STAGING(_device, BUFFER, stagingBuffer, wimg);
          TRACK_STAGING(_device, DEVICE_MEMORY, stagingBufferMemory, wimg);