Take the plane/aspect into account when copying/blitting

Bug: b/132437008
Tests: dEQP-VK.*ycbcr*
Change-Id: Ic09d04c63f38247d0f0d6acbd6969f62aa4405ce
Reviewed-on: https://swiftshader-review.googlesource.com/c/SwiftShader/+/31615
Presubmit-Ready: Nicolas Capens <nicolascapens@google.com>
Tested-by: Nicolas Capens <nicolascapens@google.com>
Reviewed-by: Alexis Hétu <sugoi@google.com>
Kokoro-Presubmit: kokoro <noreply+kokoro@google.com>
diff --git a/src/Device/Blitter.cpp b/src/Device/Blitter.cpp
index a97fc3d..f105190 100644
--- a/src/Device/Blitter.cpp
+++ b/src/Device/Blitter.cpp
@@ -80,7 +80,7 @@
 
 		for(; subresLayers.mipLevel <= lastMipLevel; subresLayers.mipLevel++)
 		{
-			VkExtent3D extent = dest->getMipLevelExtent(subresLayers.mipLevel);
+			VkExtent3D extent = dest->getMipLevelExtent(aspect, subresLayers.mipLevel);
 			if(!renderArea)
 			{
 				area.extent.width = extent.width;
@@ -203,7 +203,7 @@
 		{
 			int rowPitchBytes = dest->rowPitchBytes(aspect, subresLayers.mipLevel);
 			int slicePitchBytes = dest->slicePitchBytes(aspect, subresLayers.mipLevel);
-			VkExtent3D extent = dest->getMipLevelExtent(subresLayers.mipLevel);
+			VkExtent3D extent = dest->getMipLevelExtent(aspect, subresLayers.mipLevel);
 			if(!renderArea)
 			{
 				area.extent.width = extent.width;
@@ -1712,14 +1712,13 @@
 			std::swap(region.dstOffsets[0].y, region.dstOffsets[1].y);
 		}
 
-		VkExtent3D srcExtent = src->getMipLevelExtent(region.srcSubresource.mipLevel);
+		VkImageAspectFlagBits srcAspect = static_cast<VkImageAspectFlagBits>(region.srcSubresource.aspectMask);
+		VkImageAspectFlagBits dstAspect = static_cast<VkImageAspectFlagBits>(region.dstSubresource.aspectMask);
+		VkExtent3D srcExtent = src->getMipLevelExtent(srcAspect, region.srcSubresource.mipLevel);
 
 		int32_t numSlices = (region.srcOffsets[1].z - region.srcOffsets[0].z);
 		ASSERT(numSlices == (region.dstOffsets[1].z - region.dstOffsets[0].z));
 
-		VkImageAspectFlagBits srcAspect = static_cast<VkImageAspectFlagBits>(region.srcSubresource.aspectMask);
-		VkImageAspectFlagBits dstAspect = static_cast<VkImageAspectFlagBits>(region.dstSubresource.aspectMask);
-
 		float widthRatio = static_cast<float>(region.srcOffsets[1].x - region.srcOffsets[0].x) /
 		                   static_cast<float>(region.dstOffsets[1].x - region.dstOffsets[0].x);
 		float heightRatio = static_cast<float>(region.srcOffsets[1].y - region.srcOffsets[0].y) /
@@ -1955,7 +1954,7 @@
 
 		void(*cornerUpdateFunction)(const CubeBorderData *data) = (void(*)(const CubeBorderData*))cornerUpdateRoutine->getEntry();
 
-		VkExtent3D extent = image->getMipLevelExtent(subresourceLayers.mipLevel);
+		VkExtent3D extent = image->getMipLevelExtent(aspect, subresourceLayers.mipLevel);
 		CubeBorderData data =
 		{
 			image->getTexelPointer({ 0, 0, 0 }, posX),
@@ -1994,7 +1993,7 @@
 		int bytes = image->getFormat(aspect).bytes();
 		int pitchB = image->rowPitchBytes(aspect, srcSubresourceLayers.mipLevel);
 
-		VkExtent3D extent = image->getMipLevelExtent(srcSubresourceLayers.mipLevel);
+		VkExtent3D extent = image->getMipLevelExtent(aspect, srcSubresourceLayers.mipLevel);
 		int w = extent.width;
 		int h = extent.height;
 		if(w != h)
diff --git a/src/Vulkan/VkImage.cpp b/src/Vulkan/VkImage.cpp
index 1d434cb..fb9c43c 100644
--- a/src/Vulkan/VkImage.cpp
+++ b/src/Vulkan/VkImage.cpp
@@ -181,8 +181,8 @@
 	int dstRowPitchBytes = dst->rowPitchBytes(dstAspect, pRegion.dstSubresource.mipLevel);
 	int dstSlicePitchBytes = dst->slicePitchBytes(dstAspect, pRegion.dstSubresource.mipLevel);
 
-	VkExtent3D srcExtent = getMipLevelExtent(pRegion.srcSubresource.mipLevel);
-	VkExtent3D dstExtent = dst->getMipLevelExtent(pRegion.dstSubresource.mipLevel);
+	VkExtent3D srcExtent = getMipLevelExtent(srcAspect, pRegion.srcSubresource.mipLevel);
+	VkExtent3D dstExtent = dst->getMipLevelExtent(dstAspect, pRegion.dstSubresource.mipLevel);
 	VkExtent3D copyExtent = imageExtentInBlocks(pRegion.extent, srcAspect);
 
 	bool isSinglePlane = (copyExtent.depth == 1);
@@ -259,15 +259,21 @@
 
 void Image::copy(VkBuffer buf, const VkBufferImageCopy& region, bool bufferIsSource)
 {
-	if(!((region.imageSubresource.aspectMask == VK_IMAGE_ASPECT_COLOR_BIT) ||
-	     (region.imageSubresource.aspectMask == VK_IMAGE_ASPECT_DEPTH_BIT) ||
-	     (region.imageSubresource.aspectMask == VK_IMAGE_ASPECT_STENCIL_BIT)))
+	switch(region.imageSubresource.aspectMask)
 	{
-		UNIMPLEMENTED("imageSubresource");
+	case VK_IMAGE_ASPECT_COLOR_BIT:
+	case VK_IMAGE_ASPECT_DEPTH_BIT:
+	case VK_IMAGE_ASPECT_STENCIL_BIT:
+	case VK_IMAGE_ASPECT_PLANE_0_BIT:
+	case VK_IMAGE_ASPECT_PLANE_1_BIT:
+	case VK_IMAGE_ASPECT_PLANE_2_BIT:
+		break;
+	default:
+		UNSUPPORTED("aspectMask %x", int(region.imageSubresource.aspectMask));
+		break;
 	}
 
 	auto aspect = static_cast<VkImageAspectFlagBits>(region.imageSubresource.aspectMask);
-
 	Format copyFormat = getFormat(aspect);
 
 	VkExtent3D imageExtent = imageExtentInBlocks(region.imageExtent, aspect);
@@ -306,7 +312,7 @@
 	int srcRowPitchBytes = bufferIsSource ? bufferRowPitchBytes : imageRowPitchBytes;
 	int dstRowPitchBytes = bufferIsSource ? imageRowPitchBytes : bufferRowPitchBytes;
 
-	VkExtent3D mipLevelExtent = getMipLevelExtent(region.imageSubresource.mipLevel);
+	VkExtent3D mipLevelExtent = getMipLevelExtent(aspect, region.imageSubresource.mipLevel);
 	bool isSinglePlane = (imageExtent.depth == 1);
 	bool isSingleLine  = (imageExtent.height == 1) && isSinglePlane;
 	bool isEntireLine  = (imageExtent.width == mipLevelExtent.width) &&
@@ -481,7 +487,7 @@
 	return adjustedExtent;
 }
 
-int Image::borderSize(VkImageAspectFlagBits aspect) const
+int Image::borderSize() const
 {
 	// We won't add a border to compressed cube textures, we'll add it when we decompress the texture
 	return (isCube() && !format.isCompressed()) ? 1 : 0;
@@ -491,31 +497,50 @@
 {
 	VkImageAspectFlagBits aspect = static_cast<VkImageAspectFlagBits>(subresource.aspectMask);
 	VkOffset3D adjustedOffset = imageOffsetInBlocks(offset, aspect);
-	int border = borderSize(aspect);
+	int border = borderSize();
 	return adjustedOffset.z * slicePitchBytes(aspect, subresource.mipLevel) +
 	       (adjustedOffset.y + border) * rowPitchBytes(aspect, subresource.mipLevel) +
 	       (adjustedOffset.x + border) * getFormat(aspect).bytesPerBlock();
 }
 
-VkExtent3D Image::getMipLevelExtent(uint32_t mipLevel) const
+VkExtent3D Image::getMipLevelExtent(VkImageAspectFlagBits aspect, uint32_t mipLevel) const
 {
 	VkExtent3D mipLevelExtent;
 	mipLevelExtent.width = extent.width >> mipLevel;
 	mipLevelExtent.height = extent.height >> mipLevel;
 	mipLevelExtent.depth = extent.depth >> mipLevel;
 
-	if(mipLevelExtent.width == 0)
+	if(mipLevelExtent.width == 0)  mipLevelExtent.width = 1;
+	if(mipLevelExtent.height == 0) mipLevelExtent.height = 1;
+	if(mipLevelExtent.depth == 0)  mipLevelExtent.depth = 1;
+
+	switch(aspect)
 	{
-		mipLevelExtent.width = 1;
+	case VK_IMAGE_ASPECT_COLOR_BIT:
+	case VK_IMAGE_ASPECT_DEPTH_BIT:
+	case VK_IMAGE_ASPECT_STENCIL_BIT:
+	case VK_IMAGE_ASPECT_PLANE_0_BIT:  // Vulkan 1.1 Table 31. Plane Format Compatibility Table: plane 0 of all defined formats is full resolution.
+		break;
+	case VK_IMAGE_ASPECT_PLANE_1_BIT:
+	case VK_IMAGE_ASPECT_PLANE_2_BIT:
+		switch(format)
+		{
+		case VK_FORMAT_G8_B8_R8_3PLANE_420_UNORM:
+		case VK_FORMAT_G8_B8R8_2PLANE_420_UNORM:
+			ASSERT(mipLevelExtent.width % 2 == 0 && mipLevelExtent.height % 2 == 0);  // Vulkan 1.1: "Images in this format must be defined with a width and height that is a multiple of two."
+			// Vulkan 1.1 Table 31. Plane Format Compatibility Table:
+			// Half-resolution U and V planes.
+			mipLevelExtent.width /= 2;
+			mipLevelExtent.height /= 2;
+			break;
+		default:
+			UNSUPPORTED("format %d", int(format));
+		}
+		break;
+	default:
+		UNSUPPORTED("aspect %x", int(aspect));
 	}
-	if(mipLevelExtent.height == 0)
-	{
-		mipLevelExtent.height = 1;
-	}
-	if(mipLevelExtent.depth == 0)
-	{
-		mipLevelExtent.depth = 1;
-	}
+
 	return mipLevelExtent;
 }
 
@@ -524,7 +549,8 @@
 	// Depth and Stencil pitch should be computed separately
 	ASSERT((aspect & (VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT)) !=
 	                 (VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT));
-	return getFormat(aspect).pitchB(getMipLevelExtent(mipLevel).width, borderSize(aspect), true);
+
+	return getFormat(aspect).pitchB(getMipLevelExtent(aspect, mipLevel).width, borderSize(), true);
 }
 
 int Image::slicePitchBytes(VkImageAspectFlagBits aspect, uint32_t mipLevel) const
@@ -532,22 +558,16 @@
 	// Depth and Stencil slice should be computed separately
 	ASSERT((aspect & (VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT)) !=
 	                 (VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT));
-	VkExtent3D mipLevelExtent = getMipLevelExtent(mipLevel);
+
+	VkExtent3D mipLevelExtent = getMipLevelExtent(aspect, mipLevel);
 	Format usedFormat = getFormat(aspect);
 	if(usedFormat.isCompressed())
 	{
 		sw::align(mipLevelExtent.width, usedFormat.blockWidth());
 		sw::align(mipLevelExtent.height, usedFormat.blockHeight());
 	}
-	return getFormat(aspect).sliceB(mipLevelExtent.width, mipLevelExtent.height, borderSize(aspect), true);
-}
 
-int Image::bytesPerTexel(VkImageAspectFlagBits aspect) const
-{
-	// Depth and Stencil bytes should be computed separately
-	ASSERT((aspect & (VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT)) !=
-	                 (VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT));
-	return getFormat(aspect).bytes();
+	return usedFormat.sliceB(mipLevelExtent.width, mipLevelExtent.height, borderSize(), true);
 }
 
 Format Image::getFormat(VkImageAspectFlagBits aspect) const
@@ -578,6 +598,27 @@
 			return memoryOffset + getStorageSize(VK_IMAGE_ASPECT_DEPTH_BIT);
 		}
 		break;
+
+	case VK_FORMAT_G8_B8_R8_3PLANE_420_UNORM:
+		if(aspect == VK_IMAGE_ASPECT_PLANE_2_BIT)
+		{
+			return memoryOffset + getStorageSize(VK_IMAGE_ASPECT_PLANE_1_BIT)
+			                    + getStorageSize(VK_IMAGE_ASPECT_PLANE_0_BIT);
+		}
+		// Fall through to 2PLANE case:
+	case VK_FORMAT_G8_B8R8_2PLANE_420_UNORM:
+		if(aspect == VK_IMAGE_ASPECT_PLANE_1_BIT)
+		{
+			return memoryOffset + getStorageSize(VK_IMAGE_ASPECT_PLANE_0_BIT);
+		}
+		else
+		{
+			ASSERT(aspect == VK_IMAGE_ASPECT_PLANE_0_BIT);
+
+			return memoryOffset;
+		}
+		break;
+
 	default:
 		break;
 	}
@@ -602,7 +643,7 @@
 
 VkDeviceSize Image::getMipLevelSize(VkImageAspectFlagBits aspect, uint32_t mipLevel) const
 {
-	return getMipLevelExtent(mipLevel).depth * slicePitchBytes(aspect, mipLevel);
+	return getMipLevelExtent(aspect, mipLevel).depth * slicePitchBytes(aspect, mipLevel);
 }
 
 VkDeviceSize Image::getMultiSampledLevelSize(VkImageAspectFlagBits aspect, uint32_t mipLevel) const
@@ -646,12 +687,22 @@
 
 VkDeviceSize Image::getStorageSize(VkImageAspectFlags aspectMask) const
 {
-	if (aspectMask == (VK_IMAGE_ASPECT_DEPTH_BIT|VK_IMAGE_ASPECT_STENCIL_BIT))
+	if((aspectMask & ~(VK_IMAGE_ASPECT_COLOR_BIT | VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT |
+	                   VK_IMAGE_ASPECT_PLANE_0_BIT | VK_IMAGE_ASPECT_PLANE_1_BIT | VK_IMAGE_ASPECT_PLANE_2_BIT)) != 0)
 	{
-		ASSERT(!format.isCompressed());
-		return arrayLayers * (getLayerSize(VK_IMAGE_ASPECT_DEPTH_BIT) + getLayerSize(VK_IMAGE_ASPECT_STENCIL_BIT));
+		UNSUPPORTED("aspectMask %x", int(aspectMask));
 	}
-	return arrayLayers * getLayerSize(static_cast<VkImageAspectFlagBits>(aspectMask));
+
+	VkDeviceSize storageSize = 0;
+
+	if(aspectMask & VK_IMAGE_ASPECT_COLOR_BIT)   storageSize += getLayerSize(VK_IMAGE_ASPECT_COLOR_BIT);
+	if(aspectMask & VK_IMAGE_ASPECT_DEPTH_BIT)   storageSize += getLayerSize(VK_IMAGE_ASPECT_DEPTH_BIT);
+	if(aspectMask & VK_IMAGE_ASPECT_STENCIL_BIT) storageSize += getLayerSize(VK_IMAGE_ASPECT_STENCIL_BIT);
+	if(aspectMask & VK_IMAGE_ASPECT_PLANE_0_BIT) storageSize += getLayerSize(VK_IMAGE_ASPECT_PLANE_0_BIT);
+	if(aspectMask & VK_IMAGE_ASPECT_PLANE_1_BIT) storageSize += getLayerSize(VK_IMAGE_ASPECT_PLANE_1_BIT);
+	if(aspectMask & VK_IMAGE_ASPECT_PLANE_2_BIT) storageSize += getLayerSize(VK_IMAGE_ASPECT_PLANE_2_BIT);
+
+	return arrayLayers * storageSize;
 }
 
 const Image* Image::getSampledImage(const vk::Format& imageViewFormat) const
@@ -854,7 +905,7 @@
 	{
 		for(; subresourceLayers.mipLevel <= lastMipLevel; subresourceLayers.mipLevel++)
 		{
-			VkExtent3D mipLevelExtent = getMipLevelExtent(subresourceLayers.mipLevel);
+			VkExtent3D mipLevelExtent = getMipLevelExtent(static_cast<VkImageAspectFlagBits>(subresourceLayers.aspectMask), subresourceLayers.mipLevel);
 
 			int pitchB = decompressedImage->rowPitchBytes(VK_IMAGE_ASPECT_COLOR_BIT, subresourceLayers.mipLevel);
 
diff --git a/src/Vulkan/VkImage.hpp b/src/Vulkan/VkImage.hpp
index 7356419..9872585 100644
--- a/src/Vulkan/VkImage.hpp
+++ b/src/Vulkan/VkImage.hpp
@@ -54,7 +54,7 @@
 	uint32_t                 getLastLayerIndex(const VkImageSubresourceRange& subresourceRange) const;
 	uint32_t                 getLastMipLevel(const VkImageSubresourceRange& subresourceRange) const;
 	VkSampleCountFlagBits    getSampleCountFlagBits() const { return samples; }
-	VkExtent3D               getMipLevelExtent(uint32_t mipLevel) const;
+	VkExtent3D               getMipLevelExtent(VkImageAspectFlagBits aspect, uint32_t mipLevel) const;
 	int                      rowPitchBytes(VkImageAspectFlagBits aspect, uint32_t mipLevel) const;
 	int                      slicePitchBytes(VkImageAspectFlagBits aspect, uint32_t mipLevel) const;
 	void*                    getTexelPointer(const VkOffset3D& offset, const VkImageSubresourceLayers& subresource) const;
@@ -79,10 +79,9 @@
 	VkExtent3D imageExtentInBlocks(const VkExtent3D& extent, VkImageAspectFlagBits aspect) const;
 	VkOffset3D imageOffsetInBlocks(const VkOffset3D& offset, VkImageAspectFlagBits aspect) const;
 	VkExtent2D bufferExtentInBlocks(const VkExtent2D& extent, const VkBufferImageCopy& region) const;
-	int bytesPerTexel(VkImageAspectFlagBits flags) const;
 	VkFormat getClearFormat() const;
 	void clear(void* pixelData, VkFormat pixelFormat, const vk::Format& viewFormat, const VkImageSubresourceRange& subresourceRange, const VkRect2D& renderArea);
-	int borderSize(VkImageAspectFlagBits aspect) const;
+	int borderSize() const;
 	void decodeETC2(const VkImageSubresourceRange& subresourceRange) const;
 
 	const Device *const      device = nullptr;
diff --git a/src/Vulkan/VkImageView.cpp b/src/Vulkan/VkImageView.cpp
index 1560d0b..c8e1f65 100644
--- a/src/Vulkan/VkImageView.cpp
+++ b/src/Vulkan/VkImageView.cpp
@@ -178,7 +178,8 @@
 		resolveAttachment->subresourceRange.layerCount
 	};
 	region.dstOffset = { 0, 0, 0 };
-	region.extent = image->getMipLevelExtent(subresourceRange.baseMipLevel);
+	region.extent = image->getMipLevelExtent(static_cast<VkImageAspectFlagBits>(subresourceRange.aspectMask),
+	                                         subresourceRange.baseMipLevel);
 
 	image->copyTo(*(resolveAttachment->image), region);
 }
@@ -219,7 +220,8 @@
 
 VkExtent3D ImageView::getMipLevelExtent(uint32_t mipLevel) const
 {
-	return image->getMipLevelExtent(subresourceRange.baseMipLevel + mipLevel);
+	return image->getMipLevelExtent(static_cast<VkImageAspectFlagBits>(subresourceRange.aspectMask),
+	                                subresourceRange.baseMipLevel + mipLevel);
 }
 
 void *ImageView::getOffsetPointer(const VkOffset3D& offset, VkImageAspectFlagBits aspect, uint32_t mipLevel, uint32_t layer, Usage usage) const
@@ -233,6 +235,7 @@
 		subresourceRange.baseArrayLayer + layer,
 		subresourceRange.layerCount
 	};
