align readback copies between backends

Author: EmilioLaiso

include/Support/Pipeline.h

@@ -157,6 +157,19 @@ struct CPUBuffer {
       return Stride;
     return getSingleElementSize() * Channels;
   }
+
+  // The natural per-row byte size of this buffer when interpreted as a 2D
+  // image (no padding).
+  uint32_t getImageRowBytes() const {
+    return OutputProps.Width * getElementSize();
+  }
+
+  // Copy a 2D image readback from a GPU mapping into Data[0]. The host
+  // buffer is tightly packed with top-left origin. SrcRowPitch is the
+  // source's per-row stride in bytes; pass `getImageRowBytes()` when the
+  // source is tightly packed (Vulkan / Metal), or the GPU's reported pitch
+  // when the source has row padding (e.g., D3D12's 256-byte aligned rows).
+  void copyFromTexture(const void *Src, size_t SrcRowPitch);
 };
 
 struct Result {

lib/API/DX/Device.cpp

@@ -1581,18 +1581,14 @@ class DXDevice : public offloadtest::Device {
     if (!IS.RTReadback)
       return llvm::Error::success();
 
-    // Map readback and copy into host buffer, accounting for row pitch and
-    // flipping vertical orientation. DirectX render target origin is top-left,
-    // while our image writer expects bottom-left.
-    const CPUBuffer &B = *P.Bindings.RTargetBufferPtr;
     void *Mapped = nullptr;
     auto &Readback = llvm::cast<DXBuffer>(*IS.RTReadback);
     if (auto Err = HR::toError(Readback.Buffer->Map(0, nullptr, &Mapped),
                                "Failed to map render target readback"))
       return Err;
 
-    // Query the copy footprint to get the actual padded row pitch used by the
-    // copy operation.
+    // Query the copy footprint to get the actual padded row pitch used by
+    // the copy operation (D3D12 requires 256-byte aligned rows).
     auto &RT = llvm::cast<DXTexture>(*IS.RT);
     const D3D12_RESOURCE_DESC RTDesc = RT.Resource->GetDesc();
     D3D12_PLACED_SUBRESOURCE_FOOTPRINT Placed = {};
@@ -1602,23 +1598,8 @@ class DXDevice : public offloadtest::Device {
     Device->GetCopyableFootprints(&RTDesc, 0u, 1u, 0u, &Placed, &NumRows,
                                   &RowSizeInBytes, &TotalBytes);
 
-    const uint32_t RowPitch = Placed.Footprint.RowPitch;
-    const uint32_t RowBytes =
-        static_cast<uint32_t>(B.getElementSize() * B.OutputProps.Width);
-    const uint32_t Height = static_cast<uint32_t>(B.OutputProps.Height);
-
-    const uint8_t *SrcBase = reinterpret_cast<uint8_t *>(Mapped);
-    uint8_t *DstBase =
-        reinterpret_cast<uint8_t *>(P.Bindings.RTargetBufferPtr->Data[0].get());
-
-    // Copy rows in reverse order.
-    for (uint32_t Y = 0; Y < Height; ++Y) {
-      const uint8_t *SrcRow = SrcBase + static_cast<size_t>(Y) * RowPitch;
-      uint8_t *DstRow =
-          DstBase + static_cast<size_t>(Height - 1 - Y) * RowBytes;
-      memcpy(DstRow, SrcRow, RowBytes);
-    }
-
+    P.Bindings.RTargetBufferPtr->copyFromTexture(Mapped,
+                                                 Placed.Footprint.RowPitch);
     Readback.Buffer->Unmap(0, nullptr);
     return llvm::Error::success();
   }

lib/API/MTL/MTLDevice.cpp

@@ -705,23 +705,9 @@ class MTLDevice : public offloadtest::Device {
       }
     }
     if (P.isGraphics()) {
-      CPUBuffer *RTarget = P.Bindings.RTargetBufferPtr;
-      const uint64_t Width = RTarget->OutputProps.Width;
-      const uint64_t Height = RTarget->OutputProps.Height;
-      const size_t ElemSize = RTarget->getElementSize();
-      const size_t RowBytes = Width * ElemSize;
-
-      // Read from the readback buffer. The blit copied the texture data in
-      // GPU layout order, so we flip rows here to produce an upright image.
       auto &FBReadback = llvm::cast<MTLBuffer>(*IS.FrameBufferReadback);
-      const unsigned char *Src =
-          reinterpret_cast<const unsigned char *>(FBReadback.Buf->contents());
-      unsigned char *Buf =
-          reinterpret_cast<unsigned char *>(RTarget->Data[0].get());
-      for (uint64_t R = 0; R < Height; ++R) {
-        const uint64_t SrcRow = (Height - 1) - R;
-        memcpy(Buf + R * RowBytes, Src + SrcRow * RowBytes, RowBytes);
-      }
+      auto *RT = P.Bindings.RTargetBufferPtr;
+      RT->copyFromTexture(FBReadback.Buf->contents(), RT->getImageRowBytes());
     }
     return llvm::Error::success();
   }

