From b0854454fab54b2c641d67116f19ea7bd673eb4a Mon Sep 17 00:00:00 2001
From: Robert Quill <robert.quill@imgtec.com>
Date: Tue, 27 Aug 2024 21:36:50 +0100
Subject: [PATCH] Add a check that the image sizes match when copying images

Add a supporting test and remove an incorrect TODO.
(Size checking for copying from Tensor to Tensor, Tensor to Image and
Image to Tensor is covered by the total size check in
Memory::recordCopyFrom.

Signed-off-by: Robert Quill <robert.quill@imgtec.com>
---
 src/Image.cpp            |  8 ++++++--
 src/Tensor.cpp           |  1 -
 test/TestOpCopyImage.cpp | 21 +++++++++++++++++++++
 3 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/src/Image.cpp b/src/Image.cpp
index 7d4cf56e..6b3efd6a 100644
--- a/src/Image.cpp
+++ b/src/Image.cpp
@@ -93,7 +93,12 @@ Image::recordCopyFrom(const vk::CommandBuffer& commandBuffer,
     layer.layerCount = 1;
     vk::Offset3D offset = { 0, 0, 0 };
 
-    // FIXME: Check the size of the dest and source images match
+    if (this->getWidth() != copyFromImage->getWidth() ||
+        this->getHeight() != copyFromImage->getHeight())
+    {
+        throw std::runtime_error("Kompute Image recordCopyFrom image sizes do not match");
+    }
+
     vk::Extent3D size = { this->mWidth, this->mHeight, 1 };
 
     vk::ImageCopy copyRegion(layer, offset, layer, offset, size);
@@ -133,7 +138,6 @@ Image::recordCopyFrom(const vk::CommandBuffer& commandBuffer,
     layer.layerCount = 1;
     vk::Offset3D offset = { 0, 0, 0 };
 
-    // FIXME: Check the size of the dest and source images match
     vk::Extent3D size = { this->mWidth, this->mHeight, 1 };
 
     vk::BufferImageCopy copyRegion(0, 0, 0, layer, offset, size);
diff --git a/src/Tensor.cpp b/src/Tensor.cpp
index d4374d60..3b5c80d7 100644
--- a/src/Tensor.cpp
+++ b/src/Tensor.cpp
@@ -113,7 +113,6 @@ Tensor::recordCopyFrom(const vk::CommandBuffer& commandBuffer,
     layer.layerCount = 1;
     vk::Offset3D offset = { 0, 0, 0 };
 
-    // FIXME: Check the size of the dest and source images match
     vk::Extent3D size = { copyFromImage->getWidth(),
                           copyFromImage->getHeight(),
                           1 };
diff --git a/test/TestOpCopyImage.cpp b/test/TestOpCopyImage.cpp
index 045e41d8..33731090 100644
--- a/test/TestOpCopyImage.cpp
+++ b/test/TestOpCopyImage.cpp
@@ -350,3 +350,24 @@ TEST(TestOpCopyImage, CopyImageThroughStorageViaAlgorithmsUninitialisedOutput)
     // Making sure the GPU holds the same vector
     EXPECT_EQ(ImageIn->vector(), ImageOut->vector());
 }
+
+TEST(TestOpCopyImage, CopyDeviceToDeviceImage2DMismatchedSizes)
+{
+    kp::Manager mgr;
+
+    std::vector<float> testVecA;
+    std::vector<float> testVecB;
+
+    for (int i = 0; i < 256; i++) {
+        testVecA.push_back(i);
+        testVecB.push_back(0);
+    }
+
+    std::shared_ptr<kp::ImageT<float>> imageA = mgr.image(testVecA, 16, 1, 1);
+    std::shared_ptr<kp::ImageT<float>> imageB = mgr.image(testVecB, 1, 16, 1);
+
+    EXPECT_TRUE(imageA->isInit());
+    EXPECT_TRUE(imageB->isInit());
+
+    EXPECT_THROW(mgr.sequence()->eval<kp::OpCopy>({ imageA, imageB }), std::runtime_error);
+}