+
 	return getImage(usage)->getTexelPointer(offset, imageSubresourceLayers);
 }
 
diff --git a/src/Vulkan/VkImageView.hpp b/src/Vulkan/VkImageView.hpp
index f857f4d..73ee4ab 100644
--- a/src/Vulkan/VkImageView.hpp
+++ b/src/Vulkan/VkImageView.hpp
@@ -45,6 +45,7 @@
 
 	VkImageViewType getType() const { return viewType; }
 	Format getFormat(Usage usage = RAW) const;
+	Format getFormat(VkImageAspectFlagBits aspect) const { return image->getFormat(aspect); }
 	int rowPitchBytes(VkImageAspectFlagBits aspect, uint32_t mipLevel, Usage usage = RAW) const;
 	int slicePitchBytes(VkImageAspectFlagBits aspect, uint32_t mipLevel, Usage usage = RAW) const;
 	int layerPitchBytes(VkImageAspectFlagBits aspect, Usage usage = RAW) const;
diff --git a/src/WSI/XlibSurfaceKHR.cpp b/src/WSI/XlibSurfaceKHR.cpp
index 3bf3022..c512245 100644
--- a/src/WSI/XlibSurfaceKHR.cpp
+++ b/src/WSI/XlibSurfaceKHR.cpp
@@ -58,7 +58,7 @@
 	XWindowAttributes attr;
 	libX11->XGetWindowAttributes(pDisplay, window, &attr);
 
-	VkExtent3D extent = vk::Cast(image->image)->getMipLevelExtent(0);
+	VkExtent3D extent = vk::Cast(image->image)->getMipLevelExtent(VK_IMAGE_ASPECT_COLOR_BIT, 0);
 
 	int bytes_per_line = vk::Cast(image->image)->rowPitchBytes(VK_IMAGE_ASPECT_COLOR_BIT, 0);
 	char* buffer = static_cast<char*>(vk::Cast(image->imageMemory)->getOffsetPointer(0));
@@ -89,7 +89,7 @@
 
 		if(xImage->data)
 		{
-			VkExtent3D extent = vk::Cast(image->image)->getMipLevelExtent(0);
+			VkExtent3D extent = vk::Cast(image->image)->getMipLevelExtent(VK_IMAGE_ASPECT_COLOR_BIT, 0);
 			libX11->XPutImage(pDisplay, window, gc, xImage, 0, 0, 0, 0, extent.width, extent.height);
 		}
 	}