lib/API/VK/Device.cpp

@@ -2448,8 +2448,8 @@ class VulkanDevice : public offloadtest::Device {
       vkMapMemory(Device, Readback.Memory, 0, VK_WHOLE_SIZE, 0, &Mapped);
       vkInvalidateMappedMemoryRanges(Device, 1, &Range);
 
-      const CPUBuffer &B = *P.Bindings.RTargetBufferPtr;
-      memcpy(B.Data[0].get(), Mapped, B.size());
+      auto *RT = P.Bindings.RTargetBufferPtr;
+      RT->copyFromTexture(Mapped, RT->getImageRowBytes());
       vkUnmapMemory(Device, Readback.Memory);
     }
     return llvm::Error::success();

lib/Image/Image.cpp

@@ -159,10 +159,10 @@ static llvm::Error writePNGImpl(ImageRef Img, llvm::StringRef OutputPath) {
   png_bytepp Rows =
       (png_bytepp)png_malloc(PNG, Img.getHeight() * sizeof(png_bytep));
   const uint64_t RowSize = Img.getWidth() * Img.getChannels() * Img.getDepth();
-  // Step one row back from the end
-  const uint8_t *Row = reinterpret_cast<const uint8_t *>(Img.data()) +
-                       (RowSize * Img.getHeight()) - RowSize;
-  for (uint32_t I = 0; I < Img.getHeight(); ++I, Row -= RowSize)
+  // The host buffer has top-left origin (matching every backend's readback
+  // and the standard PNG row order), so we walk forward.
+  const uint8_t *Row = reinterpret_cast<const uint8_t *>(Img.data());
+  for (uint32_t I = 0; I < Img.getHeight(); ++I, Row += RowSize)
     Rows[I] = const_cast<png_bytep>(Row);
 
   png_write_image(PNG, Rows);

lib/Support/Pipeline.cpp

@@ -26,6 +26,21 @@ void PushConstantBlock::getContent(
     memcpy(Output.data() + V.OffsetInBytes, V.Data.data(), V.Data.size());
 }
 
+void CPUBuffer::copyFromTexture(const void *Src, size_t SrcRowPitch) {
+  const uint32_t Height = OutputProps.Height;
+  const uint32_t RowBytes = getImageRowBytes();
+  assert(SrcRowPitch >= RowBytes && "Source row pitch is smaller than image");
+  uint8_t *Dst = reinterpret_cast<uint8_t *>(Data[0].get());
+  if (SrcRowPitch == RowBytes) {
+    memcpy(Dst, Src, static_cast<size_t>(Height) * RowBytes);
+    return;
+  }
+  const uint8_t *S = reinterpret_cast<const uint8_t *>(Src);
+  for (uint32_t Y = 0; Y < Height; ++Y)
+    memcpy(Dst + static_cast<size_t>(Y) * RowBytes,
+           S + static_cast<size_t>(Y) * SrcRowPitch, RowBytes);
+}
+
 uint32_t PushConstantBlock::size() const {
   uint32_t Size = 0;
   for (const PushConstantValue &V : Values)